Echo：基于编译器的 LSTM RNN 训练 GPU 内存占用缩减

Bojian Zheng; Abhishek Tiwari; Nandita Vijaykumar; Gennady Pekhimenko

Echo：基于编译器的 LSTM RNN 训练 GPU 内存占用缩减

机器学习 2019-12-02 v5 人工智能机器学习

作者: Bojian Zheng , Abhishek Tiwari , Nandita Vijaykumar , Gennady Pekhimenko

摘要

长短期记忆循环神经网络（Long-Short-Term-Memory Recurrent Neural Networks, LSTM RNNs）是一类用于分析序列数据的流行机器学习模型。然而，其在现代 GPU 上的训练受限于 GPU 内存容量。我们基于 LSTM RNN 的神经机器翻译（Neural Machine Translation, NMT）模型的性能分析结果显示，注意力层与 RNN 层中的特征图构成了内存瓶颈，且在 GPU 上训练时运行时间在不同层间分布不均。基于这两点观察，我们提议重新计算特征图而非将其持久暂存于 GPU 内存中。尽管特征图重计算的思想此前已被考虑，现有方案未能提供令人满意的内存占用缩减，因为它们未解决两个关键挑战。为使每次特征图重计算有效且高效，必须仔细估计其对（1）总内存占用与（2）总执行时间的影响。为此，我们提出 *Echo*，一种新颖的基于编译器的优化方案，以实用机制估计整个计算图上重计算的内存收益来解决第一重挑战，并利用层特性非保守地估计重计算开销来解决第二重挑战。*Echo* 自动且透明地缩减 GPU 内存占用，无需对训练源代码做任何改动，且对 LSTM RNN 之外的模型亦有效。我们在配备现代 GPU 的真实系统上对众多最先进机器学习负载评估了 *Echo*，观察到平均 1.89 倍、最大 3.13 倍的内存占用缩减比。此类缩减可转化为以更大批量更快训练、节省 GPU 能耗（例如，单 GPU 训练快如四 GPU）和/或在相同 GPU 内存预算下增加最大层数。

关键词

LSTM RNN GPU 内存编译器优化特征图重计算神经机器翻译训练

引用

@article{arxiv.1805.08899,
  title  = {Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training},
  author = {Bojian Zheng and Abhishek Tiwari and Nandita Vijaykumar and Gennady Pekhimenko},
  journal= {arXiv preprint arXiv:1805.08899},
  year   = {2019}
}

Echo：基于编译器的 LSTM RNN 训练 GPU 内存占用缩减

摘要

关键词

引用

相关论文