TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

arXiv:2611.04897v1 [cs.DC] 3 Nov 2026

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training Kaihui Ren∗

Yifan Zhao†

Zejian Wang∗

Yingtong Xiong

renkaihui@kuaishou.com KlingAI Research Beijing, China

zhaoyifan14@kuaishou.com KlingAI Research Shanghai, China

wangzejian@kuaishou.com KlingAI Research Shanghai, China

xiongyingtong@kuaishou.com KlingAI Research Shanghai, China

Xiaoling Zhou∗

Daohong Xu

Chenhao Guan∗

Zhikang Mo

zhouxiaoling03@kuaishou.com KlingAI Research Beijing, China

xudaohong@kuaishou.com KlingAI Research Beijing, China

guanchenhao@kuaishou.com KlingAI Research Shanghai, China

mozhikang@kuaishou.com KlingAI Research Shanghai, China

Jihu Guo∗

Zhuoru Ma

Yufeng Li

Lin Zhang

guojihu@kuaishou.com KlingAI Research Shanghai, China

mazhuoru@kuaishou.com KlingAI Research Shanghai, China

liyufeng06@kuaishou.com KlingAI Research Shanghai, China

zhanglin22@kuaishou.com KlingAI Research Shanghai, China

Yiqiao Liao‡

Chengru Song‡

liaoyiqiao@kling.ai KlingAI Research Beijing, China

songchengru@kuaishou.com KlingAI Research Beijing, China recomputation time

Abstract Selective activation checkpointing (SAC) decides which activations to save for the backward pass and which to recompute: each choice is a point in the plane of peak memory and recomputation time, and the goal is that plane’s Pareto frontier. Current practice falls short of this goal. Hand-tuned policies fix a single point that does not adapt when the model, batch or kernels change; PyTorch’s SAC reruns each block’s forward pass in program order, recomputing values the backward pass never reads; and the memory budget of torch.compile treats recomputation as free or forbidden and selects saved tensors by bytes alone, an assumption that fails for the custom kernels that produce many of a large model’s activations. We present TapeSAC, which automates SAC in eager PyTorch without editing the model’s forward code. TapeSAC records each block’s forward pass at the dispatcher and, during backward, replays only the producers of the values autograd consumes, so saving a tensor eliminates every producer needed only for it. A planner selects the saved tensors from measured recomputation costs under a single limit on step-peak memory. On Qwen3.8-27B and Qwen3.635B-A3B, replay alone reduces step time by 3.0% and 1.7% versus full checkpointing at no greater memory cost; memory-aware planning raises these gains to 6.3% and 13.9% under a 90%-of-devicememory limit. The evaluated hand-tuned policies fall short of this memory–recomputation frontier and, on Qwen3.6-35B-A3B, run out of memory at larger batches. Over 500 steps on real data, TapeSAC reduces training time by 8–11%; on Qwen3.5-397B, it is 2.7% faster than a model-specific hand-tuned policy while using 14% less memory.

∗ This work was conducted during the author’s internship at KlingAI. † Project leader. ‡ Corresponding authors.

memory limit

full checkpointing hand-tuned policy

same policy, larger batch OOM

TapeSAC

frontier

peak memory

Figure 1: Schematic of the peak-memory–recomputationtime trade-off in activation checkpointing. Full checkpointing and a hand-tuned policy lie above the Pareto frontier; TapeSAC targets a frontier point under the memory limit.

Keywords Selective activation checkpointing, Tensor-level replay, Memoryaware planning, Eager PyTorch, Large-model training, Mixture-ofexperts

1

Introduction

Training large models relies on several forms of parallelism. Dataparallel sharding distributes parameters, gradients and optimizer state across ranks [19, 27], while tensor, sequence, context and expert parallelism partition the computation and, with it, the activations held by each rank [11, 12, 14, 21]. Even with all of them enabled, activation checkpointing [4] remains necessary: parallelism reduces per-rank activation memory, but not enough to fit, so intermediate tensors are discarded after the forward pass and recomputed when the backward pass requires them. Full checkpointing reruns each block’s forward pass during backward, adding

Kling Infra Team

Table 1: Comparison of activation-checkpointing mechanisms: who decides what to save, what is recomputed, how time and memory are estimated, and what each requires. Mechanism

Who decides what to save

What is recomputed

Time and memory estimates

Requires

Full checkpointing

nobody (block inputs only) a hand-written policy

the whole block, from its input, in program order the whole block, from its input, except operators whose outputs are saved operators the partitioner does not save, inside the compiled backward the operators of a planned schedule

none

eager execution

none

eager execution

fixed rules, FLOPs or compile-time timing; saved bytes modeled or profiled time; modeled memory

a captured graph

PyTorch SAC torch.compile memory budget Static and trace-based planners

automatic

TapeSAC

automatic

automatic

only the producers of values the backward pass reads, starting from the nearest saved tensors

roughly one forward pass per step; selective activation checkpointing (SAC) [10] instead saves the intermediates that are cheap to store and expensive to regenerate, and recomputes the rest. Each choice of what to save determines a point in a twodimensional space: the peak memory of the rank and the recomputation time added to the backward pass. The objective is the Pareto frontier of this space—for any amount of available memory, the choice that minimizes recomputation (Figure 1). In practice, however, SAC is configured with a fixed policy, hand-tuned for a particular model, sequence length and kernel set. Such a policy yields one operating point per configuration and offers no means of trading memory for recomputation time; when the model, sequence length or kernels change, that point moves in unplanned ways and may run out of memory (OOM). When a policy fails to carry over, the fallback is full checkpointing, which forgoes the trade altogether (§4.4). The automatic option built into PyTorch is the memory budget of torch.compile’s partitioner [1]. It handles time and memory in separate stages rather than optimizing them jointly: ban rules reduce each operator’s recomputation time to zero or infinity, a minimum cut then selects saved tensors by bytes alone, and when the result exceeds the budget an additive knapsack decides which banned operators to re-allow. This pricing is accurate only if every allowed recomputation fuses into a backward kernel for free, which does not hold for the custom kernels and unfusable operators that produce many of a large model’s activations; the budget, moreover, bounds saved bytes rather than the step’s peak (§2.3). In real training workloads, graph capture also fails before partitioning, because TorchDynamo encounters graph breaks. Eager-mode checkpointing avoids graph capture, but still leaves two problems unresolved. First, the save set must be chosen by hand: PyTorch’s selective checkpointing executes a hand-written policy. Second, PyTorch’s checkpointing recomputes each block as a whole, in program order. When backward first requests a saved tensor, it reruns the block’s forward pass, skipping calls only if their outputs were cached. Consequently, saving a tensor avoids recomputing its own operator but not its upstream producers. Operators whose outputs are never consumed by backward also run again, and recomputed values remain live together instead of being materialized only when needed.

measured time; the step’s peak memory, from measured must-only steps

a known graph or trace, and a runtime that executes the schedule eager execution

Avoiding this unnecessary work requires two pieces of information: which forward values the backward pass consumes, and which operators produce those values. We found that eager PyTorch already exposes both through the extension points used by its selective checkpointing mechanism. The dispatcher observes every operator a block executes, and autograd’s saved-tensor hooks report which forward values the backward pass may consume. Accordingly, TapeSAC uses them to record each checkpointed block’s forward pass, as it executes, into a tape in static single-assignment (SSA) form. During backward it replays only the producers of the values autograd requests, when it requests them, stops at saved tensors, and frees each replayed value after its last use. Recomputation thus becomes tensor-level: what is replayed is determined by the tensors autograd consumes, not by the operators of the block, and saving a tensor eliminates the recomputation of every producer needed only for it. Moreover, no changes to the model’s forward code are needed; kernels that bypass the dispatcher require only one-time registration as custom operators (§3.3). On top of replay, a planner chooses what to save, and two properties make this choice difficult. The value of saving a tensor depends on what else is saved, because recomputation chains overlap, so candidates must be evaluated together on the graph rather than summed one by one; and the step’s peak is set in the backward pass, largely by model states, communication buffers and allocator fragmentation rather than by activations. TapeSAC therefore prices candidates on the recorded graph, with recomputation costs measured on the actual kernels, and constrains the plan by the step’s peak as the caching allocator sees it. The user specifies a single memory limit, as a fraction of device memory; sweeping it traces an empirical frontier. We evaluate TapeSAC on Qwen3.8-27B and Qwen3.6-35B-A3B and validate it on Qwen3.5-397B. Two findings stand out (Figure 2). First, tensor-level replay recomputes less than block recomputation: it replays only the minimal subgraph that the consumed values need, skipping dead nodes (operators whose outputs no consumed value depends on) and stopping at saved tensors. With the save set of full checkpointing, where little is saved, it shortens the step by 1.7–3.0% at the same memory. Replaying values on demand can also lower peak memory: reserved memory is 1.5% lower on the 27B model with the save set of full checkpointing, and 2.1%

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

25.5 full checkpointing

step time (s)

31

−3.0% time −1.5% memory

time

hand-tuned

hand-tuned

