Almost every write-up of RLHF is about the algorithm: the reward model, the policy-gradient objective, the KL leash back to the reference. That is not the part that makes it hard to run. On a cluster, post-training is a systems problem, and a strange one: it is the only large-scale training workload where the data does not exist until the model under training generates it, and where several distinct models must be resident inside a single optimizer step. Pretraining is one model chewing a static shard of tokens. A post-training loop is an inference service and a trainer welded together, exchanging weights and samples every step. This article walks that layout — who lives in memory, which phase owns the clock, how the halves are placed on hardware, and the seams where the pipeline quietly goes idle.

Post-training changes the shape of the job

A pretraining step is beautifully rigid. One set of weights, a fixed sequence length, a batch drawn from disk, a forward and backward pass, an optimizer update. Every step costs the same, every rank does the same work, and the whole thing is a dense-compute problem you can reason about with a roofline.

A reinforcement-learning post-training step breaks all of those assumptions at once. It begins by sampling completions from the current policy, which is autoregressive decoding, not batched training. Those completions have data-dependent lengths, so the batch reaching the trainer is ragged and its cost varies step to step. The samples must then be scored, which means extra models sweeping the same tokens. Only then does a familiar forward-backward-update happen: two execution regimes glued together, repeated thousands of times.

RLHF pipelineStage 1: SFTdemonstrationsStage 2: RMpreference pairsStage 3: PPORL vs rewardStage 3 is the systems problem: several models resident at once
The post-training stages. Stage 3 is where the GPU cost lives.
Advertisement

Four models, one memory budget

The classic PPO-style loop keeps up to four models in play: the policy being trained, a frozen reference copy used for the KL term, a reward model that scores completions, and a critic or value network. Their costs are wildly asymmetric. A trainable model under mixed-precision Adam runs about 16 bytes per parameter once you add the working copy, the gradient, the fp32 master weights and two optimizer moments. A frozen, forward-only model costs roughly 2 bytes per parameter plus activations.

So the reference and reward models are comparatively cheap lodgers; the critic is the expensive one, because it is trained and therefore carries the full optimizer tail. That is why critic-free variants are attractive from a hardware standpoint before you say anything about their statistics. The trainable models are then sharded exactly as they would be in pretraining — see gpu_zero_sharding and gpu_pipeline_parallel; nothing about RLHF changes those schemes.

The memory pretraining never has to pay: the KV cache

There is a fifth consumer with no analogue in pretraining at all. During sampling the policy is decoding autoregressively, and decoding needs a KV cache proportional to batch size times sequence length times layers times head dimension. For a large rollout batch with long generations that cache is not a rounding error — it can rival the weights themselves.

This is where the arithmetic gets uncomfortable. The generation peak and the training peak are different peaks, and if both live on the same devices the job is sized by whichever is larger, not by the average. A naive implementation that pre-allocates a rectangular cache for the maximum sequence length wastes most of it, because most generations stop early. Block-based cache management is the fix, the mechanism described in gpu_paged_kv_cache. Post-training is one of the workloads where borrowing a serving-side memory manager is not optional.

Generation and training are two different machines

The two phases stress opposite parts of the GPU. Training is dense and compute-bound: large GEMMs, tensor cores near their throughput limit, arithmetic intensity high enough to sit on the compute roof. Decoding is the opposite. Each token requires a full forward pass that reads the entire weight set to produce one row of activations per sequence, so it is bound by HBM bandwidth rather than FLOPs, and the tensor cores idle.

That asymmetry is why sampling so often dominates wall clock even though the trainer does far more arithmetic. A generation of length L costs L sequential forward passes that cannot be batched along the time axis, while the training pass over those same tokens is one parallel sweep. Optimizing the trainer first is usually optimizing the wrong half.

Co-location versus dedicated pools

Two placement strategies exist and the choice shapes everything downstream. In the co-located design all roles share the same GPUs and time-slice: the policy generates, the sampler is torn down, the same devices run the training step, and idle models are offloaded to host memory in between. No device sits unused, and weight sync is nearly free because both copies are local. The cost is memory pressure plus offload traffic across PCIe on every phase transition.

In the disaggregated design, each role gets its own pool: a generation cluster, a trainer cluster, perhaps a reward-model service. Each is sized and parallelised for its own workload and scales independently — you can throw GPUs at sampling without touching the trainer. The cost is that whichever pool is not the current bottleneck burns money idle, and weight sync now crosses the network.

Borrowing the inference engine for rollouts

The largest speedup available in a post-training loop is usually to stop generating with the training framework. A trainer’s built-in generate() typically decodes a padded rectangular batch, holds every sequence until the slowest finishes, and allocates cache for the worst case. A purpose-built engine — vLLM, TensorRT-LLM and the like — does iteration-level scheduling, admits new prompts as others retire, and manages the cache in blocks; see gpu_continuous_batching and gpu_llm_serving_architecture.

