[graph_trainer] FX memory/runtime/transfer estimators + naive outer ILP for joint AC + offload

[nnurlan008]

[graph_trainer] FX memory / runtime / transfer estimators + naive outer ILP for joint AC + offloading

Summary

Milestone 1 of the joint activation-checkpointing + offloading project. This PR lands
(1) the three FX-graph estimators the solver's cost model needs, and (2) a first
outer ILP that allocates a per-block keep / recompute / offload budget under a
peak-memory constraint, extending torch's two-level SAC MILP (sac_ilp.py) with
activation offloading (D2H/H2D copy-engine modeling).

All estimators operate on the joint fwd+loss+bwd FX graph produced by minimal_fx_tracer.
The outer ILP currently computes and logs the per-block plan; tagging /
materialization is the next PR (see Limitations).

Files

file	status	what
memory_estimator.py	updated	static peak-memory estimator (storage-keyed liveness, per-category breakdown, per-node live_bytes, per-tensor gaps)
runtime_estimator.py	new	per-op runtime in 3 modes: cost-model (roofline), benchmark (per-op CUDA timing), interpreter (real graph execution)
transfertime_estimator.py	updated	D2H/H2D transfer-time model + one-shot measured link bandwidth
naive_autoAC_outer.py	new	outer ILP: per-block keep/recompute/offload fractions under peak budget + shared copy-engine windows (renamed from the WIP new_autoAC_full_by_hand.py)

The estimators

• MemoryEstimator — walks the graph in execution order, tracks per-untyped_storage
  liveness (birth/death), categorizes each storage (param / grad / activation / temporary /
  optimizer-state), reports peak + the schedule point + per-category breakdown, and exposes
  live_bytes[t] and per-tensor [last_fwd_use, first_bwd_use) gaps for the solver.
• RuntimeEstimator — per-node time; roofline max(compute, HBM) by default, with
  benchmark / interpreter alternatives.
• TransferEstimator — per-tensor D2H/H2D time from size + measured PCIe/NVLink bandwidth
  (get_transfer_bw, benchmarked once per device, pinned host memory).

The outer ILP (naive_autoAC_outer.py)

Per transformer block b, continuous fractions r_b (recompute), o_b (offload); keep is
the residual k_b = 1 - r_b - o_b. Minimize added runtime under the peak budget:

