[论文解读] Data Movement Is All You Need: A Case Study on Optimizing Transformers

Andrei Ivanov, Nikoli Dryden+3 more|arXiv (Cornell University)|Jun 30, 2020

Advanced Neural Network Applications参考文献 96被引用 24

一句话总结

本文 identifies data movement 为训练 Transformer 模型的主要瓶颈，证明尽管计算能力不断提升，但由于 Amdahl's Law 的影响，训练已变为内存受限。通过系统性地分析数据流并进行全局数据布局与移动优化，作者将数据移动减少了高达 22.91%，在 BERT 编码器层上实现 1.30× 的加速，在完整 BERT 模型上实现 1.19× 的加速，优于当前最先进的框架。

ABSTRACT

Transformers are one of the most important machine learning workloads today. Training one is a very compute-intensive task, often taking days or weeks, and significant attention has been given to optimizing transformers. Despite this, existing implementations do not efficiently utilize GPUs. We find that data movement is the key bottleneck when training. Due to Amdahl's Law and massive improvements in compute performance, training has now become memory-bound. Further, existing frameworks use suboptimal data layouts. Using these insights, we present a recipe for globally optimizing data movement in transformers. We reduce data movement by up to 22.91% and overall achieve a 1.30x performance improvement over state-of-the-art frameworks when training a BERT encoder layer and 1.19x for the entire BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

研究动机与目标

• 识别训练 Transformer 模型的主要性能瓶颈。
• 分析为何现有深度学习框架尽管内核高度优化，仍未能有效利用 GPU 计算能力。
• 研究数据移动与次优数据布局对训练性能下降的影响。
• 开发一种全局优化策略，系统性地减少整个计算图中的数据移动。
• 证明数据移动优化带来的性能提升优于仅依赖算法或硬件层面改进的效果。

提出的方法

• 作者将 Transformer 计算建模为有状态的数据流多图（SDFG），以系统化分析跨操作的数据移动。
• 他们进行详细的数据流分析，识别出内存受限的操作，特别是在多头注意力中，37% 的运行时间消耗在数据移动上。
• 该方法引入一种全局布局转换策略，通过重新排序和重构操作间的数据布局，以最小化冗余数据移动。
• 该方法使用参数化、扁平化的数据流表示，以支持在多个粒度层级上的粗粒度与细粒度布局及融合优化。
• 作者使用自研的基于 SDFG 的编译器栈评估其优化效果，与 PyTorch、TensorFlow with XLA 和 cuDNN 进行性能对比。
• 他们应用布局转换以减少内存访问开销，在矩阵-矩阵乘法（MMM）上实现最高 52% 的性能提升。

实验结果

研究问题

• RQ1 训练现代 Transformer 模型的主要性能瓶颈是什么？
• RQ2 为何现有深度学习框架尽管内核高度优化，仍无法高效利用 GPU 计算能力？
• RQ3 数据布局如何影响数据移动以及 Transformer 模型的整体训练性能？
• RQ4 全局端到端的数据移动优化能否超越局部化或内核级优化？
• RQ5 在 BERT 及其组件中，数据移动减少能在多大程度上提升训练吞吐量？

主要发现

• 在 BERT 训练迭代中，尽管张量收缩操作占了超过 99% 的算术运算，但数据移动仍占 37% 的运行时间。
• 通过全局布局与融合优化，作者将数据移动减少了高达 22.91%，显著提升了性能。
• 所提出的方法在单个 BERT 编码器层上实现 1.30× 的加速，在完整 BERT 模型上实现 1.19× 的加速，优于当前最先进的框架。
• 性能提升主要源于优化后的数据布局，其将矩阵-矩阵乘法（MMM）性能最高提升了 52%。
• 该研究证实，由于 Amdahl's Law 及计算能力与内存带宽增长的不均衡，训练已变为内存受限。
• 该方法具有可扩展性，可推广至其他深度神经网络，为优化其他网络的数据移动提供了一套通用框架。