Skip to main content
Chapter 13 arranged the components of an RL system so that the hardware is busy. This chapter is about what the busy hardware spends its time and memory on. The answers are less glamorous than objectives and more often the reason a run fails: padding tokens that are computed and thrown away, a single 60k-token trajectory that runs a microbatch out of memory, a logits tensor larger than the model, a gradient that silently changes when the microbatch size does, and rollouts spent on tasks the policy always solves. Each of them is a correctness question as well as a cost question, because each fix (packing, token budgets, chunking, accumulation, filtering) can change the gradient if it is done wrong. We build and save four tools: sequence packing with a correct block-diagonal attention mask, token-budget microbatches, checkpointed chunked log-probabilities, and pass-rate filtering for the task distribution. Each is checked against the straightforward computation it replaces, and the chapter ends with the arithmetic that turns token counts into GPU-hours.

​
14.1 Padding and its waste

collate pads every episode of a batch to the longest one. The cost is visible from the lengths alone: a padded batch computes B⋅Tmax⁡B \cdot T_{\max} positions, of which ∑iTi\sum_i T_i are real.
Output:
TinyAgent’s episodes range from 10 tokens (an immediate submit) to 149 (a long episode with several observations), and the median is 18. Padding them all to 149 computes 9,536 positions for 2,472 real tokens, so 74% of the forward and backward pass is spent on padding. Agent trajectories have this shape at every scale: most episodes are short, a few are long, and the long ones set the width of the batch. Sequence packing concatenates several sequences into one row of a fixed length and fills the rows nearly full. It removes the waste, and it introduces two ways to compute the wrong thing. In a packed row, a token of the second sequence sits after the tokens of the first. A plain causal mask lets it attend to them, so its log-probability is conditioned on an unrelated episode: an attention leak. And with absolute or rotary positions, its position index continues from the first sequence instead of starting at 0. Correct packing needs a block-diagonal causal mask (a token attends to earlier tokens of its own sequence only) and position ids that restart at each sequence boundary. Packed training kernels such as FlashAttention’s variable-length interface receive the boundaries as cumulative sequence lengths instead of a mask; the computation is the same. pack_sequences places trajectories into rows of max_tokens with first-fit decreasing: longest first, each into the first row with room. It returns the packed ids, per-token position ids and segment ids (the index of the sequence within its row, −1-1 for padding), and the action masks and sampler log-probabilities in the same [R, L-1] layout as collate. A sequence’s last position predicts nothing of its own, so the mask entry that would predict the next sequence’s first token stays 0.
TinyAgent’s forward builds its own causal mask from the sequence length, so a packed forward pass has to be assembled from its parts: token embeddings plus position embeddings at the given position ids, the transformer blocks with an explicit mask, the final norm and head. packed_logits builds the mask from the segment ids: position jj may attend to position ii if i≤ji \le j and both belong to the same segment. PyTorch’s attention takes a per-row mask as a float tensor of shape [R * n_heads, L, L], with −∞-\infty where attention is forbidden. Padding positions attend to themselves only, so no row of the softmax is empty.
We pack the 64 episodes into rows of 256 tokens, TinyAgent’s maximum length, and compare the packed log-probabilities of every policy token with those of each episode run alone. Three variants show the two bugs and the fix.
Output:
Packing puts the 64 episodes into 10 rows that are 96.6% full, against 26% real tokens in the padded batch. With a plain causal mask and positions running across the row, policy-token log-probabilities differ from the unpacked ones by up to 6.4 nats, a factor of about 600 in probability. Restarting the positions fixes part of it; the leak through attention alone still moves some log-probabilities by 1.7 nats. With the block mask and restarted positions, the packed computation matches the unpacked one to 1.4×10−61.4 \times 10^{-6}, float32 rounding. Nothing about the leaky versions fails loudly: the shapes are right, the loss is finite, and training proceeds on the wrong conditional distributions. Figure 14.1: The same 64 TinyAgent episodes laid out as a padded batch (left, one row per episode, padded to the longest) and packed into rows of 256 tokens (right). Blue: policy tokens; aqua: prompt, template and tool tokens; grey: padding. Figure 14.1 draws both layouts at the same scale, so area is compute. The padded batch is mostly grey; the packed rows are nearly solid. The dark ticks in the packed rows are the sequence boundaries that the block mask must respect.

