一句话：提出调度语言 Slapo，将 PyTorch 模型的执行与定义解耦，通过高层调度原语渐进式优化大模型训练。
大规模深度学习模型训练面临可用性与效率的权衡：动态图框架易用但性能次优，静态图或定制分布式方案则牺牲灵活性。Slapo 受深度学习编译器将算子算术定义与平台相关优化解耦的启发，提出一种调度语言，在 PyTorch 模型上将执行与定义分离。它提供一组调度原语，可按需渐进地改造模型，支持高性能 kernel、有效 3D 并行和高效 activation checkpointing 等优化，从而在较大程度上保留可编程性与可调试性。评估将已有手工优化系统化地调度，相比 DeepSpeed 和 Megatron-LM 的开箱性能，在单机 8 张 V100 上吞吐最高提升 2.92×，在多机最多 64 GPU 上最高提升 1.41×。
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.92× on a single machine with 8 NVIDIA V100 GPUs, and by up to 1.41× on multiple machines with up to 64 GPUs, when compared to the out-of-the-box performance of DeepSpeed and Megatron-LM.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@inproceedings{ChenYZZZ024,
title = {Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training},
author = {Hongzheng Chen and Cody Hao Yu and Shuai Zheng and Zhen Zhang and Zhiru Zhang and Yida Wang},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620665.3640399}
}