Easy FunctionGemma finetuning with Tunix on Google TPUs

This post demonstrates finetuning FunctionGemma with Tunix, a JAX-based LLM post-training library, on Google TPUs.

• •Tunix supports SFT, PEFT, preference tuning, RL, and distillation for Gemma, Qwen, and LLaMA models
• •LoRA is applied via Qwix to attention layers, running on free-tier Colab TPU v5e-1
• •Completion-only loss is handled via a custom dataset class using the Mobile Action dataset
• •One epoch of training yields a significant accuracy boost with high TPU utilization
• •The merged LoRA model is exported as safetensors for on-device deployment via LiteRT