arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2302.08005cs.LG
收藏

Slapo:一种用于大型深度学习模型训练渐进式优化的调度语言

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

Cornell University(康奈尔大学);Boson AI, Inc(玻色人工智能公司);Amazon Web Services(亚马逊云服务)

Hongzheng Chen, Cody Hao Yu, Shuai Zheng, Zhen Zhang, Zhiru Zhang, Yida Wang

AI总结:

提出调度语言Slapo,将PyTorch模型执行与定义解耦,通过高级原语渐进式优化训练,提升吞吐量达2.92倍(单机)和1.41倍(多机)。

摘要:

近年来,大型深度学习(DL)模型的开发日益增多,这使得训练效率变得至关重要。常见的做法在可用性与性能之间难以权衡。一方面,诸如PyTorch之类的DL框架使用动态图来方便模型开发者,但代价是模型训练性能欠佳。另一方面,实践者提出了各种方法来提高训练效率,但牺牲了一定的灵活性,从使图静态化以进行更彻底的优化(例如XLA),到针对大规模分布式训练定制优化(例如DeepSpeed和Megatron-LM)。在本文中,我们旨在通过关注点分离来解决可用性与训练效率之间的紧张关系。受DL编译器将张量级运算符的平台特定优化与其算术定义解耦的启发,本文提出了一种调度语言Slapo,将模型执行与定义解耦。具体而言,Slapo作用于PyTorch模型,并使用一组调度原语将模型转换为常见的模型训练优化,例如高性能内核、有效的3D并行和高效的激活检查点。与现有优化解决方案相比,Slapo通过高级原语按需渐进式优化模型,从而在很大程度上保留了用户的可编程性和可调试性。我们的评估结果表明,通过使用Slapo以系统化的方式调度现有手工优化的优化,与DeepSpeed和Megatron-LM的开箱即用性能相比,我们在配备8个NVIDIA V100 GPU的单台机器上将训练吞吐量提高了最多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.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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