The catch is that you have now introduced a second runtime holding a second copy of the model, with its own parallelism layout, memory pool and possibly its own numeric precision. It is a good trade, but it converts a simple loop into a two-system distributed problem — and creates the hardest engineering issue in the pipeline.

Advertisement

The weight-sync problem

The policy changes every step, and the inference engine holds a copy of it. The inference copy is therefore stale the instant the optimizer runs, and must be refreshed before the next rollout — every step, forever. This is not a one-time load cost that amortises; it is a fixed per-step tax.

Worse, it is rarely a copy. The trainer holds parameters in its own layout: flattened shards across data-parallel ranks under ZeRO or FSDP, or split by tensor and pipeline dimensions. The inference engine holds them in its layout, usually a different tensor-parallel split. Refreshing weights therefore means gathering, reshaping and re-scattering — a resharding operation. Done well it is a set of NCCL broadcasts between peer ranks, or on co-located GPUs a pointer handoff via CUDA IPC with no copy at all. Done badly it gathers the full model to rank zero, stages it through host memory, and stalls every GPU in the job. If the rollout engine runs lower precision, add a cast on the same critical path.

Variable-length generations and load imbalance

Sampling produces sequences that stop whenever the model emits an end token. Within one rollout batch some completions finish in twenty tokens and some run to the cap. In a synchronous rectangular decode the batch then runs at the length of its longest member while nearly every lane sits masked — the classic tail, where a handful of live sequences occupy the whole device.

Spread that across data-parallel ranks and it becomes a collective problem: every rank waits at the next barrier for whichever rank drew the longest generations. The mitigations are scheduling, not kernels — iteration-level batching so retired slots are refilled, a hard cap on new tokens, grouping prompts by expected length, or over-generating and discarding the overflow. More aggressively, overlap the phases: generate the next batch while training on the current one. That hides the tail entirely, at the price of training on slightly stale samples — a real algorithmic consequence to decide deliberately.

The raggedness then follows the data into the trainer, where padding sequences into a rectangle burns tensor-core cycles on pad tokens and makes per-step activation memory data-dependent — a source of out-of-memory failures that surface randomly, thousands of steps in. Scoring adds a quiet multiplier on top: the reward model, the frozen reference and a recomputation of the policy’s own log-probabilities are three extra sweeps over the same tokens before any backward pass. Run serially they are dead time; batched together, or with the reward model as a concurrent service, most of it is reclaimed.

Checkpointing and restarting a multi-model job

Checkpointing here is not ‘save the model.’ A resumable state includes policy weights and optimizer state, critic weights and optimizer state, sampler and dataloader positions, and the RNG streams — the last matters far more than in pretraining, because generation is stochastic. The frozen reference and reward models never change, so reference them by identity rather than copying them into every checkpoint; doing otherwise multiplies checkpoint size for no benefit.

The subtle requirement is consistency across the two runtimes. A checkpoint must land on a clean phase boundary, or you restart with a rollout engine holding weights from one step and a trainer holding another. Checkpoint after the sync, never mid-rollout. And accept that bitwise-identical resumption is out of reach: sampling and ragged batching make the loop statistically reproducible at best.

Where the pipeline actually stalls

Instrument the loop by phase before optimizing anything and the same short list of culprits appears. The generation tail, where a few long sequences hold a whole device. The weight sync, where every GPU waits on a gather-and-broadcast written as an afterthought. The scoring passes, run serially when they could overlap. Offload traffic across PCIe on every phase transition in a co-located job. And the reward path itself, which is sometimes not a GPU model at all but a CPU-bound verifier, a sandboxed execution step, or an external API — in which case the whole cluster idles behind a single-threaded bottleneck.

The ordering follows from the phase split: fix generation throughput first, then weight sync, then the scoring passes, and only then the trainer. Most post-training jobs run at a fraction of their achievable throughput not because any kernel is slow, but because the GPUs spend much of each step waiting for a different phase to finish.

RLHF is an inference service and a trainer sharing a cluster, and almost every difficulty follows from that. Budget memory by role — roughly 16 bytes per parameter for anything trained, 2 for anything frozen, plus a rollout KV cache pretraining never pays. Expect the sampling phase to own the wall clock, because decoding is bandwidth-bound and sequential while training is compute-bound and parallel; profile the phase split before optimizing a kernel. Co-locate for cheap weight sync, dedicate pools for independent scaling. Drive rollouts with a real inference engine, then treat the weight refresh as a first-class resharding problem, not a copy. Schedule around variable-length generations rather than padding through them, and checkpoint at phase boundaries with RNG state included.