HiRemate: Hierarchical Approach for Efficient Re-materialization of Neural Networks
Authors: Julia Gusak
Xunyi Zhao
Théotime Le Hellard
Zhe LI
Lionel Eyraud-Dubois
Olivier Beaumont
Venue: ICML 2025 poster (2025)
Keywords: Rematerialization
Checkpointing
Memory-Efficient Training
Neural Networks
PyTorch
Integer Linear Programming
Training
Primary area: general_machine_learning->hardware_and_software

Abstract: Training deep neural networks (DNNs) on memory-limited GPUs is challenging,  as storing intermediate activations often exceeds available memory. Re-materialization, a technique that preserves exact computations, addresses this by selectively recomputing activations instead of storing them. However, existing methods either fail to scale, lack generality, or introduce excessive execution overhead. We introduce ${\mbox{HiRemate}}$ a ${\textit hierarchical}$ re-materialization framework that recursively partitions  large computation graphs, applies optimized solvers at multiple levels, and merges solutions into a global efficient training schedule. This enables scalability to significantly larger graphs than prior ILP-based methods while keeping runtime overhead low. Designed for single-GPU models and activation re-materialization, HiRemate extends  the feasibility of training networks with thousands of graph nodes, surpassing  prior methods in both efficiency and scalability. Experiments on various types of networks yield up to 50-70%  memory reduction with only 10-15% overhead, closely matching optimal solutions while significantly reducing solver time. Seamlessly integrating with PyTorch Autograd, HiRemate requires almost no  code change to use,  enabling broad adoption in memory-constrained deep learning.

