← All transcripts

Scale JAX models to multi-GPU systems Transcript, AI Summary & Key Points

Google Cloud Tech · 2 hours ago · Science & Technology · 03:28 · EN-US

AI Summary

JAX can scale training across multiple GPUs by declaring how arrays are placed rather than rewriting training code with communication logic. Mesh identifies devices, PartitionSpec describes array splits, NamedSharding combines the plan, and device put applies it. JAX then inserts parallel operations such as gradient averaging. Automatic sharding is recommended for experimentation, while shard map provides explicit per-device control for debugging and fine-grained behavior. The same approach is applied to a small causal language model trained on byte-level tiny Shakespeare data, with NNX state management and Orbax checkpointing.

Key Points

  • Multi-GPU JAX training focuses on placing arrays correctly while keeping the training code familiar.
  • With data parallelism, each GPU receives a slice of the batch and gradients are averaged so model copies remain synchronized.
  • Mesh indicates the devices, PartitionSpec describes the split, NamedSharding combines the plan, and device put applies it.
  • JAX uses array placement information to insert parallel work, including gradient averaging.
  • JAX debug visualizer array sharding can show where the batch is split and where weights are replicated.
  • Throughput can improve when there is enough work per device, but small workloads can lose performance to communication overhead.
  • Automatic sharding is useful for experimentation, while shard map provides explicit control and makes local shards visible to each device.
  • The tiny transformer uses pre-norm blocks, layer normalization, causal self-attention, residual connections, and feed-forward layers.

🔒 3 more in the full analysis

AI in practice

Used for

Scale deep learning model training across multiple GPUs
What
Distribute batches and synchronize replicated model copies through gradient averaging.
Automatically insert multi-device parallel operations
What
Use array placement information so JAX manages parallel work without manually rewriting the training step.
Train a causal language model
What
Predict the next byte from tiny Shakespeare data, optimize cross-entropy, track perplexity, and generate text.

🔒 6 more in the full analysis

From this video

4 products

Flax NNX JAX jax.debug.visualize_array_sharding Orbax

Links mentioned

🔒 Full analysis locked

Unlock more videos and the full analysis

Buy credits to process more videos. Each run includes the full analysis, not just the summary — and you get access to the locked analysis across the library.

Inquire for details →

Transcript

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.