6.7 Colocated vs Disaggregated RL & Weight Synchronization¶

Every reinforcement-learning-for-LLM system is, at heart, a tug-of-war between two very different programs fighting over the same scarce resource: GPUs. On one side is the trainer — a backward-pass-heavy, memory-hungry process that holds the policy weights, the optimizer state, gradients, and activations, and runs at full tensor-core utilization in short, dense bursts. On the other side is the rollout engine (also called the generator or actor) — an autoregressive decoder that is memory-bandwidth bound, runs for a long time per step, and spends most of its life waiting on the KV cache. These two workloads have nearly opposite hardware appetites, and how you place them on your cluster is the single most consequential infrastructure decision in an RL run.

This chapter is about that decision. We cover the two dominant placement strategies — colocated (trainer and rollout share GPUs) and disaggregated (they live in separate GPU pools) — and the family of tricks in between (time-slicing, offloading). Then we get to the hard part that ties both worlds together: weight synchronization. After every (or every few) optimizer steps, the freshly updated policy weights must reach the inference engine, which is laid out in a completely different parallelism scheme. We dissect the three transport mechanisms (NCCL broadcast, CUDA IPC, checkpoint reload), the resharding problem between training and inference layouts, and the latency/utilization math that decides which design wins.

This is the chapter where the anatomy of an RL system and the generation–training loop become a concrete engineering problem. We assume you know what GRPO and PPO are (GRPO, RLOO & Critic-Free RL, Policy Gradients & PPO for Language Models); here we care only about where the bytes live and how they move.

The Two Workloads and Why Placement Matters¶

Start with a single RL iteration. Ignoring the reward step, the loop is:

The rollout and the trainer touch the same model, but they want it in different forms. Let us be precise about the differences, because the entire chapter follows from them.

Memory profile. The trainer holds, for a model with \(P\) parameters trained in bf16 with the Adam optimizer:

That well-known \(\approx 16P\) bytes (plus activations) is why training a 7B model needs on the order of 112 GB of optimizer-related state. The rollout engine, by contrast, holds only weights and a KV cache:

A bf16 7B policy is 14 GB of weights; the rest of an 80 GB GPU is free for KV cache, which is exactly what an inference engine like vLLM or SGLang wants — more KV cache means larger rollout batches and higher decode throughput.

Compute profile. Training is compute bound: a forward+backward pass is \(\approx 6P\) FLOPs per token and saturates the tensor cores. Decoding is memory-bandwidth bound: each generated token reads the entire weight matrix from HBM but does only \(\approx 2P\) FLOPs, so the GPU’s arithmetic units idle while the memory bus is the bottleneck. (This is the prefill-vs-decode distinction from The Anatomy of LLM Inference, and the roofline picture from The Roofline Model & Performance Engineering.)

Temporal profile. Generation is long and bursty per sample but parallelizable across samples; a long chain-of-thought rollout can take tens of seconds. Training is short and dense: once the rollouts and advantages are ready, a handful of minibatch gradient steps finish in seconds.

These three mismatches are the crux. If you put both workloads on the same GPU at the same time, they contend for HBM and you must shrink both. If you put them on separate GPUs, one pool sits idle while the other works (unless you pipeline). The art of RL infrastructure is hiding that idle time.

Colocated: Sharing GPUs Between Training and Rollout¶

In a colocated design, the same physical GPUs run both the trainer and the rollout engine. This is the default in TRL, the classic configuration in veRL’s “hybrid engine,” and the simplest thing to reason about. The fundamental constraint is that an 80 GB GPU cannot simultaneously hold \(16P\) bytes of training state and \(2P\) bytes of inference weights and a useful KV cache. Something has to give, and the three strategies for making it give are time-slicing, offloading, and memory partitioning.

Time-slicing (the hybrid engine)¶

The cleanest colocated pattern runs the two phases sequentially on the same GPUs and swaps which one owns memory. During the rollout phase, the inference engine is “live”: weights are resident, KV cache is allocated, the trainer’s optimizer state and gradients are pushed out of the way (to CPU or freed). During the training phase, the inference engine releases its KV cache and the trainer reclaims the memory for activations and gradients.

veRL calls this the hybrid engine: a single set of workers that wears two hats, flipping between a training role (FSDP or Megatron) and a generation role (vLLM/SGLang) on command from the single controller. The key implementation detail is that the inference engine must support sleep/wake: it tears down and rebuilds its KV-cache memory pool around the training phase rather than holding it for the whole iteration.