HIREMATE:
Hierarchical Approach for Efficient Re-materialization of Neural Networks
Julia Gusak" 1
Xunyi Zhao *1
Théotime Le Hellard 2 Zhe Li Lionel Eyraud-Dubon Olivier Beaumont
Abstract
Training deep neural networks (DNNs) on memory-limited GPUs challenging. as storing intermediate activations often exceeds available memory. Re-materialization, technique that pre- serves exact computations, addresses this by se- lectively recomputing activations instead of stor- ing them. However, existing methods either fail to scale, lack generality, or introduce excessive execution overhead. We introduce HIREMATE, hierarchical re-materialization framework that recursively partitions large computation graphs, applies optimized solvers at multiple levels, and merges solutions into global efficient train- ing schedule. This enables scalability to signifi- cantly larger graphs than prior ILP--baed meth- ods while keeping runtime overhead low. De- signed for single-GPt models and activation re- materialization, HiRemate extends the feasibil- ity of training networks with thousands of graph nodes, surpassing prior methods in both efficiency and scalability. Experiments on various types of networks yield up to 50-70% memory reduc- tion with only 10-15% overhead, closely match- ing optimal solutions while significantly reduc- ing solver time. Seamlessly integrating with Py- Torch Autograd, HiRemate requires almost no code change use, enabling broad adoption in memory-constrainee deep learning.
Firect solvers
3. Solving bottom-level nevel subgraphs with different solvers bunlgets
hieracchian
Recusivedy solving higher-level subgraphs
Partitioned graph
Merging schodules into final solution
Original graph
Figure 1. HIREMATE takes a PyTorch nn. Module e and creates M.MModuee,, e, which provides outputs while satisfying peak memory budget constraints B during training. (1) Obtain data-flow graph G from the PyTorch module. (2) Recursively partition G with H-Partition (Section 3.1) into small-size subgraphs. (3)-(5) Obtain a re-materialization schedule from G with budget using H-Solver (Section 3.2). (6) Produce new nn. Module whose execution follows the schedule (Appendix).
like ResNet-50 (Wu et al., 2019) Natural Language Pro- cessing transformer-basee models (Vaswani et al., 2017) like GPT. On the one hand, the size of the models and the resolution of the data are increasing, which raises problems for the storage of both weights and activations. On the other hand, the first models had chain-like structures (se- quence of convolution layers) and then moved chains of complex blocks (chains of Transformer blocks for GPT-like models). Finally, recent NN exhibit arbitrarily complex dependency graph structures between layers of the neural network: for example, UNO (Wen et al., 2022) is struc- tured as a U-Net of complex blocks, and encoder-decoder transformers (Vaswani et al., 2017) feature very long skip connections.
1. Introduction
Modern Neural Networks (NN) undergo several important evolutions which have consequences on the computation and memory requirements, from the first vision networks
'qqual contribution IIniaa Center at the University of Bordcaux "École Normale Supérieure, PSL University, Paris. Correspon- dence to: Julia Gusak Xyulia.gusa @ in Lionel Eyraud- Dubois wewionnl.yrrudddubleeo @ in
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Re-materialization well known and efficient technique to limit the memory requirements related to activations dur- ing training, The idea is to avoid storing all the necessary activations because of the subsequent dependencies during the Forward phase and the Backward phase (those related to the Stochastic Gradient Descent mechanism). Some ac- tivations are computed and then deleted to make room in memory, and they will have to be re-computed later when needed. In the case of simple chains (Beaumont al., 2019) and in the case of chains of complex blocks (Zhao et al., 2023), re-materialization has shown its efficiency: it is of- ten possible to save 50% of memory for computational overhead of about 10 to 15%.
trol the memory usage of their neural network: model HRockmate (model, sample, memory budget).
To achieve this result, we rely on the following main contri- butions:
A data-flow graph decomposition algorithm H- Partition that builds a hicrarchy of blocks of reason- able sizes (to keep an acceptable computational com- plexity) while minimizing the memory size of the in- terfaces between the blocks
new linear programming solver H-ILP adapted to this hierarchical decomposition.
From a theoretical point of view, the problem is to minimize the computational time required execute the forward and backward passes under predefined memory budget. It has been proven NP-Hard in (Naumann, 2008) for the case of general dependency graphs represented as general data-flow graphs. Some solutions such as TW- W-RREMTT (Kumar et al., 2019) or CHECKMATE (Jain et al., 2020) have nevertheless been designed to deal with the case of general graphs, but both suffer from a number of limitations. These limitations stem either from the computational cost and the scalability of the algorithm that generates the ce-materialization strat- egy, or from the overhead of that re- re-materialization strategy during the training phase, as discussed in Section 2. Another line of research is to propose re- materialization strategies whose computational cost and overhead are controlled, as in ROTOR (Beaumont et al., 2019) and ROCKMATE (Zhao et al., 2023), however at the cost of generality, by propos- ing solutions only for limited classes of dependency graphs where chain structure (of potentially complex blocks) can be identified.
general framework for integrating any existing (or future) re-materialization strategy at any level of the hierarchy, and combining their strengths.
The rest of the paper is organized as follows. In Section 2, we review the related work memory saving strategies for DNN training. HIREMATE is presented in Section 3 which covers graph decomposition, partial problems resolution and global re- ce-matrrialization strategy reconstruction. Section 4 demonstrates the efficiency of the proposed method on a large number of networks, and provides an evaluation of the computational overhead induced by re-materialization. Finally, concluding remarks are proposed in Section 5,
Forward
Backwaar
Forward subgraph
Backward subgraph
Loss, between
The work presented here is at the convergence of these re- search lines and we propose a computationally efficient, low-overhead solution that can address general graphs. Our framework HIREMATE is based on hierarchical decom- position approach of the computation graph, to find re- materialization strategy for any graph of dependencies between layers. In the case where the graph is too large to be addressed directly by an approach based on Integer Lin- ear Programming, we propose to decompose it into a graph of complex blocks, with potentially several levels in the hier- archy handle very large graphs. Efficient solutions for dif- ferent memory budgets are generated for each of the blocks at the bottom of the hierarchy, using different approaches from the literature. Then, we provide a new Integer Linear Programming (ILP) formulation to efficiently recombine these low-level solutions into candidate solutions for the is
forward bacleward
Data
Figure 2. HIREMATE recursively finds schedules for each sub- graph. Different schedules correspond to differrent memory budget constraints and hence have different values for memory/time ratio, peak memory, size allocated tensors, and execution time.
2. Related Work
We mostly cover the contributions that rely on the exclu- sive use of re-materialization. Memory footprint of ac- tivations might also be reduced via data parallelism (Das et al., 2016; Zhang et al., 2013) (by distributing mini-batches across several computational resources) and offloading (Rhu et al., 2016; Wang al., 2018) (which offloads and later retrieves activations from GPU memory to CPU memory). Gradient accumulation can also reduce memory usage by splitting large batch into smaller sub-batches and accu- mulating gradients across them, but is limited by the low GPU utilization because of smaller sub- batches and re- quires that one smaller sub-batch fits in memory. Model
higher levels of the hierarchy. fully compatible with the autograd mechanism of Py- Torch, that modification of the code is required to use it. With a single line, the user can automatically con-
2

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Table Comparison of Different Re- Re-Mattrialization Strategies in terms of (i) Generality (ability handle complex models with arbitrary dependencies), (ii) Scalability (ability to handle models with many operations), (iii) Quality (ability to generate a me-materialization strategy with low overhead during training), and Ease use terms of modifications to be made the original Py PoToch model). Scores range from "-" (true limitation) to "++" (excellent).
GENERALITY
SCALABILITY
QUALITY
EASE OF USE
CHECKMATE (JAIN ET AL., 2020)
0
TW-REMAT (KUMAR ET AL.., 2019)
++
MOCCASIN (BARTAN ET AL., 2023)
++
++
ROCKMATE (ZHAO ET AL.,, 2023)
HIREMATE (PRESENT PAPER)
++
++
parallelism (Huang et al., 2019; Narayanan et al., 2019), in particular in its pipelined version, can be used to mini- mize the memory requirements by storing weights across several devices. The optimal combination of all these strate- gies with re- materialization, discussed by (Beaumont et al., 2021) in the context of sequential models, is left for future work.
or encoder-decoder Transformer (Vaswani et al., 2017)), ROCKMATE degenerates to CHECKMATE and inherits its difficulties handle large data-flow graphs.
Such static rematerialization strategies are very efficient for models with fixed computation graphs, where memory usage can be optimized offline for efficient execution. Nev- ertheless, they struggle to handle models with dynamic or input-dependent control flow. In contrast, dynamic remateri- alization strategies (Kirisame et al., 2020) adapt runtime behavior by making decisions on-the-fly the-ffy. allowing them to support broader range of model architectures. However, this flexibility often comes at the cost of higher runtime overhead and less predictable performance. Notably, dy- namic approaches such as POET (Patil et al., 2022) and TaiChi (Hu al., 2022) combine rematerialization with additional mechanisms like paging.
TW-Remat (Kumar et al., 2019), based on a tree-width de- composition of the dependency graph, was initially designed for the case where all computational costs and activation sizes have unitary weight. A greedy heuristic was later pro- posed to generalize the algorithm to the cases of non-unitary weights, but without guarantee. (Kusumoto et al., 2019) proposed an optimal dynamic programming algorithm re- stricted to a special class of solutions, where each node is computed at most twice. The XLA framework (xla) contains re-materialization feature based on greedy heuristic de- scribed in Kumar et al. (2019), and shown to be less efficient than Te-Remat. CHECKMATE (Jain et al., 2020) is based on the solution of an Integer Linear Program (ILP), and com- putes a generic optimal solution under a set of reasonable assumptions, but the computation cost becomes prohibitive as soon as the number of nodes in the network exceeds 70- 90 nodes. Moccasin (Bartan et al., 2023) is constraint programming based solution we show that HIREMATE does actually outperform Moccasin, both in terms of solution time and the quality of the solution produced.
In this paper, focus specifically static strategies and we design an algorithm and a framework to address any kind of network with static data-flow graph, with reasonable solving time on networks with up to 2000 nodes.
In the experimental evaluation of Section 4, we compare HIREMATE with the state-of-the-ae strategies: CHECK- MATE, MOCCASIN, TW-REMAT, and ROCKMATE. Table I summarizes the comparison of these re-materialization tech- niques.
In the experimental evaluation of Section 4, we compare HIREMATE with the state-of-the-art strategies: CHECK- MATE, TW-Remat, and ROCKMATE. Table summarizes the comparison of these e--ateriiaizztion techniques with several metrics. Generality denotes the ability handle complex models with arbitrary dependencies, scalability denotes the ability to handle models with many operations. Quality measures the ability to generate re-materialization strategy with low overhead during training, and ease of use is expressed in terms of modifications to be made to the original Py PyTorch model.
Other approaches have been proposed to find efficient solu- tions for limited classes of dependency graphs. (Beaumont et al., 2024 (accepted for publication) relies on dynamic programming in the case of networks whose graph is se- quence of basic operations, covering ResNet-like networks (considering one ResNet block as one operation). This extends in a proven framework the heuristic techniques pro- posed by (Chen et al., 2016). Then, the case of graphs for which the forward graph is a sequence of complex blocks has been addressed in the ROCKMATE framework in (Zhao et al., 2023), which covers most Transformer-based networks (where one Transformer block is considered as complex block). However, for architectures with long skip connections (like U-Net (Ronneberger et al., 2015)
3

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
3. HIREMATE
Algorithm H-Partition Bottom-to-Top algorithm
1: Input: data-flow graph G
Problem statement We associate each neural network with a data-flow graph, where each node represents a tensor- level computation. schedule a sequence of elementary computations and tensor deletions that performs the forward and backward passes for one training iteration.
Result: a recursive partition of G
3: Parameters: max high-level size M max lower-level
size M', score parameter
4: while G' has more than M* nodes do
5: C + UseG candidate group containing all nodes
Let FÚ and B; denote the forward pass and backward pass associated with layer i and D; denotes the deletion of the output of F.. Then, for a model that is a sequence of 3 layers, a schedule corresponding to the standard au- todiff training can be written as F F2/ ByllgBBBBl. B1, while is another valid sched- ule, yielding smaller peak memory since it does not require storing all intermediate activations simultaneously. The re-materialization optimization problem (2) is to find a schedule whose peak memory is below a given budget and whose execution time as small as possible.
between x and a(x)
6: while is non empty and G has more than M nodes
do
Select candidate C' which minimizes so (eq. 1)
8:
Wrap the nodes of C' into group
Fi F2F2D,F33 Fy
9:
Update C
10: end while
11: Consider all groups as subgraphs
12: Update G so that each subgraph is considered as a node
13: end while
14: return partitioned graph
Figure 1 describes the main steps of the HIREMATE ap- proach. In the first step, using the graph building tool adapted from (Zhao et al., 2023), data flow graph of the input module is obtained. In the second step, the H- partition algorithm (see Section 3. recursively partitions the graph into subgraphs of manageable sizes, This parti- tioning done recursively to ensure that all parent graphs also of manageable size. Steps 3, 4, and 5 of Figure describe the H-Solver algorithm (see Section 3.2), which builds training schedule. Starting at the lowest level of the decomposition, the algorithm computes schedules for each subgraph (Figure 2) with different memory budgets to explore different time-memory tradeoffs, providing several options (i.e ways to (re)compute operations and store ac- tivations during forward and backward passes through the subgraph) for the nodes at higher levels. This procedure continues until the top-level graph is solved (i.e schedule is found) for a single memory budget corresponding to the overall memory available for activations. It is important to emphasize that the current implementation of the general scheme described above and Figure 1 is fully modular. Thus, HIREMATE allows the injection of custom parti- tioners as well as solvers at any level of the graph.
parameters: M' denotes the maximum number of nodes in a lower-level subgraph, and M denotes the maximum number of nodes in the top-level graph. Since it is advanta- geous to allow longer solution time for the top-level graph, we use M 2 M Our partitioning algorithm is a greedy bottom-to-toy heuristic described in Algorithm 1. Each iteration has three main steps: forming candidate groups, selecting the best candidate according to our evaluation cri- terion, merging the selected candidate, and updating the candidates. Each step is described below.
Forming candidate groups For each node x in G, we consider a(z), the closest common ancestor of all direct predecessors of x: (a common global ancestor is added in case G has multiple entries). We create four candidate groups with all nodes on all paths from azz) to x, depending on whether x and/or axx) are included. Any candidate group with more than A nodes is discarded.
Selecting the best candidate When selecting the best group among all candidates, our goal is avoid incurring too much memory pressure when the group is used sub- graph. The memory pressure depends directly on the size of the input and output values. It also depends, although not as directly, on the length of the schedule during which they will be alive, which we evaluate by the number of compute nodes in the group. We use the following score function s sCC) for candidate C:
3.1. H-Partition
The goal of HIREMATE'S partitioning step is to reduce the size of the problems to be solved without compromising the quality of the overall solution. The result of this step is decomposition into hierarchy of subgraphs, where a node at a given level represent an entire subgraph at the level below (see Figure 7 in Appendix). Note that the orig- inal graph is partitioned without requiring a topological order; topological sorting only performed later, within each subgraph, during the schedule generation phase (see Section 3.2). The subgraph sizes are bounded by two main
8o (C)
Σ input output value of C
memary size of x
* of compute) O nodes C.
(1)
where is a hyperparameter whose default value is 0.5.
4

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Updating the candidates Once the best candidate C is chosen, it becomes group: all its nodes are considered together from now on. However, it is not yet a subgraph: if it is not too large, it can be merged with other groups later in the same phase to reduce the number of groups. The remaining candidates are updated: in each candidate C' whose intersection with C is not empty, we add all other vertices of C ensure that all vertices of C' are kept together. If this union contains more than M nodes, this candidate C" is no longer acceptable and is removed.
different solvers, which further improves the quality of the solutions. All solvers can be used any level, but those not acting hierarchically ignore the underlying partition. Each solver comes with an applicability criterion to decide whether it is suited for given graph. For each subgraph, all applicable solvers are used to produce schedules.
The HIREMATE framework currently includes two ad- ditional algorithms. First, H-TWREMAT is wrapper around the TW-REMA implementation (Shepperd, 2021) of heuristic based on treewidth decomposition ap- proach (Kumar et al., 2019). This wrapper and the HIREMATE framework allow this heuristic to be used with PyTorch where previously it was only available for Tensor- Flow. Second, RK-ROTOR is the dynamic programming algorithm from (Zhao et al., 2023), which also provides very good schedules in the case of (forward) graphs consisting of a sequence of possibly complex subgraphs. It is therefore limited in genericity, but its computational time is low. The RK-ROTOR solver is only activated when we detect that the graph can be decomposed into a sequence.
Section A in Appendix contains a proof that this procedure always leads to a valid decomposition, in the sense that no subgraph contains a directed cycle. After partitioning, HIREMATE identifies subgraphs that correspond to the exe- cution of the same piece of code avoid solving the same optimization problem multiple times.
3.2. H-Solver
3.2.1. SOLVING FRAMEWORK
The idea of hierarchical re-materialization is that re- materialization schedules be computed for each sub- graph independently, with multiple possible memory bud- gets. The resulting re- re-materialization schedules for given subgraph are called options, and each corresponds to a differ- ent tradeoff between computation time and required mem- ory. Once all subgraphs at a given level of hierarchy have been resolved, schedule for the upper level can be computed with our proposed H-II.P. H-II. is In- teger Linear Program (ILP) formulation that provides optimal schedule within given memory budget. It can be used on a subgraph of arbitrary structure, but using it on large subgraphs can lead to unreasonably long solution times. Therefore, it only applied to subgraphs with a sufficiently small number of nodes. This algorithm is in- spired by RK- CH-CHECKMATE (Zhao et al., 2023) and CHECK- MATE (Jain et al., 2020). The general idea of extending IL.P-based approach to a graph of arbitrary size and struc- ture, which is hierarchically decomposed into subgraphs of manageable size, is one of the major contributions of the present work and is described in detail in Section 3.2.2.
3.2.2. HIERARCHICAL ILP FORMULATION
Consider an arbitrary graph H. where each computation node represents a subgraph (Figure 2), and where dependen- cies are carried by data nodes representing values that can be stored in memory. The H-ILP formulation computes the minimum runtime schedule whose peak memory remains below a given memory budget. The computation nodes are numbered in topological order.
In the linear programming formulation of H-ILP, the sched- ule is divided into phases, and the goal of phase t is to compute node t for the first time. For the sake of brevity, we describe here only the main ideas that allow H-ILP to be used in a hierarchical setting, but refer the reader to the Section B of Appendix for a complete description of the H-ILP formulation.
Compute options and phantom nodes The innovation of H-ILP is that each compute node can represent a subgraph of the original graph. Such a compute node can be com- puted with one of several options. Each of these options represents a possible schedule for the forward and backward phases of the associated subgraph. There is a tight coupling between the forward computation and its corresponding backward computations, and each backward computation should be performed with the same option as its correspond- ing forward computation. The H-ILP formulation contains additional variables and constraints to specify which option is used for each computation node within each phase.
For significantly large graphs, partitioning with only two levels would result in the top-level graph still being too large to be solved with this technique. Our hierarchical approach makes it possible to introduce more than two levels into the hierarchy and compute schedules for the levels from bottom to top, handling graphs of arbitrary size while keeping each subproblem of manageable size.
The H-ILP solver we propose is very generic and efficient, and it provides very good quality solutions on all types of neural networks. Thanks to the genericity and extensibil- ity of the HIREMATE framework, it allows combining
In H-ILP, we also introduce an explicit representation of the data stored in memory between a forward computation and its corresponding backward computation. We call them
5

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
phantom nodes, and we update the formulation by consider- ing them as special data nodes, with two specificities. First, phantom node is always created by its forward computa- tion, always deleted after its backward computation, and is not used by any other computation node. In the formulation, we take advantage of this by not including additional vari- ables expressing whether the phantom node is deleted or not. Second, the values stored a phantom node (and thus the associated memory size) depend on the option used for the forward and backward computations. For this reason, the formulation includes additional variables that specify which option of each phantom node is in memory at each stage.
the GPU memory bound. For each computation node k, the number of these constraints for node k is bounded in the worst case by min(ks, 2 of mode )), where (x,o is the number of operations of the schedule associated to option o for node k. In practice, this number of constraints remains small enough so that the H-ILP formulation can be solved in reasonable time.
Schedule selection The number of binary variables in the H-ILP formulation depends linearly the total number of options of all nodes. To avoid wasting resources when several very similar options are available for a given node, we include in H-ILP hyperparameter No that imposes a limit on the total number of options. To stay within this limit, we may need to select only few options from all the schedules generated the lower level. To do so, we split No to assign number of options to each node proportional to the number of basic operations within the associated subgraph. Then, we greedily select the schedules whose memory peaks are farther apart from each other: starting with the schedule with the highest memory peak, then the one with the lowest memory peak, then the one closest to the middle, and so on.
Objective Let E equals 1 if node i is computed with option o during phase and 0 otherwise. Then the objective minimize the total execution time expressed
min Σ Ri, time of computing option for node i (2) 1,,00
subject to constraints on auxiliary variables specified in Section B.3 of Appendix.
phase
3.3. Complexity Analysis
phase
a compute operation i during phase t using option o
The complexity of using the H-ILP solver depends mostly on the number N of nodes in the graph. Obtaining the dataflow graph with RK-GB has complexity O(N) and is very fast in practice. With our current implementation, recursively partitioning the graph with H-partition has a complexity of O(N" log N). This step is also fast for graph sizes up to N = 1000, but handling very large graphs would require more work on the graph algorithms: recursively partitioning a graph of size N 10 takes 2 hours on an Intel Xeon Gold processor.
onm
phase?
phase
operation
Figure 3. Example of the solution. Ri. equals 1 if node i is com- puted with option o during phase f.
Correction terms for memory usage In (Zhao et al., 2023), to achieve good performance, the input values of a block are deleted as soon as they are no longer needed, instead of being kept until the end of the block. This is possible due to the sequential structure, where each block is the only user of its input values. In the more general context of this paper, a subgraph may share its input values with other subgraphs. To achieve good performance, H-ILP allows subschedules to delete their input values and the gradients associated to their outputs; however, if another subgraph needs these values, the higher-level algorithm adjusts the subschedules. This decision has an impact on the memory consumption of the corresponding subgraphs.
The most time-consuming step is the computation of re- materialization schedules with H-ILP. Thanks to our hierar- chical approach, this step actually has linear complexity. In fact, the hierarchical dccomposition has a logarithmic depth, and the total number of subgraphs is ONN M Solving a subgraph of size M for a number O of budget options only requires solving O Integer Linear Programs, whose sizes and solving times do not depend on N. Moreover, all ILPs at the same level are completely independent and can be run in parallel. As detailed in the Appendix, even without parallelization, our current implementation is able to han- dle the largest versions of neural networks used in practice: the hawwardbbccwaard graphs of GPT with 96 layers and Transformer with 36 encoders and decoders have 2500 nodes and are solved in 15 and 150 minutes, respectively.
We have added constraints to the formulation to ensure that the memory peak is correctly evaluated in all cases. These constraints are introduced only when solving the top-level graph, since they make the ILP harder to solve, and only the top-level graph solution is required to fulfill accurately
6

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
performance of ROCKMATE, CHECKMATE, MOCCASIN, and H-IIP. Since solving an ILP has exponential complex- ity, it is important to limit the size of the problem. A clear disadvantage of ROCKMATE is that it needs to see the graph as a sequence of blocks, which for general neural networks means that some blocks are very large, leading to unac- ceptable solving times for graphs that do not follow this structure. This is typically the case in the encoder-decoder transformer (Vaswani al., 2017), shown on the bottom of Figure 4. Specifically, ROCKMATE takes 15 hours to obtain a solution for a 4-layer encoder-decoder structure transformer, while H-ILP achieves better results within 12 minutes. Note that H-ILP is run once before the entire model training, so 12 minutes is clearly acceptable. By controlling the total number of nodes in each graph with the H-partition algorithm, H-ILP is able to efficiently limit the solving time without compromising solution quality. It is important to note that though MOCCASIN claims to beat CHECKMATE, it does not provide an optimal solution in general, because the number of rematerializations is bounded in practice (to 2, to limit complexity), and the time allocated to the constraint programming solver is also bounded (to 1h, because the time reach a solution can be arbitrarily long).
Overhead for Transformer
1.1-
1.2-
Algonthm
Checkmate
Hip
Meccast
Netaave
11-
Acckmatt
Number of Layers Dogscale)
Solving time for Transformer
1000
Algorilm
Dhecknats
10-
Rescain
Number of Layers (logstale)
Figure 4. Experiments for varying graph size: (top) iteration time
Performance Figure 5 shows that H-ILP consistently out- performs TW-REMAT in iteration time for a given bud- get. On the encoder-decoder (Figure 5(c)). TW-REMAT is able to find schedules for smaller memory budgets than H-ILP, but for UNet the situation is the reverse. The HIREMATE algorithm refers to the complete framework, including the RK-ROTOR and TW-REMAT solvers. By inte- grating TE-REMAT, HIREMATE overcomes this limitation of H-ILP and provides efficient solutions over the entire range of memory budgets. Figures 5(a) and 5(b) show that while H-II.P may perform slightly worse than ROCKMATE ATT, H-ILP achieves similar performance to ROCKMATE, a base- line approach specifically tailored for such architectures.
(bottom) solving time
4. Experimental Evaluation
Experimental settings The HIREMATE framework is de- signed to be used for end-to-end training. Since re- materialization guarantees that all computed values (includ- ing the gradients) are the same as without any recomputation, focus in our experiments on measuring peak memory usage and iteration time for a single training iteration. We perform a warm-up phase consisting of five initial runs; the subsequent ten runs are used to evaluate the peak memory and computation time, providing reliable estimates of per- formance. Measurements show that the standard error of the iteration time is at least two orders of magnitude smaller than the mean. As a result, error bars are not shown in NVIDIA
Overall, HIREMATE acts as general solution that includes H-ILP, ROTOR, ROCKMATE, and CHECKMATE. Our exper- imental results consistently show that HIREMATE performs on par with these baseline algorithms, demonstrating its ver- satility and effectiveness in optimizing memory utilization for a wide range of network architectures.
the plots. The experiments were performed Quadro RTX8000 GPU with 48 GB of memory and an NVIDIA V100 GPU with 16 GB of memory, using PyTorch 2.0.1, CUDA 11.6, and Gurobi 9.5.0. We intentionally report performance on settings where the experimental plat- form can run the original model, so that we can compare our results with the training time obtained with regular PyTorch Autodiff. All experiments can be scaled up by increasing image or batch size, to point where training requires using HIREMATF. Additional experiments including an ablation study, varying batch sizes and sequence lengths, as well as a dozen different architectures are available in the Appendix.
Controlling Sub-Optimality With our hierarchical ap- proach, the H-ILP solver does not compute an optimal so- lution to the re-materialization problem. We compare our results with RK- K-CHECKMAAE,, which provides an optimal solution for a given topological order of operations. The selected results in Table 2 show that H-ILP achieves perfor- mance close to this optimal solver on a network size where it remains tractable to compute. In another experiment, we investigate the impact of hierarchy depth in the H-partition by restricting the size of each subgraph, thereby increas-
Effficieccy In Figure we compare the solving time and
7

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Resnet101
Relative Peak Memory
Relative Peak Memory
6.
Itieemate mly)
1.25
E.Ramaa OB oolp
Redenate
-15
Itrratinn Time (m1))
X
Autiddif
MiRemmate
WWillemat
Www.eAAeee Arorage 11:0e
Autudddr
EiRoraate
1.15
I'arration Time
1.10
Relative
1.00
1.0
10
Peak Memory (GIB)
Pank Memory (GIB)
Encoder-Decoder Transformer
Unet
Relative Memorry
Relative Peak Memory
0.2
84
0.6
1.
1300
tly solo
2.4
MiRenate Of Upp only)
1.5
Hitemnte
30Remate
Demal
2.2
TW: TW-Stama
Autodilf
1100
-.0
(mo)
Averaape
Itsration Time
1.5
900
1.6
700
1.0
500
Memory (GiB)
Memory (GiD)
Figure 5. Experiments on different network architectures: (top) sequential-like neural networks (bottom) non-sequential neural networks, which ROCKMATE cannot solve efficiently.
ing the depth. As shown in Figure 6, the solution quality achieved by H-ILP shows only slight degradation as the hierarchy becomes deeper. This indicates that HIREMATE effectively preserves performance despite deeper partition- ing, It important note, however, that increasing the depth significantly affects the complexity and tractability of the problem as shown in Section 3.3, making this sult particularly notable: H-ILP manages to avoid quality degradation in settings where solving the problem becomes considerably more challenging. The HIREMATE framework allows the user to use several trade-off's between solution quality and solving time, by adjusting the partition gran- ularity, the number O of memory budget levels used in lower-level subgraphs, or the number N, of options retained when solving one subgraph.
Table 2. Comparison of H-ILP and RK-CHECKMATE overhead in terms iteration time
Network
Budget (GB)
Checkmate
H-ILP
UNet
7.1
5.32%
5.44%
2-layer Tformer
7.0
5.32%
5.34%
MLP-Mixer
6.0
2.47%
2.44%
5-layer GPT2
1.25
6.02%
9.24%
of overhead during training. comparable to those of the literature for problems of small size or with a particular graph structure, but without these limitations. On the prac- tical side, HIREMATE integrates seamlessly with PyTorch, improving the memory-time tradeoff and efficiency, and is fully compatible with PyTorch Autograd. The theoret- ical contributions include hierarchical approach and an original linear programming formulation H-ILP. Although HIREMATE focuses on computational graphs with primitive operations, the hierarchical approach with linear complexity is potentially applicable to other intensive tasks targeting
5. Discussion and Conclusion
This paper introduces the HIREMATE framework, which offers both theoretical and practical advances in re- materialization for PyTorch models. HIReMATE provides solution that can find very effective solutions in terms
8

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Bartan, B., Li, H., Teague, H., Lott, C., and Dilkina, B. Moccasin: Efficient tensor rematerialization for neural networks. In International Conference Machine Learn- PMLR, 2023.
Secoder-Decoder Transformer encoders/6 decoders)
410
12.0emate G-bew10
12itemate D-aewio)
12ieemate (5-a0ei0
ing. pp.
1aiaemate (-aaeio
evea Namory budgata
Beaumont, O., Byraud-Dubois, I.., Hermann, J., Joly, A., and Shilova, A. Rotor, 2019. URL inria. wt/hiepacs//aae rotor.
360
Iteration
270
360
Beaumont, O., Eyraud-Dubois, L., and Shilova, A. Efficient combination of rematerialization and offloading for train- ing dans. Advances in Neural Information Processing Systems, 34:23844 23225744-3387,, 2021.
19 Peak Memory (GIB)
Beaumont, O., Eyraud-Duboiss I.., Hermann, J., Joly, A., and Shilova, A. Optimal checkpointing for heterogeneous chains: how train deep neural networks with limited memory. ACM Transactions Mathematical Software (TOMS), 2024 (accepted for publication). URL https:
Figure 6. Experiments 6-layer encoder-decoder Transformer with different depth of hierarchical structure.
graphs resulting from tiling. The framework also incorpo- rates previous approaches, ensuring state-of-the-arr results. The remaining limitation of HIREMATE is that it not adapted to dynamic neural network architectures, where the the structure of the computational graph changes based on runtime inputs. However, HIREMATE be extended to support mildly dynamic graphs by unrolling fixed-length control flow, precomputing schedules for small number of known variants, or applying scheduling locally to sub- modules with static behavior. critical point is that the code is fully modular, and possible for external con- tributors to introduce a new graph partitioner solver at any level of the hierarchical decomposition. We hope that this modularity will stimulate research and further improve- ment of HIREMATE. Future research may include exploring how to integrate offloading into the optimization problem to reduce the need for recomputation, and how to optimally combine rematerialization with pipelined model parallelism. Advances these areas hold promise for improving the performance and efficiency of deep learning systems.
/ //arxiv...o///////// org/ /ttpg//ab//911111/11111 13214.
Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.00 1604.06141100000 77,, 2016.
Das, D., Avancha, S., Mudigere, D., Vaidynathan, K., Srid- haran, S., Kalamkar, D., Kaul, B., and Dubey, P. Dis- tributed deep learning using synchronous stochastic gra- dient descent. arXiv preprint arXiv: 162.:0/60220609 2016.
Hu, Z., Xiao, J., Deng, Z., Li, M., Zhang, K., Zhang, X., Meng, K., Sun, N., and Tan, G. Megtaichi: Dynamic tensor-based memory management optimization for dnn training. In Proceedings of the 36th ACM International Conference on Supercomputing, pp. 1-13, 2022.
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in Neural Information Process- ing Systems, pp. 103-112. 2019.
Acknowledgments
The research has been partially supported by the EUPEX (European Pilot for Exascale) project, which received fund- ing from the European High-Performant Computing Joint Undertaking (JU) under grant agreement No 101033975.
Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K., and Stoica, Checkmate: Breaking the memory wall with optimal tensor remateri- alization. Proceedings of Machine Learning and Systems, 2:497-511, 2020.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must specifically highlighted here.
Kirisame, M., Lyubomirsky, S., Haan, A., Brennan, J., He, M., Roesch, J., Chen, T., and Tatlock, Z. Dynamic ten- sor rematerialization. arXiv preprint arXiv:2006.094: 0961,, 2020.
Kumar, R., Purohit, M., Svitkina, Z., Vee, E.,, and Wang, J. Efficient rematerialization for deep networks. Advances Neural Information Processing Systems, 32, 2019.
References
XLA. http://ww.tennorrloo.orr/////

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Kusumoto, M., Inoue, T., Watanabe, G., Akiba, T., and Koyama, M. A graph theoretic framework of recompu- tation algorithms for memory-efficient backpropagation. Advances in Neural Information Processing Systems, 32, 2019.
Wen, G., Li, Z.., Azizzadenesheli, K., Anandkumar, A., and Benson, S. M. U-fno-an enhanced fourier neural operator-based deep-leaming model for multiphase flow. Advances in Water Resources, 163:104180 2022.
Wu, Z., Shen, C., and Van Den Hengel, A. Wider deeper: Revisiting the resnet model for visual recognition. Patlern Recognition, 90:1 19-133, 2019.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.0885 2020.
Zhang, S., Zhang, C.,, You, Z., Zheng. R., and Xu, Asyn- chronous stochastic gradient descent for dnn training, In 2013 IEEE International Conference on Acoustics, Speech and Signal Processing, pp. 6660-6663. IEEE, 2013.
Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, B., and Za- haria, M. PieaDream: generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Sym- posium Operating Systems Principles, pp. 1-15, 2019.
Zhao, X., Le Hellard, T., Eyraud-Dubois, L., Gusak, J., and Beaumont, O. Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re- Re-materialization in PyTorch. In ICML. 2023, Honolulu (HI), United States, July 2023. URL https: /hal..a.. scwence hhaeeneeeealeeeeee 04095305.9555.
Naumann, U. Call tree reversal is np-complete. In Advances in automatic differentiation, pp. 13-22. Springer, 2008.
Patil, G., Jain, P., Dutta, P., Stoica, I., and Gonzalez, Poet: Training neural networks on tiny devices with inte- grated rematerialization and paging. In International Con- ference Machine Learning, pp. 17573-17583, PMLR, 2022.
Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-no: U-shaped neural operators. arXiv preprint arXiv:2204.11127, 2022.
Rhu, M., Gimelshein, N., Clemons, J., Zulfiqar, A., and Keckler, S. W. vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design. In The 49th Annual IEEE/ACM International Symposium on Microarchitecture, pp. 18. IEEE Press, 2016.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Confer- ence, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234-241. Springer, 2015.
Shepperd, N. Fine tuning on custom datasets, 2021. URL https:// hwtp://ttttu...oo/oosaoeeeeeemme wwwpnshepperd/
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Wang, L.,, Ye, J., Zhao, Y., Wu, W., Li, A., Song. S. L.,, Xu, Z., and Kraska, T. Superncurons: Dynamic gpu memory management for training deep neural networks. SIGPLAN Not., 53(1):41 February 2018. ISSN 0362-1340, doi: 10.1 1 45/3200691.3178444
10

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
HIREMATE: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks
Appendix
A. H-Partition algorithm
Algorithm 2 H-Partition algorithm
1: Input: data-flow graph G
2: Result: recursive partition of G
3: Parameters: max high-level size M max lower-level size M score parameter a
4: while has more than M nodes do
5: C + UZEG candidate group containing all nodes between and a(x)
6: while C is non empty and G has more than MM nodes do
7: Select candidate C which minimizes Sa (eq.
8: Wrap the nodes of C into group
9: Update C
10: end while
11: Consider all groups as subgraphs
Update so that each subgraph is considered as a node
end while
14: return partitioned graph
Build small groups
Build small groups
Original graph
Hierarchical decomposition
Consider groups nodes
High-level graph
Figure 7. Visualization of recursive partitioning of the forward graph with 3 levels of hierarchy
In this section, we analyze the H-Partition algorithm, whose description is reproduced in Algorithm 2. We prove that the partition computed by this algorithm is always valid, in the sense that the resulting subgraphs do not contain any cycle. When merging a subgraph into a node for the higher level, all edges entering or exiting a vertex of the subgraph are attached the resulting node. To ensure that this does not result in cycle, we guarantee that of the subgraphs are convex the graph theoretic sense, defined below.
Convexity
We first provide some graph notations. Given two nodes and we write - b there a direct edge from and ~4 b if there is path of any length from a to b. Paths of length are also valid, so that a always true. With these notations, can define the convexity of subgraph:
Definition A.1. subgraph C' of graph convex if for any two elements a, b C, C' contains all nodes G on any

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
path from b. This be written as:
Va, b E C, G, u and - uEC
Merging a convex subgraph C' into node n does not create new cycles into the graph: such new cycle would be path starting at n and going back to n, going through another node u & C. If C is convex, any path from a node of C to another node of C' only goes through nodes of C, which ensures the absence of cycles.
Candidate groups The candidate groups C₂ formed on line 5 of Algorithm 2 contain all nodes on all paths from a(z) to a, where az)) is the common ancestor to all direct predecessors of They have the following property, where h(C2) = a(r): Property subset C' of nodes valid candidate group, and only if there exists a head h(C) such that:
If u C' and + u, If u E C,
then C or h(C)
(3)
then h(C²) ~~
(4)
This property ensure their convexity:
Lemma A.2. Any candidate group C which satisfy Property is convex.
Proof. Consider a and b in C, and u in G such that a ~ u and u -> b. There are two cases:
u +f h(C), then since b C' and + b, according to (3) we have u E C.
If u = hC),, then since a E C, by (4), we have ~ a. Since G' is acyclic, this implies = E C.
Update of candidates Once the best candidate C' is chosen, becomes group: all its nodes will be considered together from now on. The remaining candidates are updated: in any candidate C" whose intersection with C is nonempty, we add all the other nodes of C to ensure that all nodes of C' remain together. The following results show that the resulting set of nodes is still a valid candidate group: in particular it is also convex.
Lemma A.3. IfC and C' are valid candidate groups with Cnc' * 0, then h(C) ~ li(C) or I(C") ~y I(C).. Furthermore, (f((C) # h(C",, then the first case implies h(C") E C and the second case implies h(C) C².
Proof. Let E C C', and consider G such that - w. If no such exists, then is the source of G and h(C ((C).. If C'CC", we can start over with u
We now have E C n C", and o C C" with v + u. We have three cases:
If u C and v C²: from (3) applied C", have M(C") E C, and from (4) applied to C get h(C) h(C").
Symmetrically, if v & C and e C", we get h(C") h(C).
I C' a'dnd v & C': from (3) applied to both C and C", we get = h(C))())))))))) = h(C").
Theorem A.4. C and C' are valid candidate groups with C C' + 0, then D CU is a valid candidate group.
Proof. From Lemma A.3, we know that h(C)) ~ h((C") or ((C") ~ h(C)). We detine the head of D h(D) = ((C)) in the first case, and h(D)= h(C") otherwise. For simplicity, we assume in the following that h(C) ~4 h(C"); the other case is symmetrical.
It is clear that D satisfies (4): consider any u E D. If E C, then h(D) = -((C) ~y u by (4) applied to C. If u E C", then h(D) ~ (CC)) by assumption and ((C") ~ u by (4), so that both cases h(D) ~ u.
We now show that D satisfies (3). Let u. E 7 and - a with E D. We distinguish two cases:
12

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
u C", then since & C', by (3) applied to C" we get l((");; and since & C, the contrapositive of Lemma A.3 yields h(C')) h(C) h(D). Thus h(D).
u E C, then e & C' and (3) applied to C yield directly h(C) = h()).
This completes the validity proof of Algorithm 2: all candidate groups in C satisfy Property 1 all along the execution of the algorithm, both when they are created (line 5) and when they are updated (line 9). This implies that all subgraphs created line 11 are convex.
Identification of identical subgraphs To further improve efficiency of HIREMATE, we rely on and improve an idea from (Zhao et al., 2023): once the graph is partitioned, we identify all identical subgraphs that correspond to the execution of exactly the same code on values with the same shape'. A schedule computed for one of these identical subgraphs can be used for any of them, that significantly reduces the solving time on networks with a large number of identical blocks, such as GPT. This is performed in a more efficient way than in ROCKMATE, by expressing each graph in canonical way and by relying on a hash table.
B. H-ILP hierarchical formulation
In this section, we provide details on H-ILP formulation, the hierarchical re-materialization approach based on solving linear programming problem.
B.1. Context
We assume that we have an arbitrary graph H, where each compute node represents a subgraph. Similar to RK KK-CHECKMATE from (Zhao et al., 2023), dependencies are carried by data nodes, that represent values that can be saved memory. A value said to be alive at some time in a schedule if is stored in memory at that time. The memory usage at a given time in schedule is the sum of the memory sizes of all values alive at that time, and the peak memory of schedule is the largest memory usage over the length of the schedule. The H-ILP formulation computes the schedule with minimum rumning time whose peak memory remains below specified memory budget B. We denote by T the number of compute nodes, by 7 the number of data nodes. Compute nodes are numbered in topological order.
B. I. COMPUTE OPTIONS AND PHANTOM NODES
The novelty of H-ILP compared to EK-CHICCKAAE is that each compute node can represent a subgraph of the original graph. Such a compute node can be computed with one of several options. Each of these options represents a possible schedule for the forward and backward phases of the subgraph (Figure 8). There is strong link between the forward and the respective backward computations, and each backward computation should be performed with the same option as its corresponding forward computation. A pair of forward and the corresponding backward nodes are called a layer. For a layer/ j. its forward and backward compute nodes are denoted F, and B, respectively.
In H-ILP, we also introduce an explicit representation of the data saved in memory between a forward computation and its corresponding backward. We call them phantom nodes, and we update the formulation by considering them as special data nodes, with two specificities:
phantom node is always created by its forward computation, can only be deleted by its backward computation, and is not required by any other computing node. In the formulation, we can take advantage of this by not including additional variables expressing whether the phantom node is deleted or not.
values saved in an phantom node (and thus the associated memory size) depend on the option used for the forward and backward computations. For this reason, the formulation contains additional variables that specify which option of each phantom node exists in memory during each phase.
'This does not require to solve the difficult graph isomorphism problem, since we can order nodes according their execution in the original code.
13

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Forward
Backward
O Data
Forward
Backward
subgraph
subgraph
Loss, between forward backward
Figure 8. HIREMATE recursively finds schedules for each subgraph. Different schedules correspond to different memory budget constraints (options) and hence have different values for memory/time ration, peak memory, size of the saved data, and time for the backward computation.
In this formulation, we consider schedules which for a given phantom node, only one option is present in memory at a given time. However, it is possible that a phantom node produced several times with different options during the course of the schedule.
B.1.2. NOTE ABOUT INPUT DEPENDENCIES
Option-specifie dependencies An output value of the forward computation (i.e., data node which is computed during forward and used by another compute node) is never included in the phantom node. However, it happens that an output value is also used within the forward computation to produce other results. An example could be:
compuse (2):
Faal
In this example, the value x is both an output of the layer and used to produce y. that case, the backward schedule might choose either to use x as input be able to perform the backward g () (if having it in memory between forward and backward fits in the budget), to recompute it during backward. The implication is that for given layer, each option leads to specific dependencies for the backward compute node, depending which inputs is used by the corresponding schedule. If option of a computation node k depends on value d, we denote this as d S k.
Multiple predecessors data node can have several predecessors. This happens in backward when computing gradients: each computation is a contribution the same memory slot (gradients are accumulated), successor of such data node only be processed if all its contributions have been computed.
B.2. Formulation
The schedule is divided into T' phases. The goal of phase is compute node for the first time. In the following, we denote compute nodes with index k. data nodes with index d, options with index o, phases with index and layers (a pair of forward and corresponding backward nodes) with index j.
F is the set of final data nodes. The graph contains specific loss node which represents the computations that takes place between the forward and backward passes of our graph. If is the main highest-level graph, this represents the computations the loss for the training, if G is any subgraph (Figure 8), this also contains other computations from the rest of the graph. The index of the loss node
14

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Variables
The H-ILP formulation only contains binary variables, which take the value 0 or
R2 is if and only if node 1: computed with option o during phase t.
Pá is 1 if and only if data node d is present in memory before phase t.
S for k predecessor of d is if and only if the contribution of compute node l has been included data node d before phase t.
Spj.0 is and only if the phantom node of layer j is saved with option o before phase t.
G,a is if and only if data node d created when computing node k during phase
Da is 1 if and only if data node d is deleted after computing node k during phase t
Spt..: keep saved data of op. between stage - and stage with option
: compute operation during stage using option
ttme
operation
operation i
Figure 9.
Objective
The objective is to minimize the total running time, expressed as
min Σ R time of computing option for node
1,500
B.3. Constraints of H-ILP
Constraints for Options phantom nodes. Namely:
H-ILP formulation contains constraints relative the choice of options and the management of
Vt, Vk,
Σ
Rh. 1
at most one option for each computation
(5)
vt, Σ Spi 1
only option of phantom node is in memory (6)
Vt. Vo,
Sp < Sp + R/,,0
phantom node is only be created by its F,
(7)
Vt, Vo,
Sp 2 Sp + R4,,0
R R
phantom node is only deleted by its B, computing B, requires the phantom node
(8)
Vt, Vo,
Ry S Sp/.0 + R2p
(9)
For any compute node & and any phase t, we denote by CR = Σ. Rt the equivalent of the R variable of RK CHECKMATE (Zhao et al., 2023), which equal if node computed during phase (with any option).
15

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Validity constraints
Vt, vk > t, CR = 0
Node k: >> t can not be computed in phase t
14, s.t. F3 > t, Vo, Spj.0 0
No phantom j from node Fj > t is saved before phase t
Ak- + d, k S 0
No result of node lc can be saved in any phase before phase is Data node is not in memory before any of its predecessors
AA < minnk//k- + d],
Pa =0 0 AA E F, 44 + d, Si + CR =1
After the last phase, all final data nodes should be in memory
Vt, CR =1 Σ CR =1
Node t is executed in phase
The loss node is executed only once
Data dependencies
Vt, vk + d Sk S P'J
Data node d with at least one contribution l: is alive
V1 < vk - d. Skd S Sh + CR
New contribution k to node d only appears by being computed Computing node k' requires all contributions k: to input node d Option-specific dependencies
1I, vk + d + K, CR2 S CR + Skid
14
14 SA Vo, Ak d 2 Bj R S CR + Sh.d
Alive status of values A computing node k is related to a data node d if k - d or d + k. We denote this with k <> d. For a data node d, only computing nodes & that are related to d can affect its alive status. For any t, if k related to d, we denote with A the alive status of node d during phase after computing node lo (and also after performing all deletions mandated by variables D). In phase i after computing node data node is alive if stored before phase created in phase before node k and not deleted until then, so that we can write:
vt, + d, Ak.a = Pt
Σ k's
CKa-
D's keedkkkk
Σ
In the above equation and in the following, we use = denote an alias definition, so that A dd can be replaced by the right-hand side in any constraint, whereas the = sign is used to denote constraint that is added to the formulation.
Constraints relative to liveness
Vt, vk: + d, 0 S A $ 1 Vt, Vk dd, Ald 2 CR² D
Data node d is either alive not d is alive if computed and not deleted value can only be created by a node & that is really computed Value d is alive after phase t iff it is alive after its last related node k
VL, Vk: +d d, Cha= < C.R
Vt < T. vd,
Plt A (
One additional constraint states that a value d is deleted after computing node k in phase t if it is not used afterwards: neither by later computing nodes k' > 1 in the same phase I, nor in the next phase /. 4 1. This can be stated as:
Vt, vk + d, D'a= = if and only if CR and p 0 and
E d=
CR =0
However, this constraint not linear. can be linearized in the similar way as in the original CHECKMATE paper (Jain
al., 2020, Section 4.5): if we denote by hg.d = 2 + [[ + k, K, k}] the number of equalities in the above statement, it is equivalent to:
Vt, Vk ++ d,
D'.d CR P
Σ
CR
d bokkRS
Vt, Vk + d. h1 D(.a) 2 CR + Pt
Σ
CR
doRkk
16

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Evaluating memory usage In contrast RK-CHECKMATE, the memory budget constraints of H-ILP are updated to take into account the memory sizes of the phantom nodes where appropriate. For this purpose, we denote by U/ the memory usage after computing node k in phase t. For any value d, let sa be the amount of memory required to store d: and for any layer/ and option o, let Pj.o be the amount of memory required to store option o of the phantom node of layer j.
Then U can be expressed as a linear combination of the formulation variables. Indeed, the memory usage before starting phase is:
M =
Σ 8d Pd + Σ Pio Sp
Then, the increment when computing a forward node k Fj is:
IF = Σ 88 Ok,d Σ 88 D + Pisoo
k-dd
k+dd
When computing a backward node k Bj. the increment is:
IBx =
Σ Sd Ck k-d
Σ sa Sa.Dii - k++d
-Epjo Ps.o o (Rt Spj.o -Sp
The expression within the parenthesis is equal to 1 if the allocated node j is deleted, and 0 otherwise. Indeed, since By is the only compute node that can use it, Spfr0 [+1
means that phantom node j can be deleted right after Bj. Constraint (8)
ensures that if R 0, then the expression within the parenthesis is also O.
Finally, we can express Ut iteratively (similar CHECKMATE and RK- CHECKMATE formulations):
Vt, Up = M, + TFo
VK = FJ, U% = U 1FR
AI, = Bj, Uh U + IB
Memory budget constraints Thanks the UL definitions, we can express constraints ensure that the memory usage always within the memory budget B. If we detail single step & of some phase t. it corresponds to: (a) allocating memory for the newly created values (according to Ch, variables), (b) computing node k, (c) freeing the memory of the deleted values (according variables aa.)). The highest memory usage in this sequence is during (b), but the memory usage Uj corresponds to after (c). In addition, the computation of node k with some option might incur memory overhead (by allocating temporary values), which we denote by miço- In total, in RK- AK-CHECKMATE, the memory budget constraints are written as:
Vt,ak, Uk +
Σ m R +
sa Sa.Di D k.d S B > ketd
(10)
However, in H-ILP each compute node l represents not a single basic computation, but a sequence of basic computations (defined by the corresponding schedule). After the H-IIP formulation is solved, the actual schedule is modified: the memory deallocations for the values freed in step (c) above are performed as early possible, possibly during the schedule of node I (in the middle of step b).
This means that the memory overhead ma,o during the computation of option o of node k might depend on whether some values are alive before or after computing node k. For example, the corresponding schedule deletes a value in the middle of computation, its memory overhead mako assumes that the deletion is delayed until the end of the schedule. that value actually not needed later in the higher-level schedule computed by H-ILP, it will be deleted within the schedule, which may or may not change the memory overhead.
In the following, we present how we modify the memory budget constraint to account for this kind of situation. Consider a specific phase f, and an option (and thus a schedule) for node k. For simplicity of presentation, let us consider only inputs; the situation with outputs is similar and symmetric. Consider a sub-step i of the schedule. We compute the memory overhead at this sub-step as m assuming that value deletion happens after the computation this schedule of node k. We denote by F; the set of values which are not used in the following sub-steps of that schedule, Within the schedule computed
17

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
by H-ILP, values in Fg are deleted after sub-step if and only if they are deleted after step kt Hence, the actual memory usage of sub-step is m Eder, Sa" D k The corresponding memory constraint is:
Vk, Vi, U + m R +
Σ ktdd
88 Didd
Σ dEF.
88'
D S B
We write one constraint for each option and each sub-step. Since all the correction terms negative, and all mggo are least 0, if R is O, this constraint is weaker than (10). We write such constraint for each sub-step of the schedule, and this provides a more precise assessment of the memory usage of the solution. The case of output values is the same, except that we care whether the output value is created during the computation of node k which is represented with variable ktdt
An interesting remark is that it is not necessary to write one constraint for each sub-step: the set of inputs not needed after sub-steps and j are the same (F; = Fj), we can keep only one of both constraints (the one with the larger memory usage ). The number of constraints is thus bounded by min (numbee of sub-steps, 2 (iipuus) | [ouputs)). ). In practice, the number of different constraints remain low enough. In addition, these constraints only introduced when solving the top-level graph, where the constraint to remain under budget B is required to be accurate possible.
C. Implementation details of the HIREMATE framework
We provide here some implementation details about our framework, and in particular on the differences compared to ROCKMATE. addition to the partitioning and solving steps, described in detail above, there are two other steps in HIREMATE: the graph building procedure, and the execution process.
C.1. Graph builder
In HIREMATE, we use the same RK-GB module as in ROCKMATE to obtain a data-flow graph from an arbitrary Py PyTorch nn Module. This module has four steps:
1. obtain the forward basic graph with torch. jit, which each node represents exactly one computational operation
2. obtain a simplified graph in which all operation which not produce a new Tensor are merged with the operation that produced the involved Tensor. Typically, all view and size calls, and all in-place operations have this behavior.
3. building the backward graph, and
4. measuring the time and memory requirements of each operation.
In ROCKMATE, there is another step where the simplified forward graph decomposed into a sequence of blocks (so that each block can be solved with RK- CHECKMATE, and the sequence can be solved with RK-ROTOR). Once all blocks are identified, ROCKMATE performs pairwise comparisons identify the blocks with similar structure, so that EK-CHICCMAAEE is only applied once for all identical blocks.
In EIAE, the H-Partition step is replacement for this "sequence-be buueencebuilding" step. In the following, we provide technical details about three parts of RK-GB which were adapted (and sometimes improved) for the HIREMATE framework: the identification of identical subgraphs, the management of "soft" dependencies, and the dependencies on the inputs of the model.
C.I.1. EFFICIENT IDENTICAL GRAPHS RECOGNITION
Instead of pairwise comparison is done in ROCKMATE, in the HIREMATE framework the identical subgraphs are identified by using canonical representation and hash table. As reminder, in both cases, identifying identical graphs is made possible by the fact that we can rely on the order of execution of the code; and the main use case of this feature is to identify when the same code is executed several times in the model. RK-GB is designed to be fully deterministic, so that it provides the same result on the same executed code.
We provide now a description of the identical subgraph recognition in HIREMATE:
18

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
We first identify identical compute nodes from the original (not partitioned) graph. Two nodes identical if they have the same code (modulo variable renaming) and if all inputs and parameters have the same characteristics (shape, type, ...). Each compute node is thus associated with equivalence class.
Subgraphs are determined by the list of nodes they contain. To obtain the canonical representation of a subgraph, we go through this list of nodes, in the topological ordering inferred from the code execution. Each node is represented by its equivalence class, and by the indices (in the subgraph) of their predecessor nodes. The canonical representation of the subgraph is the concatenation of the representation of all its nodes.
The obtained representation is inserted in a hash table: if subgraph with the same representation exists, they will be identified as identical. Only one of these identical subgraphs solved within H-Solver; it is called the representative.
HIREMATE also features a "translation" function that renames variables appropriately convert a schedule for the representative into schedule for each of the subgraphs in the class.
This procedure can be performed in linear time, and is thus very efficient in practice.
C. SOFT DEPENDENCIES
Consider situation where a computation node k does not depend on the Tensor z computed by previous node ks but it depends on the characteristics x (typically, the result of a s ze (x) call). In the simplification process, the size () call merged within the node k. But having node k' depend on node would result in undesired behavior a situation where has been computed once and removed from memory: would recomputed to enable computing k', whereas k' only requires the size of k.
The solution introduced by ROCKMATE is to consider this particular situation soft dependency during the simplification process. These soft dependencies are used to produce consistent topological orderings. After the simplification process, for solving the blocks with RK- SE-CCHECKMAT, all soft dependencies are removed from the graph. The dependency of k' on the size of x is ensured by the topological ordering: the first computation of k will take place before any computation of k'. Since the data carried by a soft dependency is not a Tensor, it does not occupy any memory space, and is never removed.
Soft dependencies with H-p H-partition However, in HIREMATE, the H-Partition algorithm has take the soft dependencies into account for the convexity considerations. Without the soft dependency from i to K', the predecessor j of l might be placed in the same subgraph as K,, but without k. This subgraph would be executed before the subgraph of k (since - R), which would break the soft dependency. For this reason, the soft dependencies are retained all through the H-Partition algorithm, and ignored only during the H-Solver part.
C. DEPENDENCIES ON INPUTS OF THE MODEL
The Tensor values provided as input to the model have to remain in memory throughout the whole execution, for any possible schedule. It is not useful to take into account their memory sizes, since they cannot be managed by HIREMATE. In addition, if such an input value has requi res_ses____ res__gad=Faaa 1 se, then no backward computation is associated either. In HIREMATE, unlike ROCKMATE, the corresponding dependency is ignored. This enables HIREMATE to have a more efficient discovery of sequential sub-parts of the model.
Consider for example the torch. nn. Trans former r model, with 6 encoder and 6 decoder layers. This model has two inputs, src used at the first encoder layer, and tgt used at the first decoder layer. Since ROCKMATE takes into account the dependency carried by the tgt value, it is unable to identify the sequential structure of the encoding part: there is a dependency from the input node the first node of the decoder. As a result, ROCKMATE sees this model as one large block of 219 computing nodes.
In contrast, with this change in the HIREMATE framework, the same model can be decomposed into 24 blocks, the biggest of which contains 131 nodes (corresponding to the decoder part). This is still too large to be solved directly with RK-CHEC CHECKKAAT,, but this change enables HIREMATE to run the ROCKMATE algorithm on a nn. Trans fo former with 4 layers. Beyond the benefit for ROCKMATE, having fewer dependencies also provides more opportunities for partitioning in the H-Partition algorithm.
19

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
C.2. Execution of the schedule
Once a schedule has been computed by the H-Solver, the HIREMATE framework generates a new nn Module which follows that execution schedule. For performance and compatibility reasons, this part of HIREMATE is completely new and not based on the corresponding ROCKMATE package RK-EXEC
C.2.1. THE EXECUTION PROCEDURE
The RK-EXEC package converts the complete schedule produced by RK-ROTOR into single very large st r nng of Python commands, which is then executed with the exec () built-in function. This has significant complexity in the code, and results significant overhead the running time. For example, running the nn. Trans former model with this approach results in almost doubling the execution time.
We different approach in HIREMATE, where the schedule is interpreted the fly and each operation is performed in controlled environment. This enables HIREMATE efficiently execute schedules for large models, as demonstrated in the experimental section.
C.2.2. AUTOGRAD COMPATIBILITY
The result of ROCKMATE does not respect the PyTorch convention for r.n n.. Module: the backward phase must be explicitly triggered, and the produced module can not be safely re-used with different inputs. In addition, ROCKMATE schedules have the possibility delete the output Tensor during the backward pass, which might break the code of the user if they need to re-use it afterwards.
This possibility is removed in HIREMATE, and the implementation provided by HIREMATE is completely compatible with the autograd mechanism of PyTorch. For example, the following code works as expected:
rone:Mor
BRookmate e (node), budget
Lmp, Lig-
y
Linear (inp)
x1
remat Mod
22
semaLMod (x)
loes 1osx__urrtt Funnt (ry Egt))
bussbbackwe:d))
20

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Incader-Derodde Transformer
UNet
400
Iphha
alpha
440
0.2
470
0.5
0.3
0.5
0.5
Reratton Tire
460
0.6
0.5
400
250
0.9
440
430
360
6000
7000
8000
9000
10000
11000
12000
13000
6000
7000
8000
9000
10000
11000
12000
Memory
Memory
(a) Encoder Decoder Transformer, batch size 64
Encoder-Decoder
(b) UNet, batch size 64
Figure 10. Experiments with H-ILP for different graph partitioning defined hyper- paramettr
D. Ablation Study
Through this ablation study, aim shed light on the key factors influencing the performance of HIREMATE and highlight its superiority over existing approaches in terms of peak memory training time trade-off across different batch sizes and input data resolutions. The findings of this study contribute to deeper understanding of the underlying mechanisms and efficacy of HIREMATE, further solidifying its position as a leading solution in the field of re-materialization
D.I. HIREMATE typer-parameters
In this section, we conduct an ablation study to examine the influence on the performance of HIREMATE of the hyper- parameter, which impacts the quality of the partitioning. By systematically varying the value of a, we aim to gain insights into its effects on the overall performance of HIREMATE. This analysis provides valuable information for optimizing the hyper-parametes configuration achieve the best possible results in different scenarios.
On Figure 10, observe that the parameter has small but measurable impact the quality of the solution returned by HIREMATE. However, the value which provide the most efficient solution depends the graph of the model, and we can not conclude on optimal value for a. In the rest of the plots in this paper, we use the default = 0.5 which provide reasonable results in most cases.
D.2. Performance for Different Input Sizes
Furthermore, we show performances of our module HIREMATE, across different batch sizes and input data resolutions. We focus on the Transformer and UNet models, and we vary the batch sizes and the input size (in terms of sequence length for Transformer, and in terms of image resolution for UNet). Figure reports results for sequence length 200 and varying batch sizes. Figure 12 reports results for batch size 64 and varying sequence lengths. Figure reports results for fixed resolution 256x256, and varying batch sizes. Figure 14 reports results for fixed batch size 64 and varying resolutions. In addition, also provide on Figure 15 the results for the MLPMixer model with batch size 64 and resolution 256x256. The behavior of HIREMATE is consistent across all cases, and significantly outperforms TW-Remat
21

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Seeooder-Decode Transformer Reiative Bosk Momicy
Aaccode--Decoder Transormer Beletive Pook Mamory
Encoder Decoder Transformer Relariv Peok Memory
12.0
lalaaaee Sant Jon #ondd
libeue bech de 114 abiCE
1200
Itweadom Sirw Tmaa)
Iucuton Sitoo Itea)
1000
Tocotan
160
125
12
39
Memory
Memory
Mammory
(a) Batch size 64
(b) Batch size 128
(c) Batch size 256
Figure Experiments Encoder-Decoder Transformer with diffcrent batch sizes
Encoder-Decoder Transformer Seladive Peak Memory
Enonder-Decader Transformer
Encomder-eecdde Transformer Marmry
Llleuttereeeee 158 atsee
1104
Tterstton Time
Iaeesion Tima
30t
Poak Memory JGBBI
Posk Memory (GL3)
Poal Memooy (CIB)
(a) Sequence length 100
(b) Sequence length 200
(c) Sequence length 400
Figure 12. Experiments Encoder- Transformer with different sequence lengths
Unet Relativo Poak Menoty
Unet Relative Peak Memary
Unet
Memory
Edtoc
1.00
Tallumal lacc use 228
Deralion Tiare (mxd
Buration Ttaso (taa))
Teration
D80
Poak Mumory (GIB)
Poak Menory (GL3)
Posk Mertoyy (CIB)
(a) Batch size 32
(b) Batch size 64
(c) Batch size 128
Figure 13. Experiments UNet with different batch sizes and resolution 256x256
22

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Rolatice
C4
commlriell 123
Htrrtton Time
Hemation Tree
Paak (RMI)
Pwst Mettery (GUI
(a) Resolution 128x 128
(b) Resolution 256x256
Figure 14, Experiments UNet with batch size 64 and different resolutions
MIP Mixer
Relative Peak Memory
0.2
0.6
HiRanate Enlly
TW-Remat
320
X
Asto
Relative Average Time
Iteration Time (tn))
3.0
1.44
2.00
1.2
240
1.2
-111
200
Peak Memory (GIB)
Figure 15. Experiments MLPMixer
23

