MaxText vs Tunix: fine-tune Llama 3.1 with JAX on GPUs

This tutorial compares two notebook workflows for fine-tuning Llama 3.1-8B with JAX on NVIDIA GPUs. The speaker demonstrates GPU detection and explains that XLA-compiled training can run on GPU hardware, addressing the assumption that JAX requires TPUs.

The MaxText demo uses full supervised fine-tuning, which trains all model weights. It targets multi-GPU systems, with eight GPUs carrying at least 80 GB each given as an example. Setup includes distributed initialization for multiple machines, Hugging Face authentication for the gated checkpoint, and conversion to MaxText format. The speaker recommends doing that conversion on CPU to save GPU memory. An SFT YAML config launches a 100-step UltraChat run with 1,024 tokens and BFloat16; TensorBoard tracks loss.

The Tunix workflow keeps base weights frozen and adds adapters to attention and MLP layers. The speaker gives single-GPU memory targets of 16 GB or more for QLoRA and 24 GB or more for LoRA. The demo adapts its device mesh to the GPU count, converts the downloaded model to JAX, and distributes it across the mesh. QLoRA uses NF4 4-bit quantization with tile size 32, a small UltraChat subset, and evaluation every 25 steps.

For adapting a local LLM, the speaker recommends LoRA or QLoRA for style, format, and lightweight domain changes. Full SFT is presented as an option for broader behavior changes when the available compute and evaluation budget allow it.