# Optimized Forward-Backward Rematerialization for Memory-Efficient Pipeline Parallel Training

Adrien Aguila-Multner, Olivier Beaumont, Lionel Eyraud-Dubois, Julia Gusak

<figure><img src="image_1.png" /></figure>

To cite this version:

Adrien Aguila-Multner, Olivier Beaumont, Lionel Eyraud-Dubois, Julia Gusak. Optimized Forward-Backward Rematerialization for Memory-Efficient Pipeline Parallel Training. 2025. hal-05151601v2

HAL Id: hal-05151601

https://inria.hal.science/hal-05151601v2

Preprint submitted on 12 Nov 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
OPTIMIZED FORWARD-BACKWARD REMATERIALIZATION FOR
MEMORY-EFFICIENT PIPELINE PARALLEL TRAINING

Adrien Aguila–Multner<fnref n="1" /> Olivier Beaumont<fnref n="1" /> Lionel Eyraud-Dubois<fnref n="1" /> Julia Gusak<fnref n="1" />

ABSTRACT

Pipeline parallelism is a key technique for scaling deep network training across multiple devices. Recent works have significantly reduced pipeline idle time by improving scheduling efficiency. Decoupling the computation of gradients with respect to weights and activations led to the development of schedules with almost no idle time. However, these methods still require substantial memory, limiting their applicability on resource-constrained hardware. Our first contribution is to introduce recomputation to the backward pass, extending rematerialization beyond the forward pass. This enables executing schedules with decoupled gradient computations under much tighter memory constraints. Our second contribution is a unified optimization approach that, given a model and hardware memory constraints, formulates and solves an Integer Linear Programming (ILP) problem to determine the optimal per-microbatch, per-GPU rematerialization strategy for a given schedule, applicable to both one-wave and multi-wave pipeline schedules. Our third contribution shows that, as device memory constraints vary, the relative advantages of different pipeline schedules also change in the presence of rematerialization. We provide corresponding insights and a PyTorch framework that enables finding and executing the optimal combination of pipeline scheduling and rematerialization strategies. Experiments demonstrate the effectiveness of all three contributions, showing that our approach enables efficient training of larger models under tight memory budgets, adapts optimally to varying memory capacities, and reduces recomputation overhead compared to existing recomputation solutions.

# 1 INTRODUCTION

Training large neural networks requires distributing the computation across multiple devices. Pipeline parallelism is a widely adopted solution for large-scale training that enables different model layers to run concurrently across GPUs. Scheduling strategies such as Gpipe (Huang et al., 2019), 1F1B (Harlap et al., 2018), Megatron (Narayanan et al., 2021), and Hanayo (Liu et al., 2023b) substantially improve device utilization by overlapping the execution of forward and backward passes across microbatches. However, achieving pipeline efficiency requires a sufficiently large number of microbatches, which leads to a significant memory pressure. Even with efficient pipelining strategies, large models or long sequences cannot be accommodated.

Rematerialization (Jain et al., 2020; Gusak et al., 2025; Zhao et al., 2023a; Korthikanti et al., 2023), also known as activation checkpointing, is a well-established technique that

reduces the memory footprint by discarding intermediate activations computed during the forward pass and recomputing them during the backward pass. It is commonly used in single-device and pipeline-parallel training, where it typically targets forward activations and where strategies are adapted to individual pipeline stages (Liu et al., 2023a; 2025). Nevertheless, at the level of a stage, most existing approaches apply the same rematerialization policy to all microbatches regardless of their position in the pipeline schedule or the memory pressure present at different steps of execution. Importantly, these approaches ignore memory usage during backward computation, especially when gradient computations of activations and weights are decoupled.

Gradient decoupling, introduced in recent work such as Zero-Bubble (Qi et al., 2024b), enables more efficient pipelining by splitting the backward pass into separate computations for gradients with respect to activations and weights. While this technique improves throughput, it also increases memory pressure since some backward intermediate values must be retained for a longer time (before computing gradients with respect to weights).

The main contributions of this paper are as follows:

* Extending the rematerialization approach to pipeline-

<footnote n="1">¹Inria. Correspondence to: Adrien Aguila–Multner <adrien.aguila-multner@inria.fr>, Olivier Beaumont <olivier.beaumont@inria.fr>, Lionel Eyraud-Dubois <lionel.eyraud-dubois@inria.fr>, Julia Gusak <yulia.gusak@inria.fr>.</footnote>
parallel training by adding recomputations for backward operations through the use of decoupled gradient (see Section 3.2). This enables memory reductions beyond forward-only approaches.

* Optimization-based rematerialization strategies: we formulate the memory optimization problem as an Integer Linear Program, which outputs per-microbatch rematerialization decisions that are optimal with respect to the pipeline schedule and the memory budget constraint (see Section 3.3).

* Analysis and framework support: we show how varying device memory constraints alters the relative benefits of different pipeline schedules in the presence of rematerialization, and introduce ELF, a PyTorch-based framework integrating automated model profiling, scheduling, rematerialization planning and execution for efficient mid-scale, memory-limited training.

## 2 RELATED WORK

To address memory usage, many pipeline systems employ rematerialization (also called activation checkpointing), primarily to recompute forward activations. For example, GPipe (Huang et al., 2019) applies a uniform checkpointing of every few layers, Megatron-LM (Narayanan et al., 2021) allows users to manually designate checkpointed layers, and others apply simple heuristics at the stage level. However, these methods use a uniform or greedy policy across microbatches and do not consider memory usage from backward computations, which can become dominant in decoupled schedules such as those proposed in (Qi et al., 2024b). In contrast, our work performs schedule-aware rematerialization of both forward and backward computations, optimized globally across microbatches.

Rematerialization has also been studied extensively outside of pipelining, typically in the context of single-device execution (Jain et al., 2020; Zhao et al., 2023a; Korthikanti et al., 2023; Patil et al., 2022; Bartan et al., 2023). Static approaches such as Checkmate (Jain et al., 2020) and Rotor (Beaumont et al., 2024) (and follow-ups as Rockmate (Zhao et al., 2023a), Hiremate (Gusak et al., 2025)) formulate checkpoint selection as an optimization problem to minimize training time under a memory constraint for single-batch execution. These methods are powerful, but assume sequential execution and do not generalize to pipelined settings with overlapping microbatches and stage-level parallelism. DTR (Dynamic Tensor Rematerialization) (Kirisame et al., 2020; Hu et al., 2022; Patil et al., 2022) is a runtime approach that evicts tensors and recomputes them on-the-fly when needed. While DTR can adapt to runtime memory pressure, it lacks a global view of the entire pipeline schedule, and cannot guarantee an optimal plan or full utilization of available memory budgets.

The gradient decoupling technique for pipelines was introduced by Zero-Bubble (Qi et al., 2024b), with the observation that backward computation (namely, gradients with respect to activations and weights) can be decoupled. Zero-Bubble proposed schedules for microbatches that benefit from this decoupling to reduce or eliminate the idle time (called pipeline bubbles). These schedules are focused on maximizing throughput by eliminating pipeline stalls. Our work builds on this decoupled backward execution model, but switches the focus towards memory optimization. We integrate decoupled backward passes with rematerialization, something that was not explored in the above works.

Other parallelization strategies like tensor parallelism (Wang et al., 2022), ZeRO3 (Ren et al., 2021), Zero-infinity (Rajbhandari et al., 2021) or FSDP (Zhao et al., 2023b) are alternative approaches to address the memory pressure from increased model size, by splitting optimizer, parameters or gradients across devices. However, these approaches typically introduce additional inter-device communications, that increase with the number of devices. In contrast, our approach targets activation and intermediate gradient memory and only uses local recomputation (no extra communication), making it complementary to sharding-based memory approaches. Likewise, sequence parallelism (Li et al., 2021) partitions the sequence length across devices to train longer sequences, but requires model modifications and per-token communication. Our method is model-agnostic, requires no changes to model internals, and can achieve longer-sequence training benefits through recomputation.

## 3 PROPOSED METHOD

Our approach builds on the observation that every scheduling algorithm has a different memory footprint that depends on various factors such as the number of GPUs, the number of micro-batches, the model architecture and hyperparameters, and others. As a consequence, finding an optimal rematerialization strategy for a given configuration is not trivial and can lead to significant performance gains. In this section, we first outline memory bottlenecks in pipeline parallel training, then describe the space of rematerialization options and extend it to decoupled schedules with backward rematerialization. We then present our ILP formulation for optimal checkpoint planning.

### 3.1 Concepts and Notations

It is common for pipeline parallelism to model a neural network as a sequence of *stages*, where the output of each stage is the input to the next. These stages are then each assigned to a specific GPU. Furthermore, to allow parallel execution, the input batch is divided into *m* microbatches: once the first GPU completes the forward pass on the first microbatch, the second GPU can begin processing that mi-
crobatch while the first GPU moves to the next one. In the following, we denote with *s* the number of stages, *p* the number of GPUs, and with *m* the number of microbatches.

The overall execution is defined by a schedule, which specifies for each GPU the order of execution of the different (forward and backward) computations on the different microbatches. Several schedule families have been proposed in the literature, with trade offs between computational efficiency and memory usage. In single-wave schedules (Harlap et al., 2018; Huang et al., 2019; Qi et al., 2024b), there is exactly one stage per GPU, but multi-wave schedules (Liu et al., 2023b; Li & Hoefler, 2021; Narayanan et al., 2021; Fan et al., 2021; Qi et al., 2024b) allow for the possibility to have several stages on each GPU. The Integer Linear Programming (ILPs) formulations (Chvatal, 1983) proposed in this paper are valid for any single-wave or multi-wave schedule.

In the classical setting, backward consists of a single operation; however, recent work like Zero-Bubble (Qi et al., 2024b) has decomposed it into two distinct steps: **B (computation of gradients w.r.t. activations)** and **W (computation of gradients w.r.t. parameters)**. As such, schedules in this family contain three operation types: *F*, *B*, and *W*, where *F* denotes the forward computations. The benefit of this decomposition is that the backward operations of a stage only depends on the *B* operation of the next stage, and the *W* operation can be performed later when the GPU is otherwise idle. However, this has an impact on memory usage since the data required to compute *W* needs to be kept in memory for a longer time between *B* and *W*.

With no recomputation, for a given stage and a given microbatch, the order of operations is *F*, *B*, *W*. We denote by $M_F$ and $M_B$ the amount of data kept in memory for one microbatch after operations *F* and *B* respectively* for a given stage. Since the tensors stored after *B* are partially activations and partially gradients, we further denote with $M_B^{act}$ (respectively $M_B^{grad}$) the memory usage for activations (respectively gradients) kept between *B* and *W*. Similarly, we denote with $T_F$, $T_B$ and $T_W$ the respective times required for operations *F*, *B* and *W*.

## 3.2 Rematerialization Scenarios and Complexity

In many pipeline scheduling algorithms, even for homogeneous stages, the memory consumption is imbalanced between GPUs. For instance, if all stages have the same values of $M_F$ and $M_B$, the 1F1B algorithm (Harlap et al., 2018) has a peak memory of $(p-1)M_F$ on the first GPU, but only $M_F$ on the last one. Similarly, ZB-H2 from Zero-Bubble (Qi et al., 2024b) induces a peak memory of $(2p-1)M_F$ on the first GPU, and $(2p-1)M_B$ on the last one.

The latter is usually smaller because $M_B$ tends to be below $M_F$. In general, the peak memory usage on one GPU is $\max_t a(t)M_F + b(t)M_B$, where $a(t)$ and $b(t)$ are respectively the number of *B* and *W* operations waiting for computation at time *t*. For most classical schedules, the peak is reached due to large $a(t)$ values for the first GPUs in the pipeline, and due to large $b(t)$ values for the last ones.

Traditional (forward only) rematerialization allows one to lower the first term $(a(t)M_F)$ down to zero if necessary, by deleting every activation after *F* during the forward step and recomputing them before the corresponding *B*. However, this approach cannot easily be extended to the second term $b(t)M_B$. In particular, for ZB-H2, the model cannot be trained as soon as $(2p-1)M_B$ exceeds the memory capacity.

