Revisiting Parameter Server in LLM Post-Training

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.

💡 Research Summary

The paper addresses a critical bottleneck in large‑language‑model (LLM) post‑training: the high variance of input sequence lengths creates severe workload imbalance across GPUs. Modern data‑parallel training for LLMs typically relies on Fully Sharded Data Parallel (FSDP), which shards parameters, gradients, and optimizer states across devices and uses collective all‑gather and reduce‑scatter operations at every layer. These per‑layer collectives assume balanced workloads; when this assumption is violated, faster GPUs must idle while waiting for the slowest device, leading to up to 50 % idle time in long‑sequence fine‑tuning.

The authors revisit the classic parameter‑server (PS) paradigm, which naturally tolerates heterogeneous workloads because workers pull parameters and push gradients asynchronously. Rather than building a separate PS, they integrate PS‑style on‑demand communication directly into FSDP, creating a “decentralized PS” where each GPU simultaneously acts as a server (owning a shard of the model and optimizer state) and as a worker (processing its local data). The key innovation, called On‑Demand Communication (ODC), replaces the all‑gather and reduce‑scatter collectives with point‑to‑point gather and scatter‑accumulate primitives. A device fetches only the parameter shards it needs when it is ready, and pushes its computed gradients directly to the devices that own the corresponding shards. Consequently, synchronization is relaxed from once per layer to once per minibatch (the optimizer step), eliminating the fine‑grained barriers that cause stragglers.

Implementation-wise, ODC avoids MPI or NCCL, which require ordered participation and can deadlock under on‑demand patterns. Instead, it leverages RDMA‑based CUDA IPC for intra‑node transfers and NVSHMEM for inter‑node communication. These mechanisms allow transparent data movement without a dedicated server thread; only gradient accumulation needs a lightweight daemon. The communication kernel is built on Triton‑Distributed, exposing RDMA directly in Python Triton kernels, thus keeping the codebase simple and portable.

A second contribution is a simplified load‑balancing strategy. Traditional FSDP relies on sophisticated micro‑batch packing to mitigate imbalance, but memory constraints limit the number of samples per micro‑batch, especially for long sequences. ODC’s decoupled execution permits each GPU to independently pack its locally assigned samples into micro‑batches based solely on its memory budget. The global assignment is performed at the minibatch level to equalize total computational load across devices, which is a much coarser and easier optimization problem. This shift yields better balance without complex packing algorithms.

The authors evaluate ODC on two representative LLM post‑training tasks: supervised fine‑tuning (SFT) using the LongAlign dataset and reinforcement learning (RL) for LLM reasoning. Experiments span models from 7 B to 70 B parameters and sequence lengths from 512 to 4096 tokens. Across all settings, ODC achieves 20 %–36 % higher throughput than vanilla FSDP, with device utilization improvements of 10 %–25 %. The gains are most pronounced for workloads with extreme sequence‑length variance, where the synchronization overhead of FSDP is largest. Profiling shows that ODC reduces the number of synchronization points, improves overlap of communication and computation, and lowers network bandwidth consumption.

The paper concludes that the PS architecture is not obsolete; when combined with modern sharding techniques, it provides the robustness needed for imbalanced LLM workloads while preserving memory efficiency and scalability. The authors open‑source their implementation (https://github.com/sail‑sg/odc), enabling immediate adoption. Limitations include reliance on GPU‑to‑GPU RDMA (performance on CPU‑only clusters remains to be explored) and the fact that the current design still uses synchronous optimizer updates; extending ODC to fully asynchronous training is left as future work. Overall, ODC offers a compelling alternative to collective‑communication‑heavy DP for LLM post‑training, delivering substantial speedups and simplifying workload balancing.