PaperlayerView exact revision on arXiv

Source-linked AI summary

When Do Larger Batches Help Scale LLM Reinforcement Learning?

Ziniu Li, Jinbo Wang, Guanhua Huang, Feiyuan Zhang, Pengbo Li, Alex Chen

arXiv:2608.29296v1cs.LGcs.AI

TL;DR

The paper asks whether larger batches actually reduce wall-clock time-to-target in LLM reinforcement learning despite consuming more samples per update. It separates sample-indexed learning from execution throughput, finding that tuned batches preserve learning approximately over a bounded range while improving throughput and reducing GRPO time-to-target.

  • Problem

    Whether larger batches reduce wall-clock time-to-target remains unclear because they trade lower gradient variance against greater sample consumption and execution cost.

  • Method

    The paper compares batch configurations at equal cumulative samples with retuned hyperparameters, then measures rollout-generation and training throughput separately.

  • Results

    Larger batches improve fixed-hardware generation throughput by up to 2.29×, while tuned GRPO configurations reduce time-to-target by up to 29%.

  • Takeaways & Limitations

    Larger batches should be aligned for learning first and accelerated only while their throughput gain exceeds their samples-to-target penalty.

  • Takeaways & Limitations

    The experiments retune only learning rate, leaving other batch-dependent hyperparameters and broader invariance rules for future work.

Abstract

from arXiv · show

Larger batches reduce the variance of stochastic gradients per update and are therefore often expected to accelerate training. Yet whether this statistical benefit translates into lower wall-clock time-to-target remains unclear, because each update consumes more samples and may take longer to execute. We study this tradeoff in reinforcement learning for large language models. We separate its algorithmic and systems effects by comparing learning and execution along their natural axes. At the algorithmic level, we compare configurations at equal cumulative sample counts while retuning batch-dependent hyperparameters. Over a bounded range of batch sizes, this procedure yields an approximately batch-size-invariant family whose members follow similar sample-indexed learning trajectories. At the systems level, we exploit the computational asymmetry between rollout generation and training: autoregressive generation is often memory-bandwidth-bound at low concurrency, whereas training work scales approximately with the number of processed tokens. Combining these two views yields a direct decision rule: a larger-batch configuration reduces time-to-target only when its throughput gain exceeds its samples-to-target penalty. Experiments with GRPO and PPO support both sides of this decomposition. At the algorithmic level, square-root learning-rate scaling with Adam produces approximately batch-size-invariant learning curves over a bounded range of batch sizes. At the systems level, larger batches improve generation throughput by up to 2.29x on fixed hardware. In GRPO, combining higher throughput with learning-rate retuning reduces time-to-target by up to 29%, whereas increasing the batch without retuning is slower despite its higher throughput.

1 Introduction

The paper asks when larger batches reduce wall-clock time-to-target in LLM reinforcement learning, separating learning-per-sample effects from systems throughput. It proposes retuning batch-dependent hyperparameters before exploiting rollout-generation parallelism.

  • Larger batches can improve hardware utilization and reduce gradient variance, but consume more samples per update and may take longer to execute.
  • The comparison fixes cumulative samples and performance targets, separating algorithmic sample efficiency from realized end-to-end throughput.
  • A larger batch helps only when its throughput gain exceeds its samples-to-target penalty.
  • Square-root learning-rate scaling with Adam produces approximately batch-size-invariant learning curves across GRPO and PPO over a bounded batch-size range.The invariance begins to break down at extremely large batches and depends on batch-dependent learning-rate retuning.
  • 2.29× generation throughput and up to 29% lower GRPO time-to-target demonstrate the systems and end-to-end benefits of tuned larger batches.A fixed-learning-rate control is slower despite higher throughput.

2 Problem Formulation

The formulation distinguishes samples used for learning from execution throughput. It defines training and generation batch quantities while covering PPO, ReMax, and GRPO within a common rollout-and-update framework.

  • At a rollout boundary, the current policy samples G responses for each of P prompts, producing PG rollout samples.
  • PPO trains an actor and value model, while ReMax and GRPO use reward-based baselines without the same learned-critic structure.
  • One optimizer update uses nominal training batch B_train = PG, while generation batch B_gen may be larger under asynchronous streaming.Micro-batching and gradient accumulation affect execution only.
  • The algorithmic curve J(N) determines samples-to-target N⋆, whereas a systems realization determines throughput q = N⋆/T and time-to-target T.

3 Training-Side Comparability: Invariance and the Critical Batch

