# Revisiting Parameter Server in LLM Post-Training

arXiv: https://arxiv.org/abs/2601.19362 · HTML: https://mlsanity.com/paper/2601.19362/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 46 references that are in the map.
- First posted on arXiv 27 Jan 2026.

## Builds on (most cited first, 23 of 23)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [SWE-smith: Scaling Data for Software Engineering Agents](https://mlsanity.com/paper/2504.21798/) (2025)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://mlsanity.com/paper/2505.24298/) (2025)
- [Revisiting Distributed Synchronous SGD](https://mlsanity.com/paper/1702.05800/) (2017)
- [LongAlign: A Recipe for Long Context Alignment of Large Language Models](https://mlsanity.com/paper/2401.18058/) (2024)
- [Efficient Training of Large Language Models on Distributed Infrastructures: A Survey](https://mlsanity.com/paper/2407.20018/) (2024)
- [Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler](https://mlsanity.com/paper/2504.19442/) (2025)
- [WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training](https://mlsanity.com/paper/2503.17924/) (2025)
- [Enhancing Training Efficiency Using Packing with Flash Attention](https://mlsanity.com/paper/2407.09105/) (2024)
- [Hierarchical Balance Packing: Towards Efficient Supervised Fine-tuning for Long-Context LLM](https://mlsanity.com/paper/2503.07680/) (2025)

## Built on by (most cited first, 2 of 2)

- [Rethinking the Trust Region in LLM Reinforcement Learning](https://mlsanity.com/paper/2602.04879/) (2026)
- [NestPipe: Large-Scale Recommendation Training on 1,500+ Accelerators via Nested Pipelining](https://mlsanity.com/paper/2604.06956/) (2026)

## Abstract (arXiv)

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 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.