Tunix

Post-training on JAX: an RL cluster over meshes, learners for GRPO and PPO, rollouts through the sampler or vLLM or SGLang, agentic trajectories with tool calls, and PEFT.

Explainers

• Running a Tunix Recipe: From a YAML File to an RLCluster — There is no tunix command. You run a module, hand it a YAML file as the first positional argument, and a few hundred lines of config code turn that file into a tokenizer, a dataset, a device mesh per role, three or four models, a rollout config, a cluster and a learner, in that order, before a single token is generated. This reads the whole launch surface at a pinned commit of google/tunix: the four sources config is merged from and the exact rule that makes setting one key by both an environment variable and the command line an error; the two base YAML files and the eleven keys that exist in only one of them; the five sections where a partial override silently deletes every sibling key you did not restate; the unknown-key check that catches a typo at the top level and cannot see one a single level down. A recipe turns out not to be a YAML file at all but a Python module exposing create_dataset, and a reward function is any module-level function in a file you name. The signature exhibit walks one real launch command from the file to the objects, with the arithmetic the code actually does at each step.
• 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.
• The RL Learner Loop, Line by Line — One file in google/tunix, 834 lines, read whole. It never touches a model and never computes a gradient. What it does is arithmetic on batch sizes and a handoff between two threads: a producer that pulls prompts, generates completions and computes advantages, and a trainer that pulls finished examples off a queue and hands each one to the optimizer. The page derives the six batch sizes from the config, walks the producer and the consumer region by region, and ends on a computed step timeline where you set the sizes and watch rollouts and gradient steps either overlap or take turns, depending on one line that compares two meshes.
• The GRPO Learner, Line by Line — Six hundred lines of google/tunix that add group-relative policy optimization to a training loop they do not contain, and that compute neither the advantage nor the loss the file is named for. Both arrive as function pointers, fetched from a process-global table by a string that lives on the config, filled by an import the file never mentions again. The page walks the whole file at a pinned commit: the config whose docstring documents a field it does not declare, the constructor that hands the actor trainer its loss, the one method that turns prompts into a differentiable record, the seventy lines that measure how far the rollout sampler has drifted from the trainer, and the four abstract slots that are the entire delta over the base learner. Then a group of four completions you score yourself, with the real estimator and the real clipped loss ported line for line, and the tied group that costs a full step and teaches nothing.
• The Tunix Sampler, Line by Line: Generation as Two Compiled Functions — Every RL post-training step needs completions, and in Tunix the in-tree sampler is what produces them. It is one file, and its whole design falls out of one constraint: JAX wants fixed shapes. So the prompt gets left-padded to a power of two, the output lives in a buffer sized before the first token is generated, the KV cache is allocated once and donated to the compiler, and the decode loop is a lax.while_loop that runs until the slowest sequence in the batch stops. We read tunix/generate/sampler.py top to bottom at one pinned commit, then compute the decode timeline and the cache filling for a batch you dial yourself.
• 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.
• The Trajectory Collect Engine, Line by Line — A normal RL rollout is one prompt in, one completion out. An agentic rollout is a conversation: the model writes a tool call, an environment runs it, the result comes back as a message, and the model writes again, for however many turns the task takes. This walk reads the 763-line file in Google’s tunix that owns exactly one of those episodes: how it counts steps, where the only deadline in the file actually sits, which tokens end up carrying gradient, and how dozens of these run at once and stream their finished trajectories to a learner that never stops training.
• The Tunix PEFT Trainer, Line by Line — A LoRA fine-tune freezes the base model by leaving it out. This walk reads tunix/sft/peft_trainer.py at one pinned commit: the type tag that decides which parameters an optimizer is even built over, the two update paths inside one train step, the gradient accumulator that carries a denominator, and the checkpoint that writes only the adapter.