기술 요약: PACI를 통한 버블(Bubble)의 타파

문제 정의

파이프라인 병렬성(Pipeline parallelism)은 대규모 신경망, 특히 트랜스포머 모델을 학습할 때 모델을 여러 장치에 분할하고 연산을 중첩시키기 위한 핵심적인 기술입니다. 그러나 기존의 스케줄링 전략은 하드웨어 활용도, 메모리 효율성, 최적화 일관성 사이에서 지속적인 트레이드오프(trade-off) 문제에 직면해 있습니다:

• 동기식 스케줄 (예: 1F1B-Flush): 이는 순전파/역전파 가중치 일관성(순전파 시 사용된 파라미터를 역전파에서도 동일하게 사용하는 것을 보장)을 유지하지만, "파이프라인 버블"(유휴 시간)로 인해 처리량(throughput)이 저하됩니다.
• 비동기식 스케줄: 이는 버블을 제거하여 완전한 하드웨어 활용도를 달달성할 수 있지만, **순전파/역전파 가중치 버전 불일치(forward/backward weight-version inconsistency)**를 초래합니다. 이러한 스케줄에서는 특정 마이크로 배치(micro-batch)의 역전파가 순전파 이후 파라미터가 상당히 진화한 버전의 파라미터를 사용하여 실행될 수 있습니다.
• 기존 완화 방법: 현재의 비동기 방식들은 가중치 스태싱(weight stashing)(여러 개의 파라미터 복사본 저장), 예측(prediction)(미래 가중치 근사), 또는 **글로벌 동기화(global synchronization)**를 통해 이러한 불일치를 해결하려고 시도합니다. 이러한 접근 방식들은 상당한 메모리 오버헤드, 추가 연산, 시스템 복잡성 또는 비표준적인 최적화 의미론(semantics)을 수반합니다.

저자들은 순전파/역전파 불일치를 완전히 제거하지 않고도, 이를 허용할 수 있을 만큼 작게 만들 수 있는지, 즉 기존의 완화 전략 비용을 피하면서 비동기 실행의 효율성을 유지할 수 있는지를 질문합니다.

방법론: PACI

본 논문은 가중치 스태싱, 예측 또는 글로벌 동기화 없이 버전 드리프트(version drift)를 제한하는 비동기 1F1B 파이프라인 스케줄인 **PACI (Pipeline Asynchronous training with Controlled Inconsistency)**를 소개합니다.

핵심 메커니즘: 버전 제어로서의 그래디언트 누적

PACI는 불일치가 비동기성 그 자체에 의해서가 아니라, 마이크로 배치의 순전파와 역전파 사이에 발생하는 옵티마이저 업데이트 횟수에 의해 결정된다는 관찰을 활용합니다.

1. 업데이트와 역전파의 분리: 표준 비동기 1F1B에서는 모든 역전파 후에 파라미터가 업데이트됩니다. PACI는 이를 분리합니다. 각 스테이지는 입력이 도착하는 즉시 역전파 연산을 수행하지만, 단 한 번의 옵티마이저 단계를 적용하기 전까지 $a$번의 역전파 윈도우 동안 그래디언트를 로컬하게 누적합니다.
2. 버전 드리프트 제한: 파이프라인 지연에 비해 파라미터 버전의 진화 속도를 늦춤으로써, PACI는 각 마이크로 배치가 통과하는 업데이트 횟수를 제한합니다. 최대 불일치($\Delta_{max}$)는 파이프라인 깊이($N$)와 누적 계수($a$)에 의해 다음과 같이 제한됩니다:
   $$\Delta_{max} \le \left\lfloor \frac{N-1}{a} \right\rfloor$$
3. 로컬 흐름 제어 (Local Flow Control): 비균일한 스테이지 시간을 처리하기 위해, PACI는 로컬 백프레셔(backpressure) 규칙을 채택합니다. 각 스테이지는 미결 상태인 순전파 횟수($q_i$)를 추적하며, $q_i < N + 1 - i$인 경우에만 새로운 순전파를 허용합니다. 이는 상위 스테이지가 글로벌 동기화 배리어(barrier)를 도입하지 않고도 하위 스테이지보다 지나치게 앞서 나가는 것을 방지합니다.
4. 리소스 효율성: PACI는 추가적인 가중치 메모리를 전혀 요구하지 않습니다. 단일 파라미터 복사본을 스테이지당 하나씩 사용하며, 옵티마이저에 이미 존재하는 표준 그래디언트 버퍼를 재사용합니다. 가중치 스태싱, 예측 버퍼 또는 글로벌 동기화가 필요하지 않습니다.

마이크로 배칭과의 차이점

동기식 스케줄에서 마이크로 배치 수($m$)를 늘리는 것은 주로 버블 오버헤드를 상쇄하기 위함입니다. 반면 PACI에서 마이크로 배칭(누적 계수 $a$를 통한)은 주로 버전 드리프를 제어하기 위한 **일관성 조절 노브(consistency knob)**로 사용됩니다. 이를 통해 PACI는 높은 활용도라는 시스템 목표와 낮은 불일치라는 최적화 목표를 분리할 수 있습니다.

