What does RLCluster do in Tunix?

It is the single object an RL learner talks to, and its job is placement. It takes the models by role, actor and rollout and reference and reward and critic, plus one cluster config, and decides where each of them physically lives: which device mesh, in accelerator memory or pinned host memory, and whether two roles get their own copies of the weights or share one. After construction it exposes the operations a learner needs, generate for completions, per-token log-probabilities from the reference or the actor, an update for the actor, and a weight sync, and every one of those opens by entering the mesh belonging to the role it is about to use. The class is named RLEngine at this commit; RLCluster is a one-line alias on the last line of the file, so both names refer to the same object.

How does Tunix decide whether the actor and the rollout share one model?

With a single equality test between two mesh objects. A private method builds what the file calls the backbone sharing map, a dictionary from a role to the list of roles it may share weights with, and its first rule is that if the actor mesh and the rollout mesh are the same object, the two roles are linked in both directions. The constructor then reads that map: if the rollout is linked to the actor it assigns the actor model to the rollout attribute, one assignment and no second copy in memory. If it is not linked and the engine is the built-in vanilla sampler, the actor is loaded a second time onto the rollout mesh, optionally cast to a smaller dtype. Any other engine, vLLM or SGLang-JAX or your own class, is handed the trainer copy as an initial value and manages its own weights from there. A second rule links the reference model to the actor, but only when the actor is LoRA-tuned and the two meshes match, because that is the case where the frozen base weights can genuinely be shared with the adapters filtered out.

What is the anchor policy state in an RL cluster?

It is a snapshot of the actor parameters, pinned to host memory, taken once at the end of construction and retaken at every weight sync. It exists because of a timing problem: by the time a policy-gradient loss is computed, the actor has usually already taken optimizer steps, so scoring the sampled completions with the live weights would compare the policy against itself rather than against the policy that produced the tokens. The actor scoring method therefore splits the graph definition off the live trainer model and runs the forward pass over the anchor state instead, moving it into device memory first if it is still pinned to the host. A comment in the file names the two conditions that make this necessary rather than merely tidy: a mini-batch smaller than the full batch, or more than one inner iteration over the same data. The method guards against a missing snapshot and raises with a message telling you to sync first, though in practice the constructor has already seeded it, so that guard only fires if construction was cut short.

What happens on sync_weights when the actor and rollout are colocated?

The full sync still runs. The method does not consult the sharing map at all. It filters the actor state down to the LoRA parameters if the actor is LoRA-tuned, or to every parameter otherwise, and hands both the filtered state and the filter to the rollout engine, which reshards it onto whatever sharding the rollout copy uses. When the two roles are one model object on one mesh, the source and destination shardings are identical, so the reshard has nothing to move and the call costs almost nothing. When they are on different meshes, the identical call becomes a genuine transfer of every synced parameter across the interconnect. The mesh assignment alone decides the cost; the code path is the same either way. The method then re-pins the anchor snapshot and increments the global step counter, because a completed sync is what a finished step means here.

Can a rollout engine override the mesh set in the cluster config?

Yes, and the ordering is worth knowing. The sharing map is built early in the constructor from the meshes the config carries, before any engine exists. The rollout engine is constructed much later, and immediately afterwards the file checks whether that engine exposes a mesh of its own and, if it does, overwrites the rollout entry in the role-to-mesh mapping with it. The stated reason is that an engine may rearrange devices for better performance and the rest of the run should use the arrangement it chose. At this commit only the vLLM rollout exposes such a property; the vanilla sampler and the SGLang-JAX rollout do not, even though the latter is handed a mesh at construction. So a colocated run served by vLLM can end up with the actor and rollout registered as sharing a model while the recorded rollout mesh is one the engine picked.

rl_cluster.py, Line by Line

Reinforcement learning on a language model means running several copies of that model at once, each doing a different job: one samples completions, one scores them, one is a frozen snapshot you measure drift against, and one is taking the gradient steps. This file is the object that owns all five roles and decides where each of them physically lives. The answer is one dictionary from role to mesh, compared for object identity in exactly two places, and everything else follows: whether the sampler is a second copy of the policy or a second name for the same buffer, whether a weight sync moves bytes across the interconnect or nothing at all, and whether a LoRA reference model costs any memory. We read all 1,252 lines of the construction and step path at a pinned commit, then dial the three settings that decide placement and watch the file’s own branches fire.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of rl_cluster.py, Line by Line — built to be understood, not skimmed.