​
14.2 Token budgets and length-balanced microbatches

Packing fixes waste within a row; the next question is how many sequences go into a microbatch. A fixed number of sequences per microbatch makes memory proportional to the longest sequence that happens to land in it. Agent trajectories have long tails, so a fixed count either wastes memory on most microbatches or runs out on the unlucky one. A token budget fixes the quantity memory depends on instead: each microbatch holds as many sequences as fit in max_tokens padded positions. Sorting by length first (length-balanced microbatches) keeps sequences of similar length together, which cuts padding further.
TinyAgent’s episodes are too uniform to show the effect, so we use a long-tailed length distribution typical of agent trajectories: a log-normal with a median of 8k tokens, capped at a 64k context. The three schemes see the same 512 trajectories. For each we report the number of microbatches (a proxy for steps of the backward pass), the share of padded positions, and the largest microbatch in padded tokens, which sets peak activation memory.
Output:
With 8 trajectories per microbatch in random order, 59% of the computed positions are padding, and the largest microbatch is 524k padded tokens: it contains a trajectory at the 64k cap, and 8 times 64k is what it costs. Sorting by length before grouping removes most of the padding (3.6%), because neighbors have similar lengths, but leaves the largest microbatch at 524k, since the 8 longest trajectories still share one. The token budget keeps the padding low (3.0%) and caps every microbatch at 131k tokens, putting the longest trajectories into microbatches of one or two. It also needs fewer microbatches, 51 instead of 64, because short trajectories are grouped many to a microbatch. Peak memory is now set by the budget, a number the engineer chooses, rather than by the longest trajectory the policy happens to produce.

​
14.3 The memory of logits

Section 3.6 introduced chunked_log_probs and noted that chunking alone saves forward memory but not backward memory. Here is the arithmetic for a training microbatch. For NtokN_{\text{tok}} tokens and vocabulary VV, the full logits tensor holds NtokVN_{\text{tok}} V numbers. Qwen3’s VV is 151,936, so one token’s logits are 0.6 MB in fp32, and a 32k-token microbatch has 20 GB of logits before the log-softmax output and the gradient, each of the same size. The model’s own per-token activations at hidden width DD are a few times DD per layer, and VV is 37 to 148 times larger than DD for Qwen3’s sizes.
Output:
The fix has two parts. Chunk the time axis so that only one chunk’s logits exist at a time in the forward pass, and recompute each chunk’s logits in the backward pass instead of storing them. torch.utils.checkpoint does the second part: it runs the chunk without saving intermediates and replays it during backward. checkpointed_log_probs is Chapter 3’s chunked_log_probs with that one change.
We measure the peak memory of a forward and backward pass through Qwen3-0.6B’s output layer on the GPU, with random bf16 hidden states standing in for the final layer’s output. Only the logits computation is measured: the transformer’s own activations are not included.
Output:
Peak allocator memory varies by a few hundredths of a GB with whatever ran on the GPU earlier, so it is printed to one decimal. All three methods return identical log-probabilities. At 4,096 tokens, the full computation peaks at 7.5 GB for an output layer alone, about 12 bytes per logit: the bf16 matmul output, its fp32 copy, and the fp32 log-softmax, some of them alive at once during backward. Chunking halves that, and it grows with TT just the same (11.2 GB at 16k tokens), because autograd keeps each chunk’s log-softmax for the backward pass: chunking only removed the transient full tensor. Recomputation makes the peak independent of sequence length: 2.2 GB at both lengths, one chunk’s worth plus the weight gradient. The cost is a second forward pass through the output layer, a few percent of the step’s FLOPs. Figure 14.2: Per-GPU memory of one RL training microbatch for an 8B model with full sharding over 8 GPUs, by component (arithmetic under stated assumptions). Logits are shown for the full tensor, chunked, and chunked with recomputation. Figure 14.2 puts the logits next to everything else on the GPU for an 8B model with fully sharded parameters, gradients, and Adam states on 8 GPUs, at 32k tokens per microbatch. With the full logits tensor, the logits are the largest item, larger than the sharded model and optimizer together. Chunking reduces them to second place; recomputation to a small one. This is why logit handling is the first memory decision for long-context RL, before activation checkpointing or sharding strategy.

