How LiteScale Cuts Wait Times in Large‑Model Post‑Training with Gradient Accumulation

The waiting problem in large‑model RL post‑training

Practitioners of large‑model reinforcement learning often face a "waiting" bottleneck: inference (rollout) must finish before training can update weights, and the updated model must wait for the next inference, creating a single‑lane pipeline.

Why synchronous rollout is suboptimal

Traditional pipelines treat inference and training as two separate systems with identical pacing, forcing each rollout batch to be fully collected before any parameter update. This mirrors a factory line where processing and assembly cannot overlap, severely limiting throughput.

Partial (cross‑step) rollout and its limits

Partial rollout (or cross‑step rollout) overlaps inference and training by consuming only a portion of rollout data per update. However, the remaining rollout data become off‑policy for the newly updated model, potentially degrading training accuracy.

Gradient accumulation as a more intuitive solution

Gradient accumulation, a known technique for handling large batches that exceed GPU memory, is repurposed for RL post‑training. Each incoming rollout batch is used to compute gradients immediately, but parameter updates are deferred until all rollout data for the current iteration are collected, then a single update is performed.

Ensuring loss equivalence

The key is twofold:

Loss normalization factor : In Megatron's forward_step(), loss is divided by num_microbatches. In asynchronous scenarios, batch sizes vary, breaking this normalization. LiteScale adds a new argument global_num_microbatches to forward_backward_pipelining_without_interleaving() and to forward_step(), using the total number of micro‑batches for the whole iteration instead of the per‑call count. The modified code replaces output_tensor /= num_microbatches with

if global_num_microbatches is not None:
    output_tensor /= global_num_microbatches
else:
    output_tensor /= num_microbatches

Gradient storage : Megatron's DistributedOptimizer originally overwrites shard_main_param.grad each step, discarding previous gradients. LiteScale introduces accumulate_grad_step() that adds new gradients to the existing ones, preserving accumulated information until the final update. The method checks for an existing gradient and either clones the first or adds to it.

# accumulate_grad_step() behavior
existing = shard_main_param.grad
shard_val = shard_model_grad.float()
if existing is None:
    shard_main_param.grad = shard_val.clone()
else:
    existing.add_(shard_val)

After all micro‑batches are processed, step_with_accumulated_grads() performs gradient clipping ( clip_grad_norm) and a single parameter update ( step_with_ready_grads).

Online Knowledge Distillation (GKD) and LogitsExpress

Generalized Knowledge Distillation (GKD) uses KL divergence between a teacher and a student model, requiring transmission of full‑vocabulary logits for each token—a prohibitive bandwidth cost. LiteScale applies a top‑K truncation, sending only the highest‑probability logits, which reduces data volume by an order of magnitude with minimal accuracy loss.

LogitsExpress solves the remaining communication challenge across heterogeneous Megatron instances (e.g., teacher 4DP×8TP vs. student 8DP×4TP). Instead of aggregating all logits on a single node, it performs direct point‑to‑point NCCL transfers based on a bucketed neighbor mapping derived from the smaller TP dimension. The process:

Each GPU broadcasts its DP/TP coordinates and IP, forming a global topology table.

Rank 0 on teacher and student exchanges TP/DP sizes to verify divisibility.

TP axes are bucketed; GPUs sharing a bucket become neighbors.

Neighbors establish NCCL process groups without a central aggregator.

During each round, teachers send sliced logits ( (batch//dp_t, seq_len, vocab//tp_t)) to the corresponding student, which reassembles them into (batch//dp_s, seq_len, vocab//tp_s).

This design dramatically improves communication efficiency and scalability for large‑scale distillation.

Why LiteScale instead of existing frameworks

Frameworks like OpenRLHF or veRL embed heavy control planes (e.g., Ray) and complex scheduling, making debugging difficult. LiteScale deliberately retains Megatron's native training pipeline, decouples inference from training, and keeps each module replaceable, offering a lightweight, transparent environment for research.

Modular rollout architecture

Rollout runs in an independent thread communicating with the training process via two Python queues ( input_queue and output_queue), eliminating shared memory and locks. The architecture consists of:

Service layer : Manages long‑lived connections to external systems (e.g., AsyncSGLangService wraps SGLang's HTTP API). Services are stateless and handle only request/response.

Worker layer : Implements the actual data processing. Each sample gets a dedicated Worker instance with a single run() method returning a MultiResponseSample. Workers are selected via a service_dict and can be mixed within a batch.

Extending the system requires only adding a new Worker subclass, registering it, and declaring the dataset_type in the YAML configuration—no changes to the training loop or Service code.

Conclusion

LiteScale demonstrates that asynchronous rollout, gradient accumulation with proper loss normalization, and a lightweight modular design can alleviate the "waiting" bottleneck in large‑model post‑training while preserving training fidelity. Its open‑source implementation (github.com/st01tyy/LightScale) provides a concise, extensible platform for researchers focusing on algorithmic innovation rather than framework complexity.