Skip to content

Commit 417feb3

Browse files
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

File tree

0 commit comments

Comments
 (0)