Questions

What does RLCluster do in Tunix?: It is the single object an RL learner talks to, and its job is placement. It takes the models by role, actor and rollout and reference and reward and critic, plus one cluster config, and decides where each of them physically lives: which device mesh, in accelerator memory or pinned host memory, and whether two roles get their own copies of the weights or share one. After construction it exposes the operations a learner needs, generate for completions, per-token log-probabilities from the reference or the actor, an update for the actor, and a weight sync, and every one of those opens by entering the mesh belonging to the role it is about to use. The class is named RLEngine at this commit; RLCluster is a one-line alias on the last line of the file, so both names refer to the same object.
How does Tunix decide whether the actor and the rollout share one model?: With a single equality test between two mesh objects. A private method builds what the file calls the backbone sharing map, a dictionary from a role to the list of roles it may share weights with, and its first rule is that if the actor mesh and the rollout mesh are the same object, the two roles are linked in both directions. The constructor then reads that map: if the rollout is linked to the actor it assigns the actor model to the rollout attribute, one assignment and no second copy in memory. If it is not linked and the engine is the built-in vanilla sampler, the actor is loaded a second time onto the rollout mesh, optionally cast to a smaller dtype. Any other engine, vLLM or SGLang-JAX or your own class, is handed the trainer copy as an initial value and manages its own weights from there. A second rule links the reference model to the actor, but only when the actor is LoRA-tuned and the two meshes match, because that is the case where the frozen base weights can genuinely be shared with the adapters filtered out.
What is the anchor policy state in an RL cluster?: It is a snapshot of the actor parameters, pinned to host memory, taken once at the end of construction and retaken at every weight sync. It exists because of a timing problem: by the time a policy-gradient loss is computed, the actor has usually already taken optimizer steps, so scoring the sampled completions with the live weights would compare the policy against itself rather than against the policy that produced the tokens. The actor scoring method therefore splits the graph definition off the live trainer model and runs the forward pass over the anchor state instead, moving it into device memory first if it is still pinned to the host. A comment in the file names the two conditions that make this necessary rather than merely tidy: a mini-batch smaller than the full batch, or more than one inner iteration over the same data. The method guards against a missing snapshot and raises with a message telling you to sync first, though in practice the constructor has already seeded it, so that guard only fires if construction was cut short.
What happens on sync_weights when the actor and rollout are colocated?: The full sync still runs. The method does not consult the sharing map at all. It filters the actor state down to the LoRA parameters if the actor is LoRA-tuned, or to every parameter otherwise, and hands both the filtered state and the filter to the rollout engine, which reshards it onto whatever sharding the rollout copy uses. When the two roles are one model object on one mesh, the source and destination shardings are identical, so the reshard has nothing to move and the call costs almost nothing. When they are on different meshes, the identical call becomes a genuine transfer of every synced parameter across the interconnect. The mesh assignment alone decides the cost; the code path is the same either way. The method then re-pins the anchor snapshot and increments the global step counter, because a completed sync is what a finished step means here.
Can a rollout engine override the mesh set in the cluster config?: Yes, and the ordering is worth knowing. The sharing map is built early in the constructor from the meshes the config carries, before any engine exists. The rollout engine is constructed much later, and immediately afterwards the file checks whether that engine exposes a mesh of its own and, if it does, overwrites the rollout entry in the role-to-mesh mapping with it. The stated reason is that an engine may rearrange devices for better performance and the rest of the run should use the arrangement it chose. At this commit only the vLLM rollout exposes such a property; the vanilla sampler and the SGLang-JAX rollout do not, even though the latter is handed a mesh at construction. So a colocated run served by vLLM can end up with the actor and rollout registered as sharing a model while the recorded rollout mesh is one the engine picked.

rl_cluster.py, Line by Line

Questions

Related explainers