---

HIREMATE: Hierarchical Approach for Efficient Re-materialization of Neural Networks
Network ResNer34
Budget
Checkmate
H-ILP
5.0e9
1.0138
1.0184
4.0e9
1.0360
1.0394
3.5e9
1.0517
1.0554
3.0e9
1.0712
1.0761
UNet
9.Ic9
1.0090
1.0090
8.1e9
1.0293
1.0317
7.1c9
1.0532
1.0544
6.1e9
1.0776
1.0830
2-layer Tformer
9.0c9
1.0195
1.0194
8.0e9
1.0355
1.0356
7.0e9
1.0532
1.0534
6.0e9
1.0700
1.0702
MLP-Mixer
8.0e9
1.0050
1.0044
7.0e9
1.0137
1.0137
6.0e9
1.0247
1.0244
5.0c9
1.0466
1.0545
5-layer GPT2
1.4c9
1.0113
1.0117
1.25e9
1.0602
1.0924
1.1e9
1.1515
1.1929
1.0e9
1,2593
Table 3. Comparison of H-II. and the optimal solution of EK-CHECKMATE: overhead in terms of iteration time compared to Autodiff for different models different budget values
E. Comparison with optimal solvers
We evaluate the quality of the solution returned by H-ILP by comparing with the optimal solver CK-CHECKMATE small graphs, where the optimal solver is tractable to compute. We provide in Table 3 broad comparison over several different models and different budget values. We observe on a wide variety of networks and budgets that the quality of the solution produced by ILP remains very close to the optimal value provided by RK- CHECKMATE when possible to compute
24

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
F. Broad study variety of models
We assess robustness our approach by analyzing its behavior on a variety of models. Table 4 provides results on GPT-medium and GaT-largest (which have 24 and 96 layers), UNet, MLPMixer, RegNet32, ResNet101, Transformer with encoder and decoder layers, Transformer with 36 encoder and decoder layers, ID and 3D FNO (Li et al., 2020) models, UFNO (Wen et al., 2022) and UNO (Rahman et al., 2022).
Table 4(a) describes the computational graphs of each model:
the number of nodes the computational graph before simplification:
the number of nodes in the forward computational graph, after simplification, which is the input to the H-Partition algorithm:
the number of nodes in the forward-backe bawkwrd computational graph;
the total number of computation and data nodes in the final computational graph, which is the input to the H-Solver algorithm.
Table 4(b) describes the result of the H-Partition algorithm:
the depth of the hierarchical decomposition. The top-level count the first level, so CHECKMATE, which works directly takes the whole graph, corresponds to depth I, and ROCKMATE, that takes chain of blocks, corresponds to depth of
the number of subgraphs:
the number of unique subgraphs after identifying identical subgraphs;
the size of the largest subgraph (in terms of number of nodes).
Table 4(c) describes the result of HIREMATE on each model:
the total execution time of the HIREMATE framework (including building the graph, partitioning and solving);
the budget provided HIREMATE;
the relative memory usage (compared to the peak memory of the autodi ff solution) of the resulting nn. on.Mo Modu le
created by HIREMATE, as measured in real execution;
the "predicted" value of the rumning time of the computed schedule, which is the sum of the measured execution times of all operations in the schedule, expressed relatively to the execution time of the autodi ff solution;
the execution time measured from running the nn Moouue created by HIREMATE, also expressed relatively the execution time of the autodife ff solution.
can see that HIREMATE obtains robust results on all these very different models: able to reduce memory usage by a factor 2, for cost in execution time that varies between 7% (for RegNet32) and 22% (for UNO). All graphs except for UNet, FNO 1d and FNO 3d are much too large for CHECKMATE solve. The Transformer, U-FNO and UNO models that ROCKMATE also fails these graphs. In contrast, the solving time of
not be meaningfully sequentialized, HIREMATE remains below one hour for most graphs, and below three hours for the largest Transformer network, which is totally acceptable.
25

