One-Step Gradient Delay is Not a Barrier for Large-Scale Asynchronous Pipeline Parallel LLM Pretraining

기술 요약: 1단계 그래디언트 지연은 대규모 비동기 파이프라인 병렬 방식의 LLM 사전 학습에 있어 장벽이 아니다

문제 정의

현대적인 대규모 언어 모델(LLM) 사전 학습은 메모리 제약을 관리하기 위해 파이프라인 병렬성(Pipeline Parallelism, PP)에 점점 더 많이 의존하고 있다. 그러나 동기식 PP 구현은 다른 스테이지로부터 데이터를 기다리는 동안 GPU가 유휴 상태로 머무는 '파이프라인 버블(pipeline bubbles)' 현상을 겪으며, 이는 계산 효율성을 저하시킨다. 비동기 파이프라인 병렬성(Async PP)은 동기화 배리어를 제거함으로써 이러한 버블을 없애고 처리량(throughput)을 극대화한다.

Async PP 채택의 주요 장벽은 그 결과로 발생하는 그래디언트 신선도 저하(stale gradient, 즉 오래된 파라미터를 사용하여 그래디언트를 계산하는 현상)가 최적화를 근본적으로 불안정하게 만든다는 믿음이었다. Ajanthan 등(2025)의 선행 연구와 같은 기존 작업들은 원래의 PipeDream 스케줄을 사용했는데, 이는 파이프라인 깊이에 따라 증가하는 **가변적 그래디언트 지연(variable gradient delays)**을 유발하여 심각한 수렴 저하를 초도한다. PipeDream-2BW(Narayanan et al., 2021a)는 파이프라인 깊이에 관계없이 일정한 1단계 지연(one-step delay)을 제공하지만, 단 1단계의 지연조차도 특히 AdamW와 같은 지배적인 옵티마이저에 대해 본질적으로 불안정하다는 가정 때문에 채택이 제한되어 왔다.

방법론

저자들은 고정된 1단계 지연 하에서의 옵티마이저 동작에 대한 포괄적인 경험적 및 이론적 분석을 수행함으로써, 1단계 지연이 근본적인 장벽이라는 가설에 도전한다.

1. 옵티마이저 벤치마킹: 본 연구는 135M에서 360M 파라미터 규모의 모델에 대해 다양한 옵티마이저(AdamW, Muon, SOAP, Adan, Lion, MARS 포함)를 평가한다. 이들은 일정한 1단계 지연($w_{t+1} = w_t - u_{t-1}(g_{t-1})$)을 보장하는 PipeDream-2BW 스케줄을 사용하여 동기식 훈련과 비동기식 훈련을 비교한다.
2. 하이퍼파라미터 민감도: 저자들은 하이퍼파라미터, 특히 모멘텀 감쇠 계수($\beta_1$ 또는 $\mu$)가 신선도 저하(staleness)에 대한 강건성에 어떻게 영향을 미치는지 조사한다. 또한 배치 크기, 학습률, 가중치 감쇠(weight decay)의 영향도 분석한다.
3. 완화 전략: 동기식 훈련과 비동기식 훈련 사이의 성능 격차를 메우기 위해 여러 전략을 테스트한다:
   • 동기식 시작(Synchronous Start): 비동기 업데이트로 전환하기 전에 동기식으로 훈련을 시작함.
   • 동기식 쿨다운(Synchronous Cooldown): 훈련 종료 시점에 다시 동기식 업데이트로 전환함.
   • 테일러 기반 지연 보상(Taylor-based Delay Compensation): 테일러 스타일의 보정 항을 사용하여 그래디언트를 수정함.
   • 에러 피드백(Error Feedback, EF): Seide 등(2014)에서 영감을 얻은 옵티마이저 불가지론적(optimizer-agnostic) 보정 방식으로, 업데이트 레벨에서 적용된다. 저자들은 구체적인 업데이트 규칙인 $x_{t+1} = x_t - 2u_{t-1}(g_{t-1}) + u_{t-2}(g_{t-2})$를 도출한다.
4. 이론적 분석: 본 논문은 고정된 그래디언트 지연 하에서 선형 최소화 오라클(Linear Minimization Oracle, LMO) 알고리즘(특히 Muon)에 대한 최초의 수렴 분석을 제공하며, 제안된 에러 피드백 보정이 있는 경우와 없는 경우를 모두 다룬다.
5. 대규모 검증: 200B 토큰으로 훈련된 10B 파라미터 Mixture-of-Experts(MoE) 모델을 통해 연구 결과를 검증하며, 에러 피드백을 적용한 경우와 적용하지 않은 경우의 Async PP를 동기식 베이스라인과 비교한다.