• (C1') per-position peak — memory is evaluated at every block's backward step, not a
  single point:
  m[b] = fixed + act_b + sum_{k<b} (1-r_k-o_k)*act_k <= eff_budget.
  fixed = peak - sum(act) is the global non-freeable baseline; act_b is unconditional
  (block b is re-materialized at its own backward), so the last layer's discard yields no
  peak benefit and is correctly kept. Reduces to the single boundary point at full-keep.
• (C2') shared copy-engine windows — D2H (evict) and H2D (reload) are each one serial
  engine. For every block j, the offload of blocks k>=j must fit the time still available
  after j: sum_{k>=j} o_k*act_k <= bw*window[j]. D2H is hard (a slow evict leaves the
  tensor resident at the peak -> memory, not runtime); H2D is soft with a stall variable
  (a late reload only delays backward -> runtime). This one family is the exact single-machine
  feasibility condition and subsumes both per-block windows and global bandwidth.
• Objective — sum r_b*fwd_rt[b] + h2d_stall + keep_eps*(freed). Offload that fits its
  window is free; keep is preferred when the budget allows.

This mirrors sac_ilp.py's per-position m_i + convex structure, adding the offload option
and its non-monotonic, timing-sensitive profile.

Results

Estimator validation (llama3 1B, via measure_traced)

• Memory: static peak 41.277 GB == measured liveness peak 41.277 GB (ratio 1.0000);
  real allocator peak 43.78 GB (+6%, workspace/fragmentation); activation-at-peak
  static 21.86 vs measured 20.75 GB.
• Runtime (roofline vs measured serialized): fwd 101.9 vs 113.4 ms (-10%), bwd 202.2 vs
  233.9 ms (-14%), total 304 vs 344 ms (-12%) -- consistent underestimate, order preserved.
• Transfer (H100 PCIe): measured D2H 53.9 GB/s, H2D 57.5 GB/s. Offloading 11 GiB of
  activation lowered the measured peak by 10.8 GiB and added +58 ms (~0% overlap at that op
  count) -- see Limitations.

Outer ILP: model x budget x runtime mode (GiB)

Single H100, AUTOAC_MODE=autoac_byhand, b4/s2048. Dense models only (deepseek MoE has
dynamic shapes; out of scope). Runtime cost model: cost-model (roofline) and benchmark
(per-op CUDA timing). (The interpreter mode is omitted here -- its per-node times include
Python/dispatch overhead that inflates the overlap windows; use cost-model or benchmark for
the solver.)

Memory accounting. Optimizer state (opt) is resident the whole step and is not part
of the keep/recompute/offload decision, so it's subtracted up front: eff_budget = budget - opt.
The ILP constrains modeled peak = fixed + kept-activation <= eff_budget. The real GPU peak
= modeled peak + opt = the budget (last column). opt = 9.21 GiB (llama3_1b),
12.82 GiB (qwen3_1.7b); fixed (params+grads+temp+non-block acts) = 7.67 / 10.06 GiB.

model / budget	mode	keep	recompute	offload	modeled peak (fixed+kept)	opt	real peak (+opt)
llama3_1b / 20	cost-model	3.12	8.15	1.83	10.79	9.21	20.00
llama3_1b / 20	benchmark	3.12	8.22	1.77	10.79	9.21	20.00
llama3_1b / 24	cost-model	7.12	4.15	1.83	14.79	9.21	24.00
llama3_1b / 24	benchmark	7.12	4.22	1.77	14.79	9.21	24.00
qwen3_1.7b / 30	cost-model	7.12	10.41	3.29	17.18	12.82	30.00
qwen3_1.7b / 30	benchmark	7.12	10.14	3.56	17.18	12.82	30.00

Findings:

1. cost-model ~= benchmark (totals within ~0.3 GiB) -- both give kernel-level times, so
   the copy-engine windows and the recompute/offload split agree.
2. keep is runtime-mode-independent (3.12 / 7.12 / 7.12 per budget) -- it's set by the
   peak/budget (memory) constraint; only the freed split (recompute vs offload) depends on
   runtime.
3. offload is bandwidth-capped; on this H100 PCIe link only ~1.8 GiB (llama3) / ~3.3 GiB
   (qwen3) of activation can be evicted within the forward window, so the remainder recomputes.

Modeled peak lands exactly on eff_budget in every case (per-position constraint is honest).

How to run

These modules are standalone/importable; the results below were produced with a
local trainer integration (AUTOAC_MODE=autoac_byhand) that lands in a follow-up
PR. This PR is the estimators + solver module only.

AUTOAC_MODE=autoac_byhand AUTOAC_BUDGET_GIB=20 \
NGPU=1 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_1b ./run_train.sh \
  --compile.mode aot_fx_trace --compile.memory_policy default \
  --training.local_batch_size 4 --training.seq_len 2048 \
  --compile.disable_passes cudagraph_pass,tag_with_memory_policy_pass,apply_cpu_offload_pass,selective_activation_remat_pass \
  --hf_assets_path ./tests/assets/tokenizer --dataloader.dataset c4_test --training.steps 3

Knobs: AUTOAC_BUDGET_GIB (target peak), AUTOAC_RT_MODE = operator-level-{cost-model,benchmark,interpreter}.

Known limitations / follow-ups

• Outer ILP only. It logs the per-block plan and returns None; wiring the fractions to
  per-tensor MUST_SAVE/RECOMPUTE/CPU_OFFLOAD tags + a real-peak validate-and-tighten loop is
  the next PR. Nothing is applied to the graph yet.
• Backward recompute working set is approximated by act_b (a lower bound); should use the
  full re-materialized footprint (sac_ilp's ACM/IA split).
• Identical-block degeneracy: for homogeneous layers the per-block distribution is an
  arbitrary tie-break (only totals are pinned); a symmetry-break / convex per-block curve would
  make it deterministic and handle heterogeneous (MoE) models.
• Offload overlap gap: the model assumes offload hides when it fits its window, but the
  current apply_cpu_offload_pass achieves ~0% overlap at high op counts (prefetch depth 1);
  the pass needs improvement before offload-heavy plans realize their modeled benefit.
• Minor cleanup pending in naive_autoAC_outer.py: unused params (mem_est,
  time_fwd/time_bwd, t_fwd_total/t_bwd_total), dead d2h_stall, cpu_budget_bytes naming
  (GiB not bytes).

Per-layer allocation (keep / recompute / offload fractions)

Full per-block plan for every run (cost-model and benchmark runtime modes).
k/r/o per layer. Note: for identical layers the per-block distribution is a
tie-break (only column totals are pinned); read the totals and the structural
patterns (keep = budget headroom; last layer kept). The peak in each totals row
is the modeled peak (fixed + kept); the real GPU peak = modeled peak + opt =
the budget (opt = 9.21 GiB llama3_1b, 12.82 GiB qwen3_1.7b).

llama3_1b / 20 GiB

layer	cost-model	benchmark
0	0.26/0.44/0.30	0.00/0.86/0.14
1	0.00/1.00/0.00	0.00/0.00/1.00
2	0.00/0.00/1.00	0.00/1.00/0.00
3	0.00/1.00/0.00	0.00/1.00/0.00
4	0.00/1.00/0.00	0.56/0.00/0.44
5	0.00/1.00/0.00	1.00/0.00/0.00
6	0.00/1.00/0.00	0.00/1.00/0.00
7	0.00/1.00/0.00	0.00/1.00/0.00
8	0.00/0.51/0.49	0.00/1.00/0.00
9	1.00/0.00/0.00	0.00/1.00/0.00
10	1.00/0.00/0.00	0.00/1.00/0.00
11	0.00/1.00/0.00	0.43/0.00/0.57
12	0.55/0.00/0.45	0.00/1.00/0.00
13	0.00/1.00/0.00	0.00/1.00/0.00
14	0.00/1.00/0.00	0.82/0.18/0.00
15	1.00/0.00/0.00	1.00/0.00/0.00
total (k/r/o, peak)	3.12/8.15/1.83, 10.79	3.12/8.22/1.77, 10.79

llama3_1b / 24 GiB

layer	cost-model	benchmark
0	0.85/0.00/0.15	0.00/0.00/1.00
1	0.85/0.00/0.15	0.00/0.00/1.00
2	0.00/0.00/1.00	0.85/0.15/0.00
3	0.00/1.00/0.00	1.00/0.00/0.00
4	0.93/0.07/0.00	0.00/1.00/0.00
5	0.00/1.00/0.00	0.00/1.00/0.00
6	1.00/0.00/0.00	0.00/1.00/0.00
7	0.51/0.00/0.49	1.00/0.00/0.00
8	1.00/0.00/0.00	1.00/0.00/0.00
9	1.00/0.00/0.00	0.00/1.00/0.00
10	1.00/0.00/0.00	0.84/0.00/0.16
11	0.00/1.00/0.00	1.00/0.00/0.00
12	0.55/0.00/0.45	1.00/0.00/0.00
13	0.00/1.00/0.00	1.00/0.00/0.00
14	0.00/1.00/0.00	0.00/1.00/0.00
15	1.00/0.00/0.00	1.00/0.00/0.00
total (k/r/o, peak)	7.12/4.15/1.83, 14.79	7.12/4.22/1.77, 14.79

qwen3_1.7b / 30 GiB

layer	cost-model	benchmark
0	0.00/0.84/0.16	0.00/0.00/1.00
1	0.00/0.84/0.16	0.00/0.00/1.00
2	0.00/0.84/0.16	0.00/1.00/0.00
3	0.00/0.51/0.49	0.00/1.00/0.00
4	1.00/0.00/0.00	0.00/1.00/0.00
5	0.00/1.00/0.00	0.00/1.00/0.00
6	0.00/0.84/0.16	1.00/0.00/0.00
7	0.00/0.00/1.00	0.00/1.00/0.00
8	0.00/0.00/1.00	0.00/1.00/0.00
9	0.00/1.00/0.00	0.00/1.00/0.00
10	1.00/0.00/0.00	0.00/1.00/0.00
11	1.00/0.00/0.00	0.00/0.00/1.00
12	0.37/0.00/0.63	0.00/0.00/1.00
13	0.00/1.00/0.00	0.21/0.00/0.79
14	0.00/1.00/0.00	0.00/1.00/0.00
15	0.00/1.00/0.00	1.00/0.00/0.00
16	0.00/1.00/0.00	0.36/0.64/0.00
17	0.86/0.14/0.00	1.00/0.00/0.00
18	1.00/0.00/0.00	0.00/1.00/0.00
19	1.00/0.00/0.00	1.00/0.00/0.00
20	1.00/0.00/0.00	1.00/0.00/0.00
21	1.00/0.00/0.00	1.00/0.00/0.00
22	0.00/1.00/0.00	0.00/1.00/0.00
23	0.34/0.00/0.66	1.00/0.00/0.00
24	0.00/1.00/0.00	0.00/1.00/0.00
25	0.00/1.00/0.00	1.00/0.00/0.00
26	0.00/1.00/0.00	0.00/1.00/0.00
27	1.00/0.00/0.00	1.00/0.00/0.00
total (k/r/o, peak)	7.12/10.41/3.29, 17.18	7.12/10.14/3.56, 17.18