graph_trainer: budget-aware activation memory policy (auto_perf_maxing)

[nnurlan008]

Stack from ghstack (oldest at bottom):

• -> graph_trainer: budget-aware activation memory policy (auto_perf_maxing) #4636

What

Adds a memory policy that takes a peak memory budget and decides, per tensor,
whether to keep, recompute, or offload it — maximizing throughput subject to that
budget.

Currently, the choices are full (recompute everything: safe, slow) and
sac_and_offload (fast, but has no notion of a budget — on llama3-8b it settles
at 120 GiB), etc. This policy lets you state the peak you can actually afford and
keeps most of the throughput.

As the offloaded activations are prefetched earlier (preferably 1 layer ahead), the
whole chain depending on the offloaded activations become available earlier compared
to their original recompute order. So, the inductor pass can move the whole ready
chain earlier and create idle tensors, which unintentionally increases the peak memory.
This commit also changes the order of passes. Regional inductor happens before tagging
passes.

How

Two levels, then the existing passes do the work:

1. Outer LP, once over all layers: splits the budget across transformer
   blocks into per-layer keep/recompute/offload fractions.
2. Inner solver, per layer: turns those fractions into per-tensor decisions
   and tags the graph. Layers with identical candidate signatures are solved once
   and the pattern is reused.
3. apply_cpu_offload_pass and selective_activation_remat_pass act on the
   tags — both unchanged.

New files:

file	lines	contents
auto_sac_offload_solver.py	1980	the two-level solver and graph tagging
auto_sac_offload_solver_helpers.py	399	peak-memory estimator, optimizer-state sizing, runtime and bandwidth benchmarks
auto_sac_offload_solver_utils.py	127	shared constants and predicates

New flags

flag	default	what it does
--compile.memory_policy auto_perf_maxing	default	new value on the existing flag; turns the policy on
--compile.memory_budget_gb	1000.0	target peak GPU memory per rank, in GiB, covering the graph plus optimizer state. The default is effectively unbounded, so the policy is a no-op until you set it
--compile.solver_type	greedy	inner solver. greedy ranks tensors by recompute cost per byte and fills the offload/recompute budgets in that order; ilp solves the same split exactly with CBC
--compile.runtime_est_mode	benchmark	where per-node runtimes come from. Only benchmark is implemented today for now
--compile.cpu_offload_bw	10000	D2H/H2D bandwidth in GB/s used to size offload windows. The default is a sentinel meaning "measure it at startup and divide for local ranks"; pass a real number to override
--compile.host_memory_fraction	0.80	share of free host RAM that may be pinned for offloaded activations
--compile.debug_memory_policy_solver	False	log the solver's inputs and per-layer decisions; silent otherwise

Results

GB200, llama3-8b, batch size 32, seq len 2048.

budget	greedy tps	peak GiB	vs full	ILP tps	peak GiB	vs full
80	19,600	77.73	+9.2%	19,656	79.01	+9.5%
81	20,113	79.48	+12.1%	20,108	80.01	+12.0%
82	20,346	81.23	+13.4%	20,274	81.26	+13.0%
83	20,484	81.23	+14.1%	20,404	82.01	+13.7%
84	20,562	82.98	+14.6%	20,458	83.01	+14.0%
88	20,618	86.93	+14.9%	20,251	86.95	+12.8%

Baselines:

policy	tps	peak GiB	vs full
full	17,947	78.77	—
sac_and_offload	20,581	120.58	+14.7%

At an 84 GiB budget the greedy solver reaches 20,562 tps — within 0.1% of
sac_and_offload — while peaking at 82.98 GiB instead of 120.58 GiB. The two
solvers track each other closely; greedy is the default because it solves faster
where the ILP takes several seconds.

Every measured peak lands at or under its budget.

Known limitations

• runtime_est_mode honours benchmark only. cost_model and interpreter
  log a warning and fall back to benchmark currently.
• The budget models the graph plus optimizer state. It might not model gradient
  buffers retained by zero_grad(set_to_none=False), which cost an extra copy of
  the sharded params if cudagraphs are disabled via --compile.disable_passes
  rather than --training.disable_cuda_graphs.