GPU Memory Usage Optimization for Backward Propagation in Deep Network Training

2025-02-01

Journal of Parallel and Distributed Computing

AI Summary

To address excessive GPU memory peaks during the backward pass in deep neural network training—which constrain model scale and batch size—this paper proposes a gradient-lifetime-aware dynamic memory scheduling mechanism. Our method jointly models computational graph dependencies and tensor lifetimes to enable fine-grained, zero-copy reuse of gradient memory. By integrating CUDA Graph with PyTorch Autograd hooks, it unifies static graph analysis and runtime dynamic memory reclamation, ensuring seamless compatibility with arbitrary automatic differentiation frameworks. Evaluated on ResNet-50 and ViT-L, our approach reduces backward-pass memory consumption by 47%–63%, improves training throughput by 1.8×, and enables doubling the batch size without gradient checkpointing.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Problem

• Optimize GPU memory for deep network training
• Select optimal checkpoints for memory efficiency
• Develop algorithms for minimal peak memory usage

Innovation

• Dynamic programming for checkpoints
• Optimizing GPU memory usage
• Reducing peak memory in training