[论文解读] Survey on Large Scale Neural Network Training

Julia Gusak, Daria Cherniuk+9 more|arXiv (Cornell University)|Feb 21, 2022

Advanced Neural Network Applications被引用 7

一句话总结

本综述全面分析了在单张或多张GPU系统上高效训练大规模深度神经网络的技术。内容涵盖通过重计算（rematerialization）和显存卸载（offloading）实现的内存优化，数据并行、模型并行和流水线并行等并行训练策略，以及结合梯度压缩和低精度计算的优化器，从而在不改变模型或硬件的前提下支持更大的批量大小并提升计算效率。

ABSTRACT

Modern Deep Neural Networks (DNNs) require significant memory to store weight, activations, and other intermediate tensors during training. Hence, many models do not fit one GPU device or can be trained using only a small per-GPU batch size. This survey provides a systematic overview of the approaches that enable more efficient DNNs training. We analyze techniques that save memory and make good use of computation and communication resources on architectures with a single or several GPUs. We summarize the main categories of strategies and compare strategies within and across categories. Along with approaches proposed in the literature, we discuss available implementations.

研究动机与目标

• 识别通用且非侵入性的技术，以在不改变模型或硬件的前提下提升大规模DNN的训练效率。
• 解决单张GPU上因大模型参数、激活值和优化器状态导致的内存瓶颈问题。
• 优化多GPU训练环境中的通信与计算效率。
• 评估大规模DNN训练中内存、计算与通信之间的权衡。
• 对现有技术进行系统性比较，并分析其在现代深度学习框架中的集成方式。

提出的方法

• 通过反向传播过程中重新计算激活值（rematerialization）来减少GPU内存使用。
• 将激活值和权重卸载至CPU或系统内存，以扩展GPU显存容量。
• 采用混合精度和低比特量化技术处理梯度与优化器状态，以降低内存与通信开销。
• 采用数据并行、模型并行和流水线并行策略，将内存与计算负载分发至多张GPU。
• 应用梯度压缩与通信高效的优化器（如LAMB、LARS），以减少GPU间数据交换。
• 集成检查点机制与混合策略（如卸载+重计算），以平衡内存与计算开销。

实验结果

研究问题

• RQ1 在单张GPU上，针对大规模DNN，哪些内存优化技术最为有效？
• RQ2 不同的并行策略（数据并行、模型并行、流水线并行）如何影响通信、计算与内存效率？
• RQ3 梯度压缩与低精度算术能在多大程度上提升训练的可扩展性？
• RQ4 当与并行训练结合时，重计算与卸载策略如何相互作用？
• RQ5 在使用这些优化技术时，训练速度、内存使用与模型收敛性之间的权衡是什么？

主要发现

• 重计算技术可在单张GPU上实现显著的内存节省，且计算开销极低，尤其在与卸载技术结合时效果更佳。
• 激活值与权重的卸载可支持超出GPU显存容量的模型训练，但会增加通信开销。
• 梯度压缩与低精度算术可显著减少多GPU训练中的通信量，提升可扩展性，且对收敛性影响极小。
• 流水线并行与模型并行可降低每张GPU的内存占用，但会增加通信开销，而梯度压缩可有效缓解此问题。
• 重计算与卸载技术的结合可实现比单独使用任一技术更高的内存效率。
• 如LAMB与LARS等优化器可支持更大的批量大小，并提升训练稳定性和收敛性，尤其在与混合精度训练结合时效果更显著。