−4.8% time −11% memory

24.5 −2.7% time −14% memory 24

9 70

80

90

reserved memory, rank 0 (% of device)

(a) Qwen3.8-27B

hand-tuned

9.5

29 60

−1.8% time

25

−4.7% time

10 30

full checkpointing

full checkpointing

−1.7%

10.5

100

20

30

40

50

60

70

reserved memory, rank 0 (% of device)

(b) Qwen3.6-35B-A3B

70

75

80

85

90

reserved memory, rank 0 (% of device)

(c) Qwen3.5-397B

Figure 2: Measured step time and reserved memory under PyTorch SAC and TapeSAC. Light arrows keep the save set of full checkpointing and replace block recomputation with replay; dark arrows replace the hand-tuned policy with a planned save set at no more memory. Faint lines are the TapeSAC frontiers. lower on Qwen3.5-397B with that of its hand-tuned policy. Second, automatic planning replaces the hand-tuned policy: because saving a tensor cuts off the replay of its producers, the planner prices each candidate by the subgraph it removes, and as the memory limit varies, its plans trace an empirical frontier of diminishing returns; at 90% of device memory they shorten the step by 6.3% and 13.9% relative to full checkpointing. The hand-tuned policies we evaluate lie off this frontier, by up to 5.1% at the same memory, and every hand-tuned policy we evaluate on the 35B model runs out of memory when the batch grows, while TapeSAC at the same memory limit keeps its gain. Over 500 steps on real data, TapeSAC shortens training by 8–11%, and on Qwen3.5-397B it is 2.7% faster than its hand-tuned policy while using 14% less memory. Table 1 contrasts TapeSAC with existing mechanisms, including static and trace-based planners (§2.4). Overall, we summarize our main contributions as follows:

• Automatic solving of the save set. A hand-tuned policy fixes which tensors to keep, so its memory moves with the configuration and can run out. TapeSAC needs only a memory limit, as a fraction of device memory: its planner solves for the saved tensors automatically and solves again when the configuration changes, and varying the limit traces the memory–recomputation frontier (§3.4, §4.4). • Recompute only what the backward pass needs, when it needs it. PyTorch’s checkpointing recomputes the whole block from its input as soon as backward requests its first saved tensor. TapeSAC recomputes at the tensor level instead, which brings two gains. First, a value is recomputed only when autograd requests it, starting from the nearest saved tensors rather than from the block’s input, and is freed after its last use; recomputed values no longer stay alive together, which can lower the step’s peak memory. Second, TapeSAC skips dead nodes: saving a tensor makes every producer needed only for it unnecessary, and some operators, such as the down projection of a multilayer perceptron (MLP), are recomputed by PyTorch only to re-register their saved inputs. Skipping them shortens the step at no memory cost, even with the save set of full checkpointing.

TapeSAC reproduces the recomputed values bitwise for any plan (§3.2, §4.2, §4.3). • Planning on measured costs and the step’s peak. Under replay, saving a tensor eliminates every producer needed only for it, so the values of candidates do not add up; and the step’s peak is set largely by model states, communication buffers and fragmentation rather than by activations. TapeSAC therefore evaluates candidates jointly on the recorded graph, prices them with measured kernel times, and constrains the step’s peak as the allocator sees it rather than the saved bytes (§3.4, §4.4). • Directly usable in eager PyTorch. The dispatcher and the saved-tensor hooks of eager PyTorch already expose the executed operator graph and the values the backward pass consumes. Building on them, TapeSAC needs no graph capture, no change to PyTorch and no edit to the model’s forward code; a kernel that bypasses the dispatcher becomes replayable after a one-time registration (§3.1, §3.3).

2

Background and Related Work

This section first formalizes the checkpointing problem (§2.1). It then examines selective checkpointing in PyTorch (§2.2) and the memory budget of torch.compile (§2.3), stating their limitations one at a time. Finally, it reviews other rematerialization systems (§2.4) and identifies the remaining gaps (§2.5).

2.1

SAC as a Pareto Problem

Activation checkpointing is applied to a module, typically a transformer layer; we call each forward call of such a module a checkpointed block. For a block, let 𝐺 be its forward dataflow graph, 𝐿 its inputs, including its parameters, and 𝑃 its demand set, the forward values that the backward pass may consume. A save set 𝑆 is the set of values saved across the forward/backward boundary; every other value in 𝑃 is regenerated during backward from 𝑆 ∪ 𝐿. A choice of save sets for all blocks determines two quantities: 𝑇 , the recomputation time added to the backward pass, and 𝑀, the peak device memory of the training step. The goal is the Pareto frontier of these choices: for a given memory limit, the feasible save sets

Kling Infra Team

with minimal 𝑇 . Neither quantity is a sum over individual values, as the rest of this subsection explains. Which calls must be recomputed depends on 𝑆. To regenerate a value in 𝑃 that is not saved, its producer must run, and so must the producers of that call’s inputs, and so on back, stopping at inputs that are saved or are block inputs, that is, values in 𝑆 ∪𝐿. Every other call need not be recomputed; we call it a dead node. A call is dead in two cases. In the first, nothing the backward pass reads depends on its output: for example, a matrix multiplication whose output only enters an addition, since the gradients of the multiplication need its inputs and those of the addition need neither operand. In the second, its output serves only to compute saved values: if 𝑎 = 𝜙 (𝑥) and 𝑏 = 𝜓 (𝑎), and the backward pass reads only 𝑏, saving 𝑏 makes both 𝜙 and 𝜓 dead. A save set therefore removes the recomputation of every call that it makes dead, and 𝑇 is the time of the calls that remain. Whether a call is dead depends on the set as a whole. The time that a save set removes is therefore not the sum of the times that its values would remove one at a time, and once some values are saved, saving one more removes a different amount of time than it would have at the start. We define the return of saving a value as the recomputation time it removes, given the values already saved, divided by the bytes it keeps. In the example above, saving 𝑎 first removes 𝑓 ; but once 𝑏 is saved, 𝑓 is already dead, and saving 𝑎 removes nothing. Conversely, if 𝑎 also feeds a value 𝑐 = 𝑘 (𝑎) that the backward pass reads, saving 𝑏 first removes only 𝑔, because 𝑐 still needs 𝑓 ; once 𝑐 is saved, saving 𝑏 removes both 𝑔 and 𝑓 . The return of a value thus changes as the save set grows, and the save set must be evaluated as a whole. Memory does not add up either. At the moment of the peak 𝑀, the device holds much more than saved activations: model states (parameters, gradients and optimizer state), communication buffers, allocator fragmentation, and the temporary tensors of the block being recomputed. A saved value raises 𝑀 only if it is still alive at that moment; a value already freed by then does not. Whether saving a value raises the peak, and by how much, therefore cannot be read from its size alone; it depends on the step as it actually runs.

2.2

Selective Checkpointing in PyTorch

Selective checkpointing in PyTorch works in two steps. In the forward pass, a dispatch mode (TorchDispatchMode) intercepts every operator call in the block and asks a hand-written policy (create_ selective_checkpoint_contexts) whether to cache the call’s output. In the backward pass, when autograd first needs one of the block’s saved values, the block’s forward pass is run again from its input: cached calls return the outputs saved in the forward pass, and all other calls execute again. We call this way of recomputing block recomputation. The design has three limitations. Limitation 1: The save set is not solved automatically. The policy must be written by hand, deciding operator by operator which outputs to cache. The annotation library torch_remat [15] replaces the per-operator policy with save or recompute annotations on regions of the model’s source code, but the choice is still made by hand, and the regions marked for recomputation are still rerun in full. A hand-written policy fixes which tensors are kept, not how much memory they take: when the configuration changes,

its memory changes with it and can run out (§4.4), and every new configuration or memory limit requires rewriting the policy. Limitation 2: Recomputation always starts from the block’s input. Block recomputation runs from the first operator onward and regenerates, in one pass, every value that is not cached; all of them then stay in memory until autograd reads them. Yet the saved values already provide starting points inside the block: each value could be regenerated from the nearest saved values upstream of it, only when autograd needs it, and freed after its last use. Values that autograd reads only near the end of the block’s backward pass would then not occupy memory from its start, which can lower the peak (§2.1). Limitation 3: Dead nodes are recomputed. Block recomputation runs every call that is not cached, whether or not the backward pass needs its result, so it recomputes both kinds of dead node defined in §2.1. Calls that the backward pass does not need at all still run: on Qwen3.6-35B-A3B, for example, block recomputation reruns the experts’ down projection in every mixture-of-experts (MoE) layer, although the backward pass does not read its output (§4.3). Calls needed only by cached values also still run, because block recomputation starts from the block’s input (Limitation 2); caching a value therefore removes only the call that produces it, not the calls upstream of it. The early_stop option removes only the calls after the last value that the backward pass needs has been regenerated, at the end of the block, and no dead node before that point. Limitations 2 and 3 share a cause: block recomputation has neither the block’s dataflow graph 𝐺 nor the set 𝑃 of values the backward pass consumes (§2.1), so it can tell neither where a value could be regenerated from nor whether a call is needed at all. Replay, which executes on demand only the producers of the consumed values and stops at saved values, needs both; §3.2.3 formalizes the difference in cost.

