Wiring up offload_opt_states

[pengdurice]

Make DeepCompile's optimizer-state offloading work under inductor and reachable from config

Summary of Changes

• Registered ops instead of closures — four ops in the existing dc namespace. The graph
  carries a tensor anchor and an integer index; live tensors stay in module state.
• ORDERED effect registration so stock inductor keeps the ops and preserves their order.
  Reload-before-sync is correctness-critical. A test fails without it.
• Config wiring, capacity-first schedule [(0,[z3]), (1,[for_init, z3, move_opt_states])]:
  states are emptied to host before profiling, so the plan is made against the floor and a job
  that only fits with offloading never runs a step with everything resident. (The pass author's
  own ordering from their test harness; the budget formula is unchanged.)
• Three stream races fixed, all silent if wrong: copies wait for the optimizer's writes;
  record_stream protects reload buffers from early reuse; and the copy stream waits for the
  compute stream before writing a reload buffer — without it a mid-backward reload overwrites a
  live activation, which showed up as NaN losses. All are stream dependencies, no host waits.
• Reload buffers allocate from the pool with room (compute stream, just freed by backward).
  The wrong pool costs an allocator retry plus a device-wide sync per step.
• Multi-graph gating so graph breaks do not duplicate copies; empty_cache once per compile
  phase rather than per step (per-step measured +28%); mutual exclusion with offload_parameters.

Results

Qwen3-14B, 8×H200 (141 GB), ZeRO-3, micro-batch 4, fp32 states (22.2 GB/rank),
expandable_segments:True. Medians over each run's final phase, single campaign.

configuration	seq 2048	seq 2400
plain eager ZeRO-3, no compile	OOM	OOM
ZeRO-Offload (CPU Adam)	10.11	OOM
DeepCompile, no offload (z3 only)	1.89	dies at the recompile
DeepCompile, no offload (default schedule)	1.82	dies at the recompile
this pass, blocking variant (eager-only)	4.06	4.42
this pass (async), eager	2.43	4.32
this pass (async), inductor	2.01	3.70
this pass at micro-batch 1	0.76 — planner offloads nothing	0.85

Limitations

• Runs in place of prefetch and selective gather (worth 3–27%); combining them is future work.
• Designed for gradient_accumulation_steps=1: the graph runs per micro-batch, so accumulation
  repeats the whole cycle. Documented in the config docstring.
• Mutually exclusive with offload_parameters.

Tests

tests/unit/v1/compile/test_offload_opt_states.py — op and ORDERED-effect registration; a
mechanism test compiling side-effect ops through stock inductor and asserting program order;
budget planning and node placement; re-run and multi-graph gating; once-per-phase empty_cache;
and a 2-GPU end-to-end loss-parity test whose op counters prove the ops ran in the compiled graph
(reloads are skipped while profiling, so a nonzero reload count is the proof).