[論文レビュー] Revisiting Parameter Server in LLM Post-Training

Xinyi Wan, Penghui Qi+4 more|arXiv (Cornell University)|Jan 27, 2026

Cloud Computing and Resource Management被引用数 0

ひとこと要約

要約: 論文は On-Demand Communication (ODC) を紹介します。これは Fully Sharded Data Parallel (FSDP) の各層ごとの集約を点対点転送に置き換える分散パラメータサーバーのようなスキームで、不均衡な LLM ポストトレーニング負荷（SFT および RL）に対してデバイス利用率とスループットを改善し、最大36%の高速化を実現します。

ABSTRACT

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.

研究の動機と目的

• LLM ポストトレーニングにおける負荷不均衡耐性の必要性を動機づける（シーケンス長のばらつきが同期障害を引き起こす）。
• PS の概念を FSDP に適用してメモリ効率を損なわず On-Demand Communication (ODC) を提案する。
• ODC が SFT および RL タスク全体でデバイス利用率とトレーニングスループットを向上させることを実証する。
• ノード間通信オーバーヘッドを緩和するロードバランシングと実装に関する実践的なガイダンスを提供する。

提案手法

• FSDP における層ごとの全集約とリダースを直接の点対点パラメータ取得と勾配プッシュに置換する。
• 層ごとの同期を minibatch 粒度に緩和しつつ、デバイスの進行を分離して同期化最適化の意味論を保持する。
• FSDP を分散パラメータサーバーとして、サーバーとワーカーロールを全デバイスで共置する。
• ODC を RDMA ベースのインターフェース（同一ノード内は CUDA IPC、ノード間は NVSHMEM）と Triton-Distributed カーネルを用いてデータ転送を実装する。
• ODC を FSDP に統合し、集約的操作を置換し、 minibatch 境界で蓄積勾配を収集する。
• packing の決定をマイクロバッチからミニバッチレベルへと移すロードバランシング戦略を提案し、バランスを単純化・改善する。

Figure 1: Collective communications introduces per-layer synchronization barriers in FSDP.

実験結果

リサーチクエスチョン

• RQ1 ODC は LLM ポストトレーニングにおける負荷不均衡が引き起こす同期障害とアイドル時間を削減できるか。
• RQ2 PS のようなデカップリング通信を FSDP に統合して、負荷不均衡のある workloads でもメモリ効率を保持しつつスループットを向上させるか。
• RQ3 長い文脈の LLM 訓練において minibatch 粒度でどのロードバランシング戦略が ODC に最も適しているか。
• RQ4 ODC は 1.5B 〜 32B パラメータのモデル規模で、監視付きファインチューニングと強化学習タスクでどう性能を示すか。
• RQ5 ノード間 ODC 通信の制限と潜在的な緩和策は何か。

主な発見

• ODC は SFT と RL タスク全体でデバイス利用率とエンドツーエンドのスループットを一貫して改善する。
• ODC は長文脈の SFT シナリオで標準 FSDP に対して最大 36% の高速化を達成する。
• 従来の FSDP では長いシーケンスの監督付きファインチューニング時に負荷不均衡によるアイドル時間が最大 50% に達することがある。
• FSDP をオンデマンドの点対点転送を備えた分散型 PS として再定義することでストラグラーを緩和し、マイクロベンチマークの同期を緩和する。
• LB-Mini および LB-Micro ロードバランシングのバリアントはミニバッチレベルの効果的な平衡を実現し、RL および SFT 設定でしばしばベースラインを上回る。
• ODC は同一ノード内の集約手法と競合する水準を維持するが、ノード間通信オーバーヘッドを示す場合があり、ハイブリッドシャーディングや計算と通信の重畳など設計選択で緩和できる。

Figure 2: On-demand communications relaxes the synchronization barriers to minibatch end.