2.3

The Memory Budget of torch.compile

The memory budget of torch.compile chooses the save set automatically. Its partitioner decides what to save on the joint forward– backward graph [1]. It was designed for a compiler that fuses cheap operators into backward kernels, and it assumes that recomputing such operators is free and that other operators should not be recomputed at all. By default, it therefore does not recompute matrix multiplications, attention or custom operators, and among the remaining operators it uses a minimum cut, as in fusion-aware checkpointing [7], to pick the save set with the fewest saved bytes, weighted by heuristics. The memory budget (activation_memory_budget, between 0 and 1) asks for fewer saved bytes than this default. The partitioner first relaxes some of its restrictions; if the saved bytes are still above the budget, a 0-1 knapsack that weighs each operator’s size against its estimated runtime chooses which expensive operators to recompute, and the bytes-only cut is rerun with those operators allowed [18]. The budget mode thus applies a cut built for free recomputation to exactly the case where recomputation is not free: time enters only through the knapsack, one operator at a time. The design has four limitations.

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

query copy matmul gate query copy matmul gate query copy matmul gate saved for backward

(a) Save the gate: matmul recomputed, 2.0 ms removed.

(b) Save the query copy: matmul recomputed, 1.1 ms removed.

(c) Save both: matmul dead, 6.0 ms removed.

recomputed

dead, not recomputed

Figure 3: A joint effect in a full-attention block of Qwen3.635B-A3B. One matrix multiplication produces both the query and the output gate; it is dead only when both are saved. Times are the replay time removed per block.

Limitation 4: The graph of real training is not captured whole. The budget applies only to a graph that torch.compile has captured. On the three models we train, the graph breaks before the partitioner runs, so the budget cannot be applied at all (§4.6). Where a break splits a model into several graphs, each graph is partitioned on its own, and no decision covers the training step as a whole. Limitation 5: Recomputation time is not measured on the training step. By default, the knapsack uses an operator’s count of floating-point operations (FLOPs) in place of its time. FLOPs cannot stand in for time: how long a kernel takes in training depends on much more than its arithmetic, for example on the machine, on some devices being slower than others, and on communication that overlaps the kernel and slows it down. The other estimator does time operators, but during compilation, outside the training step. Recomputation time has to be measured on the training step itself. Limitation 6: The memory limit is not the step’s peak. The budget is a fraction of the bytes that one graph saves for backward by default, not an amount of device memory. Whether a step fits, however, depends on the peak memory of the whole step (§2.1), and much of that peak lies outside the budget: model states, communication buffers, the allocator’s fragmentation and the activations of other graphs. The peak also depends on when values are alive, not only on how many bytes are saved. It falls inside the backward window of one block: by then, the blocks whose backward pass has finished have freed their saved values, and the block itself adds the temporaries of its own recomputation. A count of saved bytes sees neither. A user therefore cannot derive the budget from the memory available; on controlled models, even a budget of zero leaves the peak unchanged (§4.6). Planning needs the peak measured on the training step, including where in the backward pass it falls. Limitation 7: Dead nodes and joint effects are not captured. The knapsack scores each operator by its own estimated time and adds the scores up. The return of saving a value, however, is neither limited to the call that produces it nor fixed (§2.1): it includes the time of the upstream calls that saving the value makes dead, and it depends on what else is saved. Two values saved together can remove more time than the sum of what each removes alone. In Figure 3, from Qwen3.6-35B-A3B, saving the query copy or the gate alone removes 1.1 or 2.0 ms of replay time per block, whereas

saving both removes 6.0 ms, because the matrix multiplication that produces them becomes dead. A fixed, additive score can express neither effect, and the cut, although it sees the whole graph, counts only bytes. No part of the partitioner therefore evaluates the time that the finally chosen save set removes. Limitations 5–7 matter most when memory is scarce enough that matrix multiplications and custom operators must be recomputed, which is the regime of the models studied here; §4.6 quantifies their effect on Qwen3-8B, which the compiler can capture, and on controlled models.

2.4

Other Rematerialization Systems

Static-graph planners. Revolve and sublinear-memory checkpointing place checkpoints on chains [4, 5], and Rotor solves heterogeneous chains by dynamic programming [3]. Checkmate, Moccasin and XEngine schedule rematerialization on general graphs with integer or constraint programming [2, 8, 20], while MODeL reorders operators and assigns tensor addresses to lower the peak [22]. For large language models, Yuan et al. enumerate save sets within a transformer layer to obtain a memory–computation frontier and search it jointly with the parallel configuration [24]. These planners assume a known graph or layer structure and a runtime that executes the chosen schedule. In contrast, TapeSAC records the concrete graph of each eager invocation and keeps autograd’s backward. Trace-based and runtime planners. Rockmate traces the forward pass on an example input, partitions the trace, and generates a module that drives the backward pass according to its schedule [26]; HiRemate extends this approach with hierarchical planning [6]. Mimose predicts per-layer memory online and caches a checkpoint plan for each input size [13]; however, it plans whole layers and recomputes each layer from its input (Limitations 2 and 3). TapeSAC, by contrast, needs no ahead-of-time trace, keeps autograd’s backward, and records and replays at the tensor level. Estimation-based planners. Mist traces a symbolic-shape graph, constructs a backward graph to estimate time and memory, and co-optimizes activation checkpointing with the parallel configuration [28]. PyTorch’s experimental SAC tools estimate operator costs on fake tensors, either with a roofline model of peak device throughput or by benchmarking each operator that the dispatcher sees [17]. Under an estimated peak, they decide which modules to checkpoint and how much of their activations to discard; within each module, a knapsack whose value is the sum of operator runtimes selects the operators to save. PyTorch’s SAC then executes the result with block recomputation and thus inherits Limitations 2 and 3. Like the memory budget, these planners estimate costs rather than measure them on the training step (Limitation 5), plan under an estimated peak rather than the measured one (Limitation 6), and score each operator by its own runtime (Limitation 7). In contrast, TapeSAC measures the concrete invocation, evaluates saved tensors jointly on the recorded graph, and replays rather than recomputing whole blocks. Reactive rematerialization. DTR and Capuchin evict tensors under memory pressure and regenerate them from recorded lineage [9, 16], and Coop evicts contiguous tensors together to limit the allocator fragmentation that tensor-wise eviction causes [25]. These systems

Kling Infra Team

decide at run time what to evict. TapeSAC instead plans a repeatable operating point from measured costs under the step’s peak.

2.5

Summary

Table 1 contrasts these mechanisms. PyTorch’s SAC leaves the choice to a hand-written policy (Limitation 1) and recomputes whole blocks (Limitations 2 and 3). The compiler chooses automatically, but only on captured graphs (Limitation 4), without measuring time or the peak on the training step (Limitations 5 and 6), and without knowing the return of what it saves (Limitation 7). Rematerialization planners require a static graph or trace, estimate costs rather than measure them, take over the backward pass, or plan whole layers. To the best of our knowledge, no prior mechanism for eager PyTorch chooses save sets at the tensor level from measured kernel costs under the step’s peak memory while replaying only the producers that autograd’s demand requires.

3

Methodology

TapeSAC has two layers. The recording layer captures each checkpointed block as a tape of the operator calls that actually ran and, during backward, replays only what autograd needs; whichever tensors are saved, backward receives the values the forward pass produced. The planning layer chooses those tensors: under a memory limit, it finds the plan with the least recomputation time, from costs and peak memory measured on the training step itself. Recording makes every choice safe; planning makes the choice good. Sections 3.2 and 3.3 describe recording, and Section 3.4 planning.

3.1

Overview

Eager PyTorch routes every ATen operator, every functional collective and every library-registered custom operator through the dispatcher, where a TorchDispatchMode observes each call with its concrete arguments and results; PyTorch’s own SAC is built on this mechanism. In the same forward pass, saved_tensors_hooks report which tensors autograd saves for backward. Together they provide, for each block invocation, the forward graph 𝐺 and the demand set 𝑃 of §2.1, for the computation that actually executed and without ahead-of-time capture, except for fused kernels that are not registered custom operators (§3.3). Figure 4 illustrates recording and planning on a gated MLP block, and Algorithm 1 gives one training step. To adopt TapeSAC, the training code swaps its checkpoint wrapper for TapeSAC’s, and a model adapter reports each block’s region type, the kind of block that the planner budgets separately (§3.4.1), and its token count. In every step, each block records its tape and saves the planned and mandatory values during forward, and the executor replays on demand during backward (§3.2). Plans come from profiling a few early steps (Figure 4b): after 𝑊 warm-up steps, the next 𝑅 steps (𝑊 = 𝑅 = 2 in all experiments) time the operators (reference steps) and measure peak memory (sample steps), and rank 0 then solves a plan for each range of token counts that occurs. Profiling is repeated when token counts move far from the profiled ones (§3.4).

3.2

