PaperlayerView exact revision on arXiv

Source-linked AI summary

Revisiting Parameter Server in LLM Post-Training

Xinyi Wan, Penghui Qi, Guangxing Huang, Chaoyi Ruan, Min Lin, Jialin Li

arXiv:2601.19362v1cs.DCcs.AI

TL;DR

LLM post-training has highly variable sequence lengths, making FSDP’s fine-grained collective synchronization inefficient under imbalanced workloads. ODC replaces those collectives with point-to-point communication and achieves up to a 36% speedup over conventional FSDP across diverse post-training tasks.

  • Problem

    Variable sequence lengths create persistent computational imbalance, while FSDP’s per-layer collective synchronization can under-utilize devices with smaller workloads.

  • Method

    ODC adapts FSDP into a decentralized parameter server by replacing per-layer collectives with independent point-to-point parameter fetches and gradient pushes, synchronizing at minibatch boundaries.

  • Results

    Up to 36% speedup over conventional FSDP was achieved, with consistent throughput and utilization improvements across diverse LLM post-training tasks.

  • Takeaways & Limitations

    ODC is a strong fit for imbalanced LLM post-training workloads, improving device utilization and throughput while retaining FSDP’s memory and scaling advantages.

  • Takeaways & Limitations

    ODC forgoes hierarchical interconnect optimizations, potentially increasing communication overhead in multinode settings.

Abstract

from arXiv · show

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.

1 INTRODUCTION

LLM post-training violates the balanced-workload assumption underlying collective communication because variable sequence lengths create persistent device imbalance. ODC adapts parameter-server principles to FSDP with point-to-point communication, reducing synchronization and preserving sharded data-parallel scaling benefits.

  • Motivation: Variable sequence lengths create persistent computational imbalance because attention cost grows quadratically with sequence length.Activation memory grows linearly, while attention computation grows quadratically.
  • Motivation: Packing strategies can reduce but not eliminate workload skew, especially when memory constraints split minibatches into microbatches.Microbatch splitting also increases the number of synchronization points.
  • Problem: FSDP is memory-efficient and standard for LLM post-training, but sharded data parallelism suffers severely from workload imbalance.FSDP shards parameters, gradients, and optimizer states across devices, enabling scaling to trillion-parameter models.
  • Contribution: ODC replaces per-layer collective operations with point-to-point parameter fetches and gradient pushes, reframing FSDP as a decentralized parameter server.Server and worker roles are colocated, preserving FSDP’s memory and scaling advantages.
  • Contribution: ODC brings parameter-server workload tolerance into FSDP instead of building a standalone parameter server.The paper argues that parameter-server architecture is better suited to heterogeneous workloads in LLM post-training.

2 BACKGROUND

LLM training uses gradient accumulation over multiple microbatches when a desired minibatch exceeds device memory. In FSDP, per-layer collective communication reconstructs parameters and aggregates gradients, creating synchronization barriers that make imbalanced workloads inefficient.

  • Gradient Accumulation: A minibatch is split into M microbatches when the desired batch exceeds memory, with gradients accumulated before one optimizer update.For microbatch m, per-parameter gradients g(m) are computed and aggregated using weights w_m.
  • FSDP Communication: FSDP partitions parameters and gradients across devices, using all-gather to materialize layer parameters and reduce-scatter to aggregate gradients.Parameters are reconstructed before each layer’s forward and backward computation, then discarded after use to save memory.
  • Synchronization Bottleneck: Per-layer collective operations impose synchronization barriers, forcing faster devices to idle until the slowest device completes each communication step.All devices must finish all-gather before layer computation and reduce-scatter before gradient accumulation proceeds.
  • Batching Limitations: Batching research seeks an optimal assignment P⋆ that minimizes minibatch runtime, but such approaches face fundamental limitations.Runtime is bounded by the slowest device at each per-layer step under the batching solution P_M.

3 ON-DEMAND COMMUNICATIONS

