기술 요약: OctoPipe

문제 정의

파이프라인 병렬성(Pipeline Parallelism, PP)은 분산 GPU 클러스터에서 대규모 언어 모델(LLM)을 학습시키기 위한 표준적인 접근 방식입니다. 그러나 PP의 효율성은 장치 간 의존성으로 인해 발생하는 유휴 시간인 "파이프라인 버블(pipeline bubbles)"에 의해 빈번하게 저해됩니다. 기존 연구들이 동질적 모델(동일한 레이어로 구성된 모델)의 버블 문제를 다루어 왔으나, 이질적 모델(Heterogeneous models, 예: Jamba, Nemotron)의 등장은 이 문제를 더욱 악화시켰습니다. 이러한 모델들은 서로 다른 계산 및 메모리 특성을 가진 다양한 레이어 유형(예: Mamba, Self-Attention, MoE, FFN)을 포함합니다. 이러한 이질성은 심각한 스테이지 불균형을 초래하며, 이로 인해 느린 스테이지가 스트래글러(straggler) 역할을 하여 빠른 스테이지를 유휴 상태로 만듭니다.

기존의 접근 방식들은 일반적으로 파이프라인 스케줄의 단일 단계만을 최적화합니다:

1. 파티셔닝(Partitioning): 스테이지당 레이어 수를 조정함.
2. 배치(Placement): 스테이지를 장치에 매핑함 (예: 가상 스테이지 사용).
3. 스케줄링(Scheduling): 마이크로 배치 실행 순서를 재정렬함 (예: 1F1B, ZeroBubble).

본 논문은 이러한 단계들이 서로 결합되어 있기 때문에 단일 단계만을 최적화하는 것은 불충분하며, 한 영역에서의 개선이 다른 영역의 최적화되지 않은 설정에 의해 상쇄될 수 있다고 주장합니다. 또한, 이질적 모델을 위해 세 가지 단계를 공동 최적화하는 과정에는 다음과 같은 세 가지 구체적인 과제가 따릅니다:

• 복잡한 성능 모델링: 복잡한 교차 단계 상호작용과 다양한 레이어 특성으로 인해 엔드 투 엔드(end-to-end) 지연 시간과 메모리 사용량을 정확하게 추정하기 어렵습니다.
• 조합 탐색 공간: 레이어, 스테이지, 마이크로 배치의 수에 따라 파티셔닝, 배치, 스케줄링의 결합 설계 공간이 기하급수적으로 증가하여 전수 조사가 불가능해집니다.
• 불규칙한 실행 순서: 공동 최적화는 기존 파이프라인 실행기가 가정하는 고정된 실행 패턴을 위반하는 불규칙한 계산 및 통신 시퀀스를 생성하며, 이는 데드락(deadlock)을 유발하거나 버블을 다시 발생시킬 수 있습니다.

방법론

저자들은 이질적 모델을 위해 파티셔닝, 배치, 스케줄링을 공동으로 최적화하도록 설계된 파이프라인 병렬성 시스템인 OctoPipe를 제안합니다. 이 시스템은 세 가지 핵심 구성 요소로 이루어져 있습니다.

1. 그래프 기반 파이프라인 시뮬레이터

복잡한 성능 모델링 문제를 해결하기 위해, OctoPipe는 파이프라인 실행을 **유향 비순환 그래프(Directed Acyclic Graph, DAG)**로 모델링하는 시뮬레이터를 채택합니다.

• 모델링: 시뮬레이터는 이질적 레이어에 대한 오프라인 계산, 통신 및 메모리 비용을 프로파일링합니다. 후보 스케줄이 주어지면, 시뮬레이터는 연산(순전파/역전파 계산, 송신/수신 통신, 메모리 할당)을 노드로 하고 의존성을 에지로 하는 DAG를 구축합니다.
• 시뮬레이션: DAG를 순회함으로써 시뮬레이터는 계산, 통신 및 메모리 사용량에 대한 미세한 타임라인을 생성합니다. 이를 통해 파티셔닝 변경이 통신 패턴을 어떻게 바꾸는지, 또는 스케줄링이 메모리 피크에 어떤 영향을 미치는지와 같은 교차 단계 상호작용을 정확하게 포착합니다.
• 정확도: 시뮬레이터는 처리량(throughput)에 대해 5.06%, 피크 메모리 사용량에 대해 5.29%의 평균 예측 오차를 달립니다.

2. 반복적 버블 인지 파이프라인 튜너 (Iterative Bubble-Aware Pipeline Tuner)

조합 탐색 공간을 효율적으로 탐색하기 위해, OctoPipe는 경험적 관찰을 바탕으로 버블 감소를 우선시하는 반복적 튜너를 사용합니다. 튜너는 현재의 병목 현상을 줄일 가능성이 가장 높은 단계를 조정하는 루프를 통해 작동합니다:

1. 파티셔닝 (우선순위 1): 부하가 가장 많은 장치에서 가장 적은 장치로 레이어를 이동시켜 워크로드를 균형 있게 맞춥니다. 이는 주요 버블 원인인 스테이지 불균형을 해결합니다.
2. 배치 (우선순위 2): 워크로드가 균형을 이루면, 스테이지 분산(stages를 더 세밀한 단위로 분할) 및 레이어 교환(계산 시간을 줄이는 대신 버블 시간을 줄이기 위해 장치 간 레이어를 교체)과 같은 전략을 사용하여 경계 버블(파이프라인 시작 및 종료 시의 유휴 시간)을 줄입니다.
3. 스케줄링 (우선순위 3): 마지막으로, 튜너는 마이크로 배치 연산을 재정렬하여 잔여 버블을 제거합니다. 전략으로는 간극을 채우기 위해 순전파(F) 및 역전파(B) 패스를 앞당기는 것, 의존성이 적은 가중치(W) 그래디언트 계산을 지연시키는 것, 그리고 통신과 중첩될 수 있도록 독립적인 계산을 삽입하는 것 등이 포함됩니다.

