大規模言語モデルのポストトレーニングにおけるパラメータサーバーの再検討

Revisiting Parameter Server in LLM Post-Training

January 27, 2026

著者: Xinyi Wan, Penghui Qi, Guangxing Huang, Chaoyi Ruan, Min Lin, Jialin Li

cs.AI

要旨

現代のデータ並列（DP）訓練は、バランスの取れたワークロード下ではシンプルで効率的であることから、パラメータサーバ（PS）よりも集団通信を重視している。しかし、大規模言語モデル（LLM）の事後訓練では、系列長の高い分散により、バランスの取れたワークロードという前提が成り立たなくなる。不均衡なワークロード下では、集団通信は同期の障壁を生み出し、ワークロードが小さいデバイスの利用効率低下を招く。この訓練動態の変化により、不均衡に対する頑健性という観点からPSパラダイムの再評価が求められている。本研究では、オンデマンド通信（ODC）を提案する。ODCは、集団的なall-gatherおよびreduce-scatterを直接的なポイントツーポイント通信に置き換えることで、PSを完全分散データ並列（FSDP）に適応させる。FSDPと比較して、ODCは同期の障壁を層ごと1回からミニバッチごと1回に削減し、各デバイスのワークロードを分離するため、処理の速いワーカーが待機状態にならない。また、ミニバッチレベルでのよりシンプルで効果的な負荷分散を可能にする。多様なLLM事後訓練タスクにおいて、ODCは一貫してデバイス使用率と訓練スループットを向上させ、標準的なFSDPと比べて最大36%の高速化を達成した。これらの結果は、ODCがLLM事後訓練において広く見られる不均衡なワークロードに極めて適合していることを示している。ODCの実装およびFSDPとの統合は、https://github.com/sail-sg/odc で公開されている。

English

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 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.