Repository navigation
perf(pt): pass is_inference to compile_fx_inner for compiled eval graphs - #6056
long-yi-2019 wants to merge 2 commits into
Conversation
The custom inference compiler entry calls compile_fx_inner without the is_inference flag, so Inductor plans memory under training assumptions: outputs are allocated early and autograd-visible intermediates stay alive. torch.compile's own default path passes is_inference=True through its inference compiler; this change matches that for the eval graph (training keeps False). Scheduling-only change: outputs agree to rounding level. Measured at N=4096 (SeZM eval compile, DP_COMPILE_INFER=1): peak VRAM -12.5% (23647 -> 20696 MB, RTX 5090); independently -13.9% on an H800 baseline.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely identifies the main change: passing is_inferencetocompile_fx_innerfor compiled evaluation graphs.Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
- Fix all pre-merge checks with AI
✨ Finishing Touches: 💡 1
-
🧪 Generate unit tests (beta)
-
- Create a new PR
🛠️ Fix failing CI checks: 💡
- Commit to this branch
- Create a new PR
-
- Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed the exact one-file diff and its inference/training integration. No new actionable correctness defect found.
The flag is appropriate for this specific graph, rather than being justified by .eval() alone: make_fx has already materialized edge_energy_deriv's energy-to-edge/spin derivatives and force/virial assembly, with create_graph=False. Both aot_module_simplified lowering and compiled eval execution are inside torch.no_grad(). Training uses the separate torch.compile branch with create_graph=True, and the cache distinguishes training/eval. The changed callback therefore does not suppress the first derivatives needed to return forces or alter force-loss double-backward training.
Checked upstream compile_fx_inner and default inference dispatch in PyTorch tags 2.11.0, 2.12.1 and 2.13.0, covering this repository's compile allowlist: all accept is_inference and propagate it to lowering; default inference dispatch supplies True.
Validation remains incomplete. The reported default test_descriptor_sezm.py run does not establish compiled-model coverage. Please run the existing TestSeZMModelCompile eval cache/shape-change and first-frame tests in test_sezm_model.py, plus its force-loss double-backward test, on a supported PyTorch build. The standard CUDA workflow currently pins torch 2.10.0, which skips these compile tests, so a green CUDA badge alone would not establish this coverage.
At 11:00 UTC, the exact-head test/build/CodeQL workflows report action_required and have not provided test results; pre-commit passed and Read the Docs failed. I ran source/AST checks only: PyTorch is not installed in this review environment, and I did not reproduce the GPU parity or memory measurements.
Agent: dot
Reviewed head: 1e3e40c
|
Following the review, ran the requested compile coverage on this branch head (1e3e40c), on a supported build with CUDA available (RTX 5090, torch 2.12.1+cu130):
For completeness, also on this branch:
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #6056 +/- ##
==========================================
- Coverage 77.74% 77.61% -0.14%
==========================================
Files 1155 1170 +15
Lines 139649 140912 +1263
Branches 5056 5056
==========================================
+ Hits 108573 109363 +790
- Misses 29191 29666 +475
+ Partials 1885 1883 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@OutisLi @wanghan-iapcm @njzjz requesting a review, please. Following the 2026-10-13 discussion we validated the combined stack of our six performance PRs (#6056, #6057, #6060, #6070, #6071, #6072) with a multi-scale gate: diamond + bcc-Cu at ~500 / ~2000 / ~4400 atoms, three columns (pure torch.compile with all custom kernels off / current master / our changes on). All accuracy cells sit within 2.2x the pure baseline's own run-to-run band (criterion: 3x), and no tested atom count regresses in time or memory — full numbers in the PR descriptions. Suggested review order for the six: #6056 -> #6057 -> #6060 (small and self-contained) -> #6070 -> #6071 -> #6072. |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Verified at head 1e3e40c (torch 2.11.0+cpu, CPU only):
_inductor_inference_compileris only reached from the eval branch (aot_module_simplifiedundertorch.no_grad()), so the new kwarg is alwaysis_inference=True; the training graph still goes throughtorch.compileand is unaffected.- With the flag, Inductor's post-grad
reorder_for_localitypass now runs on the eval graph (instrumented: it runs on this head and does not run on the parent commit);config.memory_planningstays off, andreorder_for_peak_memory=Falsefrombuild_inductor_compile_options(inference=True)is still applied. compile_fx_inneracceptsis_inferenceas a keyword in torch 2.11 (checked locally; it defaults toFalse).- Local runs:
TestSeZMModelCompile9 passed;TestSeZMEmbeddingForward::test_compiled_matches_eagerandTestSeZMSpinModel::test_compile_matches_eagerpassed. Compiled eval matches eager for energy, force and virial on the tiny test system. - The CPU CI jobs install torch 2.11.0+cpu, so the compile tests in
test_sezm_model.pyran in CI and passed. The VRAM reduction was not reproduced here because no GPU was available.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One non-blocking clarity suggestion for _inductor_inference_compiler in deepmd/pt/model/model/sezm_model.py:2370-2388: please express this callback's inference-only contract directly with is_inference=True.
The callback is defined only in the evaluation branch, and aot_module_simplified is invoked under torch.no_grad(). Training takes the separate torch.compile branch above. Consequently, not self.training adds a dependency on mutable model state where the graph's inference-only role already determines the value. The added sentence about the training graph keeping False also suggests that training passes through this callback, which it does not.
Please keep the hardware-specific memory measurements and before/after validation results in the PR description rather than in the implementation comment. A short explanation that this callback lowers the forward-only AOTAutograd inference graph is sufficient release-grade documentation.
No functional blocker found in this change. Local validation: 61 tests and 122 subtests passed; actual CPU Inductor inference matched eager energy, force, and virial for 3- and 5-atom inputs; the evaluation cache remained usable after a train/eval switch; training backward produced finite gradients for 82 parameters. GPU memory savings were not independently measured.
|
Thank you both for the reviews! @wanghan-iapcm thanks for the verification and approval. @OutisLi addressed in 31067de: the callback now states its contract directly with Validation on my side (CUDA, torch 2.12.1+cu130): |
Motivation
The custom inference compiler entry in
sezm_model.pycallscompile_fx_inner(fx_gm, fx_inputs)without theis_inferenceflag, soInductor plans memory for the eval graph under training assumptions:
outputs are allocated early and autograd-visible intermediates stay
alive, inflating peak VRAM of compiled inference.
torch.compile's own default path already passesis_inference=Truethrough its inference compiler (see
compile_fx'sinference_compiler = functools.partial(fw_compiler_base, is_inference=True)).This PR makes the custom entry match that behavior for eval graphs;
training graphs are unchanged (
is_inference=not self.training).Change
One-line semantic change plus a comment in
deepmd/pt/model/model/sezm_model.py::_inductor_inference_compiler.Measurements
SeZM model, N=4096 atoms (~646k edges), RTX 5090, torch 2.12.1+cu130,
DP_COMPILE_INFER=1, throughput mode (16 evals, P50):Accuracy (same input, master vs this PR, compiled eval): energy rel diff
2.2e-10; max force diff 4.5e-6 with max |f| ~ 3.9 — outputs agree to
rounding level; this is a scheduling-only change.
The gain grows when more of the peak is scheduler-visible intermediates:
on a heavier memory-optimized kernel stack the same flag measures −12.5%
peak VRAM (23647→20696 MB), and −13.9% on an H800.
Tests
source/tests/pt/model/test_descriptor_sezm.pyon this branch, defaultconfiguration: 61 passed, 122 subtests passed.
Summary by CodeRabbit