---

HIREMATE: Hicrarchical Approach for Efficient Re-materialization of Neural Networks
Table Test HIREMATE over models, with a memory budget around half the autograd memory usage. models passed sanity check: both forward and backward passes produce the exact same result as the original module. Experiments are done NVIDIA P100 GPU with 16GB. See column headers and model definitions on previous page.
(a) Size of the computation graphs
SIZE BEFORE
SIZE FWD
SIZE FWD
SIZE COMPUTATION
SIMPLIFICATION
NODES
GPT 24 (MEDIUM)
1858
367
734
1614
GPT 96 (LARGEST)
7402
1447
2894
6366
UNET
73
50
101
205
MLPMIXER REGNET32 RESNETIO1
203
125
250
549
245
173
347
721
346
211
123
851
TRANSFORMER 6-6
1030
221
117
991
TRANSFORMER 36-36
6160
1331
2457
5851
FNO ID
42
82
167
FNO 3D
317
64
121
257
U-FNO
567
118
221
467
UNO
997
129
229
495
(b) Results H-partition
NUMBER OF LEVELS DEPTH
NUMBER OF SUBGRAPHS
UNIQUE SUBGRAPHS
LARGEST
SUBGRAPH
GPT 24 (MEDIUM) GPT (LARGEST) UNET MLPMIXER REGNET32 RESNETIO1 TRANSFORMER 6-6
28
8
15
9
9
15
15
6
10
15
10
15
TRANSFORMER 36-36
16
128
67
19
FNO ID
10
5
8
FNO 3D
12
12
10
U-FNO
UNO
12
12
19
(c) Solving time and performances
TOTAL
MEMORY
OBSERVED
PREDICTED RELATIVE AVERAGE TIME
OBSERVED
SOLVING TIME
BUDGET
RELATIVE
RELATIVE
(MIN)
(GB)
PEAK MEMORY
AVERAGE TIME
GPT 24 (MEDIUM)
7.9
1.7
36.6%
122.3%
120.3%
GPT 96 (LARGEST)
4.5
47.2%
119.6%
126.0%
UNET
5
51.3%
113.5%
113.3%
MLPMIXER
1.5
4.5
53.3% %
107_3%
109.0%
REGNET32
5
3.5
41.1% %
107.0%
105.8%
R0SNET101
5.5
4.5
42.6%
115.3%
115.3%
TRANSFORMER TRANSFORMER 36-36
3.5
46.9%
117.6%
114.8
147
4.9
44.6%
118.4%
FNO ID
2.5
5
53.1%
114.7%
111.0%
FNO 3D
5.3
4
47.2%
117.4%
104. %
U-FNO
16.3
5.2
56.7%
116.3%
116.5%
UNO
5.2
58.5%
121.9%
116.7%
26