주요 기여

• 옵티마이저 강건성 발견: 저자들은 1단계 지연 하에서의 성능 저하가 지연 자체의 본질적인 특성이 아니라 옵티마이저에 크게 의존한다는 것을 입증한다. AdamW는 심각한 품질 저하를 겪는 반면, Muon과 같은 현대적인 옵티마이저는 강력한 강건성을 보여 동기식 훈련과의 성능 격차를 작게 유지한다.
• 모멘텀의 역할: 모멘텀 감쇠와 신선도 저하 강건성 사이의 결정적인 관계를 식별했다. 높은 모멘텀 계수는 일관되게 지연에 의한 노이즈를 억제함으로써 1단계 지연으로 인한 손실 페널티를 줄여준다.
• 에러 피드백 보정: 본 논문은 동기-비동기 격차를 지속적으로 좁히는 업데이트 레벨의 에러 피드백 메커니즘을 도입한다. Muon과 같은 강건한 옵티마이저의 경우 성능 저하의 5070%를 회복하며, AdamW의 경우 8590%를 회복한다.
• 스케줄 비교: 연구는 PipeDream-2BW(일정한 지연)가 원래의 PipeDream 스케줄(가변적 지연)보다 우수함을 경험적으로 검증한다. 강건한 옵티마이저들이 얕은 깊이에서는 원래 스케줄의 가변적 지연을 부분적으로 완화할 수는 있지만, 더 깊은 파이프라인 깊이에서는 저하를 막지 못하며, 이는 확장 가능한 Async PP를 위해 일정한 지연 스케줄이 필수적임을 뒷받침한다.
• 이론적 보장: 본 논문은 그래디언트 지연 하에서 에러 피드백 유무에 따른 Muon의 수렴 보장을 확립하여, 기존 이론을 LMO 기반 방법론으로 확장한다.

결과

• 소규모 모델 (135M/360M):
  • AdamW는 1단계 지연 하에서 심각하게 저하된다 (검증 손실 > 0.2 증가).
  • Muon은 아무런 완화 조치 없이도 최소한의 격차(예: 360M 모델에서 +0.012)를 유지한다.
  • 에러 피드백은 Muon의 격차를 더욱 줄이고 AdamW를 효과적으로 구제하여, 성능을 동기식 베이스라인에 가깝게 끌어올린다.
• 대규모 모델 (10B MoE):
  • Muon과 에러 피드백을 사용하여 200B 토큰 동안 10B MoE 모델을 훈련한 결과, 동기식 베이스라인과 동일한 1.906의 최종 검증 손실을 달성했다.
  • 표준 Async PP(에러 피드백 미적용)는 1.911의 손실을 기록하여 격차가 미미했다.
  • 다운스트림 벤치마킹은 동일한 검증 손실이 동등한 태스크 성능으로 이어진다는 것을 확인해주었다.
• 확장성: 50B에서 200B 토큰으로 훈련이 진행됨에 따라 동기-비동기 격차가 커지지 않았으며, 이는 해당 방법이 효과적으로 확장됨을 시사한다.

의의 및 주장

본 논문은 품질 저하 없이 10B 파라미터 규모에서 비동기 파이프라인 병렬성을 성공적으로 입증한 첫 사례라고 주장한다. 일정한 지연 스케줄(PipeDream-2BW), 강건한 옵티마이저(Muon), 그리고 에러 피드백 보정을 결 조합함으로써, 저자들은 동기식 훈련과의 성능 격차를 완전히 해소했다.

저자들은 1단계 그래디언트 지연이 대규모 LLM 사전 학습의 근본적인 장벽이 아니다라고 결론짓는다. 대신, 인지된 불안정성은 주로 부적절한 옵티마이저(AdamW 등)와 가변 지연 스케줄을 사용한 결과였다. 이 연구는 적절한 옵티마이저와 완화 전략이 갖춰진다면, Async PP가 파이프라인 버블을 제거하고 대규모 분산 훈련에서 하드웨어 활용도를 극대화할 수 있는 실질적인 잠재력을 가지고 있음을 강조한다.

참고: 본 논문은 높은 모멘텀이 왜 강건성을 향상시키는지에 대한 완전한 메커니즘적 설명의 부족과, 배치 크기 민감도 분석이 소규모 모델로 제한되었다는 점을 한계로 인정한다. 또한 WPipe 스케줄이 유망함을 보여주었으나, 부록의 제한된 실험에서만 연구되었음을 명시한다.