Wiring up offload_opt_states (#8186)
### 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).
---------
Signed-off-by: pengdurice <pengduhit@gmail.com>