In the present paper, we propose to delete some values (activations and/or gradients) between *B* and *W* also and to recompute them before the *W* operation, in order to lower the $b(t)M_B$ term. Since the tensors stored after *B* are partially activations, and partially gradients, there are two possibilities: recomputing activations and recomputing gradients. In many cases, recomputing gradients requires the knowledge of all the forward activations, so we only consider this opportunity when combined with the rematerialization of activations.

In addition to the basic solution with no recomputations at all, we can therefore use **full recomputation** between *F* and *B*, and either recomputation of **activations** or recomputation of **activations and gradients** between *B* and *W*. These options yield new memory tradeoffs detailed in Tables 1 and 2. A detailed example of how these options can be applied in the case of a standard Transformer block is provided in Section A.5 of the Appendix.

<table><caption>Table 1: Forward Recomputation Options</caption><thead><tr><td>Type</td><td>Memory Kept (F → B)</td><td>Time Overhead</td></tr></thead><tbody><tr><td>None</td><td>$M_F$</td><td>0</td></tr><tr><td>Full</td><td>0</td><td>$T_F$</td></tr></tbody></table>

<table><caption>Table 2: Backward Recomputation Options</caption><thead><tr><td>Type</td><td>Memory Kept (B → W)</td><td>Time Overhead</td></tr></thead><tbody><tr><td>None</td><td>$M_B^{act} + M_B^{grad}$</td><td>0</td></tr><tr><td>Acts</td><td>$M_B^{act}$</td><td>$T_F$</td></tr><tr><td>Acts & Grads</td><td>0</td><td>$T_F + T_B$</td></tr></tbody></table>

## 3.3 ILP-Based Rematerialization Planning

We formalize pipeline rematerialization as an optimization problem. We consider a sequential model partitioned into *S* stages, and a schedule with a decoupled two-phase backward pass. The problem is to choose, for each microbatch and each stage, a rematerialization option from Tables 1 and 2. The objective is to minimize the training time over-

<footnote>*The amount of data corresponding to *W* is $M_W = 0$.</footnote>
head while satisfying a memory budget constraint per stage. Selecting more recomputation (deleting more intermediate values) saves memory at the cost of extra compute. Selecting less recomputation saves compute but requires more memory.

**ILP formulation.** We define binary decision variables for each stage *s* and microbatch *m* indicating whether the intermediate forward results at that stage for that micro-batch will be recomputed (as opposed to stored). Additional binary variables are used to model the recomputation of backward-phase intermediates separately from forward activations. We then introduce linear constraints to model pipeline execution dependencies and memory usage: for each stage and time step, the total memory occupied by all live activations and gradients (those not discarded) cannot exceed the GPU memory budget. These constraints incorporate the schedule — i.e. which microbatch outputs are live at a given time — making the problem schedule-aware. The objective is to minimize the total extra computation time due to recomputation. Since recomputation time for a layer is known from profiling, this can be computed as a weighted sum of the decision variables. We emphasize that the per-microbatch decisions of the ILP formulation enable non-uniform rematerialization. In addition, since the formulation minimizes the total computation time under a memory budget constraint, it is fully adaptive: if the available memory is large enough so that recomputation is not needed, then the resulting solution will contain no rematerialization and incur no overhead.

**Solving ILP formulation** The ILP formulation is solved offline (before training) given the models profiled memory/computation characteristics and the target schedule. In practice, we solve this ILP using an open-source solver (PuLP) with a time limit, yielding an optimal or near-optimal plan. The impact of this time limit on the quality of the solution is studied in section 4.4. The output specifies, for each microbatch and each stage, whether to store or recompute that stages activations and gradients, fully or partially. An example for a strategy obtained by solving the ILP is given Figure 1.

All ILP details, including complete variable definitions, constraints, and solver parameters, are provided for reference in Section A.2 of the appendix.

# 4 EXPERIMENTS

We evaluate our ILP-based approach on decoder-only Transformer models as described in (Korthikanti et al., 2023), with FlashAttention-2 from (Dao, 2023). For model hyper-parameters, we use three model sizes: 9.6B, 21B and 38.5B, as described in (Qi et al., 2024a).

All experiments are run on a fixed number of GPUs (4, 8, or 16), using NVIDIA H100 80GB GPUs. All parameters, activations and gradients are stored in BF16 precision. Computations are also performed using BF16. We use PyTorch 2.7.1 with CUDA 12.6 and NCCL 2.26.2.

Pipeline parallelism is implemented in our library for distributed training, named ELF. ELF uses static lists of operations — either computations or communications — that are executed at runtime in the given order on each GPU. We have implemented state-of-the-art pipeline schedules in this library, allowing for fair and easy comparison of all approaches. Rematerialization is implemented by inserting additional forward and/or backward operations just before the computation that needs them, based on the solution of the ILP.

The ELF runtime automatically manages communications, while respecting the schedule order. Communication operations are first introduced in the schedule just before the corresponding computation, then reordered by a static analysis that optimizes concurrency while avoiding deadlocks. At runtime, they are performed asynchronously and overlapped with computations when possible.

We compare against state-of-the-art pipeline parallel baselines and measure two key aspects: (1) the maximum sequence length that can be trained for a given model size and given number of GPUs, and (2) the overhead in iteration time introduced by our rematerialization strategies. We consider the following pipeline training setups for comparison:

**No Rematerialization:** A baseline pipeline run with no rematerialization at all (stores all activations/activation gradients). This maximizes throughput but quickly runs out of memory (OOM) for large models.

**Uniform Rematerialization:** The standard, naive approach that consists in recomputing every stage and micro-batch in the same manner. For non-decoupled schedules, we only recompute forward (**Uni-F**). Otherwise we recompute both forward before *B*, and forward + backward before *W* (**Uni-FB**).

**Greedy Algorithm:** A simple algorithm as proposed in (Korthikanti et al., 2023) that decides whether to recompute each micro-batch greedily. We use it as-is (**Greedy-F**) and also extend it to backward rematerialization (**Greedy-FB**). A more precise description of this algorithm can be found in A.4.

**ILP-based Rematerialization (ours):** The rematerialization strategy based on our proposed ILP approach. We denote **ILP-F** the variant that only allows forward rematerialization, and **ILP-FB** the one that allows both forward and backward rematerialization.

We evaluate the proposed methods across a set of five dis-
<figure><img src="image_3.png" /><figcaption>Figure 1: Example of a rematerialization strategy given by solving the ILP on 4 GPUs. The base schedule is ZB-H1 and profiling statistics from model 9.6B with sequence length 4096 are used. The memory budget was set to 25120 MB.</figcaption></figure>

tinct base schedules in order to determine the most appropriate scheduling strategy for each scenario. The schedules considered are the decoupled variants **ZB-H1**, **ZB-H2**, and **ZB-V**, together with two configurations of the 1F1B approach: one stage per GPU (**1F1B**) and two stages per GPU (**1F1B-2W**). We also use ZB-V with two stages per GPU.

## 4.1 Backward Rematerialization

Table 3 details the benchmarked iteration times on 8 GPUs for different models and sequence lengths, depending on the rematerialization strategy and base scheduler used. We measure the average time per iteration over 15 consecutive iterations of the pipeline, and repeat this measurement 10 times to ensure statistical significance. The reported time is the median of the 10 measurements, after discarding the single highest and single lowest values.

More experimental data for different number of GPUs, as well as MFU values, are available in Appendix A.3.

We first examine the benefit of backward rematerialization. For decoupled schedules, we comparing variants with forward-only and forward+backward recomputation options for Uniform, Greedy and ILP strategies. Table 3 shows that the forward-only Uniform strategy (Uni-F) frequently runs out of memory — recomputing every forward step is insufficient to fit the model. In contrast, its backwardaugmented counterpart (Uni-FB) always succeeds. The same pattern holds for the ILP strategy: when ILP-F fails to find a feasible solution, ILP-FB always yields a valid one. Moreover, for some configurations the best solution among all tested strategies and schedules, is a decoupled schedule which would not fit with forward-only recomputation: for instance 9.6B with sequence length 16384, 21B with sequence length 8192 or 38.5B with sequence length 4096. In those cases backward recomputation effectively reduces the best iteration time achievable among our methods.

These results demonstrate the effectiveness of backward rematerialization in increasing the range of possible strategies with efficient decoupled schedules.

## 4.2 ILP effectiveness evaluation

In Figure 2, we present a comparison between all strategies for our 3 decoupled base schedules. We use model 21B with increasing sequence lengths to test cases with different memory needs.

From these plots we can make a first observation: when training without rematerialization fits in memory for low sequence length, our ILP solution has the same performance (or nearly so) as NoRemat. Indeed, the ILP formulation only introduces rematerialization when it is required to fit into memory; this ensures that using our approach does not have any downside when enough memory is available.

In the presence of stronger memory constraints (when rematerialization is needed to fit in memory because of larger sequence length), we can see that using per micro-batch rematerialization (Greedy or ILP) yields a significantly reduced training time compared to the standard baseline of uniform rematerialization. Furthermore, the stronger optimization provided by our ILP formulation provides further reduction in training time for cases with more limited memory compared to the Greedy approach. In extreme cases, the Greedy algorithm fails to find a valid solution, whereas our ILP-based approach is still able to provide significant speedup compared to the uniform rematerialization baseline.

## 4.3 Scheduling and Rematerialization

These experiments reveal an interesting observation: the optimal scheduling algorithm for a given scenario is not always straightforward. In particular, ZB-H2 and ZB-V perform best for shorter sequence lengths where no rematerialization is required; however, the recomputation overhead of ZBH2 increases rapidly as memory demands grow. In extreme cases — for instance, when $s = 16384$ on 16 GPUs — 1F1B can even outperform ZeroBubble schedules in the presence of recomputations.

We examine these observations in this Section by comparing the objective value achieved by the ILP solver for decreasing memory budgets on Figure 3. We use profiling data from
<figure><img src="image_4.png" /><figcaption>Figure 2: Comparison of our method with baselines for different sequence lengths on model 9.6B. The missing bars correspond to a solution not being found or an Out of Memory error during execution.</figcaption></figure>

model 9.6B with sequence length 4096 as an example. The solving time limit was set to 128s for faster solving. The plot spans from the last memory budget that is sufficient to run all schedules without rematerialization, to the last memory budget where one of the schedule can be executed. Note that these values are only obtained by solving the ILP, and never actually measured with real execution.

These figures illustrate an important observation: the ranking of scheduling strategies is different between very high memory budgets, and tight memory limits. In the first case Zero-Bubble schedules clearly outperform others, whereas under more severe constraints they are surpassed.

To understand this contrast, consider the extreme case where only one forward (or backward) pass can be stored in memory. In that case, everything needs to be recomputed for every micro-batch.

For 1F1B and other non-decoupled schedules, this requires performing an additional forward computation before every backward. Assuming $T_F = 1$ and $T_B = 2$, this resulting overhead is $\frac{4}{3}$.

For schedules that decouple gradient computation, recomputing *F* and *B* is also needed before *W*. The overhead then rises to $\frac{6}{3} = 2$ assuming $T_W = T_B = 1$.

These theoretical values align with the observation shown in Figure 3. For instance with 8 GPUs, the makespan for

the ILP-F (1F1B) strategy grows from 1.21 to 1.59 seconds between the first and last data points, a ×1.32 increase. For ILP-FB (ZB-H2), the makespan increases from 1.01 to 2.06 seconds, which is a ×2.03 overhead. ZB-V is a special case because some *W* are performed directly after their corresponding *B*, while some others are not. As a consequence the overhead factor lies in between $\frac{4}{3}$ and 2 (in this case, ×1.54).

From that observation we can conclude that the optimal schedule for any given case depends on the memory budget, the number of GPUs, the profiling values, etc. The strength of our method is that we can solve the ILP for different base schedules, and use the best solution found among all algorithms.

## 4.4 ILP Solving Time

When solving an ILP problem, usually some stopping criterion is applied to be able to control the maximum solving time. For all experiments in the paper, the ILP solutions are found within a time limit of 10 minutes unless otherwise specified.