ODC adapts FSDP into a decentralized, modern parameter-server paradigm by replacing fine-grained collective synchronization with coarser-grained, on-demand point-to-point communication. This preserves FSDP’s memory layout and computational graph while improving tolerance to imbalanced workloads.

  • Communication design: By replacing per-layer collectives with on-demand point-to-point communication, ODC reduces idle time caused by devices waiting for the slowest worker.Standard FSDP’s fine-grained synchronization barriers violate per-device computational independence and directly cause idle time under imbalanced workloads.
  • Communication design: ODC replaces FSDP’s synchronous collectives with targeted gather requests and scatter-accumulate operations, relaxing synchronization without altering training semantics.It preserves FSDP’s memory layout and computational graph while replacing collective communication with point-to-point operations.
  • Implementation: Non-intrusive transfers let devices gather or scatter-accumulate data without interrupting computation on target devices.This property is essential because colocated workers may communicate with servers while those servers concurrently compute.
  • Decentralized parameter server: ODC reframes FSDP as a decentralized parameter server by colocating server and worker roles on every device.Devices own parameter and optimizer-state shards while simultaneously executing forward and backward computation on assigned data.
  • Decentralized parameter server: ODC retains FSDP’s memory efficiency, decentralization, scalability, and simplicity while gaining the imbalance tolerance of a parameter server.This follows from replacing FSDP’s per-layer collective communication with on-demand point-to-point operations.
  • Implementation: ODC uses CUDA IPC for intra-node and NVSHMEM for inter-node RDMA communication, with gradient accumulation handled by a lightweight daemon.Its integration into FSDP requires replacing collective calls with ODC primitives and retrieving accumulated gradients at minibatch end.

4 SIMPLIFIED LOAD BALANCING WITH ODC

ODC simplifies load balancing by decoupling microbatch execution across devices, removing FSDP’s uniform-microbatch requirement and shifting balancing from microbatches to minibatches. This addresses the limitations of microbatch-level sequence packing under memory constraints and highly variable sequence lengths.

  • Sequence Packing: Sequence packing concatenates samples with attention masks to reduce padding waste and balance workloads across microbatches.It was introduced to improve utilization under sequence-length variation and has been broadly adopted and extended.
  • Sequence Packing: Microbatch-level packing remains limited because device memory bounds microbatch size, leaving substantial workload variance across devices.The limitation is amplified in long-sequence training regimes, where compute alignment can become infeasible.
  • Sequence Packing: Compute alignment can be infeasible when a microbatch contains one maximum-length sample that shorter samples cannot match in runtime.In this case, no feasible packing of shorter samples can achieve equivalent runtime.
  • ODC Load Balancing: ODC decouples microbatch execution across devices, eliminates FSDP synchronization barriers, and removes the requirement for equal microbatch counts per device.This enables shifting workload balancing from the fine-grained microbatch level to the coarser minibatch level.

5 EVALUATIONS

ODC is evaluated across SFT and RL post-training tasks on DeepSeek-R1-Distill-Qwen models from 1.5B to 32B, using varied communication and load-balancing schemes. It consistently improves throughput over collective communication, with gains depending on packing, minibatch size, sequence length, and communication scope.

  • Evaluation Setup: Evaluations cover SFT on LongAlign and SWE-Smith trajectories, plus RL with GRPO on AIME prompts.These tasks span context-window extension, software engineering, and mathematical reasoning workloads.
  • Evaluation Setup: Models range from 1.5B to 32B and run on up to 32 NVIDIA A100 80G GPUs.RL experiments use models up to 14B on 16 GPUs because 32B inference would take too long.
  • SFT Results: ODC consistently improves SFT throughput over collectives in unpacked and packed settings, reaching up to a 36% speedup under packing.At minibatch size one, all methods perform similarly because ODC synchronizes after every sample.
  • RL Results: ODC achieves up to 10% speedup over collectives on RL tasks, with smaller gains attributed to implementation constraints and less long-tailed sequence lengths.verl requires identical sample counts per device, limiting LB-Mini; AIME distributions are also less long-tailed than SFT datasets.
  • Load Balancing: LB-Mini often outperforms LB-Micro at small minibatch sizes, while LB-Micro narrows the gap as minibatches grow.Minibatch-level balancing lets devices process different numbers of microbatches, whereas larger minibatches give LB-Micro more balancing flexibility.
  • Parametric Study: ODC acceleration peaks at moderate minibatch sizes, increases with sequence length, and decreases with packing ratio.Larger batches give collectives more balancing flexibility; longer sequences amplify quadratic compute cost and imbalance, while packing improves balance.
  • Communication Primitives: Within a node, ODC primitives achieve bandwidth comparable to collectives, but ODC lags significantly when communication spans multiple nodes.The comparison uses gather and scatter-accumulate versus NCCL all-gather and reduce-scatter with synchronized launches and barriers.

