Why does RL post-training need a separate rollout model at all?

Because generating completions and computing a gradient are different workloads that want different shapes. Training runs one large forward-and-backward pass over a whole batch and keeps the matrix units busy; decoding produces one token at a time and streams the entire model out of memory for each token, so it wants paged attention, continuous batching and a prefix cache instead. A model laid out for gradients has none of that. So a serious setup keeps a second copy of the same weights inside an inference engine shaped for decoding, and pays a per-batch cost to keep that copy current. The cost is real but bounded: only trainable parameters cross, never optimizer state, and under LoRA only the adapter parameters cross. In tunix that copy is a rollout worker, and the three shipped ones wrap the in-tree JAX sampler, vLLM through a TPU adapter, and SGLang-JAX.

What actually happens during a weight sync between a trainer and a sampler?

Four things, and only one of them is a copy. First the trainer state is filtered down to the parameters that matter, LoRA adapter leaves if the policy has adapters and plain parameters otherwise, which is what keeps Adam moment estimates out of the transfer. Second that filtered pytree is handed to the rollout worker, whose update method in all three backends is the single line that forwards to its sampler. Third the sampler places the arrays: the in-tree one reshards, casts to its own dtype, merges the new keys over the old ones and rebuilds the module from the graph definition it kept; the two foreign engines instead rename and reshape every leaf against a per-model mapping table and reshard once at the end. Fourth the cluster pins a host-memory snapshot of the actor weights, which is what old-policy log probabilities are later computed against, and increments the global step counter.

What does resharding a pytree between two device meshes actually do?

It reads the destination layout and asks JAX to place the source arrays there. Concretely the reshard function takes a source tree and a target tree, and the target is a template rather than a destination: every leaf is turned into a named sharding carrying three facts, a mesh of devices, a partition spec saying which mesh axis each array dimension is cut along, and a memory kind distinguishing device memory from pinned host memory. Those shardings are collected into a tree mirroring the target, an implementation is picked by trying Pathways first and falling back to jax.device_put, and the whole tree is handed over in one call. Nothing anywhere compares the source and destination layouts. The equal case is cheap for a mechanical reason: device_put onto the sharding an array already carries returns a new handle to the same buffer rather than copying it.

Is weight sync free when the trainer and sampler share a mesh?

Cheaper than free, because it usually does not run at all. The learner decides once in its constructor whether the actor and rollout models are sharing weights, and the test is object identity: it flattens both states and compares leaves with Python is, not by value and not by shape. That is only true when the cluster handed the same module instance to both roles, which it does when the two roles are assigned the same mesh. When the test passes, the learner skips the sync entirely at each mini-batch boundary and just advances the step counter the sync would have advanced. When the roles are split across separate meshes the identity test fails, the full path runs every batch, and for an 8B policy at bf16 that is about fifteen gibibytes crossing the interconnect per batch.

What happens when a rollout backend is handed weights it cannot place?

It depends which backend, and the two answers are opposite. The in-tree sampler wraps its reshard in a try block that catches AttributeError and ValueError, and the handler assigns the unresharded source weights and continues. Those are exactly the two failures that path produces: reading a sharding off a leaf that has none, and handing device_put a sharding tree that is not a prefix of the value tree. So an impossible sync downgrades silently to a merge of trainer-laid-out arrays, with no warning and no counter. The mapped backends refuse instead. A shape the aligner cannot satisfy raises a dedicated error saying padding and repetition are only supported for attention weights, and a model with no mapping table and no MaxText config is rejected before anything moves. If rewards go flat after a mesh change on the in-tree path, a silently skipped reshard is the first thing to rule out.

Rollout Backends and Weight Sync: Getting New Weights Into a Sampler

An RL loop trains one copy of the policy and samples from another, so every batch ends with a call that moves the updated weights across a device mesh, renames them for whichever inference engine is on the far side, and quietly does nothing at all when the two already agree. We read Google's tunix rollout backends and its reshard function at one pinned commit, then compute the reshard plan for a real Llama-3.1-8B parameter tree: which arrays cross, how each is cut on the other side, and what a batch boundary actually costs.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Rollout Backends and Weight Sync: Getting New Weights Into a Sampler — built to be understood, not skimmed.