​
14.4 Gradient accumulation and loss normalization

A batch that does not fit in memory is split into microbatches whose gradients are summed before the optimizer step. The sum must equal the gradient of the full batch’s loss, and with token_mean it usually does not. token_mean divides by the number of policy tokens in the batch. A microbatch loop that computes token_mean per microbatch and divides by the number of microbatches gives each microbatch equal total weight, so a token’s weight becomes 1/(Knk)1 / (K n_k) instead of 1/N1/N, where nkn_k is its microbatch’s token count, KK the number of microbatches, and N=∑knkN = \sum_k n_k. Tokens in short microbatches are up-weighted. The fix is to divide every microbatch’s token sum by the global NN, which aggregate_loss already supports as mode='constant' with norm_const=N. The same bug appears across data-parallel ranks when each rank normalizes by its own token count before the gradient all-reduce.
The reference is one backward pass over all 64 episodes. We split into 4 microbatches in random order, and in length-sorted order, as a token-budget scheduler would.
Output:
With the global normalizer, both splits reproduce the full-batch gradient to 1.3×10−71.3 \times 10^{-7}, rounding error. With per-microbatch token_mean, a random split already has a relative error of 8% (cosine 0.997), because the microbatches hold 159 to 208 policy tokens. A length-sorted split, which is what a token-budget scheduler produces, holds 96 to 355 tokens per microbatch, and the error rises to 52% (cosine 0.88): the gradient points in a measurably different direction. Since the short microbatch consists of short episodes, which in this gym are mostly immediate submits, the bug systematically over-weights one kind of behavior. The loss value printed during training gives no hint of this.

​
14.5 Filtering tasks by pass rate

The default run of Chapter 9 spent most of its late rollouts on groups with zero advantage: by steps 60 to 80, about 80% of sampled groups passed in every attempt or failed in every attempt. Those episodes cost a full rollout each and contribute no gradient. The cheapest fix is upstream of the loss: do not sample tasks whose outcome is already certain. With a running estimate p^j\hat p_j of each task’s pass rate, a pass-rate filter keeps task jj for sampling only if low<p^j<high\text{low} < \hat p_j < \text{high}. With low=0\text{low} = 0 and high=1\text{high} = 1, it drops exactly the tasks whose last group was all-fail or all-pass. A softer variant samples tasks with probability proportional to p^j(1−p^j)\hat p_j (1 - \hat p_j), the variance of a Bernoulli reward, which is the expected size of a group’s advantage signal.
train_curriculum is train_rl with a task pool that changes. Each task’s estimate is its last group’s mean reward, initialized at 0.5 so that every task is sampled at first; the estimates cost no extra rollouts. Each step samples prompts_per_step tasks from those inside the filter, plus revisit tasks drawn at random from those outside it, so that an excluded task’s estimate is refreshed now and then and a task the policy forgets can come back. When few tasks pass the filter, the step samples fewer episodes: that is the point. A large task set would keep the batch full with informative tasks; our 20 training tasks show the saving directly as rollouts not spent.
With low=-1 and high=2, every task passes the filter, the pool is the whole training set, and the loop samples 8 tasks per step exactly as train_rl does, so the unfiltered arm runs through the same code. We train 60 steps per run, two seeds per arm, and evaluate on held-out tasks with 16 episodes per task.
Output:
For the first 20 steps the two arms are the same: every estimate starts at 0.5, the policy is weak, and almost every task produces mixed groups. After that, tasks start passing in all 8 attempts and leave the filtered pool, and by step 60 only one to three tasks remain inside it. The filtered arm sampled 2,088 and 2,184 episodes in 60 steps, against 3,840 for the unfiltered one, 43% to 46% fewer, and its zero-advantage share in the last 20 steps was 0.38 to 0.47 against 0.61 to 0.68. Its held-out pass rate is 0.73 and 0.77 against 0.76 and 0.82: lower by about 0.04 in both seeds, a difference of about one standard error of an 160-episode evaluation. The honest reading is: about the same held-out pass rate for a bit more than half the rollouts, with a possible small cost that two seeds cannot resolve. The cost has a plausible mechanism: late in training, the filtered arm’s steps contain only a handful of tasks, so each gradient is about fewer tasks, the same trade-off as group size against prompts in Section 9.7. Filtering on the last group alone is the crudest estimator. A task that passed 8 of 8 once may pass 6 of 8 next time; with G=8G = 8 and a true pass rate of 0.9, a group is all-pass 43% of the time. Production curricula use running averages over several visits, keep a small probability of resampling excluded tasks, and filter offline before an epoch when the dataset is large.
Evidence. Measured here, two seeds per arm, 60 steps on 20 training tasks. The rollout saving is a direct consequence of the filter and is robust; the comparison of held-out pass rates is inside seed noise. Figure 14.3 adds two more seeds.
Figure 14.3: Pass-rate filtering on the BugFix gym, four seeds per arm. Left: held-out pass rate (8 episodes per task) against the cumulative number of episodes sampled. Right: episodes sampled per step. Figure 14.3 shows the same comparison with four seeds per arm. On the left, the cumulative number of episodes: the arms separate after about step 20, when tasks begin to saturate, and the filtered arm’s curve flattens as its pool empties. On the right, each run is a point: the filtered runs sit at about 2,100 to 2,600 episodes and the unfiltered ones at 3,840. Their held-out pass rates overlap (0.72 to 0.81 against 0.76 to 0.85), with the filtered arm slightly lower on average, as in the two seeds above.

