Tunix, From the Inside
Tunix is post-training in JAX: SFT, DPO, GRPO and PPO, distillation, and agentic RL, on Flax NNX models, with rollouts served by its own sampler or by vLLM and SGLang. The track starts with a recipe you run on a Colab TPU and ends inside the cluster object that wires meshes, trainers, and rollout workers together.
Phase 1 — Run it first
Two upstream notebooks that run on a Colab TPU, and the article on what a recipe is made of.
- GRPO on Gemma, upstream notebook ↗ — The reference GRPO run: a Gemma model, a math reward, and the cluster brought up in a notebook.
- QLoRA on Gemma, upstream notebook ↗ — The SFT side: a quantised LoRA fine-tune through the PEFT trainer.
- Running a Tunix Recipe: From a YAML File to an RLCluster — The CLI surface: the config merge order, recipes and reward functions, and every object a launch builds before the first step.
Phase 2 — The cluster
RLCluster owns the meshes, the trainer, the reference and reward models, and the rollout workers.
- The RL Cluster: Five Roles, One Mesh Dial — The design of an RL post-training cluster: what has to exist, where the bytes move, and where it breaks.
- rl_cluster.py, Line by Line — The orchestration hub read whole: one dictionary from role to mesh decides which models share a buffer and what a weight sync costs.
- tunix/rl/common.py ↗ — The shared types and helpers every RL component imports.
Phase 3 — The learner
The loop that turns completions into a training signal, and the GRPO learner that specialises it.
- GRPO Advantage: Z-Score Your Siblings, Line by Line — algo_core.py line by line: how GRPO scores a completion against its siblings.
- The RL Learner Loop, Line by Line — rl_learner.py read whole: six batch sizes, the queue between the producer and the trainer, and one micro-batch from prompts to an optimizer step.
- The GRPO Learner, Line by Line — The algorithm file read whole: four abstract slots filled, and the loss and advantage fetched by name from another module.
- tunix/rl/ppo/ppo_learner.py ↗ — The PPO sibling: a value function, GAE, and the clipped surrogate loss.
Phase 4 — Rollouts
Three ways to generate completions, and how the trainer’s weights reach the rollout backend after every update.
- The Tunix Sampler, Line by Line: Generation as Two Compiled Functions — generate/sampler.py read whole: the fixed token buffer, the donated KV cache, the jitted prefill, and the while_loop that decodes.
- Rollout Backends and Weight Sync: Getting New Weights Into a Sampler — Three rollout engines behind one interface, and the reshard that places trainer weights into each after every batch.
- tunix/generate/vllm_sampler.py ↗ — The vLLM-backed sampler and its weight-sync entry point.
- tunix/generate/sglang_jax_sampler.py ↗ — The SGLang-JAX-backed sampler.
Phase 5 — Agentic RL
Rollouts that call tools: an async producer of trajectories and the learner that consumes them.
- The Trajectory Collect Engine, Line by Line — trajectory_collect_engine.py read whole: one episode of model and tool call, the five ways it can end, and the queue that feeds a learner while rollouts keep running.
- tunix/rl/agentic/agentic_rl_learner.py ↗ — The learner that consumes trajectories instead of plain completions.
- tunix/rl/agentic/pipeline/rollout_orchestrator.py ↗ — The orchestrator that keeps producers and consumers in step.
Phase 6 — SFT and PEFT
The supervised side: the PEFT trainer, LoRA through qwix, and DPO.
- The Tunix PEFT Trainer, Line by Line — peft_trainer.py read whole: the type tag that decides the trainable set, the two update paths in one step, and the LoRA-only checkpoint.
- How Qwix Quantizes Any Flax Model Without Touching Its Code — How qwix intercepts ops to quantise without touching model code, which is how QLoRA enters Tunix.
- tunix/sft/dpo/dpo_trainer.py ↗ — DPO on top of the same trainer scaffolding.