6 DISCUSSION

ODC preserves communication-computation overlap and offers hybrid sharding to mitigate communication costs, but its direct point-to-point topology forgoes hierarchical collective optimizations. The discussion identifies further opportunities in communication optimization, relaxed synchronization, elasticity, and fault tolerance.

  • Communication trade-offs: ODC uses point-to-point RDMA without increasing communication volume, but forfeits hierarchical interconnect optimizations available to collective primitives.The discussion notes that larger DP scale may amplify straggler-related benefits, although the supplied passage is truncated before completing that argument.
  • Overlapping Communication with Computation: For long sequences, ODC overlaps communication with computation because communication per microbatch is constant in sequence length while computation scales as O(s^2).This overlap hides communication latency and yields no significant slowdown in long-context workloads despite non-hierarchical communication.
  • Hybrid Sharding: Hybrid sharding addresses microbatches too small to hide communication costs by restricting parameter and gradient sharding within nodes while retaining cross-node optimizer-state sharding.This removes cross-node parameter gathering and gradient scatter-accumulate at the cost of higher per-node memory usage.
  • ODC-specific Optimizations: ODC’s communication graph could be optimized by fetching cached parameter shards from same-node peers, creating a topology-aware hierarchical path.This is presented as a future optimization for the current direct point-to-point implementation.
  • Relaxing Synchronization Guarantees: Relaxing the minibatch-boundary synchronization barrier could reduce idle time and improve utilization through asynchronous or bounded-staleness updates, especially in heterogeneous environments.The current barrier preserves identical training semantics, while relaxing it would require addressing the resulting changes in update guarantees.
  • Elasticity and Fault Tolerance: Integrating elasticity and fault tolerance would improve ODC’s resilience and flexibility for large-scale, long-running LLM training jobs.The discussion contrasts PS-style architectures’ natural support for these capabilities with the brittleness and resizing difficulty of collective-based systems.

7 CONCLUSION

The paper identifies fine-grained FSDP synchronization as a bottleneck under imbalanced LLM post-training workloads and proposes ODC, which replaces collectives with point-to-point communication to improve execution decoupling, load balancing, throughput, and utilization.

  • 7 CONCLUSION: FSDP’s per-layer all-gather and reduce-scatter collectives create synchronization barriers that amplify straggler effects from workload imbalance.This bottleneck arises in modern sharded data-parallel training for LLM post-training.
  • 7 CONCLUSION: ODC replaces these collectives with point-to-point operations, relaxing synchronization from the layer level to the minibatch level.The approach reframes FSDP as a decentralized parameter server.
  • 7 CONCLUSION: ODC decouples device execution and enables more effective load balancing across imbalanced workloads.These design effects address the straggler problem caused by fine-grained collective synchronization.
  • 7 CONCLUSION: ODC delivers consistent throughput and utilization gains across long-sequence supervised fine-tuning and reinforcement-learning tasks.The reported gains span a range of LLM post-training workloads.

B IMPLEMENTATION DETAILS OF ODC

ODC uses CUDA-IPC for intra-node communication and Triton-Distributed primitives with a custom kernel for inter-node communication. Its gather and scatter-accumulate operations use RDMA get/put mechanisms, with transfer limiting and server-side polling to support concurrent communication without observable slowdown of the colocated compute process.

  • Communication mechanisms: CUDA-IPC provides native remote-GPU tensor reads and writes for intra-node communication without custom GPU kernels.For inter-node communication, ODC uses a custom kernel built with Triton-Distributed’s put mem and get mem primitives.
  • Gather: Gather uses get mem for each rank to pull data from all other ranks.Limiting the communication payload per transfer helps stabilize RDMA traffic when servers receive requests from multiple clients.
  • Scatter-accumulate: Scatter-accumulate combines put mem transfers with same-channel notifications that trigger server-side gradient accumulation.A lightweight daemon polls for notifications without occupying GPU SMs, causing no observable slowdown of the colocated compute process.

C SEQUENCE PACKING STRATEGIES USED IN EXPERIMENT … D COMMUNICATION VOLUME COMPARISON

