Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, et al.2019arXiv 2019
Read it on arxiv.org(opens in a new tab)Why this one
Read this after Attention Is All You Need and before ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. The useful idea here is not “use more GPUs,” it is exactly where to cut a Transformer layer so each GPU does real matrix work, then pays the smallest possible communication bill. People often lump tensor parallelism together with data parallelism or pipeline parallelism, but Megatron-LM is the paper that makes the intra-layer split feel mechanical: column split here, row split there, all-reduce only where the math forces it. If you are building or serving large models, this is worth an evening because it turns vague cluster mysticism into shapes, collectives, and bandwidth costs you can reason about.
What to take away
- Tensor parallelism splits individual Transformer matrices across GPUs, not just batches or layers.
- The MLP and attention blocks can be arranged so communication happens at a few predictable boundaries.
- Large model scaling is limited as much by interconnect placement as by raw GPU FLOPs.
Reads with
- Attention Is All You Need
the layer shapes being split across gpus come from the transformer block
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
solves the other scaling bottleneck by sharding optimizer state instead of layer math
- Mixed Precision Training
cuts the bandwidth and memory costs that make tensor parallelism viable