```
├── .gemini/
├── styleguide.md (1000 tokens)
├── .github/
├── CODEOWNERS
├── ISSUE_TEMPLATE/
├── bug_report.md (100 tokens)
├── feature_request.md (100 tokens)
├── pull_request_template.md (200 tokens)
├── workflows/
├── build_and_test_tunix.yml (600 tokens)
├── build_and_test_tunix_nightly_regression.yml (400 tokens)
├── build_docs.yml (200 tokens)
├── build_package.yml (300 tokens)
├── cpu-tests.yml (1500 tokens)
├── issue-auto-assign.js (400 tokens)
├── issue-auto-assign.yaml (100 tokens)
├── pypi_release.yml (600 tokens)
├── tpu-nightly-regression.yml (3.7k tokens)
├── tpu-tests.yml (2.6k tokens)
├── .gitignore (100 tokens)
├── .pre-commit-config.yaml (300 tokens)
├── .pylintrc (3k tokens)
├── .pylintrc-notebooks (100 tokens)
├── Dockerfile (1000 tokens)
├── LICENSE (omitted)
├── README.md (1600 tokens)
├── build_docker.sh (700 tokens)
├── docs/
├── _static/
├── custom.css
├── img/
├── sphx_glr_local_module_thumb.png
├── tunix.png
├── agentic_rl.md (1200 tokens)
├── algorithms.md (1600 tokens)
├── api/
├── api_distillation.rst
├── api_generation.rst
├── api_rl.rst (100 tokens)
├── api_sft.rst (100 tokens)
├── code-of-conduct.md (900 tokens)
├── conf.py (800 tokens)
├── contributing.md (1400 tokens)
├── design.md (1600 tokens)
├── examples.md (700 tokens)
├── gallery.rst (300 tokens)
├── guides/
├── Makefile (200 tokens)
├── images/
├── agentic_agent:env.png
├── agentic_rollout_pipeline.png
├── batch_vs_async_rollout.png
├── collocated_mode.png
├── disagged_mode.png
├── global_batch_flow.png
├── mini_batch_flow.png
├── model_id_huggingface.png
├── model_id_kaggle.png
├── perf_metrics_TB.png
├── perf_metrics_perfetto.png
├── tunix_architecture.png
├── tunix_in_jax_ecosystem.png
├── tunix_reshard_algorithm.png
├── tunix_rl_pipeline.png
├── tunix_sft_pipeline.png
├── index.md (100 tokens)
├── intro.md (1100 tokens)
├── launching.md (3.9k tokens)
├── make.bat (200 tokens)
├── metrics.md (4.6k tokens)
├── models.md (2.1k tokens)
├── performance.md (2.1k tokens)
├── quickstart.md (2.8k tokens)
├── reliability.md (900 tokens)
├── rollout.md (4.6k tokens)
├── talks.md (400 tokens)
├── examples/
├── README.rst (1000 tokens)
├── deepscaler/
├── math_eval_nb.py (4.8k tokens)
├── run_deepscaler_disagg_v5p16.sh (1500 tokens)
├── train_deepscaler_nb.py (4k tokens)
├── deepswe/
├── action_guard.py (3k tokens)
├── deepswe_data.py (600 tokens)
├── eval_deepswe.py (4.6k tokens)
├── guarded_swe_env.py (300 tokens)
├── openhands_utils.py (1200 tokens)
├── r2e_gym_helper.py (500 tokens)
├── run_deepswe_disagg_v5p_32.sh (1600 tokens)
├── swe_agent.py (1300 tokens)
├── swe_env.py (4.6k tokens)
├── system_prompts.py (3.7k tokens)
├── template.py (5.6k tokens)
├── train_deepswe_nb.py (7.3k tokens)
├── train_maxtext_nb.py (8.1k tokens)
├── dpo_gemma.ipynb (1932.3k tokens)
├── frozenlake/
├── agent.py (1600 tokens)
├── configs/
├── gemma4_e2b.yaml (700 tokens)
├── data.py (1200 tokens)
├── env.py (2000 tokens)
├── run_gemma4_e2b.sh (200 tokens)
├── train_frozenlake.py (4.9k tokens)
├── train_frozenlake_qwen3.py (4.3k tokens)
├── grpo_gemma.ipynb (11.5k tokens)
├── logit_distillation.ipynb (3.8k tokens)
├── math_gsm8k/
├── gemma_grpo_demo_nb.py (3.7k tokens)
├── qwen3_grpo_demo.py (3.9k tokens)
├── qwen_agentic_maxtext.py (8.5k tokens)
├── model_load/
├── from_safetensor_load/
├── gemma2_model_load.ipynb (1500 tokens)
├── gemma3_model_load.ipynb (1500 tokens)
├── qlora_gemma.ipynb (61.9k tokens)
├── qlora_llama3_gpu.ipynb (6.2k tokens)
├── rl/
├── README.md (600 tokens)
├── grpo/
├── gsm8k/
├── configs/
├── gemma2_2b.yaml (400 tokens)
├── gemma3_12b.yaml (400 tokens)
├── gemma3_1b.yaml (400 tokens)
├── gemma3_4b.yaml (400 tokens)
├── gemma4_e2b.yaml (400 tokens)
├── llama3.1_8b.yaml (400 tokens)
├── llama3.2_1b.yaml (400 tokens)
├── qwen3_0.6b.yaml (400 tokens)
├── run_gemma2_2b.sh (500 tokens)
├── run_gemma3_12b.sh (400 tokens)
├── run_gemma3_1b.sh (400 tokens)
├── run_gemma3_4b.sh (400 tokens)
├── run_gemma_7b.sh (800 tokens)
├── run_llama3.1_8b.sh (400 tokens)
├── run_llama3.2_1b.sh (400 tokens)
├── run_qwen3.sh (900 tokens)
├── run_qwen3_8b_disagg.sh (1400 tokens)
├── run_qwen3_8b_disagg_maxtext.sh (1500 tokens)
├── run_qwen3_simplereward.sh (900 tokens)
├── run_qwen3_vllm_disagg.sh (800 tokens)
├── verl_compatible/
├── README.md (200 tokens)
├── run_llama3.2_1b.sh (800 tokens)
├── ppo/
├── gsm8k/
├── configs/
├── gemma2_2b.yaml (500 tokens)
├── gemma3_1b.yaml (500 tokens)
├── gemma_2b.yaml (500 tokens)
├── llama3.1_8b.yaml (500 tokens)
├── qwen3_1.7b.yaml (500 tokens)
├── run_gemma2_2b.sh (600 tokens)
├── run_gemma3_1b.sh (600 tokens)
├── run_gemma_2b.sh (600 tokens)
├── run_llama3.1_8b.sh (500 tokens)
├── run_qwen3_1.7b.sh (600 tokens)
├── sft/
├── mtnt/
├── README.md (600 tokens)
├── configs/
├── gemma2_2b.yaml (200 tokens)
├── run_gemma2_2b.sh (300 tokens)
├── run_gemma3_4b.sh (400 tokens)
├── run_gemma_2b.sh (400 tokens)
├── run_llama3.2_3b.sh (300 tokens)
├── run_qwen2.5_0.5b.sh (300 tokens)
├── vlm_training.py (2.9k tokens)
├── tpu_utils.sh (500 tokens)
├── vl_dpo_gemma3.ipynb (3.3k tokens)
├── pyproject.toml (700 tokens)
├── readthedocs.yml (100 tokens)
├── requirements/
├── maxtext_requirements.txt (100 tokens)
├── requirements.txt
├── special_requirements.txt (100 tokens)
├── scripts/
├── build_raiden_wheel.sh (400 tokens)
├── forward_pass_logit_checker.py (2.3k tokens)
├── grpo_demo_llama3_qwen2.py (7.7k tokens)
├── install_tunix_vllm_requirement.sh (300 tokens)
├── llama3_example.py (400 tokens)
├── run_remote_pw.sh (2.5k tokens)
├── setup_cli_gpu_single_host.sh (200 tokens)
├── setup_cli_tpu_single_host.sh (100 tokens)
├── setup_notebook_tpu_single_host.sh (1300 tokens)
├── tests/
├── cli/
├── base_rl_pipeline_test.py (4.5k tokens)
├── config_test.py (5.1k tokens)
├── grpo_main_test.py (4.5k tokens)
├── ppo_main_test.py (3.3k tokens)
├── utils/
├── data_test.py (2.1k tokens)
├── model_test.py (2.2k tokens)
├── common/
├── configs_test.py (700 tokens)
├── conftest.py (300 tokens)
├── diffusion/
├── types_test.py (1700 tokens)
├── distillation/
├── distillation_trainer_test.py (1700 tokens)
├── feature_extraction/
├── pooling_test.py (1300 tokens)
├── projection_test.py (1100 tokens)
├── sowed_module_test.py (2.7k tokens)
├── strategies/
├── attention_test.py (1000 tokens)
├── feature_pooling_test.py (1500 tokens)
├── feature_projection_test.py (800 tokens)
├── logit_test.py (1700 tokens)
├── examples/
├── data/
├── math_dataset_test.py (1200 tokens)
├── deepswe_mllog_utils_test.py (3.9k tokens)
├── swe_env_test.py (4.1k tokens)
├── template_test.py (1200 tokens)
├── experimental/
├── common/
├── batch_utils_test.py (500 tokens)
├── datatypes_test.py (2.5k tokens)
├── lineage_test.py (900 tokens)
├── logging_utils_test.py (900 tokens)
├── rpc_utils_test.py (400 tokens)
├── distributed/
├── deployment/
├── yaml_generator_test.py (1300 tokens)
├── examples/
├── worker_transport/
├── worker_transport_test.py (1100 tokens)
├── runtime/
├── context_test.py (300 tokens)
├── contexts/
├── k8s_context_test.py (1000 tokens)
├── local_context_test.py (500 tokens)
├── discovery/
├── discovery_test.py (700 tokens)
├── executors/
├── executor_test.py (400 tokens)
├── main_test.py (900 tokens)
├── examples/
├── common/
├── run_trainer_node_test.py (4.4k tokens)
├── frozenlake_dist/
├── frozenlake_test.py (1400 tokens)
├── math_gsm8k_dist/
├── gsm8k_test.py (1700 tokens)
├── orchestrator/
├── algorithm_adapter_test.py (4.6k tokens)
├── batch_assembly_test.py (15.7k tokens)
├── distributed_rl_engine_test.py (11k tokens)
├── health_monitor_test.py (1100 tokens)
├── lifecycle_test.py (700 tokens)
├── orchestrator_test.py (2.9k tokens)
├── rl_program_test.py (21.2k tokens)
├── startup_validation_test.py (1300 tokens)
├── worker_registry_test.py (1300 tokens)
├── queue_manager/
├── trajectory_queue_manager_test.py (1200 tokens)
├── rl/
├── agentic/
├── registry_test.py (2.4k tokens)
├── rollout/
├── collector_test.py (5.4k tokens)
├── inprocess_vllm_sampler_adapter_test.py (2.8k tokens)
├── manager_test.py (2.6k tokens)
├── raiden_weight_sync_mixin_test.py (700 tokens)
├── rollout_test.py (2.1k tokens)
├── router_replay_end_to_end_test.py (1500 tokens)
├── sampler_test.py (1200 tokens)
├── vanilla_sampler_adapter_test.py (1700 tokens)
├── vllm_sampler_adapter_test.py (2.3k tokens)
├── train/
├── peft_trainer_v2_test.py (9.7k tokens)
├── peft_trainer_v2_weight_sync_test.py (600 tokens)
├── trajectory/
├── action_converter_test.py (7.4k tokens)
├── async_writer_test.py (6.6k tokens)
├── benchmarks/
├── benchmark_lib_test.py (600 tokens)
├── data_generator_test.py (400 tokens)
├── run_benchmark_test.py (400 tokens)
├── converter_test.py (8.9k tokens)
├── file_store_test.py (3.4k tokens)
├── in_memory_store_test.py (700 tokens)
├── schema_test.py (1100 tokens)
├── trajectory_test.py (6.6k tokens)
├── weight_sync/
├── raiden_handler_test.py (5.9k tokens)
├── raiden_integration_test.py (800 tokens)
├── raiden_preload_test.py (300 tokens)
├── raiden_synchronizer_test.py (4.1k tokens)
├── raiden_weight_sync_delegate_test.py (1500 tokens)
├── weight_sync_coordinator_test.py (13.7k tokens)
├── weight_sync_test.py (4.9k tokens)
├── worker/
├── abstract_worker_test.py (700 tokens)
├── examples/
├── agentic_remote_execution_demo_test.py (300 tokens)
├── rl_loop_remote_execution_demo_test.py (400 tokens)
├── inference_worker_test.py (1500 tokens)
├── mock_worker_test.py (500 tokens)
├── remote_execution_test.py (8.5k tokens)
├── rollout_worker_test.py (1200 tokens)
├── rollout_worker_weight_sync_test.py (1000 tokens)
├── traffic_controller_test.py (800 tokens)
├── trainer_worker_test.py (1000 tokens)
├── trainer_worker_weight_sync_test.py (600 tokens)
├── generate/
├── beam_search_test.py (2.3k tokens)
├── mappings_config_test.py (500 tokens)
├── sampler_test.py (6.2k tokens)
├── sglang_jax_lora_test.py (1600 tokens)
├── sglang_jax_sampler_test.py (1100 tokens)
├── tokenizer_adapter_test.py (500 tokens)
├── utils_test.py (17k tokens)
├── vllm_driver_test.py (1900 tokens)
├── vllm_sampler_qwen_test.py (900 tokens)
├── vllm_sampler_test.py (4.4k tokens)
├── model_alignment/
├── qwen_align_test.py (1300 tokens)
├── models/
├── automodel_test.py (3.1k tokens)
├── cache_utils_test.py (3.7k tokens)
├── dummy_model_creator_test.py (700 tokens)
├── flash_attention_cache_test.py (1500 tokens)
├── gemma3/
├── merge_embeddings_test.py (800 tokens)
├── utils_test.py (500 tokens)
├── gemma4/
├── attention_test.py (8k tokens)
├── gemma4_params_test.py (4.7k tokens)
├── model_test.py (7.4k tokens)
├── params_safetensors_test.py (600 tokens)
├── gemma_all/
├── gemma_params_test.py (5.6k tokens)
├── llama3/
├── params_test.py (900 tokens)
├── naming_test.py (2.5k tokens)
├── qwen3/
├── qwen_params_test.py (2.1k tokens)
├── registry_test.py (600 tokens)
├── safetensors_loader_test.py (2k tokens)
├── perf/
├── experimental/
├── export_v2_test.py (1400 tokens)
├── timeline_test.py (2.1k tokens)
├── timeline_utils_test.py (1900 tokens)
├── trace_writer_test.py (3.2k tokens)
├── tracer_test.py (4.9k tokens)
├── export_test.py (2.7k tokens)
├── metrics_test.py (600 tokens)
├── perfetto_test.py (1900 tokens)
├── span_test.py (1600 tokens)
├── trace_test.py (2.3k tokens)
├── processors/
├── audio_processor_test.py (1900 tokens)
├── image_processor_test.py (2.4k tokens)
├── rl/
├── agentic/
├── agentic_grpo_learner_test.py (16.5k tokens)
├── agentic_rl_learner_test.py (1600 tokens)
├── agentic_utils_test.py (2.3k tokens)
├── agents/
├── agent_types_test.py (1300 tokens)
├── parser/
├── chat_template_parser/
├── chat_template_parser_test.py (2.5k tokens)
├── pipeline/
├── rollout_orchestrator_test.py (1500 tokens)
├── queue_manager/
├── group_queue_manager_test.py (1600 tokens)
├── rewards/
├── reward_test.py (800 tokens)
├── tools/
├── agentic_tool_test.py (1500 tokens)
├── calculator_test.py (500 tokens)
├── trajectory/
├── trajectory_collect_engine_test.py (5.2k tokens)
├── algo_core_test.py (1100 tokens)
├── algorithm_config_test.py (1600 tokens)
├── common_test.py (7.6k tokens)
├── function_registry_test.py (800 tokens)
├── grpo/
├── dapo_learner_test.py (2k tokens)
├── drgrpo_learner_test.py (1100 tokens)
├── grpo_learner_test.py (9.5k tokens)
├── packing_test.py (1900 tokens)
├── ppo/
├── ppo_helpers_test.py (1100 tokens)
├── ppo_learner_test.py (4.3k tokens)
├── reward_manager_test.py (3k tokens)
├── rl_cluster_test.py (6.8k tokens)
├── rl_learner_test.py (600 tokens)
├── rl_utils_test.py (6.4k tokens)
├── rollout/
├── mock_rollout_test.py (2.8k tokens)
├── router_replay_maxtext_test.py (1300 tokens)
├── sft/
├── checkpoint_manager_test.py (4.9k tokens)
├── checkpoint_options_test.py (1300 tokens)
├── diffusion_sft_test.py (1800 tokens)
├── dpo/
├── dpo_trainer_test.py (2.5k tokens)
├── orpo_trainer_test.py (3.6k tokens)
├── metrics_logger_test.py (2.7k tokens)
├── otel_wandb_test.py (1900 tokens)
├── peft_trainer_test.py (16k tokens)
├── profiler_test.py (2.6k tokens)
├── progress_bar_test.py (700 tokens)
├── sft_tpu_smoke_test.sh (300 tokens)
├── sft_utils_test.py (1100 tokens)
├── sharding_utils_test.py (1200 tokens)
├── system_metrics_calculator_test.py (1000 tokens)
├── test_data/
├── checkpoints/
├── 1/
├── _CHECKPOINT_METADATA (100 tokens)
├── model_params/
├── _METADATA (200 tokens)
├── _sharding (200 tokens)
├── array_metadatas/
├── process_0 (100 tokens)
├── d/
├── 8f2ed9d45ec81eba9c5612065ee66958
├── descriptor/
├── descriptor.pbtxt
├── uuid-913517f8-082d-4d69-925b-eeac7abe4017
├── manifest.ocdbt
├── ocdbt.process_0/
├── d/
├── 47dc145b2db6c5634d657942d0feb7ef
├── 485463e063b2af866bef815972b04a48
├── e9b70cbab93d7f1ad394c1158af2cc7b
├── fb0a3f9c92c7441d7a25f29326befa36
├── manifest.ocdbt
├── smoke_tests/
├── model_creation_test.py (800 tokens)
├── utils/
├── compat_test.py (300 tokens)
├── gsm8k_vtc_test.py (1000 tokens)
├── math_utils_test.py (500 tokens)
├── maxtext_utils_test.py (2.7k tokens)
├── mesh_test.py (8.4k tokens)
├── token_sanitization_test.py (600 tokens)
├── topology_test.py (1600 tokens)
├── torch_utils_test.py (300 tokens)
├── trajectory_logger_test.py (500 tokens)
├── tunix/
├── __init__.py (600 tokens)
├── cli/
├── README.md (400 tokens)
├── base_agentic_config.yaml (2.1k tokens)
├── base_config.yaml (2k tokens)
├── base_rl_pipeline.py (6.5k tokens)
├── config.py (8.1k tokens)
├── grpo_main.py (1400 tokens)
├── grpo_vllm_main.py (200 tokens)
├── peft_main.py (700 tokens)
├── ppo_main.py (800 tokens)
├── recipes/
├── __init__.py (100 tokens)
├── deepscaler_data.py (500 tokens)
├── reward_fn/
├── gsm8k.py (900 tokens)
├── gsm8k_verl.py (500 tokens)
├── simple_math.py (400 tokens)
├── utils/
├── data.py (1600 tokens)
├── model.py (1600 tokens)
├── common/
├── configs.py (3.3k tokens)
├── datatypes.py (200 tokens)
├── diffusion/
├── README.md (600 tokens)
├── __init__.py (200 tokens)
├── interfaces.py (400 tokens)
├── types.py (1400 tokens)
├── distillation/
├── __init__.py (200 tokens)
├── distillation_trainer.py (1100 tokens)
├── feature_extraction/
├── __init__.py (300 tokens)
├── pooling.py (900 tokens)
├── projection.py (1000 tokens)
├── sowed_module.py (1300 tokens)
├── strategies/
├── __init__.py (200 tokens)
├── attention.py (900 tokens)
├── base_strategy.py (1000 tokens)
├── feature_pooling.py (1400 tokens)
├── feature_projection.py (1300 tokens)
├── logit.py (1200 tokens)
├── examples/
├── __init__.py
├── data/
├── __init__.py
├── math_dataset.py (1400 tokens)
├── translation_dataset.py (1400 tokens)
├── experimental/
├── common/
├── batch_utils.py (400 tokens)
├── datatypes.py (3.9k tokens)
├── lineage.py (800 tokens)
├── logging_utils.py (300 tokens)
├── rpc_utils.py (300 tokens)
├── test_utils.py (2.9k tokens)
├── distributed/
├── deployment/
├── yaml_generator.py (1200 tokens)
├── yamls/
├── jobset.cpu.yaml (800 tokens)
├── jobset.mcjax.yaml (1200 tokens)
├── jobset.pathways.qwen3.5-397b.yaml (2.5k tokens)
├── jobset.pathways.yaml (2.5k tokens)
├── jobset.tpu.yaml (900 tokens)
├── leaderworkerset.mcjax.ray.yaml (700 tokens)
├── examples/
├── README.md (2.6k tokens)
├── basics/
├── basic.py (100 tokens)
├── door.py (200 tokens)
├── flag.py (100 tokens)
├── knocker.py (100 tokens)
├── tpu.py (400 tokens)
├── rl/
├── launcher.sh (1300 tokens)
├── orchestrator.py (800 tokens)
├── rollout.py (500 tokens)
├── service.proto (200 tokens)
├── trainer.py (600 tokens)
├── vllm_rollout/
├── launcher.sh (1400 tokens)
├── orchestrator.py (1700 tokens)
├── remote.py (500 tokens)
├── rollout.py (800 tokens)
├── worker_transport/
├── README.md (800 tokens)
├── remote_worker_server.py (500 tokens)
├── transport.py (300 tokens)
├── worker.py (200 tokens)
├── runtime/
├── context.py (500 tokens)
├── contexts/
├── k8s_context.py (1400 tokens)
├── local_context.py (1000 tokens)
├── discovery/
├── discovery.py (900 tokens)
├── discovery_service.proto (300 tokens)
├── executor.py (100 tokens)
├── executors/
├── k8s_executor.py (300 tokens)
├── local_executor.py (300 tokens)
├── main.py (1600 tokens)
├── docs/
├── orchestrator_plan_v2.md (7.9k tokens)
├── examples/
├── common/
├── enter_kube_context.sh (400 tokens)
├── models.py (600 tokens)
├── run_inference_node.py (1600 tokens)
├── run_rollout_node.py (4.4k tokens)
├── run_trainer_node.py (4k tokens)
├── deepswe_dist/
├── README.md (300 tokens)
├── deepswe.py (1600 tokens)
├── k8s_launcher.sh (2.9k tokens)
├── launcher.sh (2.8k tokens)
├── run_deepswe_dist.py (2.9k tokens)
├── frozenlake_dist/
├── README.md (500 tokens)
├── frozenlake.py (900 tokens)
├── launcher.sh (2.2k tokens)
├── run_frozenlake_dist.py (2.7k tokens)
├── math_gsm8k_dist/
├── gsm8k.py (1400 tokens)
├── k8s_launcher.sh (3.9k tokens)
├── launcher.sh (5.3k tokens)
├── run_gsm8k_dist_grpo.py (3.2k tokens)
├── recipes/
├── trellis_gsm8k_qwen3p5_35b.sh (600 tokens)
├── metrics/
├── metrics.py (400 tokens)
├── orchestrator/
├── algorithm_adapter.py (2.9k tokens)
├── batch_assembly.py (7.5k tokens)
├── distributed_rl_engine.py (5.3k tokens)
├── health_monitor.py (1100 tokens)
├── lifecycle.py (700 tokens)
├── orchestrator.py (2.9k tokens)
├── rl_engine_interface.py (1400 tokens)
├── rl_program.py (7.5k tokens)
├── simple_orchestrator_nb.py (1100 tokens)
├── startup_validation.py (1100 tokens)
├── worker_registry.py (1000 tokens)
├── queue_manager/
├── trajectory_queue_manager.py (1100 tokens)
├── rl/
├── agentic/
├── registry.py (1700 tokens)
├── rollout/
├── __init__.py (300 tokens)
├── collector.py (1800 tokens)
├── inprocess_vllm_sampler_adapter.py (2.4k tokens)
├── manager.py (2.9k tokens)
├── raiden_weight_sync_mixin.py (1300 tokens)
├── sampler.py (1600 tokens)
├── vanilla_sampler_adapter.py (2.3k tokens)
├── vllm_sampler_adapter.py (3.5k tokens)
├── train/
├── abstract_trainer.py (1400 tokens)
├── peft_trainer_v2.py (11.4k tokens)
├── trajectory/
├── action_converter.py (2.4k tokens)
├── async_writer.py (2.9k tokens)
├── benchmarks/
├── benchmark_lib.py (1500 tokens)
├── data_generator.py (600 tokens)
├── run_benchmark.py (1000 tokens)
├── converter.py (3.8k tokens)
├── file_store.py (2.1k tokens)
├── in_memory_store.py (1000 tokens)
├── schema.py (900 tokens)
├── schema_testing.py (2.1k tokens)
├── store.py (800 tokens)
├── store_testing.py (3k tokens)
├── testdata/
├── sample_atif_v1_7.json (300 tokens)
├── trajectory.py (4.8k tokens)
├── trajectory_testing.py (800 tokens)
├── weight_sync/
├── raiden_handler.py (4.1k tokens)
├── raiden_preload.py (400 tokens)
├── raiden_synchronizer.py (6.1k tokens)
├── raiden_weight_sync_delegate.py (1200 tokens)
├── weight_sync.py (4.1k tokens)
├── weight_sync_coordinator.py (14.2k tokens)
├── worker/
├── __init__.py (300 tokens)
├── abstract_worker.py (500 tokens)
├── examples/
├── agentic_remote_execution_demo.py (1100 tokens)
├── rl_loop_remote_execution_demo.py (2000 tokens)
├── sample_agent_and_env.py (500 tokens)
├── inference_worker.py (1800 tokens)
├── mock_worker.py (700 tokens)
├── remote_execution.py (9.3k tokens)
├── rollout_worker.py (3.2k tokens)
├── traffic_controller.py (1000 tokens)
├── trainer_worker.py (2000 tokens)
├── generate/
├── base_sampler.py (600 tokens)
├── beam_search.py (2.1k tokens)
├── mappings.py (1100 tokens)
├── sampler.py (7.7k tokens)
├── sglang_jax_sampler.py (3k tokens)
├── tokenizer_adapter.py (2.1k tokens)
├── utils.py (14.8k tokens)
├── vllm_async_driver.py (2.8k tokens)
├── vllm_sampler.py (5.1k tokens)
├── models/
├── automodel.py (4.8k tokens)
├── cache_utils.py (3k tokens)
├── dummy_model_creator.py (600 tokens)
├── gemma/
├── model.py (6.2k tokens)
├── params.py (500 tokens)
├── params_safetensors.py (1400 tokens)
├── sampler.py (4.4k tokens)
├── gemma3/
├── merge_embeddings.py (400 tokens)
├── model.py (7.5k tokens)
├── params.py (2.7k tokens)
├── params_safetensors.py (2.6k tokens)
├── utils.py (1000 tokens)
├── vision.py (3.9k tokens)
├── gemma4/
├── __init__.py (200 tokens)
├── attention.py (6.1k tokens)
├── audio.py (7.6k tokens)
├── config.py (2.9k tokens)
├── layers.py (2.3k tokens)
├── mapping_vllm_jax.py (1700 tokens)
├── model.py (4.9k tokens)
├── moe.py (1400 tokens)
├── params.py (2.6k tokens)
├── params_safetensors.py (7k tokens)
├── sampling_example.ipynb (2.2k tokens)
├── vision.py (5.2k tokens)
├── llama3/
├── __init__.py (200 tokens)
├── mapping_sglang_jax.py (1400 tokens)
├── mapping_vllm_jax.py (900 tokens)
├── model.py (4k tokens)
├── params.py (700 tokens)
├── naming.py (2.4k tokens)
├── qwen2/
├── __init__.py (200 tokens)
├── mapping_sglang_jax.py (1400 tokens)
├── mapping_vllm_jax.py (1000 tokens)
├── model.py (5.9k tokens)
├── params.py (700 tokens)
├── qwen3/
├── __init__.py (200 tokens)
├── mapping_sglang_jax.py (1600 tokens)
├── mapping_vllm_jax.py (900 tokens)
├── model.py (9.2k tokens)
├── params.py (1200 tokens)
├── registry.py (3.4k tokens)
├── safetensors_loader.py (3.3k tokens)
├── safetensors_saver.py (900 tokens)
├── oss/
├── utils.py (600 tokens)
├── perf/
├── experimental/
├── constants.py (200 tokens)
├── export.py (1100 tokens)
├── timeline.py (2.7k tokens)
├── timeline_utils.py (1300 tokens)
├── trace_writer.py (2.6k tokens)
├── tracer.py (2.3k tokens)
├── export.py (4.2k tokens)
├── metrics.py (1500 tokens)
├── perfetto.py (1600 tokens)
├── span.py (1700 tokens)
├── trace.py (2.7k tokens)
├── processors/
├── audio_processor.py (1500 tokens)
├── image_processor.py (3.4k tokens)
├── rl/
├── __init__.py
├── agentic/
├── agentic_grpo_learner.py (7.9k tokens)
├── agentic_rl_learner.py (9.7k tokens)
├── agents/
├── agent_types.py (1900 tokens)
├── base_agent.py (1600 tokens)
├── model_agent.py (400 tokens)
├── tool_agent.py (1200 tokens)
├── environments/
├── base_environment.py (2.4k tokens)
├── task_environment.py (800 tokens)
├── tool_environment.py (1800 tokens)
├── parser/
├── chat_template_parser/
├── parser.py (3.1k tokens)
├── tool_parser/
├── gemini_parser.py (200 tokens)
├── qwen_parser.py (800 tokens)
├── tool_parser_base.py (600 tokens)
├── tool_parser_registry.py (200 tokens)
├── pipeline/
├── rollout_orchestrator.py (3k tokens)
├── queue_manager/
├── group_queue_manager.py (1400 tokens)
├── rewards/
├── reward.py (1400 tokens)
├── reward_types.py (500 tokens)
├── tools/
├── base_tool.py (1500 tokens)
├── calculator_tool.py (900 tokens)
├── tool_manager.py (1700 tokens)
├── trajectory/
├── trajectory_collect_engine.py (5.7k tokens)
├── utils.py (1800 tokens)
├── algo_core.py (4.9k tokens)
├── algorithm_config.py (1700 tokens)
├── common.py (9.4k tokens)
├── function_registry.py (1200 tokens)
├── grpo/
├── dapo_learner.py (1200 tokens)
├── drgrpo_learner.py (300 tokens)
├── grpo_learner.py (3.9k tokens)
├── inference/
├── inference_worker.py (600 tokens)
├── packing.py (1900 tokens)
├── ppo/
├── ppo_learner.py (4.1k tokens)
├── queue/
├── data_queue.py (400 tokens)
├── reshard.py (1600 tokens)
├── reward_manager.py (3.1k tokens)
├── rl_cluster.py (9.3k tokens)
├── rl_learner.py (6k tokens)
├── rollout/
├── base_rollout.py (600 tokens)
├── mock_rollout.py (2.5k tokens)
├── rollout_engine_group.py (500 tokens)
├── rollout_traffic_router.py (200 tokens)
├── sglang_jax_rollout.py (1000 tokens)
├── vanilla_rollout.py (600 tokens)
├── vllm_rollout.py (1100 tokens)
├── trainer.py (800 tokens)
├── utils.py (4.5k tokens)
├── sft/
├── checkpoint_manager.py (2.4k tokens)
├── checkpoint_options.py (1900 tokens)
├── diffusion.py (600 tokens)
├── dpo/
├── dpo_trainer.py (5.7k tokens)
├── hooks.py (500 tokens)
├── inflight_throttler.py (400 tokens)
├── metrics_logger.py (2.9k tokens)
├── otel_wandb.py (1100 tokens)
├── peft_trainer.py (8.9k tokens)
├── profiler.py (1300 tokens)
├── progress_bar.py (600 tokens)
├── sharding_utils.py (700 tokens)
├── system_metrics_calculator.py (500 tokens)
├── utils.py (1500 tokens)
├── tests/
├── lora_params_test_base.py (2.7k tokens)
├── test_common.py (3.3k tokens)
├── utils/
├── compat.py (400 tokens)
├── env_utils.py (300 tokens)
├── gsm8k_vtc.py (1700 tokens)
├── immutable.py (400 tokens)
├── math_rewards.py (700 tokens)
├── math_utils.py (4k tokens)
├── maxtext_utils.py (2.4k tokens)
├── mesh.py (12.3k tokens)
├── mllog_utils.py (5.5k tokens)
├── script_utils.py (1000 tokens)
├── sharding_utils.py (300 tokens)
├── token_sanitization.py (500 tokens)
├── topology.py (4.3k tokens)
├── torch_utils.py (100 tokens)
├── trajectory_logger.py (1600 tokens)
```
## /.gemini/styleguide.md
# Tunix Code Review Style Guide
This guide defines the architectural standards and review criteria for Tunix. **Gemini Code Assist must use this guide to provide high-level, context-aware reviews.**
---
## 1. Core Philosophy: JAX-Native & NNX-First
Tunix is a JAX-native post-training library for LLMs. Tunix principally uses NNX, a neural network library built on top of JAX. Please follow all common JAX and NNX patterns.
* **Readability:** Code should be easy to understand for all maintainers and users.
* **Maintainability:** Code should be easy to modify and extend.
* **Modularity**
* **Consistency**: Adherence to a consistent style across all projects. For example, adherence to naming and file structure conventions is crucial for predictability and maintainability.
* **Reusability**: Try re-using components instead of re-writing.
---
## 2. Reference Documentation (Read First)
**Gemini:** Before reviewing any code, **read the following files** to understand the architectural context. Your review should be grounded in these documents.
* **Contributing guidelines**: [`contributing.md`](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/contributing.md)
* **Models**: [`models.md`](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/models.md) - `AutoModel` patterns and naming.
* **RL Algorithms**: [`algorithms.md`](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/algorithms.md) - Registry and Config patterns for RL.
* **Agentic RL**: [`agentic_rl.md`](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/agentic_rl.md)
* [**JAX**](https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html)
* [**NNX**](https://flax.readthedocs.io/en/latest/nnx_basics.html)
In general, you can look into the [documentation files](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/).
---
## 3. Review Categories
Identify the type of contribution and apply the corresponding checklist.
### A. Model Contributions (`tunix/models/`)
* **Naming:** Must follow the strict pattern `<family><ver>p<min>_<size>` (e.g., `gemma2p0_9b`, `qwen2p5_1p5b`).
* **AutoModel:** New models must be integrated into `AutoModel.from_pretrained` and support all sources (`HUGGINGFACE`, `KAGGLE`, `GCS`).
* **Pattern:**
* `ModelConfig`: Dataclass for architecture params.
* `ShardingConfig`: separate Dataclass for partition specs.
* `Module`: Pure NNX module implementation.
### B. RL Algorithms (`tunix/rl/`)
* **Pattern:** Logic must be split between a **Configuration** and a **Learner**.
* **Config:** Must inherit from `AlgorithmConfig` (e.g., `class PPOConfig(AlgorithmConfig)`).
* **Learner:** Must inherit from `RLLearner` (e.g., `class PPOLearner(RLLearner)`).
* **Registry:** Loss functions and advantage estimators *must* be registered (e.g., `@register_policy_loss_fn`) to allow hot-swapping via config.
* **Reward Managers:** Complex reward logic should live in a `RewardManager`, not the Learner loop.
### C. Stand-alone Algorithms (`tunix/sft/`)
For non-RL algorithms, you can follow a pattern similar to `PeftTrainer` or `DpoTrainer`.
* **Pattern:** SFT uses `TrainingConfig` (not `AlgorithmConfig`) and `PeftTrainer`.
* **Trainer:** `PeftTrainer`
### D. Bug Fixes
* **Reproduction:** Critical bug fixes should include a reproduction script or a link to a Colab notebook demonstrating the issue.
* **Regression Test:** A matching test case in `_test.py` is mandatory.
### E. Notebooks & Examples
Prioritise readability. Have explanatory text cells to explain the code. Mention the hardware/accelerator on which the example will run.
---
## 4. Formatting & Linting
**Instructions for Contributors:**
If the formatting is off, please instruct the user to run the linter/formatter. You can pull instructions from [contributing.md](https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/contributing.md).
## 5. Other generic advice
### Type hints
* **Use type hints:** Type hints improve code readability and help catch errors early.
### Comments
* **Write clear and concise comments:** Explain the "why" behind the code, not just the "what".
* **Comment sparingly:** Well-written code should be self-documenting where possible.
* **Use complete sentences:** Start comments with a capital letter and use proper punctuation.
### Logging
* **Use absl for logging**
* **Log at appropriate levels:** DEBUG, INFO, WARNING, ERROR, CRITICAL
* **Provide context:** Include relevant information in log messages to aid debugging.
### Error Handling
* **Use specific exceptions:** Avoid using broad exceptions like `Exception`.
* **Handle exceptions gracefully:** Provide informative error messages and avoid crashing the program.
## /.github/CODEOWNERS
```github/CODEOWNERS path="/.github/CODEOWNERS"
* @tianshub @wang2yn84 @lc5211 @hgao327 @sizhit2 @abheesht17 @jiangyangmu @s-noghabi
# Trajectory Store
/tunix/experimental/trajectory/ @niting @fineguy @shz-google
/tests/experimental/trajectory/ @niting @fineguy @shz-google
```
## /.github/ISSUE_TEMPLATE/bug_report.md
---
name: Bug report
about: Create a report to help us fix/improve something
title: 'A brief description of the bug'
labels: 'bug'
assignees: ''
---
**Expected Behavior**
**Actual Behavior**
**Steps to Reproduce the Problem**
1.
2.
3.
<!---
If possible, provide a link to a [Colab Notebook](https://research.google.com/colaboratory/),
which reproduces the bug.
-->
**Environment**
- **OS:** [e.g., Ubuntu, etc.]
- **Project Version:** [e.g., 0.0.1]
**Checklist**
- [ ] I have searched the existing issues for a similar bug report.
- [ ] I have provided all the required information in the "Environment" section.
- [ ] I have provided a minimal, reproducible example.
**Would you like to help us fix it?**
## /.github/ISSUE_TEMPLATE/feature_request.md
---
name: Feature request
about: Suggest an idea for this project
title: 'A brief description of the feature'
labels: 'enhancement'
assignees: ''
---
**Is your feature request related to a problem? Please describe.**
<!---
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-->
**Describe the solution you'd like**
<!---
A clear and concise description of what you want to happen.
-->
**Additional context**
<!---
Add any other context or screenshots about the feature request here.
-->
**Checklist**
- [ ] I have searched the existing issues for similar feature requests.
- [ ] This is not a support question (please use the "bug template" for that).
## /.github/pull_request_template.md
Resolves #\<issue_number_goes_here\>
> It's a good idea to open an issue first for discussion.
<!--- Describe your changes in detail. -->
**Reference**
<!--- Link to the reference implementation, research paper, and GitHub issue. -->
**Colab Notebook**
<!-- If adding any new API, attach a Colab notebook showing the high-level usage.-->
**Checklist**
<!--- Please make sure all checkboxes are ticked before submitting this PR for review. -->
- [ ] I have added all the necessary unit tests for my change.
- [ ] I have verified that my change does not break existing code and all unit tests pass.
- [ ] I have added all appropriate doc-strings/documentation.
- [ ] My PR is based on the latest changes of the main branch (if unsure, rebase the code).
- [ ] I have signed the [Contributor License Agreement](https://cla.developers.google.com/about).
- [ ] I have followed [Contribution Guidelines](https://github.com/google/tunix/blob/main/docs/contributing.md).
> **Note**: Standard CPU unit tests, package builds, and documentation checks will run automatically on pull requests. Once the PR is approved and ready for submission, maintainers will add the `ready-to-submit` label to trigger full TPU testing.
## /.github/workflows/build_and_test_tunix.yml
```yml path="/.github/workflows/build_and_test_tunix.yml"
# Copyright 2025 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This workflow will build tunix python package and run tests.
name: Tunix Package Tests
# Trigger the workflow only when code is merging to the 'main' branch.
on:
pull_request: # Automatically trigger on pull requests affecting this file
branches:
- main
types: [opened, synchronize, reopened, labeled]
workflow_dispatch:
schedule:
# Run the job every 4 hours
- cron: '0 */4 * * *'
concurrency:
# Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else
group: >
${{
github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) ||
github.event_name == 'schedule' && format('{0}-schedule', github.workflow) ||
github.run_id
}}
cancel-in-progress: true
permissions:
contents: read
jobs:
build_tunix_package:
name: Build tunix package
uses: ./.github/workflows/build_package.yml
tunix_cpu_unit_tests:
needs: build_tunix_package
uses: ./.github/workflows/cpu-tests.yml
secrets:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
tunix_tpu_unit_tests:
needs: build_tunix_package
if: ${{ github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'ready-to-submit') }}
uses: ./.github/workflows/tpu-tests.yml
secrets:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
notify_failure:
name: Notify failed build # creates an issue or modifies last open existing issue for failed build
needs: [build_tunix_package,tunix_cpu_unit_tests, tunix_tpu_unit_tests]
if: ${{ always() }}
runs-on: ubuntu-latest
permissions:
issues: write
steps:
- name: Check whether one of the jobs failed
if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
```
## /.github/workflows/build_and_test_tunix_nightly_regression.yml
```yml path="/.github/workflows/build_and_test_tunix_nightly_regression.yml"
# Copyright 2025 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This workflow will build tunix python package and run tpu regression tests.
name: Tunix Nightly Regression Tests
on:
workflow_dispatch:
schedule:
# Run the job every day at 2am
- cron: '0 10 * * *' # 10:00 AM UTC, 2:00 AM PST
concurrency:
# Dedup scheduled runs but nothing else
group: >
${{
github.event_name == 'schedule' && format('{0}-schedule', github.workflow) ||
github.run_id
}}
cancel-in-progress: false
permissions:
contents: read
jobs:
build_tunix_package:
name: Build tunix package
uses: ./.github/workflows/build_package.yml
tunix_tpu_nightly_regression:
needs: build_tunix_package
uses: ./.github/workflows/tpu-nightly-regression.yml
secrets:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
notify_failure:
name: Notify failed build # creates an issue or modifies last open existing issue for failed build
needs: [build_tunix_package, tunix_tpu_nightly_regression]
if: ${{ always() }}
runs-on: ubuntu-latest
permissions:
issues: write
steps:
- name: Check whether one of the jobs failed
if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }}
uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
```
## /.github/workflows/build_docs.yml
```yml path="/.github/workflows/build_docs.yml"
name: Check the Documentation Build
on:
pull_request:
types: [opened, synchronize]
workflow_dispatch:
permissions:
contents: read
jobs:
build-docs:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
with:
python-version: '3.12'
cache: 'pip' # caching pip dependencies
- name: Install dependencies
run: pip install .[docs]
- name: Build documentation
run: |
pushd docs # XXX _collection/examples not found if run from root
sphinx-build -W -b html . ./_build/html
```
## /.github/workflows/build_package.yml
```yml path="/.github/workflows/build_package.yml"
# Copyright 2025 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file defines a module for building and uploading a tunix pacakge
# based on the pyproject.toml at the current github workspace.
name: Build and Upload Tunix Package
on:
workflow_call:
permissions:
contents: read
jobs:
build_and_upload:
name: Build Tunix wheel
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
with:
python-version: '3.11'
- name: Install build tools
run: |
python -m pip install --upgrade pip build
- name: Build tunix wheel
run: |
python -m build --wheel
- name: Upload wheel artifact
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
with:
name: tunix-wheel
path: dist/*.whl
```
## /.github/workflows/cpu-tests.yml
```yml path="/.github/workflows/cpu-tests.yml"
# Copyright 2025 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file defines a module for running tests against the built tunix package.
name: Run Tests Against Tunix Package
on:
workflow_call:
secrets:
HF_TOKEN:
description: 'HuggingFace token for model downloads'
required: false
KAGGLE_USERNAME:
description: 'Kaggle Username'
required: false
KAGGLE_KEY:
description: 'Kaggle API Key'
required: false
permissions:
contents: read
jobs:
run_core:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
with:
python-version: '3.11'
- name: Download the tunix wheel
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
with:
name: tunix-wheel
- name: Install the tunix wheel
run: |
python -m pip install --upgrade pip
WHEEL=$(ls -1 google_tunix-*-py3-none-any.whl | head -n1)
python -m pip install "${WHEEL}[prod,test,experimental]"
- name: Verify Tunix imports from installed package
run: |
python3 -c "
import tunix
import tunix.models
import tunix.generate
import tunix.sft
import tunix.distillation
import tunix.rl
assert tunix.__version__ != '0.0.0.dev0', 'Tunix version not set correctly'
print('All tunix modules imported successfully and version is', tunix.__version__)
"
- name: Run agentic RL tests
run: |
python -m pytest tests/rl/agentic/ -v --tb=short
- name: Run RL unit tests
run: |
python -m pytest tests/rl/ -v --tb=short -k "not test_grpo_with_lora_model" --ignore=tests/rl/agentic --ignore=tests/rl/rl_cluster_test.py
- name: Run Cli utils tests
run: |
python -m pytest tests/cli/utils/ -v --tb=short
- name: Run shared mesh and topology tests
run: |
python -m pytest tests/utils/mesh_test.py tests/utils/topology_test.py -v --tb=short
- name: Run perf tests
run: |
python -m pytest tests/perf/ -v --tb=short
- name: Run model tests
run: |
python -m pytest tests/models/ -v --tb=short \
--ignore=tests/models/naming_test.py
- name: Compile gRPC Protobuf definitions
run: |
python -m pip install grpcio-tools
find tunix/experimental/distributed -name "*.proto" -exec \
python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. {} +
- name: Run experimental tests
run: |
# TODO(tunix-dev): add back raiden tests after tpu_sync lib is ready.
python -m pytest tests/experimental/ -v --tb=short \
--ignore-glob='*/raiden_*.py' \
--ignore=tests/experimental/weight_sync/weight_sync_test.py
run_vllm:
runs-on: ubuntu-latest
container:
image: vllm/vllm-tpu:nightly-20260911-4b80060-51da0ca # zizmor: ignore[unpinned-images]
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- name: Download the tunix wheel
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
with:
name: tunix-wheel
- name: Install the tunix wheel
run: |
python3 -m pip install --upgrade pip setuptools wheel
WHEEL=$(ls -1 google_tunix-*-py3-none-any.whl | head -n1)
python3 -m pip install "${WHEEL}[prod,test,experimental]"
# tpu-inference/Numba needs NumPy 2.3 or less.
python3 -m pip install numpy==2.3.5 --force-reinstall
- name: Run vllm driver tests
run: |
python3 -m pytest tests/generate/vllm_driver_test.py -v --tb=short
# ---------------------------------------------------------------------------
# Job: run_dev
# Purpose: Model download, CLI integration, naming, and alignment tests on CPU.
# Hardware: Standard Ubuntu runner.
# Secrets: Requires HF_TOKEN and KAGGLE_* secrets (skipped on external fork PRs).
# Scope: Smoke model creation, CLI config tests, HF naming, and PyTorch alignment.
# ---------------------------------------------------------------------------
run_dev:
if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
with:
python-version: '3.11'
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Download the tunix wheel
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
with:
name: tunix-wheel
- name: Install the tunix wheel and dependencies
run: |
python -m pip install --upgrade pip
WHEEL=$(ls -1 google_tunix-*-py3-none-any.whl | head -n1)
python -m pip install "${WHEEL}[prod,test,experimental]"
python -m pip install --upgrade kagglehub kagglesdk torch
- name: Run Smoke tests (model creation)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
run: |
echo "Running Smoke tests..."
python -m pytest tests/smoke_tests/model_creation_test.py -v --tb=short
- name: Run tunix cli integration tests
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
run: |
python -m pytest tests/cli/ -v --tb=short \
--ignore=tests/cli/utils/model_test.py
- name: Run tunix model naming tests
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pytest tests/models/naming_test.py -v --tb=short
- name: Run model alignment tests
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
JAX_PLATFORMS=cpu python -m pytest tests/model_alignment/ -v --tb=short
```
## /.github/workflows/issue-auto-assign.js
```js path="/.github/workflows/issue-auto-assign.js"
/**
* Copyright 2026 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/**
* @param {{github: !Object, context: !Object}} params
* @returns {!Promise<void>}
*/
module.exports = async ({github, context}) => {
// Define the list once to keep things clean and easy to update
const assigneesList = [
"tianshub",
"wang2yn84",
"lc5211",
"hgao327",
"sizhit2",
"abheesht17",
"jiangyangmu"
];
let issueNumber;
// Grab the correct number depending on the event type
if (context.payload.issue) {
issueNumber = context.payload.issue.number;
} else if (context.payload.pull_request) {
issueNumber = context.payload.pull_request.number;
} else {
console.log('Not an issue or PR');
return;
}
console.log('Assignee list:', assigneesList);
console.log('Entered auto assignment for this issue/PR:', issueNumber);
// Calculate total weeks since Unix Epoch (Jan 1, 1970)
const now = new Date();
const weekCount = Math.floor(now.getTime() / 604800000);
// Pick the assignee
const selection = weekCount % assigneesList.length;
const assigneeForIssue = assigneesList[selection];
console.log(`Issue/PR Number = ${issueNumber}, assigning to: ${assigneeForIssue}`);
// Return the API call so the YAML 'await' works correctly
return github.rest.issues.addAssignees({
issue_number: issueNumber,
owner: context.repo.owner,
repo: context.repo.repo,
assignees: [assigneeForIssue],
});
};
```
## /.github/workflows/issue-auto-assign.yaml
```yaml path="/.github/workflows/issue-auto-assign.yaml"
name: Auto-assign issues and PRs
on:
issues:
types: [opened, reopened]
pull_request_target:
types: [opened, reopened]
permissions:
contents: read
issues: write
pull-requests: write
jobs:
assign:
runs-on: ubuntu-latest
steps:
- name: Checkout Repository
uses: actions/checkout@v4
- name: Run Assignment Script
uses: actions/github-script@v7
with:
script: |
const script = require('./\.github/workflows/issue-auto-assign.js')
await script({github, context})
```
## /.github/workflows/pypi_release.yml
```yml path="/.github/workflows/pypi_release.yml"
# GitHub Actions workflow for building, testing, and publishing a Python package to PyPI
# using Trusted Publishing (OIDC) and a secure Environment.
name: Build, Test, and Publish Python Package to PyPI
# This triggers when a new "release" is published in the GitHub UI.
# This is safer than publishing on every push to 'main'.
on:
release:
types: [published]
permissions:
contents: read
jobs:
build:
name: Build package
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: '3.12'
- name: Install build tools
run: |
python -m pip install --upgrade pip build
- name: Build distributions
run: |
python -m build
- name: Store package artifacts
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.4.0
with:
name: python-package
path: dist/*
test:
name: Run unit tests
needs: [build]
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4 # v4.2.2
- uses: actions/setup-python@v5
with:
python-version: '3.12'
- name: Download package artifacts
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
with:
name: python-package
path: dist
- name: Install built package
run: |
python -m pip install --upgrade pip
WHEEL=$(ls -1 dist/google_tunix-*.whl | head -n1)
python -m pip install "${WHEEL}[test]"
- name: Run smoke tests
run: |
python3 -c "
import tunix
import tunix.models
import tunix.generate
import tunix.sft
import tunix.distillation
import tunix.rl
print('All tunix modules imported successfully')
"
publish-to-pypi:
name: Publish package to PyPI
needs: [test]
runs-on: ubuntu-latest
# 1. IMPORTANT: Specify the environment
# This tells GitHub to apply the 'release' environment's rules
# and to include 'environment: release' in the OIDC token.
environment: release
# 2. IMPORTANT: Set permissions for OIDC
# This gives the workflow write access to the OIDC token,
# which is required for trusted publishing.
permissions:
id-token: write # Required for trusted publishing
steps:
- name: Download package artifacts
# Retrieve the .whl and .tar.gz files from the 'build' job
# Note: We download again here; jobs run on fresh VMs.
uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
with:
name: python-package
path: dist/
- name: Publish package to PyPI
# This is the official action for PyPI Trusted Publishing
uses: pypa/gh-action-pypi-publish@release/v1
```
## /.github/workflows/tpu-nightly-regression.yml
```yml path="/.github/workflows/tpu-nightly-regression.yml"
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
name: Tunix Nightly Regression
on:
workflow_call:
secrets:
HF_TOKEN:
required: true
description: 'HuggingFace token for model downloads'
KAGGLE_USERNAME:
required: true
description: 'Kaggle username for model downloads'
KAGGLE_KEY:
required: true
description: 'Kaggle API key for model downloads'
concurrency:
# Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else
group: ${{ github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || github.run_id }}
cancel-in-progress: true
env:
HF_HOME: ~/.cache/huggingface
HF_HUB_ENABLE_HF_TRANSFER: "1"
jobs:
run_prod:
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: vllm/vllm-tpu:nightly-20260911-4b80060-51da0ca # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu,cpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Install tunix dependencies
run: |
python3 -m pip install --upgrade pip setuptools wheel
# Install Tunix with dev and test dependencies without overwriting the vLLM dependencies.
python3 -m pip install -e .[dev,test,cli]
# tpu-inference/Numba needs NumPy 2.3 or less.
python3 -m pip install numpy==2.3.5 --force-reinstall
- name: Verify TPU availability
run: |
python -c "
import jax
print(f'JAX version: {jax.__version__}')
print(f'JAX devices: {jax.devices()}')
# Check if we have TPU devices specifically
devices = jax.devices()
has_tpu = len(devices) > 0 and all(device.platform == 'tpu' for device in devices)
print(f'TPU available: {has_tpu}')
if not has_tpu:
print('ERROR: No TPU devices found! Expected TPU devices but got:', [device.platform for device in devices])
exit(1)
else:
print(f'SUCCESS: Found {len(devices)} TPU device(s)')
"
- name: Run SFT shell scripts
id: sft_tests
shell: bash
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
run: |
SCRIPT_DIR="./examples/sft/mtnt"
MAX_STEPS=5
EVAL_EVERY_N_STEPS=1
# Check if directory exists
if [ ! -d "$SCRIPT_DIR" ]; then
echo "Directory $SCRIPT_DIR does not exist"
exit 1
fi
echo "### SFT Scripts Summary" >> $GITHUB_STEP_SUMMARY
echo "| Script | Status |" >> $GITHUB_STEP_SUMMARY
echo "|---|---|" >> $GITHUB_STEP_SUMMARY
echo "🔍 Finding scripts in $SCRIPT_DIR"
FINAL_EXIT_CODE=0
FAILED_SCRIPTS=()
for script in "$SCRIPT_DIR"/*.sh; do
if [ -f "$script" ]; then
echo "::group::📦 Executing: $script"
echo "MAX_STEPS=$MAX_STEPS, EVAL_EVERY_N_STEPS=$EVAL_EVERY_N_STEPS"
chmod +x "$script"
if bash "$script" \
training_config.max_steps="$MAX_STEPS" \
training_config.eval_every_n_steps="$EVAL_EVERY_N_STEPS" \
rl_training_config.max_steps="$MAX_STEPS" \
rl_training_config.eval_every_n_steps="$EVAL_EVERY_N_STEPS"; then
echo "✅ Successfully completed: $script"
echo "| $script | ✅ Pass |" >> $GITHUB_STEP_SUMMARY
else
exit_code=$?
echo "❌ Failed to complete: $script (Exit Code: $exit_code)" >&2
echo "| $script | ❌ Fail |" >> $GITHUB_STEP_SUMMARY
FINAL_EXIT_CODE=1
FAILED_SCRIPTS+=("$script")
fi
echo "::endgroup::"
fi
done
if [ "$FINAL_EXIT_CODE" -ne 0 ]; then
echo "🚫 One or more SFT scripts failed:" >&2
for failed_script in "${FAILED_SCRIPTS[@]}"; do
echo " - $failed_script" >&2
done
exit 1
fi
echo "🎉 All SFT scripts completed execution."
- name: Run RL shell scripts
id: rl_tests
if: success() || failure()
shell: bash
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
KAGGLE_USERNAME: ${{ secrets.KAGGLE_USERNAME }}
KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
run: |
SCRIPT_DIRS=(
"./examples/rl/grpo/gsm8k"
"./examples/rl/ppo/gsm8k"
)
EXCLUDE_DIR="verl_compatible"
MAX_STEPS=5
EVAL_EVERY_N_STEPS=1
for dir in "${SCRIPT_DIRS[@]}"; do
if [ ! -d "$dir" ]; then
echo "Directory $dir does not exist" >&2
exit 1
fi
done
echo "### RL Scripts Summary" >> $GITHUB_STEP_SUMMARY
echo "| Script | Status |" >> $GITHUB_STEP_SUMMARY
echo "|---|---|" >> $GITHUB_STEP_SUMMARY
echo "🔍 Finding scripts in ${SCRIPT_DIRS[*]}, excluding $EXCLUDE_DIR"
FINAL_EXIT_CODE=0
FAILED_SCRIPTS=()
for SCRIPT_DIR in "${SCRIPT_DIRS[@]}"; do
while IFS= read -r script; do
if [ -f "$script" ]; then
# TODO(b/510383570)
if [[ "$script" == *"run_qwen3_8b"* ]]; then
echo "⏭️ Skipping: $script"
echo "| $script | ⏭️ Skipped |" >> $GITHUB_STEP_SUMMARY
continue
fi
# Fix run_qwen3_vllm_disagg.sh
if [[ "$script" == *"run_qwen3_vllm_disagg"* ]]; then
# Issue 1: HBM OOM
export SKIP_JAX_PRECOMPILE=1
# Issue 2: FAILED_PRECONDITION: The program continuator has halted unexpectedly.
export JAX_COMPILATION_CACHE_DIR=/tmp/jax_cache
rm -rf /tmp/jax_cache
else
unset SKIP_JAX_PRECOMPILE
unset JAX_COMPILATION_CACHE_DIR
fi
echo "::group::📦 Executing: $script"
echo "MAX_STEPS=$MAX_STEPS, EVAL_EVERY_N_STEPS=$EVAL_EVERY_N_STEPS"
chmod +x "$script"
if bash "$script" \
rl_training_config.max_steps="$MAX_STEPS" \
rl_training_config.eval_every_n_steps="$EVAL_EVERY_N_STEPS"; then
echo "✅ Successfully completed: $script"
echo "| $script | ✅ Pass |" >> $GITHUB_STEP_SUMMARY
else
exit_code=$?
echo "❌ Failed to complete: $script (Exit Code: $exit_code)" >&2
echo "| $script | ❌ Fail |" >> $GITHUB_STEP_SUMMARY
FINAL_EXIT_CODE=1
FAILED_SCRIPTS+=("$script")
fi
echo "::endgroup::"
fi
done < <(find "$SCRIPT_DIR" -name "*.sh" -type f | grep -v "$SCRIPT_DIR/$EXCLUDE_DIR/" | grep -E -v "7b|8b|12b")
done
if [ "$FINAL_EXIT_CODE" -ne 0 ]; then
echo "🚫 One or more RL scripts failed:" >&2
for failed_script in "${FAILED_SCRIPTS[@]}"; do
echo " - $failed_script" >&2
done
exit 1
fi
echo "🎉 All RL scripts completed execution."
- name: Run regression scripts
if: success() || failure()
shell: bash
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
id: regression_tests
run: |
echo "### Regression Scripts Summary" >> $GITHUB_STEP_SUMMARY
echo "| Script | Status |" >> $GITHUB_STEP_SUMMARY
echo "|---|---|" >> $GITHUB_STEP_SUMMARY
FAILED=0
FAILED_SCRIPTS=()
function run_and_log() {
local name="$1"
shift
echo "::group::📦 Executing: $name"
if "$@"; then
echo "✅ Successfully completed: $name"
echo "| $name | ✅ Pass |" >> $GITHUB_STEP_SUMMARY
else
echo "❌ Failed to complete: $name" >&2
echo "| $name | ❌ Fail |" >> $GITHUB_STEP_SUMMARY
FAILED=1
FAILED_SCRIPTS+=("$name")
fi
echo "::endgroup::"
}
# Download GSM8K dataset
mkdir -p /tmp/grpo_test/rl/grpo/data
run_and_log "grpo_demo_llama3_qwen2.py (colocated)" python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20
run_and_log "grpo_demo_llama3_qwen2.py (disaggregated-2-way)" python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --cluster-setup=disaggregated-2-way
run_and_log "grpo_demo_llama3_qwen2.py (disaggregated-3-way)" python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --cluster-setup=disaggregated-3-way
pip install 'jax[tpu]<0.11.0'
run_and_log "math_eval_nb.py" bash -c "NUM_BATCHES=1 python examples/deepscaler/math_eval_nb.py"
# TODO(b/508252632): Re-enable once the bug is fixed.
# vLLM Tests
# run_and_log "vllm grpo_demo_llama3_qwen2.py colocated mode" \
# python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=vllm
# run_and_log "vllm grpo_demo_llama3_qwen2.py 2 way disaggregated mode" \
# python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=vllm --cluster-setup=disaggregated-2-way
# run_and_log "vllm grpo_demo_llama3_qwen2.py 3 way disaggregated mode" \
# python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=vllm --cluster-setup=disaggregated-3-way
# run_and_log "vllm grpo_demo_llama3_qwen2.py vllm server mode in 2 way disaggregated mode" \
# python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=vllm --rollout-server-mode=True --cluster-setup=disaggregated-2-way
# run_and_log "vllm grpo_demo_llama3_qwen2.py with vllm server mode in 3 way disaggregated mode" \
# python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=vllm --rollout-server-mode=True --cluster-setup=disaggregated-3-way
# SGLang Tests
echo "::group::📦 Setting up SGLang"
unset JAX_PLATFORMS
pip list | egrep 'jax|flax|libtpu'
cd ..
git clone https://github.com/sgl-project/sglang-jax.git && cd sglang-jax/python && pip install -e . && cd ../..
pip list | egrep 'jax|flax|libtpu'
cd tunix
echo "::endgroup::"
run_and_log "sglang_jax colocated" bash -c "unset JAX_PLATFORMS && python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=sglang_jax"
run_and_log "sglang_jax disaggregated-2-way" bash -c "unset JAX_PLATFORMS && python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=sglang_jax --cluster-setup=disaggregated-2-way"
run_and_log "sglang_jax disaggregated-3-way" bash -c "unset JAX_PLATFORMS && python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=20 --rollout-engine=sglang_jax --cluster-setup=disaggregated-3-way"
if [ "$FAILED" -ne 0 ]; then
echo "🚫 One or more regression scripts failed:" >&2
for failed_script in "${FAILED_SCRIPTS[@]}"; do
echo " - $failed_script" >&2
done
exit 1
fi
echo "🎉 All regression scripts completed execution."
run_latest:
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: vllm/vllm-tpu:nightly # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu,cpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Install tunix dependencies and latest LKG vLLM
run: |
# Install compiler and build dependencies for compiling vLLM from source
apt-get update && apt-get install -y build-essential cmake ninja-build
# Ensure build requirements are upgraded and present
pip install packaging setuptools wheel setuptools-rust kagglehub kagglesdk --upgrade
pip install -e .[prod,test,cli]
# Fetch latest vllm lkg commit hash from tpu-inference repo
echo "Fetching latest vLLM LKG version..."
VLLM_COMMIT=$(curl -sSL https://raw.githubusercontent.com/vllm-project/tpu-inference/main/.buildkite/vllm_lkg.version)
if [ -z "$VLLM_COMMIT" ]; then
echo "ERROR: Failed to fetch vLLM LKG commit hash."
exit 1
fi
echo "Fetched vLLM LKG commit: $VLLM_COMMIT"
# Install vllm from lkg commit with no build isolation
export VLLM_TARGET_DEVICE=tpu
pip install --no-build-isolation "vllm @ git+https://github.com/vllm-project/vllm.git@${VLLM_COMMIT}"
# Install tpu-inference from main with no build isolation
pip install --no-build-isolation "tpu-inference @ git+https://github.com/vllm-project/tpu-inference.git@main"
# reinstall qwix==0.1.6 after tpu-inference is installed
pip install --no-deps "qwix==0.1.6" "flax>=0.12.5"
- name: Verify TPU availability
run: |
python -c "
import jax
print(f'JAX version: {jax.__version__}')
print(f'JAX devices: {jax.devices()}')
devices = jax.devices()
has_tpu = len(devices) > 0 and all(device.platform == 'tpu' for device in devices)
print(f'TPU available: {has_tpu}')
if not has_tpu:
print('ERROR: No TPU devices found!')
exit(1)
else:
print(f'SUCCESS: Found {len(devices)} TPU device(s)')
"
- name: Run latest tests
shell: bash
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p /tmp/grpo_test/rl/grpo/data
FAILED=0
FAILED_SCRIPTS=()
echo "📦 Executing Non-Agentic CLI (vLLM Batch Mode, Colocated) ..."
python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=2 --rollout-engine=vllm || { FAILED=1; FAILED_SCRIPTS+=("Non-Agentic CLI (vLLM Batch Mode, Colocated)"); }
pkill -9 -f "vllm" || true
sleep 5
echo "📦 Executing Non-Agentic CLI (vLLM Server Mode, Disaggregated 2-way) ..."
python scripts/grpo_demo_llama3_qwen2.py --root-dir=/tmp/grpo_test --num-batches=2 --rollout-engine=vllm --rollout-server-mode=True --cluster-setup=disaggregated-2-way || { FAILED=1; FAILED_SCRIPTS+=("Non-Agentic CLI (vLLM Server Mode, Disaggregated 2-way)"); }
pkill -9 -f "vllm" || true
sleep 5
echo "📦 Executing Agentic CLI (Qwen2.5 0.5B, vLLM Server Mode, Disaggregated 4:4) ..."
model_name="qwen2.5-0.5b" model_id="Qwen/Qwen2.5-0.5B-Instruct" checkpoint_dir="/tmp/qwen3_ckpts_server" num_batches=2 num_train_epochs=1 total_tpus=8 train_mesh="(4,1)" rollout_mesh="(4,1)" \
bash examples/rl/grpo/gsm8k/run_qwen3_8b_disagg.sh || { FAILED=1; FAILED_SCRIPTS+=("Agentic CLI (Qwen2.5 0.5B, vLLM Server Mode, Disaggregated 4:4)"); }
if [ "$FAILED" -ne 0 ]; then
echo "❌ One or more run_latest tests failed:" >&2
for failed_script in "${FAILED_SCRIPTS[@]}"; do
echo " - $failed_script" >&2
done
exit 1
else
echo "✅ All run_latest tests completed successfully."
fi
# TODO(tunix-dev): Re-enable once the bug is fixed.
# - name: Run Agentic CLI MaxText Test (Qwen3 0.6B, Disaggregated)
# id: maxtext_test
# if: success() || failure()
# shell: bash
# env:
# HF_TOKEN: ${{ secrets.HF_TOKEN }}
# run: |
# # Install MaxText if not present
# python3 -c "import maxtext" 2>/dev/null || python3 -m pip install git+https://github.com/AI-Hypercomputer/maxtext.git
# # Install necessary dependencies for MaxText
# python3 -m pip install aqtp tokamax
# pkill -9 -f "vllm" || true
# sleep 5
# echo "📦 Executing Agentic CLI (Qwen3 0.6B, MaxText, vLLM Server Mode, Disaggregated 4:4) ..."
# model_name="qwen3-0.6b" model_id="Qwen/Qwen3-0.6B" tokenizer_path="Qwen/Qwen3-0.6B" base_emb_dim=1024 \
# checkpoint_dir="/tmp/qwen3_maxtext_ckpts_server" num_batches=2 num_train_epochs=1 total_tpus=8 train_mesh="(4,1)" rollout_mesh="(4,1)" \
# bash examples/rl/grpo/gsm8k/run_qwen3_8b_disagg_maxtext.sh
```
## /.github/workflows/tpu-tests.yml
```yml path="/.github/workflows/tpu-tests.yml"
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
name: TPU Tests
on:
workflow_call:
secrets:
HF_TOKEN:
description: 'HuggingFace token for model downloads'
KAGGLE_USERNAME:
description: 'Kaggle Username'
KAGGLE_KEY:
description: 'Kaggle API Key'
concurrency:
# Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else
group: ${{ github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || github.run_id }}
cancel-in-progress: true
env:
HF_HOME: ~/.cache/huggingface
HF_HUB_ENABLE_HF_TRANSFER: "1"
jobs:
# ---------------------------------------------------------------------------
# Job: run_core
# Purpose: Multi-device TPU unit test suite on candidate JAX stable stack.
# Hardware: Live TPU v6e-8 accelerator.
# Secrets: None required.
# Scope: Sharded model, generation, SFT, distillation, and RL unit tests.
# ---------------------------------------------------------------------------
run_core:
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Install tunix dependencies
run: |
pip install --upgrade pip
pip install -e ".[test]"
- name: Verify TPU availability
run: |
python -c "
import jax
print(f'JAX version: {jax.__version__}')
print(f'JAX devices: {jax.devices()}')
# Check if we have TPU devices specifically
devices = jax.devices()
has_tpu = len(devices) > 0 and all(device.platform == 'tpu' for device in devices)
print(f'TPU available: {has_tpu}')
if not has_tpu:
print('ERROR: No TPU devices found! Expected TPU devices but got:', [device.platform for device in devices])
exit(1)
else:
print(f'SUCCESS: Found {len(devices)} TPU device(s)')
"
- name: Run tunix generation tests (PASSED only)
run: |
# tokenizer_adapter_test requires access to gated repo
python -m pytest tests/generate/ -v --tb=short \
--ignore=tests/generate/vllm_sampler_test.py \
--ignore=tests/generate/vllm_sampler_qwen_test.py \
--ignore=tests/generate/vllm_driver_test.py \
--ignore=tests/generate/tokenizer_adapter_test.py \
--ignore=tests/generate/sglang_jax_sampler_test.py \
--ignore=tests/generate/sglang_jax_lora_test.py
- name: Run tunix SFT tests
run: |
python -m pytest tests/sft/ -v --tb=short
- name: Run tunix distillation tests
run: |
python -m pytest tests/distillation/ -v --tb=short
- name: Run tunix RL cluster tests
run: |
python -m pytest tests/rl/rl_cluster_test.py -v --tb=short
- name: Run tunix tests not covered by the above categories
run: |
# This category is to catch tests added but not covered by CI yet. Whenever you add new folders under tests/, please add a new category above and skip those tests here.
python -m pytest tests/ -v --tb=short --ignore=tests/perf/ --ignore=tests/model_alignment/ --ignore=tests/models/ --ignore=tests/cli/ --ignore=tests/utils/mesh_utils_test.py --ignore=tests/utils/topology_test.py --ignore=tests/generate/ --ignore=tests/sft/ --ignore=tests/distillation/ --ignore=tests/rl/ --ignore=tests/smoke_tests/ --ignore=tests/experimental/ || code=$?
if [ "${code:-0}" = "5" ]; then
echo "No tests collected (expected)."
exit 0
else
exit "${code:-0}"
fi
# ---------------------------------------------------------------------------
# Job: run_dev
# Purpose: Multi-device TPU training and PEFT integration tests.
# Hardware: Live TPU v6e-8 accelerator.
# Secrets: Requires HF_TOKEN and KAGGLE_* secrets (skipped on external fork PRs).
# Scope: End-to-end GRPO RL demo training and SFT PEFT 2x2 mesh integration tests.
# ---------------------------------------------------------------------------
run_dev:
if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: vllm/vllm-tpu:nightly-20260911-4b80060-51da0ca # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu,cpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Setup Tunix , tpu-inference and dependencies
run: |
echo "Current directory:"
pwd
python3 -m pip install --upgrade pip setuptools wheel
# Install Tunix with dev and test dependencies without overwriting the vLLM dependencies.
python3 -m pip install -e .[test]
# tpu-inference/Numba needs NumPy 2.3 or less.
python3 -m pip install numpy==2.3.5 --force-reinstall
- name: GRPO Integration Test (Llama3 / Qwen2 multi-device RL training demo)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
# Run GRPO demo script with minimal configuration
python3 scripts/grpo_demo_llama3_qwen2.py \
--root-dir=/tmp/grpo_test \
--num-batches=2 \
--num-test-batches=1 \
--global-batch-size=2 \
--train-mini-batch-size=2 \
--train-micro-batch-size=2 \
--rollout-engine=vanilla
- name: Run tunix SFT integration tests (PEFT with 2x2 mesh on TPU)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
# Update dependencies
python3 -m pip install --upgrade kagglehub kagglesdk
# Loading tfds requires tensorflow.
python3 -m pip install tensorflow
export JAX_PLATFORMS=tpu,cpu
./tests/sft/sft_tpu_smoke_test.sh
# ---------------------------------------------------------------------------
# Job: run_vllm
# Purpose: Live inference sampler tests with vLLM engine on TPU hardware.
# Hardware: Live TPU v6e-8 accelerator with vLLM TPU container.
# Secrets: Requires HF_TOKEN for downloading gated models (Llama-3.2, Qwen3).
# Scope: vLLM Sampler generation, LoRA mapping, and checkpoint loading on TPU.
# ---------------------------------------------------------------------------
run_vllm:
if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: vllm/vllm-tpu:nightly-20260911-4b80060-51da0ca # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu,cpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Setup Tunix , tpu-inference and dependencies
run: |
echo "Current directory:"
pwd
python3 -m pip install --upgrade pip setuptools wheel
# Install Tunix with dev and test dependencies without overwriting the vLLM dependencies.
python3 -m pip install -e .[test]
# tpu-inference/Numba needs NumPy 2.3 or less.
python3 -m pip install numpy==2.3.5 --force-reinstall
- name: Run vllm tests
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
unset JAX_PLATFORMS
python3 -m pytest tests/generate/vllm_sampler_qwen_test.py -v --tb=short
python3 -m pytest tests/generate/vllm_sampler_test.py --collect-only -q --no-header --no-summary --disable-warnings | grep '::' > test_collections.txt
while read -r test; do
python3 -m pytest -s "$test" -v --tb=short
done < test_collections.txt
# ---------------------------------------------------------------------------
# Job: run_sglang
# Purpose: SGLang-JAX integration and sampler tests on TPU hardware.
# Hardware: Live TPU v6e-8 accelerator with isolated candidate TPU container.
# Secrets: Requires HF_TOKEN for model downloads.
# Scope: SGLang sampler, LoRA, and RL cluster integration tests.
# ---------------------------------------------------------------------------
run_sglang:
if: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository }}
runs-on: [linux-x86-ct6e-180-8tpu]
environment: testing
container:
image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest # zizmor: ignore[unpinned-images]
options: --privileged
env:
CLOUD_TPU_ACCELERATOR: v6e-8
JAX_PLATFORMS: tpu,cpu
steps:
# Cache Hugging Face hub
- name: Cache HF hub
uses: actions/cache@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
with:
path: ~/.cache/huggingface
key: hf-${{ runner.os }}-${{ hashFiles('pyproject.toml', 'requirements*.txt', 'constraints*.txt') }}
restore-keys: |
hf-${{ runner.os }}-
- name: Checkout code
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
fetch-depth: 0
persist-credentials: false
- name: Install tunix dependencies
run: |
pip install --upgrade pip
pip install -e ".[test]"
- name: Run install sglang-jax && test
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
## because sglang-jax has codes like jax.local_devices('cpu')
# TODO(lancewang): Re-enable this test once the bug is fixed.
unset JAX_PLATFORMS
python3 -m pip list | egrep 'jax|flax|libtpu'
cd ..
git clone https://github.com/sgl-project/sglang-jax.git && cd sglang-jax/python && python3 -m pip install -e . --force-reinstall && cd ../..
# SGLang Jax removed qwix dependency in https://github.com/sgl-project/sglang-jax/pull/734
python3 -m pip install qwix --force-reinstall
python3 -m pip install 'jax[tpu]<0.11.0' --force-reinstall
# TODO(b/470113586): Remove this test once the bug is fixed.
python3 -m pip list | egrep 'jax|flax|libtpu'
cd tunix
python3 tests/generate/sglang_jax_sampler_test.py
python3 tests/generate/sglang_jax_lora_test.py
python3 -m pytest tests/rl/rl_cluster_test.py -v --tb=short -k "not vllm"
```
## /.gitignore
```gitignore path="/.gitignore"
*.pyc
*.so
*.egg-info
*.whl
/dist/
.ipynb_checkpoints
.DS_Store
.mypy_cache/
.pytype/
.idea
.vscode
.envrc
# virtualenv/venv directories
/venv/
/bin/
/include/
/lib/
/share/
# Documentation
/docs/_build/
/docs/_collections/
/docs/jupyter_execute
/docs/sg_execution_times.rst
/docs/generated
build/*
```
## /.pre-commit-config.yaml
```yaml path="/.pre-commit-config.yaml"
repos:
- repo: https://github.com/google/pyink
rev: '25.12.0'
hooks:
- id: pyink
types_or: [python, jupyter]
files: '\.(py|ipynb){{contextString}}#39;
additional_dependencies:
- "pyink[jupyter]"
- repo: https://github.com/pylint-dev/pylint
rev: v4.0.5
hooks:
- id: pylint
name: pylint
types: [python]
verbose: true
args:
- --rcfile=.pylintrc
# Prevents pylint from failing the pre-commit hook. This is necessary
# because pylint is very strict.
- --exit-zero
exclude: ^tests/
- repo: https://github.com/nbQA-dev/nbQA
rev: 1.9.1
hooks:
- id: nbqa-pylint
stages: [manual]
additional_dependencies: [pylint==3.3.8]
args: [--rcfile=.pylintrc-notebooks, --fail-under=10.0]
files: '\.ipynb{{contextString}}#39;
- repo: https://github.com/executablebooks/mdformat
rev: 1.0.0
hooks:
- id: mdformat
args: ["--wrap", "80"]
additional_dependencies:
- mdformat-gfm==0.3.5
- "markdown-it-py[linkify]"
# (optional) other plugins:
# - mdformat-frontmatter
# - mdformat-tables
files: '\.md{{contextString}}#39;
- repo: https://github.com/PyCQA/isort
rev: 6.0.1
hooks:
- id: isort
name: isort
types: [python]
exclude: ^tests/
```
## /.pylintrc
```pylintrc path="/.pylintrc"
# This pylintrc file contains a best-effort configuration to uphold the
# best-practices and style described in the Google Python style guide:
# https://google.github.io/styleguide/pyguide.html
#
# Its canonical open-source location is:
# https://google.github.io/styleguide/pylintrc
[MASTER]
# Add files or directories to the blacklist. They should be base names, not
# paths.
ignore=third_party
# Add files or directories matching the regex patterns to the blacklist. The
# regex matches against base names, not paths.
ignore-patterns=
# Pickle collected data for later comparisons.
persistent=no
# List of plugins (as comma separated values of python modules names) to load,
# usually to register additional checkers.
load-plugins=
# Use multiple processes to speed up Pylint.
jobs=4
# Allow loading of arbitrary C extensions. Extensions are imported into the
# active Python interpreter and may run arbitrary code.
unsafe-load-any-extension=no
# A comma-separated list of package or module names from where C extensions may
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code
extension-pkg-whitelist=
[MESSAGES CONTROL]
# Only show warnings with the listed confidence levels. Leave empty to show
# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED
confidence=
# Enable the message, report, category or checker with the given id(s). You can
# either give multiple identifier separated by comma (,) or put this option
# multiple time (only on the command line, not in the configuration file where
# it should appear only once). See also the "--disable" option for examples.
#enable=
# Disable the message, report, category or checker with the given id(s). You
# can either give multiple identifiers separated by comma (,) or put this
# option multiple times (only on the command line, not in the configuration
# file where it should appear only once).You can also use "--disable=all" to
# disable everything first and then reenable specific checks. For example, if
# you want to run only the similarities checker, you can use "--disable=all
# --enable=similarities". If you want to run only the classes checker, but have
# no Warning level messages displayed, use"--disable=all --enable=classes
# --disable=W"
disable=apply-builtin,
arguments-differ,
attribute-defined-outside-init,
backtick,
bad-option-value,
buffer-builtin,
c-extension-no-member,
chained-comparison,
cmp-builtin,
cmp-method,
coerce-builtin,
coerce-method,
consider-iterating-dictionary,
consider-merging-isinstance,
consider-refactoring-into-while-condition,
consider-ternary-expression,
consider-using-any-or-all,
consider-using-augmented-assign,
consider-using-dict-comprehension,
consider-using-dict-items,
consider-using-dict-literal,
consider-using-enumerate,
consider-using-f-string,
consider-using-from-import,
consider-using-generator,
consider-using-in,
consider-using-join,
consider-using-list-comprehension,
consider-using-max-with-key,
consider-using-min-with-key,
consider-using-namedtuple-or-dataclass,
consider-using-set-comprehension,
consider-using-sys-exit,
consider-using-tuple,
consider-using-walrus,
consider-using-with,
delslice-method,
div-method,
duplicate-code,
eq-without-hash,
execfile-builtin,
file-builtin,
filter-builtin-not-iterating,
fixme,
getslice-method,
global-statement,
hex-method,
idiv-method,
implicit-str-concat-in-sequence,
import-error,
import-outside-toplevel,
import-self,
import-star-module-level,
input-builtin,
intern-builtin,
invalid-str-codec,
locally-disabled,
long-builtin,
long-suffix,
map-builtin-not-iterating,
metaclass-assignment,
next-method-called,
next-method-defined,
no-absolute-import,
no-else-break,
no-else-continue,
no-else-raise,
no-else-return,
no-member,
no-self-use,
nonzero-method,
oct-method,
old-division,
old-ne-operator,
old-octal-literal,
old-raise-syntax,
parameter-unpacking,
print-statement,
raising-string,
range-builtin-not-iterating,
raw_input-builtin,
rdiv-method,
redefined-argument-from-local,
reduce-builtin,
relative-import,
reload-builtin,
round-builtin,
setslice-method,
signature-differs,
standarderror-builtin,
suppressed-message,
sys-max-int,
too-few-public-methods,
too-many-ancestors,
too-many-arguments,
too-many-boolean-expressions,
too-many-branches,
too-many-instance-attributes,
too-many-locals,
too-many-public-methods,
too-many-return-statements,
too-many-statements,
too-many-positional-arguments,
trailing-newlines,
unichr-builtin,
unicode-builtin,
unnecessary-comprehension,
unnecessary-lambda-assignment,
unnecessary-pass,
unpacking-in-except,
use-dict-literal,
useless-else-on-loop,
useless-suppression,
using-cmp-argument,
wrong-import-order,
xrange-builtin,
zip-builtin-not-iterating,
[REPORTS]
# Set the output format. Available formats are text, parseable, colorized, msvs
# (visual studio) and html. You can also give a reporter class, eg
# mypackage.mymodule.MyReporterClass.
output-format=text
# Tells whether to display a full report or only the messages
reports=no
# Python expression which should return a note less than 10 (10 is the highest
# note). You have access to the variables errors warning, statement which
# respectively contain the number of errors / warnings messages and the total
# number of statements analyzed. This is used by the global evaluation report
# (RP0004).
evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
# Template used to display messages. This is a python new-style format string
# used to format the message information. See doc for all details
#msg-template=
[BASIC]
# Good variable names which should always be accepted, separated by a comma
good-names=main,_
# Bad variable names which should always be refused, separated by a comma
bad-names=
# Colon-delimited sets of names that determine each other's naming style when
# the name regexes allow several styles.
name-group=
# Include a hint for the correct naming format with invalid-name
include-naming-hint=no
# List of decorators that produce properties, such as abc.abstractproperty. Add
# to this list to register other decorators that produce valid properties.
property-classes=abc.abstractproperty,cached_property.cached_property,cached_property.threaded_cached_property,cached_property.cached_property_with_ttl,cached_property.threaded_cached_property_with_ttl
# Regular expression matching correct function names
function-rgx=^(?:(?P<exempt>setUp|tearDown|setUpModule|tearDownModule)|(?P<camel_case>_?[A-Z][a-zA-Z0-9]*)|(?P<snake_case>_?[a-z][a-z0-9_]*))$
# Regular expression matching correct variable names
variable-rgx=^[a-z][a-z0-9_]*$
# Regular expression matching correct constant names
const-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
# Regular expression matching correct attribute names
attr-rgx=^_{0,2}[a-z][a-z0-9_]*$
# Regular expression matching correct argument names
argument-rgx=^[a-z][a-z0-9_]*$
# Regular expression matching correct class attribute names
class-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$
# Regular expression matching correct inline iteration names
inlinevar-rgx=^[a-z][a-z0-9_]*$
# Regular expression matching correct class names
class-rgx=^_?[A-Z][a-zA-Z0-9]*$
# Regular expression matching correct module names
module-rgx=^(_?[a-z][a-z0-9_]*|__init__)$
# Regular expression matching correct method names
method-rgx=(?x)^(?:(?P<exempt>_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase|tearDownTestCase|setupSelf|tearDownClass|setUpClass|(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next)|(?P<camel_case>_{0,2}[A-Z][a-zA-Z0-9_]*)|(?P<snake_case>_{0,2}[a-z][a-z0-9_]*))$
# Regular expression which should only match function or class names that do
# not require a docstring.
no-docstring-rgx=(__.*__|main|test.*|.*test|.*Test)$
# Minimum line length for functions/classes that require docstrings, shorter
# ones are exempt.
docstring-min-length=10
[TYPECHECK]
# List of decorators that produce context managers, such as
# contextlib.contextmanager. Add to this list to register other decorators that
# produce valid context managers.
contextmanager-decorators=contextlib.contextmanager,contextlib2.contextmanager
# Tells whether missing members accessed in mixin class should be ignored. A
# mixin class is detected if its name ends with "mixin" (case insensitive).
ignore-mixin-members=yes
# List of module names for which member attributes should not be checked
# (useful for modules/projects where namespaces are manipulated during runtime
# and thus existing member attributes cannot be deduced by static analysis. It
# supports qualified module names, as well as Unix pattern matching.
ignored-modules=
# List of class names for which member attributes should not be checked (useful
# for classes with dynamically set attributes). This supports the use of
# qualified names.
ignored-classes=optparse.Values,thread._local,_thread._local
# List of members which are set dynamically and missed by pylint inference
# system, and so shouldn't trigger E1101 when accessed. Python regular
# expressions are accepted.
generated-members=
[FORMAT]
# Maximum number of characters on a single line.
max-line-length=80
# TODO(https://github.com/PyCQA/pylint/issues/3352): Direct pylint to exempt
# lines made too long by directives to pytype.
# Regexp for a line that is allowed to be longer than the limit.
ignore-long-lines=(?x)(
^\s*(\#\ )?<?https?://\S+>?$|
^\s*(from\s+\S+\s+)?import\s+.+$|
.*pytype:.*$
)
# Allow the body of an if to be on the same line as the test if there is no
# else.
single-line-if-stmt=yes
# Maximum number of lines in a module
max-module-lines=99999
# String used as indentation unit. The internal Google style guide mandates 2
# spaces. Google's externaly-published style guide says 4, consistent with
# PEP 8. Here, we use 2 spaces, for conformity with many open-sourced Google
# projects (like TensorFlow).
indent-string=' '
# Number of spaces of indent required inside a hanging or continued line.
indent-after-paren=4
# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
expected-line-ending-format=
[MISCELLANEOUS]
# List of note tags to take in consideration, separated by a comma.
notes=TODO
[VARIABLES]
# Tells whether we should check for unused import in __init__ files.
init-import=no
# A regular expression matching the name of dummy variables (i.e. expectedly
# not used).
dummy-variables-rgx=^\*{0,2}(_$|unused_|dummy_)
# List of additional names supposed to be defined in builtins. Remember that
# you should avoid to define new builtins when possible.
additional-builtins=
# List of strings which can identify a callback function by name. A callback
# name must start or end with one of those strings.
callbacks=cb_,_cb
# List of qualified module names which can have objects that can redefine
# builtins.
redefining-builtins-modules=six,six.moves,past.builtins,future.builtins,functools
[LOGGING]
# Logging modules to check that the string format arguments are in logging
# function parameter format
logging-modules=logging,absl.logging
[SIMILARITIES]
# Minimum lines number of a similarity.
min-similarity-lines=4
# Ignore comments when computing similarities.
ignore-comments=yes
# Ignore docstrings when computing similarities.
ignore-docstrings=yes
# Ignore imports when computing similarities.
ignore-imports=no
[SPELLING]
# Spelling dictionary name. Available dictionaries: none. To make it working
# install python-enchant package.
spelling-dict=
# List of comma separated words that should not be checked.
spelling-ignore-words=
# A path to a file that contains private dictionary; one word per line.
spelling-private-dict-file=
# Tells whether to store unknown words to indicated private dictionary in
# --spelling-private-dict-file option instead of raising a message.
spelling-store-unknown-words=no
[IMPORTS]
# Deprecated modules which should not be used, separated by a comma
deprecated-modules=regsub,
TERMIOS,
Bastion,
rexec,
sets
# Create a graph of every (i.e. internal and external) dependencies in the
# given file (report RP0402 must not be disabled)
import-graph=
# Create a graph of external dependencies in the given file (report RP0402 must
# not be disabled)
ext-import-graph=
# Create a graph of internal dependencies in the given file (report RP0402 must
# not be disabled)
int-import-graph=
# Force import order to recognize a module as part of the standard
# compatibility libraries.
known-standard-library=
# Force import order to recognize a module as part of a third party library.
known-third-party=enchant, absl
# Analyse import fallback blocks. This can be used to support both Python 2 and
# 3 compatible code, which means that the block might have code that exists
# only in one or another interpreter, leading to false positives when analysed.
analyse-fallback-blocks=no
[CLASSES]
# List of method names used to declare (i.e. assign) instance attributes.
defining-attr-methods=__init__,
__new__,
setUp
# List of member names, which should be excluded from the protected access
# warning.
exclude-protected=_asdict,
_fields,
_replace,
_source,
_make
# List of valid names for the first argument in a class method.
valid-classmethod-first-arg=cls,
class_
# List of valid names for the first argument in a metaclass class method.
valid-metaclass-classmethod-first-arg=mcs
[EXCEPTIONS]
# Exceptions that will emit a warning when being caught. Defaults to
# "Exception"
overgeneral-exceptions=builtins.Exception,
builtins.BaseException
```
## /.pylintrc-notebooks
```pylintrc-notebooks path="/.pylintrc-notebooks"
# pylintrc-notebooks
# Derived from Google's pylintrc, adjusted for Jupyter notebooks.
[MESSAGES CONTROL]
disable = missing-module-docstring,
missing-function-docstring,
wrong-import-position,
ungrouped-imports,
redefined-outer-name,
wrong-import-order,
import-error,
[FORMAT]
indent-string=' '
```
## /Dockerfile
``` path="/Dockerfile"
# Base image with Python 3.12
FROM python:3.12-slim
# Set environment variables to non-interactive to avoid prompts during installation
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC
# Install system dependencies, including Python 3 and pip
RUN apt-get update && \
apt-get install -y build-essential curl git python3 python3-pip && \
rm -rf /var/lib/apt/lists/*
# Upgrade pip
RUN python3 -m pip install --upgrade pip
# Create a virtual environment
RUN python3.12 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
# Upgrade pip
RUN pip install --upgrade pip
# Install uv
RUN pip install uv
RUN pip install git+https://github.com/ayaka14732/jax-smi.git
# If you encounter a checkpoint issue, try using following old version of pathways-utils.
# RUN pip install git+https://github.com/AI-Hypercomputer/pathways-utils.git@b72729bb152b7b3426299405950b3af300d765a9#egg=pathwaysutils
RUN pip install gcsfs
RUN pip install wandb
# Set the working directory
WORKDIR /app
# Copy scripts and requirements first to leverage Docker cache
COPY scripts/install_tunix_vllm_requirement.sh scripts/
COPY requirements/ requirements/
RUN bash scripts/install_tunix_vllm_requirement.sh
# Copy pyproject.toml and README.md to install dependencies first
COPY pyproject.toml README.md /app/
RUN mkdir /app/tunix && touch /app/tunix/__init__.py
RUN uv pip install .
# Install SFT/MaxText dependencies (unconditional)
RUN uv pip install --upgrade flax && \
uv pip install torchax aqtp tokamax math_verify drjax && \
uv pip install --no-deps git+https://github.com/google/maxtext.git
# Build argument to conditionally install Kubernetes tools
ARG INSTALL_K8S_TOOLS=false
# Install gcloud, kubectl, k9s
RUN if [ "$INSTALL_K8S_TOOLS" = "true" ]; then \
apt-get update && \
apt-get install -y vim lsof procps apt-transport-https ca-certificates gnupg && \
(echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list) && \
(curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --batch --yes --no-tty --dearmor -o /usr/share/keyrings/cloud.google.gpg) && \
apt-get update && apt-get install -y google-cloud-cli google-cloud-cli-gke-gcloud-auth-plugin kubectl && \
(curl -sS https://webinstall.dev/k9s | bash) && \
rm -rf /var/lib/apt/lists/*; \
fi
# Build argument to conditionally install DeepSWE evaluation dependencies
ARG INSTALL_DEEPSWE_DEPS=false
# Install DeepSWE specific dependencies and apply runtime patches conditionally
RUN if [ "$INSTALL_DEEPSWE_DEPS" = "true" ]; then \
uv pip install kubernetes gym swebench==3.0.2 && \
uv pip install --no-deps git+https://github.com/kubernetes-sigs/agent-sandbox.git#subdirectory=clients/python/agentic-sandbox-client && \
uv pip install --no-deps git+https://github.com/kubernetes-sigs/agent-sandbox.git#subdirectory=examples/agent-sandbox-rl && \
uv pip install --no-deps git+https://github.com/r2e-gym/r2e-gym.git@0d94c4eb9431cd195c55a7ea3abd54006c9a1735 && \
sed -i 's/create_repo, upload_folder, HfFolder/create_repo, upload_folder/' /opt/venv/lib/python3.12/site-packages/r2egym/agenthub/utils/utils.py && \
sed -i 's/self.commit = ParsedCommit(\*\*json.loads(self.commit_json))/self.commit = ParsedCommit(\*\*(json.loads(self.commit_json) if isinstance(self.commit_json, str) else self.commit_json))/' /opt/venv/lib/python3.12/site-packages/r2egym/agenthub/runtime/docker.py; \
fi
# Build argument to conditionally install MaxText dependencies
ARG INSTALL_MAXTEXT=false
# Install MaxText specific dependencies conditionally
RUN if [ "$INSTALL_MAXTEXT" = "true" ]; then \
uv pip install -r /app/requirements/maxtext_requirements.txt --torch-backend=cpu; \
fi
# Build argument to conditionally install Raiden weight sync dependencies
ARG INSTALL_RAIDEN=false
ARG RAIDEN_WHEEL_DIR=/app/raiden_wheels
# Install Raiden specific dependencies conditionally
COPY raiden_wheels/ ${RAIDEN_WHEEL_DIR}/
RUN if [ "$INSTALL_RAIDEN" = "true" ]; then \
if [ -d "$RAIDEN_WHEEL_DIR" ] && ls "$RAIDEN_WHEEL_DIR"/*.whl 1>/dev/null 2>&1; then \
pip install --force-reinstall --no-deps "$RAIDEN_WHEEL_DIR"/*.whl; \
else \
pip install keyrings.google-artifactregistry-auth && \
pip install tpu-raiden-jax --extra-index-url https://us-python.pkg.dev/cloud-tpu-inference-test/tpu-raiden/simple/; \
fi; \
fi
# Force install numpy version to avoid version conflicts.
RUN uv pip install numpy==2.3.5
# Copy the rest of the project files
COPY . .
# Compile proto buffer for discovery service, this has to be the last step.
RUN uv pip install grpcio-tools
RUN cd /app && find tunix/experimental/distributed -name "*.proto" -exec python -m grpc_tools.protoc -I/app --python_out=/app --grpc_python_out=/app {} +
# Install Tunix in editable mode
RUN uv pip install --no-deps -e .
# Set the default command to bash
CMD ["bash"]
```
## /README.md
# Tunix: A Lightweight LLM Post-Training Library
<div align="left">
<a href="https://tunix.readthedocs.io/en/latest/index.html"><img src="https://img.shields.io/badge/documentation-blue"></a>
</div>
**Tunix (Tune-in-JAX)** is a JAX based library designed to streamline the
post-training of Large Language Models. It provides efficient and scalable
support for:
- **SOTA Training performance on TPUs**
- **Supervised Fine-Tuning**
- **Reinforcement Learning (RL)**
- **Agentic RL**
Tunix leverages the power of JAX for accelerated computation and seamless
integration with JAX-based modeling frameworks like
[Flax NNX](https://flax.readthedocs.io/en/latest/nnx_basics.html), and
integrates with high-performance inference engines like vLLM and SGLang-JAX for
rollout. **For our detailed documentation, please refer to the [Tunix Website](https://tunix.readthedocs.io/en/latest/index.html)**.
**Current Status: V2 Release**
Tunix is under active development. Our team is actively working on expanding its
capabilities, usability and performance. Stay tuned for upcoming updates and new
features! See [Talks and Announcements](https://tunix.readthedocs.io/en/latest/talks.html) for latest updates, talks, and blog posts.
## High Level Architecture
Tunix serves as a state-of-the-art post-training library within the JAX training
stack, positioned to leverage foundational tools like Flax, Optax, Orbax, etc.
for efficient model refinement. It sits as an intermediate layer between these
core utilities and optimized models like MaxText and MaxDiffusion, streamlining
tuning workflows on top of the XLA and JAX infrastructure. See [Design Overview](https://tunix.readthedocs.io/en/latest/design.html) for more details on the architecture.
## Key Features
- **[Supervised Fine-Tuning (SFT)](https://tunix.readthedocs.io/en/latest/algorithms.html)**:
- Full Weights Fine-Tuning
- [PEFT](https://tunix.readthedocs.io/en/latest/performance.html#peft-with-lora) (Parameter-Efficient
Fine-Tuning)
- [DPO](https://arxiv.org/abs/2305.18290) (Direct Preference Optimization)
- [ORPO](https://arxiv.org/abs/2403.07691) (Odds Ratio Preference Optimization)
- **[Reinforcement Learning (RL)](https://tunix.readthedocs.io/en/latest/algorithms.html)**:
- [PPO](https://arxiv.org/abs/1707.06347) (Proximal Policy Optimization)
- [GRPO](https://arxiv.org/abs/2402.03300) (Group Relative Policy
Optimization)
- [GSPO-Token](https://arxiv.org/abs/2507.18071) (Token-level Group
Sequence Policy Optimization)
- [DAPO](https://arxiv.org/abs/2503.14476) (Direct Alignment via Preference
Optimization)
- [Dr.GRPO](https://arxiv.org/abs/2503.20783) (Distributionally Robust
GRPO)
- **[Agentic RL](https://tunix.readthedocs.io/en/latest/agentic_rl.html)**:
- Multi-turn tool use
- Asynchronous rollout for high-throughput trajectory collection
- Trajectory batching and grouping
## News
- [2026/04] Gemma4 models are supported in Tunix! Stay tuned for upcoming training recipes.
- [2026/01] Tunix model now supports efficient kernel execution ([splash attn](https://github.com/google/tunix/blob/main/tunix/models/qwen3/model.py#L150-L151), [GMM MoE](https://github.com/google/tunix/blob/main/tunix/models/qwen3/model.py#L638)).
- [2025/12] [Agentic RL Training](https://github.com/google/tunix/tree/main/tunix/rl/agentic) has been released, with efficient support of multi-turn agent-env interaction, tool usage, async rollout, etc.
## Framework & Infra Highlights
- **Modularity**:
- Components are designed to be reusable and composable
- Easy to customize and extend
- **Performance & Efficiency**:
- Native [vLLM](https://tunix.readthedocs.io/en/latest/rollout.html#vllm) and
[SGLang-JAX](https://tunix.readthedocs.io/en/latest/rollout.html#sglang) on TPU integration for performant
rollout
- Native [MaxText](https://github.com/AI-Hypercomputer/maxtext) model
integration for high performance kernels and model execution
- [Micro-batching](https://tunix.readthedocs.io/en/latest/performance.html#batching-config) support for component
level efficient execution
- **Stability**
- Seamless multi-host distributed training with Pathways which can scale
up to thousands of devices
- [Checkpointing and Fault Tolerance](https://tunix.readthedocs.io/en/latest/reliability.html)
## Getting Started
**Installation:** Jump to [Installation](https://tunix.readthedocs.io/en/latest/quickstart.html#installation) to install Tunix and run your first training
job.
For TPU users integrating `vllm` and `tpu-inference`, there are two supported
setup paths:
- Docker image builds use [Dockerfile](https://github.com/google/tunix/blob/main/Dockerfile) and install
the pinned dependencies directly from `requirements/requirements.txt` and
`requirements/special_requirements.txt`.
- Local TPU VM or developer-machine installs can use
[scripts/install_tunix_vllm_requirement.sh](https://github.com/google/tunix/blob/main/scripts/install_tunix_vllm_requirement.sh),
which installs the same requirement files outside Docker.
These are separate entry points. If you are building the Docker image, you do
not need to run the install script inside the container build.
**Examples:** To get started, we have a number of detailed examples and tutorials. You can see [Quick Start](https://tunix.readthedocs.io/en/latest/quickstart.html) for a great set of starting examples and [Examples and Guides](https://tunix.readthedocs.io/en/latest/examples.html) for a comprehensive list of all the notebooks and examples we have.
## Supported Models
Tunix supports a growing list of models including Gemma, Llama, and Qwen families.
See [Models](https://tunix.readthedocs.io/en/latest/models.html) for a full list and details on how to add new ones.
## Contributing and Feedback
We welcome contributions! As Tunix is in early development, the contribution
process is still being formalized. The detailed contribution process is outlined
[here](https://tunix.readthedocs.io/en/latest/contributing.html). In
the meantime, you can make feature requests, report issues and ask questions in
our
[Tunix GitHub discussion forum](https://github.com/google/tunix/discussions).
## Collaborations and Partnership
[GRL](https://github.com/lmgame-org/GRL/blob/tunix_integration_dev/README.md)
(Game Reinforcement Learning), developed by
[Hao AI Lab](https://hao-ai-lab.github.io/) from UCSD, is an open-source
framework for post-training large language models through multi-turn RL on
challenging games. In collaboration with Tunix, GRL integrates seamless TPU
support—letting users quickly run scalable, reproducible RL experiments (like
PPO rollouts on Qwen2.5-0.5B-Instruct) on TPU v4 meshes with
[minimal setup](https://github.com/lmgame-org/GRL/blob/tunix_integration_dev/README.md#5-launch-the-quick-test-defaults-to-qwen2505b-supports-4-tpu-v4-with-mesh-22).
This partnership empowers the community to push LLM capabilities further,
combining Tunix’s optimized TPU runtime with GRL’s flexible game RL pipeline for
cutting-edge research and easy reproducibility.
## Citing Tunix
```bibtex
@misc{tunix2025,
title={Tunix (Tune-in-JAX)},
author={Bao, Tianshu and Carpenter, Jeff and Chai, Lin and Gao, Haoyu and Jiang, Yangmu and Noghabi, Shadi and Sharma, Abheesht and Tan, Sizhi and Wang, Lance and Yan, Ann and Yu, Weiren and others},
year={2025},
howpublished={\url{https://github.com/google/tunix}},
}
```
## Acknowledgements
Thank you to all our wonderful contributors!
[](https://github.com/google/tunix/graphs/contributors)
## /build_docker.sh
```sh path="/build_docker.sh"
# This scripts takes a docker image that already contains the GRL dependencies, copies the local source code in and
# uploads that image into GCR. Once in GCR the docker image can be used for development.
# Each time you update the base image via a "bash docker_build_dependency_image.sh", there will be a slow upload process
# (minutes). However, if you are simply changing local code and not updating dependencies, uploading just takes a few seconds.
# Script to buid a GRL base image locally, example cmd is:
# bash build_docker.sh
set -e
INSTALL_MAXTEXT=false
INSTALL_RAIDEN=false
RAIDEN_WHEEL_DIR=/app/raiden_wheels
INSTALL_DEEPSWE_DEPS=false
usage() {
cat <<'MSG'
Usage: bash build_docker.sh [--maxtext] [--raiden] [--raiden-wheel-dir PATH] [--deepswe]
Options:
--maxtext Install MaxText-specific dependencies.
--raiden Install Raiden-specific dependencies.
--raiden-wheel-dir PATH Use prebuilt Raiden wheels from PATH inside the Docker build context.
--deepswe Install DeepSWE evaluation dependencies.
MSG
}
while [[ "$#" -gt 0 ]]; do
case $1 in
--maxtext) INSTALL_MAXTEXT=true; shift ;;
--raiden) INSTALL_RAIDEN=true; shift ;;
--raiden-wheel-dir)
if [[ -z "$2" ]]; then
echo "Error: --raiden-wheel-dir requires a path"
usage
exit 1
fi
RAIDEN_WHEEL_DIR="$2"
shift 2
;;
--deepswe) INSTALL_DEEPSWE_DEPS=true; shift ;;
--help|-h)
usage
exit 0
;;
*) echo "Unknown parameter: $1"; usage; exit 1 ;;
esac
done
DOCKERFILE=./Dockerfile
if [ ! -f "$DOCKERFILE" ]; then
echo "Error: Dockerfile not found at $DOCKERFILE"
exit 1
fi
export LOCAL_IMAGE_NAME=tunix_base_image
echo "Building base image: $LOCAL_IMAGE_NAME"
echo "Using Dockerfile: $DOCKERFILE"
# Use Docker BuildKit so we can cache pip packages.
export DOCKER_BUILDKIT=1
echo "Starting to build your docker image. This will take a few minutes but the image can be reused as you iterate."
build_ai_image() {
COMMIT_HASH=$(git rev-parse --short HEAD)
echo "Building Tunix Image at commit hash ${COMMIT_HASH}..."
DOCKER_COMMAND="docker"
if docker info >/dev/null 2>&1; then
DOCKER_COMMAND="docker"
else
# Avoid invoking sudo interactively which can prompt for a password.
# Check whether non-interactive sudo would work (no password).
if sudo -n docker info >/dev/null 2>&1; then
DOCKER_COMMAND="sudo docker"
else
cat <<'MSG'
Docker does not appear usable from this account and the build would prompt for a password.
Run the build with sufficient privileges (will prompt): sudo bash build_docker.sh
On Linux, add your user to the docker group so sudo isn't required (you must re-login):
sudo usermod -aG docker "$USER" && newgrp docker
MSG
exit 1
fi
fi
$DOCKER_COMMAND build \
--network=host \
--build-arg INSTALL_MAXTEXT=${INSTALL_MAXTEXT} \
--build-arg INSTALL_RAIDEN=${INSTALL_RAIDEN} \
--build-arg RAIDEN_WHEEL_DIR=${RAIDEN_WHEEL_DIR} \
--build-arg INSTALL_DEEPSWE_DEPS=${INSTALL_DEEPSWE_DEPS} \
-t ${LOCAL_IMAGE_NAME} \
-f ${DOCKERFILE} .
}
build_ai_image
echo ""
echo "*************************
"
echo "Built your docker image and named it ${LOCAL_IMAGE_NAME}.
It now installs Tunix and the pinned vLLM and tpu-inference dependencies from requirements/requirements.txt. "
```
## /docs/_static/custom.css
```css path="/docs/_static/custom.css"
.wy-side-nav-search .wy-nav-top {
background: transparent;
}
```
## /docs/_static/img/sphx_glr_local_module_thumb.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/_static/img/sphx_glr_local_module_thumb.png
## /docs/_static/img/tunix.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/_static/img/tunix.png
## /docs/agentic_rl.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Agentic RL
## Architecture
## Core Components
The framework consists of several key components:
* **Agent**: Interacts with the environment by generating actions based on
observations and conversation history.
* **Environment**: Represents the task or problem to be solved, processes
agent actions, and returns observations, rewards, and termination signals.
* **Tool**: A reusable component that provides specific functionalities (e.g.,
calculation, search) that an agent can invoke.
* **Parser**: Translates between natural language model responses and
structured data like tool calls, and formats conversation history into
model-specific input formats.
* **TrajectoryCollectEngine**: Manages the interaction loop for a single
agent-environment pair to produce a complete trajectory.
* **RolloutOrchestrator**: Manages multiple `TrajectoryCollectEngine`
instances for parallel trajectory collection.
### Agents
Agents inherit from `ConversationAgentBase`, which provides common functionality
for maintaining conversation history (`chat_completions`) and recording
interaction steps in a `Trajectory`.
* **`ModelAgent`**: A simple agent for single-turn tasks where the model's
response is treated as the final answer.
* **`ToolAgent`**: A more complex agent that can parse model responses to
detect and invoke tool calls. It uses a `ToolManager` to manage available
tools and a `ToolParser` (e.g., `QwenToolParser`, `GeminiToolParser`) to
understand model outputs and format tool schemas for the model prompt.
### Environments
Environments inherit from `BaseTaskEnv`, which handles episode lifecycle
management like `max_steps`.
* **`TaskEnvironment`**: Designed for single-turn tasks. The environment
terminates after the first agent action and computes a reward based on the
final response using a provided `reward_fn`.
* **`ToolEnvironment`**: Designed for multi-turn, tool-using tasks. It
receives actions from the `ToolAgent`, and if they are tool calls, it uses
its internal `ToolManager` to execute them via `execute_calls`. The results
of tool execution are returned to the agent as a new observation in
`{"tool_outputs": ...}` format. The episode terminates when the agent
invokes a special `finish` function or `max_steps` is reached, at which
point `reward_fn` is called on the final answer.
### Tool Integration
Tools inherit from `BaseTool` and must implement `get_json_schema()` to define
their interface (parameters, description) and either `apply()` (synchronous) or
`apply_async()` (asynchronous) to define their logic. The `ToolManager`
discovers, registers, and executes tools by name. It can execute multiple tool
calls in parallel for efficiency.
### Agent/Environment interaction
--------------------------------------------------------------------------------
## Key Features and Optimizations
### Multi-turn Tool Use
Tunix fully supports multi-turn interactions involving tool use. The typical
flow is:
1. The `ToolAgent` sends the conversation history (including user query and
prior tool results) to the LLM.
2. The LLM responds with a tool call, e.g., `calculator(a=1, b=1)`.
3. The `ToolAgent` uses its `ToolParser` to parse this into an `Action` object.
4. The `ToolEnvironment` receives this action, uses its `ToolManager` to
execute `calculator`, and receives the result "2".
5. The `ToolEnvironment` returns an observation like `{"tool_outputs":
{"call_id_123": "Tool returned result: 2"}}`, reward=1, and done=False.
6. The `ToolAgent` adds the tool result to its history as a `role: tool`
message.
7. The loop continues until the agent calls `finish(answer=...)` or `max_steps`
is reached.
### Asynchronous Rollouts
To accelerate trajectory collection, Tunix supports asynchronous rollouts via
the `RolloutOrchestrator`. It leverages Python's `asyncio` to manage multiple
concurrent agent-environment interactions using `TrajectoryCollectEngine`
instances, with parallelism controlled by `max_concurrency`. The
`run_producers_from_stream` method manages a pool of workers that draw
agent-environment pairs from a stream, run full episodes via `collect()`, and
queue the resulting trajectories. The `yield_batches` method allows a consumer
(like an RL learner) to receive trajectories as they are generated. This
parallel execution significantly speeds up data collection, especially when
interacting with external models or tools with high latency.
Furthermore, Tunix provides a `RolloutSyncLock` to manage concurrency between
rollouts and model weight synchronization in distributed training setups. This
lock ensures that rollouts (`acquire_rollout`) are temporarily paused when a
weight sync (`acquire_weight_sync`) is requested, preventing agents from
generating trajectories with stale parameters.
### Trajectory Batching and Grouping
Tunix supports batching of agentic trajectories through the `GroupQueueManager`.
This component, used within the `RolloutOrchestrator`, collects `TrajectoryItem`
instances into buckets based on a configurable `group_key` (e.g., prompt ID via
`env.task["group_id"]`) and `episode_id`. Once a bucket reaches a predefined
`group_size` (e.g., `num_generations` in GRPO), it is marked as a "ready group"
and made available for downstream processing by `yield_batches`. This mechanism
is essential for algorithms like GRPO which require multiple trajectory samples
for each prompt, and improves efficiency by yielding full groups of trajectories
in batches. The `max_open_buckets` parameter can be used to limit memory usage
by controlling the number of groups that can be populated simultaneously.
## /docs/algorithms.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Algorithms
Tunix supports a wide array of SOTA algorithms for RL and SFT. Its modular design
also allows users to easily extend Tunix with custom algorithms, as described
further below.
## Supported Algorithms
* **Supervised Fine-Tuning (SFT) & Preference**
* **[PEFT](performance.md#peft-with-lora)** (Parameter-Efficient Fine-Tuning)
* **[DPO](https://arxiv.org/abs/2305.18290)** (Direct Preference Optimization)
* **[ORPO](https://arxiv.org/abs/2403.07691)** (Odds ratio Preference Optimization)
* **Reinforcement Learning (RL)**
* **[PPO](https://arxiv.org/abs/1707.06347)** (Proximal Policy Optimization)
* **[GRPO](https://arxiv.org/abs/2402.03300)** (Group Relative Policy Optimization)
* **[GSPO-Token](https://arxiv.org/abs/2507.18071)** (Token-level Group Sequence Policy Optimization)
* **[DAPO](https://arxiv.org/abs/2503.14476)** (Direct Alignment via Preference Optimization)
* **[Dr.GRPO](https://arxiv.org/abs/2503.14476)** (Distributionally Robust GRPO)
## Add a New RL Algorithm
Tunix is designed to be highly extensible. You can introduce new algorithms by
subclassing `AlgorithmConfig` (or its descendants) and implementing a
corresponding Learner.
The system uses a **parallel inheritance** pattern: extending a Configuration
often requires extending a Learner to consume it.
### Class Hierarchy & Interaction Diagram
<!-- TODO(b/475597805): Add better formatted diagram. -->
```text
CONFIGURATION LEARNER (The Engine)
(Defines Params) (Orchestrates Execution)
+-------------------+ +-----------------------+
| AlgorithmConfig | <---(binds)-- | RLLearner |
+-------------------+ +-----------------------+
^ ^ |
| | +---(Uses)---> [Function Registry]
(Inheritance) (Inheritance) (Loss, Advantage, Reward)
| |
+-------------------+ +-----------------------+
| GRPOConfig | <---(binds)-- | GRPOLearner |
+-------------------+ +-----------------------+
^ ^
| |
+-------------------+ +-----------------------+
| DAPOConfig | <---(binds)-- | DAPOLearner |
+-------------------+ +-----------------------+
```
--------------------------------------------------------------------------------
### 1. Defining the Configuration & Learner
To add a new algorithm, you typically define a config (to hold your params) and
a learner (to use them).
**Step 1: The Configuration** Inherit from `AlgorithmConfig` (or a specific
child like `GRPOConfig` if your algorithm is a variant of it). Use
`__post_init__` to validate your new settings.
```python
@dataclasses.dataclass(slots=True, kw_only=True)
class MyNewAlgoConfig(AlgorithmConfig):
# 1. Identity
algo_variant: str = "my_new_algo"
# 2. Components (References strings in FunctionRegistry)
advantage_estimator: str = "gae"
policy_loss_fn: str = "my_custom_loss"
reward_manager: str = "sequence-level"
# 3. Custom Hyperparameters
my_hyperparam: float = 0.5
def __post_init__(self):
...
```
**Step 2: The Learner** Inherit from `RLLearner` (or `GrpoLearner` etc). This is
where you inject specific execution logic, such as modifying the training loop
or injecting custom reward functions.
```python
class MyNewAlgorithmLearner(RLLearner):
def __init__(self, rl_engine, algo_config: MyNewAlgoConfig, reward_fns, ...):
# Custom initialization (e.g., adding specific reward shaping)
if algo_config.my_hyperparam > 0.1:
reward_fns.append(my_custom_reward_fn)
super().__init__(
rl_engine=rl_engine,
algo_config=algo_config,
reward_fns=reward_fns,
...
)
```
--------------------------------------------------------------------------------
### 2. Custom Loss & Advantage (The Registry)
Tunix uses a **Function Registry** to manage mathematical components. This
allows you to hot-swap loss functions or advantage estimators in your config
without changing the Learner code.
**How to Register a New Loss Function:** Define your loss function and decorate
it with `@register_policy_loss_fn`.
```python
from tunix.registry import register_policy_loss_fn
@register_policy_loss_fn("my_custom_loss")
def compute_my_custom_loss(log_probs, advantages, **kwargs):
"""
Args:
log_probs: Tensor of log probabilities.
advantages: Tensor of calculated advantages.
Returns:
Scalar loss tensor.
"""
return -torch.mean(log_probs * advantages)
```
**Usage:** Once registered, simply reference it in your config: `policy_loss_fn:
str = "my_custom_loss"`
--------------------------------------------------------------------------------
### 3. Custom Reward Management
Rewards are handled by a **Manager** pattern. The Learner delegates to the
Manager to compute rewards from model output and log the results.
* **Reward Function:** A simple callable that calculates scores based on
completion texts (e.g., regex matching, length constraints, keyword
presence).
* **Reward Manager:** The orchestrator that calls reward functions, formats
the output, and handles logging.
**When to use what?**
* **Simple Case:** Just add a new function to the `reward_fns` list.
* **Complex Case:** Subclass `AbstractRewardManager` if you need custom
aggregation (e.g., weighted sums) or specialized logging strategies.
**Example: Custom Manager**
Below is an example of a manager that performs custom aggregation and injects
specific intermediate logs.
```python
class MyCustomRewardManager(AbstractRewardManager):
def __call__(self, prompts: List[str], completions: List[str], **kwargs) -> Dict[str, Any]:
"""
Orchestrates reward calculation.
"""
# 1. Run all reward functions
raw_scores = [fn(prompts, completions) for fn in self.reward_fns]
# 2. Custom Aggregation (e.g., Multi-objective weighted sum)
final_rewards = self.aggregate_logic(raw_scores)
# 3. Calculate log metrics
log_metrics = self._prepare_log_metrics(prompts, completions, raw_scores, final_rewards)
# 4. Return format required by Learner
return {
"rewards": final_rewards,
"log_metrics": log_metrics
}
def _prepare_log_metrics(
self,
prompts: List[str],
completions: List[str],
rewards: np.ndarray,
sum_rewards: np.ndarray,
) -> Dict[str, Any]:
"""
Logs individual and summed rewards.
"""
# 1. Standard Logs (prompts, completions, sum, min, max)
metrics_to_log = super()._prepare_log_metrics(prompts, completions, rewards, sum_rewards)
# 2. Custom Intermediate Logging
# User may freely add intermediate reward logging results here.
# Example: Logging specific components of the reward signal separately
if hasattr(self, "reward_fns"):
for i, fn in enumerate(self.reward_fns):
# Log the specific contribution of each function (e.g. rewards/grammar_score)
name = getattr(fn, "__name__", f"fn_{i}")
metrics_to_log[f"rewards/{name}"] = (rewards[:, i], np.mean)
return metrics_to_log
```
## /docs/api/api_distillation.rst
```rst path="/docs/api/api_distillation.rst"
Distillation
============
.. currentmodule:: tunix
.. autosummary::
DistillationTrainer
DistillationTrainingConfig
----
.. autoclass:: DistillationTrainer
----
.. autoclass:: DistillationTrainingConfig
```
## /docs/api/api_generation.rst
```rst path="/docs/api/api_generation.rst"
Generation
==========
.. currentmodule:: tunix
.. autosummary::
Sampler
CacheConfig
----
.. autoclass:: Sampler
----
.. autoclass:: CacheConfig
```
## /docs/api/api_rl.rst
```rst path="/docs/api/api_rl.rst"
Reinforcement learning (RL)
===========================
.. currentmodule:: tunix
.. autosummary::
GRPOConfig
GRPOLearner
RewardFn
PPOConfig
PPOLearner
ClusterConfig
RLEngine
RLTrainingConfig
Role
RolloutConfig
-------
.. autoclass:: GRPOConfig
-------
.. autoclass:: GRPOLearner
-------
.. autoclass:: RewardFn
-------
.. autoclass:: PPOConfig
-------
.. autoclass:: PPOLearner
-------
.. autoclass:: ClusterConfig
-------
.. autoclass:: RLEngine
-------
.. autoclass:: RLTrainingConfig
-------
.. autoclass:: Role
-------
.. autoclass:: RolloutConfig
```
## /docs/api/api_sft.rst
```rst path="/docs/api/api_sft.rst"
Supervised fine-tuning (SFT)
============================
.. currentmodule:: tunix
.. autosummary::
PeftTrainer
TrainingConfig
DPOTrainer
DPOTrainingConfig
MetricsLogger
MetricsLoggerOptions
-------
.. autoclass:: PeftTrainer
-------
.. autoclass:: TrainingConfig
-------
.. autoclass:: DPOTrainer
-------
.. autoclass:: DPOTrainingConfig
-------
.. autoclass:: MetricsLogger
-------
.. autoclass:: MetricsLoggerOptions
```
## /docs/code-of-conduct.md
# Code of Conduct
## Our Pledge
In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to making participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, gender identity and expression, level of
experience, education, socio-economic status, nationality, personal appearance,
race, religion, or sexual identity and orientation.
## Our Standards
Examples of behavior that contributes to creating a positive environment
include:
- Using welcoming and inclusive language
- Being respectful of differing viewpoints and experiences
- Gracefully accepting constructive criticism
- Focusing on what is best for the community
- Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
- The use of sexualized language or imagery and unwelcome sexual attention or
advances
- Trolling, insulting/derogatory comments, and personal or political attacks
- Public or private harassment
- Publishing others' private information, such as a physical or electronic
address, without explicit permission
- Disrespecting the community's time by sending spam or other unsolicited
commercial messages
- Other conduct which could reasonably be considered inappropriate in a
professional setting
## Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, or to ban temporarily or permanently any
contributor for other behaviors that they deem inappropriate, threatening,
offensive, or harmful.
## Scope
This Code of Conduct applies both within project spaces and in public spaces
when an individual is representing the project or its community. Examples of
representing a project or community include using an official project e-mail
address, posting via an official social media account, or acting as an appointed
representative at an online or offline event. Representation of a project may be
further defined and clarified by project maintainers.
This Code of Conduct also applies outside the project spaces when the Project
Steward has a reasonable belief that an individual's behavior may have a
negative impact on the project or its community.
## Conflict Resolution
We do not believe that all conflict is bad; healthy debate and disagreement
often yield positive results. However, it is never okay to be disrespectful or
to engage in behavior that violates the project’s code of conduct.
If you see someone violating the code of conduct, you are encouraged to address
the behavior directly with those involved. Many issues can be resolved quickly
and easily, and this gives people more control over the outcome of their
dispute. If you are unable to resolve the matter for any reason, or if the
behavior is threatening or harassing, report it. We are dedicated to providing
an environment where participants feel welcome and safe.
Reports should be directed to *\[PROJECT STEWARD NAME(s) AND EMAIL(s)\]*, the
Project Steward(s) for *\[PROJECT NAME\]*. It is the Project Steward’s duty to
receive and address reported violations of the code of conduct. They will then
work with a committee consisting of representatives from the Open Source
Programs Office and the Google Open Source Strategy team. If for any reason you
are uncomfortable reaching out to the Project Steward, please email
opensource@google.com.
We will investigate every complaint, but you may not receive a direct response.
We will use our discretion in determining when and how to follow up on reported
incidents, which may range from not taking action to permanent expulsion from
the project and project-sponsored spaces. We will notify the accused of the
report and provide them an opportunity to discuss it before any action is taken.
The identity of the reporter will be omitted from the details of the report
supplied to the accused. In potentially harmful situations, such as ongoing
harassment or threats to anyone's safety, we may take action without notice.
## Attribution
This Code of Conduct is adapted from the Contributor Covenant, version 1.4,
available at https://www.contributor-covenant.org/version/1/4/code-of-conduct/
## /docs/conf.py
```py path="/docs/conf.py"
"""Configuration file for the Sphinx documentation builder.
For the full list of built-in configuration values, see the documentation:
https://www.sphinx-doc.org/en/master/usage/configuration.html
"""
import logging
from sphinx.util import logging as sphinx_logging
# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
project = "Tunix"
copyright = "2025, Tunix Developers" # pylint: disable=redefined-builtin
author = "Tunix Developers"
# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
extensions = [
"myst_nb",
"sphinx_gallery.gen_gallery",
"sphinx_collections",
# api docs
"sphinx.ext.autodoc",
"sphinx.ext.autosummary",
"sphinx.ext.intersphinx",
"sphinx.ext.napoleon",
]
templates_path = ["_templates"]
exclude_patterns = [
"_build",
"Thumbs.db",
".DS_Store",
(
"_collections/examples/model_load/from_safetensor_load/*"
"_collections/examples/rl/README.md"
),
"_collections/examples/sft/**",
"_collections/examples/deepscaler/**",
]
source_suffix = [".rst", ".md", ".ipynb"]
# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
html_theme = "sphinx_book_theme"
html_static_path = ["_static"]
html_logo = "_static/img/tunix.png"
html_css_files = [
"custom.css",
]
html_theme_options = {
"show_toc_level": 2,
"repository_url": "https://github.com/google/tunix",
"use_repository_button": True, # add a "link to repository" button
"navigation_with_keys": False,
}
# -- Options for sphinx-gallery ----------------------------------------------
sphinx_gallery_conf = {
"examples_dirs": "_collections/examples", # path to your example scripts
"gallery_dirs": (
"_collections/gallery"
), # path to where to save gallery generated output
"filename_pattern": "*.py",
"ignore_pattern": r"rl/|sft/|deepscaler/",
}
# -- Options for myst -------------------------------------------------------
myst_heading_anchors = 3
myst_enable_extensions = [
"amsmath",
"dollarmath",
]
nb_execution_mode = "off"
nb_execution_allow_errors = False
nb_render_image_options = {}
nb_execution_excludepatterns = [
"*.ipynb",
]
# -- Options for sphinx-collections
collections = {
"examples": {
"driver": "copy_folder",
"source": "../examples",
"ignore": [
"model_load",
"rl",
"sft",
"deepscaler",
],
}
}
suppress_warnings = ["misc.highlighting_failure"]
# -- Options for the API reference
default_role = "py:obj"
napoleon_include_init_with_doc = False
autodoc_default_options = {
"members": True,
"imported-members": True,
"undoc-members": True,
}
intersphinx_mapping = {
"optax": ("https://optax.readthedocs.io/en/latest/", None),
"flax": ("https://flax.readthedocs.io/en/stable/", None),
"jax": ("https://docs.jax.dev/en/latest/", None),
}
class FilterSphinxWarnings(logging.Filter):
"""Filter autosummary 'duplicate object description' warnings.
These warnings are unnecessary as they do not cause missing documentation
or rendering issues, so it is safe to filter them out.
"""
def __init__(self, app):
self.app = app
super().__init__()
def filter(self, record: logging.LogRecord) -> bool:
msg = record.getMessage()
filter_out = ("duplicate object description",)
return not msg.strip().startswith(filter_out)
def setup(app):
"""Set up custom logging filters."""
logger = logging.getLogger("sphinx")
warning_handler, *_ = [
h
for h in logger.handlers
if isinstance(h, sphinx_logging.WarningStreamHandler)
]
warning_handler.filters.insert(0, FilterSphinxWarnings(app))
```
## /docs/contributing.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Contributing
Welcome! We appreciate your interest in contributing to Tunix. This guide
details how to contribute to the project in a way that is efficient for
everyone.
We follow
[Google's Open Source Community Guidelines](https://opensource.google/conduct/).
## Contributing code
### 1. Propose Changes in an Issue
Before starting on your contribution, please
[check for an existing issue](https://github.com/google/tunix/issues).
For significant changes, please
[open an issue](https://github.com/google/tunix/issues/new) to discuss your
proposal first. This allows the team to provide feedback and ensure the change
aligns with the project's goals.
For minor changes, such as documentation updates or simple bug fixes, you can
open a pull request directly.
All bug fixes must include a link to a
[Colab](https://colab.research.google.com/) notebook or a new unit test case
that clearly reproduces the error.
### 2. Make code changes
To begin coding, fork the repository and create a new branch from main.
#### Setting up a development environment
We recommend creating an isolated virtual environment before installing Tunix's
development dependencies. From the repository root you can run:
```bash
python3 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -e .[dev]
```
The `dev` extra pulls in the bleeding-edge dependencies we rely on during local
development. If you prefer to stick with released packages, skip the extra and
run `pip install -e .` instead.
### 3. Create a pull request
Once your changes are ready, open a pull request from your branch to the main
branch of the upstream Tunix repository. Please provide a clear title and
description, linking to the relevant issue if one exists.
### 4. Sign the Contributor License Agreement
If this is your first contribution, you will be prompted to sign the Google CLA
after submitting your pull request. You can review the agreement
[here](https://cla.developers.google.com/clas).
### 5. Code review
A project maintainer will review your pull request. Be prepared for one or more
rounds of comments and requested changes as we work with you to refine the
contribution.
During review and development, standard CPU unit tests, package builds, and
documentation checks run automatically on every pull request push.
### 6. TPU CI Validation (`ready-to-submit` label)
To conserve shared hardware resources and prevent long queue times on live TPU
runners, multi-device TPU integration tests do not run automatically on every
pull request update.
Once your pull request has completed code review, addressed all reviewer comments,
and is ready for final verification:
1. A project maintainer will apply the `ready-to-submit` label to the pull
request.
2. The `ready-to-submit` label automatically triggers the comprehensive TPU test
suite (`tunix_tpu_unit_tests`) running against live TPU accelerators.
3. If further commits are pushed to the pull request while the `ready-to-submit`
label is present, TPU tests will automatically re-run.
### 7. Internal Import and Merging
Once the pull request has been approved and all CI checks (including the TPU test
suite triggered by `ready-to-submit`) have passed:
1. The PR will be converted/imported into Google's internal repository.
2. The change is verified and approved from the internal review workflow.
3. The change is submitted internally, which automatically syncs to GitHub and
resolves the pull request.
**Important Note for Maintainers & Reviewers**: Do **not** click the "Merge" button
directly on the GitHub pull request UI. All PR merges must go through the internal
import and sync pipeline to keep internal and open-source repositories consistent.
Thank you for your contribution!
## Formatting and linting
We use [Pyink](https://github.com/google/pyink) and
[Pylint](https://github.com/pylint-dev/pylint) for formatting and linting,
respectively.
For the first time you are setting up the repo, please run `pre-commit install`.
Note that this needs to be done only once at the beginning.
Now, you go through the usual flow of pushing code:
```
git add .
git commit -m "<message>"
```
Whenever you run `git commit -m "<message>"`, the code is automatically
formatted, and lint error messages are displayed.
If there's any error, the commit will not go through. Most of the times, the
errors are fixed automatically.
Note: Pylint errors are not binding, i.e., your commit will not fail if you have
linting errors. This is because Pylint is very strict. It is, however, advised
that you address most of the errors.
Once you are done fixing your errors, re-run the following:
```
git add .
git commit -m "<message>" # This will not get logged as a duplicate commit.
```
In case you want to run the above manually on all files, you can do the
following:
```
pre-commit run --all-files
```
If you want to opt out of pre-commit, you can always do the following (but make
sure you run the pre-commit hooks manually):
```
git commit -m "<message>" --no-verify
```
## Documentation
The Tunix documentation website is built using
[Sphinx](https://www.sphinx-doc.org) and
[MyST](https://myst-parser.readthedocs.io/en/latest/). Documents can be written
in
[MyST Markdown syntax](https://myst-parser.readthedocs.io/en/latest/syntax/typography.html#syntax-core)
or
[reStructuredText](https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html).
### Building the documentation locally (optional)
If you are writing documentation for Tunix, you may want to preview the
documentation site locally to ensure things work as expected before a
deployment.
First, make sure you install the necessary dependencies. You can do this by
navigating to your local clone of the Tunix repo and running:
```bash
pip install ".[docs]"
```
Once the dependencies are installed, you can navigate to the `docs/` folder and
run:
```bash
make html
```
This will generate the documentation in the `docs/_build/html` directory. These
files can be opened in a web browser directly, or you can use a simple HTTP
server to serve the files. For example, you can run:
```bash
python -m http.server -d docs/_build/html
```
Then, open your web browser and navigate to `http://localhost:8000` to view the
documentation.
### Adding new documentation files
If you are adding a new document, make sure it is included in the `toctree`
directive corresponding to the section where the new document should live. For
example, if adding a new page, make sure it is listed in the `toctree` directive
in `docs/index.md`.
<!-- ### Documentation deployment
The Tunix documentation is deployed to [https://tunix.readthedocs.io](https://tunix.readthedocs.io) on any successful merge to the main branch. -->
## /docs/design.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Design Overview
This document provides an overview of Tunix (Tune-in-JAX), a JAX-native
framework for post-training large language models. It covers the overall layered
architecture and the training loop architectures for Supervised Fine-Tuning
(SFT), Reinforcement Learning (RL), and Agentic RL. Understanding this design
should help you better utilize and extend Tunix for model tuning.
## High Level Architecture
The diagram below illustrates the layered architecture of Tunix:
- **User Interface & Application Layer**: This top layer provides the primary
access points for users through CLI tools, interactive notebooks, and
configuration files to define and manage tuning experiments. See
[Launching Jobs](launching.md) for more details.
- **Algorithm & Workflow Layer**: This layer houses the core training
paradigms, including supervised fine-tuning, reinforcement learning, and
agentic workflows, to support diverse model optimization strategies. See
[Algorithms Supported](algorithms.md) for more details.
- **Core Components Layer**: Responsible for the inner workings of the system,
this layer handles training loop orchestration, metrics tracking, and the
complex distributed state management required for large-scale runs.
- **Foundation Frameworks & Integration Layer**: This layer integrates
powerful underlying libraries like JAX, Flax, and Optax to provide
high-performance computation, neural network primitives, and efficient
gradient optimization.
- **Hardware Layer**: The base of the stack manages the physical or virtual
compute resources, specifically optimized for Google Cloud TPUs, multi-host
GPU clusters, and CPU hosts.
## SFT
Supervised Fine-Tuning (SFT) is a fundamental technique in machine learning used
to adapt pre-trained models to specific downstream tasks by training them on a
labeled dataset. In Tunix, the SFT pipeline is designed to efficiently fine-tune
large language models (LLMs) using various datasets and optimization strategies,
leveraging the underlying JAX ecosystem.
The following diagram illustrates the Supervised Fine-Tuning (SFT) pipeline in
Tunix, showing the flow of data and control.
- **Config**: Manages configurations for all pipeline components, such as
dataset details (e.g., URL), model specifics (e.g., URL, tokenizer), trainer
settings (e.g., max steps, hyperparameters), and checkpointing (e.g., save
frequency). See [Config Explanation](launching.md#config-explanation) for
more details.
- **Dataset Iterator**: Manages the stream of data from the external Training
Dataset, ensuring the model receives processed batches during the training
process. Tunix supports various data sources like TFDS and Parquet, and also
allows for the integration of custom datasets.
- **Model**: Manages the initialized LLMs (e.g., nnx.Graph) loaded with
weights from external Model Params, serving as the core entity being
trained. Tunix supports a set of predefined models (e.g. gemma, llama, qwen)
and also allows for the integration of new custom models. See
[Models](models.md) for more details.
- **Trainer**: Orchestrates the training steps by coordinating interactions
between the model, data, and optimizer.
- **Optimizer**: Applies optimization algorithms (e.g., AdamW) to update model
parameters based on the calculated loss.
- **Checkpoint Manager**: Handles periodic saving of model states to external
storage for training recovery or model deployment. Tunix leverages the Orbax
library for robust and efficient checkpoint management.
- **Metrics Logger**: Captures performance data during training and exports it
to external databases for monitoring and analysis. See [Metrics](metrics.md)
for more details.
## RL
Reinforcement Learning (RL) is a paradigm where an agent learns to make
decisions by interacting with an environment to maximize a cumulative reward.
Tunix provides a framework for Reinforcement Learning, designed to support
various RL algorithms and optimization strategies, leveraging the underlying JAX
ecosystem.
The following diagram illustrates a typical Reinforcement Learning (RL) pipeline
in Tunix. The exact pipeline details can differ based on the specific RL
algorithm being used.
- **RL Config**: Provides the central hyperparameters and algorithmic settings
that initialize the entire pipeline and define the training objectives. See
[Config Explanation](launching.md#config-explanation) for more details.
- **Orchestrator**: Manages the global workflow, coordinating **Resource
Control**, monitoring progress via the **Metrics Logger**, and directing the
execution of the chosen RL algorithm (such as PPO or GRPO). See
[Algorithms Supported](algorithms.md) for more details.
- **Rollout Workers**: These workers generate sample trajectories from the
current model using optimized runtimes like **vLLM** or **SGLang**. This
ensures high-throughput data collection, crucial for efficient RL training.
See [Rollout](rollout.md) for more details.
- **Inference Workers**: These workers host inference models (e.g., critic,
reference, and reward models in PPO) to evaluate collected samples. They
compute rewards, reference log probabilities, and value estimates.
- **Train Data Queue**: This in-memory buffer collects evaluated samples from
the inference and rollout stages and streams them into the trainers.
- **Trainers**: Performs weight updates for the **Actor** and **Critic**
models. A subsequent **Weight Sync** pushes these updated parameters back to
the Rollout Workers, ensuring they use the latest model for the next
iteration.
The **Orchestrator** drives the overall RL training loop. Within each iteration,
**Rollout Workers** generate sample trajectories by interacting with the
environment. These trajectories are then passed to the **Inference Workers** for
evaluation, computing necessary values like rewards and log probabilities. The
evaluated samples are collected and used by the **Trainers** to update the Actor
and Critic models. Finally, a **Weight Sync** step propagates the updated model
parameters back to the Rollout Workers, preparing for the next iteration. This
entire process can be orchestrated in both synchronous and asynchronous manners
(e.g. async data loading, async rollout).
## Agentic RL
Agentic Reinforcement Learning in Tunix provides a framework for
training agents that can perform multi-turn reasoning and interact with external
tools. The design follows a standard RL paradigm where an **Agent** interacts
with an **Environment** over multiple steps to complete a task.
The core design supports agents that engage in **multi-turn conversations**,
breaking down complex problems into sequential steps of reasoning, tool
invocation, and response generation. Agents can leverage external **tools** (e.g.,
search, code execution, APIs) to gather information or perform actions; the
environment manages tool execution and feeds results back to the agent for
subsequent steps.
A key design focus is performance and scalability, achieved through a highly
**asynchronous architecture for both rollout and training**. The trajectory
collection pipeline is designed for high throughput, allowing many
agent-environment interactions to run concurrently. This design efficiently
overlaps model inference latency, I/O-bound tool execution, and reward
computations, maximizing hardware utilization (e.g., TPUs) and enabling
efficient online RL training.
The framework naturally supports **trajectory batching and grouping**, making it
compatible with various RL algorithms, including those like GRPO that require
multiple trajectory samples per prompt for robust learning.
## /docs/examples.md
# Examples and Guides
This section provides a high-level overview of the Colab notebooks, scripts, and
example directories.
All examples are located in this
[directory](https://github.com/google/tunix/tree/main/examples).
<table>
<thead>
<tr>
<th align="center">Category</th>
<th align="center">Name/Path</th>
<th align="center">Description</th>
</tr>
</thead>
<tbody>
<tr>
<td rowspan="4" align="center" valign="middle" style="text-align: center; vertical-align: middle;"><b>Colab Notebook</b></td>
<td><a href="https://github.com/google/tunix/tree/main/examples/qlora_gemma.ipynb"><code>qlora_gemma.ipynb</code></a></td>
<td>End-to-end tutorial on fine-tuning (SFT) Gemma 270M model for English-French translation using parameter-efficient LoRA and QLoRA techniques.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/grpo_gemma.ipynb"><code>grpo_gemma.ipynb</code></a></td>
<td>Reinforcement learning tutorial using Group Relative Policy Optimization (GRPO) to train the Gemma 3 1B IT model for math reasoning on the GSM8K benchmark.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/dpo_gemma.ipynb"><code>dpo_gemma.ipynb</code></a></td>
<td>Preference tuning using Direct Preference Optimization (DPO) to tune the Gemma 3 1B-IT model on the GSM8K dataset.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/logit_distillation.ipynb"><code>logit_distillation.ipynb</code></a></td>
<td>Demonstrates knowledge distillation from a Gemma 7B-IT teacher to a Gemma 2B-IT student for translation task.</td>
</tr>
<tr>
<td rowspan="6" align="center" valign="middle" style="text-align: center; vertical-align: middle;"><b>Script</b></td>
<td><a href="https://github.com/google/tunix/tree/main/examples/rl/grpo/gsm8k/"><code>rl/grpo/gsm8k/</code></a></td>
<td>Bash scripts for fine-tuning different models and presets (Gemma, Llama, etc.) on the GSM8K mathematical reasoning task using GRPO.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/rl/grpo/gsm8k/verl_compatible/"><code>rl/grpo/gsm8k/verl_compatible/</code></a></td>
<td>Bash scripts for GRPO-training on the GSM8K dataset to train with a verl-compatible setup.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/deepscaler/"><code>deepscaler/</code></a></td>
<td>Scripts and notebooks for reproducing the <a href="https://pretty-radio-b75.notion.site/DeepScaleR-Surpassing-O1-Preview-with-a-1-5B-Model-by-Scaling-RL-19681902c1468005bed8ca303013a4e2">Deepscaler experiment</a> (<code>train_deepscaler_nb.py</code>) and math evaluation.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/sft/mtnt/"><code>sft/mtnt/</code></a></td>
<td>Bash scripts for SFT examples on the MTNT translation task for Gemma, Llama, and Qwen models.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/model_load/"><code>model_load/</code></a></td>
<td>Examples for loading Gemma2 and Gemma3 models from safetensors format.</td>
</tr>
<tr>
<td><a href="https://github.com/google/tunix/tree/main/examples/agentic/"><code>agentic/</code></a></td>
<td>Examples and scripts for agentic workflows, with async rollout.</td>
</tr>
</tbody>
</table>
## /docs/gallery.rst
```rst path="/docs/gallery.rst"
.. _gallery:
🖼️ Example gallery
==================
.. toctree::
:glob:
:hidden:
:maxdepth: 1
_collections/examples/README
.. raw:: html
<div class="sphx-glr-thumbnails">
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="GRPO Demo">
.. only:: html
.. image:: _static/img/sphx_glr_local_module_thumb.png
:alt: GRPO Demo.
:doc:`_collections/examples/grpo_gemma`
.. raw:: html
<div class="sphx-glr-thumbnail-title">GRPO Demo</div>
</div>
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="Logit Distillation">
.. only:: html
.. image:: _static/img/sphx_glr_local_module_thumb.png
:alt: Knowledge Distillation with Tunix: Gemma 7B to Gemma 2B
:doc:`_collections/examples/logit_distillation`
.. raw:: html
<div class="sphx-glr-thumbnail-title">Knowledge Distillation with Tunix: Gemma 7B to Gemma 2B</div>
</div>
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="QLoRA Demo">
.. only:: html
.. image:: _static/img/sphx_glr_local_module_thumb.png
:alt: QLoRA Demo
:doc:`_collections/examples/qlora_gemma`
.. raw:: html
<div class="sphx-glr-thumbnail-title">QLoRA Demo</div>
</div>
.. raw:: html
<div class="sphx-glr-thumbcontainer" tooltip="DPO Demo with math (gsm8k)">
.. only:: html
.. image:: _static/img/sphx_glr_local_module_thumb.png
:alt: DPO Demo
:doc:`_collections/examples/dpo_gemma`
.. raw:: html
<div class="sphx-glr-thumbnail-title">DPO Demo with math (gsm8k)</div>
</div>
.. raw:: html
</div>
```
## /docs/guides/Makefile
``` path="/docs/guides/Makefile"
# Minimal makefile for Sphinx documentation
#
# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS ?=
SPHINXBUILD ?= sphinx-build
SOURCEDIR = .
BUILDDIR = _build
# Put it first so that "make" without argument is like "make help".
help:
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
.PHONY: help Makefile
clean:
rm -rf $(BUILDDIR)
rm -rf _collections/
html:
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
```
## /docs/images/agentic_agent:env.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/agentic_agent:env.png
## /docs/images/agentic_rollout_pipeline.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/agentic_rollout_pipeline.png
## /docs/images/batch_vs_async_rollout.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/batch_vs_async_rollout.png
## /docs/images/collocated_mode.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/collocated_mode.png
## /docs/images/disagged_mode.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/disagged_mode.png
## /docs/images/global_batch_flow.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/global_batch_flow.png
## /docs/images/mini_batch_flow.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/mini_batch_flow.png
## /docs/images/model_id_huggingface.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/model_id_huggingface.png
## /docs/images/model_id_kaggle.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/model_id_kaggle.png
## /docs/images/perf_metrics_TB.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/perf_metrics_TB.png
## /docs/images/perf_metrics_perfetto.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/perf_metrics_perfetto.png
## /docs/images/tunix_architecture.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/tunix_architecture.png
## /docs/images/tunix_in_jax_ecosystem.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/tunix_in_jax_ecosystem.png
## /docs/images/tunix_reshard_algorithm.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/tunix_reshard_algorithm.png
## /docs/images/tunix_rl_pipeline.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/tunix_rl_pipeline.png
## /docs/images/tunix_sft_pipeline.png
Binary file available at https://raw.githubusercontent.com/google/tunix/refs/heads/main/docs/images/tunix_sft_pipeline.png
## /docs/index.md
```{include} intro.md
:start-line: 2
```
```{toctree}
:maxdepth: 2
quickstart.md
design.md
agentic_rl.md
performance.md
reliability.md
launching.md
rollout.md
algorithms.md
models.md
metrics.md
examples.md
talks.md
contributing.md
code-of-conduct.md
gallery.rst
```
```{eval-rst}
.. toctree::
:hidden:
:caption: 📖 Reference
:maxdepth: 2
api/api_sft
api/api_rl
api/api_distillation
api/api_generation
```
## /docs/intro.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Tunix: A Lightweight LLM Post-Training Library
**Tunix (Tune-in-JAX)** is a JAX based library designed to streamline the
post-training of Large Language Models. It provides efficient and scalable
support for:
- **SOTA Training performance on TPUs**
- **Supervised Fine-Tuning**
- **Reinforcement Learning (RL)**
- **Agentic RL**
Tunix leverages the power of JAX for accelerated computation and seamless
integration with JAX-based modeling frameworks like
[Flax NNX](https://flax.readthedocs.io/en/latest/nnx_basics.html), and
integrates with high-performance inference engines like vLLM and SGLang-JAX for
rollout.
**Current Status: V2 Release**
Tunix is under active development. Our team is actively working on expanding its
capabilities, usability and performance. Stay tuned for upcoming updates and new
features! See [Talks and Announcements](talks.md) for latest updates, talks, and blog posts.
## High Level Architecture
Tunix serves as a state-of-the-art post-training library within the JAX training
stack, positioned to leverage foundational tools like Flax, Optax, Orbax, etc.
for efficient model refinement. It sits as an intermediate layer between these
core utilities and optimized models like MaxText and MaxDiffusion, streamlining
tuning workflows on top of the XLA and JAX infrastructure.
See [Design Overview](design.md) for more details on the architecture.
## Key Features
- **[Supervised Fine-Tuning (SFT)](algorithms.md)**:
- Full Weights Fine-Tuning
- [PEFT](performance.md#peft-with-lora) (Parameter-Efficient Fine-Tuning)
- [DPO](https://arxiv.org/abs/2305.18290) (Direct Preference Optimization)
- [ORPO](https://arxiv.org/abs/2403.07691) (Odds Ratio Preference
Optimization)
- **[Reinforcement Learning (RL)](algorithms.md)**:
- [PPO](https://arxiv.org/abs/1707.06347) (Proximal Policy Optimization)
- [GRPO](https://arxiv.org/abs/2402.03300) (Group Relative Policy
Optimization)
- [GSPO-Token](https://arxiv.org/abs/2507.18071) (Token-level Group
Sequence Policy Optimization)
- [DAPO](https://arxiv.org/abs/2503.14476) (Direct Alignment via
Preference Optimization)
- [Dr.GRPO](https://arxiv.org/abs/2503.20783) (Distributionally Robust
GRPO)
- **[Agentic RL](agentic_rl.md)**:
- Multi-turn tool use
- Asynchronous rollout for high-throughput trajectory collection
- Trajectory batching and grouping
## Framework & Infra Highlights
- **Modularity**:
- Components are designed to be reusable and composable
- Easy to customize and extend
- **Performance & Efficiency**:
- Native [vLLM](rollout.md#vllm) and [SGLang-JAX](rollout.md#sglang) on
TPU integration for performant rollout
- Native [MaxText](https://github.com/AI-Hypercomputer/maxtext) model
integration for high performance kernels and model execution
- [Micro-batching](performance.md#batching-config) support for component
level efficient execution
- **Stability**
- Seamless multi-host distributed training with Pathways which can scale
up to thousands of devices
- [Checkpointing and Fault Tolerance](reliability.md)
## Getting Started
**Installation:** Jump to [Installation](https://tunix.readthedocs.io/en/latest/quickstart.html#installation) to install Tunix and run your first training
job.
For TPU users integrating `vllm` and `tpu-inference`, there are two supported
setup paths:
- Docker image builds use [Dockerfile](https://github.com/google/tunix/blob/main/Dockerfile) and install
the pinned dependencies directly from `requirements/requirements.txt` and
`requirements/special_requirements.txt`.
- Local TPU VM or developer-machine installs can use
[scripts/install_tunix_vllm_requirement.sh](https://github.com/google/tunix/blob/main/scripts/install_tunix_vllm_requirement.sh),
which installs the same requirement files outside Docker.
These are separate entry points. If you are building the Docker image, you do
not need to run the install script inside the container build.
**Examples:** To get started, we have a number of detailed examples and tutorials. You can see [Quick Start](https://tunix.readthedocs.io/en/latest/quickstart.html) for a great set of starting examples and [Examples and Guides](https://tunix.readthedocs.io/en/latest/examples.html) for a comprehensive list of all the notebooks and examples we have.
## Supported Models
Tunix supports a growing list of models including Gemma, Llama, and Qwen families.
See [Models](models.md) for a full list and details on how to add new ones.
## Citing Tunix
```bibtex
@misc{tunix2025,
title={Tunix (Tune-in-JAX)},
author={Bao, Tianshu and Carpenter, Jeff and Chai, Lin and Gao, Haoyu and Jiang, Yangmu and Noghabi, Shadi and Sharma, Abheesht and Tan, Sizhi and Wang, Lance and Yan, Ann and Yu, Weiren and others},
year={2025},
howpublished={\url{https://github.com/google/tunix}},
}
```
## Acknowledgements
Thank you to all our wonderful contributors!
[](https://github.com/google/tunix/graphs/contributors)
## /docs/launching.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Launching Jobs
Tunix supports several ways to launch training jobs, depending on your workflow:
* **[Tunix CLI](#tunix-cli)**: The default choice. A simple CLI launching tool with comprehensive configuration options.
* **[Interactive and Custom Launch](#interactive-and-custom-launch)**: While the Tunix CLI is preferred for simplicity, this approach provides full control. This is ideal for:
* **Early experimentation**: Getting familiar with the framework.
* **Advanced customization**: Complex cases requiring flexibility beyond the CLI.
## Tunix CLI
Tunix offers a configurable CLI to launch SFT and RL job directly from the command line. There a number of knobs and parameters allowing full customization of the job (described below), as well as a number of pre-defined examples to get you started.
### Configuration Hierarchy
You can tune CLI parameters in one of three ways. Configurations are merged from these sources in the following order of precedence (later sources override earlier ones):
1. **Base Config File** (Lowest Priority):
The default settings found in `base_config.yaml`.
2. **Config File Override**:
An optional secondary config file specified via the `override_config_file` argument (e.g., `override_config_file=/path/to/override.yaml`). Values here override the base config.
*(See an [example script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/configs/gemma2_2b.yaml)).*
3. **CLI Arguments** (Highest Priority):
Individual `key=value` pairs provided as command-line arguments. These override values from both the base config and the override file.
### Example CLI Scripts
This collection includes command-line interface (CLI) scripts designed to handle various tasks.
> NOTE: 🔑 Required Credentials
>
> Before running the scripts, please ensure you have the following environment variables configured to access the necessary model repositories:
>
> * **Hugging Face Access:**
> * `HF_TOKEN`: Required to authenticate and download models from Hugging Face.
>
> * **Kaggle Access:**
> * `KAGGLE_USERNAME`: Your Kaggle username.
> * `KAGGLE_KEY`: Your Kaggle API key.
<section class="zippy">
#### Guide to Setup Credential
TL;DR: This is a guide that covers how to get your keys, authorize them for restricted access, and save them permanently in a `.env` file.
##### Hugging Face (Token & Access)
**Step A: Generate the Token**
1. Log in to [Hugging Face](https://huggingface.co/).
2. Click your **Profile Picture** > **Settings** > **Access Tokens**.
3. Click **Create new token**.
* **Name:** `CLI-Access`
* **Permissions:** **Read** (sufficient for downloading).
4. **Copy** the token string (starts with `hf_`).
**Step B: Authorize Restricted Models (Critical)**
*Your token will fail if you skip this.*
1. Go to the specific model page (e.g., [`meta-llama/Llama-3.2-1B`](https://huggingface.co/meta-llama/Llama-3.2-1B)).
2. Find the **"Access this model"** banner at the top.
3. Review the license and click **Agree and access repository**.
4. Wait for permissions to sync to your token.
##### Kaggle (Key & Access)
**Step A: Generate the API Key**
1. Log in to [Kaggle](https://www.kaggle.com/).
2. Click your **Profile Picture** > **Settings**.
3. Scroll to the **API** section and click **Create New Token**.
4. Open the downloaded `kaggle.json` file to find your username and key.
**Step B: Authorize Access (Critical)**
1. Go to the specific **Models** page on Kaggle (e.g., [`gemma model family`](https://www.kaggle.com/models/google/gemma)).
2. Click the **Request Access** tab.
3. Review and Sign **Consent Form**.
##### Create the `.env` File
Instead of typing `export` every time, we will save these credentials in a file that sits in your project folder.
1. Create a new file in your project root named `.env` (no filename, just the extension).
2. Paste the following content into it, replacing the placeholders with your actual keys:
```bash
# .env
# Hugging Face Access
HF_TOKEN=hf_12345exampletokenstring
# Kaggle Access
KAGGLE_USERNAME=your_kaggle_username
KAGGLE_KEY=example_kaggle_key
```
>**Security Warning:** If you are using Git, you **must** add `.env` to your `.gitignore` file immediately. This prevents you from accidentally uploading your passwords to GitHub.
</section>
* **Supervised Fine Tuning**
*Refer to the [hardware requirement](https://github.com/google/tunix/blob/main/examples/sft/mtnt/README.md) before proceeding.*
* **Peft Training on MTNT dataset** ([Source Folder](https://github.com/google/tunix/blob/main/examples/sft/mtnt))
| Model Variant | Script Name | Link |
| :---------------- | :---------------------- | :------------------------------------------------------------------------------------------------------------------------------- |
| **Gemma 2B** | `run_gemma_2b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/sft/mtnt/run_gemma_2b.sh) |
| **Gemma 2 2B** | `run_gemma2_2b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/sft/mtnt/run_gemma2_2b.sh) |
| **Gemma 3 4B** | `run_gemma3_4b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/sft/mtnt/run_gemma3_4b.sh) |
| **Llama 3.2 3B** | `run_llama3.2_3b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/sft/mtnt/run_llama3.2_3b.sh) |
| **Qwen 2.5 0.5B** | `run_qwen2.5_0.5b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/sft/mtnt/run_qwen2.5_0.5b.sh) |
* **Reinforcement Learning**
*Refer to the [hardware requirement](https://github.com/google/tunix/blob/main/examples/rl/README.md) before proceeding.*
* **GRPO Training on GSM8K dataset** ([Source Folder](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/))
| Model Variant | Script Name | Link |
| :---------------- | :---------------------- | :------------------------------------------------------------------------------------------------------------------------------- |
| **Gemma 7B** | `run_gemma_7b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_gemma_7b.sh) |
| **Gemma 2 2B** | `run_gemma2_2b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_gemma2_2b.sh) |
| **Gemma 3 1B** | `run_gemma3_1b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_gemma3_1b.sh) |
| **Gemma 3 4B** | `run_gemma3_4b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_gemma3_4b.sh) |
| **Llama 3.1 8B** | `run_llama3.1_8b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_llama3.1_8b.sh) |
| **Llama 3.2 1B** | `run_llama3.2_1b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/run_llama3.2_1b.sh) |
* **PPO Training on GSM8K dataset** ([Source Folder](https://github.com/google/tunix/blob/main/examples/rl/PPO/gsm8k/))
| Model Variant | Script Name | Link |
| :---------------- | :---------------------- | :------------------------------------------------------------------------------------------------------------------------------- |
| **Gemma 2B** | `run_gemma_7b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/ppo/gsm8k/run_gemma_2b.sh) |
| **Gemma 2 2B** | `run_gemma2_2b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/ppo/gsm8k/run_gemma2_2b.sh) |
| **Gemma 3 1B** | `run_gemma3_1b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/ppo/gsm8k/run_gemma3_1b.sh) |
| **Llama 3.1 8B** | `run_llama3.1_8b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/ppo/gsm8k/run_llama3.1_8b.sh) |
| **Qwen 3 1.7B** | `run_llama3.2_1b.sh` | [View Script](https://github.com/google/tunix/blob/main/examples/rl/ppo/gsm8k/run_qwen3_1.7b.sh) |
### CLI Scripts Overview
CLI Core code and CLI example launch scripts reside separately. This section provides a high level overview of the CLI structure.
* [CLI Core](https://github.com/google/tunix/blob/main/cli)
* `base_config.yaml`: Define all the configurations that could be tuned when launching a job.
* `peft_main.py`: Main entry point to trigger Parameter-Efficient Fine-Tuning (PEFT) Trainer from the CLI configs.
* `grpo_main.py`: Main entry point to trigger Group Relative Policy Optimization (GRPO) Trainer from the CLI configs.
* `ppo_main.py`: Main entry point to trigger Proximal Policy Optimization (PPO) Trainer from the CLI configs.
* `config.py`: Logic to read and process the config passed by command line or environment variable.
* `reward_fn/...`: Predefined reward functions for reinforcement learning jobs. You could modify your shell scripts to use your own reward function.
* `gsm8k.py`: Predefined reward functions running on [gsm8k dataset](https://www.tensorflow.org/datasets/catalog/gsm8k).
* `gsm8k_verl.py`: Predefined reward functions on gsm8k dataset compatible with verl, refer to [this](https://github.com/google/tunix/blob/main/examples/rl/grpo/gsm8k/verl_compatible/README.md) for details.
* [CLI Examples](https://github.com/google/tunix/blob/main/examples)
* `SFT/..`: All available SFT CLI shell scripts.
* `RL/.. `: All available RL CLI shell scripts.
### Usage
#### Setup Cloud VM Environment
**TL;DR:** To automate the setup process for a single host, run the provided shell script matching your hardware accelerator.
**Automated Setup:**
* **For TPU:**
```shell
source scripts/setup_cli_tpu_single_host.sh
```
* **For GPU:**
```shell
source scripts/setup_cli_gpu_single_host.sh
```
<section class="zippy">
**Manual Setup (Detailed Breakdown)**
The following steps explain the environment setup process performed by the scripts above. You may skip this if you ran one of the automated scripts.
* **1. Create a project specific environment.**
```shell
python3 -m venv .venv
# Or simply `python -m venv .venv` depending on your system configuration.
```
* **2. Activate the Environment**
```shell
source .venv/bin/activate
```
* **3. Install Tunix dependency**
Make sure you have an updated pip version installed:
```shell
pip install --upgrade pip
```
* **Option A: TPU Only**
If you only require TPU, install with `[prod]` extra:
```shell
pip install -e .[prod]
```
* **Option B: Other Accelerators (e.g., GPU)**
First, install the core Tunix dependency:
```shell
pip install -e .
```
Then, install your accelerator-specific dependency separately, for example:
```shell
# Example for GPU
pip install jax[gpu]
```
</section>
### Config Explanation
This section provides a detailed explanation of the configuration parameters available in `base_config.yaml`. These parameters allow you to customize model selection, training dynamics, hardware utilization (mesh), and reinforcement learning specific settings.
#### Model Configuration (`model_config`)
These parameters define the base model, where to download it from, and how to shard it across TPUs/GPUs. Note that `actor_model_config`, `reference_model_config`, and `rollout_model_config` typically inherit from this base configuration.
* **`model_name`**: The unique full name identifier of the model. This
corresponds to the full name and should match exactly with the model name
used in Hugging Face or Kaggle. It is typically all lowercase and formatted
as `<model-family>-<model-version>`.
* *Example*: `gemma-2b`, `llama-3.1-8b`, `gemma2-2b-it`.
Refer to [models documentation](models.md#naming-conventions) for model naming.
* **`model_source`**: The source repository for the model. Options: `"huggingface"`, `"kaggle"`, `"gcs"`, or empty string `""` for local paths.
* **`model_id`**: The exact repository ID (case sensitive) as it appears on Hugging Face or Kaggle (e.g., `"meta-llama/Llama-3.1-8B"`).
* *Example for Hugging Face*: `meta-llama/Llama-3.1-8B` is extracted as shown belows
{: width="75%"}
* *Example for Kaggle*: `google/gemma-2/flax/gemma2-2b-it` is extracted as shown belows
{: width="75%"}
* **`model_path`**: Used if `model_source` is GCS or local. Specifies the direct file path to the model.
* **`model_download_path`**: Local directory where downloaded checkpoints will be cached (e.g., `"/tmp/models"`).
* **`rng_seed`**: Integer seed for initializing the `nnx.Rngs` state to manage randomness (e.g., `0`).
* **`model_display`**: Boolean flag. If set to `true`, prints the model structure/summary.
* **`intermediate_ckpt_dir`**: Directory for temporary storage when converting specific formats like Kaggle Gemma/Gemma2 to NNX (e.g., `"/tmp/intermediate_ckpt/"`).
* **`lora_config`**: Configuration for Low-Rank Adaptation (LoRA).
* `module_path`: Regex identifying layers to adapt (e.g., `".*q_einsum|.*kv_einsum|.*gate_proj..."`). Refer to [Lora](performance.md#peft-with-lora) for more details.
* `rank`: The rank of the low-rank approximation (e.g., `16`).
* `alpha`: Scaling factor for LoRA weights (e.g., `2.0`).
* `weight_qtype`: Quantization type for the weights (e.g., `"nf4"`).
* `tile_size`: Tile size for efficient computation (e.g., `256`).
* **`mesh`**: Defines the hardware mesh layout for distributed training.
* `shape`: Tuple string defining mesh dimensions (e.g., `"(2,2)"` for a 2x2 grid).
* `axis_names`: Names for mesh axes, often used for parallelism strategies (e.g., `"('fsdp','tp')"` for Fully Sharded Data Parallelism and Tensor Parallelism).
#### Tokenizer Configuration (`tokenizer_config`)
* **`tokenizer_path`**: Path or ID of the tokenizer. Usually matches `model_id`.
* **`tokenizer_type`**: The library to use for tokenization. Options: `"huggingface"`, `"sentencepiece"`.
* **`add_bos`** / **`add_eos`**: Boolean. Whether to automatically prepend Beginning of Sentence (BOS) or append End of Sentence (EOS) tokens.
#### Dataset Configuration
* **`dataset_name`**: The identifier for the dataset (e.g., `"Helsinki-NLP/opus-100"` for SFT or `"gsm8k"` for RL).
* **`batch_size`**: Global batch size per training step.
* **`max_target_length`**: Maximum length of the target sequence (in tokens).
* **`num_train_epochs`**: Number of complete passes through the training dataset.
* **`tfds_download`**: Boolean. Controls download behavior for TensorFlow Datasets.
#### Optimizer Configuration (`optimizer_config`)
Controls the gradient descent algorithm and learning rate scheduling. Tunix uses [optax](https://optax.readthedocs.io/en/latest/getting_started.html) to for optimizer.
* **`opt_type`**: The optimizer algorithm to use (e.g., `"adamw"`). Refer to [optax.optimizer](https://optax.readthedocs.io/en/latest/api/optimizers.html#optimizers) for available optimizer.
* **`learning_rate`**: A global scaling factor, either fixed or evolving along iterations with a scheduler if `schedule_type` is set.
* **`schedule_type`**: The learning rate schedule function (e.g., `"warmup_cosine_decay_schedule"`). Refer to [optax.schedule](https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html#) for available schedulers.
* **`warmup_steps`**: Number of steps to linearly increase the learning rate from `init_value` to `peak_value`.
* **`decay_steps`**: Number of steps for the decay phase of the schedule.
* **`max_grad_norm`**: Gradient clipping threshold. Essential for preventing exploding gradients, especially in RL.
#### Training Configuration (`training_config`)
General settings for the training loop, logging, and checkpointing.
* **`max_steps`**: Total number of training steps to run.
* **`eval_every_n_steps`**: Frequency of running evaluation steps.
* **`gradient_accumulation_steps`**: Number of steps to accumulate gradients
before performing a parameter update (simulates larger batch sizes).
* **`checkpointing_options`**:
* `max_to_keep`: Number of recent checkpoints to retain.
* `save_interval_steps`: How often to save a checkpoint.
* `enable_async_checkpointing`: Boolean to toggle asynchronous checkpointing execution.
* `timeout_secs`: Maximum time permitted for asynchronous writes natively.
* **`metrics_logging_options`**: Settings for logging. Includes project name, run name, and flush frequency.
* **`data_sharding_axis`**: Specifies which mesh axis is used for data sharding (e.g., `["fsdp"]`).
#### GRPO Configuration (`grpo_config`, `rollout_config`)
Specific parameters for Group Relative Policy Optimization (GRPO).
* **`num_generations`**: (GRPO specific) The number of responses generated per prompt in a single step (corresponds to $\varepsilon$ in the paper).
* **`beta`**: Coefficient for the KL divergence penalty. Keeps the trained model close to the reference model.
* **`epsilon`**: Clipping parameter for the loss function (similar to PPO) to ensure stable updates.
* **`temperature`**: Sampling temperature for rollouts. Higher values (e.g., 0.9) encourage diversity, which is critical for GRPO.
* **`total_generation_steps`**: Maximum tokens to generate during the rollout phase.
* **`reward_functions`**: List of python file paths containing the reward logic (e.g., checking math answers for GSM8K).
#### PPO Configuration (`ppo_config`, `rollout_config`)
Specific parameters for Proximal Policy Optimization (PPO).
* **`num_iterations`**: (PPO specific) The number of iterations per batch (corresponds to $\varmu$ in the paper)
* **`lambda`**: The lambda parameter for Generalized Advantage Estimation (GAE).
* **`gamma`**: The discount factor for future rewards in GAE.
* **`beta`**: Coefficient for the KL divergence penalty. Keeps the trained model close to the reference model.
* **`epsilon`**: Clipping parameter for the loss function to ensure stable updates.
* **`reward_functions`**: List of python file paths containing the reward logic (e.g., checking math answers for GSM8K).
## Interactive and Custom Launch
For interactive development or custom cluster setup, refer to our notebooks, examples, and guides, which demonstrate how to build the cluster and launch jobs programmatically:
* **[Quick Start Guides](quickstart.md)**: Step-by-step guides for SFT, RL, Agentic workflows, and more.
* **[Examples & Colabs](examples.md)**: A comprehensive list of interactive notebooks and scripts.
These resources are an excellent starting point for learning Tunix core concepts. Additionally, they demonstrate how to use Tunix in more complex scenarios that require full control. However, note that the **[Tunix CLI](#tunix-cli)** is recommended option for most standard workflows due to its simplicity.
## /docs/make.bat
```bat path="/docs/make.bat"
@ECHO OFF
pushd %~dp0
REM Command file for Sphinx documentation
if "%SPHINXBUILD%" == "" (
set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=.
set BUILDDIR=_build
%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
echo.
echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
echo.installed, then set the SPHINXBUILD environment variable to point
echo.to the full path of the 'sphinx-build' executable. Alternatively you
echo.may add the Sphinx directory to PATH.
echo.
echo.If you don't have Sphinx installed, grab it from
echo.https://www.sphinx-doc.org/
exit /b 1
)
if "%1" == "" goto help
%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end
:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
:end
popd
```
## /docs/metrics.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Metrics
Tunix provides a comprehensive observability stack for training LLMs,
encompassing everything from basic training metrics to detailed execution
traces. This section is composed of three main pillars:
* **[Collected Metrics](#collected-metrics)**: Describes a rich
set of system, model, and RL-specific performance metrics out of the box.
* **[Metric Loggers](#metric-loggers)**: Describes a flexible,
protocol-based logging system that allows you to seamlessly integrate with
your preferred logging service (e.g., TensorBoard, Wandb, CLU) or create
custom backends.
* **[Performance Metric Tracing](#performance-metric-tracing)**: Describes a
built-in, lightweight tracing system that generates detailed execution
timelines for deep performance analysis and visualization in Perfetto.
## Collected Metrics
Tunix automatically collects a rich set of metrics during training to help you
monitor performance, convergence, and resource utilization.
### Common Metrics (SFT & RL)
These metrics are collected for both Supervised Fine-Tuning (SFT) and
Reinforcement Learning (RL) jobs:
* **`loss`**: The training loss for the current step.
* **`perplexity`**: The perplexity of the model on the training batch
(`exp(loss)`).
* **`learning_rate`**: The current learning rate from the optimizer.
* **`step_time_sec`**: The time taken to execute a single training step (in
seconds).
* **`steps_per_sec`**: The training speed, measured in steps per second.
* **`tflops_per_step`**: The estimated Trillion Floating Point Operations
(TFLOPs) performed per step (if supported by the hardware/backend).
### DPO & ORPO-Specific Metrics
For Direct Preference Optimization (DPO) and Odds Ratio Preference Optimization (ORPO) training, Tunix collects preference-specific alignment and loss metrics:
* **`rewards/chosen`**: The average implicit reward for the chosen responses (calculated using policy and reference log probs for DPO, or using length-averaged policy log probs for ORPO).
* **`rewards/rejected`**: The average implicit reward for the rejected responses.
* **`rewards/margin`**: The difference between chosen and rejected rewards (`rewards/chosen - rewards/rejected`). A positive margin indicates the model prefers the chosen responses.
* **`rewards/accuracy`**: The frequency with which the chosen response has a higher reward than the rejected response (should approach 1.0 as training converges).
* **`log_probs/chosen`**: The average log probability of the chosen responses under the policy model.
* **`log_probs/rejected`**: The average log probability of the rejected responses under the policy model.
* **`odds_ratio`** *(ORPO only)*: The average odds ratio between the chosen and rejected responses.
* **`sft_loss`** *(ORPO only)*: The supervised fine-tuning loss component (NLL) of the ORPO loss.
* **`or_loss`** *(ORPO only)*: The odds ratio preference loss component of the ORPO loss.
### RL-Specific Metrics (PPO/GRPO)
For Reinforcement Learning jobs, Tunix collects additional metrics related to
the RL algorithm (e.g., PPO), reward modeling, and generation.
#### Rewards & Scores
* **`rewards/sum`**: The sum of rewards for a trajectory.
* **`rewards/mean`**, **`rewards/max`**, **`rewards/min`**: Statistics of the
rewards across the batch.
* **`score/mean`**, **`score/max`**, **`score/min`**: Statistics of the raw
scores from the reward model (before any algorithm-specific modifications
like KL penalty).
* **`reward_kl_penalty`**: The KL divergence penalty applied to the reward
(if applicable).
* **`rewards/<reward_fn_name>`**: If using multiple reward functions,
individual reward components are logged by name.
#### Policy & Value (PPO)
* **`advantages/mean`**, **`advantages/max`**, **`advantages/min`**:
Statistics of the advantages.
* **`returns/mean`**, **`returns/max`**, **`returns/min`**: Statistics of
the returns.
* **`values/mean`**, **`values/max`**, **`values/min`**: Statistics of the
value function estimates.
* **`pg_clipfrac`**: The fraction of the batch where the policy gradient was
clipped.
* **`vf_clipfrac`**: The fraction of the batch where the value function update
was clipped.
* **`loss/entropy`**: The entropy of the policy (if entropy regularization is
enabled).
#### Generation & Data
* **`prompts`**: The input prompts used for generation.
* **`completions`**: The text completions generated by the model.
* **`completions/mean_length`**, **`completions/max_length`**,
**`completions/min_length`**: Statistics on the length of generated
completions.
* **`trajectory_ids`**: Unique identifiers for the trajectories.
* **`actor_dequeue_time`**: Time spent waiting for data from the rollout
workers (if async rollout is enabled).
## Metric Loggers
Tunix provides a flexible, protocol-based logging system that allows you to
integrate any logging service or library.
The primary interface for logging is the `MetricsLogger`. It is configured
using `MetricsLoggerOptions`. Below is an example of how to configure the
`MetricsLogger`. **Note**: The exact fields that need to be configured depend
on the backend, which typically defaults based on the execution environment. See
[Logging Backends Supported](#logging-backends-supported) for details on
backend-specific configurations.
```python
from tunix.sft import metrics_logger
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
project_name="my-project",
run_name="experiment-1",
)
logger = metrics_logger.MetricsLogger(metrics_logger_options=options)
```
With the above, agentic_grpo_learner will by default start an async trajectory
logger which logs the trajectories including prompts, responses, etc. to the
specified `log_dir`.
### Enabling Metrics in Jobs
Once you have your `MetricsLoggerOptions` configured, you can pass it to your
SFT or RL job via the training configuration.
#### Supervised Fine-Tuning (SFT)
For SFT, pass the `metrics_logging_options` to the `TrainingConfig`.
```python
from tunix.sft import metrics_logger
from tunix.sft import peft_trainer
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
project_name="my-sft-project",
)
training_config = peft_trainer.TrainingConfig(
eval_every_n_steps=100,
metrics_logging_options=options,
# ... other configurations
)
trainer = peft_trainer.PeftTrainer(
model=model,
optimizer=optimizer,
training_config=training_config,
)
```
#### Reinforcement Learning (RL)
For RL, pass the `metrics_logging_options` to the `RLTrainingConfig`, which is
then used in `ClusterConfig`.
```python
from tunix.rl import rl_cluster as rl_engine_lib
from tunix.sft import metrics_logger
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
project_name="my-rl-project",
)
training_config = rl_engine_lib.RLTrainingConfig(
actor_optimizer=optimizer,
metrics_logging_options=options,
# ... other configurations
)
cluster_config = rl_engine_lib.ClusterConfig(
role_to_mesh=role_to_mesh,
training_config=training_config,
rollout_config=rollout_config,
# ... other configurations
)
engine = rl_engine_lib.RLEngine(
actor=actor_model,
tokenizer=tokenizer,
cluster_config=cluster_config,
)
```
### Logging Backends Supported
Tunix supports several logging backends out of the box, powered by `metrax`
[link](https://github.com/google/metrax/ ). The default backend selection
depends on the execution environment.
#### Wandb
[Weights & Biases](https://wandb.ai/) is a supported backend for experiment tracking. ([Backend Code](https://github.com/google/metrax/blob/main/src/metrax/logging/wandb_backend.py))
* **Availability**: *Enabled by default* in external environments (if `wandb`
is installed).
* **Configuration**:
* `project_name`: Sets the Wandb project name (default: "tunix").
* `run_name`: Sets the specific run name. If not provided, it defaults to
a timestamp (e.g., `2025-01-14_08-40-01`). **Note:** Wandb distinguishes
between a run name and a run id. Runs with the same name are tracked as
separate entities differentiated by their run id.
#### TensorBoard
[TensorBoard](https://www.tensorflow.org/tensorboard) is supported for visualizing metrics. ([Backend Code](https://github.com/google/metrax/blob/main/src/metrax/logging/tensorboard_backend.py))
* **Availability**: *Enabled by default* in external environments.
* **Configuration**:
* `log_dir`: Directory where event files are written.
* `flush_every_n_steps`: Frequency of flushing logs to disk (default: 100).
### Experimental: OpenTelemetry double-write
Tunix is evaluating [OpenTelemetry](https://opentelemetry.io/) as a
vendor-neutral instrumentation layer. As an opt-in first step, `MetricsLogger`
supports a **double-write** mode: every scalar passed to `log()` is emitted both
through the existing Metrax/`jax.monitoring` backends (unchanged, still the
default) and as an OpenTelemetry gauge measurement. This lets you compare the
two pipelines side by side before any switch is made.
Enable it with the `enable_opentelemetry` flag (requires `pip install
'google-tunix[otel]'`):
```python
# OTLP exporter: pip install opentelemetry-exporter-otlp
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import (
OTLPMetricExporter,
)
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from tunix.sft import metrics_logger
reader = PeriodicExportingMetricReader(OTLPMetricExporter())
provider = MeterProvider(metric_readers=[reader])
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs", # Existing backends keep working unchanged.
enable_opentelemetry=True, # Additionally emit OpenTelemetry gauges.
)
logger = metrics_logger.MetricsLogger(options, otel_meter_provider=provider)
```
Key properties of the double-write path:
* **Off by default.** With the flag unset, behavior is identical to previous
releases; the OpenTelemetry API does not even need to be installed.
* **Application-owned providers.** Tunix never configures exporters or shuts
down providers. Configure the global `MeterProvider` (or inject one via the
keyword-only `otel_meter_provider` argument) and own its flush/shutdown.
`MetricsLogger.close()` only closes the Metrax backends.
* **Process policy.** Like the Metrax backends, only JAX process 0 emits.
* **Naming.** Known metrics map to stable instrument names (`loss` →
`tunix.training.loss`, `perplexity` → `tunix.training.perplexity`,
`learning_rate` → `tunix.training.learning_rate`, `grad_norm` →
`tunix.training.gradient.norm`); other metric keys are normalized into the
`tunix.*` namespace (e.g. `rewards/score mean` →
`tunix.rewards.score.mean`). The legacy prefix and mode become the
low-cardinality attributes `tunix.metrics.prefix` and `tunix.training.mode`,
and the logical step is emitted as a separate `tunix.training.step` gauge.
#### OpenTelemetry → Weights & Biases
W&B ingests OpenTelemetry traces natively (via the
[Weave OTLP endpoint](https://docs.wandb.ai/weave/guides/tracking/otel)), but
has no OTLP endpoint for run metrics. Tunix therefore ships
`tunix.sft.otel_wandb.WandbMetricsExporter`, an OpenTelemetry SDK metric
exporter that forwards the gauges above to `wandb.log`, using the familiar
`{prefix}/{mode}/{name}` chart keys (e.g. `actor/train/tunix.training.loss`) and
the `tunix.training.step` gauge as the W&B step:
```python
import wandb
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from tunix.sft import metrics_logger
from tunix.sft import otel_wandb
run = wandb.init(project="my-project")
reader = PeriodicExportingMetricReader(
otel_wandb.WandbMetricsExporter(run), export_interval_millis=10_000
)
provider = MeterProvider(metric_readers=[reader])
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
enable_opentelemetry=True,
)
logger = metrics_logger.MetricsLogger(options, otel_meter_provider=provider)
```
Because the Metrax `WandbBackend` also remains active by default, you can point
both pipelines at the same W&B project and compare the charts directly.
### Custom metric logger
You can integrate any logging service by creating a custom backend that conforms
to the `metrax.logging.LoggingBackend` protocol.
#### 1. The Protocol
Your custom backend class only needs to implement `log_scalar` and `close`.
Explicit inheritance from a base class is not required since Metrax uses
Python's structural typing (duck typing).
```python
from typing import Protocol
import numpy as np
class LoggingBackend(Protocol):
def log_scalar(self, event: str, value: float | np.ndarray, **kwargs):
"""Logs a scalar value.
Args:
event: The name of the metric/event (e.g., "train/loss").
value: The scalar value of the metric.
**kwargs: Additional arguments, typically including 'step' (int).
"""
...
def close(self):
"""Closes the logger and flushes any pending data."""
...
```
#### 2. Creating a Custom Backend
Here is an example of a backend that simply prints metrics to stdout:
```python
class SimplePrintBackend:
def log_scalar(self, event, value, **kwargs):
print(f"Logged {event}: {value}")
def close(self):
print("Closing backend.")
```
#### 3. Using Your Custom Backend
To use your custom backend, you must pass a **factory** (a callable that returns
an instance) to `MetricsLoggerOptions`. This ensures configuration objects
remain serializable and safe to copy.
##### Case A: Simple Backend (No Arguments)
If your backend class requires no arguments in its `__init__`, you can simply
pass the class itself.
```python
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
backend_factories=[SimplePrintBackend],
)
logger = metrics_logger.MetricsLogger(metrics_logger_options=options)
```
##### Case B: Backend with Arguments
If your backend requires arguments, use a `lambda` to create a factory.
```python
class FileBackend:
def __init__(self, filename):
self.file = open(filename, 'w')
def log_scalar(self, event, value, **kwargs):
self.file.write(f"{event},{value}\n")
def close(self):
self.file.close()
# Create a factory using a lambda
my_file_factory = lambda: FileBackend(filename="/tmp/metrics.csv")
options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/logs",
backend_factories=[my_file_factory]
)
logger = metrics_logger.MetricsLogger(metrics_logger_options=options)
```
## Performance Metric Tracing
Tunix provides a highly lightweight performance tracing and metrics collection
system designed specifically for RL workflows. Unlike detailed profiling tools
(e.g., xprof or standard JAX profiling), which collect exhaustive low-level
details but incur significant overhead and are typically only used for short
debugging intervals (e.g., 10s of seconds), this tracing system is designed
with minimal overhead. It can safely be left enabled for the entire duration
of your training run. It allows you to monitor the execution time of different
stages (e.g., rollout, actor training, reference inference) across both host
and device timelines.
> **Note:** Performance metric tracing is currently only supported for the GRPO
> main entry point.
There are currently two versions of the performance metrics system:
* **Original Version (v1)** ([Code](https://github.com/google/tunix/blob/main/perf/export.py)):
Uses `PerfSpanQuery` to extract spans and compute metrics (e.g., rollout
time, wait time).
* **Experimental Version (v2)** ([Code](https://github.com/google/tunix/blob/main/perf/experimental/export.py)):
A more flexible version that tracks `Timeline` objects and is planned to
replace v1.
Both versions can export metrics to your custom export function and write
detailed Perfetto trace files that can be visualized at
[ui.perfetto.dev](https://ui.perfetto.dev/).
### Using Performance Metrics via CLI
When running Tunix via the CLI, you can configure performance metrics by
providing a `perf_metrics_options` dictionary inside your `rl_training_config`.
```yaml
rl_training_config:
perf_metrics_options:
enable_perf_v1: true # Enable v1 (default: true)
enable_perf_v2: false # Enable v2 (default: false)
enable_trace_writer: true # Enable writing Perfetto trace files (default: true)
trace_dir: "/tmp/perf_trace" # Directory to write the trace files to
custom_export_fn_path: "path.to.my.custom_fn" # Optional path to a custom v1 export function
custom_export_fn_path_v2: "path.to.my.custom_fn_v2" # Optional path to a custom v2 export function
```
The CLI automatically parses these options, initializes the appropriate export
functions, and registers them with the training cluster.
Note that `enable_perf_v1` and `enable_perf_v2` can be toggled independently,
allowing you to use one or both systems simultaneously. If you wish to use a
custom export function instead of the defaults, you must provide the fully
qualified import path to your function via `custom_export_fn_path` (for v1) and
`custom_export_fn_path_v2` (for v2).
### Using Performance Metrics via Code
If you are initializing the `RLEngine` programmatically, you must construct a
`PerfMetricsConfig` and pass it to the engine manually.
#### Original Version (v1)
In v1, use `PerfMetricsExport.from_cluster_config()` to generate a default
export function. This function automatically computes various duration metrics
based on the cluster's mesh topology (e.g., whether the rollout and actor models
are collocated or on different TPU meshes
[code](https://github.com/google/tunix/blob/main/perf/export.py;l=102).
The metrics are aggregated per `global_step` and use the
[Metric Logger](#metric-loggers) to log to the desired output. For example,
if TensorBoard is activated:
By default, v1 also writes detailed execution traces to Perfetto Proto formatted
file. It reads `perf_metrics_options` from the cluster configuration to
initialize the trace writer. You can specify the output directory by configuring
`trace_dir` within `PerfMetricsOptions` inside your `RLTrainingConfig`.
```python
from tunix.perf import metrics as perf_metrics
from tunix.perf import export as perf_export
from tunix.rl import rl_cluster as rl_engine_lib
from tunix.sft import metrics_logger
# 1. Define metric logger options (for collecting aggregate perf metrics).
metric_logger_options = metrics_logger.MetricsLoggerOptions(
log_dir="/tmp/tensorboard/grpo",
project_name="my-rl-project",
)
training_config = rl_engine_lib.RLTrainingConfig(
metrics_logging_options=metric_logger_options,
# ... other configurations
)
cluster_config = rl_engine_lib.ClusterConfig(
training_config=training_config,
# ... other configurations
)
# 2. Create a PerfMetricsConfig object.
perf_config = perf_metrics.PerfMetricsConfig()
perf_config.custom_export_fn = (
perf_export.PerfMetricsExport.from_cluster_config(cluster_config)
)
# 3. Pass the config to the RLEngine.
engine = rl_engine_lib.RLEngine(
actor=actor_model,
tokenizer=tokenizer,
cluster_config=cluster_config,
perf_config=perf_config,
)
```
#### Experimental Version (v2)
For the experimental version, you can use the default export function, which
writes the raw timelines to a Perfetto trace file by using the
`PerfMetricExport` class. The trace files can be written to a directory
by defining `trace_dir`. If `trace_dir` is not provided, it defaults to
`/tmp/perf_traces`. The v2 version supports local files and remote endpoints
supported by `etils.epath` including GCS such as `gs://your-bucket/path/`.
Note that v2 is still experimental and additional
capabilities, such as exporting aggregated metrics to TensorBoard, are WIP. Once
the functionality is complete, v2 will be replacing the original version.
```python
from tunix import PerfMetricsConfig
from tunix.perf.experimental.export import PerfMetricsExport
from tunix.rl import rl_cluster as rl_engine_lib
# 1. Define cluster_config (as you would for your RL job)
cluster_config = rl_engine_lib.ClusterConfig(
training_config=training_config,
# ... other configurations
)
# 2. Create a PerfMetricsConfig with custom_export_fn_v2 defined.
perf_metrics_config = PerfMetricsConfig(
custom_export_fn_v2=PerfMetricsExport.from_cluster_config(
cluster_config=cluster_config,
trace_dir="/tmp/agentic_perf",
).export_metrics
)
# 3. Pass the config to the RLEngine.
engine = rl_engine_lib.RLEngine(
actor=actor_model,
tokenizer=tokenizer,
cluster_config=cluster_config,
perf_config=perf_config,
)
```
### Custom Export Functions
If you want to compute custom metrics from the collected spans instead of using
the defaults, you can define and provide your own export function.
**Custom Export Function for v1:**
A v1 export function takes a `PerfSpanQuery` and returns a dictionary of
metrics.
```python
from tunix.perf import metrics as perf_metrics
def my_custom_export_fn(query: perf_metrics.PerfSpanQuery) -> perf_metrics.MetricsT:
# Example: query main thread for the latest 'global_step' group
global_steps = query().main().last_group("global_step").get()
if global_steps:
# MetricsT maps metric names to (value, optional_aggregation_fn)
return {"perf/custom_step_time": (global_steps[0].duration, None)}
return {}
perf_config.custom_export_fn = my_custom_export_fn
```
**Custom Export Function for v2:**
A v2 export function takes a mapping of timeline IDs to their respective
`Timeline` objects.
```python
from tunix.perf import metrics as perf_metrics
from tunix.perf.experimental import tracer
def my_custom_export_fn_v2(timelines: dict[str, tracer.Timeline]) -> perf_metrics.MetricsT:
# Example: iterate over host and device timelines
for tl_id, timeline in timelines.items():
pass # Analyze timeline.root span
return {}
perf_config.custom_export_fn_v2 = my_custom_export_fn_v2
```
### Visualizing with Perfetto
If you have enabled the trace writer (by setting `enable_trace_writer: true` via
the CLI or by specifying `trace_dir` in your configuration), a proto-formatted
file (e.g., `perfetto_trace_1771973518.pb`) containing the raw spans and
timelines will be saved to the specified directory (which defaults to
`/tmp/perf_traces`). To view the trace, download the file to your local
machine and drag-and-drop it into the
[Perfetto UI](https://ui.perfetto.dev/). The interface allows you to
interactively zoom, pan, and query the execution trace, as shown below:
## /docs/models.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Models
## Models supported
Tunix supports the following models:
| Model | Sizes |
|:---|:---|
| Gemma | 2B, 7B, 9B |
| Gemma 2 | 2B, 9B |
| Gemma 3 | 270M, 1B, 4B, 12B, 27B |
| Gemma 4 | E2B, E4B, 12B, 26B-A4B, 31B |
| Llama 3 | 70B, 405B |
| Llama 3.1 | 8B, 70B, 405B |
| Llama 3.2 | 1B, 3B |
| Qwen 2.5 | 0.5B, 1.5B, 3B, 7B |
| Qwen 3 | 0.6B, 1.7B, 4B, 8B, 14B, 30B, 32B |
| Qwen 3.5 | 35B, 397B |
### Model Sources
#### Huggingface & Kaggle
The model configurations and checkpoints should be accessible from Huggingface and Kaggle.
For example, following snippets shows how to load the Gemma 2B model from Huggingface:
```python
ignore_patterns = [
"*.pth", # Ignore PyTorch .pth weight files
]
MODEL_PATH = snapshot_download(repo_id="google/gemma-2-2b-it", ignore_patterns=ignore_patterns)
```
#### GCS
You can also store model checkpoints to GCS. So if you have GCS bucket resources
and have uploaded the model checkpoints there, you can access them as well.
```python
MODEL_PATH = "gs://<your-bucket-dev>/your-model-checkpoints"
```
#### Maxtext
Tunix also allows loading fully optimized models from [Maxtext](https://github.com/AI-Hypercomputer/maxtext).
Once you have an accessible model path from one of the above approach, you are able to load it through Tunix model loading API as following:
```python
config = model_lib.ModelConfig.gemma2_2b()
mesh = jax.make_mesh((1, 1), ("fsdp", "tp"), axis_types=(jax.sharding.AxisType.Auto,) * 2)
with mesh:
gemma = params_lib.create_model_from_safe_tensors(MODEL_PATH, config, mesh)
```
## Fully optimized models
Model optimization is critical for efficient model execution. This includes optimal shardings on TPUs, optimization with Pallas kernels, etc. Tunix provides a lightweight suite of models which is only optimized to some extent. Integration of Tunix and [Maxtext](https://github.com/AI-Hypercomputer/maxtext) enables users to run the RL workloads with fully optimized models. Refer to the [single-host](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/tutorials/posttraining/rl.md) and [multi-host](https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/tutorials/posttraining/rl_on_multi_host.md) tutorial on how to run an optimized model RL workload with Maxtext and Tunix.
## Adding a new model
You can add new models to Tunix codebase by following the Tunix convention.
### Model Family
If the new model falls into one of the existing model families (e.g. Gemma, Llama, etc.) then adding a new model doesn't need to create new files. You just need to add the model specs to the corresponding model family. Take a look at the [Llama examples](https://github.com/google/tunix/blob/main/models/llama3/model.py;l=98-135).
If the new model is from a new model family that Tunix hasn't supported yet. You will need to follow the design and APIs as the existing model families to create the model implementation.
### Naming
Adding the new model needs to following the naming convention that Tunix supports so that `AutoModel`(as described below) could work correctly. We use the pattern of `<model_family><major_version>p<minor_version>_<model_size>`to name a model. For example, the `Llama3.2 1b` model is named as `llama3p2_1b` while a `Qwen2.5 1.5b` model is named as `qwen2p5_1p5b`.
## AutoModel
`AutoModel` provides a unified interface for instantiating Tunix models from
pretrained checkpoints, similar to the Huggingface `AutoModel` API. It allows
you to load a model simply by providing its `model_id`, handling the download
and initialization for you.
### Basic Usage
To load a model, use the `AutoModel.from_pretrained` method with the model
identifier and your JAX sharding mesh. By default this will download the model
from Huggingface.
```python
from tunix.models.automodel import AutoModel
import jax
# 1. Define your mesh
mesh = jax.make_mesh((1, 1), ("fsdp", "tp"), axis_types=(jax.sharding.AxisType.Auto,) * 2)
# 2. Load the model
# By default, this downloads from Huggingface.
model, model_path = AutoModel.from_pretrained(
model_id="google/gemma-2-2b-it", # Using HF id as model_id
mesh=mesh
)
print(f"Model loaded from: {model_path}")
```
### Specifying Model Source
You can load models from different sources (e.g., Kaggle, GCS, etc.) using the
`model_source` argument.
#### From Huggingface:
This is the default choice (`ModelSource.HUGGINGFACE`) as shown in the
example above.
#### From Kaggle:
For Kaggle, you must provide the `model_id` which is the Huggingface identifier or model_config_id (see [Naming Conventions](models.md#naming-conventions)) to determine the model configuration and the `model_path` which is the Kaggle
Hub model identifier (used to download the model from Kaggle).
```python
model, model_path = AutoModel.from_pretrained(
model_id="gemma2_2b_it", # Using model_config_id as model_id
mesh=mesh,
model_source=ModelSource.KAGGLE,
model_path="google/gemma-2/flax/gemma2-2b-it",
)
```
For example the `model_path` for the `google/gemma-2/flax/gemma2-2b-it` is extracted on Kaggle as shown below
{: width="75%"}
#### From GCS:
For GCS, you must provide the `model_id` which is the Huggingface identifier or model_config_id (see [Naming Conventions](models.md#naming-conventions)) to determine the model configuration and the `model_path` (the actual GCS
location).
```python
model, model_path = AutoModel.from_pretrained(
model_id="gemma2_2b_it", # Using model_config_id as model_id
mesh=mesh,
model_source=ModelSource.GCS,
model_path="gs://my-bucket/gemma-2-2b-it"
)
```
#### From Maxtext:
For Maxtext, you must provide the `model_id`, the `model_name` (to specify the exact architecture expected by MaxText), and the `model_path` (which can be a parameter path or omitted if loading base configuration).
```python
model, model_path = AutoModel.from_pretrained(
model_id="llama3p1_8b",
mesh=mesh,
model_source=ModelSource.MAXTEXT,
model_name="llama3.1-8b",
model_path="gs://my-bucket/maxtext-checkpoint",
base_emb_dim=4096,
sparse_matmul=True,
remat_policy="minimal"
)
```
To load a Maxtext model when launching training via a shell script, append the corresponding override arguments directly to the script execution:
```bash
bash examples/rl/grpo/gsm8k/run_qwen3.sh \
--model_config.model_name="Qwen3-0.5B" \
--model_config.model_source="maxtext" \
--model_config.model_path="gs://my-bucket/maxtext-checkpoint" \
--model_config.kwargs.base_emb_dim=4096 \
--model_config.kwargs.sparse_matmul=true \
--model_config.kwargs.remat_policy="minimal"
```
### Model Download Path
Optionally, you can also provide the `model_download_path` argument, which
specifies where the model is to be downloaded to. Depending on the
`model_source` the effect of specifying this variable is different:
* **Huggingface**: Files are downloaded directly to this directory.
* **Kaggle**: Sets the `KAGGLEHUB_CACHE` environment variable to this path.
* **GCS**: No-op.
* **Internal**: Files are copied to this directory. If omitted, the model is loaded directly from the `model_path`. This mode (Internal) is not supported in OSS version.
## Naming Conventions
This section outlines the naming conventions used within Tunix for model
identification and configuration. These conventions ensure consistency when
loading models from various sources like Huggingface or Kaggle.
The `ModelNaming` dataclass handles the parsing and standardization of model names.
* **`model_id`**: This is a unique identifier used to identifty the model in mind and extract the family, version, and desired config from. Tunix support two identifiers as the `model_id`:
1. **Huggingface (HF) IDs:** The full model name identifier (case sensitive), as it appears
on Huggingface, including the parent directory.
* **Extracting model_id from HF**: For example,
`meta-llama/Llama-3.1-8B` is extracted as shown below:
{: width="75%"}
2. **Native Tunix model_configs:** the `model_config_id` representing the exact config from the model class can be used directly as the `model_id`. In this case it will also be treated as the `model_name`.
* **Extracting model_id from model_config_id**: In this case, you would need to refer to the source code (`model.py`) for each model family and select the config id from the `ModelConfig` class, for example `llama3p1_8b` from the llama [model code](https://github.com/google/tunix/blob/main/models/llama3/model.py;bpv=1;bpt=1;l=138).
* **`model_name`**: The unique full name identifier of the model. This
corresponds to the full name and should match exactly with the model name
used in Hugging Face or Kaggle. It is typically all lowercase and formatted
as `<model-family>-<model-version>` (when HF is used for model_id) or `<model-family>_<model-version>` (when model_config_id is used for model_id) .
* *Example for HF as model_id*: `gemma-2b`, `llama-3.1-8b`, `gemma-2-2b-it`.
* *Example for model_config_id as model_id*: `gemma_2b`, `llama3p1_8b`, `gemma2_2b_it`.
* **`model_family`**: The standardized model family. Unnecessary hyphens are
removed, and versions are standardized (e.g., replacing dot with `p`).
* *Example*: `gemma`, `gemma2`, `qwen2p5`.
* *Conversion*: `gemma-2` -> `gemma2`, `qwen2.5` -> `qwen2p5`.
* **`model_version`**: The standardized version of the model family (lowercase,
hyphens to underscores, dots to `p`). This is usually the second portion of
the `model_name` and includes size information or tuning variants (e.g., "it"
for instruction tuned).
* *Example*: `2b_it`.
* *Conversion*: `2b-it` -> `2b_it`
* **`model_config_category`**: The Python class name of the `ModelConfig` class. This groups models that share the same configuration structure.
* *Example*: Both `gemma` and `gemma2` models fall under the `gemma` category, with the `ModelConfig` class defined in `models/gemma/model.py`.
* **`model_config_id`**: The standardized configuration ID used within the `ModelConfig` class. It is composed of the `model_family` and `model_version`.
* *Example*: `gemma_2b_it` or `qwen2p5_0p5b`.
You can initialize `ModelNaming` by providing either the `model_id` or the
`model_name`. If `model_id` is provided, the `model_name` is inferred as the
last segment of the `model_id`. If `model_name` is provided, it is used
directly. All other naming attributes are then automatically derived and
validated.
## /docs/performance.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Performance Considerations
This guide provides an overview of performance optimization techniques in Tunix
that help maximize hardware utilization, alleviate memory pressure or OOMs, and
increase training throughput.
## PEFT with LoRA
Tunix provides native support for LoRA (Low-Rank Adaptation), enabling parameter
efficient fine-tuning for LLMs. By injecting low-rank trainable matrices into
specific transformer layers, Tunix allows you to update the model’s parameters
without the massive HBM overhead of full-parameter training.
Key benefits:
* Drastically reduces HBM overhead by training only a fraction of the total
parameters, enabling large model training with constrained capacity.
* Optimizes disk space by persisting only the lightweight LoRA parameters
during checkpointing.
* In RL workflows, it allows the Actor and Reference models to share a single
frozen base model, effectively reducing the HBM footprint.
* Enhances weight synchronization speed between the Actor and Rollout engines
by transferring only the LoRA parameters instead of the full model
parameters.
* Tunix automatically handles LoRA layer detection and processing, ensuring a
zero-config experience for the user.
You can apply LoRA/QLoRA to Tunix models using
[qwix](https://github.com/google/qwix) as shown below. If the base model is
sharded, then sharding strategy will be automatically propagated to LoRA layers.
Under the hood, Tunix will automatically detect if a model contains LoRA layers
and configures the training pipeline accordingly, requiring no manual
intervention from the user.
```python
import qwix
lora_provider = qwix.LoraProvider(
module_path=".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj",
rank=16,
alpha=2.0,
weight_qtype="nf4", # enable qlora, you can omit this if normal lora is preferred
)
lora_model = qwix.apply_lora_to_model(
model, lora_provider, **model.get_model_input()
)
# rest code with Tunix is not changed
```
## Batching Config
Tunix supports three hierarchical levels of batch size configuration for RL
training: **global batch size**, **mini-batch size**, and **micro-batch size**.
Let's deep dive into each of them and understand their performance implications
and when to tune them.
### Global Batch Size
Global batch size represents the total amount of data processed in one full RL
iteration, which is the product of dataloader batch size and the number of
generations per prompt. For example, if your input batch contains 256 prompts
and number of generations is set to 4 (for algorithms like GRPO, DAPO, etc),
then the effective global batch size will be 1024. For algorithms that don't
require group generation (e.g. PPO), global batch size equals dataloader's batch
size. **Note**: you don't directly tune the global batch size in Tunix, but
instead this is derived from the dataloader batch size and
[num_generation](https://github.com/google/tunix/blob/04f17e0b842901c13339874a25db77ba920adbc5/tunix/rl/grpo/grpo_learner.py#L52-L55)
configurations.
### Mini Batch Size
[Mini batch size](https://github.com/google/tunix/blob/04f17e0b842901c13339874a25db77ba920adbc5/tunix/rl/rl_cluster.py#L102)
determines how frequently the model is updated within a global step (commonly
known as mini-batch gradient update), which is used to reduce memory pressure on
the trainer for a large global batch and stabilize training. For a global batch
size of 1024 and mini batch size of 256, the model will be updated 4 times in a
global step. **NOTE**: Global batch size needs to be divisible by mini batch
size; if not specified, mini batch size defaults to global batch size.
### Micro Batch Size
Micro batching is a practical solution to further reduce HBM pressure while
enabling efficient, component-specific computation. Within the RL training
pipeline, HBM requirements vary significantly across different stages, such as
rollout, inference, and training. For example, the trainer typically has a much
higher memory footprint than the inference worker (which only handles forward
passes) because it must maintain optimizer states, gradients, and activations in
HBM. This means the effective batch size for the trainer is typically smaller than
the inference worker.
Tunix allows per-component micro-batching to maximize throughput across
different stages:
1. [rollout_micro_batch_size](https://github.com/google/tunix/blob/04f17e0b842901c13339874a25db77ba920adbc5/tunix/rl/rl_cluster.py#L104)
controls the input batch size to rollout engine. **NOTE**: with the new
agentic rollout, this is deprecated as the system automatically decomposes
batches into individual prompts to leverage continuous batching.
2. [compute_logps_micro_batch_size](https://github.com/google/tunix/blob/04f17e0b842901c13339874a25db77ba920adbc5/tunix/rl/rl_cluster.py#L105):
Controls the batch size for the inference worker during log-probability
calculation. Since this is forward-only, you can typically set this 2–4x
higher than the trainer's micro-batch size.
3. [train_micro_batch_size](https://github.com/google/tunix/blob/04f17e0b842901c13339874a25db77ba920adbc5/tunix/rl/rl_cluster.py#L103C3-L103C25):
Controls the trainer's input batch size. Setting this automatically enables
**gradient accumulation**, where the number of accumulation steps is
calculated as `mini_batch_size / train_micro_batch_size`. This value is
usually determined by your model size, sharding strategy, and available HBM.
### Tuning Tips
After determining your global batch size and mini batch size based on your
algorithm and training targets, follow these steps to optimize your hardware
utilization:
1. Based on the actor model's size, sharding strategy, and HBM capacity, decide
the effective `train_micro_batch_size`. Follow
[this guide](https://jax-ml.github.io/scaling-book/training/) for detailed
effective batch size estimation.
2. Set `compute_logps_micro_batch_size` to be 2x the size of
`train_micro_batch_size`. You can also try bumping it up to 4x for higher
throughput, but if you encounter OOM issues, lower it back to 2x.
3. Generally, you can set `rollout_micro_batch_size` to the same value as
`compute_logps_micro_batch_size` (you can try to push it higher as well). If
you are using the new agentic rollout, then you can omit this configuration.
## Collocated vs Disaggregated Training
Tunix provides two fundamental mechanisms for executing RL workflows: Collocated
and Disaggregated execution. Choose between them based on your available
resources and training requirements.
### Collocated Execution
In collocated mode, the entire TPU cluster is shared by all components (rollout,
inference, and trainer) throughout the workflow.
This setup is ideal for resource-constrained scenarios where the model must be
sharded across every available chip to fit into memory. In collocated mode, each
component will be executed in sequential order, which means the cluster will
finish rollout generation before shifting to inference and training. To further
maximize the hardware utility, you can consider enabling
[host_offloading](https://github.com/google/tunix/blob/main/rl/rl_cluster.py?q=offload_to_cpu)
with `offload_to_cpu`, which saves HBM by moving non-active models to CPU RAM
when a different component is occupying the TPU.
Enabling collocated mode is straightforward; you simply provide the same mesh to
every component when configuring the `role_to_mesh` mapping for your
`rl_engine`.
```python
import numpy as np
import jax
from jax.sharding import Mesh
from tunix.rl.rl_engine import ClusterConfig, Role
devices = jax.devices()
devices_mesh = np.array(devices).reshape(len(devices), 1)
mesh = Mesh(devices_mesh, axis_names=('fsdp', 'tp'))
ClusterConfig(
role_to_mesh={
Role.ACTOR: mesh,
Role.REFERENCE: mesh,
Role.ROLLOUT: mesh,
},
...
)
```
### Disaggregated Execution
Disaggregated mode partitions the TPU cluster into distinct "sub-meshes",
assigning dedicated hardware resources to specific components.
In the disaggregated setup, the rollout, inference worker, and trainer can execute
concurrently. This mode unlocks advanced performance optimization techniques. For example,
the inference worker and trainer can begin executing in parallel with the
rollout engine as soon as the first [mini-batch](performance.md#mini-batch-size) or [micro-batch](performance.md#micro-batch-size) of data is
produced. This is the preferred mode for large-scale runs where maximizing
global throughput is more critical.
To enable disaggregated mode, simply assign distinct meshes to each component
when defining the role_to_mesh mapping in your rl_engine configuration.
```python
import numpy as np
import jax
from jax.sharding import Mesh
from tunix.rl.rl_engine import ClusterConfig, Role
devices = jax.devices()
split = int(len(devices) / 2)
rollout_devices = np.array(devices[:split]).reshape(split, 1)
train_devices = np.array(devices[split:]).reshape(split, 1)
rollout_mesh = Mesh(rollout_devices, axis_names=('fsdp', 'tp'))
train_mesh = Mesh(train_devices, axis_names=('fsdp', 'tp'))
ClusterConfig(
role_to_mesh={
Role.ACTOR: train_mesh,
Role.REFERENCE: train_mesh,
Role.ROLLOUT: rollout_mesh,
},
...
)
```
## Weight transfer/sync
Weight transfer and sync is critical to keeping rollout models up-to-date with
the policy model. This process is controlled by RLEngine regardless of which
rollout engine is used (vanilla, vllm, or sglang-jax), and requires no user
intervention.
Rollout models often have different mesh or sharding specifications than policy
models, meaning that when syncing weights, resharding is required to place the
new weights onto the right devices with the right shards. Efficient weight sync
improves end-to-end performance and reduces HBM and host memory capacity
requirements.
The Tunix resharding algorithm eliminates expensive all-gather collective operations
between different meshes when needed. It resolves the OOM issue where all-gather
replicates the model weights to each TPU device. Please see
[this section](https://github.com/google/tunix/blob/main/rl/reshard.py;l=70-93)
for more details.
The following diagram illustrates how the Tunix resharding algorithm avoids OOM caused by
all-gather ops.
## /docs/quickstart.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Quick Start
This page contains several quickstart guides and is a great place to understand
how to get started with Tunix. It covers installation and provides several
hands-on examples across the board for SFT, RL, and Agentic RL training.
Additionally, it shows how to enable to multi-node training.
## Installation
Tunix is written in Python and **requires Python 3.11** or later. We recommend
installing Tunix in a Python virtual environment.
1. Create a project specific environment.
```sh
python3 -m venv .venv
# Or simply `python -m venv .venv` depending on your system configuration.
```
2. Activate the Environment
```sh
source .venv/bin/activate
```
3. Install Tunix dependency
Make sure you have an updated pip version installed:
```sh
pip install --upgrade pip
```
There are several ways to install Tunix. Please select one from below.
### Option A: From PyPI (**Recommended**)
You can install the latest stable release of Tunix from PyPI. Tunix relies on
JAX for computation, which must be installed with support for your specific
hardware (TPU, GPU, or CPU).
**TPU**
Tunix is optimized for execution on TPUs. If you have TPU hardware, you can
install Tunix and JAX with TPU support by specifying the `[prod]` extra:
```sh
pip install "google-tunix[prod]"
```
**GPU**
If you are using GPUs, first install Tunix, then install JAX with GPU (CUDA)
support. You may need to adjust the CUDA version based on your system setup.
Refer to the
[JAX installation guide](https://github.com/google/jax#installation) for more
details.
```sh
pip install google-tunix
# Install JAX with CUDA 13 support
pip install -U "jax[cuda13]"
```
**CPU**
To run Tunix in a CPU-only environment:
```sh
pip install google-tunix "jax[cpu]"
```
### Option B: From GitHub
You can install the latest development version directly from GitHub:
```sh
# For TPU
pip install "git+https://github.com/google/tunix#egg=google-tunix[prod]"
# For GPU/CPU
pip install git+https://github.com/google/tunix
# Then install JAX for GPU or CPU as described above.
```
### Option 3: From Source
If you plan to modify Tunix, you can perform an editable installation from a
local clone of the repository:
```sh
git clone https://github.com/google/tunix.git
cd tunix
pip install -e ".[dev]"
# Then install JAX for your hardware as described above.
```
For TPU development, you can use:
```sh
pip install -e ".[prod]"
```
### Optional Dependencies
For accelerated inference, Tunix supports integration with vLLM and SGLang-Jax.
These need to be installed manually.
**vLLM on TPU**
The TPU-inference supported version of `vllm` is not always available as a
single PyPI release, and installing the TPU build sometimes requires extra pip flags
so that `libtpu` wheels (hosted by the JAX project) can be resolved. You can
install the pinned vLLM + TPU requirements from this repository using one of
the raw requirement-file URLs below.
Install from remote:
```sh
pip install -r https://github.com/google/tunix/raw/main/requirements/requirements.txt
pip install -r https://github.com/google/tunix/raw/main/requirements/special_requirements.txt
```
Or (direct raw.githubusercontent URL):
```sh
pip install -r https://raw.githubusercontent.com/google/tunix/main/requirements/requirements.txt
pip install -r https://raw.githubusercontent.com/google/tunix/main/requirements/special_requirements.txt
```
If you prefer a single-line install that directly overrides `tpu-inference`, you can also run:
```sh
pip install vllm @git+https://github.com/vllm-project/vllm.git@<commit>
pip install --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \\
--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \\
--pre \\
tpu-inference@git+https://github.com/vllm-project/tpu-inference.git@<commit>
```
Or install from source:
```
bash scripts/install_tunix_vllm_requirement.sh
```
**SGLang-Jax**
After installing Tunix, you can install SGLang-Jax from source:
```sh
git clone git@github.com:sgl-project/sglang-jax.git
cd sglang-jax/python
pip install -e .
```
**GCS File System**
If you need to access models or data stored in Google Cloud Storage (GCS), e.g.,
this is commonly used as the default option for Gemma3 models when using Tunix
CLI, you may need to install `gcsfs`:
```sh
pip install gcsfs
```
## Quick start: GRPO
To get started with the library, let's walk through an example of training (full
, LoRA and QLoRA fine-tuning) the Gemma 3 270M model on the English-to-French
translation dataset. We will use Tunix's `PeftTrainer` for this task.
Note: This example is meant to be a quick-start. For the complete example, refer
to
[this](https://github.com/google/tunix/blob/main/examples/qlora_gemma.ipynb)
notebook.
### Load the model
First up, let's load the model:
```python
from huggingface_hub import snapshot_download
import jax
from tunix.models.gemma3 import model as gemma_lib
from tunix.models.gemma3 import params_safetensors as params_safetensors_lib
# Define sharding mesh for the model (assuming 1 TPU).
MESH = [(1, 1), ("fsdp", "tp")]
mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0]))
# Load the model.
model_id = "google/gemma-3-270m-it"
model_path = snapshot_download(
repo_id=model_id, ignore_patterns=["*.pth"]
)
config = gemma_lib.ModelConfig.gemma3_270m()
with jax.set_mesh(mesh):
model = params_safetensors_lib.create_model_from_safe_tensors(
model_path, config, mesh
)
```
Note: we could have simply used Tunix's `AutoModel` class, but don't use it here
since Gemma 3 isn't supported for now. `AutoModel` is the preferred way of
loading models.
### Load and preprocess the dataset
Next, we load the English-French translation dataset. Note you can use your own
datasets too (PyGrain, Hugging Face dataset, TFDS, etc.).
```sh
gcloud storage cp gs://gemma-data/tokenizers/tokenizer_gemma3.model .
```
```python
from tunix.generate import tokenizer_adapter
from tunix.examples.data import translation_dataset as data_lib
tokenizer = tokenizer_adapter.Tokenizer(
tokenizer_path="./tokenizer_gemma.model")
train_ds, val_ds = data_lib.create_datasets(
'mtnt/en-fr',
global_batch_size=64,
max_target_length=256,
num_train_epochs=3,
tokenizer=tokenizer,
)
```
We need to process the inputs to make sure we are feeding the data to the model
in the right format.
```python
from tunix.sft import utils
def input_fn(x):
mask = x.input_tokens != tokenizer.pad_id()
return {
'input_tokens': x.input_tokens, 'input_mask': x.input_mask,
'positions': utils.build_positions_from_mask(mask),
'attention_mask': utils.make_causal_attn_mask(mask),
}
```
### Train the model
#### Full fine-tuning
We can now train our model. We need to pass the `input_fn` defined above here:
```python
import optax
from tunix.sft import peft_trainer
training_config = peft_trainer.TrainingConfig(
eval_every_n_steps=20,
max_steps=100,
)
trainer = peft_trainer.PeftTrainer(
model=model,
optimizer=optax.adamw(learning_rate=1e-4),
training_config=training_config,
).with_gen_model_input_fn(input_fn)
trainer.train(train_ds=train_ds, eval_ds=val_ds)
```
#### LoRA/QLoRA fine-tuning
The above case handles the full SFT case where all model parameters are updated.
We can choose to use LoRA. In this case, we just need to use Qwix, like so:
```python
import qwix
lora_provider = qwix.LoraProvider(
module_path=".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj",
rank=RANK,
alpha=ALPHA,
# for QLoRA, uncomment the lines below.
# weight_qtype="nf4",
# tile_size=128,
)
model_input = model.get_model_input()
lora_model = qwix.apply_lora_to_model(
model, lora_provider, **model_input
)
with jax.set_mesh(mesh):
state = nnx.state(lora_model)
pspecs = nnx.get_partition_spec(state)
sharded_state = jax.lax.with_sharding_constraint(state, pspecs)
nnx.update(lora_model, sharded_state)
```
The rest of the flow remains the same.
### Evaluate the model
To evaluate the model, we can use the `Sampler` API to generate outputs.
```python
from tunix.generate import sampler as sampler_lib
sampler = sampler_lib.Sampler(
transformer=lora_model,
tokenizer=tokenizer,
cache_config=sampler_lib.CacheConfig(
cache_size=256,
num_layers=model_config.num_layers,
num_kv_heads=model_config.num_kv_heads,
head_dim=model_config.head_dim,
),
)
input_batch = [
"Translate this into French:\nHello, my name is Morgane.\n",
"Translate this into French:\nThis dish is delicious!\n",
]
out_data = sampler(
input_strings=input_batch,
max_generation_steps=10, # number of generated tokens
)
```
### Trajectory Logging
During reinforcement learning (RL) training, it is often useful to analyze the
generated trajectories (prompts, responses, rewards, etc.). Tunix provides an
`AsyncTrajectoryLogger` to log this data asynchronously to CSV files without
blocking the training loop. It's enabled in agentic_grpo_learner by default, if
you provide a log directory in your cluster configuration training config.
```python
# In your cluster configuration setup
cluster_config.training_config.metrics_logging_options.log_dir = "./logs"
# GCS paths are also supported
```
When enabled, the learner will automatically log trajectories during the
training process. Users can then consume the logged data by loading the CSV
files into a pandas DataFrame or other query engine.
## Quick Start: Multi-Node Training
Tunix supports running on a multi-node setup using Pathways in GKE ([more details](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). This is a
transparent change that simply requires you to submit your job through Pathways
instead of running directly on a VM. To run Tunix in a multi-node Pathways
cluster basically requires 3 steps: 1. create a Pathways cluster, 2. Build a
docker image, 3. launch a Tunix job. The following sections cover each step in
further detail.
### 1. Create a Pathways cluster in GKE
#### Install xpx
We will use XPK to create a Pathways cluster in GKE.
```sh
pip install xpk
```
#### Install gcloud cli
For Debian or Ubuntu, install gcloud via apt. Make sure prerequisites are met:
```sh
sudo apt-get update
sudo apt-get install apt-transport-https ca-certificates gnupg curl
```
Import the Google Cloud public key:
```sh
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg
```
Add the Google Cloud CLI distribution URI as a package source:
```sh
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
```
Update and install:
```sh
sudo apt-get update && sudo apt-get install google-cloud-cli
```
#### Create a Pathways cluster
Then we will create the Pathways cluster.
```sh
# install gcloud beta commands
gcloud components install beta
# create pathways cluster
export CLUSTER_NAME='your-cluster-name'
export ZONE='your-tpu-zones'
export TPU_TYPE='your-tpu-type' # e.g. v5p-16
export CLUSTER_CPU_MACHINE_TYPE=n2d-standard-32 # you can adjust this to use beefier CPU node
export PROJECT='your-gke-projec'
NETWORK_NAME=${CLUSTER_NAME}-mtu9k-wx
NETWORK_FW_NAME=${NETWORK_NAME}-fw-wx
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
# run `gcloud auth application-default login` and
# `gcloud auth login --update-adc` if you encounter permission issue when creating the network.
# Check if this is the service account you want to use.
gcloud auth list
gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT} \
--subnet-mode=auto \
--bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
--network ${NETWORK_NAME} \
--allow tcp,icmp,udp \
--project=${PROJECT}
xpk cluster create-pathways \
--cluster $CLUSTER_NAME \
--cluster-cpu-machine-type=$CLUSTER_CPU_MACHINE_TYPE \
--num-slices=1 \
--tpu-type=$TPU_TYPE \
--zone $ZONE \
--project $PROJECT \
--custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
### 2. Build a Tunix Docker Image
Build local docker image. We will be using the `build_docker.sh`
[script](https://github.com/google/tunix/blob/main/build_docker.sh).
in the `tunix` directory. If you do not have docker on glinux, please install it
before building the docker image.
```sh
# cleanup unused docker images and caches if disk is not enough
sudo docker system prune
bash ./build_docker.sh
# It will default to generate a local docker image
export LOCAL_IMAGE_NAME=tunix_base_image
# You can also optionally push to GKE's artifact registry for faster download in the future
```
### 3. Launch the job
Now you are ready to submit your Tunix workload. You will use `xpk` to do this,
similar to the cmd below.
```sh
xpk workload create-pathways \
--cluster=$CLUSTER_NAME \
--workload=$WORKLOAD_NAME \
--command="TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE='1' source your-script-to-launch-job.sh" \
--num-slices=1 \
--tpu-type=$TPU_TYPE \
--base-docker-image docker.io/library/tunix_base_image \
--priority=medium
```
## Next Steps
Now that you've completed the quick start, you can explore other training
techniques and models. In particular, the following would be worth exploring:
- [SFT and PEFT](https://github.com/google/tunix/blob/main/examples/qlora_gemma.ipynb)
- [Agentic RL](https://github.com/google/tunix/blob/main/examples/math_gsm8k/gemma_grpo_demo_nb.py)
A complete list is given [here](examples.md).
## /docs/reliability.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Reliability
## Checkpoint Support
Tunix provides robust checkpointing capabilities to save and resume training
progress, leveraging Orbax as the backend. This includes saving model parameters
(supporting full state or only LoRA parameters for PEFT) and optimizer state.
Checkpointing is managed by [`CheckpointManager`](https://github.com/google/tunix/tree/main/tunix/sft/checkpoint_manager.py?q=symbol:CheckpointManager)
and integrated into [`PeftTrainer`](https://github.com/google/tunix/tree/main/tunix/sft/peft_trainer.py?q=symbol:PeftTrainer).
SFT uses `PeftTrainer` directly, while RL uses [`rl.Trainer`](https://github.com/google/tunix/tree/main/tunix/rl/trainer.py?q=symbol:Trainer),
a subclass of `PeftTrainer`, inside of the [`RLLearner`](https://github.com/google/tunix/tree/main/tunix/rl/rl_learner.py?q=symbol:RLLearner).
Therefore, both SFT and RL share the same checkpointing mechanism. Checkpointing
and restarting are built-in features that require no special setup beyond
configuration. To enable checkpointing, users simply need to set
`checkpoint_root_directory` in `SFTConfig` or `RLConfig`; if this path is
provided, Tunix automatically saves checkpoints and resumes training from the
most recent one if interrupted, restoring model weights, optimizer state, and
training step count. By default, checkpointing is disabled if
`checkpoint_root_directory` is not specified. Users can further customize
checkpointing behavior via `checkpointing_options` in the config.
Users customize background preservation behavior granularly using components
defined inside `checkpoint_options`:
* **Save Decision Policies**: Dictates when to initiate a checkpoint based on
defined steps or intervals. Supported configurations include
`FixedIntervalPolicy` and `ContinuousCheckpointingPolicy`. The default is
`ContinuousCheckpointingPolicy(minimum_interval_secs=180)` (saves every 180
seconds). See Orbax v1 [`save_decision_policies.py`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/experimental/v1/_src/training/save_decision_policies.py)
for the complete interface contracts.
* **Preservation Policies**: Sets specifications regarding tracking
checkpoints over bounded timelines (e.g., `LatestN`). The default is
`LatestN(n=3)` (keeps the latest 3 checkpoints). See Orbax v1
[`preservation_policies.py`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/experimental/v1/_src/training/preservation_policies.py)
for the complete interface contracts.
* **Step Name Format**: Defines the representation of directory names for step
checkpoints. The default is `ocp.path.step.standard_name_format()` (uses
simple integer step names).
* **Asynchronous Processing**: Manage asynchronous behavior by specifying:
* `enable_async_checkpointing`: Whether to use async checkpointing.
Defaults to `True`. **It is recommended to keep this enabled** to
prevent the main thread from blocking during training runs while
checkpoints are written to storage.
* `timeout_secs`: The timeout for asynchronous operations.
Defaults to `1200` seconds.
## Fault Tolerance
Tunix ensures fault tolerance primarily through its checkpointing mechanism,
allowing training to resume after interruptions such as machine restarts or
pre-emptions.
Additionally, to prevent out-of-memory (OOM) errors due to excessive HBM usage,
Tunix includes an [`InflightThrottler`](https://github.com/google/tunix/tree/main/tunix/sft/inflight_throttler.py?q=symbol:InflightThrottler). This mechanism limits the
number of TPU computations that can be scheduled concurrently, as configured by
`max_inflight_computations` in `TrainingConfig`, thus providing more stable
training runs on memory-constrained hardware.
## Determinism Guarantee
Tunix supports deterministic training runs through careful management of random
number generation and data handling:
* **Model Initialization**: Models can be initialized with a specific random
seed (`rng_seed` or `random_seed`) to ensure consistent initial weights
across runs.
* **Data Shuffling**: RL learners accept a `data_shuffle_seed` parameter,
which ensures that dataset shuffling is deterministic.
* **Dropout and Stochastic Layers**: JAX and Flax RNG handling ensures that
stochastic operations can be made deterministic if RNGs are correctly seeded
and managed.
By providing explicit seeds for these components, users can ensure
reproducibility of training experiments.
## /docs/talks.md
<!-- DO NOT REMOVE! Placeholder for TOC. -->
# Talks and Announcements
Track the latest developments in the Tunix ecosystem, including technical deep
dives, launch news, and collaborations.
## Announcements and Blogs
| Date | Type | Title / Description | Link |
| :--- | :--- | :--- | :--- |
| **2025-12-16** | Blog | **GRL**<br>Turning verifiable games into a post-training suite for LLM agents<br>with Tunix on TPUs | [Read](https://opensource.googleblog.com/2025/12/grl-turning-verifiable-games-into-a-post-training-suite-for-llm-agents-with-tunix-on-tpus.html) |
| **2025-12-11** | Blog | **Empowering App Developers**<br>Fine-tuning Gemma 3 for mobile with Tunix in Google Colab. | [Read](https://opensource.googleblog.com/2025/12/empowering-app-developers-fine-tuning-gemma-3-for-mobile-with-tunix-in-google-colab.html) |
| **2025-11-14** | Hackathon | **Google Tunix Hackathon (Kaggle)**<br>Challenge to train reasoning models **(active)**. | [View](https://www.kaggle.com/competitions/google-tunix-hackathon) |
| **2025-09-30** | Launch | **Introducing Tunix: A JAX-Native Library**<br>Official launch post. | [Read](https://developers.googleblog.com/introducing-tunix-a-jax-native-library-for-llm-post-training/) |
## Talks
| Date | Event | Title / Description | Link |
| :--- | :--- | :--- | :--- |
| **2025-12-11** | JAX/OpenXLA DevLab | **Tunix: Tune in JAX**<br>Core team deep dive into the latest framework features<br>followed by SFT, GRPO and PPO demos. | [Watch](https://www.youtube.com/watch?v=NMrL3ZMrqA0), [Slides](https://docs.google.com/presentation/d/1M5_2B8bUmaxCL6QVwPq0Yf8YEErihu_i2179ClwmhkY/edit?usp=sharing) |
| **2025-09-30** | Google for Developers | **How to fine-tune LLMs with Tunix**<br>Introductory guide to SFT and RL on TPUs. | [Watch](https://www.youtube.com/watch?v=8essLqkBsX8) |
| **2025-09-24** | JAX DevLab (Summer) | **JAX: Tunix (Lightning Talk)**<br>Brief intro to the library. | [Watch](https://www.youtube.com/watch?v=Kcw6GTqQjBE), [Slides](https://docs.google.com/presentation/d/1bl8cdoCFpNyBVCCNxLnK2Wpl6jTPwiL5Oc7-3kH3s5k/edit?usp=sharing) |
## /examples/README.rst
```rst path="/examples/README.rst"
Tuning
======
Fine-tuning examples using Google Tunix.
Notebooks
---------
The following notebooks provide comprehensive examples of different fine-tuning techniques:
* ``qlora_gemma.ipynb`` - LoRA and QLoRA fine-tuning with Gemma models. Demonstrates parameter-efficient fine-tuning techniques using low-rank adaptation.
* ``grpo_gemma.ipynb`` - GRPO (Group Relative Policy Optimization) reinforcement learning. Shows how to fine-tune models using policy optimization for improved response generation.
* ``dpo_gemma.ipynb`` - DPO (Direct Preference Optimization). Demonstrates preference-based fine-tuning to align model outputs with desired behaviors.
* ``logit_distillation.ipynb`` - Knowledge distillation from larger models. Shows how to transfer knowledge from a teacher model to a student model.
Subdirectories
--------------
``deepscaler/``
~~~~~~~~~~~~~~~
Contains scripts for training and evaluating models with DeepScaler:
* ``train_deepscaler_nb.py`` - Training script for DeepScaler models
* ``math_eval_nb.py`` - Mathematical reasoning evaluation utilities
``model_load/``
~~~~~~~~~~~~~~~
Examples for loading models from different formats:
* ``from_safetensor_load/`` - Contains notebooks for loading Gemma2 and Gemma3 models from safetensors format
* ``gemma2_model_load.ipynb``
* ``gemma3_model_load.ipynb``
``rl/``
~~~~~~~
Reinforcement learning examples and hardware resource requirements:
* ``grpo/gsm8k/`` - GRPO implementation scripts for GSM8K mathematical reasoning tasks
* Launch scripts for various models (Gemma 7b, Gemma2 2b, Llama3.2 1b/8b)
* ``README.md`` - Detailed hardware resource requirements and configuration recommendations for RL training
``sft/``
~~~~~~~~
Supervised fine-tuning examples:
* ``mtnt/`` - MTNT translation task examples with launch scripts for multiple models
* Launch scripts for Gemma 2b, Gemma2 2b, Gemma3 4b, Llama3.2 3b, Qwen2.5 0.5b
* ``README.md`` - Hardware resource requirements for SFT training
GCE VM Setup for Fine-Tuning
----------------------------
1. Create TPU VM
~~~~~~~~~~~~~~~~
Create a v5litepod-8 TPU VM in GCE:
* SW version: ``v2-alpha-tpuv5-lite``
* Name: ``v5-8``
Reference: `TPU Runtime Versions <https://docs.cloud.google.com/tpu/docs/runtimes?hl=en&_gl=1*1tpeg3j*_ga*MTk1NzE5MjMyNy4xNzYwOTEwNjk3*_ga_WH2QY8WWF5*czE3NjIxNTU1OTEkbzE3JGcwJHQxNzYyMTU1NTkxJGo2MCRsMCRoMA..#training-v5p-v5e>`_
2. Configure VM
~~~~~~~~~~~~~~~~
SSH into the VM using the supplied gcloud command, then run:
.. code-block:: bash
# Create .env file with required credentials
vim .env
# Download and install Anaconda
curl -O https://repo.anaconda.com/archive/Anaconda3-2025.06-0-Linux-x86_64.sh
bash ~/Anaconda3-2025.06-0-Linux-x86_64.sh # always input "yes"/enter
source ~/.bashrc
# Create conda environment (Python 3.12 - MUST BE 12, NOT 11!)
conda create -n colab python=3.12 -y
conda activate colab
# Install dependencies
pip install 'ipykernel<7' jupyterlab
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install --upgrade clu
Reference: `Run JAX on TPU <https://docs.cloud.google.com/tpu/docs/run-calculation-jax>`_
Exit the SSH session after setup is complete.
3. Connect from Local Machine
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
From your local machine, run the following to connect to Jupyter Lab:
.. code-block:: bash
gcloud compute tpus tpu-vm ssh v5-8 --zone=us-west1-c \
-- -L 8080:localhost:8080 -L 6006:localhost:6006 \
"source \$HOME/anaconda3/etc/profile.d/conda.sh && \
conda activate colab && \
jupyter lab \
--ServerApp.allow_origin='https://colab.research.google.com' \
--port=8080 \
--no-browser \
--ServerApp.port_retries=0 \
--ServerApp.allow_credentials=True"
Reference: `Local Runtimes in Colab <https://research.google.com/colaboratory/local-runtimes.html>`_
4. Environment Variables
~~~~~~~~~~~~~~~~~~~~~~~~
Example ``.env`` file:
.. code-block:: bash
HF_TOKEN=
KAGGLE_USERNAME=
KAGGLE_KEY=
WANDB_API_KEY=
Loading Saved Safetensors Models
--------------------------------
To load a saved safetensors model back into JAX (with a given local_path):
.. code-block:: python
import os
import jax
import jax.numpy as jnp
from tunix.models.gemma3 import params_safetensors as params_safetensors_lib
local_path = '[PLACEHOLDER]'
MESH = [(1, 1), ("fsdp", "tp")]
mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0]))
with mesh:
model = params_safetensors_lib.create_model_from_safe_tensors(
os.path.abspath(local_path), (model_config), mesh, dtype=jnp.bfloat16
)
Notes
-----
* **IMPORTANT**: Use ``%pip`` not ``!pip`` in notebooks!
* Python 3.12 is the recommended version
.. toctree::
:glob:
:maxdepth: 1
:hidden:
*
```
## /examples/deepscaler/run_deepscaler_disagg_v5p16.sh
```sh path="/examples/deepscaler/run_deepscaler_disagg_v5p16.sh"
#!/bin/bash
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# DeepScaleR training launcher using tunix/cli/base_config.yaml plus explicit
# CLI overrides.
#
# Usage:
# checkpoint_dir="" bash examples/deepscaler/run_deepscaler_disagg.sh
#
# Run from the tunix repo root.
set -euo pipefail
export SKIP_JAX_PRECOMPILE=true
num_batches="${num_batches:-312}"
num_train_epochs="${num_train_epochs:-3}"
train_fraction="${train_fraction:-1.0}"
warmup_ratio="${warmup_ratio:-0.1}"
batch_size="${batch_size:-128}"
mini_batch_size="${mini_batch_size:-128}"
max_response_length="${max_response_length:-8192}"
#TODO(b/510820709) - find the optimal mesh configuration
total_tpus="${total_tpus:-8}"
trainer_mesh="${trainer_mesh:-(4,1)}"
rollout_mesh="${rollout_mesh:-(4,1)}"
source "$(dirname "$0")/../tpu_utils.sh"
validate_mesh_allocation "$total_tpus" "$trainer_mesh" "$rollout_mesh" "null" || exit 1
checkpoint_dir="${checkpoint_dir:-gs://tunix/rl/checkpoints/01}"
checkpoint_suffix="${checkpoint_suffix:-$(printf '%04d' "$((RANDOM % 10000))")}"
if [[ -n "$checkpoint_dir" && "$checkpoint_dir" != "null" ]]; then
checkpoint_dir="${checkpoint_dir}_${checkpoint_suffix}"
fi
max_steps=$(awk "BEGIN {
value = $num_batches * $num_train_epochs * $train_fraction;
if (value < 1) value = 1;
printf \"%.0f\", value;
}")
warmup_steps=$(awk "BEGIN {
value = $warmup_ratio * $max_steps;
if (value < 1) value = 1;
printf \"%.0f\", value;
}")
python -m tunix.cli.grpo_main \
tunix/cli/base_agentic_config.yaml \
\
`# ── Model ────────────────────────────────────────────────────────────` \
model_config.model_name="deepseek_r1_distill_qwen_1_5b" \
model_config.model_id="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" \
model_config.model_source="huggingface" \
model_config.rng_seed=42 \
model_config.model_display=false \
model_config.remat_config=3 \
actor_model_config.mesh.shape="$trainer_mesh" \
actor_model_config.mesh.axis_names="('fsdp','tp')" \
reference_model_config.mesh=null \
reference_model_config.same_mesh_as="actor" \
rollout_model_config.mesh.shape="$rollout_mesh" \
rollout_model_config.mesh.axis_names="('fsdp','tp')" \
\
`# ── Data ─────────────────────────────────────────────────────────────` \
data_module="tunix.cli.recipes.deepscaler_data" \
data_config.train_data_path="gs://tunix/data/DeepScaleR-Preview-Dataset/deepscaler.json" \
data_config.eval_data_path="gs://tunix/data/HuggingFaceH4/aime_2024/train-00000-of-00001.parquet" \
data_config.shuffle=true \
data_config.seed=42 \
prompt_key="prompts" \
\
`# ── Training loop ────────────────────────────────────────────────────` \
training_config.profiler_options.log_dir="gs://lancewang-dev-supercomputer-testing/tunix/rl/grpo/profiler" \
training_config.profiler_options.profiler_steps=3 \
training_mode="agentic_grpo" \
batch_size="$batch_size" \
num_batches="$num_batches" \
num_train_epochs="$num_train_epochs" \
train_fraction="$train_fraction" \
reward_functions=["tunix/utils/math_rewards.py"] \
verl_compatible=false \
\
`# ── Rollout engine (vanilla | vllm | sglang_jax) ─────────────────────` \
rollout_engine="vllm" \
offload_to_cpu=false \
\
`# ── Rollout config ───────────────────────────────────────────────────` \
rollout_config.max_prompt_length=2048 \
rollout_config.total_generation_steps="$max_response_length" \
rollout_config.max_tokens_to_generate="$max_response_length" \
rollout_config.temperature=0.6 \
rollout_config.top_p=null \
rollout_config.top_k=null \
rollout_config.return_logprobs=true \
\
`# ── SGLang-JAX (used when rollout_engine=sglang_jax) ─────────────────` \
sglang_jax_config.mem_fraction_static=0.8 \
sglang_jax_config.init_with_random_weights=true \
sglang_jax_config.disable_radix_cache=true \
sglang_jax_config.enable_deterministic_sampling=false \
sglang_jax_config.chunked_prefill_size=2048 \
sglang_jax_config.page_size=128 \
sglang_jax_config.use_sort_for_toppk_minp=false \
\
`# ── vLLM (used when rollout_engine=vllm) ─────────────────────────────` \
vllm_config.hbm_utilization=0.4 \
vllm_config.tpu_backend_type="jax" \
vllm_config.server_mode=true \
vllm_config.async_scheduling=true \
vllm_config.max_num_seqs=768 \
vllm_config.kwargs.kv_cache_metrics=true \
vllm_config.kwargs.disable_log_stats=false \
vllm_config.kwargs.enable_prefix_caching=true \
\
`# ── Chat / agent wiring ───────────────────────────────────────────────` \
chat_parser_config.type="default" \
tokenizer_config.tokenizer_type="huggingface" \
tokenizer_config.tokenizer_path="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" \
tokenizer_config.add_bos=false \
tokenizer_config.add_eos=false \
\
`# ── GRPO algorithm ───────────────────────────────────────────────────` \
agentic_grpo_config.num_generations=8 \
agentic_grpo_config.num_iterations=1 \
agentic_grpo_config.beta=0.0 \
agentic_grpo_config.epsilon=0.2 \
agentic_grpo_config.epsilon_high=0.28 \
agentic_grpo_config.system_prompt="" \
agentic_grpo_config.max_concurrency=1024 \
agentic_grpo_config.max_response_length="$max_response_length" \
agentic_grpo_config.off_policy_steps=0 \
agentic_grpo_config.loss_agg_mode="token-mean" \
agentic_grpo_config.kl_loss_mode="low_var_kl" \
agentic_grpo_config.max_turns=1 \
\
`# ── Optimizer ────────────────────────────────────────────────────────` \
rl_training_config.actor_optimizer_config.opt_type="adamw" \
rl_training_config.actor_optimizer_config.learning_rate=1e-6 \
rl_training_config.actor_optimizer_config.schedule_type="cosine_decay_schedule" \
rl_training_config.actor_optimizer_config.init_value=1e-6 \
rl_training_config.actor_optimizer_config.end_value=0.0 \
rl_training_config.actor_optimizer_config.warmup_ratio="$warmup_ratio" \
rl_training_config.actor_optimizer_config.warmup_steps="$warmup_steps" \
rl_training_config.actor_optimizer_config.decay_steps="$max_steps" \
rl_training_config.actor_optimizer_config.b1=0.9 \
rl_training_config.actor_optimizer_config.b2=0.99 \
rl_training_config.actor_optimizer_config.weight_decay=0.01 \
rl_training_config.actor_optimizer_config.max_grad_norm=1.0 \
\
`# ── RL training ──────────────────────────────────────────────────────` \
rl_training_config.eval_every_n_steps=1000 \
rl_training_config.max_steps="$max_steps" \
rl_training_config.mini_batch_size="$mini_batch_size" \
rl_training_config.train_micro_batch_size=2 \
rl_training_config.checkpoint_root_directory="$checkpoint_dir" \
rl_training_config.checkpointing_options.save_interval_steps=500 \
rl_training_config.checkpointing_options.max_to_keep=4 \
rl_training_config.metrics_logging_options.log_dir="/tmp/tensorboard/deepscaler" \
rl_training_config.metrics_logging_options.flush_every_n_steps=20 \
\
"$@"
```
## /examples/deepswe/guarded_swe_env.py
```py path="/examples/deepswe/guarded_swe_env.py"
"""SWE environment wrapper with lightweight action guarding."""
from action_guard import ActionGuard, GuardConfig
from swe_env import SWEEnv
from tunix.rl.agentic.environments.base_environment import EnvStepResult
class GuardedSWEEnv(SWEEnv):
"""SWEEnv with failure-aware action restrictions applied in step()."""
def __init__(self, *args, guard_config=None, **kwargs):
self.guard = ActionGuard(guard_config or GuardConfig())
super().__init__(*args, **kwargs)
def _initial_observation(self):
self.guard.reset()
return super()._initial_observation()
def _step_impl(self, action):
if isinstance(action, str):
func_name, _ = self.guard._parse_action(action) # pylint: disable=protected-access
if not func_name:
return EnvStepResult(
observation=(
"[ACTION GUARD] Your previous response did not include a"
" valid function call. You must output exactly one tool call"
" in the required XML format."
),
reward=0.0,
done=False,
info={
"guard_blocked": True,
"guard_reason": "missing_function_call",
},
)
verdict = self.guard.evaluate(action)
if verdict.blocked:
return EnvStepResult(
observation=verdict.message,
reward=0.0,
done=False,
info={
"guard_blocked": True,
"guard_reason": verdict.reason,
},
)
result = super()._step_impl(action)
if isinstance(action, str):
self.guard.record_outcome(action, str(result.observation))
return result
```
## /examples/frozenlake/run_gemma4_e2b.sh
```sh path="/examples/frozenlake/run_gemma4_e2b.sh"
#!/bin/bash
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Install dependencies:
pip install -e .[frozenlake]
set -x # Enable xtrace
batch_size=${batch_size:-64}
num_batches=${num_batches:-5}
echo "Using parameters:"
echo " Batch Size: $batch_size"
echo " Num Batches: $num_batches"
python3 -m tunix.cli.grpo_main \
tunix/cli/base_agentic_config.yaml \
override_config_file=examples/frozenlake/configs/gemma4_e2b.yaml \
batch_size=$batch_size \
num_batches=$num_batches \
"$@"
```
## /tunix/examples/__init__.py
```py path="/tunix/examples/__init__.py"
```
The content has been capped at 50000 tokens. The user could consider applying other filters to refine the result. The better and more specific the context, the better the LLM can follow instructions. If the context seems verbose, the user can refine the filter using uithub. Thank you for using https://uithub.com - Perfect LLM context for any GitHub repo.