When More Memory Fails: Hidden Feasibility Holes in Dynamic Tensor Rematerialization

The gist: Dynamic Tensor Rematerialization (DTR) exhibits non-monotone feasibility holes in its reference simulator, meaning memory budgets that succeed for both smaller and larger budgets can fail for DTR <ref:2610.14369#pg6>.

How it works

Dynamic Tensor Rematerialization (DTR) trains neural networks under a memory budget by evicting tensors and recomputing them on demand <ref:2610.14369#pg2>. This process involves an online choice where when an allocation does not fit the budget, it evicts the resident tensor with the lowest score defined by hDTR(t) = c(e∗(t))m(t)·s(t) <ref:2610.14369#pg2>. This score incorporates compute, size, and staleness <ref:2610.14369#pg6>.

The study investigates whether the expectation that a budget large enough to complete training stays sufficient when increased fails for DTR <ref:2610.14369#pg2>. The researchers call these exceptions feasibility holes, which are defined as memory budgets at which an eviction policy runs out of memory although both a smaller and a larger budget complete <ref:2610.14369#pg6>.

Reproducible counterexamples to monotone feasibility for DTR in simrd

The authors provide reproducible counterexamples demonstrating these holes within the simulator, simrd <ref:2610.14369#pg3>. Specifically, they show holes on two of six confirmatory traces and on the development trace, each confirmed by repeats and by unmodified simrd <ref:2610.14369#pg5>.

The coarse grid evaluation, which sweeps the budget on a coarse grid, fails to reveal these issues because the 0.01-step grid on the confirmatory traces shows no DTR hole, although its unresolved timeouts prevent it from establishing monotonicity <ref:2610.14369#pg5>. Finer sampling, using Stages C and B, successfully reveals sample-supported failure regions on three confirmatory traces <ref:2610.14369#pg5>.

A step-by-step account of one failure, down to individual allocations, with an intervention that prevents it under the same budget

The paper details a specific failure on InceptionV4 where a larger budget avoids one eviction at an operator where the smaller budget was 173,479 bytes short <ref:2610.14369#pg6>. This extra data forces a later eviction of a tensor needed three operators afterwards, and rebuilding that tensor triggers a recomputation cascade that cannot fit <ref:2610.14369#pg6>.

The researchers trace the failure to individual allocations, showing how the smaller budget fails at operator 2,270 because the required tensor was evicted earlier <ref:2610.14369#pg6>. An intervention—preventing storage 3131 from being offered as an eviction candidate—prevents Out of Memory (OOM) in this execution under the same budget, supporting its causal involvement <ref:2610.14369#pg6>.

Changing the score does not remove the holes

The authors test two simple score changes to see if they eliminate these feasibility holes across traces <ref:2610.14369#pg7>. These changes do not remove holes across traces; both add failing budgets on some traces <ref:2610.14369#pg6>.

When testing a variant called TwoPhase, it is found that it has more failing samples than DTR on InceptionV4 (53 vs 16), U-Net (11 vs 8) and Unrolled GAN (7 vs 2) <ref:2610.14369#pg7>. This variant fails to meet the criteria for generalization, as it meets H1 and H2 on no trace <ref:2610.14369#pg7>.

Limitations

The study is limited by several factors, including the fact that it is Simulation only because simrd models no allocator, fragmentation, alignment, streams or framework memory <ref:2610.14369#pg8>. Furthermore, the results are based on Traces which are single recorded configurations from the DTR prototype <ref:2610.14369#pg8>. The paper also notes that Sampling means holes are observed at samples, and Unresolved runs include 417 valid runs that timed out <ref:2610.14369#pg8>.

The study concludes by asking whether these feasibility holes transfer to real hardware and whether a policy can guarantee monotone feasibility at acceptable cost. The author acknowledges the extensive use of AI assistants in writing, running experiments, and analyzing data.

The gist is that DTR exhibits non-monotone feasibility holes in its reference simulator, meaning memory budgets that succeed for both smaller and larger budgets can fail for DTR <ref:2610.14369#pg6>.

When testing a variant called TwoPhase, it is found that it has more failing samples than DTR on InceptionV4 (53 vs 16), U-Net (11 vs 8) and Unrolled GAN (7 vs 2) <ref:2610.14369#pg7>.

1. Identify exceptions feasibility holes in dynamic tensor rematerialization by monitoring budget behavior across different memory budgets, as memory budgets at which an eviction policy runs out of memory although both a smaller and a larger budget complete indicates a non-monotone feasibility region. This allows the system to proactively adjust its training strategy when operating near such boundaries to avoid catastrophic failure.

2. Implement an intervention mechanism that prevents specific tensor evictions, such as never offering storage 3131 as an eviction candidate, when a budget is near a known failure point, thereby ensuring that critical tensors needed three operators after its eviction remain resident and avoiding the recomputation cascade that cannot fit.

3. Develop adaptive scoring functions or scheduling policies that incorporate foresight or inclusion properties, addressing the finding that DTR’s cost-based score has no inclusion property, to ensure a policy can guarantee monotone feasibility at acceptable cost. This would prevent the system from entering regimes where overhead switches between regimes unpredictably.

4. Integrate finer sampling protocols (like Stage C) into the training sweep to specifically target and characterize failure regions, as Finer sampling finds holes on InceptionV4 (three regions) and U-Net (two regions), in addition to one on the ResNet-32 development trace. This allows for precise identification of where feasibility is not a simple threshold.

5. Enhance runtime instrumentation to track detailed memory metrics like pinned bytes (resident bytes outside the evictable pool) and maximum nesting depth, correlating them with specific tensor operations, as this signature helps distinguish between different failure modes, such as those exhibiting a deep cascade versus those that are not.

Dynamic Tensor Rematerialization (DTR) trains neural networks under a memory budget by evicting tensors and recomputing them on demand. A natural expectation is that a budget large enough to complete training stays sufficient when it is increased. Using DTR's reference simulator, simrd, to replay the public DTR training traces, we show that this expectation fails. We call the exceptions feasibility holes: memory budgets at which an eviction policy runs out of memory although both a smaller and a larger budget complete, a form of non-monotone feasibility. Under a written sweep protocol with dated amendments, a 0.01-step budget grid shows no such hole for DTR on any of six confirmatory traces, although unresolved timeouts prevent it from establishing monotonicity. Finer sampling finds holes on InceptionV4 (three regions) and U-Net (two regions), in addition to one on the ResNet-32 development trace. Every hole is reproduced by three repeats and by unmodified simrd. For one InceptionV4 pair (0.2343 versus 0.235 of peak memory) we trace the failure to individual allocations. The larger budget avoids one eviction at an operator where the smaller budget was 173,479 bytes short. The extra data it carries forward forces a later eviction of a tensor needed three operators afterwards, and rebuilding that tensor triggers a recomputation cascade that cannot fit. Keeping that one tensor resident under the same budget prevents the failure. The two simple score changes we tested do not remove holes across traces; both add failing budgets on some traces. All results are from simulation, with no allocator model; behaviour on GPU hardware is untested.