# Sketch of the time-sliced colocated loop (single-controller view, veRL-style).
# Each "worker" actually holds BOTH an FSDP-wrapped trainer and a vLLM engine
# bound to the same CUDA devices. We flip between them every iteration.

for step in range(num_steps):
    # ---- PHASE 1: ROLLOUT ----------------------------------------------------
    # Wake the inference engine: allocate the KV-cache pool on the GPUs.
    rollout_engine.wake_up()                  # cudaMalloc the paged KV blocks
    # Trainer state is offloaded so the KV cache has room (see §offloading).
    trainer.offload_optimizer_to_cpu()        # async D2H copy of Adam m, v, master
    trainer.offload_grads()

    batch = sample_prompts(dataset, n=global_batch)
    # generate() runs continuous batching over all prompts (G samples each for GRPO)
    completions = rollout_engine.generate(batch, sampling_params)

    # ---- PHASE 2: REWARD -----------------------------------------------------
    rewards = reward_fn(batch.prompts, completions)        # may call a sandbox/verifier
    advantages = compute_group_advantages(rewards)         # GRPO: r_i - mean(r_group)

    # ---- PHASE 3: TRAIN ------------------------------------------------------
    rollout_engine.sleep()                    # free the KV-cache pool back to the allocator
    trainer.reload_optimizer_to_gpu()         # bring Adam state back (async H2D)

    for micro in make_microbatches(batch, completions, advantages):
        loss = trainer.compute_loss(micro)    # fwd + bwd under FSDP
        loss.backward()
    trainer.optimizer.step()
    trainer.optimizer.zero_grad()

    # ---- PHASE 4: WEIGHT SYNC ------------------------------------------------
    # Push the just-updated weights into the (sleeping) inference engine's
    # parameter buffers so the NEXT rollout uses the new policy. See §4.
    sync_weights(trainer, rollout_engine)

Those wake_up()/sleep() calls are not pseudocode — they are the real API surface. vLLM implements sleep mode: construct the engine with enable_sleep_mode=True, then llm.sleep(level=1) discards the paged KV-cache pool and offloads the weights to host RAM (level=2 discards the weights too, for when you will overwrite them anyway), and llm.wake_up() rebuilds the pool. SGLang’s equivalent requires --enable-memory-saver at launch and exposes release_memory_occupation() / resume_memory_occupation(), both as Python calls and as HTTP endpoints on the server.

from vllm import LLM, SamplingParams

# Colocated rollout engine: cap it at a fraction of HBM so the trainer has room.
llm = LLM(model="Qwen/Qwen2.5-7B", enable_sleep_mode=True,
          gpu_memory_utilization=0.35, tensor_parallel_size=2)

outs = llm.generate(prompts, SamplingParams(n=8, max_tokens=4096))
llm.sleep(level=1)   # KV pool freed + weights to CPU; HBM returns to the allocator
# ... trainer now runs fwd/bwd/optimizer.step() using the reclaimed memory ...
llm.wake_up()        # rebuild the KV pool for the next rollout phase

The frameworks wrap this. TRL runs vLLM inside the trainer process with GRPOConfig(use_vllm=True, vllm_mode="colocate", vllm_gpu_memory_utilization=0.3); veRL’s hybrid engine is actor_rollout_ref.hybrid_engine=True together with actor_rollout_ref.rollout.free_cache_engine=True (release the KV cache after generation) and actor_rollout_ref.actor.fsdp_config.param_offload / optimizer_offload for the trainer side; OpenRLHF exposes flags in the same spirit (--colocate_all_models, --vllm_enable_sleep). Flag names drift between releases — the mechanism they toggle is the one above.

The beauty of time-slicing is GPU utilization: every GPU is busy in every phase, because there is only one pool. The cost is latency — the phases are strictly serial, so the rollout GPUs are doing inference work while the (same) training capability sits idle, and vice versa. There is no overlap. For a run where rollout dominates wall-clock (long reasoning traces), this means your expensive training-capable cluster spends most of its time doing memory-bound decode at low FLOP utilization.

Offloading¶

Time-slicing only works if you can make room. The mechanism is offloading: moving tensors that are not needed in the current phase off the GPU. There are two flavors:

1. Optimizer-state offload (host RAM). Adam’s \(m\), \(v\), and fp32 master copy are \(12P\) bytes — the largest single consumer. During rollout they are dead weight, so we copy them to pinned CPU memory over PCIe/NVLink-C2C and free the GPU tensors. Before the optimizer step we copy them back. This is exactly ZeRO-Offload / FSDP CPU-offload (Memory-Efficient Training), reused here on a per-phase cadence.