In this section we investigate how the best solution for a given solving time *t* evolves with *t*. In Figure 4, we show the best solution as well as the theoretical lower bound found by the solver after different time steps. We use model 9.6B with sequence length 16386 (the highest memory usage for
<figure><img src="image_5.png" /><figcaption>Figure 3: Estimated makespan for the ILP solution with different memory budgets.</figcaption></figure>

this model) and ZB-V as a base schedule. The solution is the estimated iteration time of the pipeline in milliseconds.

For cases up to 8 GPUs (Figure 4a & 4b), the maximum solving time does not have a significant impact, as the first solution found is very close to the best solution, and is quickly proven optimal or improved to optimality. However, when increasing the number of GPUs to 16, the problem becomes more complex and the first solutions are significantly worse than the optimal. However they quickly approach the lower bound: after 600s, the gap between the best solution found and the proven possible best is below 0.08%. This observation confirms that we can limit the solving time to 600 seconds without compromising the quality of the solution.

# 5 DISCUSSION

**Scope of Applicability** Our work targets training scenarios with a fixed number of GPUs where memory becomes the limiting factor. While pipeline parallelism can distribute compute, the memory remains a bottleneck due to activations and gradients. Our proposed algorithm **ILP-FB** can be applied to optimize rematerialization decisions for an arbitrary pipeline schedule. By being able to separately optimize forward and backward rematerialization, **ILP-FB** provides memory gains for schedules which feature gradient decoupling like the ones proposed in Zero-Bubble (Qi et al., 2024b), while still being relevant for more classical schedules. Interleaved schedules are also supported with any number of waves (stages per GPU), as well as any static assignment of stages to GPUs.

**Insights on computed solutions** A more detailed observation of the solutions provided by ILP-FB shows the following trends when the memory pressure increases (in our experiments, it corresponds to increasing model size or sequence length). The first GPUs start with forward selec-

tive recomputation if not using FlashAttention on a few micro-batches, since this gives the best trade-off for time overhead versus memory saved. This strategy then spreads to almost all GPUs. When selective recomputation is no longer enough for the first GPU, it is gradually replaced by full forward recomputation. Around the same time, the last GPUs start adding backward recomputation of activations. Finally, the last GPUs move from recomputing only activations to recomputing both activations and gradients. The pattern is clear for the first and last GPUs, however for the intermediate GPUs it is a mix that depends on their position in the pipeline. Our ILP formulation is able to find this blend automatically, whereas hand-crafting a rule for it would be complex. More details can be found in Appendix A.6.

**Other Parallelisms** Tensor parallelism splits layer-wise computations across GPUs and is effective in scaling very large models. However, it incurs substantial inter-device communication — often all-reduce operations per layer — and requires tight synchronization. In contrast, our method performs local recomputation without inter-device tensor synchronization and is optimized for memorycompute trade-offs, not just model scaling.

Sequence parallelism is widely used in transformer-based 4D training to distribute attention operations along the sequence dimension. While effective for long sequences, it introduces per-token communication, requires modifying attention kernels, and cannot be applied to non-transformer models. Our approach is model-agnostic, requires no internal model modifications, and achieves memory savings through recomputation strategies that can be applied, for example, to SSMs.

While sequence parallelism is designed for transformer architectures, Structured State Space Models (SSMs) like Mamba offer a promising direction for long-sequence mod-
<figure><img src="image_6.png" /><figcaption>Figure 4: Best solution and lower bound for our ILP formulations for increasing solving time limits. Lines end with the time limit under which an optimal solution is found.</figcaption></figure>

eling without attention. Our method could be integrated
with Mamba to reduce activation and gradient memory dur-
ing training, and enable pipeline-parallel execution with
per-microbatch rematerialization strategies for activations
and decoupled gradients.

In contrast to FSDP and Zero-3 Parallelism, our method
targets activation and intermediate gradient memory reduc-
tion by rematerializing them selectively, without any new
communication across devices.

**Per-Microbatch Rematerialization** A key distinction of our method is its ability to apply different rematerialization strategies for different microbatches during both forward and backward. Prior work whether static (e.g., Megatron (Narayanan et al., 2021) + selective recomputation (Kothikanti et al., 2023)) or dynamic (Kirisame et al., 2020; Hu et al., 2022; Patil et al., 2022) applies checkpointing decisions uniformly or based on runtime heuristics, without awareness of the pipeline structure or schedule. Dynamic approaches may incidentally produce different rematerialization behavior across microbatches, but do so without global coordination or optimization.

Our ILP-based method uses the static pipeline schedule and profiling data to assign tailored rematerialization plans per microbatch in order to minimize throughput given the memory constraints. This is, to the best of our knowledge, the first system to explicitly and deterministically optimize per-microbatch forward plus backward recomputation in pipeline-parallel training, enabling deeper memory savings and better control of the memorycompute tradeoff.

6 CONCLUSION

This work introduces a unified optimization approach to
memory-efficient pipeline-parallel training by extending

activation rematerialization into the decoupled backward
pass and formulating the rematerialization problem as an
ILP with per-microbatch, per-GPU decisions. Backward
rematerialization substantially relaxes memory constraints,
enabling the training of larger models and longer sequences
on the same hardware. The ILP-based optimization consis-
tently outperforms heuristic methods across both decoupled
(Zero-Bubble) and non-decoupled schedules. Our exper-
iments reveal that as memory constraints become tighter,
the optimal pipeline schedule systematically shifts from one
type to another, showing a structured dependency between
available memory and schedule efficiency. The proposed
method automatically identifies the best schedule – remate-
rialization combination for any given constraint. All tech-
niques are implemented in the ELF framework, providing
a practical and extensible tool for evaluating and executing
memory-aware pipeline training on multi-GPU system with
limited memory.

7 ACKNOWLEDGMENTS

This work was validated in part at the IDRIS Open
Hackathon and Poland Open Hackathon, parts of the Open
Hackathons program. The authors would like to acknowl-
edge OpenACC-Standard.org for their support.

This work was performed using HPC resources from
GENCI-IDRIS (Grant A0180601567).
REFERENCES

Bartan, B., Li, H., Teague, H., Lott, C., and Dilkina, B. Moccasin: Efficient tensor rematerialization for neural networks. In *International Conference on Machine Learning*, pp. 1826–1837. PMLR, 2023.

Beaumont, O., Eyraud-Dubois, L., Herrmann, J., Joly, A., and Shilova, A. Optimal re-materialization strategies for heterogeneous chains: How to train deep neural networks with limited memory. *ACM Trans. Math. Softw.*, mar 2024. ISSN 0098-3500. doi: 10.1145/3648633. URL https://doi.org/10.1145/3648633. Just Accepted.

Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur-Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A. M., Pillai, T. S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier-Hellstern, K., Eck, D., Dean, J., Petrov, S., and Fiedel, N. Palm: Scaling language modeling with pathways, 2022. URL https://arxiv.org/abs/2204.02311.

Chvatal, V. *Linear programming*. W. H. Freeman and Company, New York, 1983.

Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023. URL https://arxiv.org/abs/2307.08691.

Fan, S., Rong, Y., Meng, C., Cao, Z., Wang, S., Zheng, Z., Wu, C., Long, G., Yang, J., Xia, L., et al. Dapple: A pipelined data parallel approach for training large models. In *Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming*, pp. 431–445, 2021.

Gusak, J., Zhao, X., Le Hellard, T., Li, Z., Eyraud-Dubois, L., and Beaumont, O. HiRemate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks, 2025. URL https://hal.science/hal-04403844.

Harlap, A., Narayanan, D., Phanishayee, A., Seshadri, V., Devanur, N., Ganger, G., and Gibbons, P. Pipedream: Fast and efficient pipeline parallel dnn training, 2018.

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, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, Z. Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2019.

Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Keutzer, K., Stoica, I., and Gonzalez, J. E. Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2020.

Kirisame, M., Lyubomirsky, S., Haan, A., Brennan, J., He, M., Roesch, J., Chen, T., and Tatlock, Z. Dynamic tensor rematerialization. *arXiv preprint arXiv:2006.09616*, 2020.

Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. *Proceedings of Machine Learning and Systems*, 5:341–353, 2023.

Li, S. and Hoefler, T. Chimera: efficiently training large-scale neural networks with bidirectional pipelines. In *Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis*, pp. 1–14, 2021.

Li, S., Xue, F., Baranwal, C., Li, Y., and You, Y. Sequence parallelism: Long sequence training from system perspective. *arXiv preprint arXiv:2105.13120*, 2021.

Liu, W., Li, M., Tan, G., and Jia, W. Mario: Near zero-cost activation checkpointing in pipeline parallelism. In *Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming*, pp. 197–211, 2025.

Liu, Y., Li, S., Fang, J., Shao, Y., Yao, B., and You, Y. Colossal-auto: Unified automation of parallelization and activation checkpoint for large-scale models. *arXiv preprint arXiv:2302.02599*, 2023a.

Liu, Z., Cheng, S., Zhou, H., and You, Y. Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency. In *Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis*. ACM, November 2023b. doi: 10.1145/3581784.3607073. URL http://dx.doi.org/10.1145/3581784.3607073.

Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Patwary, M., Korthikanti, V. A., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Zaharia, M. Efficient large-scale language model training on gpu clusters using megatron-lm, 2021.
Patil, S. G., Jain, P., Dutta, P., Stoica, I., and Gonzalez, J. Poet: Training neural networks on tiny devices with integrated rematerialization and paging. In International Conference on Machine Learning, pp. 17573–17583. PMLR, 2022.

Qi, P., Wan, X., Amar, N., and Lin, M. Pipeline parallelism with controllable memory, 2024a. URL https://arxiv.org/abs/2405.15362.

Qi, P., Wan, X., Huang, G., and Lin, M. Zero bubble (almost) pipeline parallelism. In The Twelfth International Conference on Learning Representations, 2024b.

Rajbhandari, S., Ruwase, O., Rasley, J., Smith, S., and He, Y. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pp. 1–14, 2021.

Ren, J., Rajbhandari, S., Aminabadi, R. Y., Ruwase, O., Yang, S., Zhang, M., Li, D., and He, Y. Zero-offload: Democratizing billion-scale model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551–564, 2021.

Wang, B., Xu, Q., Bian, Z., and You, Y. Tesseract: Parallelize the tensor parallelism efficiently. In Proceedings of the 51st International Conference on Parallel Processing, pp. 1–11, 2022.

Zhao, X., Hellard, T. L., Eyraud, L., Gusak, J., and Beament, O. Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch, 2023a.

Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., Xu, M., Wright, L., Shojanazeri, H., Ott, M., Shleifer, S., et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304.11277, 2023b.

# A APPENDIX

## A.1 Environment

All experiments are executed with PyTorch 2.7.1, CUDA 12.6 and NCCL 2.26.2. The nodes used contain 4 NVIDIA H100 GPUs 80 GB each, connected with Infiniband (4 400-Gbps links per GPU node).

## A.2 Integer Linear Programming Formulation

In this Section we present the ILP formulation for a model with a given partitioning and schedule. This formulation is valid for schedules with any number of waves (number of stages per GPU).

### A.2.1 Notations

The input values of the ILP formulation use the following notations:

* $p$ is the number of GPUs, indexed by $r$;

* $s$ is the number of stages, indexed with $i$;

* $m$ is the number of microbatches, indexed with $j$;

* $c \in \{F, B, W\}$ is used to index the operation type among: forward, backward w.r.t. inputs and backward w.r.t. parameters;

* $M_{GPU}$ is the available memory of each GPU;

* $\pi_i \in [1, p]$ is the fixed assignment of stage $i$ to a GPU.

In the following, we use the notation $O = (i, j, c)$ to denote an operation of type $c$ for stage $i$ on microbatch $j$.

In order to represent the schedule (i.e. the forced order of operations), we use the following notations:

* $O' \rightsquigarrow O$ if $O$ is computed anytime after $O'$ on the same GPU;

* $O' \prec O$ if $O$ is the operation computed right after $O'$ on the same GPU;

We also denote $B_r$ and $E_r$ respectively the first and last computation executed on GPU $r$.

