# TAPAS: Fast and Automatic Derivation of Tensor Parallel Strategies for Large Neural Networks

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 58 references that are in the map.
- First posted on arXiv 1 Feb 2023.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [DAPPLE: A Pipelined Data Parallel Approach for Training Large Models](https://mlsanity.com/paper/2007.01045/) (2020)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Go Wider Instead of Deeper](https://mlsanity.com/paper/2107.11817/) (2021)
- [Supporting Very Large Models using Automatic Dataflow Graph Partitioning](https://mlsanity.com/paper/1807.08887/) (2018)
- [HET: Scaling out Huge Embedding Model Training via Cache-enabled Distributed Framework](https://mlsanity.com/paper/2112.07221/) (2021)
- [Enabling Compute-Communication Overlap in Distributed Deep Learning Training Platforms](https://mlsanity.com/paper/2007.00156/) (2020)

## Built on by (most cited first, 3 of 3)

- [Guard: Scalable Straggler Detection and Node Health Management for Large-Scale Training](https://mlsanity.com/paper/2605.17879/) (2026)
- [ROAM: memory-efficient large DNN training via optimized operator ordering and memory layout](https://mlsanity.com/paper/2310.19295/) (2023)
- [Learning to Shard: RL for Co-optimizing the Parallelism Degrees and Per-operator Sharding Dimensions in Distributed LLM Inference](https://mlsanity.com/paper/2509.00217/) (2025)

## Abstract (arXiv)

Tensor parallelism is an essential technique for distributed training of large neural networks. However, automatically determining an optimal tensor parallel strategy is challenging due to the gigantic search space, which grows exponentially with model size and tensor dimension. This prohibits the adoption of auto-parallel systems on larger models. We observe that neural networks usually contain repeated substructures, and build an automatic parallelism framework named TAPAS that eliminates redundant search efforts. TAPAS employs a divide-and-conquer approach that efficiently folds the search space by identifying those unique substructures. As a result, it runs at sub-linear complexity concerning the model size, making it a scalable solution for training large-scale networks. Our evaluations demonstrate that TAPAS outperforms the state-of-the-art automatic parallelism frameworks by up to 160× in search speed on a wide range of models, and the performance of derived strategies is competitive or even better compared with the expert-engineered Megatron-LM library.