数据移动即所需：Transformer 优化案例研究
机器学习
2021-11-09 v3 机器学习
摘要
Transformer 是当今最重要的机器学习工作负载之一。训练 Transformer 是一项计算密集型任务，常需数天或数周，且已有大量关注投向优化 Transformer。尽管如此，现有实现并未高效利用 GPU。我们发现数据移动是训练时的关键瓶颈。由于阿姆达尔定律以及计算性能的巨幅提升，训练现已变为内存受限。此外，现有框架使用次优的数据布局。利用这些见解，我们提出一种全局优化 Transformer 中数据移动的方案。我们将数据移动最多减少 22.91%，并在训练 BERT 编码器层时整体相较最先进框架实现 1.30 倍性能提升，对整个 BERT 则为 1.19 倍。我们的方法可更广泛地应用于优化深度神经网络，并为如何应对新兴性能瓶颈提供见解。
引用
@article{arxiv.2007.00072,
title = {Data Movement Is All You Need: A Case Study on Optimizing Transformers},
author = {Andrei Ivanov and Nikoli Dryden and Tal Ben-Nun and Shigang Li and Torsten Hoefler},
journal= {arXiv preprint arXiv:2007.00072},
year = {2021}
}
备注
22 pages, 8 figures; MLSys 2021 camera ready