LLM 사전 훈련 과정에서의 파라미터 서버 재검토

Revisiting Parameter Server in LLM Post-Training

최근 데이터 병렬(DP) 훈련은 단순성과 균형 잡힌 작업량 환경에서의 효율성 때문에 파라미터 서버(PS) 방식보다 집단 통신 방식을 선호합니다. 그러나, 시퀀스 길이에 큰 차이가 존재하는 대규모 언어 모델(LLM) 사전 훈련 과정에서는 균형 잡힌 작업량이라는 전제가 더 이상 유효하지 않습니다. 불균형한 작업량 환경에서 집단 통신은 동기화 병목 현상을 야기하여 작업량이 적은 장치의 활용도를 저하시킵니다. 이러한 훈련 방식의 변화는 불균형에 대한 PS 방식의 견고성을 재검토할 필요성을 제기합니다. 본 연구에서는 집단 all-gather 및 reduce-scatter 통신을 직접적인 포인트-투-포인트 통신으로 대체하여 PS 방식을 Fully Sharded Data Parallel (FSDP) 방식으로 적용하는 새로운 방식인 **On-Demand Communication (ODC)**를 제안합니다. ODC는 FSDP와 비교하여 레이어별 동기화 빈도를 줄여 미니 배치별로 동기화를 수행하고, 각 장치의 작업량을 분리하여 더 빠른 장치의 속도를 늦추는 현상을 방지합니다. 또한, 미니 배치 수준에서 더 간단하고 효과적인 로드 밸런싱을 가능하게 합니다. 다양한 LLM 사전 훈련 작업에서 ODC는 장치 활용도와 훈련 처리량을 지속적으로 향상시켜 표준 FSDP 방식 대비 최대 36%의 속도 향상을 달성했습니다. 이러한 결과는 ODC가 LLM 사전 훈련 과정에서 널리 나타나는 불균형한 작업량에 더 적합하다는 것을 보여줍니다. ODC의 구현 및 FSDP와의 통합은 https://github.com/sail-sg/odc 에서 오픈 소스로 제공됩니다.

Modern data parallel (DP) training favors collective communication over parameter servers (PS) for its simplicity and efficiency under balanced workloads. However, the balanced workload assumption no longer holds in large language model (LLM) post-training due to the high variance in sequence lengths. Under imbalanced workloads, collective communication creates synchronization barriers, leading to under-utilization of devices with smaller workloads. This change in training dynamics calls for a revisit of the PS paradigm for its robustness to such imbalance. We propose \textbf{On-Demand Communication (ODC)}, which adapts PS into Fully Sharded Data Parallel (FSDP) by replacing collective all-gather and reduce-scatter with direct point-to-point communication. Compared to FSDP, ODC reduces the synchronization barrier from once per layer to once per minibatch and decouples the workload on each device so that faster workers are not stalled. It also enables simpler and more effective load balancing at the minibatch level. Across diverse LLM post-training tasks, ODC consistently improves device utilization and training throughput, achieving up to a 36\% speedup over standard FSDP. These results demonstrate that ODC is a superior fit for the prevalent imbalanced workloads in LLM post-training. Our implementation of ODC and integration with FSDP is open-sourced at https://github.com/sail-sg/odc.