To summarize, we assume that the schedule is specified for each GPU $r$ by providing the order of the operations, viewed as:

$$B_r \prec O_1 \prec O_2 \prec \dots \prec O_k \prec E_r$$

and we denote $O' \rightsquigarrow O$ if operation $O'$ is performed at some point before $O$, ie $O' \prec \dots \prec O$.

We assume that a profiling of the model has been performed, which provides the following values:

* $M_{F_i}$ the size of the intermediate data from the $F$ operation of stage $i$;

* $M_{B_i}^{act}$ the size of intermediate activations from the $B$ operation of stage $i$;
* $M_{B_i}^{\text{grads}}$ the size of the intermediate gradients from the B operation of stage i;

* $P_{F_i}$, $P_{B_i}$ and $P_{W_i}$ the peak memory usage of each operation for stage i.

* $T_{\text{comm}_i}$ the communication time for the data sent between stage i and stage i + 1*;

* $T_{F_i}$, $T_{B_i}$ and $T_{W_i}$ the processing times of each operation for stage i.

### A.2.2 Decision variables

This formulation considers the possibility to use different rematerialization strategies with 3 binary variables for each micro batch j and stage i:

* $R_{i,j}^F$: Full forward rematerialization: keep nothing after the forward pass, then recompute everything just before the backward pass;

* $R_{i,j}^{B_{\text{act}}}$: Full rematerialization of backward activations: do not keep any activations after B, and recompute them just before W;

* $R_{i,j}^{B_{\text{grads}}}$: Full rematerialization of backward gradients: do not keep any gradients after B, and recompute them just before W.

In addition, the formulation contains a variable $E_O = E(i, j, c)$ for each stage i, each micro-batch j, and each operation c, which represents the ending time of the operation $O = (i, j, c)$.

### A.2.3 Derived values

The following notations can be computed once the above decision variables are known, and are defined here to simplify the formulation. For a specific operation c of stage i on microbatch j, denoted $O = (i, j, c)$:

* $T_O = T(i, j, c)$ is the duration of the operation;

* $M_O = M(i, j, c)$ is the memory stored after the end of the operation;

* $P_O = P(i, j, c)$ is the peak memory used during the operation. We assume that $P_{F_i} > P_{B_i} > P_{W_i}$, but the formulation can be adapted automatically if this is not the case.

These values can be written as linear expressions of the

decision variables as follows:

$$T(i, j, F) = T_{F_i} \qquad (1)$$

$$T(i, j, B) = T_{B_i} + T_{F_i} R_{i,j}^F \qquad (2)$$

$$T(i, j, W) = T_{W_i} + T_{F_i} R_{i,j}^{BA} + T_{B_i} R_{i,j}^{B_{\text{grads}}} \qquad (3)$$

$$M(i, j, F) = M_{F_i}(1 - R_{i,j}^F) \qquad (4)$$

$$M(i, j, B) = M_{B_i} - (M_{B_i}^{\text{act}} R_{i,j}^{B_{\text{act}}}) - (M_{B_i}^{B_{\text{grads}}} R_{i,j}^{B_{\text{grads}}}) \qquad (5)$$

$$M(i, j, W) = M_{W_i} = 0 \qquad (6)$$

$$P(i, j, F) = P_{F_i} \qquad (7)$$

$$P(i, j, B) = P_{B_i} + (P_{F_i} - P_{B_i}) R_{i,j}^F \qquad (8)$$

$$P(i, j, W) = P_{W_i} + (P_{F_i} - P_{W_i}) R_{i,j}^{B_{\text{act}}} \qquad (9)$$

Furthermore, we define $\Delta M$ the memory difference caused by an operation as:

$$\begin{aligned} \Delta M(i, j, F) &= M(i, j, F) \\ \Delta M(i, j, B) &= M(i, j, B) - M(i, j, F) \\ \Delta M(i, j, W) &= -M(i, j, B) \end{aligned}$$

### A.2.4 Complete formulation

We can now define the complete ILP formulation for our problem:

$$\min. \quad \max_{r \in [1, p]} E_{\mathcal{E}_r} - (E_{\mathcal{B}_r} - T_{\mathcal{B}_r}) \qquad (10)$$

$$\text{subj. to:} \qquad E_{\mathcal{B}_1} - T_{\mathcal{E}_1} = 0 \qquad (11)$$

$$\forall i \in [2, s], j \in [1, m], \qquad E(i, j, F) \ge E(i-1, j, F) + T(i, j, F) + T_{\text{comm}_{i-1}} \qquad (12)$$

$$\forall i \in [1, s-1], j \in [1, m], \qquad E(i, j, B) \ge E(i+1, j, B) + T(i, j, B) + T_{\text{comm}_i} \qquad (13)$$

$$\forall \mathcal{O}, \mathcal{O}' | \mathcal{O} \prec_{\pi} \mathcal{O}', \quad E_{\mathcal{O}} \ge E_{\mathcal{O}'} + T_{\mathcal{O}'} \qquad (14)$$

$$\forall i, j, \quad R_{i,j}^{B_{\text{grads}}} \le R_{i,j}^{B_{\text{act}}} \qquad (15)$$

$$\forall \mathcal{O} = (i, j, c), \quad M_{\text{GPU}} \ge P_{\mathcal{O}} + \sum_{\mathcal{O}' \sim \pi_i \mathcal{O}} \Delta M_{\mathcal{O}'} \qquad (16)$$

Objective (10) is the makespan. Constraint (11) asserts that the first operation starts at time 0.

Constraints (12) and (13) represent the dependencies between forward (resp. backward) passes on different GPUs for one microbatch. Constraint (14) ensures that each GPU computes only one operation at a time.

<footnote>*if i and i + 1 are placed on the same GPU, this time is 0</footnote>
As mentioned in Section 3.2, we only consider gradient recomputation if we also recompute activations, which is enforced by constraint (15).

Finally, constraint (16) ensures that the memory does not exceed the available GPU memory.

### A.2.5 How the solution is used

A solution to the above formulation provides a precise scheduling of each operation, with their starting time and end time. However, we do not attempt to enforce these timings in practical executions. We only extract from the solution the values of the binary variables $R_{i,j}^F$, $R_{i,j}^{BG}$ and $R_{i,j}^{BA}$, to specify which micro-batch is saved or recomputed on each GPU. The actual execution in ELF is performed asynchronously: each operation is processed when all required data is available and its previous operation in the schedule is finished. The recomputation is done right before the computation that needs it, i.e. B for forward recomputation, and W for backward recomputation.

## A.3 Extended Experiments

### A.3.1 Iteration time

Table 4 (resp. 5) details the benchmarked iteration times on 4 (resp. 16) GPUs for different models and sequence lengths, depending on the rematerialization strategy and base scheduler used.

### A.3.2 MFU

Tables 6, 7 and 8 show the MFU — as described in (Chowdhery et al., 2022) — corresponding to the iteration times reported in last section. Standard deviation is not displayed at all for readability.

The formulas used to compute MFU are described Equation 17. Model FLOPs are computed using PyTorch's built-in flop counter and normalized by the sequence length on a batch of size 1. GPU FLOP/s are taken from the manufacturer's specifications.