2. Weight offload / KV-cache release. The inference engine’s KV pool (tens of GB) is released between rollouts. Inference weights may be kept resident (they are the same bytes the trainer needs) or, in the most aggressive colocation, the inference engine and trainer literally alias the same weight tensors — there is one copy of the parameters and a “view” of them for each role.

The offload transfer is not free. A 7B model’s Adam state is \(12 \times 7\text{e}9 = 84\) GB. Over a 64 GB/s PCIe 4.0 x16 link that is \(84 / 64 \approx 1.3\) s each way, or 2.6 s of pure copy per iteration. Over NVLink-C2C (Grace-Hopper, hundreds of GB/s) it is a fraction of that. This copy time is the hidden tax of colocated time-slicing, and it is why fast host interconnects (NVLink-C2C, large pinned buffers, double-buffered async copies) matter so much.

Memory partitioning (true simultaneous colocation)¶

The third option is to not time-slice at all: carve the GPU’s HBM into a training partition and an inference partition that coexist. This is rare in pure RL because the combined footprint rarely fits, but it appears in two guises: (a) very small models or heavily LoRA-fied policies where \(16P\) is tiny, and (b) multiplexing via MPS (NVIDIA Multi-Process Service) or MIG, where two processes share SMs and memory under a hardware scheduler. The win is true overlap; the risk is that the two contend for HBM bandwidth and both slow down. In practice, time-slicing + offload dominates for full-parameter RL; partitioning shows up mainly with LoRA-based RL where the policy delta is small. LoRA also collapses the sync problem: only the adapter matrices change, so you move tens of megabytes instead of gigabytes, and because vLLM and SGLang can load and swap LoRA adapters at runtime (vLLM’s LoRARequest / --enable-lora path), “weight sync” degenerates to an adapter reload. That is why LoRA RL is attractive on modest hardware even when full-parameter RL would fit.

Disaggregated: Separate Pools for Training and Rollout¶

The disaggregated design gives up on sharing. You provision two distinct GPU pools: a training cluster and an inference (rollout) cluster, each with its own optimal parallelism layout, connected by a network and a weight-sync channel. This is the design of OpenRLHF, the async architecture of Prime-RL, and the throughput-optimized mode of large-scale veRL deployments. It mirrors the disaggregated prefill/decode idea from serving — split a workload whose phases have different hardware profiles onto specialized hardware.

The two pools can be sized independently. If a run is rollout-bound (long generations, small policy updates), you give the rollout pool more GPUs — say 16 inference GPUs feeding 8 training GPUs. If it is training-bound (huge model, short completions), you flip the ratio. This independent scaling is the headline advantage of disaggregation and is impossible under strict colocation, where the ratio is forced to 1:1.

The second advantage is pipelining. Because the pools are separate, the rollout cluster can be generating the next batch of experience while the training cluster is still doing gradient steps on the current batch. This overlap is exactly what colocated time-slicing cannot do. But it comes with a subtle cost: the rollouts generated by the inference pool were produced by a policy that is now one step stale relative to the trainer. We have crossed from on-policy into off-policy territory.

On-policy vs off-policy: the staleness knob¶

If the trainer waits for the rollout pool to finish before each update, and the rollout pool always uses the latest weights, the system is synchronous and on-policy — exactly equivalent to colocated time-slicing, just on separate hardware (and with one pool idle while the other works). To get overlap you must let the rollout pool run ahead, producing experience under policy \(\pi_{\theta_{t-k}}\) while the trainer is at \(\theta_t\). The lag \(k\) is the staleness.

Staleness is a knob, not a free lunch. A small lag (\(k=1\)) is usually harmless and is corrected by the importance-sampling ratio already present in PPO/GRPO (\(\rho = \pi_\theta / \pi_{\theta_\text{old}}\)); the off-policy data is reweighted toward the current policy. Large lag degrades the gradient signal and can destabilize training. Managing this — importance correction, clipping, partial rollouts, and how much asynchrony is safe — is the subject of Prime-RL, Async RL & Decentralized Training and the stability tricks in Advantage Estimation, KL Control & Stability Tricks. For our purposes the point is structural: disaggregation buys overlap at the price of off-policyness, and the weight-sync frequency is what controls the staleness.

Cost of disaggregation¶

