Revisiting Parameter Server in LLM Post-Training

📄 arXiv: 2601.19362v1 📥 PDF

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

分类: cs.DC, cs.AI

发布日期: 2026-01-27

备注: Accepted in ICLR'26

🔗 代码/项目: GITHUB

💡 一句话要点

针对LLM后训练中负载不均衡问题，提出On-Demand Communication加速训练。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 后训练 数据并行 参数服务器 负载均衡

📋 核心要点

  1. LLM后训练中序列长度差异导致负载不均衡，集体通信产生同步障碍，降低设备利用率。
  2. 提出On-Demand Communication (ODC)，用点对点通信代替集体通信，减少同步开销。
  3. 实验表明，ODC在多种LLM后训练任务中提升设备利用率和训练吞吐量，最高加速36%。

📝 摘要（中文）

现代数据并行训练由于其简单性和在平衡工作负载下的效率，通常优于参数服务器（PS）模式。然而，由于序列长度的高方差，平衡工作负载的假设在大型语言模型（LLM）后训练中不再成立。在不平衡的工作负载下，集体通信会产生同步障碍，导致工作负载较小的设备利用率不足。这种训练动态的变化要求重新审视PS模式，因为它对这种不平衡具有鲁棒性。我们提出了按需通信（ODC），通过用直接点对点通信替换集体all-gather和reduce-scatter，将PS适配到完全分片数据并行（FSDP）中。与FSDP相比，ODC将同步障碍从每层一次减少到每个小批量一次，并将每个设备上的工作负载解耦，从而不会使更快的worker停滞。它还可以在小批量级别实现更简单、更有效的负载平衡。在不同的LLM后训练任务中，ODC始终提高设备利用率和训练吞吐量，与标准FSDP相比，速度提高了高达36％。这些结果表明，ODC更适合LLM后训练中普遍存在的不平衡工作负载。ODC的实现以及与FSDP的集成已在https://github.com/sail-sg/odc上开源。

🔬 方法详解

问题定义：论文旨在解决LLM后训练中由于序列长度差异导致的数据并行训练负载不均衡问题。传统的数据并行方法，如FSDP，依赖于集体通信操作（all-gather, reduce-scatter）进行梯度同步，在负载均衡的场景下表现良好。然而，在LLM后训练阶段，不同样本的序列长度差异很大，导致每个GPU上的计算量不一致，集体通信会造成同步阻塞，降低整体训练效率。

核心思路：论文的核心思路是借鉴参数服务器（PS）的思想，将集体通信替换为点对点通信，实现按需通信（On-Demand Communication, ODC）。通过解耦每个设备上的工作负载，避免快速worker等待慢速worker，从而提高设备利用率。同时，ODC支持更灵活的负载均衡策略，可以在小批量级别动态调整任务分配。

技术框架：ODC框架基于FSDP，主要修改了梯度同步的方式。具体来说，不再使用all-gather和reduce-scatter进行梯度聚合和分发，而是采用点对点通信。每个worker在完成本地计算后，直接将梯度发送给负责更新对应参数的worker。参数更新完成后，再将更新后的参数发送给需要这些参数的worker。这种方式避免了全局同步，允许worker独立运行。

关键创新：ODC的关键创新在于将参数服务器的思想融入到FSDP框架中，通过点对点通信替代集体通信，解决了LLM后训练中负载不均衡导致的同步瓶颈问题。与传统的PS相比，ODC更加轻量级，易于集成到现有的FSDP框架中。此外，ODC还支持动态负载均衡，可以根据每个worker的计算能力和数据量，动态调整任务分配，进一步提高训练效率。

关键设计：ODC的关键设计包括：1) 使用点对点通信代替集体通信；2) 实现细粒度的负载均衡，允许在小批量级别调整任务分配；3) 与FSDP无缝集成，尽可能复用FSDP的现有功能。论文中没有明确提及具体的参数设置或损失函数修改，重点在于通信机制的优化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ODC在多种LLM后训练任务中均优于标准FSDP。例如，在某些任务上，ODC实现了高达36%的加速。这些结果证明了ODC在处理负载不均衡问题上的有效性，并表明ODC是LLM后训练的更优选择。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的后训练、微调等场景，尤其是在数据分布不均匀的情况下。通过提高设备利用率和训练吞吐量，可以缩短模型训练时间，降低训练成本，加速LLM的迭代和部署。此外，该方法也可以推广到其他深度学习任务中，例如自然语言处理、计算机视觉等。

📄 摘要（原文）

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.