기술 요약: Piper - 프로그래밍 가능한 분산 학습 시스템

문제 정의

머신러닝 모델이 수백 또는 수천 개의 가속기를 필요로 하는 파운데이션 모델 규모로 확장됨에 따라, 분산 학습 전략을 설계하는 복잡성이 병목 현상이 되고 있습니다. 현재의 접근 방식은 다음과 같은 이분법적 문제에 직면해 있습니다:

1. 배포된 시스템(Deployed Systems): 고성능 시스템(예: DeepSeek-V3에서 사용되는 시스템)은 고수준의 병렬성 전략(예: 파이프라인 병렬성(PP), 전문가 병렬성(EP))을 저수준의 실행 전략(예: GPU 내부 리소스 스케줄링, 스트림 관리)과 수동으로 공동 설계하는 데 전문 인력의 의존도가 높습니다. 이는 새로운 전략이나 결합된 병렬 차원을 적용하기 어려운 경직된 시스템을 초래합니다.
2. 범용 프레임워크(General-Purpose Frameworks): Megatron-LM, DeepSpeed, TorchTitan과 같은 프레임워크는 유연한 인터페이스를 제공하지만, 병렬 차원을 독립적인 것으로 취급하는 경우가 많습니다. 이들은 연산을 성급하게(eagerly) 디스패치하므로, 결합된 전략(예: EP All-to-all과 DP All-reduce의 중첩) 간의 통신과 연산을 공동으로 스케줄링하는 것이 어렵습니다. 또한, 레이어 훅(layer hooks)과 실행 스케줄 사이의 상호작용으로 인해 PP와 ZeRO-2 또는 ZeRO-3의 복잡한 조합을 완전히 지원하지 못하는 경우가 많습니다.

핵심 과제는 **확장성(extensibility)**입니다. 즉, 높은 성능과 메모리 효율성을 유지하면서도 임의의 분산 학습 전략을 명시하고 구현하는 데 드는 노력을 최소화하는 시스템을 만드는 것입니다.

방법론

Piper는 통합된 중간 표현(Intermediate Representation, IR)을 통해 **전략 명세(strategy specification)**와 **런타임 구현(runtime implementation)**을 분리함으로써 이러한 문제를 해결합니다. 이 시스템은 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 사용자 스케줄링 API (User Scheduling API)

Piper는 사용자가 커스텀 스케줄러를 작성하지 않고도 고수준의 배치(placement)와 저수준의 실행을 제어할 수 있도록 선언적 API를 제공합니다.

• 어노테이션(Annotations): 사용자는 컨텍스트 매니저를 사용하여 PyTorch 모델의 의미 있는 영역(예: 특정 레이어, 전문가 블록)에 태그를 지정합니다. 이 태그는 병렬 차원(예: PP, EP, DP)의 경계를 정의합니다.
• 스케줄링 지시문(Scheduling Directives): 사용자는 시스템의 IR을 변환하기 위해 지시문을 적용합니다. 주요 지시문은 다음과 같습니다:
  • 배치(Placement): Place, Replicate, Shard는 장치 매핑을 정의하고 필요한 통신 노드(P2P, All-Reduce, All-Gather, All-to-All)를 도입합니다.
  • 그래뉼러리티(Granularity): Split은 노드를 복제하여 마이크로배치(microbatches)를 생성함으로써 파이프라인 오버랩을 가능하게 합니다.
  • 순서(Ordering): Order는 서브 그래프 간의 시간적 의존성을 강제하여, 사용자가 복잡한 중첩 패턴(예: 포워드 마이크로배치 $i$와 백워드 마이크로배치 $i-1$의 중첩)을 지정하고 리소스 할당(예: 특정 논리적 스트림을 통신에 할당)을 할 수 있게 합니다.
• 리소스 제어(Resource Control): 사용자는 논리적 GPU 스트림을 특정 지시문에 할당할 수 있으며, 이를 통해 물리적 CUDA 스트림을 직접 관리하지 않고도 통신과 연산의 미세한 오버랩을 제어할 수 있습니다.

2. 통합 중간 표현 (Unified Intermediate Representation, IR)

Piper의 핵심은 **글로벌 학습 DAG (Global Training DAG)**입니다.

• 구조: DAG는 모든 연산과 통신 연산을 나타냅니다. 노드는 청크(Chunks)(조밀한 연산 단위) 또는 Comm 노드(통신 연산)로 구성됩니다.
• 변환: 컴파일러는 모델 코드에서 추출한 단일 장치 DAG에서 시작합니다. 사용자의 지시문은 그래프 재작성(graph rewrites) 역할을 수행하여, 단일 장치 DAG를 명시적인 통신 노드, 마이크로배칭, 장치 배치를 포함하는 분산 DAG로 변환합니다.
• 안전성: 시스템은 변환이 모델과 사용자 지시문에 의해 정의된 데이터 및 시간적 의존성을 준수함을 보장합니다.

3. 전략 불가지론적 분산 런타임 (Strategy-Agnostic Distributed Runtime)

Piper는 중앙 집중식 스케줄러와 Ray로 구현된 워커 기반 실행 모델을 사용합니다.

