Tunix Brings JAX-Based Fine-Tuning to FunctionGemma on Google TPUs

Google

Mar 24, 2026 · Updated Jun 5, 2026

Google published a guide for fine-tuning FunctionGemma using Tunix, a JAX library for LLM post-training on TPUs. The tutorial runs on free Colab TPU compute, covering the workflow from loading weights to exporting a LoRA-adapted model for edge deployment.

Google published a developer guide showing how to fine-tune FunctionGemma, its small language model that translates natural language into actionable API calls for edge devices. The guide uses Tunix, a JAX-based library designed for LLM post-training that supports supervised fine-tuning, LoRA, preference tuning, reinforcement learning, and model distillation across Gemma, Qwen, and Llama models.

The tutorial runs entirely on free-tier Colab TPU v5e-1, applying LoRA adapters via Qwix and demonstrating a significant accuracy improvement after a single training epoch with high TPU utilization. The fine-tuned model exports to safetensors for on-device deployment through LiteRT, making this a complete pipeline from training to edge.

If you're building function-calling agents for mobile or edge devices, the accompanying Colab notebook lets you run the full fine-tuning workflow without any GPU or TPU costs.

View the full update on developers.googleblog.com

Google AI Developers

@googleaidevsMar 23

Learn how to fine-tune FunctionGemma on TPUs in a @GoogleColab notebook. This guide uses Tunix, a lightweight JAX library, to streamline post-training your LLM🧵👇 https://t.co/QbfDFvSyd9

View on X

Every HeadsUpAI update is written based on its original source and reviewed before it's published. Read our editorial standards →

See all AI news & updates from Google →

Keep reading

Google Releases Gemma 4 QAT Checkpoints for Efficient On-Device AI

Google released new Gemma 4 Quantization-Aware Training (QAT) checkpoints, including GGUF (Q4_0) and a custom mobile schema under 1GB. These enable running Gemma 4 models locally on consumer GPUs and mobile devices with reduced memory footprint and accelerated decode speeds, while preserving reasoning quality.

Ollama Adds Google DeepMind's Gemma 4 12B for Local Agentic AI

OllamaJun 7

Ollama Adds Google DeepMind's Gemma 4 12B for Local Agentic AI

Ollama has made Google DeepMind's Gemma 4 12B model available for local execution, including support for chat and agentic applications. This expands access to a powerful, open-weight multimodal model optimized for on-device reasoning and coding, enabling private and offline AI workflows on consumer hardware.

Fireworks AI Adds Gemma 4 Training to Build Custom Reasoning Agents

Fireworks AIApr 28

Fireworks AI Adds Gemma 4 Training to Build Custom Reasoning Agents

Fireworks AI integrated Google's Gemma 4 models into its training platform, enabling full-parameter fine-tuning and DPO with a 256K context window. This allows teams to build specialized reasoning agents on a unified stack that transitions from training to production inference in seconds.

Google Gemma releases gemma-skills to accelerate agentic workflows with multi-token prediction

Google GemmaJun 2

Google Gemma releases gemma-skills to accelerate agentic workflows with multi-token prediction

Google launched the first iteration of gemma-skills, an open-source library of reusable capabilities for AI agents. By standardizing how agents select model sizes and use performance optimizations like MTP, Google is making it easier to build efficient, autonomous workflows on top of the Gemma ecosystem.