This presentation by Jacob from DeepMind offers a practical guide to large-scale Large Language Model (LLM) training using JAX. The central argument is that JAX excels at scaling from simple, NumPy-like code to massive, distributed systems with tens of thousands of devices. The session emphasizes a crucial development loop: writing code, profiling performance, analyzing it against hardware "roofline" limits, and then optimizing bottlenecks. Through a live demo on a 64-TPU system, the speaker showcases how JAX and the XLA compiler handle various parallelism strategies, including data parallelism, Fully Sharded Data Parallelism (FSDP), and model (tensor) parallelism. The talk highlights JAX's powerful APIs for explicit control, such as device meshes for organizing hardware, shard_map for writing device-local code, and sharding constraints to guide the compiler when its automatic partitioning (GSPMD) is suboptimal.
A Guide to Large-Scale LLM Training with JAX
00:00:40 The JAX Advantage for Large-Scale Training: JAX allows developers to write simple, NumPy-like code that can be scaled to tens of thousands of devices. It provides granular, explicit control over how computation and data are distributed across hardware.
00:01:25 The Core Development Loop: The speaker outlines a key workflow for performance engineering:
Write simple code.
Run it on your distributed hardware.
Profile the execution using tools like the XR profiler.
Analyze performance against the theoretical hardware limits ("roofline analysis").
Identify and fix the bottlenecks.
00:05:52 Profiling and Roofline Analysis in Practice: In an initial test on a single TPU, the presenter uses the profiler to measure the execution time of a matrix multiplication. By calculating the theoretical FLOPs required and comparing it to the TPU's specs, they verify that the XLA compiler is achieving ~98% of the theoretical peak performance, confirming it is "doing the right thing."
00:09:01 Implementing Data Parallelism: By creating a device mesh and using device_put to shard the input batch across 64 TPUs, the forward pass of the model becomes significantly faster.
00:12:23 Identifying Communication Bottlenecks: When a training step (including a backward pass) is introduced with data parallelism, the profiler reveals that the all_reduce operation, needed to synchronize gradients, becomes a major bottleneck, as computation must wait for communication to complete.
00:16:43 Using FSDP for Memory Savings: The presenter demonstrates Fully Sharded Data Parallelism (FSDP) by sharding the model's weights across devices. This saves a significant amount of memory per device but introduces all_gather operations to reconstruct the full weights for each computation. This trade-off can become a new performance bottleneck if communication cannot be hidden by computation.
00:20:57 Explicit Control Over Optimizer State: JAX requires manual management of optimizer states (e.g., moments in the Adam optimizer). These are treated as regular parameters, and the user must explicitly define how they are sharded, enabling strategies like ZeRO-3 optimization.
00:27:02 Model Parallelism and Collective Operations: The demo shows how to implement model (tensor) parallelism by sharding weights along a model dimension. XLA attempts to generate an efficient "collective matmul" that overlaps communication with computation. However, for smaller models, communication latency can still be exposed.
00:33:18 Taking Full Control with shard_map: For complex operations where the XLA compiler's automatic partitioning (GSPMD) might make poor decisions, JAX offers shard_map. This powerful API allows you to write code from the perspective of a single device, enabling explicit, device-local computations and manual control over communication primitives.
00:54:30 Sharding in Types for Predictability: A newer JAX feature, "sharding in types," makes sharding information part of the tensor's type at the JAX level, rather than relying on compiler inference. This makes behavior more predictable and forces the user to resolve sharding ambiguities explicitly.