The disadvantages are real. First, idle GPUs under perfect balance are impossible — there is always slack on one side, and you pay for both pools whether or not they are saturated. Second, you now ship two large things across the network every step: the experience (rollouts) from inference to training, and the weights from training to inference. The experience is cheap (token IDs and logprobs, megabytes). The weights are not (gigabytes), and getting them across efficiently is the whole next section. Third, the system is operationally more complex: two fleets, a queue/orchestrator (usually Ray, see OpenRLHF, NeMo-Aligner & Ray-Based Systems), and a weight-sync protocol that must survive the resharding mismatch.

Weight Synchronization Mechanisms¶

After the optimizer step, the trainer holds the new policy \(\theta_{t+1}\). The rollout engine still holds \(\theta_t\). We must move the new weights into the inference engine’s parameter buffers before (or for async, soon after) the next rollout. There are three transport mechanisms, in increasing order of speed and decreasing order of robustness.

Mechanism 1: Checkpoint reload (slow, simple, robust)¶

The trainer writes a full checkpoint to shared storage; the inference engine reloads it. This is the “obviously correct” baseline and the only mechanism that trivially handles any layout mismatch, because the checkpoint is a layout-independent state_dict and the inference engine loads it with its own sharding logic — the same path it uses at startup.

# Mechanism 1: checkpoint reload. Correct, slow, layout-agnostic.
def sync_via_checkpoint(trainer, rollout_engines, path):
    # 1. Gather the FULL (unsharded) fp32/bf16 state dict on rank 0.
    #    Under FSDP this triggers an all-gather of every shard.
    full_state = trainer.get_full_state_dict()          # heavy: materializes 2P bytes
    if trainer.is_rank0:
        torch.save(full_state, path)                    # write 2P bytes to disk/NFS
    barrier()
    # 2. Every inference engine reloads from disk and re-shards to its own TP layout.
    for engine in rollout_engines:
        engine.load_weights(path)                       # read 2P bytes, reshard

For a 7B model this is 14 GB written and re-read; on NFS that can be tens of seconds. For a 70B model (140 GB) it is minutes. Checkpoint reload is fine for occasional sync (e.g., DPO-style offline data refresh, or recovery) but far too slow to run every RL step. Its one irreplaceable virtue: it is the fallback that always works, including across nodes that share nothing but a filesystem, and it doubles as your fault-tolerance checkpoint.

Mechanism 2: NCCL broadcast (fast, in-band, the workhorse)¶

The standard fast path is to keep the weights on the GPU the entire time and move them GPU-to-GPU over NVLink/InfiniBand using NCCL collectives (Parallel Computing & Collective Communication). The trainer ranks that own each parameter shard broadcast it directly into the inference ranks’ parameter tensors, never touching disk or host memory.

The trick is that the trainer and inference processes are different process groups. To make NCCL move a tensor from one to the other you build a shared communicator that spans both — a process group whose members are (training rank 0 …) ∪ (inference rank 0 …). NCCL needs a rendezvous: rank 0 creates a unique ID, distributes it (over Ray, a file, or a socket), and every participant calls init_process_group / ncclCommInitRank with it.

# Mechanism 2: NCCL broadcast over a process group that SPANS trainer + inference.
# Run on every rank in the combined group. We assume the param NAMES match and
# we send one parameter at a time, broadcasting from the trainer rank that owns it.

import torch
import torch.distributed as dist

class WeightSyncGroup:
    def __init__(self, rank, world_size, master_addr, master_port):
        # ONE communicator that includes BOTH the training ranks and the
        # inference ranks. The trainer is the "source" of every broadcast.
        dist.init_process_group(
            backend="nccl",
            init_method=f"tcp://{master_addr}:{master_port}",
            rank=rank,
            world_size=world_size,   # = num_train_ranks + num_infer_ranks
        )
        self.src_rank = 0            # trainer rank 0 is the broadcast root

    def push(self, named_params):
        # Called on the TRAINER. Broadcasts each weight to all inference ranks.
        for name, tensor in named_params:
            # tensor must be a full (gathered) parameter in the layout the
            # inference engine expects. Resharding (§5) happens BEFORE this call.
            dist.broadcast(tensor, src=self.src_rank)

    def pull(self, param_buffers):
        # Called on each INFERENCE rank. Receives into its parameter buffers.
        for name, buf in param_buffers.items():
            dist.broadcast(buf, src=self.src_rank)     # buf filled in place