The paper defines batch-size invariance through sample-indexed learning trajectories and uses it to identify when larger batches preserve learning per sample. Experiments show approximate invariance over bounded ranges in GRPO and PPO, with alignment breaking at larger batches or under fixed learning rates.

  • Invariance framework: Batch-size invariance asks whether retuned hyperparameters keep expected evaluation trajectories aligned as cumulative training samples increase.This permits comparisons at equal sample budgets rather than equal optimizer updates.
  • Invariance framework: A larger batch reduces sequential updates at fixed samples-to-target, but diminishing variance reduction eventually limits this benefit.The critical training batch marks the onset of diminishing update-quality improvement.
  • Retuning and empirical evidence: GRPO trajectories align for P = 64–1024 at G = 8, while P = 2048 and P = 4096 deviate from the aligned family.The actor-gradient proxy follows P^-1/2 over P = 64–1024 but flattens from P = 2048 to P = 4096.
  • Retuning and empirical evidence: Fixed learning rates weaken batch scaling: the P = 256 control learns more slowly than the P = 128 reference, while scaled-LR updates reach 0.50–0.67× of reference updates.The scaled-LR interval contains the ideal 0.50× prediction for doubled batch size; the fixed-LR interval is 0.75–0.83×.
  • Retuning and empirical evidence: In GRPO, total responses P × G—not prompt count P alone—form the relevant empirical batch and sample unit over the measured group sizes.Matched P,G pairs show approximately aligned trajectories and similar actor-gradient summaries.
  • Retuning and empirical evidence: PPO trajectories align for B = 256–2048 but deviate at B = 4096, while actor and critic gradient norms exhibit different batch scaling.The actor slope is −0.471 over B = 256–2048; the critic remains comparatively batch-flat and noisier.
  • Retuning and empirical evidence: Square-root learning-rate scaling with Adam yields approximate batch-size invariance across bounded ranges in both GRPO and PPO.The protocol varies only learning rate, so failures can also reflect the restricted retuning scheme.

4 Generation–Training Asymmetry and Fixed-Hardware Efficiency

On fixed hardware, larger RL batches expose a generation–training asymmetry: generation throughput can improve sublinearly with concurrency, while training work grows approximately with processed tokens. This creates wall-clock savings only when the generation-throughput gain outweighs the added execution and sample costs.

  • Systems opportunity: Fixed hardware makes higher realized throughput the only source of wall-clock gains from larger batches.The systems analysis therefore separates rollout-collection behavior from training execution.
  • Generation scaling: 2.29×: PPO response-collection throughput rises when Btrain increases from 256 to 1024, while collection time grows only 1.74×.Responses per optimization boundary increase 4×, producing the throughput gain.
  • Generation scaling: 1.36×: GRPO response throughput improves when P reaches 512, although additional scaling shows diminishing returns.From P = 128 to 512, responses increase 4× while collection time grows 2.93×.
  • Training scaling: 2.06×: doubling GRPO P from 128 to 256 increases mean actor-update time from 101.3 to 208.6 seconds.Actor-update time is close to batch-proportional, unlike sublinear rollout-collection time.
  • Mechanism: Autoregressive generation is often memory-bandwidth-bound at low concurrency, so active sequences can amortize weight traffic and increase throughput.Training instead processes many token positions with higher-arithmetic-intensity matrix operations, causing work to grow approximately with batch.
  • Practical boundary: The critical generation batch marks where most attainable generation throughput has been realized, but it is not itself the wall-clock-optimal batch.Saturation can arise from arithmetic, KV-cache traffic, response-length imbalance, or memory capacity; deployment requires profiling feasibility and stable throughput.

5 Operating Regimes and Practical Guidance

Time-to-target combines samples required to reach a fixed target with realized throughput. The practical rule is to align learning across batches first, then increase batching only while throughput gains exceed the samples-to-target penalty.

  • Decision rule: A larger batch reduces time-to-target only when its throughput gain exceeds its samples-to-target penalty rN.The accounting uses T = N⋆/q and compares configurations at a fixed performance target.
  • Operating regimes: Generation-knee-first scaling creates a broad near-optimal region while rN remains approximately one; further scaling eventually raises time-to-target.When the training knee occurs first, the optimum can lie beyond the training critical batch if throughput continues improving sufficiently.
  • Measured GRPO sweep: 0.71: normalized GRPO time-to-target at P = 1024, before rising to 1.23 at P = 2048 and 1.18 at P = 4096.The later increases occur once rN reaches 1.60.
  • Measured GRPO sweep: 1.42×: the fixed-LR P = 256 control takes this much longer than reference because its 1.17× throughput gain does not offset its 1.67× sample penalty.Other rows use square-root learning-rate scaling, with P = 128 as the reference.
  • Practical procedure: The two-stage procedure first retunes batch-dependent hyperparameters and estimates rN, then increases generation concurrency and measures stable end-to-end throughput q.It returns the feasible configuration maximizing q/rN.

