fix freeze

Author: OutisLi

deepmd/pt/model/descriptor/sezm_nn/triton/radial_mix.py

@@ -127,13 +127,65 @@ def radial_mix_reference(
 def _radial_mix_backward_reference(
     grad_out: Tensor, compact: Tensor, x_local: Tensor, channel_basis: Tensor, lmax: int
 ) -> tuple[Tensor, Tensor]:
-    """Eager backward returning ``(grad_compact, grad_x_local)`` via autograd."""
-    with torch.enable_grad():
-        compact_req = compact.detach().requires_grad_(True)
-        x_req = x_local.detach().requires_grad_(True)
-        out = radial_mix_reference(compact_req, x_req, channel_basis, lmax)
-        grad_compact, grad_x = torch.autograd.grad(out, [compact_req, x_req], grad_out)
-    return grad_compact, grad_x
+    """Closed-form eager backward of :func:`radial_mix_reference`.
+
+    Gradients are evaluated analytically per diagonal block, mirroring the
+    contractions of the Triton backward. A closed form is required rather than a
+    nested ``autograd.grad``: this routine is the CPU backend of the
+    ``radial_mix_block_bwd`` operator, which carries no autograd formula and is
+    consequently dispatched under ``_AutoDispatchBelowAutograd`` whenever the
+    force graph is replayed without grad (the SeZM ``.pt2`` freeze does so under
+    :func:`torch.no_grad`). That guard excludes the autograd key, so a nested
+    ``autograd.grad`` would observe an output without a ``grad_fn``.
+
+    Parameters
+    ----------
+    grad_out : Tensor
+        Upstream gradient with shape ``(E, reduced_dim, C)``.
+    compact : Tensor
+        Projected radial degree kernel with shape ``(E, degree_kernel_size, R)``.
+    x_local : Tensor
+        Edge-local reduced features with shape ``(E, reduced_dim, C)``.
+    channel_basis : Tensor
+        Per-rank channel basis with shape ``(R, C)``.
+    lmax : int
+        Maximum spherical-harmonic degree.
+
+    Returns
+    -------
+    tuple[Tensor, Tensor]
+        Gradients ``(grad_compact, grad_x_local)``, matching ``compact`` and
+        ``x_local`` in shape respectively.
+    """
+    n_edge, reduced_dim, channels = x_local.shape
+    grad_x_local = torch.zeros_like(x_local)
+    grad_compact = torch.zeros_like(compact)
+    for coeff0, comp0, num_l in _block_layout(int(lmax)):
+        # Forward of this block (see ``radial_mix_reference``):
+        #   out[e, o, c] = sum_{i, r} K[e, o, i, r] * x[e, i, c] * cb[r, c]
+        #   with K[e, o, i, r] = compact[e, comp0 + i * num_l + o, r].
+        k_block = (
+            compact[:, comp0 : comp0 + num_l * num_l, :]
+            .reshape(n_edge, num_l, num_l, -1)
+            .permute(0, 2, 1, 3)
+        )  # (E, o, i, R)
+        x_block = x_local[:, coeff0 : coeff0 + num_l, :]  # (E, i, C)
+        g_block = grad_out[:, coeff0 : coeff0 + num_l, :]  # (E, o, C)
+
+        # grad_x[e, i, c] = sum_r cb[r, c] * sum_o K[e, o, i, r] * g[e, o, c].
+        gx = torch.einsum("eoir,eoc->eicr", k_block, g_block)  # (E, i, C, R)
+        grad_x_local[:, coeff0 : coeff0 + num_l, :] += torch.einsum(
+            "eicr,rc->eic", gx, channel_basis
+        )
+
+        # grad_K[e, o, i, r] = sum_c cb[r, c] * x[e, i, c] * g[e, o, c], scattered
+        # back to the compact slot comp0 + i * num_l + o. The shared m = +-1
+        # blocks address the same slots, so the in-place add accumulates both.
+        gk = torch.einsum("eoc,eic,rc->eoir", g_block, x_block, channel_basis)
+        grad_compact[:, comp0 : comp0 + num_l * num_l, :] += gk.permute(
+            0, 2, 1, 3
+        ).reshape(n_edge, num_l * num_l, -1)
+    return grad_compact, grad_x_local
 
 
 # ======================================================================

source/tests/pt/model/test_sezm_export.py

@@ -13,6 +13,7 @@
 import contextlib
 import copy
 import json
+import os
 import tempfile
 import unittest
 import zipfile
@@ -216,18 +217,24 @@ class TestSeZMExportPipeline(_ClearDefaultDeviceTestCase):
     it must reproduce the eager result exactly.  Drift here implies a
     bug in ``forward_common_lower_exportable`` or the dynamic-shape
     spec, not in AOTI.  The pipeline is built once per class because
-    ``make_fx`` and ``.pte`` round-trip dominate wall time.
+    ``make_fx`` and ``.pte`` round-trip dominate wall time.  Subclasses
+    set ``TRITON_INFER`` to drive the identical pipeline through the
+    opt-in Triton inference kernels.
     """
 
+    # ``DP_TRITON_INFER`` policy applied while the model is constructed.
+    TRITON_INFER = "0"
+
     @classmethod
     def setUpClass(cls) -> None:
         super().setUpClass()
         try:
-            cls.model = _build_tiny_sezm_model()
-            cls.sample_inputs = _make_sample(cls.model, nloc=7, start=2)
-            cls.traced, cls.loaded, cls._pte_tmp = cls._build_pipeline(
-                cls.model, cls.sample_inputs
-            )
+            with mock.patch.dict(os.environ, {"DP_TRITON_INFER": cls.TRITON_INFER}):
+                cls.model = _build_tiny_sezm_model()
+                cls.sample_inputs = _make_sample(cls.model, nloc=7, start=2)
+                cls.traced, cls.loaded, cls._pte_tmp = cls._build_pipeline(
+                    cls.model, cls.sample_inputs
+                )
         except Exception:
             super().tearDownClass()
             raise
@@ -327,6 +334,54 @@ def test_loaded_pte_matches_eager_different_shape(self) -> None:
         )
 
 
+class TestSeZMExportPipelineTritonInfer(TestSeZMExportPipeline):
+    """The same trace / ``.pte`` pipeline exercised with ``DP_TRITON_INFER=1``.
+
+    Inheriting the parity suite asserts the Triton-enabled model still traces,
+    exports, and reloads, and that the loaded ``.pte`` reproduces its eager
+    forward — including the force path, whose custom ``*_bwd`` ops run inside
+    ``_AutoDispatchBelowAutograd`` during the export's no-grad replay and must
+    therefore be closed-form.  Two checks are added: the captured graph routes
+    through the custom ops, and the Triton ``.pte`` matches the dense (Triton-off)
+    inference, proving ``DP_TRITON_INFER`` swaps the implementation without
+    changing results.
+    """
+
+    TRITON_INFER = "1"
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        try:
+            with mock.patch.dict(os.environ, {"DP_TRITON_INFER": "0"}):
+                dense_model = _build_tiny_sezm_model()
+            cls.dense_out = _eager_forward(dense_model, cls.sample_inputs)
+        except Exception:
+            super().tearDownClass()
+            raise
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        try:
+            if hasattr(cls, "dense_out"):
+                delattr(cls, "dense_out")
+        finally:
+            super().tearDownClass()
+
+    def test_force_graph_carries_triton_ops(self) -> None:
+        """``DP_TRITON_INFER=1`` must route the descriptor through the custom ops."""
+        code = self.traced.code
+        self.assertIn("radial_mix_block_bwd", code)
+        self.assertIn("rotate_to_local", code)
+
+    def test_loaded_pte_matches_dense(self) -> None:
+        """The Triton-on ``.pte`` reproduces the dense-path inference."""
+        loaded_out = self.loaded(*self.sample_inputs)
+        self._assert_dict_allclose(
+            self.dense_out, loaded_out, context="triton .pte vs dense eager"
+        )
+
+
 class _FrozenPt2Fixture(_ClearDefaultDeviceTestCase):
     """Shared setUp/tearDown: freeze a tiny SeZM checkpoint to ``.pt2`` once.