[论文解读] 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），一种去中心化的类似参数服务器的方案，用点对点传输替代 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 extbf{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 后训练中对工作负载不平衡的鲁棒性需求提供动机，其中序列长度差异导致同步障碍。
• 提出 On-Demand Communication（ODC），将 PS 概念改编到 FSDP 中，同时不牺牲内存效率。
• 证明 ODC 能在 SFT 和 RL 任务中提升设备利用率和训练吞吐量。
• 就负载均衡与实现提供实用指南，以缓解跨节点通信开销。

提出的方法

• 用直接点对点的参数获取和梯度推送替代 FSDP 中的每层 all-gather 与 reduce-scatter。
• 通过将同步粒度从层级放宽到小批量粒度来实现设备进度解耦，同时保持同步优化语义。
• 将 FSDP 视作一个去中心化的参数服务器，通过在所有设备上共 colocate 服务器和工作者角色。
• 使用基于 RDMA 的接口实现 ODC（节点内使用 CUDA IPC；节点间使用 NVSHMEM），并采用 Triton-Distributed 内核进行数据传输。
• 将 ODC 与 FSDP 集成，通过替换聚合操作并在小批量边界处收集累积梯度来实现。
• 提出负载平衡策略，将打包决策从微小批量提升到小批量级，以简化并改善平衡。

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

实验结果

研究问题

• RQ1 ODC 能否降低因工作负载不平衡而导致的 synchronization barrier 与 idle time，在 LLM 后训练中发挥作用？
• RQ2 在不牺牲内存效率的前提下，将类似 PS 的解耦通信整合到 FSDP，是否能提高在不平衡工作负载下的吞吐量？
• RQ3 在大上下文 LLM 训练中，哪些面向小批量粒度的负载平衡策略最能与 ODC 互补？
• RQ4 ODC 在从 1.5B 到 32B 参数的监督微调和强化学习任务上表现如何？
• RQ5 节点间 ODC 通信的局限性有哪些，潜在的缓解措施是什么？

主要发现

• ODC 在 SFT 和 RL 任务中始终提升设备利用率和端到端吞吐量。
• 在长上下文的 SFT 场景中，ODC 相较于标准 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.