Questions

Why does RL post-training need a separate rollout model at all?: Because generating completions and computing a gradient are different workloads that want different shapes. Training runs one large forward-and-backward pass over a whole batch and keeps the matrix units busy; decoding produces one token at a time and streams the entire model out of memory for each token, so it wants paged attention, continuous batching and a prefix cache instead. A model laid out for gradients has none of that. So a serious setup keeps a second copy of the same weights inside an inference engine shaped for decoding, and pays a per-batch cost to keep that copy current. The cost is real but bounded: only trainable parameters cross, never optimizer state, and under LoRA only the adapter parameters cross. In tunix that copy is a rollout worker, and the three shipped ones wrap the in-tree JAX sampler, vLLM through a TPU adapter, and SGLang-JAX.
What actually happens during a weight sync between a trainer and a sampler?: Four things, and only one of them is a copy. First the trainer state is filtered down to the parameters that matter, LoRA adapter leaves if the policy has adapters and plain parameters otherwise, which is what keeps Adam moment estimates out of the transfer. Second that filtered pytree is handed to the rollout worker, whose update method in all three backends is the single line that forwards to its sampler. Third the sampler places the arrays: the in-tree one reshards, casts to its own dtype, merges the new keys over the old ones and rebuilds the module from the graph definition it kept; the two foreign engines instead rename and reshape every leaf against a per-model mapping table and reshard once at the end. Fourth the cluster pins a host-memory snapshot of the actor weights, which is what old-policy log probabilities are later computed against, and increments the global step counter.
What does resharding a pytree between two device meshes actually do?: It reads the destination layout and asks JAX to place the source arrays there. Concretely the reshard function takes a source tree and a target tree, and the target is a template rather than a destination: every leaf is turned into a named sharding carrying three facts, a mesh of devices, a partition spec saying which mesh axis each array dimension is cut along, and a memory kind distinguishing device memory from pinned host memory. Those shardings are collected into a tree mirroring the target, an implementation is picked by trying Pathways first and falling back to jax.device_put, and the whole tree is handed over in one call. Nothing anywhere compares the source and destination layouts. The equal case is cheap for a mechanical reason: device_put onto the sharding an array already carries returns a new handle to the same buffer rather than copying it.
Is weight sync free when the trainer and sampler share a mesh?: Cheaper than free, because it usually does not run at all. The learner decides once in its constructor whether the actor and rollout models are sharing weights, and the test is object identity: it flattens both states and compares leaves with Python is, not by value and not by shape. That is only true when the cluster handed the same module instance to both roles, which it does when the two roles are assigned the same mesh. When the test passes, the learner skips the sync entirely at each mini-batch boundary and just advances the step counter the sync would have advanced. When the roles are split across separate meshes the identity test fails, the full path runs every batch, and for an 8B policy at bf16 that is about fifteen gibibytes crossing the interconnect per batch.
What happens when a rollout backend is handed weights it cannot place?: It depends which backend, and the two answers are opposite. The in-tree sampler wraps its reshard in a try block that catches AttributeError and ValueError, and the handler assigns the unresharded source weights and continues. Those are exactly the two failures that path produces: reading a sharding off a leaf that has none, and handing device_put a sharding tree that is not a prefix of the value tree. So an impossible sync downgrades silently to a merge of trainer-laid-out arrays, with no warning and no counter. The mapped backends refuse instead. A shape the aligner cannot satisfy raises a dedicated error saying padding and repetition are only supported for attention weights, and a model with no mapping table and no MaxText config is rejected before anything moves. If rewards go flat after a mesh change on the in-tree path, a silently skipped reshard is the first thing to rule out.

Rollout Backends and Weight Sync: Getting New Weights Into a Sampler

Questions

Related explainers