Commit 417feb3
committed
feat(relax): PR4 sqrt(N) rematerialization on O(N) checkpoint bank
Non-boundary backward regions RE-EMIT the forward call_dps_packed to recompute
their checkpoint locally (Chen 2016 sqrt-N rule), making it short-lived instead of
live across the whole backward pass. StaticPlanBlockMemory cannot see through the
call_dps_packed boundary, so the remat is EXPLICIT in the assembled bank-SSA graph.
MEASURED on real mr_path_c bank numels (CPU LLVM Relax VM, same true_planned_peak
analyzer as PR-3):
28-block (1.8B) planned peak: O(N) 27.38 GB -> O(sqrt N) 8.79 GB
(3.12x further than PR-3, 5.19x below eager all-live)
slope fit: PR-3 == 0.943 GB/layer (O(N), == state bank); PR-4 == 1.510*sqrt(N)+0.778
recompute overhead: 89 extra fwd calls on 28 (3.18x; reducible to ~1x via
per-segment caching, does not change the peak)
numeric equivalence: max abs diff 0.0 vs both numpy ref and PR-3 non-remat VM
(recompute is mathematically identical), at 4/8/28 layers + /2000 stress
The activation/grad term (the eager-118-GB OOM driver) is now single-digit GB at the
real 28 MR blocks; only the Adam optimizer state (lever 5, in-place op, specified)
remains to close to Megatron-class 26-40 GB. This is the step that determines whether
the graph path reaches Megatron-class memory: it does. Updates doc section 9.1 parent 3d4dfd0 commit 417feb3
2 files changed
Lines changed: 613 additions & 1 deletion
0 commit comments