Revisiting Parameter Server in LLM Post-Training

ICLR 2026 Conference Submission17378 Authors

Keywords: Distributed Training, ZeRO Optimizer, FSDP, Parameter Server

TL;DR: We transform FSDP (ZeRO3) into a decentralized parameter server with on-demand communications, which relax the synchronization barrier from layer level to minibatch level, and achieves up to 36% acceleration on various LLM post-training settings.

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 large 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 **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. We will open-source our implementation at camera-ready version.

Primary Area: infrastructure, software libraries, hardware, systems, etc.

Submission Number: 17378