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
- Training Models Too Big for One GPU | ZeRO & FSDPyoutube.com
- Async RL in Pure JAXdivyamakkar0.github.io
- Big Boss (@0xBADB01E) on Xx.com
- How To Scale Your Modeljax-ml.github.io
- How trillion-parameter AI models work and how to serve themyoutube.com
- How GPT, Claude, and Gemini are actually trained and served – Reiner Popeyoutube.com
- Transformers Laid Outgoyalpramod.github.io
- How to Think About TPUs | How To Scale Your Modeljax-ml.github.io
- How to Think About GPUs | How To Scale Your Modeljax-ml.github.io
- All the Transformer Math You Need to Know | How To Scale Your Modeljax-ml.github.io
- 5D parallelism in LLM training - gdymind's Bloggdymind.com
- Transformer Math 101 | EleutherAI Blogblog.eleuther.ai