# Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 71 references that are in the map.
- First posted on arXiv 16 Feb 2023.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [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)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [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)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [Horovod: fast and easy distributed deep learning in TensorFlow](https://mlsanity.com/paper/1802.05799/) (2018)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [BGL: GPU-Efficient GNN Training by Optimizing Graph Data I/O and Preprocessing](https://mlsanity.com/paper/2112.08541/) (2021)
- [Dynamic Tensor Rematerialization](https://mlsanity.com/paper/2006.09616/) (2020)
- [Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python](https://mlsanity.com/paper/2112.08429/) (2021)

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

- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [Understanding the Potential of FPGA-Based Spatial Acceleration for Large Language Model Inference](https://mlsanity.com/paper/2312.15159/) (2023)
- [HeuriGym: An Agentic Benchmark for LLM-Crafted Heuristics in Combinatorial Optimization](https://mlsanity.com/paper/2506.07972/) (2025)
- [Allo: A Programming Model for Composable Accelerator Design](https://mlsanity.com/paper/2404.04815/) (2024)
- [SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile](https://mlsanity.com/paper/2411.00284/) (2024)
- [Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization](https://mlsanity.com/paper/2503.19050/) (2025)
- [Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References](https://mlsanity.com/paper/2510.14719/) (2025)
- [Dato: A Task-Based Programming Model for Dataflow Accelerators](https://mlsanity.com/paper/2509.06794/) (2025)
- [Magellan: Autonomous Discovery of Novel Compiler Optimization Heuristics with AlphaEvolve](https://mlsanity.com/paper/2601.21096/) (2026)
- [veScale: Consistent and Efficient Tensor Programming with Eager-Mode SPMD](https://mlsanity.com/paper/2509.07003/) (2025)
- [Piper: A Programmable Distributed Training System](https://mlsanity.com/paper/2606.11169/) (2026)

## Abstract (arXiv)

Recent years have seen an increase in the development of large deep learning (DL) models, which makes training efficiency crucial. Common practice is struggling with the trade-off between usability and performance. On one hand, DL frameworks such as PyTorch use dynamic graphs to facilitate model developers at a price of sub-optimal model training performance. On the other hand, practitioners propose various approaches to improving the training efficiency by sacrificing some of the flexibility, ranging from making the graph static for more thorough optimization (e.g., XLA) to customizing optimization towards large-scale distributed training (e.g., DeepSpeed and Megatron-LM). In this paper, we aim to address the tension between usability and training efficiency through separation of concerns. Inspired by DL compilers that decouple the platform-specific optimizations of a tensor-level operator from its arithmetic definition, this paper proposes a schedule language, Slapo, to decouple model execution from definition. Specifically, Slapo works on a PyTorch model and uses a set of schedule primitives to convert the model for common model training optimizations such as high-performance kernels, effective 3D parallelism, and efficient activation checkpointing. Compared to existing optimization solutions, Slapo progressively optimizes the model "as-needed" through high-level primitives, and thus preserving programmability and debuggability for users to a large extent. Our evaluation results show that by scheduling the existing hand-crafted optimizations in a systematic way using Slapo, we are able to improve training throughput by up to 2.92x on a single machine with 8 NVIDIA V100 GPUs, and by up to 1.41x on multiple machines with up to 64 GPUs, when compared to the out-of-the-box performance of DeepSpeed and Megatron-LM.