一句话：提出Astra，利用深度学习训练的重复可预测性在线探索优化空间，动态定制执行，训练性能最高提升3倍。
深度学习训练作业的执行优化通常将计算视为通用数据流图，未充分利用领域知识，难以自动逼近手工优化性能。Astra 利用训练作业的重复性与可预测性，在训练推进的同时以 work-conserving 方式在线探索优化状态空间；通过轻量 profiling、profile 数据索引以及状态空间剪枝，为每个作业和硬件动态定制端到端基础设施，同时保持编译器简单可维护。作者在 PyTorch 与 TensorFlow 中实现了 Astra，并在先进深度学习模型上评估。结果表明，端到端训练性能最高提升 3 倍，接近 cuDNN 等手工优化实现，并在性能和鲁棒性上显著优于 TensorFlow XLA 等静态编译框架。
We present Astra, a compilation and execution framework that optimizes execution of a deep learning training job. Instead of treating the computation as a generic data flow graph, Astra exploits domain knowledge about deep learning to adopt a custom approach to compiler optimization. The key insight in Astra is to exploit the unique repetitiveness and predictability of a deep learning job, to perform online exploration of the optimization state space in a work-conserving manner while making progress on the training job. This dynamic state space exploration in Astra uses lightweight profiling and indexing of profile data, coupled with several techniques to prune the exploration state space. Effectively, the execution layer custom-wires the infrastructure end-to-end for each job and hardware, while keeping the compiler simple and maintainable. We have implemented Astra in two popular deep learning frameworks, PyTorch and Tensorflow. On state-of-the-art deep learning models, we show that Astra improves end-to-end performance of deep learning training by up to 3x, while approaching the performance of hand-optimized implementations such as cuDNN where available. Astra also significantly outperforms static compilation frameworks such as Tensorflow XLA both in performance and robustness.
DOI 原文 ·
@inproceedings{paperbot4941,
title = {Astra},
author = {Muthian Sivathanu and Tapan Chugh and Sanjay S. Singapuram and Lidong Zhou},
booktitle = {ASPLOS},
year = {2019},
doi = {10.1145/3297858.3304072}
}