Neural network training requires immense GPU memory, and memory optimization methods such as recomputation are being actively researched. Recent improvements in recomputation have reduced more than 90 % of the peak allocated size. However, because it produces complex irregular allocation patterns, PyTorch’s default caching allocator wastes up to 20 % of memory because of severe fragmentation and increases the cache management overhead. The periodic allocation patterns during training make offline memory optimization possible. Dynamic storage allocation (DSA) is the problem of minimizing the memory address ranges given the allocation pattern. This is defined as the 2D bin-packing problem, in which each rectangle can move only vertically. Although the first-fit or best-fit heuristics perform well for DSA, we propose a simulated annealing-based non-trivial heuristic algorithm that optimizes the topological ordering of allocations to further minimize fragmentation. The proposed algorithm evaluates a candidate allocation plan with O(logN) amortized time, where N is the number of allocations. We empirically tested our algorithm on both randomly generated data and allocation patterns obtained by training popular vision and text models with recomputation. The experiments showed that, on average, our algorithm reduced fragmentation caused by the PyTorch caching allocator from 29.5 % to 0.4 %, compared to 5.3 % by the first-fit method.

Training large-scale neural networks is heavily constrained by GPU memory. In order to circumvent this limitation, gradient checkpointing, or recomputation is a powerful technique. There is active research in this area with methods such as Checkmake or Moccasin. However, both Checkmate and Moccasin rely on mixed integer linear programming or constraint programming, resulting in limited scalability due to their exponentially large search space. This paper proposes a novel algorithm for recomputation (FastSA) based on a simulated annealing heuristic that achieves comparable or even better solutions than state-of-the-art alternatives. FastSA can optimize computational graphs with thousands of nodes within 3 to 30 seconds, several orders of magnitude faster than current solutions. We applied FastSA to PyTorch models and verified its effectiveness through popular large vision and text models, including recent language models with the transformer architecture. The results demonstrate significant memory reductions by 73% with extra 18% computational overheads on average. Our experiments demonstrate the practicality and effectiveness of our recomputation algorithm, further highlighting its potential for wide application in various deep learning domains.

Loop fusion—a program transformation to merge multiple consecutive loops into a single one—has been studied mainly for compiler optimization. In this paper, we propose a new loop fusion strategy, which can fuse any loops—even loops with data dependence—and show that it is useful for program verification because it can simplify loop invariants.