6 Related Work

Related work separates classical critical-batch analysis, policy batch-size invariance, rollout-scaling methods, and systems throughput mechanisms. This paper connects these threads by evaluating batch selection through samples-to-target and end-to-end throughput in fixed-hardware LLM RL.

  • Classical batch scaling: Classical critical-batch studies link useful data-parallel scaling to gradient noise, workload-dependent parallelism, and fair hyperparameter tuning.Their setting assumes examples are available and observes diminishing returns from increasing batch size.
  • Policy invariance: Policy batch-size invariance preserves behavior as a function of processed samples through batch-dependent retuning, including square-root learning-rate scaling for Adam.Prior PPO constructions use EWMA mechanisms and a single policy epoch for their invariance construction.
  • Position of this work: This work empirically transfers invariance analysis to modern LLM RL with single-pass GRPO and PPO, asynchronous partial rollouts, and streaming generation.It does not introduce the EWMA proximal-policy mechanism.
  • Systems comparison: The fixed-hardware framework compares the statistical cost of larger batches, measured by samples-to-target, with their systems benefit, measured by end-to-end throughput.Systems mechanisms such as weight-traffic amortization and KV-cache constraints affect attainable throughput but not the comparison criterion.
  • Rollout scaling: BroRL increases rollouts per prompt to improve dynamic-sampling acceptance and generation throughput, whereas this paper formally studies batch-size invariance across prompt and group dimensions.Related methods allocate prompt-dependent rollout budgets rather than uniformly increasing group size.

7 Conclusion

The paper concludes that larger batches help LLM RL when retuning preserves sample-indexed learning and throughput gains exceed sample penalties. Its experiments support this rule over a bounded regime, while identifying broader scaling and PPO actor–critic questions as open boundaries.

  • Conclusion: Batch-size invariance makes physical batches comparable through aligned sample-indexed learning curves over a bounded range.The fixed-learning-rate control and largest-batch runs show that invariance is neither automatic nor unbounded.
  • Practical takeaway: Tuned GRPO configurations convert higher generation throughput into lower time-to-target when throughput gains exceed samples-to-target penalties.The practical prescription is to retune batch-dependent hyperparameters before increasing batching.
  • Scope and open questions: The experiments retune only the learning rate, leaving scaling rules for other batch-dependent hyperparameters as future work.The paper also calls for measuring how training and generation critical batches shift across models, algorithms, hardware, and execution modes.
  • Scope and open questions: PPO requires clarification of coupled actor–critic batch scales because actor and critic optimization can respond differently to larger batches.The conclusion identifies this as an unresolved extension of the batch-scaling analysis.

A Analytical Details for Generation–Training Asymmetry

Generation and training respond differently to batch scaling: training latency grows with batch-dependent compute, while decode can amortize largely batch-independent weight traffic until compute or KV-cache limits dominate. This asymmetry explains why larger active decode batches can improve throughput, subject to concurrency and memory constraints.

  • Decode scaling: A larger logical generation batch improves throughput only insofar as scheduling produces higher active concurrency among unfinished sequences.Active concurrency depends on admission, scheduling, and the response-length distribution, so logical and kernel batch dimensions need not match.
  • Decode scaling: At small active decode batches, weight transfer dominates latency, so increasing concurrency raises token throughput nearly linearly.The same weights are applied to all active sequences, making weight traffic largely independent of active concurrency in this regime.
  • Decode scaling: At larger active batches, compute or KV-cache traffic dominates, causing step time to grow and throughput to approach a plateau.Increasing active batch amortizes weight traffic and raises arithmetic intensity until the hardware ridge point or another limit is reached.
  • Analytical comparison: Figure 9 compares analytical per-layer lower bounds for a training proxy and rollout across local training or active decode batches from 1 to 32.The panels use different meanings of B and do not imply equal training and decode batches in asynchronous deployment; the values are not end-to-end measurements.
  • Training–generation asymmetry: 32× larger local training batches raise compute-bound training-proxy latency by 32× while leaving ideal aggregate training throughput unchanged.In contrast, 32× larger active decode batches raise I/O-bound latency by only 2.32× and yield a 13.8× ideal aggregate decode throughput gain.
  • Memory constraints: KV-cache capacity can prevent aggregate generation throughput from reaching its plateau, imposing a feasibility limit that grows with context length.The feasible active batch must remain within the memory bound throughout decoding, alongside runtime admission and stability constraints.
Loading 2608.29296v1…