​
14.6 Throughput accounting

Every design decision in this chapter and the last trades one resource for another, and the trade can only be judged with a budget. The budget has two halves that behave differently: generation and training. Generation is a sequence of decode steps. Each step produces one token per sequence and must read every weight and every sequence’s KV cache from memory, so its speed is set by memory bandwidth, not FLOPs. TinyAgent shows the other half of the story first. rollout_batch has no KV cache: every decode step reruns the model over the whole prefix. We count the positions it processes for the 64 episodes of Section 14.1 and compare with a KV-cached decoder, which processes each position once.
Output:
Without a KV cache, sampling these 64 episodes processed 22,661 positions instead of 2,472, 9.2 times more, and the ratio grows linearly with episode length. TinyAgent can afford it; nothing larger can. A KV cache turns generation’s cost from quadratic in length to linear in compute, and moves the bottleneck to memory: every decode step must read the cache. For a real model, a first-order roofline estimate per GPU is enough to see where the time goes. A decode step for a batch of bb sequences with context cc reads the weights once (2N2N bytes in bf16) and each sequence’s KV cache (cc times the per-token KV size), so it takes at least (2N+b c sKV)/bandwidth(2N + b\, c\, s_{\text{KV}})/\text{bandwidth} seconds and yields bb tokens. Training processes every token of every trajectory (prompt, policy, and tool tokens alike, since all are inputs to the gradient) at about 6N6N FLOPs per token for forward and backward, at some model-FLOPs utilization (MFU). The block reads the per-token KV size from Qwen3-8B’s real configuration in the local cache and uses the H200’s published bandwidth (4.8 TB/s) and dense bf16 throughput (989 TFLOP/s, half the 1,979 TFLOP/s NVIDIA lists with sparsity). Utilization and batch sizes are assumptions.
Output:
Each token of context costs Qwen3-8B 147 KB of KV cache (36 layers, 8 KV heads of dimension 128, keys and values in bf16). At a 2k context, a batch of 64 decodes at about 8,600 tokens per second per GPU, roughly the training rate of about 8,000. At a 16k context, which is short for a coding agent, the KV reads dominate the weight reads and decode falls to about 1,800 tokens per second. A larger batch barely helps (1,935 tokens per second at 256), because each added sequence adds its own cache to read. Training, at an assumed 40% MFU, processes about 4.5 times more tokens per second per GPU than long-context decoding. This asymmetry is why RL systems give the inference engine most of the accelerators, and why each rollout token is far more expensive than each training token. Putting the two halves together gives the GPU-hours of a run. The block below estimates one step of a hypothetical agent RL run on Qwen3-8B, in the style of Appendix B, Problem 19, which works the same estimate for a 32B model. Every input is an assumption, labeled in the code.
Output:
Under these assumptions, a step costs about 1.6 GPU-hours, and 300 steps about 490 GPU-hours. Decoding is 59% of it, although policy tokens are only a quarter of all tokens, and the training pass itself is about a quarter. The two largest uncertainties are the decode rate, which depends on context length and batch size as the previous block showed, and the busy fraction, which absorbs tool latency and the straggler wait of Chapter 13: changing it from 0.5 to 0.25 adds about 1 GPU-hour per step. The Appendix B estimate for a 32B model with the same episode shape, about 7.7 GPU-hours per step, is in proportion.
Evidence. The H200 bandwidth and throughput figures are NVIDIA’s published specifications (inspected 2026-10-13); the KV size is computed from Qwen3-8B’s configuration file. Everything else in this section (episode shapes, batch sizes, MFU, busy fraction) is an assumption chosen for illustration, and the result is a first-order estimate that ignores attention FLOPs, communication, evaluation, and failed runs. Replace each assumption with a measurement from a pilot run before planning a budget.

