Revisiting Parameter Server in LLM Post-Training
Xinyi Wan, Penghui Qi, Guangxing Huang, Chaoyi Ruan, Min Lin, Jialin Li
TL;DR
The paper addresses workload imbalance in LLM post-training that undermines traditional DP training via collectives. It introduces On-Demand Communication (ODC), which fuses the robustness of a parameter-server with the memory efficiency of Fully Sharded Data Parallel by substituting point-to-point transfers for all-gather and reduce-scatter, thereby shifting synchronization from layer-level to minibatch-level. The authors implement ODC using CUDA IPC, NVSHMEM, and Triton-based RDMA kernels, integrate it with FSDP, and demonstrate up to 36% end-to-end speedups on long-context supervised fine-tuning and notable gains in RL tasks across 1.5B–32B models, while making the code open-source. The approach simplifies load balancing, improves device utilization, and offers a practical path to robustly scale LLM post-training in heterogeneous workloads. This work thus provides a foundational, scalable alternative to balanced-workload assumptions in modern distributed training of large language models.
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.