Tensor-Level Record and Replay

Each checkpointed block runs its forward pass as usual, under two hooks that eager PyTorch provides. A dispatch mode records

every operator call of the block into a tape. Autograd’s saved-tensor hooks see every tensor that autograd saves for backward: TapeSAC keeps the tensors that must be saved or that the plan saves, and replaces every other one with a value identifier, so that, as under full checkpointing, the tensor is freed after forward. In backward, when autograd unpacks an identifier, TapeSAC recomputes that value from the tape, running only the operators it depends on. The rest of this section describes the tape (§3.2.1), the values that forward keeps (§3.2.2), and how replay recomputes the others (§3.2.3). 3.2.1 Recording the tape. While a block runs forward, the dispatch mode records every operator call, in order, as a node (op, args, in_vids, out_vids): the operator, its arguments with each tensor replaced by its value identifier, and the identifiers of its tensor inputs and outputs, written in(𝑛) and out(𝑛) below. Every value is either the output of exactly one node, its producer prod(𝑣), or a leaf : an input of the block or a parameter, which the block uses but does not compute. The tape is thus the dataflow graph 𝐺 of §2.1, recorded from the calls that actually ran. We write 𝑉 for all values and 𝐿 ⊆ 𝑉 for the leaves. The block of Figure 4a gives six nodes; 𝑥 and the weights are leaves, and 𝑔 is produced by mm #0. The tape records how each value was computed, that is, which operator produced it from which inputs, but not the value’s data. Apart from saved values (§3.2.2), it holds no reference to any intermediate tensor; such a reference would keep the tensor alive past forward and defeat checkpointing. Three details keep this bookkeeping exact. An in-place update creates a new value, so a node that runs after it sees the updated tensor as a different value from the original. A view shares memory with its base, so their bytes are counted once. Parameters are recorded by name and fetched again in backward, so that under fully sharded data parallelism (FSDP2 [27]) backward does not keep the unsharded copies of forward. 3.2.2 Values that forward keeps. Backward does not need every forward value, only those that autograd saves for its gradient formulas. They form the demand set 𝑃 ⊆ 𝑉 (§2.1). In Figure 4a, 𝑃 contains ℎ, 𝑔, 𝑠, 𝑢 and 𝑚: the gradient of silu needs 𝑔, that of mul needs 𝑠 and 𝑢, and so on. It does not contain 𝑦, which no gradient formula needs. Autograd hands every value of 𝑃 to the saved-tensor hooks. For a value in the save set 𝑆, the tape holds the tensor; for a leaf, it holds the block input or the parameter’s name; for every other value, autograd keeps only the identifier, and the tensor is freed when forward no longer uses it. A few values must always be saved, because replaying them is unsafe. They form the mandatory set, which every save set contains: 𝑆 0 = 𝑉coll ∪ 𝑉dyn ∪ 𝑉alloc,

𝑆 ⊇ 𝑆0 .

(1)

𝑉coll holds the outputs of collective communication, which replay would have to communicate again. 𝑉dyn holds outputs whose size depends on the data, such as that of nonzero, which replay could allocate only after the CPU waits for the GPU. 𝑉alloc holds buffers filled by a kernel that the tape cannot see (§3.3), which replay would recreate empty. A block or step that saves only 𝑆 0 is must-only. The rest of 𝑆 comes from a plan, which the planner of §3.4 solves once and later steps apply. Those steps’ tensors do not exist yet when the plan is solved, so the plan names each value by its position

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

(b) Planning over training steps

(a) One block: record in forward, replay in backward plan: save output 0 of mm #0

mm #0 (gate)

silu

𝑔

warm-up 𝑊 steps

𝑠

ℎ 𝑥

norm

mul

𝑚

mm #2 (down)

warm-up with mandatory values only

𝑢 mm #1 (up)

backward unpacks ℎ, 𝑢, 𝑠, 𝑚 ⇒ recompute norm, mm #1, silu, mul from 𝑥 and 𝑔; mm #0 and mm #2 are not recomputed saved for backward

recomputed

profile

dead, not recomputed

solve one step

𝑅 steps

𝑦 : not saved

time the recomputed nodes of one block per type; measure the must-only step peak

train later steps

rank 0 solves under reuse the plan of the memory limit; the token bucket; solve again for a broadcast new bucket

block input

Figure 4: Execution and planning in TapeSAC: (a) recording and replay of a gated MLP block under a plan that saves 𝑔; (b) profiling and solving over training steps. in the block: coord(𝑣) = (op, 𝑘, 𝑗),

(2)

output 𝑗 of the 𝑘-th call of operator op in the block, counting from 0. Every step runs the same block code, so a coordinate points to the same place in every step. Recording applies the plan: for each operator call, it counts how many times that operator has already been called in the block, which gives 𝑘, and saves output 𝑗 if (op, 𝑘, 𝑗) is in the plan. With the plan {(mm, 0, 0)} of Figure 4a, the first mm saves its output 𝑔, and the second and third save nothing. Because a plan lists outputs rather than calls, it can save one output of a multi-output kernel and drop another, for example the log-sum-exp of an attention kernel but not its output. 3.2.3 Replay. Replay starts when autograd first unpacks a value of the block. TapeSAC first counts, for every value that replay may produce, how many times it will be read, as the input of another replayed node or by autograd. Then, for each value that autograd unpacks, it returns the value if it is available, that is, a leaf, a saved value or a value replayed earlier. Otherwise, it walks back from the value through its producers, collecting every node it passes until it reaches available values, and runs the collected nodes in dependency order. We call the nodes collected for targets 𝑄 with available values 𝐵 the producer closure of 𝑄, written closure𝐵 (𝑄). In Figure 4a, backward first asks for 𝑚. The walk goes from mul to silu, where it stops at the saved 𝑔, and to mm #1 and norm, where it stops at the input 𝑥. These four nodes run, and the values they produce, ℎ, 𝑠 and 𝑢, are ready when backward asks for them next. Replay runs whole calls: if backward needs the output of an attention kernel whose log-sum-exp is saved, the whole kernel runs again. Each replayed value is freed when its count reaches zero. A value that backward needs only late is therefore computed late and freed right after use, instead of occupying memory from the start of the block’s backward (Limitation 2). Replay never writes into a saved value or a leaf; it writes into a copy. Over the whole backward, replay runs only the nodes 𝑁 rp (𝑆) = closure𝑆∪𝐿 (𝑃),

(3)

each at most once. All other nodes never run; they are the dead nodes of §2.1. In Figure 4a, mm #0 never runs because its only output 𝑔 is saved, and mm #2 because backward needs none of its outputs.

PyTorch’s block recomputation (§2.2) instead reruns forward from the first node. It skips the cached calls and, with early_stop, stops after the last value it still needs: 𝑁 blk (𝑆) = { 𝑛𝑖 : out(𝑛𝑖 ) ⊈ 𝑆, 𝑖 ≤ 𝑖 ∗ (𝑆) },

(4)

is the 𝑖-th node of the tape and 𝑖 ∗ (𝑆) is the position of the

where 𝑛𝑖 last node with an output in 𝑃 \ 𝑆. With 𝑡 (𝑛) the measured time of node 𝑛, the recomputation times of the two are ∑︁ 𝑇blk (𝑆) = 𝑡 (𝑛), (5) 𝑛∈𝑁 blk (𝑆 )

𝑇rp (𝑆) =

∑︁

𝑡 (𝑛).

(6)

𝑛∈𝑁 rp (𝑆 )

Block recomputation runs every node that replay runs, so for any save set that both can express ∑︁ 𝑁 rp (𝑆) ⊆ 𝑁 blk (𝑆), 𝑇blk (𝑆) − 𝑇rp (𝑆) = 𝑡 (𝑛). (7) 𝑛∈𝑁 blk (𝑆 )\𝑁 rp (𝑆 )

In words, replay removes exactly the time of the dead nodes that block recomputation still runs. Because autograd may ask for only part of 𝑃, Eq. 6 is an upper bound on replay time.

3.3

Registering Kernels as Custom Operators

The tape sees only operators that go through the dispatcher. Many custom kernels, such as FlashAttention, the chunked scan of Gated DeltaNet (GDN) linear attention [23], fused normalizations and rotary position embeddings (RoPE), and expert GEMMs (general matrix multiplications), are called from an autograd.Function rather than registered as custom operators, so the dispatcher does not capture them. TapeSAC registers each such kernel once as a custom operator (torch.library.custom_op), with the original kernel as its forward and the original backward as its backward; neither the kernel nor the model code changes. The tape then records the kernel as one node, replays it, and can save each of its outputs separately, for example the attention log-sum-exp without the attention output.

3.4

Peak-Constrained Planning

Given a memory limit 𝑓 , a fraction of device memory per rank, the planner chooses the values to save that minimize the predicted replay time while keeping the predicted peak memory of the step

Kling Infra Team

