Scale and Serve JAX Models on NVIDIA GPUs

Master the fundamentals of JAX on NVIDIA GPUs and learn how to configure your GKE cluster for high-performance GPU-accelerated computing. In Part 1 of this course, Ivan Nardini, AI Engineer at Google Cloud and Ekaterina Sirazitdinova, Developer Advocate at NVIDIA will teach you how to:

• Verify your NVIDIA GPU visibility instantly using nvidia-smi and JAX device diagnostics.
• Learn how the XLA compiler traces and compiles Python code into optimized machine instructions.
• Eliminate performance bottlenecks caused by host-to-device data transfers and asynchronous dispatch latency.

This lab gets you quickly through setup and into provisioning a real Kubernetes environment with NVIDIA L4 GPUs so you can confirm JAX is talking to your hardware. Once you complete this Codelab, you'll be ready for multi-GPU transformer clusters.

JAX is fast when you are able to write code properly based on an understanding of the compilation model. This lab cracks open the tracing engine so you can see how JAX converts Python code into static XLA graphs, and teaches you how to avoid recompilation traps that can compromise performance.

Profiling is essential to truly understanding performance. Because JAX runs asynchronously, standard Python timing tools do not capture accurate hardware execution metrics. This lab shows you how to capture real hardware profiles and trace your computation all the way down to individual CUDA kernels.

Writing a standard neural network training loop in JAX requires shifting from object-oriented state management to functional programming. In Part 2 of this course, we walk through how to build a complete training loop from scratch, update model parameters using Optax, and tap into hardware-accelerated self-attention kernels.

• Structure Pure-Functional Loops: Implement model states, parameter updates, and evaluation steps without relying on hidden global variables.
• Integrate Optax Optimizers: Configure and apply modern optimizer states (like AdamW) to update parameter PyTrees cleanly.
• Optimize Attention Computations: Harness cuDNN and NVIDIA TransformerEngine to accelerate performance-heavy attention layers.

It’s time to move from basic math to deep learning. In this lab, you’ll build and compile a complete training loop for an MLP on Fashion-MNIST. Along the way, you’ll experience one of JAX’s core ideas firsthand: parameters and optimizer state flow explicitly through the training step, with each step returning a new state instead of modifying it in place.

This lab shows you how to upgrade a standard attention layer into a highly efficient, hardware-optimized implementation by leveraging fused GPU kernels that appropriately manages memory.

Scaling deep learning models across multiple GPUs used to mean rewriting hundreds of lines of complex device communication code. Part 3 of this course breaks down JAX's modern, compiler-driven sharding model, showing you how to distribute workloads automatically across physical device meshes.

• Master Sharding Concepts: Understand how Mesh, PartitionSpec, and NamedSharding declare array layouts across multiple devices.
• Implement Automatic Scaling: Write clean training code that allows the compiler to automatically manage multi-GPU gradient synchronization.
• Incorporate Flax NNX & Orbax: See how to manage state and serialize model checkpoints in a distributed training run.

One GPU is great, but eventually, you need to go bigger. This lab demystifies distributed training by introducing you to JAX's compiler-driven array sharding model, showing you how to scale workloads with almost no modifications to your math.

This is where everything comes together. You will build a complete decoder-only transformer model, scale it across a multi-GPU GKE cluster, train it on real Shakespeare text, and implement proper checkpointing to save your weights.

Deploy your models using robust JAX serving architectures. In Part 4 of this course, learn how to achieve low latency and optimize your setups specifically for web APIs.

• Deploy AOT Compilation: Use Ahead-of-Time compilation to lock down input shapes and guarantee predictable inference latency.
• Export Native Execution Graphs: Package your model code and checkpoints using jax.export for portable runtime deployment.
• Bridge JAX to TensorFlow Serving: Convert JAX graphs to standard TensorFlow SavedModels using jax2tf for corporate server integration.

This final lab focuses on production engineering, showing you how to export your completed transformer out of JAX and prepare it for highly optimized, low-latency web serving.