How to Parallelize a Transformer for Training — an explorable explanation
ezyang.github.io · 10,300 words · saved by 1 readers
An interactive TPU/GPU adaptation of the JAX Scaling Book training rooflines: DP, FSDP, TP, mixed sharding, pipeline and expert parallelism.
Model D= F= L= E=·k=·shared= Hardware C= Wici= Batch B= tokens An explorable adaptation of Part 5 of “How to Scale Your Model” by Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh, Albert Webson & Reiner Pope (Google DeepMind). ✦ We begin with the original dense TPU schemes — data parallelism, FSDP, tensor parallelism, their mixed form, and pipelining — then splice in the GPU fabric model and expert parallelism for MoEs. For each, we ask when communication becomes the bottleneck. (This summary is the…
saved by
related reading
- Big Boss (@0xBADB01E) on Xx.com
- Async RL in Pure JAXdivyamakkar0.github.io
- How To Scale Your Modeljax-ml.github.io
- irhum.github.io - Tensor Parallelism with jax.pjitirhum.github.io
- 5D parallelism in LLM training - gdymind's Bloggdymind.com
- Paradigms of Parallelism | Colossal-AIcolossalai.org
- Transformer Math 101 | EleutherAI Blogblog.eleuther.ai
- Pipeline-Parallelism: Distributed Training via Model Partitioningsiboehm.com
- Visualizing 6D Mesh Parallelism · mainmain-horse.github.io
- How to Think About GPUs | How To Scale Your Modeljax-ml.github.io
- Pretraining parallelisms and failed runs — Flashcardsflashcards.dwarkesh.com
- Transformer Inference Arithmetic | kipply's blogkipp.ly