Tunix (Tune-in-JAX)

Tunix (Tune-in-JAX) is a JAX-based library designed for the efficient and scalable post-training of Large Language Models (LLMs). It is optimized for TPU performance and integrates with the JAX ecosystem (Flax, Optax, Orbax) and high-performance inference engines (vLLM, SGLang-JAX) for rollouts.

Core Capabilities:

• Supervised Fine-Tuning (SFT): Supports Full Weights Fine-Tuning, PEFT (Parameter-Efficient Fine-Tuning), DPO (Direct Preference Optimization), and ORPO (Odds Ratio Preference Optimization).
• Reinforcement Learning (RL): Supports PPO (Proximal Policy Optimization), GRPO (Group Relative Policy Optimization), GSPO-Token, DAPO, and Dr.GRPO.
• Agentic RL: Supports multi-turn tool use, asynchronous rollout for high-throughput trajectory collection, and trajectory batching/grouping.

Tunix (Tune-in-JAX) is a JAX-based library for efficient and scalable LLM post-training. It is designed to work within the JAX ecosystem (leveraging Flax, Optax, Orbax) and sits between core utilities and optimized models like MaxText.

Key capabilities include:

Supervised Fine-Tuning (SFT)

• Full Weights Fine-Tuning
• PEFT (Parameter-Efficient Fine-Tuning) such as LoRA
• Preference Optimization: DPO (Direct Preference Optimization) and ORPO (Odds Ratio Preference Optimization)

Reinforcement Learning (RL)

• PPO (Proximal Policy Optimization)
• GRPO (Group Relative Policy Optimization)
• GSPO-Token (Token-level Group Sequence Policy Optimization)
• DAPO (Direct Alignment via Preference Optimization)
• Dr.GRPO (Distributionally Robust GRPO)

Agentic RL

• Multi-turn tool use
• Asynchronous rollout for high-throughput trajectory collection
• Trajectory batching and grouping

Tunix provides two primary methods for launching training jobs:

1. Tunix CLI: The recommended default method. It is a simple launching tool that provides comprehensive configuration options through a command-line interface.
2. Interactive and Custom Launch: A method for users who require full control. This is best suited for:
   • Early experimentation: Learning the framework's capabilities.
   • Advanced customization: Implementing complex workflows that require more flexibility than the CLI offers.

The examples/ directory contains various notebooks and scripts for different fine-tuning paradigms:

Notebooks

• qlora_gemma.ipynb: LoRA and QLoRA fine-tuning with Gemma models.
• grpo_gemma.ipynb: Group Relative Policy Optimization (GRPO) reinforcement learning.
• dpo_gemma.ipynb: Direct Preference Optimization (DPO) for alignment.
• logit_distillation.ipynb: Knowledge distillation from teacher to student models.

Subdirectories

• deepscaler/: Training and evaluating DeepScaler models (train_deepscaler_nb.py, math_eval_nb.py).
• model_load/: Loading models from different formats (e.g., from_safetensor_load/ for Gemma2/Gemma3).
• rl/: Reinforcement learning implementations (e.g., grpo/gsm8k/ for mathematical reasoning).
• sft/: Supervised fine-tuning examples (e.g., mtnt/ for translation tasks).

Tunix provides built-in support for several state-of-the-art algorithms across Supervised Fine-Tuning (SFT), Preference Optimization, and Reinforcement Learning (RL).

Supervised Fine-Tuning (SFT) & Preference

• PEFT (Parameter-Efficient Fine-Tuning)
• DPO (Direct Preference Optimization)
• ORPO (Odds ratio Preference Optimization)

Reinforcement Learning (RL)

• PPO (Proximal Policy Optimization)
• GRPO (Group Relative Policy Optimization)
• GSPO-Token (Token-level Group Sequence Policy Optimization)
• DAPO (Direct Alignment via Preference Optimization)
• Dr.GRPO (Distributionally Robust GRPO)

Tunix (Tune-in-JAX) is organized into five distinct layers that abstract the complexity of LLM post-training:

1. User Interface & Application Layer: The entry point for users via CLI tools, interactive notebooks, and configuration files.
2. Algorithm & Workflow Layer: Contains core training paradigms like Supervised Fine-Tuning (SFT), Reinforcement Learning (RL), and Agentic workflows.
3. Core Components Layer: Handles the heavy lifting of training loop orchestration, metrics tracking, and distributed state management.
4. Foundation Frameworks & Integration Layer: Leverages JAX, Flax, and Optax for high-performance computation and optimization.
5. Hardware Layer: Manages compute resources, specifically optimized for Google Cloud TPUs, multi-host GPU clusters, and CPU hosts.

Tunix allows processes to discover each other dynamically via a gRPC discovery service.

• Server Side: A process acts as a discovery server by specifying --discovery_port. It can register callbacks using context.ipc.discovery.on_register(callback) to react when new peers join.
• Client Side: A process connects to an existing discovery server using --discovery_addrs=<hostname>:<port>. It can transmit metadata to the server using context.ipc.discovery.register(metadata=...).

# Server: Register a callback for new peers
context.ipc.discovery.on_register(
    lambda hostname, _, metadata: (
        logging.info(f"{hostname} knocked and said: {pickle.loads(metadata)}")
    )
)

# Client: Register with the discovery server
context.ipc.discovery.register(metadata=pickle.dumps(args.say))

Tunix uses a layered configuration system to launch SFT (Supervised Fine-Tuning) and RL (Reinforcement Learning) jobs. Configurations are merged in the following order of precedence (later sources override earlier ones):

1. Base Config File (Lowest Priority): The default settings defined in base_config.yaml.
2. Config File Override: An optional secondary YAML file specified via the override_config_file argument. This is useful for project-specific settings.
3. CLI Arguments (Highest Priority): Individual key=value pairs passed directly in the command line. These override all previous settings.

Tunix uses a GroupQueueManager within the RolloutOrchestrator to collect TrajectoryItem instances into groups. This is essential for algorithms like GRPO that require multiple samples per prompt.

• Grouping: Items are bucketed based on a group_key (e.g., a prompt ID from env.task["group_id"]) and episode_id.
• Ready Groups: Once a bucket reaches a specified group_size (e.g., num_generations), it is marked as a "ready group" and yielded via yield_batches.
• Memory Management: Use the max_open_buckets parameter to limit the number of groups being populated simultaneously to control memory usage.

Agentic RL enables training agents capable of multi-turn reasoning and external tool interaction.

Key Concepts:

• Multi-turn Conversations: Agents break complex problems into sequential steps of reasoning, tool invocation, and response generation.
• Tool Integration: Agents can use external tools (e.g., search, code execution, APIs). The environment manages the execution and feeds results back to the agent.
• Asynchronous Architecture: To maximize hardware utilization (like TPUs), Tunix uses a highly asynchronous architecture for both rollout and training. This overlaps model inference latency, I/O-bound tool execution, and reward computations.
• Trajectory Batching: The framework supports grouping trajectories, which is essential for algorithms like GRPO that require multiple samples per prompt.