Algorithm 1 Planning and execution of a TapeSAC training step. Require: limit 𝑓 ; per region type 𝑟 : tape 𝐺𝑟 , node times 𝑡 (𝑛) and block count 𝑚𝑟 (reference steps); base(𝜏 ) (sample steps); plan cache Π Plan 1: 𝜏¯ ← maxranks 𝜏; 𝑞 ← bucket(𝜏¯ ) 2: if Π[𝑞 ] = ∅ then ⊲ solved on rank 0
3: Δ ← minranks 𝐶 − base(𝜏𝑞+ ) ⊲ Eqs. 8, 9 4: for each region type 𝑟 do
5: 𝐹𝑟 ← frontier of 𝑏 (𝑆 ), 𝑇rp (𝑆 0 ) − 𝑇rp (𝑆 0 ∪ 𝑆 ) by beam search Í Í 6: pick (𝑏𝑟 , 𝑔𝑟 ) ∈ 𝐹𝑟 maximizing 𝑟 𝑚𝑟 𝑔𝑟 s.t. 𝑟 𝑚𝑟 𝑏𝑟 ≤ Δ 7: Π[𝑞 ] ← { (𝑟, coord(𝑣) ) : 𝑣 ∈ 𝑆𝑟 }; broadcast Π[𝑞 ] Forward, each block of region type 𝑟 8: record 𝐺, 𝐿, 𝑃 ; save 𝑆 ← 𝑆 0 ∪ {𝑣 : (𝑟, coord(𝑣) ) ∈ Π[𝑞 ] } Backward, each unpack of 𝑣 ∈ 𝑃 9: if 𝑣 ∉ 𝑆 ∪ 𝐿 ∪ 𝐸 then ⊲ 𝐸: live replayed values 10: run closure𝑆∪𝐿∪𝐸 ( {𝑣 } ) in dependency order; add its outputs to 𝐸 11: return 𝑣; free each 𝑢 ∈ 𝐸 with no remaining use

within the limit. It profiles one block of each region type and derives candidate values from it (§3.4.1), computes how much memory the step leaves for saved values (§3.4.2), splits that memory among region types and blocks (§3.4.3), and shares one plan across ranks and steps (§3.4.4). 3.4.1 Region models from profiling. Every checkpointed block records its tape in every step (§3.2.1), and the model adapter labels the block with its region type (§3.1). Region types depend on the model; the models of §4 have five: the text decoder’s linearattention, full-attention and multi-token-prediction (MTP) blocks, and the vision and audio blocks. Types differ in the operators that dominate their time and in the sizes of the values those operators produce, so the same saved bytes remove different replay time in different types. The planner therefore models each type separately and lets the types compete for one capacity (§3.4.3). The text decoder’s types are primary: their token count 𝜏 drives the memory constraint. The planner profiles a few early steps (Figure 4b). The first 𝑊 steps run must-only. On the next step, each region type records the tape of one reference block, the first block with the type’s most common operator sequence, together with the size of every value, and times every node that the block’s backward replays with CUDA events. A second step repeats the timing, and the planner keeps each node’s shorter time, which filters out contention from communication. The first of these steps also runs must-only to sample memory: it records the peak allocated memory of the whole step, the device memory outside PyTorch’s caching allocator, and the allocator’s fragmentation. Sizes and times are scaled linearly with the region’s token count. Blocks whose operator sequence differs from the reference save only the mandatory set. A region type is profiled again when its token count exceeds twice that of its reference block, and memory is sampled again when the token count moves more than 15% away from the sampled one. We observe that saving values together can remove more replay time than saving them one at a time. In the full-attention block of Figure 3, saving the query copy or the gate alone removes 1.1 or 2.0 ms per block, because the matrix multiplication that produces both still runs for the other; saving both removes 6.0 ms, because

that matrix multiplication becomes dead as well. A candidate is therefore one or more output slots that are saved together. Its cost is the memory it keeps alive, counted once per storage; its benefit is the reduction in replay time (Eq. 6) that it brings. 3.4.2 A constraint on the step’s peak memory. Let base(𝜏) be the peak allocated memory of a must-only step at primary token count 𝜏, interpolated between memory samples. For region type 𝑟 , let 𝑏𝑟 be the bytes a plan saves per block beyond 𝑆 0 and 𝑚𝑟 the number of blocks of that type that share the reference’s operator sequence. The plan must satisfy ∑︁ base(𝜏) + 𝑚𝑟 𝑏𝑟 ≤ 𝐶, (8) 𝑟

with capacity 𝐶 = 𝑓 · 𝑀dev − non_torch − margin − frag,

(9)

where 𝑀dev is the total device memory, non_torch is the device memory outside PyTorch’s caching allocator, and frag is the allocator’s fragmentation; the latter two are measured. The safety reserve margin is about 2% of device memory in our experiments. If base(𝜏) alone exceeds 𝐶, the plan saves nothing beyond the mandatory set; 𝑓 = 0 forces this must-only plan. Eq. 8 constrains the whole step, not saved bytes alone (Limitation 6): base(𝜏) already contains model states, communication buffers and the temporaries of must-only replay. It charges every saved byte at once. Every saved value is still alive when forward ends, so this is nearly exact early in the backward pass, where the peak falls on our models (§4.4), and conservative later, as blocks free their values. When the reserved memory observed in training approaches the limit, the planner raises its fragmentation estimate and solves again, which can only shrink the plan. When the peak falls later in the backward pass, as it can when activations are offloaded to host memory, Eq. 8 overcharges: blocks whose backward pass has finished have freed their saved values, and saving a value in the block that sets the peak lowers that block’s replay temporaries. For this case, TapeSAC also provides a peak-aware mode, whose window model predicts the peak of every backward window from the memory measured between unpacks on a sample step; the planner then saves extra values in the block that sets the peak. On our models, the peak falls at the start of the backward pass, where Eq. 8 exceeds it by at most one block’s saved bytes (§4.4), so our experiments use Eq. 8. 3.4.3 Solver. For each region type, the planner builds an approximate Pareto frontier of sets 𝑆 of candidates saved on top of the mandatory set 𝑆 0 , in bytes per block 𝑏 (𝑆) and replay time removed per block 𝑇rp (𝑆 0 ) − 𝑇rp (𝑆 0 ∪ 𝑆), by a beam search. Starting from the empty set, each round extends every set on the frontier by one candidate, rescores the extension by walking the tape, and keeps the sets that no other set beats in both bytes and time removed; at most 24 sets, spread over bytes, go on to the next round, and a set holds at most 12 candidates. A multiple-choice knapsack then selects one frontier point (𝑏𝑟 , 𝑔𝑟 ) per region type, where 𝑔𝑟 = 𝑇rp (𝑆 0 ) −𝑇rp (𝑆 0 ∪ 𝑆𝑟 ) is the replay time its save set 𝑆𝑟 removes Í Í per block, maximizing 𝑟 𝑚𝑟 𝑔𝑟 subject to 𝑟 𝑚𝑟 𝑏𝑟 ≤ 𝐶 − base(𝜏). It is solved exactly by dynamic programming over a discretized

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

capacity, with weights rounded up so that the result stays feasible. The knapsack is needed because region types compete for one step-wide capacity. In the experiments of §4, blocks of one type may also select different frontier points, with the sums taken over blocks (Algorithm 1 shows the form with one point per type); this avoids an all-ornothing choice when a valuable candidate fits on only some blocks. Larger save sets go to blocks later in forward order: their backward runs first, so their values are freed earliest and hold memory for the shortest time. The result lists, for each block, the coordinates (Eq. 2) of the values its save set saves. 3.4.4 Shared plans and token buckets. On real data, the token count changes from step to step. Solving a new plan every step would cost solver time and, because every plan changes the sizes the allocator serves, fragment its memory. The planner therefore groups token counts into buckets and solves one plan per bucket, which later steps in the same bucket reuse. Buckets are geometric: neighboring buckets differ by a factor of 1.1 for primary regions and of 1.25 for the vision and audio regions, whose token counts vary more; for the latter, the bucket is that of the largest count seen so far, so that their swings do not change the plan. We write 𝑞 = bucket(𝜏) for the bucket of token count 𝜏 and 𝜏𝑞+ for its upper edge. More tokens need more memory, so the plan of a bucket is solved at its upper edge: each rank’s must-only peak is charged at 𝜏𝑞+ , and the plan fits every step in the bucket. All ranks run the same plan, because the step time is set by the slowest rank. At the start of each step, the ranks all-reduce their token counts (taking the maximum), their remaining capacity (taking the minimum), and whether any rank needs a new reference or memory sample. Every rank thus computes the same bucket and finds the same cached plan, or misses together; on a miss, rank 0 solves for the largest token count and the smallest capacity and broadcasts the plan.

4

Experiments

We evaluate TapeSAC on three models (Table 2). Qwen3.8-27B, a dense model, and Qwen3.6-35B-A3B, an MoE model, each trained on 16 GPUs, carry all main experiments; Qwen3.5-397B, trained on 64 GPUs, tests whether the results hold at scale. The evaluation answers the following questions. Does TapeSAC reproduce the results of plain execution exactly (§4.2)? How much does replay alone gain with the save set held fixed (§4.3)? Does planning under a memory limit outperform full checkpointing and hand-tuned policies, also when the configuration changes (§4.4)? Which design elements matter, and how does TapeSAC compare with the memory budget of torch.compile (§4.5, §4.6)? Do the results hold at scale and on real data, and what does TapeSAC cost (§4.7, §4.8)?

