[论文解读] Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory

Julien Herrmann, Olivier Beaumont+4 more|arXiv (Cornell University)|Nov 27, 2019

Stochastic Gradient Optimization Techniques参考文献 29被引用 13

一句话总结

本文提出了一种在有限内存下训练深度神经网络的最优检查点策略，通过动态选择存储哪些前向激活以及哪些进行重新计算。该方法采用一种新颖的动态规划算法，同时建模仅输入和完整操作历史的检查点机制，在支持任意内存限制和复杂顺序网络的PyTorch环境中，平均比现有方法提升17.2%的吞吐量。

ABSTRACT

This paper introduces a new activation checkpointing method which allows to significantly decrease memory usage when training Deep Neural Networks with the back-propagation algorithm. Similarly to checkpoint-ing techniques coming from the literature on Automatic Differentiation, it consists in dynamically selecting the forward activations that are saved during the training phase, and then automatically recomputing missing activations from those previously recorded. We propose an original computation model that combines two types of activation savings: either only storing the layer inputs, or recording the complete history of operations that produced the outputs (this uses more memory, but requires fewer recomputations in the backward phase), and we provide an algorithm to compute the optimal computation sequence for this model. This paper also describes a PyTorch implementation that processes the entire chain, dealing with any sequential DNN whose internal layers may be arbitrarily complex and automatically executing it according to the optimal checkpointing strategy computed given a memory limit. Through extensive experiments, we show that our implementation consistently outperforms existing checkpoint-ing approaches for a large class of networks, image sizes and batch sizes.

研究动机与目标

• 为解决在低内存设备（如物联网系统）上训练深度神经网络时的内存瓶颈问题。
• 将传统的自动微分检查点技术扩展至具有任意复杂内部层的异构深度神经网络。
• 开发一种完全自动化的、与PyTorch兼容的实现，支持任意内存限制并优化训练吞吐量。
• 在多种网络架构、图像尺寸和批量大小下，超越现有的周期性与最优检查点策略。

提出的方法

• 该方法将深度神经网络的前向与反向操作建模为一系列模块，从而实现对哪些激活进行检查点存储的动态选择。
• 引入两种检查点类型：仅存储层输入（内存高效）或记录完整的操作历史（减少重新计算次数）。
• 采用动态规划算法，在内存约束下计算最优持久化调度，利用最优解的内存持久性特性。
• 该方法支持任意顺序的PyTorch模块，包括具有复杂内部计算的模块。
• 实现与autograd无缝集成，根据用户定义的内存限制自动应用最优策略。
• 支持连续内存自适应，不同于离散分段策略，并在前向与反向阶段均实现极低开销。

实验结果

研究问题

• RQ1 如何为具有任意内部层复杂度的异构深度神经网络链优化检查点策略？
• RQ2 在仅存储输入或完整操作历史之间，内存使用与重新计算成本之间的最优权衡是什么？
• RQ3 在内存受限条件下，动态规划方法是否能超越现有的周期性与最优检查点策略，在训练吞吐量方面表现更优？
• RQ4 与PyTorch默认策略相比，该方法在不同网络深度、图像尺寸和批量大小下的可扩展性如何？
• RQ5 该方法在多大程度上可通过支持更大模型、图像或批量大小，使低内存设备具备训练能力？

主要发现

• 所提出的最优策略在所有测试配置下，平均比表现最佳的竞争对手提升17.2%的训练吞吐量。
• 对于输入尺寸为224×224、批量大小为8的ResNet-1001，最优策略达到9.18张图像/秒的吞吐量，而PyTorch默认策略因内存限制而失败。
• 在批量大小为8时，PyTorch策略需要225 GiB内存——超过最优策略峰值使用量的16倍以上。
• 最优策略即使在PyTorch默认策略失败的情况下，也能实现大模型的训练，包括在ResNet-1001上批量大小为8的情况。
• 该方法在高内存限制下，优于顺序（分段）策略与先前的revolve算法，尤其在后者无法充分利用可用内存时表现更优。
• 最优策略在不同网络规模与图像分辨率下均保持稳定性能，展现出强大的鲁棒性与灵活性。