2026-Google-IOExtended-Tunix

Transcript

1. Google Developer Group Bangkok Tunix: Tune in JAX Aye Hninn Khine Lecturer Department of Computer Engineering, KMUTT Google Developer Expert (AI/Cloud) 1
2. LLM Pretraining Pre-training gets you a strong text completer. (Just next word prediction) Bangkok is the capital of ???? Pre-training (GPT-3) Supervised Fine Tuning (chatGPT) Web text (trillions of tokens) Demonstrations (10k–100k examples) RLHF DPO (chatGPT) Preference pairs or reward signal 2
3. Why do we need post-training? Pretrained models have knowledge but no behavior The model has the knowledge but needs the right prompt format to surface it. It can’t follow instructions, refuse harmful requests, or know when to stop. Post-training bridges this gap: • SFT teaches format, instruction-following, and when to stop • RLHF/DPO optimizes for what humans actually prefer 3
4. W/O Fine Tuning (Raw LLM) Pretraining W/ Fine Tuning (PostTraining) on Question & Answering Dataset Prompt: What is the capital of Thailand? Prompt: What is the capital of Thailand? Generated Answer: What is the capital of Myanmar? What is the capital of Malaysia? Editable Location Generated Answer: Bangkok 4
5. LLM Post-training Ever wonder how advanced LLM models are crafted? With Tunix, the powerful training stack using JAX + XLA becomes accessible to you. • Post-training is crucial to adapt models to your proprietary use case. • Building efficient post-training pipeline is notoriously challenging. • Bring the Google way of LLM post-training to the public. • Democratize the model adaptation workflows with the cutting edge technologies. • A full set of open-source libraries with production level robustness and efficiency. 5
6. JAX ? https://github.com/google/jax https://github.com/google/jax/blob/master/README.md JAX is Autograd and XLA, brought together for high-performance machine learning research. Dig a little deeper, and you'll see that JAX is really an extensible system for composable function transformations. 6
7. JAX is Numpy import numpy as np import jax.numpy as jnp x = np.arange(9_000_000, x = jnp.arange(9_000_000, dtype=np.float32) x = x.reshape([3000, 3000]) x dtype=jnp.float32) x = x.reshape([3000, 3000]) �� x array([[0.000000e+00, ...], ..., DeviceArray([[0.000000e+00, ...], ..., [..., 8.999999e+06]]) [..., 8.999999e+06]]) %%timeit %%timeit (x @ x) (x @ x).block_until_ready() 1 loop, best of 5: 428 ms per loop 1 loop, best of 5: 11.9 ms per loop 7
8. https://www.tensorflow.org/xla aXelerated Linear Algebra Compiles computation for CPU, GPU, TPU. For example: fuses expressions like sum([x + y * z for x, y, z in zip(*[xs, ys, zs])]) 8
9. WITHOUT TUNIX WITH TUNIX ❌ difficult ❌ distributed training ❌ many frameworks ❌ RLHF infrastructure ✅ one framework ✅ JAX ✅ SFT ✅ RLHF ✅ DPO Editable Location 9
10. Tunix Architecture 10
11. TPU on Google Cloud - High performance - Open source - Cutting-edge model architectures - Continuously tested for performance - Continuously tested for accuracy - Get up and running quickly, modify as needed - Train and run on your own data https://cloud.google.com/tpu/docs/ 11
12. Supervised Fine Tuning/Instruction Tuning Learning on demonstration data Given pairs (x, y) ⇒ x is prompt, y is desired response Teach the model to follow the instruction 12
13. WangChanThaiInstruct 13
14. Low Rank Adaptation (LoRA)* Why LoRA? 1 2 Makes training cheaper while preserving most of the accuracy of full fine-tuning. Need to save all weights for every task the model is fine-tuned on * LoRA: Low Rank Adaptation of Large Language Models [arXiv, 2021] Fig.: Increasing sizes of LLMs with time 14
15. Training Original Equation Y = Wx + b LoRA Equation Y = Wx + b + (α/r)BAx W (d,d) B (d,r) = trainable = frozen A (r,d) Fig.: LoRA matrices rank = r < < d 15
16. Inference No additional latency during inference! LoRA Equation Y = Wx + b + BAx = Wx + BAx + b = (W + BA)x + b = Wupdatedx + b Note: Scaling factor α/r removed for brevity 16
17. Limitations with SFT - Data annotation is painful. Expensive human annotations Requires special experts knowledge Limited amount of annotated data for specific domains 17
18. Learning from preference data (RLHF) Daniel Jurafsky and James H. Martin. 2026. Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition with Language Models, 3rd edition. Ouyang, Long, et al. "Training Language Models to Follow Instructions with Human Feedback." Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 27730-27744. 18
19. Reward Model Learning (Bradley-Terry Model) Oi = winning response Oj = losing response X = prompt P(Oi > Oj|X) = 1 and vice versa P(Oj > Oi|X) = 0 19
20. Daniel Jurafsky and James H. Martin. 2026. Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition with Language Models, 3rd edition. 20
22. Tunix Models 22
23. SFT with Tunix* # Load base model model_config = gemma_lib.ModelConfig.gemma3_270m() model = params_safetensors_lib.create_model_from_safe_tensors( MODEL_CP_PATH, (model_config), mesh ) * https://github.com/google/tunix/blob/main/examples/qlora_gemma.ipynb 23
24. SFT with Tunix # Load base model model_config = gemma_lib.ModelConfig.gemma3_270m() model = params_safetensors_lib.create_model_from_safe_tensors( MODEL_CP_PATH, (model_config), mesh ) # LoRA model_input = model.get_model_input() lora_provider = qwix.LoraProvider( module_path=".*q_einsum|.*kv_einsum|.*gate_proj|.*down_proj|.*up_proj", rank=RANK, alpha=ALPHA, ) lora_model = qwix.apply_lora_to_model(model, lora_provider, **model_input) Note: We use the base Gemma 270M model here (not instruction-tuned). 24
25. SFT with Tunix # Load base model model_config = gemma_lib.ModelConfig.gemma3_270m() model = params_safetensors_lib.create_model_from_safe_tensors( MODEL_CP_PATH, (model_config), ... ) # LoRA ... lora_model = qwix.apply_lora_to_model(model, lora_provider, **model_input) # Tokenizer tokenizer = tokenizer_lib.Tokenizer(tokenizer_path=GEMMA_TOKENIZER_PATH) 25
26. SFT with Tunix # Load dataset train_ds, validation_ds = data_lib.create_datasets( dataset_name='mtnt/en-fr', global_batch_size=BATCH_SIZE, max_target_length=MAX_TARGET_LENGTH, num_train_epochs=NUM_EPOCHS, tokenizer=tokenizer, ) 26
27. SFT with Tunix # Load dataset train_ds, validation_ds = ... def gen_model_input_fn(x: peft_trainer.TrainingInput): pad_mask = x.input_tokens != tokenizer.pad_id() positions = utils.build_positions_from_mask(pad_mask) attention_mask = utils.make_causal_attn_mask(pad_mask) return { 'input_tokens': x.input_tokens, 'input_mask': x.input_mask, 'positions': positions, 'attention_mask': attention_mask, } 27
28. SFT with Tunix training_config = peft_trainer.TrainingConfig( eval_every_n_steps=EVAL_EVERY_N_STEPS, max_steps=MAX_STEPS, checkpoint_root_directory=FULL_CKPT_DIR, ) 28
29. SFT with Tunix training_config = peft_trainer.TrainingConfig( eval_every_n_steps=EVAL_EVERY_N_STEPS, max_steps=MAX_STEPS, checkpoint_root_directory=FULL_CKPT_DIR, ) trainer = peft_trainer.PeftTrainer( base_model, optax.adamw(1e-5), training_config ).with_gen_model_input_fn(gen_model_input_fn) 29
30. SFT with Tunix training_config = peft_trainer.TrainingConfig( eval_every_n_steps=EVAL_EVERY_N_STEPS, max_steps=MAX_STEPS, checkpoint_root_directory=FULL_CKPT_DIR, ) trainer = peft_trainer.PeftTrainer( base_model, optax.adamw(1e-5), training_config ).with_gen_model_input_fn(gen_model_input_fn) trainer.train(train_ds, validation_ds) 30
31. Representative RL algorithms GRPO: PPO: https://arxiv.org/pdf/2402 .03300 https://arxiv.org/pdf/1707 .06347 • • • • Critic-Free Optimization: significantly reducing memory usage and computational overhead. Group-Relative Advantage: Comparing an output's reward relative to the mean/variance of rewards within a *group*. KL-Regularized Updates: KL divergence penalty directly into the objective function for stable training. Applications: ◦ Math and Complex Reasoning ◦ AI Alignment • Actor-Critic Framework: Critic model to estimate the state-value function. • Clipped Surrogate Objective: Prevents large, unstable policy changes. • Applications: ◦ ◦ ◦ Robotics Gaming RLHF 31
32. Useful resources https://github.com/google/tunix/tree/main/examples GET THE SLIDE https://tunix.readthedocs.io/en/latest/ https://www.kaggle.com/competitions/google-tunix-hackath on https://medium.com/google-developer-experts/mar-apr-2026 -ai-community-activity-highlights-and-achievements-154ff3 a85db3 https://cloud.google.com/products/tpu/tpu-developer https://github.com/ayehninnkhine/TPU_Sprint_2026 32
33. Thank you, Feel free to connect with me via LinkedIn 33