In real systems this side channel is wrapped so the trainer makes one call per push. In vLLM you supply a worker extension (LLM(..., worker_extension_cls="my_rl_utils.WeightSyncExtension")) whose methods become invocable on every worker through llm.collective_rpc("init_weight_update_group", args=...) and llm.collective_rpc("update_weight", args=(name, dtype, shape)); each worker receives the broadcast into a scratch tensor and hands it to self.model_runner.model.load_weights([(name, tensor)]). SGLang exposes the same protocol as first-class endpoints — init_weights_update_group followed by update_weights_from_distributed(names, dtypes, shapes, group_name), with update_weights_from_tensor for the same-process case and update_weights_from_disk as the checkpoint fallback. OpenRLHF’s vllm_engine.update_weight and veRL’s resharding manager are thin layers over these primitives.

The transfer runs at NVLink/IB bandwidth: 14 GB of bf16 weights over a 900 GB/s NVLink 4 fabric (H100) is \(14 / 900 \approx 16\) ms of ideal transfer time. A real broadcast tree achieves a fraction of peak, so budget tens of milliseconds — still three orders of magnitude faster than an NFS reload. On Blackwell (B200/GB200), NVLink 5 roughly doubles the fabric to ~1.8 TB/s per GPU, halving the already-negligible sync cost. This is why NCCL broadcast is the default in-band sync for every serious RL framework.

Two practical wrinkles. First, bf16 vs fp32: broadcast in the dtype the inference engine consumes (usually bf16), halving the bytes versus fp32 master weights. Second, bucketing: broadcasting thousands of tiny tensors one at a time is latency-bound by kernel-launch and handshake overhead; frameworks flatten parameters into large contiguous buckets and broadcast a few big buffers instead, which is bandwidth-bound and far faster.

Mechanism 3: CUDA IPC (zero-copy, same-node, fastest)¶

When the trainer and inference engine are different processes on the same node sharing the same GPUs (the colocated case), you can avoid copying the weights at all. CUDA IPC (Inter-Process Communication) lets one process export a handle to a GPU allocation and another process map that exact device memory into its own address space. No bytes move over any link — the inference engine simply reads from the trainer’s weight tensors.

# Mechanism 3: CUDA IPC handle export. The trainer process exports a handle to
# its weight tensor's device memory; the inference process imports it and gets a
# tensor backed by the SAME physical GPU memory. Zero copy.

# ---- TRAINER PROCESS ----
import torch
from torch.multiprocessing.reductions import reduce_tensor

def export_handles(named_params):
    handles = {}
    for name, tensor in named_params:
        # reduce_tensor produces a picklable IPC handle (device ptr + metadata).
        # Send `handles` to the inference process over a pipe / Ray object store.
        handles[name] = reduce_tensor(tensor.detach())
    return handles

# ---- INFERENCE PROCESS ----
def import_handles(handles):
    mapped = {}
    for name, handle in handles.items():
        rebuild_fn, args = handle
        # Rebuilds a torch.Tensor whose storage points at the trainer's memory.
        mapped[name] = rebuild_fn(*args)     # NO copy; aliases trainer's weights
    return mapped

In the most tightly colocated designs the inference engine does not even re-import after every step — the parameter tensors are aliased once at startup, and because the trainer writes its updates in place (optimizer.step() mutates the same param.data), the inference engine automatically “sees” the new weights with literally zero sync work. The catch is correctness: you must ensure the trainer’s optimizer step has fully completed (CUDA stream synchronized) before the inference engine reads, and you must guarantee the two are not writing/reading the buffer concurrently. CUDA IPC is the fastest possible sync — effectively free — but only available on a single node and requiring that both processes can map the same memory.

The decision tree:

import torch

# `WeightLoader` is OUR class, passed as LLM(..., worker_extension_cls="my_rl.WeightLoader").
# Its methods run INSIDE each vLLM worker and are reachable via collective_rpc.
class WeightLoader:
    def load_from_trainer(self, named_tensors):
        # vLLM's per-architecture load_weights() already knows this model's
        # fusion rules (qkv_proj, gate_up_proj) and TP split, so it does the
        # fuse + reshard for us -- PROVIDED the names are HF-style.
        self.model_runner.model.load_weights(named_tensors)

def sync_in_process(model, llm):
    torch.cuda.synchronize()                    # the optimizer step must be complete
    named = [(k, v) for k, v in model.state_dict().items()]
    llm.collective_rpc("load_from_trainer", args=(named,))

The Resharding Problem: Training vs Inference Layouts¶

Here is the part that makes weight sync genuinely hard rather than a memcpy. The trainer and the inference engine almost never store the parameters the same way. They differ along three axes, and the sync layer must reconcile all three.

