분산 딥러닝을 위한 배치 위치 의미론 병렬 전략 통합 분석 프레임워크

초록

본 논문은 대규모 언어 모델 학습 시 필수적인 여러 병렬 전략을 “배치 위치 의미론”이라는 통일된 모델로 표현한다. 파라미터, 옵티마이저, 그래디언트, 활성값 네 가지 상태를 복제·분할·수집·구현·오프로드 등 다섯 가지 모드로 배치하고, 이 배치만으로 메모리 사용량과 통신량을 정량화한다. 제안된 의미론을 통해 ZeRO‑3가 데이터 병렬 대비 8배 적은 메모리를 사용하면서 통신 비용은 1.5배 증가한다는 기존 결과를 정확히 재현한다. 또한 그래디언트 무결성과 상태 일관성이라는 두 조건이 단일 디바이스와 동일한 학습 결과를 보장하는 필요충분조건임을 증명하고, 여러 전략을 안전하게 결합하는 구성 규칙을 제시한다. ZeRO‑1~3, FSDP, 텐서 병렬, 파이프라인 병렬을 모두 동일한 배치 선택의 특수 사례로 통합한다.

상세 분석

논문은 분산 학습에서 “어디에 무엇을 두느냐”라는 질문을 정형화함으로써 기존에 파편화된 병렬 전략들을 하나의 이론적 틀 안에 끌어들인다. 네 가지 핵심 학습 상태—파라미터, 옵티마이저 상태, 그래디언트, 활성값—를 각각 다섯 가지 배치 모드(복제, 분할, 분할‑수집, 구현, 오프로드) 중 하나에 할당한다. 이때 복제는 모든 디바이스가 동일한 복사본을 보유하는 경우, 분할은 상태를 디바이스 간에 균등하게 나누는 경우, 분할‑수집은 연산 중에 필요한 부분만을 모아 사용 후 다시 분할된 형태로 반환하는 경우, 구현은 연산 그래프 내에서 즉시 계산해 버리는 경우, 오프로드는 메인 메모리 또는 SSD 등 외부 저장소에 보관하는 경우를 의미한다.

배치 의미론만으로 메모리 요구량을 정량화하는 핵심 식은 각 상태의 크기와 선택된 모드의 메모리 계수를 곱한 합으로 표현된다. 예를 들어 파라미터를 완전 복제하면 N·|θ|(N은 디바이스 수) 만큼 메모리가 필요하지만, 완전 분할(ZeRO‑3)에서는 |θ|/N만 필요하다. 그래디언트와 옵티마이저 상태도 동일한 원리로 축소된다. 통신량은 분할된 상태를 동기화하거나 수집할 때 발생하는 데이터 전송량으로, 복제 대비 분할‑수집 모드에서는 각 단계마다 전체 파라미터 크기의 1/N을 전송하게 된다. 논문은 이러한 계산을 통해 ZeRO‑3가 데이터 병렬 대비 메모리는 1/8, 통신량은 1.5배가 된다는 실험 결과와 정확히 일치함을 보인다.

두 가지 정리—그래디언트 무결성(Gradient Integrity)과 상태 일관성(State Consistency)—은 배치 선택이 학습 결과에 미치는 영향을 이론적으로 제한한다. 그래디언트 무결성은 모든 디바이스가 동일한 전체 그래디언트를 재구성할 수 있어야 함을 의미하고, 상태 일관성은 파라미터와 옵티마이저 상태가 각 단계에서 동일한 논리적 값으로 유지되어야 함을 뜻한다. 저자는 이 두 조건이 필요충분함을 수학적으로 증명하고, 이를 위반하면 수렴 속도 저하나 정확도 손실이 발생한다는 실험적 근거를 제시한다.

또한 전략 조합을 위한 “구성 규칙”(Composition Rules)을 제시한다. 예를 들어 텐서 병렬과 데이터 병렬을 동시에 적용하려면 텐서 차원에 대한 파라미터 분할과 데이터 차원에 대한 복제가 충돌하지 않도록 배치 모드를 교차 배치해야 한다. 파이프라인 병렬과 ZeRO를 결합할 경우, 파이프라인 단계마다 로컬 메모리 한계에 맞춰 옵티마이저와 그래디언트를 오프로드하거나 분할‑수집 모드로 전환하는 것이 안전하다. 이러한 규칙은 실제 시스템 설계 시 전략 선택을 체계적으로 탐색할 수 있게 해준다.

마지막으로 논문은 기존 주요 병렬 기법들을 배치 의미론의 특수 케이스로 매핑한다. ZeRO‑1은 옵티마이저만 분할, ZeRO‑2는 옵티마이저와 그래디언트 분할, ZeRO‑3는 네 상태 모두 분할; FSDP는 파라미터와 그래디언트를 완전 분할‑수집, 텐서 병렬은 파라미터와 활성값을 차원별 분할, 파이프라인 병렬은 활성값을 단계별 오프로드와 구현으로 모델링한다. 이 통합은 연구자와 엔지니어가 새로운 혼합 전략을 설계할 때 일관된 분석 도구를 제공한다는 점에서 큰 의미가 있다.