Computer Science > Machine Learning
[Submitted on 31 Jul 2018 (v1), revised 2 Mar 2021 (this version, v5), latest version 18 Mar 2021 (v6)]
Title:Optimal Gradient Checkpoint Search for Arbitrary Computation Graphs
View PDFAbstract:Deep Neural Networks(DNNs) require huge GPU memory when training on modern image/video databases. Unfortunately, the GPU memory in off-the-shelf devices is always finite, which limits the image resolutions and batch sizes that could be used for better DNN performance. Existing approaches to alleviate memory issue include better GPUs, distributed computation and Gradient CheckPointing(GCP) training. Among them, GCP is a favorable approach as it focuses on trading computation for memory and does not require any upgrades on hardware. In GCP, during forward, only a subset of intermediate tensors are stored, which are called Gradient Checkpoints (GCs). Then during backward, extra local forwards are conducted to compute the missing tensors. The total training memory cost becomes the sum of (1) the memory cost of the gradient checkpoints and (2) the maximum memory cost of local forwards. To achieve maximal memory cut-offs, one needs optimal algorithms to select GCs.
Existing GCP approaches rely on either manual input of GCs or heuristics-based GC search on linear computation graphs (LCGs), and cannot apply to arbitrary computation graphs(ACGs). In this paper, we present theories and optimal algorithms on GC selection that, for the first time, are applicable on ACGs and achieve maximal memory cut-offs. Extensive experiments show that our approach constantly outperforms existing approaches on LCGs, and can cut off up-to 80% of training memory\footnote{Cutting off 80% of training memory means one can double the input image size or quadruple the batch size on the same GPUs.} with a moderate time overhead (around 40%) on LCG and ACG DNNs, such as Alexnet, VGG, Resnet, Densenet and Inception Net.
Submission history
From: Jianwei Feng [view email][v1] Tue, 31 Jul 2018 21:52:08 UTC (418 KB)
[v2] Sat, 5 Jan 2019 21:43:43 UTC (2,038 KB)
[v3] Thu, 23 May 2019 01:27:21 UTC (2,081 KB)
[v4] Mon, 23 Sep 2019 23:18:18 UTC (2,531 KB)
[v5] Tue, 2 Mar 2021 03:38:11 UTC (1,570 KB)
[v6] Thu, 18 Mar 2021 04:32:09 UTC (1,570 KB)
Current browse context:
cs.LG
References & Citations
Bibliographic and Citation Tools
Bibliographic Explorer (What is the Explorer?)
Litmaps (What is Litmaps?)
scite Smart Citations (What are Smart Citations?)
Code, Data and Media Associated with this Article
CatalyzeX Code Finder for Papers (What is CatalyzeX?)
DagsHub (What is DagsHub?)
Gotit.pub (What is GotitPub?)
Papers with Code (What is Papers with Code?)
ScienceCast (What is ScienceCast?)
Demos
Recommenders and Search Tools
Influence Flower (What are Influence Flowers?)
Connected Papers (What is Connected Papers?)
CORE Recommender (What is CORE?)
IArxiv Recommender
(What is IArxiv?)
arXivLabs: experimental projects with community collaborators
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.