Axis 1 — Sharding (parallelism layout). The trainer might use FSDP, which flattens every parameter and shards it into equal slices across all data-parallel ranks; the byte at offset \(i\) on rank \(r\) has no semantic meaning on its own. The inference engine uses tensor parallelism: it splits attention heads and MLP columns across its TP ranks in a structured, head-aligned way. These two shardings are incompatible. Even if both used tensor parallelism, the degree differs: the trainer might run TP=8 (to fit optimizer state) while inference runs TP=2 (to minimize decode latency and maximize KV cache). Going from TP=8 to TP=2 means each inference rank must gather the slices of four training ranks.

Axis 2 — Fusion. Inference engines fuse operations for kernel efficiency. The three separate projections \(W_Q, W_K, W_V\) that the trainer stores are concatenated into one qkv_proj tensor in vLLM; gate_proj and up_proj become one gate_up_proj. Sync must concatenate (or split) accordingly. With GQA/MQA (Multi-Head Attention, MQA, GQA & MLA) the concatenation is not even uniform — K and V have fewer heads than Q, so the fused tensor is laid out in a specific interleaved order that the sync code must replicate exactly.

Axis 3 — Dtype and quantization. The trainer holds bf16 (and fp32 master); the inference engine might run quantized weights — FP8, INT8, or INT4 (Quantization II). Sync then includes an on-the-fly quantization step: take the bf16 master, compute scales, pack to INT4, and write into the engine’s quantized buffers. This is expensive and changes the sync from “copy bytes” to “recompute the quantized representation,” which is one reason RL on quantized inference engines is fiddly.

The general recipe a correct resharding sync layer follows:

# Resharding weight sync: from a sharded trainer layout to a fused, re-sharded
# inference layout. Conceptual; real impls (veRL, OpenRLHF) are more optimized.

def reshard_and_sync(trainer, infer_engines, name_map, infer_tp):
    """
    name_map: tells us how trainer param names map to inference param names,
              including which trainer tensors fuse into one inference tensor.
    infer_tp: tensor-parallel degree of the inference engine.
    """
    for infer_name, spec in name_map.items():
        # 1) GATHER: materialize the FULL (unsharded) trainer tensor(s).
        #    Under FSDP this all-gathers the flat shards and reshapes to 2D.
        full_tensors = [trainer.gather_full_param(tn) for tn in spec.trainer_names]

        # 2) FUSE: concatenate Q/K/V (or gate/up) into the engine's fused tensor,
        #    honoring GQA head grouping if K/V have fewer heads than Q.
        fused = fuse_along(full_tensors, dim=spec.fuse_dim, head_layout=spec.heads)

        # 3) (optional) QUANTIZE: if the engine runs FP8/INT4, compute scales & pack.
        if spec.quantized:
            fused = quantize(fused, scheme=spec.quant_scheme)

        # 4) RE-SHARD: split the fused tensor into `infer_tp` slices, one per
        #    inference TP rank, along the engine's tensor-parallel dimension.
        shards = split_for_tp(fused, tp=infer_tp, dim=spec.tp_dim)

        # 5) SEND: broadcast/scatter each shard to its inference rank (Mechanism 2),
        #    or copy in place if same-node aliasing (Mechanism 3).
        scatter_to_infer_ranks(shards, infer_name, infer_engines)

Steps 1 and 2 are the expensive ones: gathering the full parameter undoes the trainer’s sharding (an all-gather of \(2P\) bytes), and you need transient memory to hold the full fused tensor. Good implementations do this layer by layer, streaming one fused tensor at a time so peak extra memory is one layer’s worth, not the whole model. veRL’s resharding manager and OpenRLHF’s update_weight path both implement a version of this name-mapped, layer-streamed gather-fuse-reshard-broadcast pipeline. The single-controller architecture of veRL is in part designed to make this resharding step expressible as one orchestrated collective rather than ad-hoc message passing.

Latency, Utilization, and a Worked Example¶

We now have all the pieces to reason quantitatively. Define one RL iteration’s wall-clock time. For a synchronous system (colocated time-sliced, or synchronous disaggregated):

For an asynchronous disaggregated system with perfect pipelining, rollout overlaps training, so the iteration is bounded by the slower stage:

The GPU utilization of the colocated synchronous design, measured as useful-work fraction on the training-capable GPUs, is roughly the share of time spent in compute-bound work:

When rollout dominates (long reasoning traces), \(U_\text{colo}\) is small — your training cluster spends most of its life doing memory-bound decode. That is the central argument for disaggregation: move the long, cheap rollout onto cheaper/more-numerous inference GPUs and keep the expensive training GPUs busy with training.

