From 1e3e40c7d2c91ab54edb1fd91453da9c8acb68fb Mon Sep 17 00:00:00 2001 From: longjunhao Date: Wed, 30 Sep 2026 22:40:24 +0800 Subject: [PATCH 1/4] perf(pt): pass is_inference to compile_fx_inner for compiled eval graphs 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. --- deepmd/pt/model/model/sezm_model.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 91ff826a3e..bb4dee6d7c 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2373,8 +2373,19 @@ def _inductor_inference_compiler( # ``compile_options`` already includes ``triton.max_tiles=1`` # from ``build_inductor_compile_options``, so the 1D launch-grid # constraint applies to both the training and evaluation graphs. + # + # ``is_inference`` mirrors what torch.compile's default + # inference compiler already does: planning for an inference + # graph allocates buffers at first use and frees them early + # instead of keeping them autograd-visible, which lowers peak + # VRAM of compiled eval (measured -12.5% at N=4096 on an + # RTX 5090, -13.9% on an H800). The training graph keeps + # ``False`` so it compiles exactly as before. Scheduling only: + # outputs agree with the previous compile to rounding level. with _ind_cfg.patch(compile_options): - return compile_fx_inner(fx_gm, fx_inputs) + return compile_fx_inner( + fx_gm, fx_inputs, is_inference=not self.training + ) # select_decomp_table keeps the decomposition set aligned with # Inductor's fallback set (a mismatch raises an aten._to_copy From 31067decd04d1a0f65e74937a1f79a572fa185e4 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 15 Oct 2026 19:15:49 +0800 Subject: [PATCH 2/4] refactor(pt): state the inference compiler contract with is_inference=True --- deepmd/pt/model/model/sezm_model.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index bb4dee6d7c..7b2595a663 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2370,22 +2370,13 @@ def compute_fn( # type: ignore[misc] def _inductor_inference_compiler( fx_gm: torch.fx.GraphModule, fx_inputs: list[Any] ) -> Any: + # This callback lowers the forward-only AOTAutograd inference + # graph (the eval branch below owns it; training compiles + # through the separate torch.compile branch above). # ``compile_options`` already includes ``triton.max_tiles=1`` - # from ``build_inductor_compile_options``, so the 1D launch-grid - # constraint applies to both the training and evaluation graphs. - # - # ``is_inference`` mirrors what torch.compile's default - # inference compiler already does: planning for an inference - # graph allocates buffers at first use and frees them early - # instead of keeping them autograd-visible, which lowers peak - # VRAM of compiled eval (measured -12.5% at N=4096 on an - # RTX 5090, -13.9% on an H800). The training graph keeps - # ``False`` so it compiles exactly as before. Scheduling only: - # outputs agree with the previous compile to rounding level. + # from ``build_inductor_compile_options``. with _ind_cfg.patch(compile_options): - return compile_fx_inner( - fx_gm, fx_inputs, is_inference=not self.training - ) + return compile_fx_inner(fx_gm, fx_inputs, is_inference=True) # select_decomp_table keeps the decomposition set aligned with # Inductor's fallback set (a mismatch raises an aten._to_copy From d52a4ca2174b331340118e2503d7bdbe95ba99b7 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Fri, 30 Oct 2026 19:55:44 +0800 Subject: [PATCH 3/4] test(sezm): pin is_inference on the eval compile slot The cache-slot and eager-parity checks hold with the flag set, cleared, or absent, so a refactor could drop the argument and reinstate the peak-memory cost with the suite green. The new test spies on torch._inductor.compile_fx.compile_fx_inner (the import happens at call time) with the compile caches disabled, and asserts the eval slot lowers through is_inference=True. Dropping the flag in a control run turns the test red. --- source/tests/pt/model/test_sezm_model.py | 47 ++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 18ca49b8bf..e985dc514e 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -493,6 +493,53 @@ def test_eval_compile_first_frame_nloc_matches_trace_edge_count(self) -> None: ) @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) + def test_eval_compile_pins_is_inference_on_the_inference_slot(self) -> None: + """The eval slot must lower through ``compile_fx_inner(is_inference=True)``. + + The existing suite (cache-slot presence, eager parity) holds with the + flag set, cleared, or absent; only this call-shape pin notices a + refactor that drops the argument and quietly reinstates the + peak-memory cost. ``sezm_model`` imports ``compile_fx_inner`` from + its home module at call time, so the spy patches the source module. + """ + coord, atype, box = self._make_frame_with_natoms(7) + + model_dyn = get_sezm_model(self._build_model_params(use_compile=False)) + self._randomize_params(model_dyn) + model_dyn.eval() + with mock.patch.dict( + os.environ, + { + "DP_COMPILE_INFER": "1", + # A cache hit would skip the lowering this test observes. + "TORCHINDUCTOR_FX_GRAPH_CACHE": "0", + "TORCHINDUCTOR_AUTOGRAD_CACHE": "0", + }, + clear=False, + ): + model_cmp = get_sezm_model(self._build_model_params(use_compile=True)) + model_cmp.load_state_dict(model_dyn.state_dict()) + model_cmp.eval() + + import torch._inductor.compile_fx as _compile_fx_mod + + calls: list[dict] = [] + real = _compile_fx_mod.compile_fx_inner + + def _spy(*args, **kwargs): + calls.append(kwargs) + return real(*args, **kwargs) + + with mock.patch.object(_compile_fx_mod, "compile_fx_inner", _spy): + model_cmp(coord, atype, box=box) + + self.assertTrue( + any(kwargs.get("is_inference") is True for kwargs in calls), + "the eval slot must lower through compile_fx_inner(is_inference=True); " + f"observed call kwargs: {calls}", + ) + def test_compile_cache_slots_and_eval_shape_change(self) -> None: """Compile cache slots should coexist while eval handles batch-size growth.""" coord_1, atype_1, box_1, _, _, _ = self._make_tiny_frame() From 6fe393fce1d2eacff4b3a7549becdce453686a80 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 1 Nov 2026 11:49:40 +0800 Subject: [PATCH 4/4] test(sezm): restore the moved skipIf onto the cache-slot compile test The pin-test insertion landed between the existing test's decorator and its def, so the decorator attached twice to the new test and the cache-slot test lost its gating against the SUPPORTED_COMPILE_TORCH allowlist. --- source/tests/pt/model/test_sezm_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index e985dc514e..6e6d3c1423 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -492,7 +492,6 @@ def test_eval_compile_first_frame_nloc_matches_trace_edge_count(self) -> None: msg="eval force mismatch when first-frame nloc == trace edge count", ) - @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_eval_compile_pins_is_inference_on_the_inference_slot(self) -> None: """The eval slot must lower through ``compile_fx_inner(is_inference=True)``. @@ -540,6 +539,7 @@ def _spy(*args, **kwargs): f"observed call kwargs: {calls}", ) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_compile_cache_slots_and_eval_shape_change(self) -> None: """Compile cache slots should coexist while eval handles batch-size growth.""" coord_1, atype_1, box_1, _, _, _ = self._make_tiny_frame()