$$ \text{Batch Size} = \text{\# of Microbatches} \times \text{Microbatch Size} $$

$$ \begin{aligned} \text{Throughput} &= \frac{\text{Batch Size}}{\text{Iteration Time}} \\ \text{MFU} &= \frac{\text{Throughput} \times \text{Model FLOPs}}{\text{\# of GPUs} \times \text{GPU FLOP/s} \times \text{Batch Size}} \end{aligned} \tag{17} $$

## A.4 Greedy Algorithm

We use the algorithm described in (Korthikanti et al., 2023) Appendix C as a baseline to compare with. We detail here

<caption>**Algorithm 1 Greedy rematerialization strategy**</caption>

**Input:** schedule, p, $M_{GPU}$, $P_F$, $\Delta M$
**for** i = 1 **to** p **do**
    limit = $M_{GPU} - \max_{s|\pi_s=i} P_{F_s}$
    memory = 0
    **for** op in schedule[i] **do**
        if memory + $\Delta M op$ ≥ limit **then**
            markAsRecomputed(op)
        **end if**
        memory += $\Delta M op$
    **end for**
    **end for**

how exactly we use it and extend it to backward recomputation.

Algorithm 1 describes the greedy method as we interpret and use it for forward only rematerialization. On each GPU, we first compute the maximum memory that could be needed to compute a forward among all stages on this GPU. We then make sure to have at least this amount of memory after every F, by marking them as recomputed if the activations they produce would bring the memory above that threshold. All input variables are the same as formulated for the ILP in A.2.

For forwardonly rematerialization the decision is binary: an operation is either recomputed in its entirety or it is not recomputed at all. When backward rematerialization is considered, we can use two distinct recomputation levels (activations only and both activations and gradients). Therefore, the algorithm is modified as follows: (1) we also consider B computations to be marked as recomputed; (2) instead of marking an operation as recomputed once, we increment its recomputation level in a loop until the memorylimit condition is satisfied, progressing from no recomputation to recomputing activations only, and finally to recomputing both activations and gradients. These simple changes allow the algorithm to support backward rematerialization.
<table><caption>Table 4: Comparing iteration time for different rematerialization strategies and schedules on 4 GPUs. All values are in seconds. Standard deviation is shown if it exceeds $5 \cdot 10^{-2}$. The empty cells correspond to problems marked infeasible by the ILP solver, or for which the Greedy algorithm fails to return a solution. OOM cells raised an out of memory error during execution. The best (lowest) value within one schedule is written in bold. The best value for the whole column is framed.</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Strategy</th><th colspan="4">9.6B</th><th colspan="4">21B</th><th colspan="4">38.5B</th></tr><tr><th>4096</th><th>8192</th><th>12288</th><th>16384</th><th>2048</th><th>4096</th><th>6144</th><th>8192</th><th>1024</th><th>2048</th><th>3072</th><th>4096</th></tr></thead><tbody><tr><td rowspan="4">1F1B</td><td>No Remat</td><td>1.1</td><td>2.9</td><td>5.2</td><td>OOM</td><td>1.0</td><td>2.3</td><td>3.8</td><td>OOM</td><td>1.0</td><td>1.9</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-F</td><td>1.5</td><td>3.8</td><td>6.7</td><td>10.5</td><td>1.4</td><td>3.0</td><td>4.9</td><td>7.3</td><td>1.4</td><td>2.5</td><td>4.0</td><td>5.3</td></tr><tr><td>Greedy-F</td><td>1.1</td><td>2.9</td><td>5.5</td><td>9.7</td><td>1.0</td><td>2.3</td><td>4.1</td><td>6.7</td><td>1.0</td><td>1.9</td><td>3.6</td><td>4.9</td></tr><tr><td>ILP-F</td><td>1.1</td><td>2.9</td><td>5.2</td><td>8.7</td><td>1.0</td><td>2.3</td><td>3.8</td><td>6.1</td><td>1.0</td><td>1.9</td><td>3.2</td><td>4.6</td></tr><tr><td rowspan="4">1F1B-2W</td><td>No Remat</td><td>1.1</td><td>2.8</td><td>5.1</td><td>OOM</td><td>1.0</td><td>2.1</td><td>3.5</td><td>OOM</td><td>0.9</td><td>1.8</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-F</td><td>1.5</td><td>3.9</td><td>7.1</td><td>11.1</td><td>1.4</td><td>3.0</td><td>5.0</td><td>7.5</td><td>1.4</td><td>2.6</td><td>4.0</td><td>5.4</td></tr><tr><td>Greedy-F</td><td>1.1</td><td>2.8</td><td>5.4</td><td>-</td><td>1.0</td><td>2.1</td><td>-</td><td>-</td><td>0.9</td><td>1.8</td><td>3.4</td><td>-</td></tr><tr><td>ILP-F</td><td>1.1</td><td>2.8</td><td>5.1</td><td>8.4</td><td>1.0</td><td>2.1</td><td>3.6</td><td>5.7</td><td>0.9</td><td>1.8</td><td>3.0</td><td>4.3</td></tr><tr><td rowspan="6">ZB-H1</td><td>No Remat</td><td>1.1</td><td>2.8</td><td>5.1</td><td>OOM</td><td>1.0</td><td>2.1</td><td>OOM</td><td>OOM</td><td>0.9</td><td>1.8</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-F</td><td>1.4</td><td>3.7</td><td>6.6</td><td>OOM</td><td>1.4</td><td>2.8</td><td>4.8</td><td>OOM</td><td>1.3</td><td>2.4</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-FB</td><td>2.0</td><td>5.3</td><td>9.5</td><td>14.9</td><td>1.9</td><td>4.0</td><td>6.8</td><td>10.2</td><td>1.8</td><td>3.4</td><td>5.4</td><td>7.3</td></tr><tr><td>Greedy-F</td><td>1.1</td><td>2.8</td><td>6.2</td><td>-</td><td>1.0</td><td>2.1</td><td>-</td><td>-</td><td>0.9</td><td>2.0</td><td>-</td><td>-</td></tr><tr><td>Greedy-FB</td><td>1.1</td><td>2.8</td><td>6.6</td><td>11.9</td><td>1.0</td><td>2.1</td><td>5.1</td><td>8.3</td><td>0.9</td><td>2.0</td><td>4.2</td><td>6.3</td></tr><tr><td>ILP-F</td><td>1.1</td><td>2.8</td><td>5.8</td><td>-</td><td>1.0</td><td>2.1</td><td>-</td><td>-</td><td>0.9</td><td>1.8</td><td>-</td><td>-</td></tr><tr><td rowspan="6">ZB-H2</td><td>No Remat</td><td>1.0</td><td>OOM</td><td>OOM</td><td>OOM</td><td>0.9</td><td>OOM</td><td>OOM</td><td>OOM</td><td>0.9</td><td>OOM</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-F</td><td>1.4</td><td>OOM</td><td>OOM</td><td>OOM</td><td>1.3</td><td>OOM</td><td>OOM</td><td>OOM</td><td>1.3</td><td>OOM</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-FB</td><td>2.0</td><td>5.1</td><td>9.3</td><td>14.6</td><td>1.8</td><td>3.9</td><td>6.6</td><td>9.9</td><td>1.7</td><td>3.3</td><td>5.2</td><td>7.1</td></tr><tr><td>Greedy-F</td><td>1.0</td><td>-</td><td>-</td><td>-</td><td>0.9</td><td>-</td><td>-</td><td>-</td><td>0.9</td><td>-</td><td>-</td><td>-</td></tr><tr><td>Greedy-FB</td><td>1.0</td><td>-</td><td>6.5</td><td>11.8</td><td>0.9</td><td>2.4</td><td>4.7</td><td>8.2</td><td>0.9</td><td>2.2</td><td>4.2</td><td>6.3</td></tr><tr><td>ILP-F</td><td>1.0</td><td>-</td><td>-</td><td>-</td><td>0.9</td><td>-</td><td>-</td><td>-</td><td>0.9</td><td>-</td><td>-</td><td>-</td></tr><tr><td rowspan="6">ZB-V</td><td>No Remat</td><td>1.0</td><td>2.6</td><td>OOM</td><td>OOM</td><td>0.9</td><td>2.0</td><td>OOM</td><td>OOM</td><td>0.9</td><td>1.7</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-F</td><td>1.4</td><td>3.5</td><td>6.4</td><td>OOM</td><td>1.4</td><td>2.8</td><td>4.6</td><td>OOM</td><td>1.4</td><td>2.4</td><td>3.7</td><td>OOM</td></tr><tr><td>Uni-FB</td><td>1.6</td><td>4.1</td><td>7.3</td><td>11.4</td><td>1.5</td><td>3.2</td><td>5.3</td><td>7.8</td><td>1.5</td><td>2.7</td><td>4.2</td><td>5.7</td></tr><tr><td>Greedy-F</td><td>1.0</td><td>2.6</td><td>5.3</td><td>-</td><td>0.9</td><td>2.0</td><td>4.0</td><td>-</td><td>0.9</td><td>1.8</td><td>-</td><td>-</td></tr><tr><td>Greedy-FB</td><td>1.0</td><td>2.6</td><td>5.7</td><td>10.9</td><td>0.9</td><td>2.0</td><td>4.2</td><td>7.6</td><td>0.9</td><td>1.8</td><td>3.6</td><td>5.6</td></tr><tr><td>ILP-F</td><td>1.0</td><td>2.6</td><td>4.9</td><td>-</td><td>0.9</td><td>2.0</td><td>3.6</td><td>-</td><td>0.9</td><td>1.7</td><td>-</td><td>-</td></tr><tr><td></td><td>ILP-FB</td><td>1.0</td><td>2.6</td><td>4.9</td><td>8.1</td><td>0.9</td><td>2.0</td><td>3.6</td><td>5.6</td><td>0.9</td><td>1.7</td><td>3.0</td><td>4.3</td></tr></tbody></table>
<table><caption>Table 5: Comparing iteration time for different rematerialization strategies and schedules on 16 GPUs. All values are in seconds. Standard deviation is shown if it exceeds 5 · 10⁻². The empty cells correspond to problems marked infeasible by the ILP solver, or for which the Greedy algorithm fails to return a solution. OOM cells raised an out of memory error during execution. The best (lowest) value within one schedule is written in bold. The best value for the whole column is framed.</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Strategy</th><th colspan="4">9.6B</th><th colspan="4">21B</th><th colspan="4">38.5B</th></tr><tr><th>4096</th><th>8192</th><th>12288</th><th>16384</th><th>2048</th><th>4096</th><th>6144</th><th>8192</th><th>1024</th><th>2048</th><th>3072</th><th>4096</th></tr></thead><tbody><tr><td rowspan="4">1F1B</td><td>No Remat</td><td><b>2.0</b></td><td>4.2</td><td><b>6.9</b></td><td>OOM</td><td><b>1.8</b></td><td><b>3.2</b></td><td>5.0</td><td>OOM</td><td><b>1.7</b></td><td><b>2.7</b></td><td>4.0</td><td><b>5.1</b></td></tr><tr><td>Uni-F</td><td>2.4</td><td>5.2</td><td>8.4</td><td>12.6</td><td>2.3</td><td>4.0</td><td>6.1</td><td>8.9</td><td>2.4</td><td>3.4</td><td>5.1±0.07</td><td>6.5</td></tr><tr><td>Greedy-F</td><td><b>2.0</b></td><td><b>4.1</b></td><td>7.0</td><td><b>10.3</b></td><td><b>1.8</b></td><td><b>3.2</b></td><td><b>4.9</b></td><td>7.2</td><td><b>1.7</b></td><td>2.8</td><td><b>3.9</b></td><td><b>5.1</b></td></tr><tr><td>ILP-F</td><td><b>2.0</b></td><td>4.2</td><td>7.0</td><td>10.4</td><td><b>1.8</b></td><td><b>3.2</b></td><td><b>4.9</b></td><td><b>7.1</b></td><td><b>1.7</b></td><td><b>2.7</b></td><td><b>3.9</b></td><td><b>5.1</b></td></tr><tr><td rowspan="4">1F1B-2W</td><td>No Remat</td><td><b>2.4</b></td><td><b>4.5</b></td><td><b>7.2±0.05</b></td><td>OOM</td><td><b>2.4</b></td><td><b>3.7</b></td><td>5.6</td><td>OOM</td><td><b>2.2</b></td><td><b>3.0</b></td><td><b>4.1</b></td><td><b>5.1</b></td></tr><tr><td>Uni-F</td><td>3.0</td><td>5.5</td><td>9.1</td><td>13.8±0.05</td><td>3.2</td><td>4.8</td><td>7.3</td><td>10.6</td><td>3.0</td><td>4.1</td><td>5.5</td><td>6.9</td></tr><tr><td>Greedy-F</td><td><b>2.4</b></td><td><b>4.5</b></td><td><b>7.2</b></td><td>–</td><td><b>2.4</b></td><td>3.8</td><td><b>5.6</b></td><td>–</td><td>2.3</td><td><b>3.0±0.05</b></td><td><b>4.1</b></td><td>5.2</td></tr><tr><td>ILP-F</td><td><b>2.4</b></td><td><b>4.5</b></td><td><b>7.2</b></td><td><b>11.2±0.09</b></td><td><b>2.4</b></td><td>3.8</td><td><b>5.6</b></td><td><b>7.9</b></td><td>2.3</td><td><b>3.0</b></td><td><b>4.1</b></td><td><b>5.1</b></td></tr><tr><td rowspan="6">ZB-H1</td><td>No Remat</td><td>1.9</td><td>3.8</td><td><b>6.3</b></td><td>OOM</td><td><b>1.6</b></td><td><b>2.8</b></td><td><b>4.5</b></td><td>OOM</td><td><b>1.6</b></td><td><b>2.4</b></td><td><b>3.5</b></td><td><b>4.5</b></td></tr><tr><td>Uni-F</td><td>2.3</td><td>4.7</td><td>8.1</td><td>OOM</td><td>2.2</td><td>3.8</td><td>5.8</td><td>8.4</td><td>2.3</td><td>3.3</td><td>4.9</td><td>6.3</td></tr><tr><td>Uni-FB</td><td>2.7</td><td>5.9</td><td>10.5</td><td>16.5</td><td>2.7</td><td>4.7</td><td>7.5</td><td>11.2</td><td>2.8</td><td>4.2</td><td>6.5</td><td>8.1</td></tr><tr><td>Greedy-F</td><td>1.9</td><td>3.8</td><td>6.6</td><td>–</td><td><b>1.6</b></td><td><b>2.8</b></td><td><b>4.5</b></td><td>–</td><td><b>1.6</b></td><td><b>2.4</b></td><td><b>3.5</b></td><td>4.9</td></tr><tr><td>Greedy-FB</td><td><b>1.8</b></td><td>3.8</td><td>6.7</td><td>–</td><td><b>1.6</b></td><td><b>2.8</b></td><td><b>4.5</b></td><td>–</td><td><b>1.6</b></td><td><b>2.4</b></td><td><b>3.5</b></td><td>5.1</td></tr><tr><td>ILP-F</td><td>1.9</td><td>3.8</td><td>6.4</td><td>–</td><td><b>1.6</b></td><td>2.9</td><td><b>4.5</b></td><td>–</td><td><b>1.6</b></td><td>2.5</td><td><b>3.5</b></td><td><b>4.5</b></td></tr><tr><td rowspan="2">ZB-H2</td><td>Greedy-F</td><td><b>1.7</b></td><td>–</td><td>–</td><td>–</td><td>1.6</td><td>–</td><td>–</td><td>–</td><td><b>1.6</b></td><td><b>2.3</b></td><td>–</td><td>–</td></tr><tr><td>Greedy-FB</td><td><b>1.7</b></td><td>–</td><td>–</td><td>–</td><td>1.6</td><td>2.8</td><td>–</td><td>–</td><td><b>1.6</b></td><td><b>2.3</b></td><td>–</td><td>–</td></tr><tr><td rowspan="2">ZB-V</td><td>ILP-F</td><td><b>1.7</b></td><td>–</td><td>–</td><td>–</td><td>1.6</td><td>2.8</td><td>–</td><td>–</td><td><b>1.6</b></td><td><b>2.3</b></td><td>–</td><td>–</td></tr><tr><td>ILP-FB</td><td><b>1.7</b></td><td><b>3.7</b></td><td><b>6.6</b></td><td>11.0</td><td><b>1.5</b></td><td><b>2.7</b></td><td><b>4.5</b></td><td><b>7.2</b></td><td><b>1.6</b></td><td><b>2.3</b></td><td><b>3.7±0.05</b></td><td><b>5.0</b></td></tr><tr><td rowspan="6">ZB-V</td><td>No Remat</td><td><b>2.4</b></td><td>4.3</td><td>OOM</td><td>OOM</td><td><b>2.2</b></td><td>3.3</td><td><b>4.6</b></td><td>OOM</td><td>2.1</td><td><b>2.8</b></td><td><b>3.6</b></td><td><b>4.4</b></td></tr><tr><td>Uni-F</td><td>3.0</td><td>5.0</td><td>8.2</td><td>OOM</td><td>2.8</td><td>4.0</td><td>5.9</td><td>8.2</td><td>2.8</td><td>3.7</td><td>4.8</td><td>6.0</td></tr><tr><td>Uni-FB</td><td>3.0</td><td>5.0</td><td>8.5</td><td>12.9</td><td>2.9</td><td>4.2</td><td>6.4</td><td>9.1±0.05</td><td>2.9</td><td>3.9</td><td>5.1±0.06</td><td>6.5</td></tr><tr><td>Greedy-F</td><td><b>2.4</b></td><td>4.3</td><td>OOM</td><td>–</td><td><b>2.2</b></td><td>3.3</td><td>4.7</td><td>–</td><td>2.1</td><td><b>2.8</b></td><td><b>3.6</b></td><td>4.6</td></tr><tr><td>Greedy-FB</td><td>2.5</td><td>4.3</td><td>OOM</td><td>–</td><td><b>2.2±0.07</b></td><td>3.3</td><td>4.7±0.08</td><td>–</td><td>2.1</td><td><b>2.8</b></td><td><b>3.6</b></td><td>4.7</td></tr><tr><td>ILP-F</td><td>2.5</td><td><b>4.2</b></td><td><b>6.8</b></td><td>–</td><td><b>2.2</b></td><td><b>3.2</b></td><td><b>4.6</b></td><td>–</td><td><b>2.0</b></td><td><b>2.8</b></td><td><b>3.6</b></td><td>4.5</td></tr><tr><td rowspan="2">ZB-V</td><td>ILP-FB</td><td><b>2.4</b></td><td>4.3</td><td><b>6.8±0.10</b></td><td>11.6</td><td><b>2.2</b></td><td><b>3.2</b></td><td><b>4.6</b></td><td><b>6.5</b></td><td>2.1</td><td><b>2.8</b></td><td><b>3.6</b></td><td>4.5±0.05</td></tr></tbody></table>
<table><caption>Table 7: Comparing MFU for different rematerialization strategies and schedules on 8 GPUs. The empty cells correspond to problems marked infeasible by the ILP solver, or for which the Greedy algorithm fails to return a solution. OOM cells raised an out of memory error during execution. The best (highest) value within one schedule is written in bold. The best value for the whole column is framed.</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Strategy</th><th colspan="4">9.6B</th><th colspan="4">21B</th><th colspan="4">38.5B</th></tr><tr><th>4096</th><th>8192</th><th>12288</th><th>16384</th><th>2048</th><th>4096</th><th>6144</th><th>8192</th><th>1024</th><th>2048</th><th>3072</th><th>4096</th></tr></thead><tbody><tr><td rowspan="4">1F1B</td><td>No Remat</td><td>34.8</td><td>35.8</td><td>34.3</td><td>OOM</td><td><strong>42.7</strong></td><td>45.0</td><td>43.4</td><td>OOM</td><td>38.5</td><td>46.1</td><td><strong>45.9</strong></td><td><strong>47.2</strong></td></tr><tr><td>Uni-F</td><td>28.9</td><td>26.7</td><td>26.9</td><td>25.4</td><td>31.1</td><td>33.8</td><td>34.2</td><td>32.3</td><td>26.8</td><td>33.7</td><td>34.5</td><td>35.9</td></tr><tr><td>Greedy-F</td><td><strong>35.5</strong></td><td>35.8</td><td>34.1</td><td>31.3</td><td>42.4</td><td><strong>45.1</strong></td><td><strong>43.6</strong></td><td>39.4</td><td><strong>38.6</strong></td><td><strong>46.2</strong></td><td>45.8</td><td>45.1</td></tr><tr><td>ILP-F</td><td>35.1</td><td><strong>35.9</strong></td><td><strong>34.4</strong></td><td><strong>32.3</strong></td><td>42.0</td><td>44.9</td><td><strong>43.6</strong></td><td><strong>41.8</strong></td><td>38.5</td><td>46.0</td><td>45.6</td><td>47.0</td></tr><tr><td rowspan="4">1F1B-2W</td><td>No Remat</td><td><strong>32.7</strong></td><td>34.0</td><td>32.9</td><td>OOM</td><td><strong>36.1</strong></td><td>43.4</td><td><strong>43.2</strong></td><td>OOM</td><td>32.1</td><td>42.6</td><td><strong>45.3</strong></td><td>OOM</td></tr><tr><td>Uni-F</td><td>24.8</td><td>25.7</td><td>24.6</td><td>23.5</td><td>26.2</td><td>32.0</td><td>31.9</td><td>30.3</td><td>22.1</td><td>31.8</td><td>33.2</td><td>34.4</td></tr><tr><td>Greedy-F</td><td>32.5</td><td><strong>34.5</strong></td><td>33.0</td><td>28.3</td><td>35.8</td><td>43.4</td><td><strong>43.2</strong></td><td>38.0</td><td><strong>32.3</strong></td><td><strong>42.8</strong></td><td>45.1</td><td>44.7</td></tr><tr><td>ILP-F</td><td>32.6</td><td>34.2</td><td><strong>33.1</strong></td><td><strong>30.3</strong></td><td>35.8</td><td><strong>43.6</strong></td><td><strong>43.2</strong></td><td><strong>40.5</strong></td><td>32.2</td><td><strong>42.8</strong></td><td>45.1</td><td><strong>46.8</strong></td></tr><tr><td rowspan="6">ZB-H1</td><td>No Remat</td><td>39.0</td><td><strong>37.1</strong></td><td><strong>35.5</strong></td><td>OOM</td><td><strong>45.2</strong></td><td><strong>48.3</strong></td><td>46.1</td><td>OOM</td><td><strong>41.9</strong></td><td><strong>49.9</strong></td><td><strong>49.4</strong></td><td>OOM</td></tr><tr><td>Uni-F</td><td>29.1</td><td>28.7</td><td>27.6</td><td>OOM</td><td>32.3</td><td>35.3</td><td>35.1</td><td>OOM</td><td>27.5</td><td>34.5</td><td>35.0</td><td>37.3</td></tr><tr><td>Uni-FB</td><td>23.7</td><td>21.3</td><td>19.7</td><td>18.3</td><td>25.6</td><td>26.6</td><td>25.3</td><td>23.7</td><td>23.0</td><td>26.7</td><td>26.3</td><td>26.9</td></tr><tr><td>Greedy-F</td><td>38.4</td><td>36.6</td><td>32.5</td><td>-</td><td>44.8</td><td>47.9</td><td>42.5</td><td>-</td><td>41.2</td><td>49.5</td><td>48.8</td><td>41.8</td></tr><tr><td>Greedy-FB</td><td>38.8</td><td>36.9</td><td>30.7</td><td>26.0</td><td><strong>45.2</strong></td><td>48.1</td><td>40.1</td><td>33.2</td><td>41.6</td><td>49.8</td><td>49.1</td><td>40.6</td></tr><tr><td>ILP-F</td><td>39.0</td><td><strong>37.1</strong></td><td>35.1</td><td>-</td><td><strong>45.2</strong></td><td>48.2</td><td>46.1</td><td>-</td><td><strong>41.9</strong></td><td><strong>49.9</strong></td><td>49.1</td><td>44.0</td></tr><tr><td rowspan="2">ZB-H2</td><td>ILP-FB</td><td><strong>39.1</strong></td><td>37.0</td><td>34.8</td><td><strong>29.8</strong></td><td>45.1</td><td>48.2</td><td><strong>46.2</strong></td><td><strong>39.6</strong></td><td>41.6</td><td>49.6</td><td>49.3</td><td><strong>46.6</strong></td></tr><tr><td>No Remat</td><td>41.9</td><td>OOM</td><td>OOM</td><td>OOM</td><td>46.7</td><td>OOM</td><td>OOM</td><td>OOM</td><td><strong>42.3</strong></td><td><strong>52.1</strong></td><td>OOM</td><td>OOM</td></tr><tr><td rowspan="6">ZB-V</td><td>Uni-F</td><td>32.8</td><td>31.7</td><td>OOM</td><td>OOM</td><td>33.3</td><td>37.2</td><td>OOM</td><td>OOM</td><td>28.3</td><td>38.0</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-FB</td><td>24.1</td><td>22.0</td><td>20.2</td><td>18.9</td><td>25.5</td><td>27.3</td><td>26.0</td><td>24.5</td><td>23.5</td><td>27.3</td><td>27.5</td><td>28.0</td></tr><tr><td>Greedy-F</td><td>41.8</td><td>-</td><td>-</td><td>-</td><td>46.3</td><td>-</td><td>-</td><td>-</td><td>41.8</td><td>48.1</td><td>-</td><td>-</td></tr><tr><td>Greedy-FB</td><td>42.1</td><td>37.5</td><td>31.9</td><td>-</td><td>46.8</td><td>46.2</td><td>-</td><td>34.1</td><td>42.0</td><td>48.4</td><td>43.8</td><td>40.2</td></tr><tr><td>ILP-F</td><td>42.1</td><td>39.4</td><td>-</td><td>-</td><td><strong>47.2</strong></td><td>48.7</td><td>-</td><td>-</td><td><strong>42.3</strong></td><td>51.2</td><td>-</td><td>-</td></tr><tr><td>ILP-FB</td><td><strong>42.4</strong></td><td><strong>39.5</strong></td><td><strong>34.7</strong></td><td><strong>28.7</strong></td><td>46.8</td><td><strong>49.1</strong></td><td><strong>44.3</strong></td><td><strong>37.7</strong></td><td>42.1</td><td>50.9</td><td><strong>46.1</strong></td><td><strong>43.9</strong></td></tr><tr><td rowspan="6">ZB-V</td><td>No Remat</td><td><strong>31.6</strong></td><td><strong>36.3</strong></td><td><strong>36.1</strong></td><td>OOM</td><td>37.4</td><td>47.5</td><td><strong>48.3</strong></td><td>OOM</td><td><strong>35.4</strong></td><td>48.0</td><td>51.3</td><td>OOM</td></tr><tr><td>Uni-F</td><td>26.1</td><td>29.0</td><td>28.0</td><td>OOM</td><td>27.6</td><td>36.0</td><td>36.2</td><td>OOM</td><td>24.4</td><td>34.8</td><td>37.3</td><td>39.1</td></tr><tr><td>Uni-FB</td><td>25.7</td><td>27.1</td><td>25.6</td><td>24.1</td><td>26.4</td><td>33.2</td><td>32.6</td><td>30.8</td><td>22.7</td><td>31.9</td><td>33.5</td><td>34.6</td></tr><tr><td>Greedy-F</td><td>31.4</td><td>35.9</td><td>34.1</td><td>-</td><td><strong>37.6</strong></td><td>47.6</td><td>46.2</td><td>-</td><td>35.0</td><td>47.6</td><td>51.0</td><td>-</td></tr><tr><td>Greedy-FB</td><td>31.4</td><td>36.0</td><td>33.9</td><td>OOM</td><td>37.2</td><td><strong>47.7</strong></td><td>44.3</td><td>34.8</td><td>35.1</td><td>48.0</td><td>51.3</td><td>44.2</td></tr><tr><td>ILP-F</td><td>31.5</td><td>36.1</td><td>35.3</td><td>-</td><td>37.4</td><td>47.5</td><td>48.2</td><td>-</td><td>35.3</td><td><strong>48.1</strong></td><td><strong>51.4</strong></td><td>-</td></tr><tr><td rowspan="2">ZB-V</td><td>ILP-FB</td><td>31.2</td><td>36.2</td><td>35.4</td><td><strong>32.9</strong></td><td>37.3</td><td><strong>47.7</strong></td><td>48.2</td><td><strong>43.8</strong></td><td>35.3</td><td>48.0</td><td>51.3</td><td><strong>50.5</strong></td></tr></tbody></table>
<table><caption>Table 8: Comparing MFU for different rematerialization strategies and schedules on 16 GPUs. The empty cells correspond to problems marked infeasible by the ILP solver, or for which the Greedy algorithm fails to return a solution. OOM cells raised an out of memory error during execution. The best (highest) value within one schedule is written in bold. The best value for the whole column is framed.</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Strategy</th><th colspan="4">9.6B</th><th colspan="4">21B</th><th colspan="4">38.5B</th></tr><tr><th>4096</th><th>8192</th><th>12288</th><th>16384</th><th>2048</th><th>4096</th><th>6144</th><th>8192</th><th>1024</th><th>2048</th><th>3072</th><th>4096</th></tr></thead><tbody><tr><td rowspan="4">1F1B</td><td>No Remat</td><td>26.5</td><td>28.6</td><td><strong>28.6</strong></td><td>OOM</td><td><strong>31.4</strong></td><td><strong>36.0</strong></td><td>36.7</td><td>OOM</td><td>28.7</td><td><strong>36.3</strong></td><td>38.7</td><td>40.7</td></tr><tr><td>Uni-F</td><td>22.3</td><td>23.0</td><td>23.5</td><td>22.9</td><td>24.1</td><td>29.0</td><td>29.7</td><td>28.6</td><td>20.7</td><td>29.2</td><td>30.2</td><td>32.2</td></tr><tr><td>Greedy-F</td><td><strong>27.3</strong></td><td><strong>29.1</strong></td><td>28.3</td><td><strong>28.1</strong></td><td>31.1</td><td>35.9</td><td><strong>37.0</strong></td><td>35.2</td><td>28.5</td><td>36.2</td><td>38.9</td><td>40.9</td></tr><tr><td>ILP-F</td><td>27.1</td><td>28.5</td><td>28.5</td><td>27.7</td><td>30.8</td><td>35.7</td><td>36.8</td><td><strong>35.8</strong></td><td><strong>28.8</strong></td><td><strong>36.3</strong></td><td><strong>39.0</strong></td><td><strong>41.0</strong></td></tr><tr><td rowspan="4">1F1B-2W</td><td>No Remat</td><td><strong>22.3</strong></td><td>26.5</td><td><strong>27.4</strong></td><td>OOM</td><td>22.9</td><td><strong>30.9</strong></td><td>32.6</td><td>OOM</td><td><strong>21.7</strong></td><td>32.9</td><td><strong>37.5</strong></td><td><strong>40.7</strong></td></tr><tr><td>Uni-F</td><td>17.7</td><td>21.7</td><td>21.8</td><td>21.0</td><td>17.4</td><td>24.0</td><td>24.8</td><td>24.0</td><td>16.2</td><td>24.3</td><td>27.7</td><td>30.3</td></tr><tr><td>Greedy-F</td><td>22.1</td><td><strong>26.7</strong></td><td><strong>27.4</strong></td><td>-</td><td>23.0</td><td>30.8</td><td><strong>32.7</strong></td><td>-</td><td>21.6</td><td>32.8</td><td>37.3</td><td>40.5</td></tr><tr><td>ILP-F</td><td>22.2</td><td>26.6</td><td><strong>27.4</strong></td><td><strong>25.9</strong></td><td><strong>23.1</strong></td><td>30.8</td><td><strong>32.7</strong></td><td><strong>32.1</strong></td><td>21.6</td><td><strong>33.1</strong></td><td>37.1</td><td><strong>40.7</strong></td></tr><tr><td rowspan="6">ZB-H1</td><td>No Remat</td><td>28.6</td><td>31.8</td><td><strong>31.3</strong></td><td>OOM</td><td><strong>33.6</strong></td><td><strong>40.8</strong></td><td><strong>40.6</strong></td><td>OOM</td><td>30.0</td><td><strong>41.3</strong></td><td>43.7</td><td><strong>45.9</strong></td></tr><tr><td>Uni-F</td><td>23.7</td><td>25.3</td><td>24.5</td><td>OOM</td><td>25.0</td><td>30.7</td><td>31.2</td><td>30.2</td><td>21.0</td><td>30.6</td><td>31.2</td><td>33.2</td></tr><tr><td>Uni-FB</td><td>19.9</td><td>20.2</td><td>18.8</td><td>17.6</td><td>20.7</td><td>24.6</td><td>24.2</td><td>22.7</td><td>17.3</td><td>23.9</td><td>23.5</td><td>25.8</td></tr><tr><td>Greedy-F</td><td>28.8</td><td>31.7</td><td>30.1</td><td>-</td><td>33.4</td><td>40.6</td><td>40.0</td><td>-</td><td>30.2</td><td>40.9</td><td>43.4</td><td>42.7</td></tr><tr><td>Greedy-FB</td><td><strong>29.2</strong></td><td>31.7</td><td>29.6</td><td>-</td><td>33.5</td><td>40.7</td><td>40.3</td><td>-</td><td>30.3</td><td>41.0</td><td>43.7</td><td>40.9</td></tr><tr><td>ILP-F</td><td>28.6</td><td>31.9</td><td>31.2</td><td>-</td><td>33.4</td><td>40.5</td><td>40.4</td><td>-</td><td>30.1</td><td>40.3</td><td>43.6</td><td><strong>45.9</strong></td></tr><tr><td rowspan="2">ZB-H2</td><td>ILP-FB</td><td>29.1</td><td><strong>32.0</strong></td><td>31.2</td><td><strong>27.9</strong></td><td>33.5</td><td>40.5</td><td>40.2</td><td><strong>36.7</strong></td><td><strong>30.4</strong></td><td>41.2</td><td><strong>43.8</strong></td><td>45.7</td></tr><tr><td>No Remat</td><td><strong>31.2</strong></td><td>OOM</td><td>OOM</td><td>OOM</td><td>35.2</td><td>OOM</td><td>OOM</td><td>OOM</td><td>30.8</td><td><strong>43.6</strong></td><td>OOM</td><td>OOM</td></tr><tr><td rowspan="6">ZB-V</td><td>Uni-F</td><td>24.4</td><td>OOM</td><td>OOM</td><td>OOM</td><td>25.3</td><td>32.9</td><td>OOM</td><td>OOM</td><td>21.8</td><td>31.4</td><td>OOM</td><td>OOM</td></tr><tr><td>Uni-FB</td><td>20.5</td><td>20.8</td><td>19.4</td><td>18.1</td><td>21.3</td><td>25.6</td><td>24.8</td><td>23.6</td><td>18.5</td><td>24.6</td><td>24.5</td><td>27.1</td></tr><tr><td>Greedy-F</td><td>30.9</td><td>-</td><td>-</td><td>-</td><td>35.2</td><td>-</td><td>-</td><td>-</td><td>30.5</td><td>43.3</td><td>-</td><td>-</td></tr><tr><td>Greedy-FB</td><td><strong>31.2</strong></td><td>-</td><td>-</td><td>-</td><td>35.3</td><td>40.8</td><td>-</td><td>-</td><td><strong>31.0</strong></td><td><strong>43.6</strong></td><td>-</td><td>-</td></tr><tr><td>ILP-F</td><td>30.9</td><td>-</td><td>-</td><td>-</td><td>35.4</td><td>41.5</td><td>-</td><td>-</td><td>30.6</td><td>43.5</td><td>-</td><td>-</td></tr><tr><td>ILP-FB</td><td>30.7</td><td><strong>32.7</strong></td><td><strong>30.1</strong></td><td><strong>26.4</strong></td><td><strong>35.6</strong></td><td><strong>42.6</strong></td><td><strong>40.2</strong></td><td><strong>35.5</strong></td><td><strong>31.0</strong></td><td><strong>43.6</strong></td><td><strong>41.4</strong></td><td><strong>41.4</strong></td></tr><tr><td rowspan="6">ZB-V</td><td>No Remat</td><td>21.9</td><td>28.0</td><td>OOM</td><td>OOM</td><td>25.1</td><td>35.1</td><td><strong>39.3</strong></td><td>OOM</td><td>23.7</td><td><strong>35.9</strong></td><td>42.6</td><td><strong>47.0</strong></td></tr><tr><td>Uni-F</td><td>18.1</td><td>23.8</td><td>24.3</td><td>OOM</td><td>19.8</td><td>28.6</td><td>30.7</td><td>31.0</td><td>17.4</td><td>27.3</td><td>32.0</td><td>35.0</td></tr><tr><td>Uni-FB</td><td>18.0</td><td>23.9</td><td>23.4</td><td>22.3</td><td>18.9</td><td>27.5</td><td>28.5</td><td>27.9</td><td>16.7</td><td>25.7</td><td>29.7</td><td>32.1</td></tr><tr><td>Greedy-F</td><td><strong>22.0</strong></td><td>28.1</td><td>OOM</td><td>-</td><td>25.2</td><td>35.2</td><td>39.0</td><td>-</td><td>23.6</td><td><strong>35.9</strong></td><td>42.4</td><td>45.0</td></tr><tr><td>Greedy-FB</td><td>21.7</td><td>27.9</td><td>OOM</td><td>-</td><td>25.1</td><td>35.5</td><td>38.4</td><td>-</td><td>23.7</td><td>35.7</td><td>42.5</td><td>44.3</td></tr><tr><td>ILP-F</td><td>21.7</td><td><strong>28.2</strong></td><td><strong>29.1</strong></td><td>-</td><td><strong>25.5</strong></td><td><strong>35.6</strong></td><td><strong>39.3</strong></td><td>-</td><td><strong>23.9</strong></td><td><strong>35.9</strong></td><td><strong>42.7</strong></td><td>46.1</td></tr><tr><td rowspan="2">ZB-V</td><td>ILP-FB</td><td><strong>22.0</strong></td><td>27.9</td><td>29.0</td><td><strong>25.0</strong></td><td>25.3</td><td><strong>35.6</strong></td><td>39.2</td><td><strong>39.3</strong></td><td>23.7</td><td>35.7</td><td><strong>42.7</strong></td><td>46.4</td></tr></tbody></table>
## A.5 Details on the Transformer Block

To help understanding our approach and notations, we detail how they are applied to a standard Transformer Block as in (Korthikanti et al., 2023). A **forward pass** through such a block, $F(\text{Block})$ corresponds to the propagation through the sequence of the following operations:

$$ 
\begin{aligned}
\xrightarrow{a_{in}} F(\text{Normalization}_{Attn}) & \xrightarrow{a_{n1}} F(\text{Linear}_{QKV}) & \xrightarrow{a_{qkv}} F(\text{SDPA}) \\
\xrightarrow{a_{sdp a}} F(\text{Linear}_{proj}) & & \xrightarrow{a_{proj}} F(\text{Dropout}_{Attn}) \\
\xrightarrow{a_{d1}} \oplus \xrightarrow{a_{attn}=a_{in}+a_{d1}} & & \\
F(\text{Normalization}_{MLP}) & \xrightarrow{a_{n2}} F(\text{Linear}_{up}) & \xrightarrow{a_{up}} F(\text{GELU}) \\
\xrightarrow{a_{gelu}} F(\text{Linear}_{down}) & & \xrightarrow{a_{down}} F(\text{Dropout}_{MLP}) \\
\xrightarrow{a_{d2}} \oplus \xrightarrow{a_{mlp}=a_{attn}+a_{d2}} & &
\end{aligned}
$$

where $F(\cdot)$ represents the forward pass through the layer, and SDPA (Scaled Dot Product Attention) refers to the computation of Dropout (Softmax($\frac{QK^T}{scale}$) · V). Together SDPA and Linear$_{QKV}$ are usually called an attention layer (Korthikanti et al., 2023).

**Forward rematerialization options** A forward rematerialization option for the block defines, which activations we forget during the forward pass $F(\text{Block})$, and which keep between $F(\text{Block})$ and activation gradients computation phase $B(\text{Block})$. For example, a decision to forget $a_{gelu}$ will require additional recomputation of the GELU layer during the backward pass.

In the current paper, on the level of the Transformer block we apply one of the following forward rematerialization options:

* no recomputation: keep all activations needed for backward

* full recomputation: forget all intermediate activations of the Transformer block (note that $a_{in}$ is kept in memory).

**Backward rematerialization options** In the backward pass, for each layer with trainable parameters whose gradients computations can be decoupled, we consider two operations: $B(\cdot)$ to compute activation gradients and $W(\cdot)$ to compute parameter gradients. For the layers which are not decoupled, we use the notation $BW(\cdot)$ to represent the combined computation of gradients with respect to activation and weights. For layers with no trainable parameters, $B(\cdot)$ states for all computations during backward.

Let $\delta_z = \frac{\partial \mathcal{L}}{\partial z}$ be the gradient of the loss with respect to tensor $z$. Taking into account this notation, we can depict the backward pass (without rematerialization) through the Transformer block as two sequences of operations corresponding to $B(\text{Block})$ and $W(\text{Block})$:

### Backward Pass: B(Block)

$$ \text{B(Residual)} \qquad \delta_{\text{attn}}^{(\text{mlp})} = \delta_{\text{mlp}}, \quad \delta_{\text{d2}} = \delta_{\text{mlp}} \quad (18) $$

$$ \text{B(Dropout}_{\text{MLP}}) \qquad \delta_{\text{down}} = \text{Dropout}'(a_{\text{down}}) \cdot \delta_{\text{d2}} \quad (19) $$

$$ \text{B(Linear}_{\text{down}}) \qquad \delta_{\text{gelu}} = \delta_{\text{down}} W_{\text{down}}^{\top} \quad (20) $$

$$ \text{B(GELU)} \qquad \delta_{\text{up}} = \text{GELU}'(a_{\text{up}}) \cdot \delta_{\text{gelu}} \quad (21) $$

$$ \text{B(Linear}_{\text{up}}) \qquad \delta_{\text{n2}} = \delta_{\text{up}} W_{\text{up}}^{\top} \quad (22) $$

$$ \text{BW(LayerNorm}_{\text{MLP}}) \qquad \delta_{\text{attn}}^{(\text{norm})} = \text{LayerNorm}'(a_{\text{attn}}) \cdot \delta_{\text{n2}}, \quad \frac{\partial \mathcal{L}}{\partial(\text{LNparams})} \quad (23) $$

$$ \text{B(Residual merge)} \qquad \delta_{\text{attn}} = \delta_{\text{attn}}^{(\text{mlp})} + \delta_{\text{attn}}^{(\text{norm})} \quad (24) $$

$$ \text{B(Residual merge)} \qquad \delta_{\text{attn}} = \delta_{\text{attn}}^{(\text{mlp})} + \delta_{\text{attn}}^{(\text{norm})} \quad (25) $$

$$ \text{B(Residual)} \qquad \delta_{\text{atn}}^{(\text{attn})} = \delta_{\text{attn}}, \quad \delta_{\text{d1}} = \delta_{\text{attn}} \quad (26) $$

$$ \text{B(Dropout}_{\text{Attn}}) \qquad \delta_{\text{a}_{\text{proj}}} = \text{Dropout}'(a_{\text{proj}}) \cdot \delta_{\text{d1}} \quad (27) $$

$$ \text{B(Linear}_{\text{proj}}) \qquad \delta_{\text{sdpa}} = \delta_{\text{a}_{\text{proj}}} W_{\text{proj}}^{\top} \quad (28) $$

$$ \text{B(SDPA)} \qquad (\delta_Q, \delta_K, \delta_V) = \text{SDPA}'(Q, K, V, \delta_{\text{sdpa}}), \quad \delta_{\text{qkv}} = \text{concat}(\delta_Q, \delta_K, \delta_V) \quad (29) $$

$$ \text{B(Linear}_{\text{QKV}}) \qquad \delta_{\text{n1}} = \delta_{\text{qkv}} W_{\text{qkv}}^{\top} \quad (30) $$

$$ \text{BW(LayerNorm}_{\text{Attn}}) \qquad \delta_{\text{atn}}^{(\text{norm})} = \text{LayerNorm}'(a_{\text{in}}) \cdot \delta_{\text{n1}}, \quad \frac{\partial \mathcal{L}}{\partial(\text{LNparams})} \quad (31) $$

$$ \text{B(Residual merge)} \qquad \delta_{\text{atn}} = \delta_{\text{attn}}^{(\text{norm})} + \delta_{\text{atn}}^{(\text{norm})} \quad (32) $$

$$ \text{B(Residual merge)} \qquad \delta_{\text{atn}} = \delta_{\text{attn}}^{(\text{norm})} + \delta_{\text{atn}}^{(\text{norm})} \quad (33) $$

### Backward Pass: W(Block)

$$ \text{W(Linear}_{\text{down}}) \qquad \frac{\partial \mathcal{L}}{\partial W_{\text{down}}} = a_{\text{gelu}}^{\top} \delta_{\text{down}} \quad (34) $$

$$ \text{W(Linear}_{\text{up}}) \qquad \frac{\partial \mathcal{L}}{\partial W_{\text{up}}} = a_{\text{n2}}^{\top} \delta_{\text{up}} \quad (35) $$

$$ \text{W(Linear}_{\text{proj}}) \qquad \frac{\partial \mathcal{L}}{\partial W_{\text{proj}}} = a_{\text{sdpa}}^{\top} \delta_{\text{a}_{\text{proj}}} \quad (36) $$

$$ \text{W(Linear}_{\text{QKV}}) \qquad \frac{\partial \mathcal{L}}{\partial W_{\text{qkv}}} = a_{\text{n1}}^{\top} \delta_{\text{a}_{\text{qkv}}} \quad (37) $$

Benefiting from presence of layers with decoupled gradient computations, namely, Linear$_{QKV}$, Linear$_{proj}$, Linear$_{up}$, and Linear$_{down}$ we propose the following **backward based rematerialization** options:

* no rematerialization

* activation recomputation: forget activations that serve as input to the $W(\text{Block})$, and
* gradient and activation recomputation: additionally forget activation gradients that serve as input to the $W$ (Block).

Thus, these options allow to reduce memory occupied between $B$ (Block) and $W$ (Block) phases. For example, within the activation recomputation option, for the layer $\text{Linear}_{\text{down}}$ activation $a_{\text{gelu}}$ (output of $F(\text{GELU})$) will be forgotten, and within the activation and gradient recomputation option $a_{\text{gelu}}$ will be forgotten along with $\delta_{\text{down}} = \frac{\partial \text{loss}}{\partial a_{\text{down}}}$ (output of $B(\text{Dropout}_{\text{MLP}})$).

## A.6 Closer Look at Recomputations

Given profiling statistics of real models, we analyze solutions of different rematerialization strategies (ILP and Greedy based) for a wide range of GPU budget constraints by estimating corresponding schedule execution time and analyzing statistics of recomputations.

On an example of 9.6B model with sequence length 4096 and GPU we will see how differs Greedy and ILP solutions.

Consider for simplicity a non decoupled 1F1B-2W schedule. Let us look at Figure 5 and Figure 6, which show how many microbatches require activation recomputations between F and B per GPU per stage. We can observe that ILP solution prioritizes recomputing first stages on the same GPU, while Greedy - second stages. Also, Greedy algorithm fails to find solutions for tighter memory budgets. Figure 7 and 8 on the example of ZB-V show that similar patterns are observed for decoupled schedules.

## A.7 Scaling Pipeline Schedules with Rematerialization

Given a schedule, ILP-based rematerialization strategy always performs better or on par with the Greedy one. Still, more insights are required to understand which schedule is more suitable to use with rematerialization given a model, sequence length and GPU memory budget. On the examples of 9.6B (Figure 10) and 38.5B (Figure 9) models we investigate relative performance of decoupled (ZB-V) and not decoupled (1F1B-2W) schedules under different memory budgets and sequence length. We can observe that given a model, sequence length and number of GPUs, the decoupled schedules are faster when there is enough memory to execute the schedule without rematerialization (the most right point of each line on subplots). They also remain better in the presence of ILP-based rematerialization strategies before the memory limit becomes very tight (intersection of blue and orange lines corresponding to the same number of GPUs). Not decoupled schedules then start outperforming them.

From the same plots we can also observe that when the sequence length increases, the fraction of memory dedicated to activations becomes larger. In that case the curves for

different number of GPUs become closer, which means that increasing the number of GPUs is less effective at reducing memory footprint. In particular, when two lines of the same colour come to cross, it means that scaling the number of GPUs increases the throughput by the exact same factor.

## A.8 Memory Decrease - Time Increase Trade-off

Another aspect we would like to discuss is the trade-off between memory limit decrease and time overhead for a given model and pipeline schedule in presence of ILP-based rematerialization strategy. Figure 12 and Figure 11 provide simulation results for models 9.6B and 38.5B. One line of the subplot corresponds to a fixed sequence length and number of GPUs, with each point corresponding to one ILP-based rematerialization strategy and showing what is the relative memory decrease and iteration time increase with respect to the strategy without any rematerialization. The closer the point of the curve to the left bottom corner, the better time overhead - memory savings trade-off the rematerialization strategy provides.

From all subplots we can observe two patterns: increasing number of GPUs for a fixed sequence length yields better trade-off. So does increasing the sequence length given a fixed number of GPUs. The first behaviour pattern can be explained by the fact that splitting the same model across more devices increases the total number of pipeline stages and therefore uses less layers per stage, which in turn yields to finer grained rematerialization strategy. The second behaviour pattern could be explained by the increasing contribution of activation memory to the total memory footprint when the sequence length increases, thus causing more benefits from applying rematerialization.
Statistics of Recomputations
Model: 9.6B, SeqLength: 4096, GPUs: 4, Schedule: Greedy-F-1f1b-2w

<figure><img src="image_7.png" /><figcaption>Figure 5: Statistics of recomputations found by Greedy algorithm for 9.6B model, sequence length 4096, trained on 4 GPUs with 1F1B-2W schedule under different memory constraints. Each subplot shows how many microbatches are recomputed on each GPU given memory limit.</figcaption></figure>

Statistics of Recomputations
Model: 9.6B, SeqLength: 4096, GPUs: 4, Schedule: ILP-F-1f1b-2w

<figure><img src="image_8.png" /><figcaption>Figure 6: Statistics of recomputations found by ILP for 9.6B model, sequence length 4096, trained on 4 GPUs with 1F1B-2W schedule under different memory constraints. Each subplot shows how many microbatches are recomputed on each GPU given memory limit.</figcaption></figure>
<figure><img src="image_9.png" /><figcaption>Figure 7: Statistics of recomputations found by Greedy algorithm for 9.6B model, sequence length 4096, trained on 8 GPUs with ZB-V schedule under different memory constraints. Each subplot shows how many microbatches are recomputed on each GPU given memory limit.</figcaption></figure>
Statistics of Recomputations
Model: 9.6B, SeqLength: 4096, GPUs: 8, Schedule: ILP-FB-zbv

<figure><img src="image_10.png" /><figcaption>Figure 8: Statistics of recomputations found by ILP for 9.6B model, sequence length 4096, trained on 8 GPUs with ZB-V schedule under different memory constraints. Each subplot shows how many microbatches are recomputed on each GPU given memory limit. We can observe that the number of recomputations between *B* and *W* increases when memory limit decreases.</figcaption></figure>
<figure><img src="image_11.png" /><figcaption>Figure 9: Estimated iteration time of ZB-V and 1F1B-2W schedules under different GPU memory constraints. When training on 8 and 16 GPUs, global batch size is 2x and 4x larger than the one used for the training on 4 GPUs.</figcaption></figure>
<figure><img src="image_12.png" /><figcaption>Figure 10: Estimated iteration time of ZB-V and 1F1B-2W schedules under different GPU memory constraints. When training on 8 and 16 GPUs, global batch size is 2x and 4x larger than the one used for the training on 4 GPUs.</figcaption></figure>
<figure><img src="image_13.png" /></figure>

<figure><img src="image_14.png" /><figcaption>Figure 11: Trade-off between memory limit decrease and time overhead for a given model and pipeline schedule in presence of ILP-based rematerialization strategy.</figcaption></figure>
<figure><img src="image_15.png" /></figure>

<figure><img src="image_16.png" /><figcaption>Figure 12: Trade-off between memory limit decrease and time overhead for a given model and pipeline schedule in presence of ILP-based rematerialization strategy.</figcaption></figure>