​
14.7 Design problems in practice

Problem: loss normalization under gradient accumulation. Symptom: changing the microbatch size, the number of data-parallel ranks, or the batch scheduler changes the learning curve, although the global batch is the same; short episodes seem to learn faster after a switch to length-sorted microbatches. Cause: per-microbatch token_mean gives each microbatch equal weight, so a token’s weight is N/(Knk)N / (K n_k) times what the global loss assigns (Section 14.4). Length-sorted microbatches make nkn_k vary the most. Options: (a) normalize by the global token count (an all-reduce of the count before the backward pass in data parallel). (b) Use a constant normalizer (Dr. GRPO’s constant mode), which needs no global count and is independent of batch composition. (c) Equal-token microbatches, which only make the error small. Default: (a) or (b), with a unit test that compares an accumulated gradient against the single-pass gradient, as Section 14.4 does. The block prints the per-token weight distortion of the length-sorted split.
Output:
The tokens in the microbatch of short episodes get 1.93 times their correct weight, and those in the microbatch of long episodes 0.52 times: a factor of 3.7 between them, chosen by the scheduler rather than by the objective.
Problem: the packing attention leak. Symptom: packed training runs without errors, but the step-0 ratio between the trainer’s log-probabilities and the sampler’s is far from 1, the clip fraction is high from the first minibatch, and the policy learns slowly or diverges. Cause: a causal mask over a packed row lets each sequence attend to the sequences packed before it, and running position ids shift every later sequence; the trainer then scores tokens under contexts the sampler never saw (Section 14.1). Options: (a) a block-diagonal mask or variable-length attention kernel with restarted position ids. (b) No packing, paying the padding. Default: (a), with a test that compares packed and unpacked log-probabilities on real trajectories. The sampler’s own log-probabilities make the bug visible without any reference run.
Output:
With leaky packing, 5.8% of policy tokens have a trainer/sampler ratio more than 20% away from 1 before any update; with correct packing, none do. On a dashboard this looks like a training–inference mismatch (Section 13.6), and a mismatch correction would hide it. The step-0 ratio distribution is a cheap test for both.
Problem: wasted rollouts on zero-variance prompts. Symptom: the zero-advantage share of groups climbs as training succeeds; rollouts per useful gradient rise; the learning curve flattens long before the held-out pass rate stops improving. Cause: a group whose rewards are all equal has zero advantage under every baseline, so each of its episodes is a full rollout with no gradient (Chapter 9). Options: (a) pass-rate filtering before sampling (Section 14.5), which needs per-task estimates and can starve tasks whose estimate is stale. (b) Dynamic sampling, which keeps the batch full by sampling more. (c) Larger groups for hard tasks, so a rare success is sampled. Default: (a) with occasional revisits of excluded tasks, plus logging the zero-advantage share. The block counts the episodes that landed in zero-advantage groups in the runs of Section 14.5.
Output:
Across both seeds, 43% of the unfiltered arm’s 7,680 episodes landed in zero-advantage groups, 3,304 rollouts that produced no gradient. The filtered arm sampled 4,272 episodes in total and wasted 1,136 of them (27%). Most of the remaining waste comes from the first 20 steps, when the estimates are still uninformative, and from the revisits of excluded tasks, which are deliberately spent on tasks believed to be solved.
Problem: one long-tail sequence dominates memory. Symptom: out-of-memory errors at random steps, long after the run started, always on the rank that received the longest trajectory; lowering the batch size makes them rarer but not impossible. Cause: with a fixed number of sequences per microbatch, padded memory is the count times the longest member, and with a long-tailed length distribution the longest member eventually reaches the context limit. A packed or padded microbatch of 8 sequences then costs 8 times the context length. Options: (a) a token budget (Section 14.2), which puts the longest trajectories in microbatches of their own. (b) Chunked and recomputed log-probs (Section 14.3), which shrink the per-token cost but not its growth. (c) A hard length cap with truncation, which changes the task. Default: (a) and (b) together, with the budget set from a measured peak at the maximum context length. The block finds the worst microbatch under a fixed count of 8 and the microbatch the token budget gives the same trajectory.
Output:
The longest trajectory reached the 64k cap, eight times the median. In a fixed microbatch of 8, it forces 524,288 padded positions, the memory of 8 maximum-length trajectories. The token budget gives it a microbatch with one other trajectory of the same length, 131,072 tokens, the budget. The run that “suddenly runs out of memory at step 2,000” is the run in which this trajectory first appeared.