The example also exposes the staleness cost we hid: in the async case, the rollouts feeding each update were generated under weights one iteration old. If that staleness hurts sample efficiency by, say, 10%, you must weigh it against the utilization gain. This is the perennial RL-infra trade and why the choice is workload-dependent rather than universal.

Putting It Together: A Decision Framework¶

There is no universally best design; there is a set of questions whose answers pick one for you.

Is the run rollout-bound or training-bound? Profile \(T_\text{rollout}\) vs \(T_\text{train}\). Long generations (reasoning, agentic multi-turn rollouts, see Agentic & Multi-Turn RL) are rollout-bound and favor disaggregation, where the cheap rollout pool can be scaled independently and the training GPUs stay saturated. Short completions with a huge model are training-bound and tolerate colocation, since rollout is a small slice of the iteration.

How much hardware do you have? On a single node (one machine, ≤8 GPUs), colocated time-slicing with CUDA-IPC or aliased weights is simplest and the sync cost is essentially zero. Disaggregation makes sense at multi-node scale where you can afford to dedicate fleets and the network supports fast NCCL sync.

How tolerant is your algorithm of staleness? Strictly on-policy methods, or runs where you have seen instability from off-policy data, push you toward synchronous designs (colocated, or synchronous disaggregated with the rollout pool idling during updates). If your importance correction and KL control (Advantage Estimation, KL Control & Stability Tricks) are solid, async disaggregation unlocks far better utilization.

What is your weight-sync budget? If sync is cheap (same-node IPC, or fast NVLink/IB NCCL), you can sync every step and keep staleness at 1. If sync is expensive (cross-cluster, must reload checkpoints), you sync rarely, which forces higher staleness — the sync mechanism and the algorithm’s off-policyness are coupled.

The frameworks map onto this flow. TRL is colocated by default (with an optional vLLM server that nudges toward disaggregation). veRL supports both: the hybrid engine for colocation and a disaggregated mode, with a single controller orchestrating the resharding sync. OpenRLHF is Ray-native and leans disaggregated, with vLLM rollout workers receiving NCCL-broadcast weights. Prime-RL takes async disaggregation to its conclusion with deliberately decoupled, possibly geographically separated, trainer and rollout fleets. The scaling tricks that make the rollout pool itself efficient — load balancing across uneven-length generations, in-flight weight updates — are the subject of Scaling RL: Throughput, Load Balancing & The Latest Tricks.

Further Reading¶

• Sheng et al., “HybridFlow: A Flexible and Efficient RLHF Framework” (veRL, 2024) — the single-controller architecture, the hybrid engine for colocation, and the resharding manager that reconciles training and inference layouts; the most directly relevant systems paper for this chapter.
• Hu et al., “OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework” — Ray-based disaggregated rollout workers with vLLM and NCCL weight broadcast; a clear reference implementation of the disaggregated pattern.
• Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training” — the optimizer-state-offload mechanism reused on a per-phase cadence in colocated time-slicing.
• Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention” (vLLM, 2023) — the rollout engine and its KV-cache memory model that determine what colocation must make room for.
• Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” and Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” — the two training-side layouts (sharded data parallel and tensor/pipeline parallel) whose mismatch with inference TP creates the resharding problem.
• NVIDIA NCCL documentation and the CUDA IPC runtime API — the collective-communication and zero-copy-handle primitives underlying weight-sync Mechanisms 2 and 3.
• DeepSeek-AI, “DeepSeek-R1” (2025) — a large-scale RLVR run whose long reasoning rollouts are the canonical rollout-bound workload motivating disaggregation.

Exercises¶

1. (Memory profile.) You want to run full-parameter GRPO on a 7B-parameter policy in bf16 with the Adam optimizer, and serve rollouts from the same 80 GB GPUs. Using the chapter’s formulas, compute (a) the trainer state \(M_\text{train}\) ignoring activations, and (b) the inference weight footprint (ignoring KV cache). Then explain, in one or two sentences, why these two numbers make simultaneous colocation impossible on an 80 GB GPU and force a time-slicing + offload strategy.

2. (Offload cost.) During the rollout phase of the time-sliced colocated loop, the 7B model’s Adam state is offloaded to host RAM and later reloaded before the optimizer step. (a) How many bytes is the Adam state (\(m\), \(v\), fp32 master)? (b) Compute the round-trip (out + in) copy time over a PCIe 4.0 x16 link at 64 GB/s. © Recompute the round-trip over an NVLink-C2C link at 450 GB/s, and state the speedup. (d) Why does the chapter insist on pinned host memory and a dedicated CUDA stream for these copies?

