Slapo：面向大规模深度学习模型训练渐进优化的调度语言
机器学习
2023-12-27 v2
摘要
近年来大规模深度学习（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 倍。
引用
@article{arxiv.2302.08005,
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},
journal= {arXiv preprint arXiv:2302.08005},
year = {2023}
}
备注
Accepted to ASPLOS'24