The paper evaluates sequence-packing strategies that address workload imbalance and memory feasibility, contrasting conventional microbatch balancing with ODC-enabled minibatch balancing. It also reports that ODC matches collective communication volume overall but increases cross-node traffic.

  • C SEQUENCE PACKING STRATEGIES USED IN EXPERIMENT: Karmarkar-Karp balances computational workloads while iteratively rejecting memory-infeasible microbatch partitions to prevent out-of-memory errors.The method extends Verl’s implementation with memory-feasibility validation before accepting a partition.
  • C.1 LB-MICRO AND LB-MINI: LB-Micro balances workloads at the microbatch level while requiring every device to process the same number of microbatches.LB-Mini instead balances at the minibatch level through ODC, removing the rigid equal-microbatch constraint.
  • C.2 VERL NATIVE TWO-LEVEL PARTITIONING STRATEGY: Verl’s native packing assumes equal samples per device and equal microbatch counts because layer-level synchronization requires synchronized processing.These constraints motivate Verl’s two-level hierarchical heuristic.
  • C.3 OPTIMIZED TWO-LEVEL PARTITIONING STRATEGY: Verl’s native strategy balances the global batch before minibatch splitting, so it does not guarantee balance within individual minibatches.The optimized strategy reverses this order by partitioning into minibatches first and balancing each minibatch across devices.
  • C.3 OPTIMIZED TWO-LEVEL PARTITIONING STRATEGY: The optimized procedure splits global data into minibatches, balances each minibatch across ranks, and then partitions each rank’s minibatch into microbatches.This ordering is reflected in the implementation pseudocode and is reported to yield substantial throughput improvements.
  • D COMMUNICATION VOLUME COMPARISON: Collectives and ODC send the same total communication volume, (D −1) ∗K, under the comparison’s device and parameter-size assumptions.Here, D is the total device count, G is devices per node, and K is per-device local parameter or gradient size.
  • D COMMUNICATION VOLUME COMPARISON: ODC increases cross-node communication because clients independently gather, scatter, and accumulate, potentially slowing end-to-end communication.The communication trade-off arises despite equal total volume between ODC and collectives.

E ZERO++ STYLE HYBRID SHARDING · F CONVERGENCY VERIFICATION

Hybrid sharding retains comparable acceleration to full sharding for shorter sequences, reaching up to 28% when ODC is compared with collectives, but uses more memory. Convergence verification shows that ODC and collectives produce almost identical loss curves in the tested setting.

  • E ZERO++ STYLE HYBRID SHARDING: Hybrid sharding is evaluated on LongAlign sequences truncated to one-eighth their original length, with maximum length 8k and average length 2k.This setting targets shorter sequence lengths, for which the hybrid strategy is described as particularly effective.
  • E ZERO++ STYLE HYBRID SHARDING: Up to 28% acceleration is achieved when comparing ODC against collectives with hybrid sharding.The reported acceleration is comparable to that of full sharding.
  • E ZERO++ STYLE HYBRID SHARDING: Hybrid sharding incurs higher memory usage than fully sharded training.A detailed memory-usage comparison is provided in Figure 13.
  • E ZERO++ STYLE HYBRID SHARDING: The memory comparison covers ODC under hybrid and full sharding.Figure 13 specifically presents memory consumption for these two sharding configurations.
  • F CONVERGENCY VERIFICATION: ODC and collectives produce almost identical loss curves on the 8k-sample LongAlign training run.The curves are shown in Figure 14.
  • F CONVERGENCY VERIFICATION: Convergence correctness is evaluated by comparing loss curves while training a 1.5B model from scratch on 8k LongAlign samples.Training from scratch is used to produce a clearer loss-descent trajectory.

G DETAILED EXPERIMENT DATA

This section presents detailed timing and bubble-rate data for SFT and RL, using packing-algorithm estimates to quantify idle time from workload imbalance. The reported ODC acceleration closely tracks predicted bubble rates, attributing gains primarily to reduced imbalance-induced idle time.

  • Timing Data: Detailed timing data are reported for both supervised fine-tuning (SFT) and reinforcement learning (RL).The data appear in Tables 5 and 3, respectively.
  • Bubble Rate Data: Bubble rate measures device idle time caused by workload imbalance as a fraction of total runtime, estimated by the packing algorithm.Bubble-rate results are reported for RL and SFT in Tables 4 and 6.
  • Performance Interpretation: ODC acceleration closely correlates with packing-predicted bubble rate, indicating that gains primarily come from reducing imbalance-related idle time.This relationship is observed across the reported SFT and RL experiment data.
Loading 2601.19362v1…