3. (Utilization and rebalancing.) A synchronous 7B run on 16 H100s measures these per-iteration stage times: \(T_\text{rollout} = 50\) s, \(T_\text{reward} = 2\) s, \(T_\text{train} = 6\) s, \(T_\text{offload} = 3\) s, \(T_\text{sync} = 0.02\) s. (a) Compute the colocated iteration time and the training-GPU utilization \(U_\text{colo}\). (b) You switch to async disaggregation, splitting the 16 GPUs into \(g_t\) training and \(g_r = 16 - g_t\) rollout GPUs. Assuming stage time scales inversely with GPU count, find the smallest \(g_t\) (integer) such that training still hides fully under rollout, and compute the resulting async iteration time. © Contrast what happened to training-GPU utilization.

4. (Sync mechanism choice + resharding.) (a) For a 70B model (140 GB in bf16), compute the weight-sync time for NCCL broadcast over a 900 GB/s NVLink fabric, and for checkpoint reload over an NFS mount sustaining 1.2 GB/s. State the ratio. (b) The chapter warns that even with a fast NCCL path, sync is “not a memcpy.” Name the three axes along which trainer and inference layouts differ, and explain concretely what a trainer running FSDP + TP=8 must do to feed an inference engine running TP=2.

5. (Implementation.) Implement the fuse + re-shard step (steps 2 and 4 of the chapter’s reshard_and_sync recipe) for a GQA attention block. Write fuse_qkv(w_q, w_k, w_v) that concatenates the three full projection weights into one qkv_proj tensor in [Q; K; V] output-row order, and shard_qkv_for_tp(w_q, w_k, w_v, tp) that produces the per-rank fused shard for each of tp inference ranks such that each rank holds a contiguous, head-aligned block of Q, K, and V heads (a naive torch.chunk of the already-fused tensor would split heads incorrectly). Verify on hidden=4096, n_heads=32, n_kv_heads=8, head_dim=128, tp=2.

import torch

def fuse_qkv(w_q, w_k, w_v):
    # Full (TP=1) fused layout: rows are [all Q heads; all K heads; all V heads].
    # w_q: [n_heads*d, hidden], w_k/w_v: [n_kv_heads*d, hidden].
    return torch.cat([w_q, w_k, w_v], dim=0)   # [(n_heads + 2*n_kv_heads)*d, hidden]

def shard_qkv_for_tp(w_q, w_k, w_v, tp, n_heads, n_kv_heads, head_dim):
    # GQA constraint: heads must divide evenly across TP ranks.
    assert n_heads % tp == 0 and n_kv_heads % tp == 0, "heads must be divisible by tp"
    q_per = (n_heads   // tp) * head_dim     # Q rows per rank
    kv_per = (n_kv_heads // tp) * head_dim   # K (and V) rows per rank
    shards = []
    for r in range(tp):
        # Split EACH projection into contiguous head groups, THEN fuse per rank.
        q_r = w_q[r * q_per : (r + 1) * q_per]
        k_r = w_k[r * kv_per : (r + 1) * kv_per]
        v_r = w_v[r * kv_per : (r + 1) * kv_per]
        shards.append(torch.cat([q_r, k_r, v_r], dim=0))  # rank-local qkv_proj
    return shards

# ---- Verify ----
hidden, n_heads, n_kv_heads, head_dim, tp = 4096, 32, 8, 128, 2
w_q = torch.randn(n_heads   * head_dim, hidden)   # [4096, 4096]
w_k = torch.randn(n_kv_heads * head_dim, hidden)  # [1024, 4096]
w_v = torch.randn(n_kv_heads * head_dim, hidden)  # [1024, 4096]

full = fuse_qkv(w_q, w_k, w_v)
assert full.shape == ((n_heads + 2 * n_kv_heads) * head_dim, hidden)  # [6144, 4096]

shards = shard_qkv_for_tp(w_q, w_k, w_v, tp, n_heads, n_kv_heads, head_dim)
# Each rank: (16 Q + 4 K + 4 V) heads * 128 = 3072 rows.
assert all(s.shape == ((n_heads//tp + 2*(n_kv_heads//tp)) * head_dim, hidden)
           for s in shards)                        # each [3072, 4096]

# Correctness: rank 0 holds the first half of Q/K/V heads, exactly.
assert torch.equal(shards[0][:16*head_dim],                  w_q[:16*head_dim])
assert torch.equal(shards[0][16*head_dim:16*head_dim+4*head_dim], w_k[:4*head_dim])
print("fused:", tuple(full.shape), "| per-rank:", tuple(shards[0].shape))