Tunix Files

This release focuses on improving performance and stability across TPU and Kaggle environments, introducing new utilities for agentic RL workflows, and adding broader model and configuration support. It also includes several important bug fixes and developer experience improvements.

Run Tunix on Kaggle TPU

We’re excited to announce that Tunix can now be launched directly in Kaggle notebooks with TPU acceleration — making it easier than ever to experiment, prototype, and run reinforcement learning workflows without complex setup.

Key highlights

First-class TPU support on Kaggle – run GRPO and other RL pipelines end-to-end in a Kaggle notebook.

Pre-configured runtime – no manual dependency juggling needed; version compatibility and performance tuning are handled automatically.

Launch the notebook here: Knowledge Distillation Demo QLoRA Demo DPO Demo GRPO Demo

New Features & Improvements

• Model & Training Options

• Added support for Gemma-3-270M model configuration.
• Enabled setting default parameter dtype for Gemma-3 models.
• Added remat options to models to improve memory efficiency.
• Created a new list container type to support both Flax ≤0.11.2 and ≥0.12.0 versions.

• Pathways & TPU Performance

• Introduced experimental pre-sharding (experimental_reshard) for Pathways on Cloud TPU.
• Improved weight synchronization logic to handle KV head duplication.
• Disabled certain profiler options by default to improve stability on Pathways backend.

• Configuration & CLI Improvements

• Enabled generic creation of optax.optimizer and optax.learning_rate_schedule directly from CLI.
• Relaxed JAX version constraints to ensure compatibility with Kaggle images.
• Added minimum resource requirements for launch scripts in the README.

• Documentation

• Added ReadTheDocs link in README.
• Expanded external notebooks with step-by-step guidance for long-running tasks.

Bug Fixes

• Fixed a bug in reward function logic causing incorrect training signals.
• Fixed a checkpoint handling issue where Colab failed to locate the final checkpoint and now cleans up intermediate directories.
• Fixed Kaggle image performance issues.
• Fixed type errors in agents/ modules.
• Optimized masked index lookups using jnp.where for better runtime efficiency.
• Resharded prompt and completion tokens to the REFERENCE mesh when rollout and reference models are distributed.

Dependency & Version Updates

• JAX pinned to 0.7.1 and libtpu downgraded to resolve Cloud TPU performance regressions.
• Relaxed JAX version requirement for Kaggle compatibility.

Full Changelog:

• Bump up the version to v0.1.0 by @wang2yn84 in https://github.com/google/tunix/pull/446
• Delete this notebook as it's redundant now. Prepare for the release. by @copybara-service[bot] in https://github.com/google/tunix/pull/445
• Add min resources requirements for launch scripts to README by @copybara-service[bot] in https://github.com/google/tunix/pull/424
• Enable generic creation of optax.optimizer, optax.learning_rate_schedule from cli by @copybara-service[bot] in https://github.com/google/tunix/pull/435
• [Tunix] Reshard prompt and completion tokens to the REFERENCE mesh before computing reference log probabilities if needed. This is needed when rollout and reference are distributed. by @copybara-service[bot] in https://github.com/google/tunix/pull/451
• Add Typed... types to ArrayLike. by @copybara-service[bot] in https://github.com/google/tunix/pull/461
• Downgrad the Jax/libtpu version to resolve performance issue on Cloud TPU by @wang2yn84 in https://github.com/google/tunix/pull/465
• Pin Jax version to 0.7.1 by @wang2yn84 in https://github.com/google/tunix/pull/468
• add a comment for version pinning by @copybara-service[bot] in https://github.com/google/tunix/pull/469
• update internal grpo notebook by @copybara-service[bot] in https://github.com/google/tunix/pull/463
• Adds experimental pre-shard to Pathways on Cloud experimental_reshard by @copybara-service[bot] in https://github.com/google/tunix/pull/473
• Relax the jax version requirement to get a working Kaggle image. by @wang2yn84 in https://github.com/google/tunix/pull/474
• add remat options to model by @copybara-service[bot] in https://github.com/google/tunix/pull/470
• Create a new list container type to support both flax<=0.11.2 and >=0.12.0. by @copybara-service[bot] in https://github.com/google/tunix/pull/476
• Enable setting default param dtype for Gemma 3 model by @copybara-service[bot] in https://github.com/google/tunix/pull/482
• new reward functions and unit tests by @copybara-service[bot] in https://github.com/google/tunix/pull/472
• add readthedoc in readme by @copybara-service[bot] in https://github.com/google/tunix/pull/485
• Refactor: Optimize masked index lookup using jnp.where by @copybara-service[bot] in https://github.com/google/tunix/pull/490
• Fix the slow kaggle image issue. by @copybara-service[bot] in https://github.com/google/tunix/pull/488
• Fix reward function bug by @copybara-service[bot] in https://github.com/google/tunix/pull/486
• Fix the colab that can't find the final checkpoint. Cleanup the intermediate checkpoint directory. by @wang2yn84 in https://github.com/google/tunix/pull/495
• Fix type errors within agents/ by @copybara-service[bot] in https://github.com/google/tunix/pull/496
• Update weight sync logic to handle KV head duplication by @wenxindongwork in https://github.com/google/tunix/pull/464
• Added grpo test to tpu-tunix tests by @mydatascience in https://github.com/google/tunix/pull/447
• Add Gemma3-270M model configuration support [#500] by @chethanuk in https://github.com/google/tunix/pull/501
• Disable setting specific profiler options on Pathways backend. by @copybara-service[bot] in https://github.com/google/tunix/pull/494
• Adds description to external notebook for steps that take long. by @wang2yn84 in https://github.com/google/tunix/pull/481

New Contributors * @chethanuk made their first contribution in https://github.com/google/tunix/pull/501

Full Changelog: https://github.com/google/tunix/compare/v0.1.0...v0.1.1