4.1

Experimental Setup

Table 2 lists the models. Each configuration we compare on them is defined by two choices: the save set, which tensors are saved for the backward pass, and the recomputation mode, whether PyTorch recomputes everything else block by block or TapeSAC replays it. Table 3 arranges the configurations along these two choices. Comparing the two columns of a row isolates replay (§4.3); comparing rows isolates the choice of tensors, that is, planning (§4.4).

Table 2: Model architectures and training configurations. All models use FSDP2 across the listed GPUs without additional parallelism. Model

Qwen3.8-27B

Qwen3.6-35B-A3B

Qwen3.5-397B

Feed-forward GDN blocks Full-attention blocks

dense 48 16

MoE 256, top-8 30 10

MoE 512, top-10 45 15

GPUs (nodes) Local batch size

16 (2) 18

16 (2) 16, 48

64 (8) 16

Table 3: Evaluated configurations by save set and recomputation mechanism. Save set

Block recomputation

Replay

Block inputs only H1–H3 Planned under limit 𝑓

full checkpointing hand-tuned SAC –

must-only ( 𝑓 = 0) replayed H1–H3 TapeSAC at 𝑓

A hand-tuned save set consists of the outputs of a few operators chosen from an offline profile, ranked by the recomputation time they remove per GB saved. H1, chosen for Qwen3.5-397B and adapted to the width of each model, saves the router’s softmax and top-𝑘, the GDN gate projections, the concatenation after RoPE and the output projections; H2 saves the attention outputs, including those of the vision tower; and H3 saves the attention outputs and the concatenation after RoPE, and on the 35B model also the router’s top-𝑘, an output projection and the GDN scan. Full checkpointing runs torch.utils.checkpoint with the stock kernels and hand-tuned SAC runs PyTorch’s SAC with the same registered kernels as TapeSAC; must-only saves only the mandatory set of §3.2.2, which is empty on our models. We report the median step time over the steps of a run (the mean on Qwen3.5-397B), unless noted otherwise, and the peak memory reserved by PyTorch’s caching allocator on rank 0, as a percentage of device memory. The memory limit 𝑓 is also a fraction of device memory, and TapeSAC runs with the bucket-edge charge and per-block points of §3.4.4 and §3.4.3.

4.2

Correctness

Because replay applies the same operators to the same inputs as the forward pass, it should reproduce every recomputed value bitwise, independently of the save set; likewise, registering a kernel as a custom operator should leave its results unchanged (§3.3). We verify both properties at three levels of granularity. At the operator level, every registered kernel agrees with its stock implementation in outputs and gradients and passes torch.library.opcheck. At the block level, a GDN block and a full-attention block of Qwen3.5397B, replayed without saving any planned value so that every node that backward needs is recomputed, reproduce the outputs and all parameter gradients of plain execution bitwise. At the model level, the first training step of each of the three models yields the same loss under full checkpointing, every hand-tuned policy in either execution mode, and every TapeSAC plan we evaluate, up to one unit in the last place of float32, a difference that also occurs between PyTorch’s own configurations. Over 500 steps of real data (§4.7.2),

Kling Infra Team

the training loss of TapeSAC follows that of full checkpointing step by step (Figure 5). 0.9

full checkpointing

0.8

Item

training loss

TapeSAC

0.8

0.7

0.7

0.6

0

200

400

300 310 320 330 340

step

(b) Steps 300–340

Figure 5: Training-loss trajectories for full checkpointing and TapeSAC on the same real-data sequence (§4.7.2). A suite of unit tests on a small network with views, transposes and in-place updates further verifies the mechanisms of §3.2: replay skips dead branches, executes fewer nodes than were recorded, and leaves saved values intact; random operators and mutations of leaves are rejected; and a module that computes indices with nonzero and repeat_interleave obtains bitwise identical gradients without host synchronization in the backward pass, since the outputs of these operators belong to the mandatory set and are saved rather than replayed.

step-time reduction (%)

Replay at a Fixed Save Set

3

full checkpointing

3.0

H1

H2

H3

2.7

2

1.8

1.7 1.3 0.9

1 0.0

0 27B

Text MLP down projection, 64 recomputations Vision MLP output projection, 27 recomputations Residual additions, 64 + 27 recomputations Recomputation skipped in total Extra copies around the attention operator Compute stream busy (measured) Step time (measured, profiled step)

GPU time (ms)

−1000.6 −97.6 −49.1 −1147.3 +79.6 −1036 −973

step

(a) Moving average over 20 steps

4.3

Table 4: GPU-time breakdown of must-only replay relative to full checkpointing for one profiled step of the 27B model, grouped by the operators launching the kernels.

1.0 0.2

0.1 35B-A3B

397B

Figure 6: Step-time reduction from tensor-level replay relative to block recomputation for each fixed save set. Replay addresses the second problem of §1: block recomputation reruns a block in program order, running operators whose outputs no backward node reads and keeping the recomputed values alive together. To separate replay from the choice of tensors, we hold the save set fixed and change only how it is recomputed (Figure 6). We use four save sets: that of full checkpointing, whose replay is must-only, and the three hand-tuned save sets H1–H3 of §4.1, which save more. For every save set and model, replay is at least as fast as block recomputation, and it gains most with the save set of full checkpointing, which leaves the most to recompute (Figure 6). With that save set, replay shortens the step by 3.0% on the 27B model, 1.7% on the 35B model and 1.8% on the 397B model. With the hand-tuned save sets, the gain is 0.9–1.3% on the 35B model and at most 0.2%

for H2 and H3 on the 27B model and for H1 on the 397B model. H1 on the 27B model is an outlier at 2.7%: it nearly fills device memory, and under block recomputation every step incurs an allocator retry. Profiling one step of full checkpointing and of must-only on rank 0 of the 27B model shows where the gain comes from (Table 4). The kernels launched by the forward pass take the same time in both modes (8.86 and 8.88 s), so recording adds no GPU time; the difference lies in the backward pass, and nearly all of it is one operator, the down projection of each text MLP, a 15.6 ms GEMM. Its backward needs its input and its weight but not its output. Autograd saves that input only when the GEMM executes, so block recomputation must run the GEMM and discard its result, whereas replay runs only the producers of the input. The output projection of each vision MLP and the residual addition at the end of each block are skipped for the same reason, while attention, the GDN core and the normalizations, whose saved tensors the backward pass reads, run the same kernels for the same time in both modes. Of the 1.15 s of recomputation skipped, 0.97 s appears as a shorter step; the rest is absorbed by copies around the attention operator, which runs FlashAttention-4 as a registered custom operator, and by longer waits in reduce-scatter on the ranks that compute faster. This is the redundancy of block recomputation described in §2.2, measured on a large model: 𝑇rp (𝑆) < 𝑇blk (𝑆) for the same 𝑆 (Eqs. 5 and 6). The 35B model shows the same mechanism: its forward kernels take the same time in both modes, and its backward kernels take 227 ms less under replay, mostly in operators whose outputs the backward pass does not read, such as the output projections of the vision and audio MLPs and the down GEMMs of the MoE layers. With H1, they take only 53 ms less: a hand-tuned save set already avoids part of the redundant recomputation and leaves less for replay to skip, which is why replay gains less with the hand-tuned save sets. Because replay materializes each recomputed value only when it is consumed and frees it after its last use, it can also lower memory. With the save set of full checkpointing, reserved memory is 1.5% lower on the 27B model and unchanged on the two MoE models, and with H1 it is 2.1% lower on the 397B model (Figures 7 and 10). The differences are small because the step’s peak is set largely by the sharded parameters, gradients, communication buffers and fragmentation (§4.4). With the save set of full checkpointing, replay thus does not increase memory: must-only, the point from which planning starts in §4.4, is already faster than full checkpointing at the same or lower memory.

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

full checkpointing

full checkpointing

10.5

must-only

step time (s)

step time (s)

31 H1

must-only

30 H1 rp H3 0.78

H1 H2

H3

10 0.45

9.5

H2 0.90

0.90

70

9

80

90

100

30

reserved memory, rank 0 (% of device)

40

50

60

70

reserved memory, rank 0 (% of device)

(a) Qwen3.8-27B

(b) Qwen3.6-35B-A3B

Figure 7: Reserved-memory–step-time trade-offs on the two main models. Curves show TapeSAC plans from must-only upward, labeled by memory limit 𝑓 ; squares denote full checkpointing, and diamonds denote hand-tuned save sets under PyTorch SAC (filled) or TapeSAC replay (open).

4.4

Planning under a Memory Limit

step-time reduction (%)

