This presentation advocates for JAX as a highly effective framework for large-scale Large Language Model (LLM) training, enabling developers to write simple NumPy-like code and distribute it efficiently across potentially tens of thousands of devices. The speaker, Jacob from DeepMind, emphasizes JAX's XLA compiler, which intelligently handles distribution and optimization, along with powerful profiling tools (XProf/XR profiler) that allow developers to verify performance against hardware roofline models.
The talk demonstrates JAX's capabilities through a "fake transformer" model, illustrating basic jax.jit compilation, data parallelism (automatic all_reduce by XLA), Fully Sharded Data Parallelism (FSDP) for memory saving (with all_gather for weights), and model parallelism (using collective matrix multiplications for communication-computation overlap). It highlights JAX's explicit control over sharding via APIs like jax.device_put, jax.sharding.Mesh, jax.lax.with_sharding_constraint, and the newer Shard Map and Sharding in Types interfaces, allowing developers to intervene when automatic optimization isn't ideal. The discussion also touches upon the co-design of LLM architectures with specific hardware (TPUs vs. GPUs) and the significant impact of low-level optimization on cost-efficiency in large-scale training.
Summarizing Large Scale LLM Training in Jax
0:07 JAX for Large-Scale LLM Training: Jacob, a DeepMind engineer who worked on Gemini, presents JAX as a superior framework for large-scale LLM training, allowing simple NumPy-like code to run on thousands of devices with granular control.
0:40 Core Principle - Profile and Optimize: The recommended development loop involves writing NumPy-like code, distributing it, profiling it using tools like XProf/XR profiler, understanding the theoretical "roofline" performance, and then optimizing bottlenecks.
2:58 JAX Ecosystem (Pathways & Device Mesh): A live demo uses a Colab connected to 64 TPUs (V5E slice) via "Pathways," which abstracts multiple JAX hosts into a single controller. A device mesh (jax.sharding.Mesh) is created to logically arrange the devices.
4:02 Initial Model Execution & Profiling: A "fake transformer" forward pass is executed using jax.jit. Profiling reveals "fusions" (matrix multiplication combined with other ops like ReLU) performing at ~98% FLOPs utilization on a single TPU. By default, without explicit sharding, computation runs only on the first device (TPU0).
8:58 Data Parallelism: The input tokens are sharded across the "data" dimension of the device mesh, enabling data parallelism. This results in a significant speedup (from 90ms to 1.4ms for the forward pass on 64 TPUs).
10:50 Training Step & Automatic All-Reduce: When a simple gradient descent training step is introduced (including loss averaging), XLA automatically inserts all_reduce operations to aggregate gradients across devices. This all_reduce is also analyzed, showing close-to-theoretical interconnect bandwidth utilization (12ms runtime for an 11ms theoretical limit).
15:20 Communication Overhead: Data parallelism can lead to communication bottlenecks (waiting for all_reduce), highlighting the need for overlapping communication with computation or exploring other parallelism strategies.
16:32 Fully Sharded Data Parallelism (FSDP): FSDP is demonstrated by sharding weight matrices along the same axis as data parallelism to save memory. XLA automatically performs "all-gather" of weights just-in-time for matrix multiplications.
18:37 FSDP Performance and Batch Size Impact: While asynchronous all_gather attempts to overlap with computation, it doesn't fully hide communication for small matrix multiplications, leading to exposed communication time. FSDP efficiency heavily depends on batch size: larger batches can hide communication, while smaller ones expose all_gather latency.
21:00 Manual Optimizer State Management: JAX requires explicit management of optimizer state sharding (e.g., Adam's first/second moments) by the user, providing fine-grained control compared to more automatic frameworks.
22:54 JAX vs. PyTorch for Parallelism: JAX's explicit sharding offers more flexibility than PyTorch's FSDP, which might flatten arrays or apply parallelism uniformly, making it challenging to mix strategies or handle small parameters.
23:43 Model Parallelism and Pipelining: JAX supports model parallelism (sharding models across devices) and pipelining (though harder to implement). A 2D device mesh (data and model axes) can be used. XLA attempts "collective matmul" to overlap communication and computation, but manual intervention might be needed for optimal performance.
30:50 Intervening on Automatic Sharding: Developers can intervene with XLA's automatic partitioning using jax.lax.with_sharding_constraint for specific intermediate tensors, or by writing device-local code with Shard Map for fine-grained control over operations like slicing.
36:17 Sharding in Types (New Feature): A newer feature allows explicit sharding types at the JAX level, enabling JAX to complain about ambiguous sharding propagations and allowing the user to resolve them with explicit communication primitives.
38:02 Real-world LLM Architectures: Real LLMs at Google largely use pure JAX, with some custom kernels (e.g., Flash Attention) written in Palace for complex operations. JAX, combined with XLA, abstracts away multi-pod and data center networking for massive scalability.
42:53 Integration with Flax and Optax: Flax and Optax work well with JAX sharding because they generate "PyTrees" (parameter/state structures) that can be easily annotated and sharded.
44:39 Memory Efficiency with Sharding: Model parallelism significantly reduces memory usage per device (e.g., from 1GB to 128MB), allowing for much larger models.
45:46 XLA Flags & Roofline Models: Optimizing JAX performance sometimes requires tuning obscure XLA flags. Understanding hardware-specific "roofline models" (e.g., FLOPs, bandwidth) is crucial to predict and debug performance.
47:03 Hardware-Software Co-design: LLM architectures are often customized to the underlying hardware's characteristics (e.g., TPU's high ICI bandwidth favors FSDP, GPU's NVLink limits often lead to more pipelining).
49:22 Knowledge Sharing & Impact: The speaker co-authored a book to share system-level knowledge, noting that a 3% efficiency improvement on large-scale LLM training can save millions of dollars.
58:17 Pipeline Parallelism Challenges: Pipeline parallelism is identified as the most difficult to implement in JAX cleanly, due to complex schedules (bubble-free pipelines) and gradient updates, often better suited for inference than training.