[论文解读] TFLMS: Large Model Support in TensorFlow by Graph Rewriting

Tung D. Le, Haruki Imai+2 more|arXiv (Cornell University)|Jul 5, 2018

Topic Modeling参考文献 11被引用 28

一句话总结

本文提出 TFLMS，一种在 TensorFlow 中的图重写方法，通过插入交换出（swap-out）和交换入（swap-in）操作，将中间激活值卸载到 CPU 内存，从而克服 GPU 内存限制，实现大规模深度神经网络的训练。该方法使用程序转换规则形式化推导这些操作，对 ResNet-50 和 3DUnet 分别实现了最高 4.7× 和 2× 的批量大小提升，包括实现无需分块处理的 $192^3$ 全尺寸 3D 图像的 3DUnet 训练。

ABSTRACT

While accelerators such as GPUs have limited memory, deep neural networks are becoming larger and will not fit with the memory limitation of accelerators for training. We propose an approach to tackle this problem by rewriting the computational graph of a neural network, in which swap-out and swap-in operations are inserted to temporarily store intermediate results on CPU memory. In particular, we first revise the concept of a computational graph by defining a concrete semantics for variables in a graph. We then formally show how to derive swap-out and swap-in operations from an existing graph and present rules to optimize the graph. To realize our approach, we developed a module in TensorFlow, named TFLMS. TFLMS is published as a pull request in the TensorFlow repository for contributing to the TensorFlow community. With TFLMS, we were able to train ResNet-50 and 3DUnet with 4.7x and 2x larger batch size, respectively. In particular, we were able to train 3DUNet using images of size of $192^3$ for image segmentation, which, without TFLMS, had been done only by dividing the images to smaller images, which affects the accuracy.

研究动机与目标

• 解决训练超大深度神经网络时超出 GPU 内存容量的挑战。
• 提供一种形式化、可证明正确的插入交换出与交换入操作的方法，用于计算图。
• 通过图重写优化内存使用，提升大规模模型训练的效率与可扩展性。
• 实现无需图像分块处理的 3DUnet 在全尺寸 $192^3$ 图像上的训练，此前该任务需依赖图像分块且影响精度。

提出的方法

• 通过将变量视为具有正式语义的一等公民，重新定义计算图，扩展 TensorFlow 的图模型。
• 使用形式化程序转换规则推导交换出与交换入操作，确保正确性与一致性。
• 引入两种策略——chain_rule 和 direct_order——静态识别控制依赖关系，以触发最优的交换入时机。
• 应用优化规则（如融合交换入操作）以减少开销并提升性能。
• 作为 TensorFlow 仓库中的拉取请求（pull request）实现该方法，命名为 TFLMS，以促进社区采纳。
• 使用可配置参数（n_tensors, lb, swap_branches）控制内存使用与训练性能之间的权衡。

实验结果

研究问题

• RQ1 如何形式化重新定义计算图，以支持一等公民变量，并实现正确的图重写以用于内存管理？
• RQ2 可采用哪些形式化规则来推导交换出与交换入操作，确保训练过程中模型正确性？
• RQ3 不同控制依赖关系识别策略对大规模模型训练性能有何影响？
• RQ4 如何配置参数（n_tensors, lb, swap_branches）以最大化批量大小与训练吞吐量？
• RQ5 TFLMS 是否能够实现无需图像分块处理的 3DUnet 在全尺寸 $192^3$ 图像上的训练？性能表现如何？

主要发现

• TFLMS 实现了 ResNet-50 最大批量大小达 832，相比无 TFLMS 时的基线 176 提升了 4.7×。
• 对于 3DUnet，TFLMS 实现了相比基线 2× 的最大批量大小提升，支持在无需分块处理的情况下对 $192^3$ 全尺寸图像进行训练。
• ResNet-50 的最高训练吞吐量出现在 lb=5 且 n_tensors=100 时，该配置在内存使用与计算-通信重叠之间达到最佳平衡。
• 融合交换入操作可提升性能，但因中间张量在 GPU 中保留时间更长，导致最大批量大小下降。
• 启用分支交换并设置阈值为 20 后，3DUnet 的最大批量大小从 3 提升至 4，显著提升了可扩展性。
• 在最优配置下，3DUnet 在 $192^3$ 输入上实现了 0.17 张像/秒的训练吞吐量，此前因内存溢出而无法实现。