From 1e3e40c7d2c91ab54edb1fd91453da9c8acb68fb Mon Sep 17 00:00:00 2001 From: longjunhao Date: Wed, 30 Sep 2026 22:40:24 +0800 Subject: [PATCH 1/2] 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/2] 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