Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.
Code (2)
Similar Papers 제목 키워드 기반
XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments
Memory efficiency is crucial in training deep learning networks on resource-restricted devices. During backpropagation, forward tensors are used to calculate gradients. Despite the option of keeping those dependencies in…
SchedulingMoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs
The pervasive "memory wall" bottleneck is significantly amplified in modern large-scale Mixture-of-Experts (MoE) architectures. MoE's inherent architectural sparsity leads to sparse arithmetic compute and also introduces…
Breaking On-device Training Memory Wall: A Systematic Survey
On-device training has become an increasingly popular approach to machine learning, enabling models to be trained directly on mobile and edge devices. However, a major challenge in this area is the limited memory availab…
NavigateSurveyFedHybrid: Breaking the Memory Wall of Federated Learning via Hybrid Tensor Management
Federated Learning (FL) emerges as a new learning paradigm that enables multiple devices to collaboratively train a shared model while preserving data privacy. However, one fundamental and prevailing challenge that hinde…
Federated LearningBare-Metal Tensor Virtualization: Overcoming the Memory Wall in Edge-AI Inference on ARM64
The deployment of Large Language Models (LLMs) on edge devices is fundamentally constrained by the "Memory Wall" the bottleneck where data movement latency outstrips arithmetic throughput. Standard inference runtimes oft…