Planning trades memory for step time from must-only upward (Figure 7). On both main models, more memory gives a shorter step, with diminishing returns: on the 27B model, the first 10% of device memory beyond must-only shortens the step by 0.67 s, and the next 10% by 0.42 s. How much planning adds depends on the room above full checkpointing. At 𝑓 = 0.90, TapeSAC is 6.3% faster than full checkpointing on the 27B model, where full checkpointing already uses 65% of device memory, and 13.9% faster on the 35B model, where it uses only 31% and the plan raises reserved memory to 75%, saving mostly in the vision and full-attention blocks.

5

replay

4.8∗

planning

4 2.9

2 0.9

1

0.6

0 H1∗

H2 27B

H3

H1

Variant both options (default) neither option bucket-edge charge only per-block points only

27B 𝑓 = 0.78 0.90 29.71 30.27 29.70 29.82

29.38 29.42 29.42 29.46

35B 0.45 0.90 9.96 10.05 10.07 9.94

9.12 9.17 9.14 9.13

5.1

4.7

3

Table 5: Effect of solver options on median step time (s) under different memory limits.

H2 35B-A3B

H3

Figure 8: Step-time reduction of TapeSAC relative to handtuned PyTorch SAC policies at the same or lower reserved memory, decomposed into replay and save-set planning gains. ∗ TapeSAC uses 11% less reserved memory. The hand-tuned policies lie off this frontier: at the same reserved memory, or less for H1 on the 27B model, TapeSAC is faster than every one of them (Figure 8). On the 35B model, it is 2.9–5.1% faster than the three policies under PyTorch’s SAC. Replaying a policy recovers about 1% of this gap; the remaining 2.0–4.1% comes from the choice of tensors rather than from their execution. On the 27B model, H2 and H3, selected for this model, come within 0.6–0.9% of TapeSAC, whereas H1, tuned for a different model, needs 97.6% of device memory, and TapeSAC, using 86.9%, is 4.8% faster. A hand-tuned policy can thus approach the frontier, but only for the configuration it was tuned for.

Such a policy fixes which tensors to save, not how much memory to use, so a change of configuration moves its memory with it. On the 35B model, the three policies shorten the step by 2.2–4.9% at a local batch size of 16, but all of them run out of memory at 48, under PyTorch’s SAC and replayed alike; TapeSAC, with the same 𝑓 = 0.90, re-plans the save set and remains 5.6% faster than full checkpointing. The planner chooses among plans by their predicted speedup, which tracks the measured one. Relative to must-only, the measured speedup is 89–95% of the predicted speedup on the 27B model and 77–79% on the 35B model. Timing each node with events in isolation overestimates replay that overlaps other kernels, but it does so consistently: on both models, the measured order of the plans matches the predicted one at every limit, which is all the planner needs to choose among them. The planner’s memory constraint is nearly exact as well. In all three models, the peak falls at the start of the backward pass, in the window of the last GDN block, which runs after only the last full-attention block has freed its saved values; Eq. 8 therefore exceeds the peak by at most one block’s saved bytes. Most of the gap between reserved and allocated memory does not depend on what is saved: on the 27B model, 7.8% of device memory at the peak is idle memory in the communication pool and in the optimizer’s side streams, each of which the caching allocator serves from its own segment.

25.5

0.8

100

no ckpt.

0.50

TapeSAC

full checkpointing

75

step time (s)

recomputation eliminated (%)

Kling Infra Team

0.30

50 25 must-only

0

TapeSAC (eager) 0.1 budget 0

full ckpt.

0

20

25 must-only

24.5

H1, SAC H1, replayed

0.84

memory budget (compiled)

40

60

80

0.99 0.90

24

100

extra memory (%)

75

80

85

90

reserved memory, rank 0 (% of device)

Figure 9: Memory–recomputation trade-off on Qwen3-8B (8 GPUs) for eager TapeSAC and the torch.compile memory budget. Each method is normalized to its own fullcheckpointing (0) and no-checkpointing (100) endpoints on both axes; labels indicate 𝑓 or the compiler budget.

4.5

Ablation Study

We remove one design element at a time: the registration that makes custom kernels replayable (§3.3) and the two solver options (§3.4.3 and §3.4.4). 4.5.1 Kernel visibility. Without registration, a custom kernel’s outputs cannot be replayed and are force-saved. The cost depends on the family. On the 27B model at 𝑓 = 0.90, leaving the GDN core unregistered force-saves its outputs in every GDN block, so that the must-only peak alone exceeds the planner’s capacity; the planner saves nothing else, and the step is 8.8% slower than with the kernel registered. Leaving the root-mean-square (RMS) normalizations unregistered force-saves both normalization outputs of every layer, and the must-only sample step runs out of memory. On the 35B model, the same two families slow the step by 5.2% and 2.9%. Leaving FlashAttention-4 unregistered costs only 0.3% on the 27B model, since its outputs are among those the planner saves anyway, and leaving the 𝑞𝑘-normalization, RoPE and attention-gate kernels unregistered force-saves nothing and costs 0.1%. Registering these kernels is thus a prerequisite for tensor-level checkpointing on these models rather than an optimization: without it, the outputs of the kernels that dominate activation memory can only be saved. 4.5.2 Solver. The beam search is not what limits the plans: on the 27B model, a search four times wider, keeping 96 sets per round with up to 16 candidates each, predicts no shorter replay time at any of ten memory limits between 𝑓 = 0.76 and 0.94. The two options, the bucket-edge charge and per-block points, matter where memory is tight, each on a different model (Table 5). On the 27B model at 𝑓 = 0.78, the FlashAttention-4 outputs of the 16 full-attention blocks need 7.4% of device memory. Without the two options, the room for saved values is 7.0%, so no block saves them; charging the must-only peak at the bucket edge raises the room to 8.3%, every block saves its output, and the step shortens by 0.55 s, with the bucket-edge charge alone giving the same. On the 35B model at 𝑓 = 0.45, the room barely changes, but per-block points let some vision blocks give up a tensor that removes about half as much recomputation time per byte as the FlashAttention-4 outputs, so

Figure 10: Step time versus rank-0 reserved memory for Qwen3.5-397B on 64 GPUs. Labels on the TapeSAC curve indicate the memory limit 𝑓 .

that all 10 full-attention blocks can save those outputs, and the step shortens by 87 ms. At 𝑓 = 0.90, both options together still give the fastest step on both models, but by at most 80 ms: they matter when a valuable group does not fit on every block of its region type, and otherwise mainly let the planner use the memory it is given.

4.6

Comparison with the Memory Budget of torch.compile

On the models we train, torch.compile is not an alternative: their graph capture breaks before the partitioner runs (§2.3), and their hand-fused kernels leave the compiler’s backend little to fuse, which is its main source of speed. Where capture does succeed, the limitations of §2.3 remain, and we measure their effect on a model that the compiler can capture. We compile torchtitan’s Qwen3-8B (text only, 36 layers) one block at a time and train it on 8 GPUs with FSDP and two sequences of 8192 tokens per GPU. The eager side runs TapeSAC from must-only to 𝑓 = 0.55, the compiled side the budget from 0 to 1, and each side also runs full checkpointing and no checkpointing. Because compilation changes the step time of every configuration, full checkpointing included, Figure 9 compares each side against its own endpoints: extra memory is the memory used beyond full checkpointing, and recomputation eliminated is the reduction in step time, each as a share of what disabling checkpointing uses or removes. With little extra memory, TapeSAC eliminates several times more recomputation than the budget: 43% with 10% of the extra memory, against 8% with 9% at budget 0.1. With none, must-only already eliminates 14%, whereas budget 0 is slower than full checkpointing (−4%). The two converge only with more than 60% of the extra memory (91% for TapeSAC at 𝑓 = 0.50, 94% for budget 0.8), so the budget is least effective when memory is tight, which is when selective checkpointing is needed. Controlled experiments on small models, compiled with Inductor on one GPU, isolate two properties behind this. First, the budget bounds the bytes saved for backward, not the peak. On a chain of 8 identical blocks compiled as one graph, budget 0 saves no activations at all, yet the peak is unchanged and the step is 38%

TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training

Table 6: Real-data training results for Qwen3.6-35B-A3B over 500 steps. Full checkpointing and TapeSAC (𝑓 = 0.90) use the same data order on each of two node pairs; TapeSAC ran first on node pair 1 and full checkpointing ran first on node pair 2. Settings Total time (min) full checkpointing TapeSAC change Throughput change (global tokens/s) Step-time change steps without a solve or sample steps with a solve or sample (16) steps from 139 on Reserved memory (%) full checkpointing TapeSAC

Node pair 1

Node pair 2

92.8 85.4 −8.0% +8.7%

92.4 82.1 −11.1% +12.5%

−10.7% +8.3 s each −12.3%

−12.6% +3.7 s each −13.0%

37.2 75.5

37.2 75.7

slower, whereas segment checkpointing, which recomputes one segment at a time, lowers the peak by 26%; at 16 and 32 blocks, the budget’s peak is 1.59× and 1.76× that of segment checkpointing at equal recomputation. Second, on top of per-block checkpointing, which is how large models are trained, the budget has almost no effect: on an 8-layer transformer, adding a budget of 0.1 to 0.7 to full checkpointing leaves the step time unchanged and lowers the peak by only 5%.

