# TFLMS: Large Model Support in TensorFlow by Graph Rewriting

arXiv: https://arxiv.org/abs/1807.02037 · HTML: https://mlsanity.com/paper/1807.02037/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 5 Jul 2018.

## Built on by (most cited first, 12 of 12)

- [Performance of a Geometric Deep Learning Pipeline for HL-LHC Particle Tracking](https://mlsanity.com/paper/2103.06995/) (2021)
- [Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs](https://mlsanity.com/paper/2505.04519/) (2025)
- [G10: Enabling An Efficient Unified GPU Memory and Storage Architecture with Smart Tensor Migrations](https://mlsanity.com/paper/2310.09443/) (2023)
- [Survey on Large Scale Neural Network Training](https://mlsanity.com/paper/2202.10435/) (2022)
- [Systems for Parallel and Distributed Large-Model Deep Learning Training](https://mlsanity.com/paper/2301.02691/) (2023)
- [Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch](https://mlsanity.com/paper/2307.01236/) (2023)
- [ProTrain: Efficient LLM Training via Memory-Aware Techniques](https://mlsanity.com/paper/2406.08334/) (2024)
- [A Survey on Memory-Efficient Transformer-Based Model Training in AI for Science](https://mlsanity.com/paper/2501.11847/) (2025)
- [Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage](https://mlsanity.com/paper/2506.06472/) (2025)
- [Saturn: An Optimized Data System for Large Model Deep Learning Workloads](https://mlsanity.com/paper/2309.01226/) (2023)
- [STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning](https://mlsanity.com/paper/2507.16274/) (2025)
- [GPU Memory Usage Optimization for Backward Propagation in Deep Network Training](https://mlsanity.com/paper/2502.12499/) (2025)

## Abstract (arXiv)

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 1923 for image segmentation, which, without TFLMS, had been done only by dividing the images to smaller images, which affects the accuracy.