Commit a7a5249
committed
feat(relax): PR5 lever-5 in-place Adam optimizer -- MEMORY side closes at 8.787 GB full-step peak
Wire the in-place Adam m/v optimizer update into the PR-4 sqrt(N)-remat full
train-step (docs/RELAX-GRAPH-MEMORY-PATH.md sec9 deliverable-4).
path_c_relax_step_optim.py: build_full_step_with_optim appends a pathc.adam_inplace
call_dps_packed after the backward pass that reads+writes the param/m/v banks with
the optimizer call as the LAST use of the input m/v/param banks. StaticPlanBlockMemory
reuses their freed storage for the updated banks -> optimizer-state stays at 1x
(m+v+param persistent), not 2x params. A NAIVE (inplace=False) variant keeps the
inputs live to PROVE the in-place SSA structure is load-bearing (no planner aliasing
across the call_dps_packed boundary).
MEASURED (CPU LLVM Relax VM, real bank numels, 28 MR blocks = 1.8B):
- full-step peak (fwd+bwd+remat+optimizer) = 8.787 GB == PR-4 remat peak (optimizer
+0.000 GB to PEAK: it runs after backward drains, never coexists with checkpoint peak)
- persistent working-set in-place 9.844 GB vs naive 10.253 GB -> 0.409 GB (~1x params)
saved, the duplicated m'/v'/param' band avoided
- numeric: in-place Adam == numpy reference Adam, max abs diff 0.0 (param', m', v') at
4/8/28 layers + /2000 stress
The graph/Relax memory path CLOSES: eager 118 GB OOM -> 8.787 GB full-step planned
peak, well inside the Megatron-class 26-40 GB target.
docs sec10 added; sec9 deliverable-4 marked WIRED. RULE #1: every check RAISES on
mismatch.1 parent cb9cffc commit a7a5249
2 files changed
Lines changed: 672 additions & 3 deletions
0 commit comments