RAF：面向深度学习模型训练的整体编译
机器学习
2023-03-09 v1 分布式、并行与集群计算
编程语言
摘要
随着深度学习在现代应用中无处不在，许多深度学习框架被提出，以帮助深度学习从业者快速开发与训练 DNN 模型。同时，随着训练大型深度学习模型近年成为趋势，训练吞吐量与内存占用变得至关重要。相应地，以编译器优化来优化训练工作负载不可避免且愈发受到关注。然而，现有深度学习编译器（DLCs）主要面向推理，并未在训练工作负载中纳入如自动微分与自动混合精度等整体优化。本文中，我们提出 RAF，一种面向训练的深度学习编译器。与现有 DLCs 不同，RAF 接受前向模型并在内部生成训练图。据此，RAF 能够系统性地整合面向性能、内存与分布式训练的图优化。此外，为赶上手工定制核库与张量编译器的最先进性能，RAF 提出一种算子方言机制，以无缝集成所有可能的核实现。我们证明，通过内部训练图生成与算子方言机制，我们能够执行整体优化，并在 GPU 上对流行 transformer 模型取得优于 PyTorch（eager 与 torchscript 模式）、XLA 与 DeepSpeed 的训练吞吐量或更大批大小。
引用
@article{arxiv.2303.04759,
title = {RAF: Holistic Compilation for Deep Learning Model Training},
author = {Cody Hao Yu and Haozheng Fan and Guangtai Huang and Zhen Jia and Yizhi Liu and Jie Wang and Zach Zheng and Yuan Zhou and Haichen Shen and Junru Shao and Mu Li and Yida Wang},
journal= {arXiv preprint arXiv:2303.04759},
year = {2023}
}