• 메모리 안전성: 튜너는 시뮬레이터를 통해 피크 메모리 사용량을 지속적으로 모니터링합니다. 만약 스케줄이 OOM(Out-Of-Memory) 오류 위험이 있다면, 메모리 피크를 낮추기 위해 선택적으로 W 계산을 앞당깁니다.

3. 통합 파이프라인 실행기 (Unified Pipeline Executor)

튜너가 생성한 불규칙한 실행 순서를 지원하기 위해, OctoPipe는 실행 로직을 특정 단계 정책으로부터 분리한 통합 실행기를 구현합니다.

• DAG 기반 실행: 실행기는 고정된 패턴(예: 1F1B)에 의존하는 대신, 파이프라인 DAG로부터 직접 실행 의존성을 도출합니다.
• 데드락 없는 재정렬: 실행기는 통신 데드락을 방-지하기 위해 재정렬 정책을 사용합니다. 순환 의존성을 유발할 수 있는 불일치하는 송수신 쌍을 식별하고, DAG에서 도출된 실행 가능한 삽입 윈도우 내에서 이를 재정렬합니다.
• 중첩 극대화: 실행기는 통신과 계산 사이의 중첩을 최대화하고 "통신 스팬(communication span)"(송신과 수신 사이의 시간 간격)을 최소화하는 방향으로 재정렬 옵션을 탐욕적(greedy)으로 선택합니다.

주요 기여

1. 시스템 설계: OctoPipe는 이질적 모델을 위해 파티셔닝, 배치, 스케줄링을 공동으로 최적화하는 최초의 시스템으로, 단일 단계 최적화의 한계를 해결합니다.
2. 시뮬레이션 프레임워크: 이질적 레이어의 복잡한 상호작용을 정확하게 모델링하여, 폐쇄형 해석 모델(closed-form analytical models) 없이도 정밀한 성능 추정을 가능하게 하는 그래프 기반 시뮬레이터를 제공합니다.
3. 최적화 알고리즘: 공동 최적화 문제를 일련의 단계별 감소 과정으로 분해하여 탐색 공간을 효과적으로 축소하는 반복적 버블 인지 튜닝 전략을 제시합니다.
4. 실행 엔진: 불규칙한 스케줄을 지원하고, 데드락 없는 동작을 보장하며, 계산-통신 중첩을 극대화하는 통합 DAG 기반 실행기를 제공합니다.

결과

저자들은 128-GPU 클러스터(NVIDIA H800)에서 다양한 모델 아키텍처(LLaMA, Jamba, Nemotron-H) 및 규모를 사용하여 OctoPipe를 평가했습니다.

• 처리량 향상: OctoPipe는 최첨단 파이프라인 병렬성 시스템(1F1B, Interleaved, Tessel, Mist) 대비 **1.15× ~ 1.44×**의 훈련 처리량(Tokens/GPU/Second) 향상을 달성했습니다.
• 이질성의 영향: 이득은 이질적 모델에서 가장 두드러졌습니다. 예를 들어, Nemotron-H의 경우 OctoPipe는 1F1B 대비 최대 **1.44×**의 속도 향상을 보였으나, 동질적 LLaMA 모델에서의 이득은 1.14×–1.17×로 더 완만했습니다. 이는 공동 최적화가 이질적 워크로드에 가장 결정적임을 확인해 줍니다.
• 버블 감소: 프로파일링 트레이스는 유의미한 유휴 시간(버블) 감소를 보여주었습니다. 예를 들어, Nemotron-H (Small)의 경우 버블 비율이 크게 감소했으며, 통신 중첩률은 1F1B의 50%에서 OctoPipe의 **91.3%**로 증가했습니다.
• 확장성: OctoPipe는 8192 GPU까지의 시뮬레이션을 포함하여, 강력한 확장성(strong scaling, GPU 증가) 및 약한 확장성(weak scaling, 배치 크기 증가) 시나리오 모두에서 성능 우위를 유지했습니다.
• 탐색 효율성: 반복적 튜너는 파티셔닝을 통해 워크로드 불균형을 먼저 해결하는 방식을 우선시함으로써, ILP 기반 접근 방식(제한된 시간 내에 실행 가능한 솔루션을 찾지 못하는 경우가 많음)보다 훨씬 빠르게 고품질 스케줄을 찾아냈으며 랜덤 탐색보다 우수한 성능을 보였습니다.

의의

본 논문은 모델 아키텍처가 점점 더 이질적으로 변함에 따라, 레이어 비용이 균일하다고 가정하는 전통적인 파이프라인 병렬성 전략은 더 이상 충분하지 않다고 주장합니다. OctoPipe의 의의는 모든 파이프라인 단계를 공동 최적화하는 것이 이질적 모델의 효율성을 극대화하는 데 필수적임을 입증했다는 데 있습니다. 복잡하고 불규칙한 파이프라인 스케줄을 정확하게 모델링하고, 효율적으로 탐색하며, 신뢰성 있게 실행할 수 있는 시스템을 제공함으로써, OctoPipe는 차세대 LLM 스케일링의 핵심 병목 현상을 해결합니다. 저자들은 이 방식이 이질적 모델에 내재된 "스트래글러" 문제를 효과적으로 완화하며, 기본 모델 아키텍처의 변경 없이도 고성능 분산 학습을 위한 견고한 솔루션을 제공한다고 주장합니다.