[论文解读] Memory Optimization for Deep Networks

Aashaka Shah, Chao-Yuan Wu+3 more|arXiv (Cornell University)|Oct 27, 2020

Advanced Neural Network Applications参考文献 34被引用 10

一句话总结

MONeT 是一个自动框架，通过联合优化检查点调度和内存高效的算子实现，减少深度学习训练中的内存使用。它通过将内存与计算权衡建模为整数规划问题，在仅增加 9–16% 计算开销的前提下，相较于 PyTorch 实现了 3 倍的内存减少，优于手工调优和自动化基线方法。

ABSTRACT

Deep learning is slowly, but steadily, hitting a memory bottleneck. While the tensor computation in top-of-the-line GPUs increased by 32x over the last five years, the total available memory only grew by 2.5x. This prevents researchers from exploring larger architectures, as training large networks requires more memory for storing intermediate outputs. In this paper, we present MONeT, an automatic framework that minimizes both the memory footprint and computational overhead of deep networks. MONeT jointly optimizes the checkpointing schedule and the implementation of various operators. MONeT is able to outperform all prior hand-tuned operations as well as automated checkpointing. MONeT reduces the overall memory requirement by 3x for various PyTorch models, with a 9-16% overhead in computation. For the same computation cost, MONeT requires 1.2-1.8x less memory than current state-of-the-art automated checkpointing frameworks. Our code is available at https://github.com/utsaslab/MONeT.

研究动机与目标

• 为解决深度学习训练中日益严重的内存瓶颈问题，过去五年间 GPU 内存仅增长 2.5 倍，而计算量却增长了 32 倍。
• 克服现有方法的局限性，这些方法仅孤立地优化局部算子实现或全局检查点调度。
• 开发一个自动框架，在固定内存预算下联合优化局部和全局内存减少技术。
• 在保持计算效率的同时，最小化前向和反向传播过程中的峰值内存消耗。
• 提供一个统一且理论基础坚实的优化框架，严格限制内存使用，并实现内存与计算之间的最优权衡。

提出的方法

• 形式化不同算子实现和检查点调度下，前向、反向和重计算阶段的内存与计算成本边界。
• 将算子实现与检查点的联合优化建模为带有线性化内存约束的 0-1 整数规划问题。
• 使用标准混合整数线性规划求解器，寻找在给定预算下最小化内存使用的最优配置。
• 支持多种优化策略：卷积算法选择、就地激活、输出激活复用以及中间激活重计算。
• 与 PyTorch 集成，通过最小的框架修改支持端到端训练。
• 采用理论分析推导出峰值内存消耗的紧致边界，实现精确优化。

实验结果

研究问题

• RQ1 联合优化局部算子实现与全局检查点调度，是否能比孤立优化任一方法获得更优的内存-计算权衡？
• RQ2 在反向传播中结合不同内存节省技术时，峰值内存消耗的理论边界是什么？
• RQ3 卷积算法选择与检查点的联合优化如何影响整体内存与计算效率？
• RQ4 自动选择内存高效算子实现，在多大程度上优于手工设计或基于启发式的方案？
• RQ5 与最先进的自动化检查点框架相比，MONeT 在内存节省和计算开销方面表现如何？

主要发现

• MONeT 在多个模型（包括 ResNet、VGG、UNet、GoogleNet 和 MobileNet-V2）上相比 PyTorch 实现了 3 倍的内存减少。
• 在相同计算成本下，MONeT 所需内存仅为当前最先进的自动化检查点框架的 1.2–1.8 倍。
• 与标准 PyTorch 训练相比，MONeT 仅带来 9–16% 的计算时间增加。
• 卷积算法与检查点的联合优化带来显著性能提升，因为卷积工作区内存特性强烈影响速度与内存的权衡。
• 消融研究显示，结合所有优化策略（卷积、输出激活、中间激活）可实现最低的计算开销。
• 对 ResNet-50 的详细案例研究证实，MONeT 将峰值内存从 PyTorch 的 14.7 GB 降低至 8.0 GB，同时保持训练稳定性。