[论文解读] Data Movement Is All You Need: A Case Study of Transformer Networks

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

Advanced Neural Network Applications被引用 3

一句话总结

本文将数据移动识别为训练Transformer网络的主要瓶颈，表明现有框架中低效的数据布局导致内存受限性能。通过全局重构数据布局并优化数据移动，作者将数据移动减少了高达22.91%，在训练BERT时相比最先进框架实现了1.30倍的加速。

ABSTRACT

Transformer neural networks have become widely used for language modeling and sequence learning tasks, and 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 BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

研究动机与目标

• 识别训练Transformer网络性能瓶颈的根本原因。
• 分析尽管计算能力持续进步，现有深度学习框架为何仍未能充分优化数据移动。
• 开发一种适用于Transformer及更广泛深度神经网络的全局优化数据移动策略。
• 通过重新思考数据布局和访问模式，减少训练中的内存受限开销。

提出的方法

• 应用阿姆达尔定律分析数据移动对整体训练性能的影响。
• 识别出现代训练工作负载因计算能力与内存带宽提升不均衡，已演变为内存受限。
• 重新审视并重构Transformer中的数据布局，以最小化跨层和注意力头的冗余数据移动。
• 实现一个全局优化的数据移动流水线，减少注意力机制和前馈层之间冗余的内存传输。
• 在标准硬件上使用BERT训练工作负载验证该方法，以最先进框架作为基线。
• 通过端到端训练时间和数据移动量指标测量性能提升。

实验结果

研究问题

• RQ1 为何现有Transformer实现尽管GPU计算能力不断提升，仍无法实现高效扩展？
• RQ2 在现代Transformer中，数据移动在训练延迟中占据多大程度的主导地位？
• RQ3 如何通过数据布局重组减少注意力和前馈子层中的数据移动？
• RQ4 在Transformer中全局优化数据移动可实现多大程度的性能提升？
• RQ5 所提出的优化策略能否推广到其他深度神经网络架构？

主要发现

• 数据移动是训练Transformer的主要性能瓶颈，导致训练过程受内存限制，尽管计算能力有显著提升。
• 现有框架使用次优的数据布局，导致跨层和注意力头之间存在冗余且低效的数据传输。
• 所提出的全局数据移动优化在BERT训练期间将数据移动量减少了高达22.91%。
• 该优化在端到端BERT训练中相比最先进框架实现了1.30倍的性能提升。
• 该方法具有通用性，为优化Transformer之外的新兴深度学习工作负载提供了洞见。
• 结果表明，对于内存受限工作负载，必须通过数据布局的架构级调整才能释放性能潜力。