Skip to content

Commit a7a5249

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

File tree

0 commit comments

Comments
 (0)