​
Summary

  • Padding TinyAgent’s 64 episodes to the longest wasted 74% of the computed positions. Packing filled rows to 97%, and it is correct only with a block-diagonal causal mask and position ids that restart per sequence: without them, policy log-probabilities moved by up to 6.4 nats, with them they matched to 10−610^{-6}.
  • A fixed number of sequences per microbatch makes memory proportional to the longest sequence; for a long-tailed length distribution, a token budget with length sorting cut the largest microbatch from 524k to 131k tokens at 3% padding.
  • Logits are the largest memory item of long-context RL at production vocabularies. Measured on Qwen3-0.6B’s output layer: 29.9 GB at 16k tokens in full, 11.2 GB chunked, 2.2 GB chunked with recomputation, independent of length.
  • Gradient accumulation with per-microbatch token_mean changes the gradient: 8% relative error for a random split and 52% for a length-sorted one on TinyAgent. Normalizing every microbatch by the global token count (mode='constant') reproduced the full-batch gradient to 10−710^{-7}.
  • Filtering tasks by their last group’s pass rate cut sampled episodes by 43% to 46% over 60 steps at a held-out pass rate within seed noise. Throughput arithmetic for an 8B agent run puts decoding at about 60% of the GPU-hours, because long-context decode is memory-bound at roughly a quarter of the training token rate.

​
Exercises

  1. Write a unit test for pack_sequences and packed_logits that fails for each of the two bugs of Section 14.1 separately. What is the smallest pair of sequences that exposes the position bug but not the mask bug?
  2. Extend token_budget_batches to packed microbatches, where the cost is the sum of lengths rather than the count times the longest. How many microbatches does the 512-trajectory example need at a 128k budget?
  3. Derive the per-token weight under per-microbatch token_mean for KK data-parallel ranks, each with its own microbatches. Show that normalizing by the global token count requires one all-reduce of a scalar before the backward pass, and implement it with aggregate_loss.
  4. Change an assumption: rerun Section 14.5 with the filter based on the mean of each task’s last three groups instead of its last group. Does the held-out gap close? How many more episodes does it cost?
  5. Implement sampling proportional to p^(1−p^)\hat p (1 - \hat p) instead of a hard filter, and compare the episodes spent and the held-out pass rate with the hard filter.
  6. Redo the GPU-hour estimate of Section 14.6 for a 64k-token context. Which term grows, and by how much? What would a 4x smaller KV cache (for example, from more aggressive grouped-query attention or an 8-bit KV cache) buy?
  7. Design review: a teammate proposes to raise the microbatch size from 4 to 16 sequences to improve training MFU. The run uses 32k-token agent trajectories with a long tail and per-microbatch loss normalization. List the three things that change besides MFU, with the measurement you would ask for on each.