4.7

Scale and Real Data

4.7.1 Qwen3.5-397B. The results hold at scale, and TapeSAC outperforms the hand-tuned policy H1 (Figure 10). At the default 𝑓 = 0.90, TapeSAC shortens the step by 4.9% relative to full checkpointing; H1 shortens it by only 2.2% while raising reserved memory by 25%, and TapeSAC is 2.7% faster than H1 with 13.6% less memory on rank 0 (79.3% against 91.8% of device memory). This memory advantage is not limited to rank 0: although vision-token counts range from 230k to 334k across ranks, the maximum reserved memory across ranks exceeds rank 0’s by only 1–3%, because the shared plan is constrained by the rank with the least available memory. 4.7.2 Real data. On real data, token counts change from step to step; each new token bucket triggers a solve, and token counts outside the measured range trigger a new memory sample. We ran the 35B model with a local batch size of 16 for 500 steps, with full checkpointing and with TapeSAC at 𝑓 = 0.90, in the same data order and on each of two pairs of nodes (Table 6). Over 500 steps, TapeSAC shortens training by 8.0% and 11.1% on the two pairs of nodes, raising throughput by 8.7% and 12.5%. The number of token buckets is bounded: after step 138, every bucket the data produces already has a plan, and from there on TapeSAC is 12.3% and 13.0% faster. The difference comes from the 16 steps that solve a plan or sample memory, all within the first 140 steps.

4.8

Overhead and Generality

Recording and replay cost host time, not GPU time: about 24 𝜇s per recorded operator, which GPU work hides from 2048 tokens per block upward, far below the token counts we train with. Running the custom kernels as registered custom operators costs 0.3–0.4%

of step time. The same TapeSAC package runs all three dense and MoE models on 16–64 GPUs, requiring only a small model-specific adapter (§3.1).

5

Conclusion

TapeSAC turns activation checkpointing from a policy that a person writes into a problem that the system solves. A hand-tuned policy fixes which tensors to keep, so its memory moves with every change of configuration; TapeSAC instead fixes the memory and solves for the tensors. Replaying a tape recorded from each block’s forward pass, it recomputes only what the backward pass consumes, and its planner prices candidates with measured kernel times and constrains the step’s peak. It runs in eager PyTorch without graph capture or edits to the model’s forward code, and reproduces the recomputed values bitwise. On Qwen3.8-27B and Qwen3.6-35B-A3B, a single limit of 90% of device memory shortens the step by 6.3% and 13.9% relative to full checkpointing; at the same or lower memory, TapeSAC is up to 5.1% faster than the hand-tuned policies we evaluate, and it keeps its gain where they run out of memory. On Qwen3.5-397B, it is 2.7% faster than the policy hand-tuned for that model with 14% less memory, and over 500 steps on real data it shortens training by 8–11%. The tape already tracks the size and lifetime of every value it records, so the planner could decide more than what to save. Recording every write the training step makes to device memory would let it also decide where each tensor is placed, in device memory or offloaded to the host, bringing allocator fragmentation into the plan; planning each rank separately, so that slow ranks recompute less, may also mitigate stragglers.

References [1] Jason Ansel, Edward Yang, Horace He, et al. 2024. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS). [2] Burak Bartan, Haoming Li, Harris Teague, Christopher Lott, and Bistra Dilkina. 2023. Moccasin: Efficient Tensor Rematerialization for Neural Networks. In International Conference on Machine Learning (ICML). [3] Olivier Beaumont, Lionel Eyraud-Dubois, Julien Herrmann, Alexis Joly, and Alena Shilova. 2024. Optimal Re-Materialization Strategies for Heterogeneous Chains: How to Train Deep Neural Networks with Limited Memory. ACM Trans. Math. Software 50, 2 (2024), 1–38. doi:10.1145/3648633 [4] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. 2016. Training Deep Nets with Sublinear Memory Cost. arXiv preprint arXiv:1604.06174 (2016). [5] Andreas Griewank and Andrea Walther. 2000. Algorithm 799: revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation. ACM Trans. Math. Software 26, 1 (2000), 19–45. [6] Julia Gusak, Xunyi Zhao, Théotime Le Hellard, Zhe Li, Lionel Eyraud-Dubois, and Olivier Beaumont. 2025. HiRemate: Hierarchical Approach for Efficient Re-materialization of Neural Networks. In Proceedings of the 42nd International Conference on Machine Learning (ICML) (Proceedings of Machine Learning Research, Vol. 267). 21418–21443. [7] Horace He and Shangdi Yu. 2023. Transcending Runtime-Memory Tradeoffs in Checkpointing by Being Fusion Aware. In Proceedings of Machine Learning and Systems (MLSys). [8] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. 2020. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization. In Proceedings of Machine Learning and Systems (MLSys). [9] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. 2021. Dynamic Tensor Rematerialization. In International Conference on Learning Representations (ICLR). [10] Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. 2023. Reducing Activation

Kling Infra Team

Recomputation in Large Transformer Models. In Proceedings of Machine Learning and Systems (MLSys). [11] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. In International Conference on Learning Representations (ICLR). [12] Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. 2023. Sequence Parallelism: Long Sequence Training from System Perspective. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (ACL). 2391–2404. doi:10.18653/v1/2023.acl-long.134 [13] Jianjin Liao, Mingzhen Li, Hailong Yang, Qingxiao Sun, Biao Sun, Jiwei Hao, Tianyu Feng, Fengwei Yu, Shengdong Chen, Ye Tao, Zicheng Zhang, Zhongzhi Luan, and Depei Qian. 2023. Exploiting Input Tensor Dynamics in Activation Checkpointing for Efficient Training on GPU. In IEEE International Parallel and Distributed Processing Symposium (IPDPS). 156–166. doi:10.1109/IPDPS54959. 2023.00025 [14] Hao Liu, Matei Zaharia, and Pieter Abbeel. 2024. RingAttention with Blockwise Transformers for Near-Infinite Context. In International Conference on Learning Representations (ICLR). [15] Meta Platforms, Inc. 2026. torch_remat: Explicit Activation Rematerialization Helpers for PyTorch. https://github.com/meta-pytorch/remat. Version 0.2.0. [16] Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, and Xuehai Qian. 2020. Capuchin: Tensor-based GPU Memory Management for Deep Learning. In Proceedings of the 25th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS). [17] PyTorch Contributors. 2026. torch/distributed/_tools/sac_ilp.py, PyTorch 2.13. https://github.com/pytorch/pytorch/blob/v2.13.0/torch/distributed/ _tools/sac_ilp.py. [18] PyTorch Contributors. 2026. torch/_functorch/partitioners.py, PyTorch 2.13. https://github.com/pytorch/pytorch/blob/v2.13.0/torch/_functorch/ partitioners.py. [19] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC).

[20] Manuela Schuler, Richard Membarth, and Philipp Slusallek. 2022. XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments. ACM Transactions on Architecture and Code Optimization 20, 1 (2022), 17:1–17:25. doi:10.1145/3568956 [21] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv:1909.08053 (2019). [22] Benoit Steiner, Mostafa Elhoushi, Jacob Kahn, and James Hegarty. 2023. MODeL: Memory Optimizations for Deep Learning. In International Conference on Machine Learning (ICML). [23] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. 2025. Gated Delta Networks: Improving Mamba2 with Delta Rule. In International Conference on Learning Representations (ICLR). [24] Tailing Yuan, Yuliang Liu, Xucheng Ye, Shenglong Zhang, Jianchao Tan, Bin Chen, Chengru Song, and Di Zhang. 2024. Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism. In USENIX Annual Technical Conference (USENIX ATC). 545–561. [25] Jianhao Zhang, Shihan Ma, Peihong Liu, and Jinhui Yuan. 2023. Coop: Memory is not a Commodity. In Advances in Neural Information Processing Systems (NeurIPS). [26] Xunyi Zhao, Théotime Le Hellard, Lionel Eyraud-Dubois, Julia Gusak, and Olivier Beaumont. 2023. Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch. In International Conference on Machine Learning (ICML). [27] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. 2023. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. Proceedings of the VLDB Endowment 16, 12 (2023), 3848–3860. [28] Zhanda Zhu, Christina Giannoula, Muralidhar Andoorveedu, Qidong Su, Karttikeya Mangalam, Bojian Zheng, and Gennady Pekhimenko. 2025. Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism CoOptimization. In Proceedings of the Twentieth European Conference on Computer Systems (EuroSys). doi:10.1145/3689031.3717461

Related documents

• Which Optimizer Leaks Most? Optimizer Implicit Bias as a Training-Data Privacy Controller#1343809
• Self-protective memory#251989
• A Solvable Theory of Pre-training Data Poisoning: Regime-Dependent Scaling Exponents#1116619
• The Price of Scale in Proactive Memory Recall#1457661
• Preisach Attention: A Hysteretic Model of Sequential Memory#222606