Revisiting Parameter Server in LLM Post-Training

2026-01-27

AI Summary

This work addresses the synchronization bottleneck in Fully Sharded Data Parallel (FSDP) during large language model (LLM) post-training, where variable sequence lengths cause inefficient All-Gather and Reduce-Scatter operations, degrading device utilization. To overcome this, the authors introduce a parameter server paradigm into FSDP and propose an On-Demand Communication (ODC) mechanism that replaces collective communication with point-to-point exchanges. ODC enables only one synchronization per mini-batch and supports micro-batch-level dynamic load balancing. Implemented on top of PyTorch FSDP with integrated dynamic load distribution and communication scheduling, ODC significantly improves device utilization and training throughput across diverse LLM post-training tasks, achieving up to a 36% speedup over standard FSDP.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Distributed Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

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.

Problem

• large language model
• post-training
• imbalanced workload
• parameter server
• collective communication

Innovation

• On-Demand Communication
• Parameter Server
• FSDP
• Load Imbalance
• LLM Post-Training