주요 기여

1. 스태싱 없는 제한된 불일치: 본 논문은 로컬 그래디언트 누적이 추가 메모리, 예측 또는 동기화 없이도 가중치 스태싱 없이 순전파/역전파 불일치를 제한하는 튜너블한 버전 제어 메커즘 역할을 할 수 있음을 입증합니다.
2. 개선된 시간-대-정확도 (Time-to-Accuracy): PACI는 동기식 1F1B-flush와 비교할 때 최종 퍼플렉서티(perplexity)가 유사하면서도 엔드 투 엔드 실행 시간을 최대 1.69배(배치 크기 128에서) 및 1.41배(배치 크기 256에서) 단축하는 안정적인 사전 학습을 달성합니다.
3. 이론과 일치하는 처리량 분석: 저자들은 동기식 1F1B-flush와 PACI 사이의 처리량 격차가 전적으로 파이프라인 버블 효율성에 의해 설명됨을 보여줍니다. PACI는 파이프라인의 이론적 최대 활용 처리량에 도달하면서도 동기식 플러싱과 동일한 피크 메모리 풋프린트를 유지합니다.

실험 결과

저자들은 8개의 NVIDIA RTX PRO 6000 GPU를 사용한 8-스테이지 파이프라인 환경에서 GPT-2 Medium을 OpenWebText 데이터셋으로 처음부터 학습시켜 PACI를 평가했습니다.

• 안정성 (Q1): 제한된 불일치 $\Delta_{max} \le 2$ (누적 계수 $a=4$ 또는 $a=8$로 달성)에 대해, PACI는 동기식 1F1B-flush와 매우 유사한 손실 역학(loss dynamics)을 보입니다. 최종 손실이 유사할 뿐만 아니라, 특히 실행 간 변동성(loss의 RMS 표준 편차)이 더 낮습니다.
• 시간-대-정확도 (Q2): PACI는 목표 퍼플렉서티 수준에 훨씬 빠르게 도달합니다.
  • 배치 크기 128에서, PACI는 가장 잘 튜닝된 1F1B-flush 베이스라인보다 1.69배 빠릅니다.
  • 매칭된 마이크로 배치 구성과 비교했을 때, PACI는 최대 2.04배의 속도 향상을 달s성합니다.
• 처리량 및 메모리 (Q3):
  • PACI는 완전하게 활용된 처리량을 달성하는 반면, 1F1B-flush의 처리량은 버블 효율성에 따라 결정됩니다(마이크로 배치가 증가함에 따라 개선되지만 완전한 활용에는 도달하지 못함).
  • PACI는 1F1B-flush와 동일한 피크 메모리 풋프린트를 가지며, 다른 비동기 방식(예: PipeDream 또는 PipeMare)은 가중치 스태싱이나 보조 상태를 위해 훨씬 더 많은 메모리를 요구합니다.
  • 최대 28.3B 파라미터까지의 대규모 외삽(extrapolation) 결과는 PACI가 Zero-Bubble (ZB-2p) 스케줄의 처리량에 도달하면서도 1F1B-flush의 메모리 효율성을 유지할 수 있음을 시사합니다.

의의 및 주장

본 논문은 순전파/역전파 불일치를 반드시 제거할 필요는 없다고 주장합니다. 대신, 그래디언트 누적을 통해 명시적으로 제한한다면, 이를 상당한 효율성 이득과 교환할 수 있습니다.

• 운영 지점 (Operating Point): PACI는 파이프라인 병렬성의 새로운 운영 지점을 확립합니다: 제로 파이프라인 버블, 제로 추가 가중치 메모리, 그리고 낮은 제한된 불일치.
• 패러다임 전환: 이 연구는 비동기 학습을 위해 가중치 스태싱이나 예측이 필수적이라는 통념에 도전합니다. 파이프라인 지연에 대한 파라미터 진화의 속도를 제어하는 것만으로도 최적화 안정성을 유지하기에 충분하다는 것을 시사합니다.
• 실질적 영향: 버블 오버헤드를 메모리 페널티 없이 제거함으로써, PACI는 특히 동기식 스케줄이 유휴 시간으로 인해 제한을 받는 상황에서 대규모 언어 모델의 더 빠른 학습 시간-대-정확도를 달성할 수 있는 실질적인 경로를 제공합니다.

저자들은 현재의 평가가 GPT-2 Medium, 8-스테이지 파이프라인, 그리고 특정 학습 설정에 국한되어 있다는 점을 언급하며 한계점에 대해 겸허한 태도를 유지하고 있습니다. 또한 액티베이션 체크포인팅(activation checkpointing) 통합 및 무효한 그래디언트(NaNs) 처리를 위한 롤백(rollback) 기능 구현은 향후 과제로 남겨두었습니다.