Wiring up offload_opt_states - #8186
Conversation
Replace the Python closures the move_opt_states pass inserted as FX graph nodes with registered custom ops (torch.library, dc namespace), making the pass compatible with the inductor backend and its compile cache. Wire the pass to the user config (compile.offload_opt_states) using the capacity-first schedule: offload everything, profile on the emptied GPU, then keep resident only what the memory budget allows. Frees are completion-driven (record_stream), empty_cache runs once per compile phase, and unit plus 2-GPU end-to-end tests cover op registration, budget planning, schedule placement, and loss parity. Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3ca19178b2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| def test_offload_ops_registered_with_ordered_effects(): | ||
| _ensure_dc_ops() | ||
| from torch._higher_order_ops.effects import SIDE_EFFECTS |
There was a problem hiding this comment.
When these tests run under PyTorch versions allowed by the module-level min_version=2.1 marker but before _register_effectful_op exists, this unconditional import/assertion fails even though the production code explicitly treats that registry as optional. This affects CPU/unit test runs on older supported torch versions; guard these tests with the same availability check or raise the pytest minimum for this file.
Useful? React with 👍 / 👎.
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
tohtana
left a comment
There was a problem hiding this comment.
Hi @pengdurice,
Thank you for submitting this PR! The adaptive offloading feature was described in the paper, but the code had remained unorganized for a long time. I’m very glad that you cleaned it up and enabled the feature.
I left a few comments about some minor issues. Can you please address them? I don’t see any issues with the core implementation of this PR.
| schedule = [] | ||
| if (compile_config.offload_parameters): | ||
| schedule.append((0, [zero3_compile.add_z3_gather_release, offload_parameters.offload_parameter_fwd])) | ||
| elif compile_config.offload_opt_states: |
There was a problem hiding this comment.
This pass can be enabled with ZeRO optimizer's offload, but they won't work together. Can we reject the combination?
There was a problem hiding this comment.
sure, just added before this line:
optimizer = engine.optimizer
use_opt = not isinstance(optimizer, DeepSpeedZeRoOffload)
Thank you!
| from unit.util import bf16_required_version_check, skip_on_arch | ||
| from unit.v1.compile.util import compare_loss | ||
|
|
||
| pytestmark = pytest.mark.skipif(not required_torch_version(min_version=2.1), |
There was a problem hiding this comment.
DeepCompile already limits the version to 2.6+. I think we should make this consistent.
There was a problem hiding this comment.
sure, updated to 2.6. thank you!
Signed-off-by: pengdurice <pengduhit@gmail.com>
Signed-off-by: pengdurice <pengduhit@gmail.com>
|
Hi @pengdurice, I found the condition
|
thank you for the comment, just fixed;-) |
tohtana
left a comment
There was a problem hiding this comment.
Thank you for the update! It looks good to me now.
Make DeepCompile's optimizer-state offloading work under inductor and reachable from config
Summary of Changes
dcnamespace. The graphcarries a tensor anchor and an integer index; live tensors stay in module state.
Reload-before-sync is correctness-critical. A test fails without it.
[(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.)
record_streamprotects reload buffers from early reuse; and the copy stream waits for thecompute 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.
The wrong pool costs an allocator retry plus a device-wide sync per step.
empty_cacheonce per compilephase 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.Limitations
gradient_accumulation_steps=1: the graph runs per micro-batch, so accumulationrepeats the whole cycle. Documented in the config docstring.
offload_parameters.Tests
tests/unit/v1/compile/test_offload_opt_states.py— op and ORDERED-effect registration; amechanism 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).