• 중앙 집중식 스케줄러: 변환된 DAG를 받아 각 파이프라인 병렬성(PP) 랭크를 위한 고유한 서브 DAG로 분해합니다. 동일한 스트림에 할당된 독립적인 작업들 사이의 부분 순서(partial orderings)를 해결하여, 스트림당 전체 순서(total order)를 생성합니다.
• 워커 실행: 워커는 자신의 모델 샤드(shard)를 로드하고 스케줄러의 계획에 따라 작업을 실행합니다.
• 리소스 관리:
  • 스트림(Streams): 작업은 물리적 CUDA 스트림으로 매핑되는 논리적 스트림에 할по됩니다. 의존성이 스트림을 가로지를 때만 동기화(CUDA 이벤트를 통해)가 삽입됩니다.
  • 통신(Communication): 데드락과 버블을 방지하기 위해 Piper는 포인트 투 포인트(point-to-point) 송수신을 위한 별도의 커뮤니케이터와 스트림을 할당하며, 이를 통해 실행 시간에 대한 글로벌 런타임 지식 없이도 일관된 순서를 보장합니다.
  • 메모리(Memory): Piper는 모델 상태(파라미터, 그래디언트, 옵티마이저 상태)와 활성화 함수 값(activations)에 대한 메모리를 관리합니다. ZeRO 스타일의 샤딩을 지원하기 위해 계산 전 상태를 실체화(materialize)하고 계산 후 다시 재샤딩(resharding)함으로써, 마이크로배치 간의 정확한 메모리 사용을 보장합니다.

주요 기여

1. 사용자 스케줄링 인터페이스: 장치 간 및 장치 내부 병렬 전략을 제어할 수 있는 유연한 API를 제공하여, 사용자가 지시문을 통해 고수준 배치와 저수준 실행 순서를 지정할 수 있게 합니다.
2. 통합 IR (글로벌 학습 DAG): 결합된 병렬 전략(예: PP + EP + ZeRO)에 의해 도입된 통신 및 연산 작업을 공동으로 스케줄링할 수 있는 표현 방식입니다.
3. 전략 불가지론적 런타임: IR에 정의된 임의의 전략을 실행할 수 있는 효율적인 분산 런타임을 제공하며, 리소스 경합과 동기화를 자동으로 처리합니다.
4. 엔드 투 엔드 평가: 표준 전략에 대해 범용 프레임워크와 성능이 대등함을 입증하는 동시에, 다른 프레임워크가 지원할 수 없는 새로운 고성능 전략을 가능하게 함을 보여줍니다.

결과

저자들은 AWS EC2 인스턴스와 NVIDIA A100 GPU 환경에서 Piper를 Megatron-LM, DeepSpeed, TorchTitan과 비교 평가했습니다.

• 공통 전략: Piper는 표준 1F1B 및 Interleaved 1F1B 스케줄에 대해 Megatron-LM 및 TorchTitan과 대등한 성능을 달야합니다.
• PP x ZeRO 호환성:
  • Megatron, DeepSpeed, TorchTitan은 PP와 ZeRO-2(파라미터 샤딩) 및 ZeRO-3(가중치 샤딩)의 결합을 지원하는 데 어려움을 겪는 반면, Piper는 모든 조합을 지원합니다.
  • 메모리 효율성: PP x ZeRO-2 및 PP x ZeRO-3 시나리오에서, Piper는 가중치와 그래디언트를 마이크로배치 간에 제대로 재샤딩하지 못해 OOM(Out-Of-Memory) 오류가 발생하는 TorchTitan에 비해 3~8배 더 큰 배치 사이즈를 허용합니다.
• 새로운 전략 (DualPipeV):
  • Piper는 PP와 EP를 결합하고 마이크로배치 오버랩을 적용한 DualPipeV 스케줄을 성공적으로 구현합니다.
  • 처리량 향상: 1B 모델에서 Piper-DualPipeV는 Piper의 베이스라인 1F1B 스케줄 대비 처리량을 13% 개선합니다. 9B 모델에서는 자체 Interleaved 스케줄 대비 10%, Megatron의 Interleaved 스케줄 대비 **6%**의 향상을 보입니다.
  • 비교: TorchTitan의 DualPipeV 구현은 스레드 기반 디스패칭으로 인한 의도치 않은 직렬화 때문에 베이스라인 대비 미미한(3%) 개선만을 보여주었습니다.

의의

본 논문은 Piper의 의의가 전략과 구현의 분리에 있다고 주장합니다. 통합 IR과 유연한 API를 제공함으로써, Piper는 현재 많은 인간 공학적 노력이 필요한 복잡한 결합 병렬 전략(예: DualPipe)을 구현하는 장벽을 낮춥니다. 이는 범용 프레임워크의 유연성과 특화된 수동 튜닝 시스템의 성능 사이의 간극을 메웁니다. 나아가, Piper는 특정 최적화를 하드코딩하지 않고도 임의의 전략을 실행할 수 있어, 미래의 자동 탐색 시스템을 위한 실행 가능한 런타임으로서 훨씬 더 넓은 범위의 분산 학습 구성 탐색을 가능하게 합니다.