Reviews:
[{"review_id": "zEm8cYP0lH", "reviewer": "Reviewer_ag6k", "summary": "This paper develops a procedure for hierarchically partitioning a data flow graph in order to find a schedule for selectively recomputing neural network activations so that memory usage during backpropagation fits within a specified budget.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful reading and detailed feedback. Below we address a question regarding gradient accumulation  comparison.\n\nGradient accumulation reduces memory usage by splitting a large batch into smaller sub-batches and accumulating gradients across them. This strategy avoids recomputation but comes with trade-offs that HiRemate addresses differently:\n\n**When batch size 1 exceeds memory** (e.g., high-resolution inputs or very deep models), gradient accumulation does not help. HiRemate enables training in these settings by recomputing activations selectively.\n\n**Low GPU utilization with small batches**: Accumulation processes sub-batches sequentially, often leaving GPU cores underutilized. HiRemate allows for larger effective batches under the same memory budget, improving throughput.\n\nIncreased GPU utilization with HiRemate is also worth applying in the case of network architectures whose computation is not dominated by large matrix multiplications, or more generally whose main computational operations are not well optimized for the hardware.\n\nWe agree that an empirical comparison with gradient accumulation would strengthen the practical case for HiRemate, and we plan to include such experiments in future work.", "claims_and_evidence": "The proposed method is a modular framework that delivers practical benefits over recent competing approaches.", "methods_and_evaluation_criteria": "Experiments measure peak memory usage and iteration time during training for HiRemate in comparison to recent related work.  Across a variety of architectures and batch sizes, HiRemate appears to consistently outperform (faster iteration time at any fixed peak memory) TW-Remate and Rockmate.", "theoretical_claims": "-", "experimental_designs_or_analyses": "See methods and evaluation above.", "relation_to_broader_scientific_literature": "Appears to be a consistent, but incremental improvement over recent work.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "With HiRemate in use, scaling batch size appears to (approximately) linearly increase both iteration time and peak memory usage (e.g., as in Figures 11 and 12).  Within a limited regime, this same behavior would be expected of a naive strategy that simply reduces batch size in order to fit within a memory budget; i.e., to achieve the same effect of one large batch, two smaller batches could be run in series, accumulating gradients for both before updating parameters.  Such a strategy would not pay any penalty for recomputation, but effectiveness would depend on whether the smaller batches fill the available parallel computational resources.  How does this naive strategy compare to HiRemate?  Specifically, at what point (what batch size) does HiRemate become more efficient than treating a larger batch as multiple smaller batches?Understanding the trade-off in comparison to this naive strategy seems crucial to making a practical argument for using HiRemate.", "other_comments_or_suggestions": "See strengths and weaknesses above.", "questions_for_authors": "See strengths and weaknesses above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "US16ff5QMD", "reviewer": "Reviewer_YaUT", "summary": "The paper presents a novel hierarchical framework to optimize memory usage during neural network training. They recursively partitions large computation graphs and apply optimized solvers at multiple levels to significantly reduces memory consumption while maintaining efficient training performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful reading and detailed feedback. Below we address each of point of the review in turn.\n\n**On Solver Overhead and Scalability**\nIndeed, H-ILP introduces a solving time overhead. However, this can be mitigated in the following ways:\n* Solving separate subgraphs can be performed independently and in parallel, providing efficient parallelization (moreover for subgraphs of the same structure we run the solver optimization once).\n* The complexity analysis of Section 3.3 shows that the total solving time of all ILP optimizations is expected to grow linearly with the number of nodes in the graph, ensuring a reasonable scaling of the framework.\n* For long training runs (e.g., months/weeks), even several hours of preprocessing is negligible.\n* For cases where fast solving is required, faster heuristics (like TW-Remat) or Dynamic Programming solvers can replace ILP optimization subproblems, keeping ILP only for the several top-level optimization problems. Configuration options to run with different solvers are already implemented in HiRemate, and its modular design is straightforward to extend with more custom solvers and partitioning strategies.\n\n\n**On Limitation to Static Computation Graphs**\nHiRemate is designed for neural networks whose forward-backward computation graph remains constant across input samples. This includes many widely used architectures, including such as GPTs, Encoder-Decoder Transformers, ResNets, UNets, ... For these networks, HiRemate extracts the full computation graph once before training and generates a schedule that meets memory constraints by selectively recomputing intermediate activations rather than storing them. The resulting schedule is reused throughout training.\n\nThis static-graph assumption covers a wide class of models used in practice. For architectures where the graph structure can change depending on runtime input, HiRemate can still be applied or extended in many useful cases:\n\n- *RNN-style architectures*: These apply the same computation block multiple times. If the number of steps is fixed, the computation can be unrolled into a static graph, which HiRemate can optimize directly. Thanks to the repeated sequential structure, the higher levels of HiRemate’s hierarchy can use a faster dynamic programming solver instead of the ILP-based one. If the number of steps varies, multiple schedules can be precomputed for different lengths and selected at runtime. In all cases, HiRemate identifies and reuses schedules for similar subgraphs, which reduces redundant computation during schedule generation.\n- *SSM-based models (e.g., S4 or Mamba like)*: These models use state-space blocks that in most cases can be expressed in either a recurrent or convolutional form (e.g., S4D, Mamba2). The convolutional version, typically used during training due to its better compatibility with modern hardware, results in a static computation graph that HiRemate can process directly. In particular, we checked that the current version of our framework is compatible with S4D from `state-spaces/s4` and Mamba2Easy from `state-spaces/mamba2`. If the recurrent form is used instead, the situation is similar to RNNs, and the same adaptations apply. \n- *Architectures with limited graph variation*: Some models include conditional branches or mode switches that result in a few possible computation graph variants. If these variations are known in advance and limited in number, HiRemate can precompute a separate schedule for each case. At runtime, the appropriate schedule is selected depending on the input or control signal.\n- *Architectures with highly dynamic graphs*: A class of models changing the graph structure at runtime in many complex ways. These models fall outside the scope of HiRemate’s current design. In such cases, dynamic heuristics  should be explored, their combination with static approaches is an interesting direction for further research.\n\nEven for models with highly dynamic structure, it is often possible to apply HiRemate locally, at the level of individual  `nn.Module` blocks whose behavior is static. This enables good block-wise memory/time trade-offs and can still lead to meaningful savings. HiRemate’s modular design supports partial integration into larger training pipelines, allowing users to optimize memory usage in parts of a model where static graphs are available. In particular, this is relevant for many kinds of mixture-of-experts architectures.\n\nIn summary, HiRemate automatically generates memory-efficient training schedules for static computation graphs and can be adapted to models with repeated or mildly varying structure. It handles large graphs through hierarchical partitioning, supports usage of diffrent solvers across subgraphs, and integrates well with PyTorch without requiring changes to model code. These properties make it practical for many real-world training pipelines and adaptable to a wide range of architecture types.", "claims_and_evidence": "The claims are well-supported by experimental results. The paper provides empirical evidence demonstrating significant memory savings with minimal computational overhead across multiple neural network architectures.", "methods_and_evaluation_criteria": "The proposed HiRemate framework and its evaluation criteria are well-aligned with the problem of optimizing memory usage in deep learning training.", "theoretical_claims": "The paper provides a rigorous theoretical analysis of its algorithm", "experimental_designs_or_analyses": "The paper includes exhaustive comparisons with state-of-the-art baselines, demonstrating its effectiveness across multiple architectures and hierarchy depths.", "relation_to_broader_scientific_literature": "HiRemate extends prior work on ILP-based re-materialization methods such as Checkmate and TW-Remat, which optimize memory usage by selectively recomputing activations. Unlike these approaches, which struggle with large computation graphs due to scalability issues, HiRemate introduces a hierarchical partitioning strategy that allows it to efficiently scale to larger models.", "essential_references_not_discussed": "The related work seems exhaustive.", "other_strengths_and_weaknesses": "Strength\n- The paper provides a rigorous theoretical analysis of its algorithm\n- The paper includes exhaustive comparisons with state-of-the-art baselines, demonstrating its effectiveness across multiple architectures and hierarchy depths.\n- The proposed method is easily integrable with PyTorch, enhancing usability and enabling practical adoption with minimal modifications.\n- The intuition behind the scheduling algorithm design is well-explained and easy to follow\n\nWeakness \n- While the hierarchical ILP solver improves scalability, it still incurs significant computational overhead for very large graphs. The reliance on multiple ILP optimizations per subgraph introduces bottlenecks, making it less practical for real-time or iterative optimization scenarios.\n\n- The framework primarily targets static computational graphs, limiting its applicability to modern AI applications that involve dynamic graphs, such as reinforcement learning, adaptive architectures, and some transformer-based models.", "other_comments_or_suggestions": "Please refer to Weaknesses.", "questions_for_authors": "Please refer to Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1dggFPcbVy", "reviewer": "Reviewer_M2sS", "summary": "The paper introduces HiRemate, a hierarchical framework for neural network re-materialization to reduce memory usage during training. The core idea involves recursively partitioning the computation graph into manageable subgraphs, solving each with optimized strategies, and merging solutions hierarchically.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the careful reading and detailed feedback. Below we address each of point of the review in turn.\n\n**On Lack of Theoretical Guarantees and Suboptimality Gap**\nWe agree that HiRemate does not provide formal guarantees on optimality. This is also true of Rockmate and TW-Remat, which, like HiRemate, are designed for scalability rather than provable optimality.\nOur focus is on solving large, complex computation graphs that are intractable for methods like Checkmate. The observed gap in small graphs (e.g., the 3% difference on GPT-2 5-layer) reflects this trade-off. HiRemate is designed to let users balance solution quality and solving time by:\n- Adjusting the number of memory budget levels used in lower-level ILPs,\n- Controlling partition granularity (larger subgraphs often improve quality but require more time),\n- Using mixed solvers (e.g., DP on sequential parts, ILP only where needed),\n- Optionally incorporating user-defined partitioning strategies.\n\nThese controls allow users to tune the solution process as needed. While we currently lack a formal analysis of suboptimality bounds, we are conducting further experiments to better characterize this trade-off in practice.\n\n**On Problem Statement and Dependency Handling**\nWe clarify that HiRemate does not assume a topological order in the original graph. We extract the forward-backward graph using `torch.export()` at the level of torch/aten operations. The graph is simplified using the **rkgb** tool (Zhao et al., 2023) to remove view operations and isolate meaningful compute and data nodes (Appendix C.1).\nThis graph forms the input to our partitioning procedure (Section 3.1), which does not require topological ordering. A valid topological order is computed only within each subgraph before solving the H-ILP problem, as required for the ILP formulation -- similar to Checkmate.\nThe generated schedule consists of forward `aten` calls, backward Autograd calls, and memory deallocations (see Appendix C.2). Dependencies are fully respected by construction and execution is sequential.\n\n**On Partition Granularity vs. Solution Quality**\nHiRemate is designed to enable trade-offs between granularity and schedule quality. We evaluate this effect on encoder-decoder Transformers in Figure 4, which shows how deeper hierarchies (finer partitions) impact both solving time and memory usage.\n\nA fully rigorous analysis would require evaluating several partitioning strategies across multiple architectures. We agree this is an important direction and are currently running additional experiments, although they were not yet complete at the time of submission.\n\n**On Comparison to Dynamic Re-materialization (e.g., DTR)**\nStatic and dynamic re-materialization strategies differ mainly in how they handle the structure of the computational graph. Static re-materialization is well-suited for models with fixed computation graphs, where memory usage can be precomputed and optimized for predictable and efficient execution. On the other hand, dynamic strategies are better adapted to models with variable or input-dependent control flow,  where runtime decisions enable on-the-fly memory management. While dynamic approaches offer more flexibility, they generally come with higher runtime overhead and less predictable behavior. Static strategies, in contrast, exceed dynamic methods in memory efficiency when applied to well-structured models. Quantifying these differences is complex because most dynamic strategies generally incorporate other algorithmic ingredients (typically paging) as in POET (https://proceedings.mlr.press/v162/patil22b/patil22b.pdf) or MegTaiChi (https://dl.acm.org/doi/pdf/10.1145/3524059.3532394), which makes comparison difficult. \nWe see integration or benchmarking against these methods as important future work.\n\n**On Checkmate and Rockmate Comparisons**\nCheckmate produces optimal schedules but scales poorly with graph size. We include comparisons where tractable (e.g., small GPT models), but for larger graphs, Checkmate does not complete in reasonable time. HiRemate is designed to provide approximate but scalable solutions for those larger models. In practice, we use Checkmate when feasible, and HiRemate otherwise.\nRockmate is designed for block-sequential architectures. When applied to non-sequential graphs (e.g., U-Nets), it treats the entire model as a single block, falling back to Checkmate behavior. As such, on non-sequential models, Rockmate effectively reduces to a limited form of Checkmate and does not offer meaningful additional comparison.\n\n**Regarding RNNs and SSMs**  Please, refer to the Answer to Reviewer YaUT", "claims_and_evidence": "- Comparisons with Checkmate on small graphs (Table 2) show close performance, but no theoretical guarantees or bounds on suboptimality. Actually for 5-layer GPT2, H-ILP already has >3% performance gap compared to the optimal one achieved by Checkmate.\n- The problem statement in Section 3 seems to be highly simplified. Do you assume a topological order of a computational graph here? How do you handle dependency between layers/operators.", "methods_and_evaluation_criteria": "- Hierarchical decomposition is sensible for scalability, but the trade-off between partition granularity and solution quality is not rigorously analyzed.", "theoretical_claims": "- As mentioned above, comparisons with Checkmate on small graphs (Table 2) show close performance, but no theoretical guarantees or bounds on suboptimality.", "experimental_designs_or_analyses": "- Please include comparisons with dynamic re-materialization methods (e.g., DTR [A]) so that readers can more effectively evaluate the benefits of the hierarchical approach, especially considering that statically planning the schedule for a Transformer model with 2500 nodes takes 2.5 hours, which may be too expensive.\n- Comparisons with Checkmate are limited to small graphs, and Rockmate is only tested on sequential models. Please consider including RNN or SSM experiments.", "relation_to_broader_scientific_literature": "It introduces a hybrid method that integrates hierarchical and solver-based approaches to address a key challenge in re-materialization.", "essential_references_not_discussed": "[A] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, Zachary Tatlock, \"Dynamic Tensor Rematerialization\", ICLR, 2021.", "other_strengths_and_weaknesses": "See the above sections.", "other_comments_or_suggestions": "See the above sections.", "questions_for_authors": "See the above sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents an approach to scheduling computations required for neural network training subject to a memory usage budget.  All reviewers lean toward accept.  The hierarchical solver proposed for this scheduling problem appears to deliver good empirical results in comparison to relevant baselines.  Reviewers ag6k and M2sS suggest additional experimental analyses that would be informative and should be considered for inclusion in a revision of the paper.
Decision: Accept (poster)