diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 35a87e3294..0e7b17c1a1 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2322,11 +2322,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. + # from ``build_inductor_compile_options``. with _ind_cfg.patch(compile_options): - return compile_fx_inner(fx_gm, fx_inputs) + 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 diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 18ca49b8bf..6e6d3c1423 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -492,6 +492,53 @@ 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) + 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}", + ) + @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."""