TFLMS: Large Model Support in TensorFlow by Graph Rewriting

By Tung D. Le, Haruki Imai, Yasushi Negishi, Kiyokuni Kawachiya. 2018

55 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “[20] uses graph rewriting to offload activations to CPU memory, but does not reduce memory required by model states.”
ProtTrans: Towards Cracking the Language of Life’s Code Through Self-Supervised Deep Learning and High Performance Computing (2020, methodology): “Another optimization technique/library crucial for our training on Summit was IBM’s large model support (LMS) [57].”
SwapAdvisor: Pushing Deep Learning Beyond the GPU Memory Limit via Smart Swapping (2020, methodology): “However, A promising approach to address the GPU memory limitation without affecting accuracy is to swap tensor data be-tween GPU and CPU memory during DNN computation [26, 30, 40, 49].”
ProtTrans: Towards Cracking the Language of Life’s Code Through Self-Supervised Learning (2021, methodology): “Another optimization technique/library crucial for our training on Summit was IBM’s large model support (LMS) [66].”
Sentinel: Efficient Tensor Migration and Allocation on Heterogeneous Memory Systems for Deep Learning (2021, background): “…existing solutions explore methods to proactively release and prefetch temporarily inactive tensors, determine inactive ten sors based on DNN topology, and find data swap time between memories by analyzing tensor access order [7]-[10] or using detailed domain knowledge [5], [6], [11].”