针对大模型微调中任务负载不均问题，提出按需通信机制提升训练效率。
Revisiting Parameter Server in LLM Post-Training
- 用点对点通信替代集体通信，减少同步开销，实现更灵活的负载分配。
- 在多种大模型微调任务中，设备利用率和吞吐量显著提升，最高提速36%。
- 适合高变长序列的大模型后训练场景，尤其适用于资源不均衡的分布式训练。
现代数据并行训练倾向于使用集体通信而非参数服务器（PS），因其在负载均衡时更简单高效。然而，在大语言模型（LLM）后训练中，序列长度差异大导致负载极不均衡，集体通信引入同步瓶颈，使处理负担轻的设备闲置。这一变化要求重新审视PS范式在不平衡负载下的鲁棒性。本文提出 extbf{按需通信（ODC）}，通过将PS与全分片数据并行（FSDP）结合，以直接点对点通信替代集体的all-gather和reduce-scatter操作。相比FSDP，ODC将同步开销从每层一次降低为每小批量一次，并解耦各设备工作负载，避免快设备被慢设备拖累。该机制还支持更简单有效的批次级负载均衡。在多种LLM后训练任务中，ODC稳定提升设备利用率与训练吞吐量，最高较标准FSDP加速36%。结果表明，ODC是当前大模型后训练中不平衡负载场景的更优选择。代码已开源：https://github.com/sail-sg/odc。
原文摘要 · Abstract (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 \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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。