Flax NNX is a neural-network API for JAX developed by the Flax team in collaboration with the JAX team. Released in 2024, it uses Python reference semantics so models can be expressed as regular Python objects with reference sharing and mutability, supporting inspection, debugging, and analysis. Its API includes layers and modules such as Linear, convolution, normalization, attention, recurrent cells, and dropout, along with patterns for replicated training, serialization, and checkpointing. In distributed training, it can manage model state while JAX handles computation across a device mesh.
JAX is a Python library for accelerator-oriented array computation and program transformation, designed for numerical computing and large-scale machine learning. It transforms Python and NumPy-style functions through automatic differentiation with `jax.grad`, automatic batching and vectorization with `jax.vmap`, and just-in-time compilation with `jax.jit`. It uses XLA to compile computations for GPUs, TPUs, and other accelerators. JAX supports compiler-driven multi-device execution: `Mesh`, `PartitionSpec`, and `NamedSharding` declare array placement across device meshes, allowing distributed training code to include automatically managed parallel work such as gradient synchronization. Its compilation tools also include ahead-of-time compilation, portable artifact export with `jax.export`, and conversion to TensorFlow SavedModels through `jax2tf` for serving integrations. The project describes itself as a research project rather than an official Google product.
A JAX debugging API function that visualizes how an array is sharded across available devices, including which portions are partitioned and which values are replicated. It is part of the open-source JAX array-computing library.
Orbax is a Google-developed collection of checkpointing and persistence utilities for JAX models and training state. Its checkpointing API saves and restores pytrees such as model weights and optimizer state, with support for asynchronous checkpointing, standard and custom types, and flexible storage formats. It is distributed through domain-specific PyPI packages including orbax-checkpoint and orbax-export, and is used by JAX-based machine-learning frameworks and model implementations.
Searchable transcript of Scale JAX models to multi-GPU systems — Google Cloud Tech (03:28). Search for a phrase, then click its timestamp to jump straight to that moment in the video.
Captions sourced from the original video on YouTube, published by Google Cloud Tech. The video, its captions and all related intellectual property remain the property of their respective owners; AINotes claims no ownership. Provided for research, accessibility and search — see the Transcript Notice and Copyright Policy.
00:00 Before we show you how to work with Transformers, allow us a quick detour. Real-life AI workloads are usually run on multiple GPUs. Running JAX on multiple GPUs is not about rewriting a training step with communication code. It's more about putting arrays in the right places, so the code stays familiar and only the placement changes. >> In the previous video, we show you how to train MLP when one GPU owns the batch, the gradients, and the update.
00:32 With data parallelism, each GPU gets a slice of the batch. Then, the gradients are averaged, so every model copy stays in sync. Data parallelism is the simplest way to scale your training. There are many other advanced techniques, which we will not cover here today. >> When you train with multiple GPUs, there are four important concepts you need to keep in mind.
00:58 Mesh indicates the devices. Partition spec describes the split. Name sharding combines the plan, and device put applies it. >> So, what happened here is we shard the batch, replicate the parameters, and call the same train step. JAX sees the placement and inserts the parallel work, including gradient averaging. >> Always verify placement. JAX debug visualizer array sharding shows where the batch is split [music] and where the weights are replicated.
01:34 >> Also, throughput should improve when there is enough work per device, but small workloads can lose communication overhead. So, scaling might not be free. >> If you need explicit control, shard map lets each device see its local shard, >> [music] >> and JAX like PMean makes the gradient average explicit. Use automatic sharding first for experimenting, but for debugging and fine-grained control, use shard map.
02:02 >> Now, we combine the pieces together into our actual language model, from attention to what we learned with MLP. The tiny transformer reuses the training loop, the causal attention, fixed shapes, throughput thinking, and the multi-GPU placement. The data set we use in the notebook is tiny Shakespeare with byte-level tokens. >> Each block is pre-norm, layer norm, causal self-attention, residual.
02:30 Then layer norm, feed forward, residual. The attention function plugs into an NNX multi-head attention. This is the causal attention from the previous video. >> So, the model predict the next byte, compare logits, optimize cross-entropy, and track perplexity. Same training story, new model shape. >> For an NNX, replicate model and optimize the state with NNX state, [music] Jax device put, and NNX update.
02:59 Shard the data and keep the NNX GT training step clean. >> And at the end, you save your weights with orbax, restore into a fresh model, verify logits match, and generate text. During generation, pad to a fixed max length, so the compiled forward path does not recompile every token. And that's it. The model is able to generate Shakespeare's poetry.