# Mesh-TensorFlow: Deep Learning for Supercomputers

arXiv: https://arxiv.org/abs/1811.02084 · HTML: https://mlsanity.com/paper/1811.02084/

## Key facts

- Cited by 169 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 5 Nov 2018.

## Built on by (most cited first, 30 of 169)

- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale](https://mlsanity.com/paper/2208.07339/) (2022)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [The Falcon Series of Open Language Models](https://mlsanity.com/paper/2311.16867/) (2023)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [Fast3R: Towards 3D Reconstruction of 1000+ Images in One Forward Pass](https://mlsanity.com/paper/2501.13928/) (2025)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Pre-Trained Models: Past, Present and Future](https://mlsanity.com/paper/2106.07139/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [Accelerating Sparse Deep Neural Networks](https://mlsanity.com/paper/2104.08378/) (2021)
- [Scaling Up Models and Data with textttt5x and textttseqio](https://mlsanity.com/paper/2203.17189/) (2022)
- [A New Generation of Perspective API: Efficient Multilingual Character-level Transformers](https://mlsanity.com/paper/2202.11176/) (2022)
- [Combined Scaling for Zero-shot Transfer Learning](https://mlsanity.com/paper/2111.10050/) (2021)
- [Scaling Laws for Fine-Grained Mixture of Experts](https://mlsanity.com/paper/2402.07871/) (2024)
- [ExT5: Towards Extreme Multi-Task Scaling for Transfer Learning](https://mlsanity.com/paper/2111.10952/) (2021)
- [Charformer: Fast Character Transformers via Gradient-based Subword Tokenization](https://mlsanity.com/paper/2106.12672/) (2021)
- [Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning](https://mlsanity.com/paper/2309.05444/) (2023)
- [A Review of Sparse Expert Models in Deep Learning](https://mlsanity.com/paper/2209.01667/) (2022)
- [PanGu-α: Large-scale Autoregressive Pretrained Chinese Language Models with Auto-parallel Computation](https://mlsanity.com/paper/2104.12369/) (2021)
- [Scaling Up Influence Functions](https://mlsanity.com/paper/2112.03052/) (2021)
- [FastMoE: A Fast Mixture-of-Expert Training System](https://mlsanity.com/paper/2103.13262/) (2021)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Accelerating Distributed MoE Training and Inference with Lina](https://mlsanity.com/paper/2210.17223/) (2022)
- [Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers](https://mlsanity.com/paper/2109.10686/) (2021)
- [Scaling Laws vs Model Architectures: How does Inductive Bias Influence Scaling?](https://mlsanity.com/paper/2207.10551/) (2022)

## Abstract (arXiv)

Batch-splitting (data-parallelism) is the dominant distributed Deep Neural Network (DNN) training strategy, due to its universal applicability and its amenability to Single-Program-Multiple-Data (SPMD) programming. However, batch-splitting suffers from problems including the inability to train very large models (due to memory constraints), high latency, and inefficiency at small batch sizes. All of these can be solved by more general distribution strategies (model-parallelism). Unfortunately, efficient model-parallel algorithms tend to be complicated to discover, describe, and to implement, particularly on large clusters. We introduce Mesh-TensorFlow, a language for specifying a general class of distributed tensor computations. Where data-parallelism can be viewed as splitting tensors and operations along the "batch" dimension, in Mesh-TensorFlow, the user can specify any tensor-dimensions to be split across any dimensions of a multi-dimensional mesh of processors. A Mesh-TensorFlow graph compiles into a SPMD program consisting of parallel operations coupled with collective communication primitives such as Allreduce. We use Mesh-TensorFlow to implement an efficient data-parallel, model-parallel version of the Transformer sequence-to-sequence model. Using TPU meshes of up to 512 cores, we train Transformer models with up to 5 billion parameters, surpassing state of the art results on WMT'14 English-to-French translation task and the one-billion-word language modeling benchmark. Mesh-Tensorflow is available at https://github.com/tensorflow/mesh .