From e334724c348fd13564b02280acf5ff30beb1d885 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sat, 10 Oct 2026 15:44:26 +0800 Subject: [PATCH 01/10] perf(pt): assemble Wigner-D behind a custom op under compile Functionalization expands the in-place block scatter into a chain of full-size zeros/index_put adds plus a realized transpose (~4 GB of HBM traffic per frame at the SeZM N=4096 workload). A wigner_d_assemble custom op keeps the eager scatter and hands the graph two contiguous tensors; the traced assembly takes it whenever a fake-tensor mode is active, so the eager path is untouched. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 163 +++++++++++++++--- 1 file changed, 139 insertions(+), 24 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index b283e3a7e5..d420dec53d 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -45,6 +45,88 @@ _TORCH_RELEASE = Version(torch.__version__).release[:2] +def _wigner_d_assemble_impl( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble the dense block-diagonal matrix and its contiguous transpose. + + In-place slice assignments are exactly what Inductor's functionalization + cannot keep in place: traced, they expand into a chain of full-size + ``zeros + index_put + add`` tensors plus a realized transpose (~4 GB of + HBM traffic at the SeZM N=4096 workload). Behind a custom op the eager + scatter stays a scatter -- one buffer, one transpose copy. + """ + ref = blocks[0] + n_edge = ref.shape[0] + d_full = torch.zeros(n_edge, dim_full, dim_full, dtype=ref.dtype, device=ref.device) + d_full[:, 0, 0] = 1.0 + off = 1 + for block in blocks[:n_regular]: + size = block.shape[1] + d_full[:, off : off + size, off : off + size] = block + off += size + if n_regular < len(blocks): + po = poly_lmin * poly_lmin + d_full[:, po:, po:] = blocks[n_regular] + return d_full, d_full.transpose(-1, -2).contiguous() + + +_wigner_d_assemble_op = torch.library.custom_op( + "sezm_nn::wigner_d_assemble", mutates_args=() +)(_wigner_d_assemble_impl) + + +@_wigner_d_assemble_op.register_fake +def _( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, +) -> tuple[torch.Tensor, torch.Tensor]: + ref = blocks[0] + d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full)) + return d_full, d_full.transpose(-1, -2).contiguous() + + +def _wigner_d_assemble_setup( + ctx: Any, inputs: Any, output: Any +) -> None: + blocks, _, poly_lmin, n_regular = inputs + ctx.blocks = blocks + ctx.poly_lmin = poly_lmin + ctx.n_regular = n_regular + + +def _wigner_d_assemble_backward( + ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor +) -> tuple[list[torch.Tensor], None, None, None]: + # Each block lands in D as itself and in Dt transposed, so its gradient + # is the matching D-slice plus the transpose of the matching Dt-slice. + # The [0, 0] constant receives no gradient. + grads = [] + off = 1 + for block in ctx.blocks[: ctx.n_regular]: + size = block.shape[1] + g = grad_d_full[:, off : off + size, off : off + size] + gt = grad_dt_full[:, off : off + size, off : off + size] + grads.append(g + gt.transpose(-1, -2)) + off += size + if ctx.n_regular < len(ctx.blocks): + po = ctx.poly_lmin * ctx.poly_lmin + g = grad_d_full[:, po:, po:] + gt = grad_dt_full[:, po:, po:] + grads.append(g + gt.transpose(-1, -2)) + return grads, None, None, None + + +_wigner_d_assemble_op.register_autograd( + _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup +) + + class CaseCoefficients(nn.Module): """ Polynomial tables for one magnitude-ordered branch of the quaternion Wigner path. @@ -546,62 +628,63 @@ def forward( eps=self.eps, ) n_edge = edge_quaternion.shape[0] - D_full = torch.zeros( - n_edge, - self.dim_full, - self.dim_full, - dtype=edge_quaternion.dtype, - device=edge_quaternion.device, - ) - D_full[:, 0, 0] = 1.0 + + # Degree blocks are computed first; the dense assembly then either + # runs in place (eager) or behind the opaque wigner_d_assemble op + # (compiled). Under tracing, the in-place slice assignments below + # would be functionalized into a chain of full-size zeros/index_put + # adds plus a realized transpose; the custom op keeps the eager + # scatter and hands the graph two plain contiguous tensors. + blocks: list[torch.Tensor] = [] + poly_block: torch.Tensor | None = None if self.lmax >= 1: with nvtx_range("WignerD/l1"): - D_full[:, 1:4, 1:4] = self._compute_l1_block(edge_quaternion) + blocks.append(self._compute_l1_block(edge_quaternion)) if self.lmax >= 2: with nvtx_range("WignerD/l2"): - D_full[:, 4:9, 4:9] = self._compute_l2_block(edge_quaternion) + blocks.append(self._compute_l2_block(edge_quaternion)) if self.lmax >= 3: if self.lmax >= 4: with nvtx_range("WignerD/l3l4"): D_l3, D_l4 = self._compute_l3l4_blocks(edge_quaternion) - D_full[:, 9:16, 9:16] = D_l3 - D_full[:, 16:25, 16:25] = D_l4 + blocks.append(D_l3) + blocks.append(D_l4) else: with nvtx_range("WignerD/l3"): - D_full[:, 9:16, 9:16] = self._compute_l3_block(edge_quaternion) + blocks.append(self._compute_l3_block(edge_quaternion)) if self.lmax >= 5: if self.lmax >= 6: with nvtx_range("WignerD/l5l6"): D_l5, D_l6 = self._compute_l5l6_blocks(edge_quaternion) - D_full[:, 25:36, 25:36] = D_l5 - D_full[:, 36:49, 36:49] = D_l6 + blocks.append(D_l5) + blocks.append(D_l6) else: with nvtx_range("WignerD/l5"): - D_full[:, 25:36, 25:36] = self._compute_l5_block(edge_quaternion) + blocks.append(self._compute_l5_block(edge_quaternion)) if self.lmax >= 7: if self.lmax >= 8: with nvtx_range("WignerD/l7l8"): D_l7, D_l8 = self._compute_l7l8_blocks(edge_quaternion) - D_full[:, 49:64, 49:64] = D_l7 - D_full[:, 64:81, 64:81] = D_l8 + blocks.append(D_l7) + blocks.append(D_l8) else: with nvtx_range("WignerD/l7"): - D_full[:, 49:64, 49:64] = self._compute_l7_block(edge_quaternion) + blocks.append(self._compute_l7_block(edge_quaternion)) if self.lmax >= 9: if self.lmax >= 10: with nvtx_range("WignerD/l9l10"): D_l9, D_l10 = self._compute_l9l10_blocks(edge_quaternion) - D_full[:, 81:100, 81:100] = D_l9 - D_full[:, 100:121, 100:121] = D_l10 + blocks.append(D_l9) + blocks.append(D_l10) else: with nvtx_range("WignerD/l9"): - D_full[:, 81:100, 81:100] = self._compute_l9_block(edge_quaternion) + blocks.append(self._compute_l9_block(edge_quaternion)) if self.lmax >= self.poly_lmin: with nvtx_range("WignerD/polynomial"): @@ -616,7 +699,7 @@ def forward( self.poly_coeffs, dtype=self.dtype, ) - D_poly = self._wigner_d_pair_to_real( + poly_block = self._wigner_d_pair_to_real( D_re, D_im, ( @@ -628,7 +711,39 @@ def forward( lmax=self.lmax, lmin=self.poly_lmin, ) - D_full[:, self.poly_offset :, self.poly_offset :] = D_poly + + # ``is_compiling`` alone is not enough here: this module is traced + # through make_fx fake tensors, where the flag stays False. A fake + # tensor carries its mode, so probe for that instead. + tracing = ( + torch.compiler.is_compiling() + or getattr(blocks[0], "fake_mode", None) is not None + if blocks + else False + ) + if blocks and tracing: + return _wigner_d_assemble_op( + blocks + ([poly_block] if poly_block is not None else []), + self.dim_full, + self.poly_lmin, + len(blocks), + ) + + D_full = torch.zeros( + n_edge, + self.dim_full, + self.dim_full, + dtype=edge_quaternion.dtype, + device=edge_quaternion.device, + ) + D_full[:, 0, 0] = 1.0 + off = 1 + for block in blocks: + size = block.shape[1] + D_full[:, off : off + size, off : off + size] = block + off += size + if poly_block is not None: + D_full[:, self.poly_offset :, self.poly_offset :] = poly_block # Consumers address the inverse rotation through explicit strides or # PyTorch strided operators, so the transpose can share D_full's storage. From aea99479029413e4c53ddeed59d0327cec86d763 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sat, 10 Oct 2026 15:58:33 +0800 Subject: [PATCH 02/10] perf(pt): return the inverse Wigner-D as a view, not a copy The flash and rotation-back kernels already address the rotation through explicit strides, so the contiguous transpose the assembly produced was pure storage: one 660 MB copy per frame that existed only to be read back. The assemble op now returns D_full alone and the caller hands out a transpose view; consumers read it in place and no transpose copy ever exists. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index d420dec53d..1c5e3d9f8c 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -50,14 +50,18 @@ def _wigner_d_assemble_impl( dim_full: int, poly_lmin: int, n_regular: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Assemble the dense block-diagonal matrix and its contiguous transpose. +) -> torch.Tensor: + """Assemble the dense block-diagonal matrix; the transpose stays a view. In-place slice assignments are exactly what Inductor's functionalization cannot keep in place: traced, they expand into a chain of full-size ``zeros + index_put + add`` tensors plus a realized transpose (~4 GB of HBM traffic at the SeZM N=4096 workload). Behind a custom op the eager - scatter stays a scatter -- one buffer, one transpose copy. + scatter stays a scatter. Only ``D_full`` is produced here; the caller + makes ``Dt_full`` a plain transpose view outside the op, and every + consumer addresses it through strides (the flash and rotation kernels + take explicit stride arguments; the eager consumers are strided torch + operators). """ ref = blocks[0] n_edge = ref.shape[0] @@ -71,7 +75,7 @@ def _wigner_d_assemble_impl( if n_regular < len(blocks): po = poly_lmin * poly_lmin d_full[:, po:, po:] = blocks[n_regular] - return d_full, d_full.transpose(-1, -2).contiguous() + return d_full _wigner_d_assemble_op = torch.library.custom_op( @@ -85,15 +89,12 @@ def _( dim_full: int, poly_lmin: int, n_regular: int, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> torch.Tensor: ref = blocks[0] - d_full = ref.new_zeros((ref.shape[0], dim_full, dim_full)) - return d_full, d_full.transpose(-1, -2).contiguous() + return ref.new_zeros((ref.shape[0], dim_full, dim_full)) -def _wigner_d_assemble_setup( - ctx: Any, inputs: Any, output: Any -) -> None: +def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: blocks, _, poly_lmin, n_regular = inputs ctx.blocks = blocks ctx.poly_lmin = poly_lmin @@ -101,24 +102,20 @@ def _wigner_d_assemble_setup( def _wigner_d_assemble_backward( - ctx: Any, grad_d_full: torch.Tensor, grad_dt_full: torch.Tensor + ctx: Any, grad_d_full: torch.Tensor ) -> tuple[list[torch.Tensor], None, None, None]: - # Each block lands in D as itself and in Dt transposed, so its gradient - # is the matching D-slice plus the transpose of the matching Dt-slice. - # The [0, 0] constant receives no gradient. + # The transpose view's gradient reaches D_full through the ordinary + # autograd of aten.transpose, so only D-slices arrive here. The [0, 0] + # constant receives no gradient. grads = [] off = 1 for block in ctx.blocks[: ctx.n_regular]: size = block.shape[1] - g = grad_d_full[:, off : off + size, off : off + size] - gt = grad_dt_full[:, off : off + size, off : off + size] - grads.append(g + gt.transpose(-1, -2)) + grads.append(grad_d_full[:, off : off + size, off : off + size]) off += size if ctx.n_regular < len(ctx.blocks): po = ctx.poly_lmin * ctx.poly_lmin - g = grad_d_full[:, po:, po:] - gt = grad_dt_full[:, po:, po:] - grads.append(g + gt.transpose(-1, -2)) + grads.append(grad_d_full[:, po:, po:]) return grads, None, None, None @@ -722,12 +719,15 @@ def forward( else False ) if blocks and tracing: - return _wigner_d_assemble_op( + d_full = _wigner_d_assemble_op( blocks + ([poly_block] if poly_block is not None else []), self.dim_full, self.poly_lmin, len(blocks), ) + # Stride-addressed consumers read the inverse rotation straight + # from the shared storage; no transpose copy ever exists. + return d_full, d_full.transpose(-1, -2) D_full = torch.zeros( n_edge, From 40fac5908d005fda5fbb90e297e30c170a25dc4d Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 15 Oct 2026 20:34:22 +0800 Subject: [PATCH 03/10] refactor(pt): share the Wigner-D assembly and inverse-layout exit, drop tensor retention --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 86 ++++++++++--------- 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index 1c5e3d9f8c..f9be1ab451 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -96,9 +96,12 @@ def _( def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: blocks, _, poly_lmin, n_regular = inputs - ctx.blocks = blocks + # Only the layout metadata is needed to slice the cotangent; storing the + # block tensors would extend every activation's lifetime past the + # forward, which the backward never reads. + ctx.block_widths = [int(block.shape[1]) for block in blocks[:n_regular]] + ctx.has_poly_tail = n_regular < len(blocks) ctx.poly_lmin = poly_lmin - ctx.n_regular = n_regular def _wigner_d_assemble_backward( @@ -109,11 +112,10 @@ def _wigner_d_assemble_backward( # constant receives no gradient. grads = [] off = 1 - for block in ctx.blocks[: ctx.n_regular]: - size = block.shape[1] + for size in ctx.block_widths: grads.append(grad_d_full[:, off : off + size, off : off + size]) off += size - if ctx.n_regular < len(ctx.blocks): + if ctx.has_poly_tail: po = ctx.poly_lmin * ctx.poly_lmin grads.append(grad_d_full[:, po:, po:]) return grads, None, None, None @@ -626,12 +628,13 @@ def forward( ) n_edge = edge_quaternion.shape[0] - # Degree blocks are computed first; the dense assembly then either - # runs in place (eager) or behind the opaque wigner_d_assemble op - # (compiled). Under tracing, the in-place slice assignments below - # would be functionalized into a chain of full-size zeros/index_put - # adds plus a realized transpose; the custom op keeps the eager - # scatter and hands the graph two plain contiguous tensors. + # Degree blocks are computed first; the dense assembly below is one + # shared implementation, reached directly in eager execution and + # behind the opaque wigner_d_assemble operator under tracing (where + # functionalization would otherwise expand the in-place scatter into + # a chain of full-size zeros/index_put adds plus a realized + # transpose). The operator hands the graph one contiguous tensor; + # the inverse rotation leaves through the shared transpose-view exit. blocks: list[torch.Tensor] = [] poly_block: torch.Tensor | None = None @@ -718,41 +721,44 @@ def forward( if blocks else False ) - if blocks and tracing: - d_full = _wigner_d_assemble_op( - blocks + ([poly_block] if poly_block is not None else []), + # One assembly implementation serves both routes: the eager call + # goes straight to the differentiable scatter (autograd flows through + # the slice assignments), the traced call hides it behind the opaque + # operator so functionalization cannot expand the scatter. + all_blocks = blocks + ([poly_block] if poly_block is not None else []) + if not all_blocks: + # lmax == 0: the rotation acts only on the scalar channel. + D_full = torch.zeros( + n_edge, + self.dim_full, + self.dim_full, + dtype=edge_quaternion.dtype, + device=edge_quaternion.device, + ) + D_full[:, 0, 0] = 1.0 + elif tracing: + D_full = _wigner_d_assemble_op( + all_blocks, self.dim_full, self.poly_lmin, len(blocks), ) - # Stride-addressed consumers read the inverse rotation straight - # from the shared storage; no transpose copy ever exists. - return d_full, d_full.transpose(-1, -2) - - D_full = torch.zeros( - n_edge, - self.dim_full, - self.dim_full, - dtype=edge_quaternion.dtype, - device=edge_quaternion.device, - ) - D_full[:, 0, 0] = 1.0 - off = 1 - for block in blocks: - size = block.shape[1] - D_full[:, off : off + size, off : off + size] = block - off += size - if poly_block is not None: - D_full[:, self.poly_offset :, self.poly_offset :] = poly_block - - # Consumers address the inverse rotation through explicit strides or - # PyTorch strided operators, so the transpose can share D_full's storage. + else: + D_full = _wigner_d_assemble_impl( + all_blocks, + self.dim_full, + self.poly_lmin, + len(blocks), + ) + + # One inverse-layout exit shared by both routes: the compatibility + # policy below is preserved as-is instead of being bypassed by the + # traced branch. Stride-addressed consumers read the inverse + # rotation from the shared storage; PyTorch 2.11 Inductor cannot + # safely lower that escaping transpose view, so it is materialized + # there exactly as the eager route does. Dt_full = D_full.transpose(-1, -2) if self._materialize_inverse_rotation: - # PyTorch 2.11 Inductor cannot safely lower this escaping transpose - # view after the slice assignments that assemble D_full. The - # materialized layout keeps the compiled graph semantically - # identical; later releases retain the shared-storage view. Dt_full = Dt_full.contiguous() return D_full, Dt_full From e98c691897acd7bd89004eb5633473847775d39d Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 15 Oct 2026 20:34:22 +0800 Subject: [PATCH 04/10] test(pt): focused coverage for the Wigner-D assemble operator --- .../pt/model/test_descriptor_sezm_triton.py | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..9094065d64 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2053,3 +2053,209 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self): if __name__ == "__main__": unittest.main() + + +class TestSeZMWignerDAssembleOp(unittest.TestCase): + """Focused coverage for the op-bound dense Wigner-D assembly.""" + + WIDTHS = (3, 5, 7) # lmax = 3 ladder + + def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64): + generator = torch.Generator(device="cuda").manual_seed(seed) + return [ + torch.randn(n_edge, w, w, device="cuda", dtype=dtype, generator=generator) + for w in widths + ] + + def _reference_assemble(self, blocks, poly, dim_full, poly_lmin): + ref = blocks[0] + dense = ref.new_zeros(ref.shape[0], dim_full, dim_full) + dense[:, 0, 0] = 1.0 + off = 1 + for block in blocks: + size = block.shape[1] + dense[:, off : off + size, off : off + size] = block + off += size + if poly is not None: + po = poly_lmin * poly_lmin + dense[:, po:, po:] = poly + return dense + + def _loss(self, dense, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + w1 = torch.randn( + dense.shape, device="cuda", dtype=dense.dtype, generator=generator + ) + w2 = torch.randn( + dense.shape, device="cuda", dtype=dense.dtype, generator=generator + ) + # cotangent contributions through both D and its transpose + return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() + + @_GPU_KERNELS + def test_op_matches_reference_with_grads(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_impl, + _wigner_d_assemble_op, + ) + + n_edge = 64 + for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): + with self.subTest(widths=widths, tail=poly_lmin is not None): + blocks = self._make_blocks(widths, n_edge, seed=5) + poly = None + n_regular = len(blocks) + if poly_lmin is not None: + poly_w = 4 + poly = self._make_blocks([poly_w], n_edge, seed=7)[0] + dim_full = sum(widths) + 1 if poly is None else poly_lmin**2 + poly_w + args = ( + blocks + ([poly] if poly is not None else []), + dim_full, + poly_lmin or 11, + n_regular, + ) + + leaves_op = [ + b.clone().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + dense_op = _wigner_d_assemble_op(leaves_op, args[1], args[2], args[3]) + leaves_ref = [ + b.clone().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + dense_ref = _wigner_d_assemble_impl( + leaves_ref, args[1], args[2], args[3] + ) + self.assertTrue(torch.equal(dense_op.detach(), dense_ref.detach())) + self.assertTrue( + torch.equal( + dense_op.detach(), + self._reference_assemble(blocks, poly, args[1], args[2]), + ) + ) + + grads_op = torch.autograd.grad( + self._loss(dense_op, 11), leaves_op, allow_unused=True + ) + grads_ref = torch.autograd.grad( + self._loss(dense_ref, 11), leaves_ref, allow_unused=True + ) + for g1, g2 in zip(grads_op, grads_ref, strict=True): + torch.testing.assert_close(g1, g2, atol=0.0, rtol=0.0) + + # fp64 gradcheck on the operator route (small edge count) + small = [ + b[:3].clone().detach().requires_grad_(True) + for b in (blocks + ([poly] if poly is not None else [])) + ] + self.assertTrue( + torch.autograd.gradcheck( + lambda *bs: _wigner_d_assemble_op( + list(bs), args[1], args[2], args[3] + ), + small, + ) + ) + + @_GPU_KERNELS + def test_opcheck_schema_fake_autograd(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_op, + ) + + for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): + with self.subTest(widths=widths, tail=poly_lmin is not None): + blocks = self._make_blocks(widths, 8, seed=9) + args_tail = [] + n_regular = len(blocks) + if poly_lmin is not None: + args_tail = self._make_blocks([4], 8, seed=13) + blocks = blocks + args_tail + dim_full = sum(widths) + 1 if not args_tail else poly_lmin**2 + 4 + saved = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device("cuda"): + torch.library.opcheck( + _wigner_d_assemble_op, + (blocks, dim_full, poly_lmin or 11, n_regular), + ) + finally: + torch.set_default_device(saved) + + @_GPU_KERNELS + def test_traced_matches_eager_and_layout_policy(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + generator = torch.Generator(device="cuda").manual_seed(21) + quat = torch.randn(32, 4, device="cuda", generator=generator) + quat = quat / quat.norm(dim=-1, keepdim=True) + calc = WignerDCalculator(lmax=3, dtype=torch.float32).cuda() + calc.eval() + + saved_device = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device("cuda"): + d_eager, dt_eager = calc.forward(quat) + d_compiled, dt_compiled = torch.compile(calc.forward)(quat) + # The assembly op itself is bitwise (see the reference test); + # the compiled ladder reassociates the degree-block math, so the + # cross-route comparison holds at rounding level. + torch.testing.assert_close(d_compiled, d_eager, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(dt_compiled, dt_eager, atol=1e-5, rtol=1e-5) + finally: + torch.set_default_device(saved_device) + + # The inverse-layout policy is one shared exit: both routes honor + # the compatibility flag instead of the traced route bypassing it. + def _run_routes(): + saved = torch.get_default_device() + torch.set_default_device(None) + try: + with torch.device("cuda"): + _, dt_eager_i = calc.forward(quat) + _, dt_compiled_i = torch.compile(calc.forward)(quat) + return dt_eager_i, dt_compiled_i + finally: + torch.set_default_device(saved) + + calc._materialize_inverse_rotation = False + dt_view, dt_view_c = _run_routes() + self.assertFalse(dt_view.is_contiguous()) + self.assertFalse(dt_view_c.is_contiguous()) + calc._materialize_inverse_rotation = True + dt_dense, dt_dense_c = _run_routes() + self.assertTrue(dt_dense.is_contiguous()) + self.assertTrue(dt_dense_c.is_contiguous()) + torch.testing.assert_close(dt_dense, dt_view, atol=0.0, rtol=0.0) + + @_GPU_KERNELS + def test_lmax_zero_and_empty_edges(self): + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + quat = torch.randn(8, 4, device="cuda") + quat = quat / quat.norm(dim=-1, keepdim=True) + calc0 = WignerDCalculator(lmax=0, dtype=torch.float32).cuda() + calc0.eval() + d0, dt0 = calc0.forward(quat) + expected = torch.zeros(8, 1, 1, device="cuda") + expected[:, 0, 0] = 1.0 + self.assertTrue(torch.equal(d0, expected)) + self.assertTrue(torch.equal(dt0, expected)) + + calc3 = WignerDCalculator(lmax=3, dtype=torch.float32).cuda() + calc3.eval() + d_empty, dt_empty = calc3.forward(torch.empty(0, 4, device="cuda")) + self.assertEqual(tuple(d_empty.shape), (0, 16, 16)) + self.assertEqual(tuple(dt_empty.shape), (0, 16, 16)) + + +if __name__ == "__main__": + unittest.main() From c5feab50604f2e6cf1a8fd639ca4db57fb2c05b8 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 22 Oct 2026 19:40:33 +0800 Subject: [PATCH 05/10] fix(pt): keep the Wigner-D assembly importable on torch < 2.4 The operator registration (torch.library.custom_op / register_fake / register_autograd) requires torch >= 2.4 while the declared minimum is 2.1, and this module is imported unconditionally through the descriptor chain: on a 2.1-2.3 install the import itself raised AttributeError before the eager fallback could ever run. The registration is now gated on the API's presence; without it both the eager and the traced route keep the autograd-through-scatter implementation, and a subprocess test reloads the module with the API hidden to pin that behavior. The make_fx fake-mode half of the tracing probe is now covered (a symbolic trace selects the opaque operator, asserted on the graph code), and the two pure-PyTorch operator tests (values/grads/gradcheck and the opcheck schema/fake/autograd contract) no longer require CUDA, so they run in the CPU CI. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 42 ++--- .../pt/model/test_descriptor_sezm_triton.py | 144 ++++++++++++++++-- 2 files changed, 156 insertions(+), 30 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index f9be1ab451..2bb27a5d62 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -78,20 +78,25 @@ def _wigner_d_assemble_impl( return d_full -_wigner_d_assemble_op = torch.library.custom_op( - "sezm_nn::wigner_d_assemble", mutates_args=() -)(_wigner_d_assemble_impl) - - -@_wigner_d_assemble_op.register_fake -def _( - blocks: list[torch.Tensor], - dim_full: int, - poly_lmin: int, - n_regular: int, -) -> torch.Tensor: - ref = blocks[0] - return ref.new_zeros((ref.shape[0], dim_full, dim_full)) +_wigner_d_assemble_op: Any = None +if hasattr(torch.library, "custom_op"): + # torch.library.custom_op / register_fake / register_autograd arrived in + # torch 2.4 while the declared minimum is 2.1: on a 2.1-2.3 install the + # operator stays unregistered (None) and both the eager and the traced + # route below keep the autograd-through-scatter implementation. + _wigner_d_assemble_op = torch.library.custom_op( + "sezm_nn::wigner_d_assemble", mutates_args=() + )(_wigner_d_assemble_impl) + + @_wigner_d_assemble_op.register_fake + def _( + blocks: list[torch.Tensor], + dim_full: int, + poly_lmin: int, + n_regular: int, + ) -> torch.Tensor: + ref = blocks[0] + return ref.new_zeros((ref.shape[0], dim_full, dim_full)) def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: @@ -121,9 +126,10 @@ def _wigner_d_assemble_backward( return grads, None, None, None -_wigner_d_assemble_op.register_autograd( - _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup -) +if _wigner_d_assemble_op is not None: + _wigner_d_assemble_op.register_autograd( + _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup + ) class CaseCoefficients(nn.Module): @@ -736,7 +742,7 @@ def forward( device=edge_quaternion.device, ) D_full[:, 0, 0] = 1.0 - elif tracing: + elif tracing and _wigner_d_assemble_op is not None: D_full = _wigner_d_assemble_op( all_blocks, self.dim_full, diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 9094065d64..ab60dcbd17 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2060,10 +2060,10 @@ class TestSeZMWignerDAssembleOp(unittest.TestCase): WIDTHS = (3, 5, 7) # lmax = 3 ladder - def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64): - generator = torch.Generator(device="cuda").manual_seed(seed) + def _make_blocks(self, widths, n_edge, seed, dtype=torch.float64, device="cuda"): + generator = torch.Generator(device=device).manual_seed(seed) return [ - torch.randn(n_edge, w, w, device="cuda", dtype=dtype, generator=generator) + torch.randn(n_edge, w, w, device=device, dtype=dtype, generator=generator) for w in widths ] @@ -2082,9 +2082,9 @@ def _reference_assemble(self, blocks, poly, dim_full, poly_lmin): return dense def _loss(self, dense, seed): - generator = torch.Generator(device="cuda").manual_seed(seed) + generator = torch.Generator(device=dense.device).manual_seed(seed) w1 = torch.randn( - dense.shape, device="cuda", dtype=dense.dtype, generator=generator + dense.shape, device=dense.device, dtype=dense.dtype, generator=generator ) w2 = torch.randn( dense.shape, device="cuda", dtype=dense.dtype, generator=generator @@ -2092,22 +2092,27 @@ def _loss(self, dense, seed): # cotangent contributions through both D and its transpose return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() - @_GPU_KERNELS + @unittest.skipUnless( + hasattr(torch.library, "custom_op"), + "the operator registration needs torch >= 2.4", + ) def test_op_matches_reference_with_grads(self): + """Pure-PyTorch operator math; runs wherever the CPU CI runs.""" from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( _wigner_d_assemble_impl, _wigner_d_assemble_op, ) + device = "cuda" if _CUDA else "cpu" n_edge = 64 for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): with self.subTest(widths=widths, tail=poly_lmin is not None): - blocks = self._make_blocks(widths, n_edge, seed=5) + blocks = self._make_blocks(widths, n_edge, seed=5, device=device) poly = None n_regular = len(blocks) if poly_lmin is not None: poly_w = 4 - poly = self._make_blocks([poly_w], n_edge, seed=7)[0] + poly = self._make_blocks([poly_w], n_edge, seed=7, device=device)[0] dim_full = sum(widths) + 1 if poly is None else poly_lmin**2 + poly_w args = ( blocks + ([poly] if poly is not None else []), @@ -2159,25 +2164,30 @@ def test_op_matches_reference_with_grads(self): ) ) - @_GPU_KERNELS + @unittest.skipUnless( + hasattr(torch.library, "custom_op"), + "the operator registration needs torch >= 2.4", + ) def test_opcheck_schema_fake_autograd(self): + """Schema/fake/autograd contract; pure PyTorch, CPU CI included.""" from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( _wigner_d_assemble_op, ) + device = "cuda" if _CUDA else "cpu" for widths, poly_lmin in ((self.WIDTHS, None), ((3, 5), 3)): with self.subTest(widths=widths, tail=poly_lmin is not None): - blocks = self._make_blocks(widths, 8, seed=9) + blocks = self._make_blocks(widths, 8, seed=9, device=device) args_tail = [] n_regular = len(blocks) if poly_lmin is not None: - args_tail = self._make_blocks([4], 8, seed=13) + args_tail = self._make_blocks([4], 8, seed=13, device=device) blocks = blocks + args_tail dim_full = sum(widths) + 1 if not args_tail else poly_lmin**2 + 4 saved = torch.get_default_device() torch.set_default_device(None) try: - with torch.device("cuda"): + with torch.device(device): torch.library.opcheck( _wigner_d_assemble_op, (blocks, dim_full, poly_lmin or 11, n_regular), @@ -2185,6 +2195,116 @@ def test_opcheck_schema_fake_autograd(self): finally: torch.set_default_device(saved) + @unittest.skipUnless( + hasattr(torch.library, "custom_op"), + "the operator registration needs torch >= 2.4", + ) + def test_make_fx_trace_selects_the_assemble_op(self): + """The fake-mode probe half: make_fx tracing picks the opaque op. + + ``torch.compile`` (covered by the CUDA test below) satisfies the + probe through ``is_compiling``; make_fx keeps that flag False and + relies on the fake-mode disjunct, which this check exercises. The + graph-node assertion locks in the mechanism: with the probe + silently false, the traced route would fall back to the raw + scatter and every value comparison would still pass. + """ + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + def _run(quat): + return calc.forward(quat) + + calc = WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + graph = make_fx(_run, tracing_mode="symbolic", _allow_non_fake_inputs=True)( + quat + ) + self.assertIn("sezm_nn.wigner_d_assemble", graph.code) + + def test_import_and_tracing_fall_back_without_custom_op(self): + """A torch 2.1-2.3 install must import the module and stay eager. + + Emulated in a subprocess: the module reloads with + ``torch.library.custom_op`` hidden for exactly the duration of its + top-level re-execution (torch's own lazily importing users of the + API stay intact), after which the operator is unregistered, the + eager forward/backward works, and make_fx tracing keeps the eager + implementation instead of the operator node. + """ + import pathlib + import subprocess + import sys + import textwrap + + import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod + + tree = pathlib.Path(wignerd_mod.__file__).resolve().parents[5] + child = textwrap.dedent( + """ + import pathlib + import sys + + tree = pathlib.Path(sys.argv[1]) + sys.meta_path = [ + f + for f in sys.meta_path + if "editable" + not in (type(f).__module__ + "." + type(f).__name__).lower() + and "editable" not in getattr(f, "__name__", "").lower() + ] + sys.path.insert(0, str(tree)) + import importlib + + import torch + import torch.distributed._functional_collectives # noqa: F401 + + import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd + + assert wd._wigner_d_assemble_op is not None, "control: op registers" + saved_custom_op = torch.library.custom_op + del torch.library.custom_op + wd = importlib.reload(wd) + torch.library.custom_op = saved_custom_op + assert wd._wigner_d_assemble_op is None, "fallback: op stays out" + + calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu() + calc.eval() + quat = torch.randn(4, 4, dtype=torch.float64, device="cpu") + quat = (quat / quat.norm(dim=-1, keepdim=True)).requires_grad_(True) + dense, _dt = calc.forward(quat) + dense.sum().backward() + + from torch.fx.experimental.proxy_tensor import make_fx + + def _run(q): + return calc.forward(q) + + graph = make_fx( + _run, tracing_mode="symbolic", _allow_non_fake_inputs=True + )(quat) + assert ( + "sezm_nn.wigner_d_assemble" not in graph.code + ), "without the API the traced route keeps the eager path" + print("fallback-ok") + """ + ) + result = subprocess.run( + [sys.executable, "-c", child, str(tree)], + capture_output=True, + text=True, + timeout=300, + ) + self.assertEqual( + result.returncode, + 0, + f"fallback child failed:\n{result.stdout}\n{result.stderr}", + ) + self.assertIn("fallback-ok", result.stdout) + @_GPU_KERNELS def test_traced_matches_eager_and_layout_policy(self): from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( From 9aeb0b910b342dba6e3cd438d1d9f1bbdc919d43 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 25 Oct 2026 13:34:56 +0800 Subject: [PATCH 06/10] fix(pt): fall back to the scatter assembly on any registration failure The schema inference for list[torch.Tensor] parameters only supports this signature on later torches still (2.4-2.6 have the custom-op API but reject this signature outright; 2.1-2.3 lack the API), so the presence gate alone was insufficient. Registration now runs under a try/except and any failure keeps the autograd-through-scatter route for both the eager and the traced call. The subprocess fallback test covers both incompatible classes with separate emulations -- the API hidden entirely (2.1-2.3) and the API replaced by a stub raising the schema-inference error (2.4-2.6) -- with the registration control asserted only when the API exists, so the test also holds on a genuinely old install. The CPU-run _loss helper creates its second weight on the operand's device instead of hardcoding CUDA. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 22 ++++-- .../pt/model/test_descriptor_sezm_triton.py | 76 ++++++++++++------- 2 files changed, 64 insertions(+), 34 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index 2bb27a5d62..42a72e557a 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -79,11 +79,14 @@ def _wigner_d_assemble_impl( _wigner_d_assemble_op: Any = None -if hasattr(torch.library, "custom_op"): +try: # torch.library.custom_op / register_fake / register_autograd arrived in - # torch 2.4 while the declared minimum is 2.1: on a 2.1-2.3 install the - # operator stays unregistered (None) and both the eager and the traced - # route below keep the autograd-through-scatter implementation. + # torch 2.4, and the schema inference for ``list[torch.Tensor]`` + # parameters only supports this signature on later releases still + # (torch 2.4-2.6 reject it outright; 2.1-2.3 lack the API). The + # declared minimum is 2.1, so any registration failure keeps the + # autograd-through-scatter implementation for both the eager and the + # traced route below instead of failing the import. _wigner_d_assemble_op = torch.library.custom_op( "sezm_nn::wigner_d_assemble", mutates_args=() )(_wigner_d_assemble_impl) @@ -97,6 +100,8 @@ def _( ) -> torch.Tensor: ref = blocks[0] return ref.new_zeros((ref.shape[0], dim_full, dim_full)) +except Exception: # pragma: no cover - depends on the installed torch + _wigner_d_assemble_op = None def _wigner_d_assemble_setup(ctx: Any, inputs: Any, output: Any) -> None: @@ -127,9 +132,12 @@ def _wigner_d_assemble_backward( if _wigner_d_assemble_op is not None: - _wigner_d_assemble_op.register_autograd( - _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup - ) + try: + _wigner_d_assemble_op.register_autograd( + _wigner_d_assemble_backward, setup_context=_wigner_d_assemble_setup + ) + except Exception: # pragma: no cover - depends on the installed torch + _wigner_d_assemble_op = None class CaseCoefficients(nn.Module): diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index ab60dcbd17..fc46922bcd 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2087,7 +2087,7 @@ def _loss(self, dense, seed): dense.shape, device=dense.device, dtype=dense.dtype, generator=generator ) w2 = torch.randn( - dense.shape, device="cuda", dtype=dense.dtype, generator=generator + dense.shape, device=dense.device, dtype=dense.dtype, generator=generator ) # cotangent contributions through both D and its transpose return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() @@ -2225,14 +2225,18 @@ def _run(quat): ) self.assertIn("sezm_nn.wigner_d_assemble", graph.code) - def test_import_and_tracing_fall_back_without_custom_op(self): - """A torch 2.1-2.3 install must import the module and stay eager. - - Emulated in a subprocess: the module reloads with - ``torch.library.custom_op`` hidden for exactly the duration of its - top-level re-execution (torch's own lazily importing users of the - API stay intact), after which the operator is unregistered, the - eager forward/backward works, and make_fx tracing keeps the eager + def test_import_and_tracing_fall_back_when_registration_fails(self): + """Unsupported torches must import the module and stay eager. + + Two emulations run in subprocesses, one per incompatible class: + ``hide`` deletes ``torch.library.custom_op`` for the module's + re-execution (torch 2.1-2.3, where the API is absent), and + ``reject`` replaces it with a stub raising the schema-inference + error (torch 2.4-2.6, where the API exists but rejects a + ``list[torch.Tensor]`` parameter). In both, torch's own lazily + importing users of the API stay intact (the change is scoped to + the reload window), the operator stays unregistered, the eager + forward/backward works, and make_fx tracing keeps the eager implementation instead of the operator node. """ import pathlib @@ -2248,7 +2252,8 @@ def test_import_and_tracing_fall_back_without_custom_op(self): import pathlib import sys - tree = pathlib.Path(sys.argv[1]) + mode = sys.argv[1] + tree = pathlib.Path(sys.argv[2]) sys.meta_path = [ f for f in sys.meta_path @@ -2264,11 +2269,25 @@ def test_import_and_tracing_fall_back_without_custom_op(self): import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd - assert wd._wigner_d_assemble_op is not None, "control: op registers" - saved_custom_op = torch.library.custom_op - del torch.library.custom_op + api_present = hasattr(torch.library, "custom_op") + if api_present: + assert wd._wigner_d_assemble_op is not None, "control: registers" + + saved_custom_op = getattr(torch.library, "custom_op", None) + if mode == "hide": + del torch.library.custom_op + elif mode == "reject": + def _rejected(*args, **kwargs): + raise TypeError( + "cannot infer schema for blocks: list[torch.Tensor]" + ) + + torch.library.custom_op = _rejected + else: + raise AssertionError(f"unknown mode {mode}") wd = importlib.reload(wd) - torch.library.custom_op = saved_custom_op + if saved_custom_op is not None: + torch.library.custom_op = saved_custom_op assert wd._wigner_d_assemble_op is None, "fallback: op stays out" calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu() @@ -2288,22 +2307,25 @@ def _run(q): )(quat) assert ( "sezm_nn.wigner_d_assemble" not in graph.code - ), "without the API the traced route keeps the eager path" + ), "without a registration the traced route keeps the eager path" print("fallback-ok") """ ) - result = subprocess.run( - [sys.executable, "-c", child, str(tree)], - capture_output=True, - text=True, - timeout=300, - ) - self.assertEqual( - result.returncode, - 0, - f"fallback child failed:\n{result.stdout}\n{result.stderr}", - ) - self.assertIn("fallback-ok", result.stdout) + for mode in ("hide", "reject"): + with self.subTest(emulation=mode): + result = subprocess.run( + [sys.executable, "-c", child, mode, str(tree)], + capture_output=True, + text=True, + timeout=300, + ) + self.assertEqual( + result.returncode, + 0, + f"fallback child ({mode}) failed:" + f"\n{result.stdout}\n{result.stderr}", + ) + self.assertIn("fallback-ok", result.stdout) @_GPU_KERNELS def test_traced_matches_eager_and_layout_policy(self): From 64696312ac758c86be71a8e89767e97a9a30763c Mon Sep 17 00:00:00 2001 From: longjunhao Date: Fri, 30 Oct 2026 19:40:23 +0800 Subject: [PATCH 07/10] fix(sezm): probe torch.compiler.is_compiling before eager Wigner use torch 2.1-2.2 lack the callable, so the eager fallback path itself raised AttributeError before reaching the operator-availability check. Resolve the callable once at import and treat it as not-compiling when absent. The operator tests now gate on registration success instead of torch.library.custom_op presence (2.4-2.6 have the API but reject the schema), the subprocess hide mode tolerates an already-absent attribute, and the successful-registration control asserts the parent-observed state instead of assuming API presence implies registration. --- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 22 ++++-- .../pt/model/test_descriptor_sezm_triton.py | 67 ++++++++++++++----- 2 files changed, 65 insertions(+), 24 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index 42a72e557a..62c7bfda0d 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -44,6 +44,15 @@ _TORCH_RELEASE = Version(torch.__version__).release[:2] +# ``torch.compiler.is_compiling`` is absent before torch 2.3. The eager +# Wigner path must stay runnable on those versions, so resolve the callable +# once at import time instead of calling it unconditionally. +_COMPILER_IS_COMPILING = getattr( + getattr(torch, "compiler", None), + "is_compiling", + None, +) + def _wigner_d_assemble_impl( blocks: list[torch.Tensor], @@ -728,13 +737,14 @@ def forward( # ``is_compiling`` alone is not enough here: this module is traced # through make_fx fake tensors, where the flag stays False. A fake - # tensor carries its mode, so probe for that instead. - tracing = ( - torch.compiler.is_compiling() - or getattr(blocks[0], "fake_mode", None) is not None - if blocks - else False + # tensor carries its mode, so probe for that instead. On torch older + # than 2.3 the callable does not exist and the fake-tensor probe + # carries the decision alone. + compiling = ( + _COMPILER_IS_COMPILING() if _COMPILER_IS_COMPILING is not None else False ) + fake_mode = bool(blocks) and getattr(blocks[0], "fake_mode", None) is not None + tracing = compiling or fake_mode # One assembly implementation serves both routes: the eager call # goes straight to the differentiable scatter (autograd flows through # the slice assignments), the traced call hides it behind the opaque diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index fc46922bcd..1c6ee2ca54 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -41,6 +41,9 @@ from deepmd.pt.model.descriptor.sezm_nn.so2 import ( DynamicRadialDegreeMixer, ) +from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + _wigner_d_assemble_op as _WIGNER_ASSEMBLE_OP, +) from deepmd.pt_expt.kernels.triton.sezm import ( TRITON_AVAILABLE, ) @@ -60,6 +63,12 @@ _CUDA = torch.cuda.is_available() +# The assembly operator registers only where the running torch accepts its +# schema (torch >= 2.7); torch 2.4-2.6 have ``torch.library.custom_op`` but +# reject a ``list[torch.Tensor]`` parameter, so tests that call the operator +# directly gate on registration success, not on API presence. +_WIGNER_OP_REGISTERED = _WIGNER_ASSEMBLE_OP is not None + # All SeZM Triton kernels ship together (every per-module availability flag # reduces to "the triton package imports"), so one gate covers every # GPU-kernel test class. Dispatch- and parsing-only classes stay ungated. @@ -2093,8 +2102,8 @@ def _loss(self, dense, seed): return (dense * w1).sum() + (dense.transpose(-1, -2) * w2).sum() @unittest.skipUnless( - hasattr(torch.library, "custom_op"), - "the operator registration needs torch >= 2.4", + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", ) def test_op_matches_reference_with_grads(self): """Pure-PyTorch operator math; runs wherever the CPU CI runs.""" @@ -2165,8 +2174,8 @@ def test_op_matches_reference_with_grads(self): ) @unittest.skipUnless( - hasattr(torch.library, "custom_op"), - "the operator registration needs torch >= 2.4", + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", ) def test_opcheck_schema_fake_autograd(self): """Schema/fake/autograd contract; pure PyTorch, CPU CI included.""" @@ -2196,8 +2205,8 @@ def test_opcheck_schema_fake_autograd(self): torch.set_default_device(saved) @unittest.skipUnless( - hasattr(torch.library, "custom_op"), - "the operator registration needs torch >= 2.4", + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", ) def test_make_fx_trace_selects_the_assemble_op(self): """The fake-mode probe half: make_fx tracing picks the opaque op. @@ -2229,15 +2238,23 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): """Unsupported torches must import the module and stay eager. Two emulations run in subprocesses, one per incompatible class: - ``hide`` deletes ``torch.library.custom_op`` for the module's - re-execution (torch 2.1-2.3, where the API is absent), and - ``reject`` replaces it with a stub raising the schema-inference - error (torch 2.4-2.6, where the API exists but rejects a - ``list[torch.Tensor]`` parameter). In both, torch's own lazily - importing users of the API stay intact (the change is scoped to - the reload window), the operator stays unregistered, the eager + ``hide`` removes ``torch.library.custom_op`` for the module's + re-execution (torch 2.1-2.3, where the API is absent -- on a torch + that really lacks it the attribute is already gone and nothing is + deleted), and ``reject`` replaces it with a stub raising the + schema-inference error (torch 2.4-2.6, where the API exists but + rejects a ``list[torch.Tensor]`` parameter). In both, torch's own + lazily importing users of the API stay intact (the change is scoped + to the reload window), the operator stays unregistered, the eager forward/backward works, and make_fx tracing keeps the eager implementation instead of the operator node. + + The child first imports the module untouched and asserts the + registration state its parent observed (passed as an argument), so + the successful-registration control is measured against reality + instead of API presence: on a torch whose registration legitimately + fails at import, the control asserts the import-time fallback, and + the operator-specific tests above are skipped by their decorator. """ import pathlib import subprocess @@ -2254,6 +2271,7 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): mode = sys.argv[1] tree = pathlib.Path(sys.argv[2]) + expected_registered = sys.argv[3] == "1" sys.meta_path = [ f for f in sys.meta_path @@ -2269,13 +2287,15 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd - api_present = hasattr(torch.library, "custom_op") - if api_present: + if expected_registered: assert wd._wigner_d_assemble_op is not None, "control: registers" + else: + assert wd._wigner_d_assemble_op is None, "control: import fallback" saved_custom_op = getattr(torch.library, "custom_op", None) if mode == "hide": - del torch.library.custom_op + if saved_custom_op is not None: + del torch.library.custom_op elif mode == "reject": def _rejected(*args, **kwargs): raise TypeError( @@ -2288,6 +2308,9 @@ def _rejected(*args, **kwargs): wd = importlib.reload(wd) if saved_custom_op is not None: torch.library.custom_op = saved_custom_op + elif mode == "reject": + # The stub was injected over a missing attribute; remove it. + del torch.library.custom_op assert wd._wigner_d_assemble_op is None, "fallback: op stays out" calc = wd.WignerDCalculator(lmax=1, dtype=torch.float64).cpu() @@ -2311,10 +2334,18 @@ def _run(q): print("fallback-ok") """ ) + registered = wignerd_mod._wigner_d_assemble_op is not None for mode in ("hide", "reject"): - with self.subTest(emulation=mode): + with self.subTest(emulation=mode, registered=registered): result = subprocess.run( - [sys.executable, "-c", child, mode, str(tree)], + [ + sys.executable, + "-c", + child, + mode, + str(tree), + "1" if registered else "0", + ], capture_output=True, text=True, timeout=300, From 9114811bfe8a3b02f8818f7f745a0fbe7cbd8e4d Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 1 Nov 2026 11:35:46 +0800 Subject: [PATCH 08/10] fix(sezm): keep export traces on the eager Wigner assembly AOTInductor lowers a plain custom_op to a proxy-executor node that the C++ package loader resolves by name, and nothing registers sezm_nn::wigner_d_assemble there, so a .pt2 frozen through the traced branch failed to load with 'Could not find schema'. The freeze seam (both the plain and the with-comm export) now traces under an explicit suppression context that keeps the eager scatter; the in-process inference path is unchanged. Regression: under the context the make_fx graph carries no operator node and reproduces the unsuppressed values bitwise. Two review hardenings on the probe: a torch-version-gated check that fails (not skips) when registration regresses on a schema-capable torch, and a mechanism pin for the is_compiling disjunct -- a truthy probe against real tensors must route through the operator. Empirical correction recorded while testing: under current dynamo the assembly graph-breaks out of torch.compile regions, so the compiled route does not reach that disjunct today; the docstrings now say so and the fake-mode half remains the live one. Also unify the module import style flagged by CodeQL. --- deepmd/pt/entrypoints/freeze_pt2.py | 61 +++++---- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 48 ++++++- .../pt/model/test_descriptor_sezm_triton.py | 118 ++++++++++++++++-- 3 files changed, 187 insertions(+), 40 deletions(-) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..f3d2384d91 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -819,15 +819,20 @@ def _export_with_comm_artifact( ) from torch._inductor import config as inductor_config - sample_inputs = _make_comm_sample_inputs(model, nloc=7, device=torch.device("cpu")) - traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) - exported = torch.export.export( - traced, - sample_inputs, - dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, ) + + sample_inputs = _make_comm_sample_inputs(model, nloc=7, device=torch.device("cpu")) + with assemble_op_suppressed_for_export(): + traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) + exported = torch.export.export( + traced, + sample_inputs, + dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) if target_device.type != "cpu": from torch.export.passes import ( @@ -1038,22 +1043,32 @@ def _freeze_sezm_to_pt2( # Each model's exportable signature matches its sample tuple positionally # (energy / native-spin edge ABI, or virtual-spin nlist ABI), so a single # splat covers all three contracts. - log.info("Tracing the lower graph on CPU (make_fx)...") - traced = model.forward_common_lower_exportable(*sample_inputs_cpu) - - # Output key order is read from the static FX output node. A CUDA-target - # trace may already contain CUDA-only custom operators and therefore must - # not be replayed on the CPU tracing inputs. - output_keys = traced_output_keys(traced) - - log.info("Exporting the traced graph (torch.export)...") - exported = torch.export.export( - traced, - sample_inputs_cpu, - dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, + # The Wigner assembly operator cannot ride an AOTInductor archive (the + # C++ package loader resolves custom ops by name and nothing registers + # sezm_nn::wigner_d_assemble there), so the export trace keeps the eager + # scatter the operator wraps. The in-process inference path is + # unaffected. + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, ) + + log.info("Tracing the lower graph on CPU (make_fx)...") + with assemble_op_suppressed_for_export(): + traced = model.forward_common_lower_exportable(*sample_inputs_cpu) + + # Output key order is read from the static FX output node. A + # CUDA-target trace may already contain CUDA-only custom operators + # and therefore must not be replayed on the CPU tracing inputs. + output_keys = traced_output_keys(traced) + + log.info("Exporting the traced graph (torch.export)...") + exported = torch.export.export( + traced, + sample_inputs_cpu, + dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) # move_to_device_pass handles FakeTensor device propagation cleanly; diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index 62c7bfda0d..f0229c1b66 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -11,15 +11,24 @@ ) import math +from contextlib import ( + contextmanager, +) from itertools import ( permutations, product, ) from typing import ( + TYPE_CHECKING, Any, ClassVar, ) +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + ) + import torch import torch.nn as nn from packaging.version import ( @@ -53,6 +62,26 @@ None, ) +# Export traces keep the eager assembly: AOTInductor lowers a plain +# custom_op to a proxy-executor node resolved by name in the C++ runtime, +# where nothing registers ``sezm_nn::wigner_d_assemble`` -- an exported +# .pt2 built through the operator fails to load. The in-process inference +# path lowers in the same process and resolves the operator through the +# dispatcher, so only the freeze/export seam raises this flag. +_assemble_op_suppressed_for_export = False + + +@contextmanager +def assemble_op_suppressed_for_export() -> Iterator[None]: + """Route export traces through the eager scatter instead of the operator.""" + global _assemble_op_suppressed_for_export + saved = _assemble_op_suppressed_for_export + _assemble_op_suppressed_for_export = True + try: + yield + finally: + _assemble_op_suppressed_for_export = saved + def _wigner_d_assemble_impl( blocks: list[torch.Tensor], @@ -735,11 +764,14 @@ def forward( lmin=self.poly_lmin, ) - # ``is_compiling`` alone is not enough here: this module is traced - # through make_fx fake tensors, where the flag stays False. A fake - # tensor carries its mode, so probe for that instead. On torch older - # than 2.3 the callable does not exist and the fake-tensor probe - # carries the decision alone. + # ``is_compiling`` alone is not enough: the production eval trace + # goes through make_fx fake tensors, where the flag stays False, and + # under current dynamo the assembly graph-breaks out of torch.compile + # regions (the flag is False in the fallback too). A fake tensor + # carries its mode, so probe for that as well; the ``is_compiling`` + # disjunct stays for compiled regions that enclose this code. On + # torch older than 2.3 the callable does not exist and the fake-tensor + # probe carries the decision alone. compiling = ( _COMPILER_IS_COMPILING() if _COMPILER_IS_COMPILING is not None else False ) @@ -760,7 +792,11 @@ def forward( device=edge_quaternion.device, ) D_full[:, 0, 0] = 1.0 - elif tracing and _wigner_d_assemble_op is not None: + elif ( + tracing + and _wigner_d_assemble_op is not None + and not _assemble_op_suppressed_for_export + ): D_full = _wigner_d_assemble_op( all_blocks, self.dim_full, diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 1c6ee2ca54..59733c158e 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -30,6 +30,9 @@ ) import torch +from packaging.version import ( + Version, +) from torch.fx.experimental.proxy_tensor import ( make_fx, ) @@ -41,9 +44,7 @@ from deepmd.pt.model.descriptor.sezm_nn.so2 import ( DynamicRadialDegreeMixer, ) -from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( - _wigner_d_assemble_op as _WIGNER_ASSEMBLE_OP, -) +import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod from deepmd.pt_expt.kernels.triton.sezm import ( TRITON_AVAILABLE, ) @@ -67,6 +68,7 @@ # schema (torch >= 2.7); torch 2.4-2.6 have ``torch.library.custom_op`` but # reject a ``list[torch.Tensor]`` parameter, so tests that call the operator # directly gate on registration success, not on API presence. +_WIGNER_ASSEMBLE_OP = wignerd_mod._wigner_d_assemble_op _WIGNER_OP_REGISTERED = _WIGNER_ASSEMBLE_OP is not None # All SeZM Triton kernels ship together (every per-module availability flag @@ -2211,12 +2213,13 @@ def test_opcheck_schema_fake_autograd(self): def test_make_fx_trace_selects_the_assemble_op(self): """The fake-mode probe half: make_fx tracing picks the opaque op. - ``torch.compile`` (covered by the CUDA test below) satisfies the - probe through ``is_compiling``; make_fx keeps that flag False and - relies on the fake-mode disjunct, which this check exercises. The - graph-node assertion locks in the mechanism: with the probe - silently false, the traced route would fall back to the raw - scatter and every value comparison would still pass. + The production eval path lowers through make_fx, whose fake + tensors carry the mode this check exercises (``is_compiling`` + stays False there; under current dynamo the assembly graph-breaks + out of torch.compile regions, so the fake-mode disjunct is the + live one today). The graph-node assertion locks in the mechanism: + with the probe silently false, the traced route would fall back to + the raw scatter and every value comparison would still pass. """ from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( WignerDCalculator, @@ -2234,6 +2237,101 @@ def _run(quat): ) self.assertIn("sezm_nn.wigner_d_assemble", graph.code) + @unittest.skipIf( + Version(torch.__version__).release < (2, 7), + "the schema is only accepted from torch 2.7", + ) + def test_operator_registers_on_schema_capable_torch(self): + """A registration regression on a capable torch must fail, not skip. + + The operator tests skip on ``_WIGNER_OP_REGISTERED``, so a regression + that silently turns the flag off would leave the suite green with + the traced route reverted to the scatter. This check is gated on + the torch release instead of the observed flag and therefore fails + when registration breaks on a schema-capable torch (CI pins 2.10 + and 2.11). + """ + self.assertIsNotNone( + wignerd_mod._wigner_d_assemble_op, + "wigner_d_assemble failed to register on a schema-capable torch", + ) + + def test_export_suppression_keeps_the_eager_assembly(self): + """Export traces keep the scatter: AOTI archives cannot carry the op. + + AOTInductor lowers a plain custom_op to a proxy-executor node + resolved by name in the C++ runtime, where nothing registers + ``sezm_nn::wigner_d_assemble`` -- an exported .pt2 built through + the operator fails to load. Under the freeze seam's suppression + context the traced graph must carry the eager scatter instead. + """ + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, + ) + + def _run(quat): + return calc.forward(quat) + + calc = wignerd_mod.WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + + plain = make_fx( + _run, tracing_mode="symbolic", _allow_non_fake_inputs=True + )(quat) + self.assertIn("sezm_nn.wigner_d_assemble", plain.code) + + with assemble_op_suppressed_for_export(): + suppressed = make_fx( + _run, tracing_mode="symbolic", _allow_non_fake_inputs=True + )(quat) + self.assertNotIn("sezm_nn.wigner_d_assemble", suppressed.code) + torch.testing.assert_close( + suppressed(quat), plain(quat), atol=0.0, rtol=0.0 + ) + + @unittest.skipUnless( + _WIGNER_OP_REGISTERED, + "the operator registration needs a torch that accepts its schema", + ) + def test_is_compiling_half_of_the_probe_selects_the_operator(self): + """The ``is_compiling`` disjunct of the probe must stay live. + + The fake-mode half has its own graph-node pin; this check drives + the other disjunct directly -- a truthy ``is_compiling`` against + real (non-fake) tensors routes the assembly through the operator, + which a dropped disjunct or a broken callable lookup silently + reverts to the scatter. (Under the *current* dynamo the assembly + sits in a graph-break region and torch.compile leaves the flag + False there, so the compiled route itself does not reach this + disjunct yet; the mock pins the mechanism, not dynamo's region + placement.) + """ + calls = [] + real_op = wignerd_mod._wigner_d_assemble_op + + def _spy(*args, **kwargs): + calls.append(1) + return real_op(*args, **kwargs) + + calc = wignerd_mod.WignerDCalculator(lmax=3, dtype=torch.float32).cpu() + calc.eval() + quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") + quat = quat / quat.norm(dim=-1, keepdim=True) + + with ( + mock.patch.object( + wignerd_mod, "_COMPILER_IS_COMPILING", lambda: True + ), + mock.patch.object(wignerd_mod, "_wigner_d_assemble_op", _spy), + ): + out = calc.forward(quat) + self.assertTrue( + bool(calls), + "a truthy is_compiling must route the assembly through the operator", + ) + def test_import_and_tracing_fall_back_when_registration_fails(self): """Unsupported torches must import the module and stay eager. @@ -2261,8 +2359,6 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): import sys import textwrap - import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod - tree = pathlib.Path(wignerd_mod.__file__).resolve().parents[5] child = textwrap.dedent( """ From 1c4882522b21212f777327a411227df217028752 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 1 Nov 2026 03:36:55 +0000 Subject: [PATCH 09/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pt/model/test_descriptor_sezm_triton.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 59733c158e..94be7ea476 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -37,6 +37,7 @@ make_fx, ) +import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod from deepmd.pt.model.descriptor.sezm_nn.indexing import ( build_m_major_index, get_so3_dim_of_lmax, @@ -44,7 +45,6 @@ from deepmd.pt.model.descriptor.sezm_nn.so2 import ( DynamicRadialDegreeMixer, ) -import deepmd.pt.model.descriptor.sezm_nn.wignerd as wignerd_mod from deepmd.pt_expt.kernels.triton.sezm import ( TRITON_AVAILABLE, ) @@ -2277,9 +2277,9 @@ def _run(quat): quat = torch.randn(8, 4, dtype=torch.float32, device="cpu") quat = quat / quat.norm(dim=-1, keepdim=True) - plain = make_fx( - _run, tracing_mode="symbolic", _allow_non_fake_inputs=True - )(quat) + plain = make_fx(_run, tracing_mode="symbolic", _allow_non_fake_inputs=True)( + quat + ) self.assertIn("sezm_nn.wigner_d_assemble", plain.code) with assemble_op_suppressed_for_export(): @@ -2287,9 +2287,7 @@ def _run(quat): _run, tracing_mode="symbolic", _allow_non_fake_inputs=True )(quat) self.assertNotIn("sezm_nn.wigner_d_assemble", suppressed.code) - torch.testing.assert_close( - suppressed(quat), plain(quat), atol=0.0, rtol=0.0 - ) + torch.testing.assert_close(suppressed(quat), plain(quat), atol=0.0, rtol=0.0) @unittest.skipUnless( _WIGNER_OP_REGISTERED, @@ -2321,9 +2319,7 @@ def _spy(*args, **kwargs): quat = quat / quat.norm(dim=-1, keepdim=True) with ( - mock.patch.object( - wignerd_mod, "_COMPILER_IS_COMPILING", lambda: True - ), + mock.patch.object(wignerd_mod, "_COMPILER_IS_COMPILING", lambda: True), mock.patch.object(wignerd_mod, "_wigner_d_assemble_op", _spy), ): out = calc.forward(quat) From f4a1a8750dfc2b62397634cef31e5e4dc5bebea9 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 5 Nov 2026 22:31:28 +0800 Subject: [PATCH 10/10] fix(sezm): thread-safe export suppression at the public seam The suppression moves from the two private freeze call sites into trace_lower_exportable, the seam behind both documented export entries, so any caller handing the traced module to torch.export inherits a loadable graph (AOTInductor resolves custom ops by name in the C++ runtime, where nothing registers sezm_nn::wigner_d_assemble); the in-process runtime trace (trace_and_compile) is untouched and the freeze wrappers are gone. Both export pipelines now assert the traced graph carries no assemble node -- the regression your round asked for at the actual seam; removing the suppression turns them red. The suppression state is thread-local: the process-global boolean corrupted interleaved contexts (thread A's finally restoring False while thread B still traced). The probe's list emptiness check uses len(blocks) > 0 instead of bool(blocks), which Dynamo cannot trace on a list (your torch-2.10 measurement); on the local 2.12 both forms show the same four pre-existing breaks elsewhere (nvtx_range among them), so the delta is not reproducible here -- the rewrite is applied verbatim as the strictly safer form. The registration subprocess warms the package while the editable finder is still alive (deepmd.lib is served by that finder under pip install -e ., and the source tree has no deepmd/lib), then drops the cached deepmd modules except the C extension so the tree path takes over the pure-Python resolution. --- deepmd/pt/entrypoints/freeze_pt2.py | 60 +++++++++---------- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 24 +++++--- deepmd/pt/model/model/sezm_model.py | 28 ++++++--- .../pt/model/test_descriptor_sezm_triton.py | 22 ++++++- source/tests/pt/model/test_sezm_export.py | 23 +++++++ 5 files changed, 106 insertions(+), 51 deletions(-) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index f3d2384d91..8ae7bd3515 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -819,20 +819,17 @@ def _export_with_comm_artifact( ) from torch._inductor import config as inductor_config - from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( - assemble_op_suppressed_for_export, - ) - sample_inputs = _make_comm_sample_inputs(model, nloc=7, device=torch.device("cpu")) - with assemble_op_suppressed_for_export(): - traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) - exported = torch.export.export( - traced, - sample_inputs, - dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, - ) + # The Wigner-assembly suppression lives in the public export seam + # (``trace_lower_exportable``), so every caller inherits it. + traced = model.forward_common_lower_exportable_with_comm(*sample_inputs) + exported = torch.export.export( + traced, + sample_inputs, + dynamic_shapes=_build_with_comm_dynamic_shapes(sample_inputs), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) if target_device.type != "cpu": from torch.export.passes import ( @@ -1048,27 +1045,24 @@ def _freeze_sezm_to_pt2( # sezm_nn::wigner_d_assemble there), so the export trace keeps the eager # scatter the operator wraps. The in-process inference path is # unaffected. - from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( - assemble_op_suppressed_for_export, - ) - log.info("Tracing the lower graph on CPU (make_fx)...") - with assemble_op_suppressed_for_export(): - traced = model.forward_common_lower_exportable(*sample_inputs_cpu) - - # Output key order is read from the static FX output node. A - # CUDA-target trace may already contain CUDA-only custom operators - # and therefore must not be replayed on the CPU tracing inputs. - output_keys = traced_output_keys(traced) - - log.info("Exporting the traced graph (torch.export)...") - exported = torch.export.export( - traced, - sample_inputs_cpu, - dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), - strict=False, - prefer_deferred_runtime_asserts_over_guards=True, - ) + # The Wigner-assembly suppression lives in the public export seam + # (``trace_lower_exportable``), so every caller inherits it. + traced = model.forward_common_lower_exportable(*sample_inputs_cpu) + + # Output key order is read from the static FX output node. A CUDA-target + # trace may already contain CUDA-only custom operators and therefore must + # not be replayed on the CPU tracing inputs. + output_keys = traced_output_keys(traced) + + log.info("Exporting the traced graph (torch.export)...") + exported = torch.export.export( + traced, + sample_inputs_cpu, + dynamic_shapes=_build_dynamic_shapes(sample_inputs_cpu), + strict=False, + prefer_deferred_runtime_asserts_over_guards=True, + ) _strip_shape_assertions(exported.graph_module) # move_to_device_pass handles FakeTensor device propagation cleanly; diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index f0229c1b66..268b802bac 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -11,6 +11,7 @@ ) import math +import threading from contextlib import ( contextmanager, ) @@ -68,19 +69,24 @@ # .pt2 built through the operator fails to load. The in-process inference # path lowers in the same process and resolves the operator through the # dispatcher, so only the freeze/export seam raises this flag. -_assemble_op_suppressed_for_export = False +# Thread-local: a process-global boolean corrupts interleaved contexts +# (thread A's finally restores False while thread B still traces). +_assemble_op_suppression = threading.local() + + +def _assemble_op_suppressed_for_export() -> bool: + return getattr(_assemble_op_suppression, "value", False) @contextmanager def assemble_op_suppressed_for_export() -> Iterator[None]: """Route export traces through the eager scatter instead of the operator.""" - global _assemble_op_suppressed_for_export - saved = _assemble_op_suppressed_for_export - _assemble_op_suppressed_for_export = True + saved = _assemble_op_suppressed_for_export() + _assemble_op_suppression.value = True try: yield finally: - _assemble_op_suppressed_for_export = saved + _assemble_op_suppression.value = saved def _wigner_d_assemble_impl( @@ -775,7 +781,11 @@ def forward( compiling = ( _COMPILER_IS_COMPILING() if _COMPILER_IS_COMPILING is not None else False ) - fake_mode = bool(blocks) and getattr(blocks[0], "fake_mode", None) is not None + # len(...) > 0, not bool(...): Dynamo cannot trace ``bool`` on a + # list and would graph-break here, splitting the hot path in two. + fake_mode = ( + len(blocks) > 0 and getattr(blocks[0], "fake_mode", None) is not None + ) tracing = compiling or fake_mode # One assembly implementation serves both routes: the eager call # goes straight to the differentiable scatter (autograd flows through @@ -795,7 +805,7 @@ def forward( elif ( tracing and _wigner_d_assemble_op is not None - and not _assemble_op_suppressed_for_export + and not _assemble_op_suppressed_for_export() ): D_full = _wigner_d_assemble_op( all_blocks, diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 91ff826a3e..c32b08958d 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2498,14 +2498,26 @@ def trace_lower_exportable( get_decompositions, ) - traced = make_fx( - fn, - tracing_mode="symbolic", - _allow_non_fake_inputs=True, - decomposition_table=get_decompositions( - [torch.ops.aten.silu_backward.default] - ), - )(*sample_inputs) + # The export seam keeps the eager Wigner assembly: AOTInductor + # resolves custom ops by name in the C++ runtime, where nothing + # registers ``sezm_nn::wigner_d_assemble``, so an export traced + # through the operator cannot load. Wiring the suppression here + # covers every caller of the documented export entries (the freeze + # flow and any direct ``torch.export`` consumer); the in-process + # runtime trace (``trace_and_compile``) is untouched. + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + assemble_op_suppressed_for_export, + ) + + with assemble_op_suppressed_for_export(): + traced = make_fx( + fn, + tracing_mode="symbolic", + _allow_non_fake_inputs=True, + decomposition_table=get_decompositions( + [torch.ops.aten.silu_backward.default] + ), + )(*sample_inputs) # make_fx can lower a reshape to an unsound aten.view when the fake # stride diverges from the eager stride; relax views back to reshapes. relax_views_to_reshapes(traced) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 94be7ea476..ebbf794dac 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2364,6 +2364,16 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): mode = sys.argv[1] tree = pathlib.Path(sys.argv[2]) expected_registered = sys.argv[3] == "1" + # Warm the package while the editable finder is still alive: a + # scikit-build editable install serves ``deepmd.lib`` (the C + # extension) through that finder and the source tree has no + # ``deepmd/lib`` -- importing first keeps the module resident + # while the tree takes over the pure-Python resolution below. + import torch + import torch.distributed._functional_collectives # noqa: F401 + + import deepmd # noqa: F401 (warms deepmd.lib, see below) + sys.meta_path = [ f for f in sys.meta_path @@ -2372,11 +2382,17 @@ def test_import_and_tracing_fall_back_when_registration_fails(self): and "editable" not in getattr(f, "__name__", "").lower() ] sys.path.insert(0, str(tree)) + # The warm import cached the installed tree of packages; drop + # every deepmd module except the C extension (location + # independent once loaded) so the tree path takes over the + # pure-Python resolution below. + for _name in [ + k for k in list(sys.modules) if k.startswith("deepmd") + ]: + if _name != "deepmd.lib": + del sys.modules[_name] import importlib - import torch - import torch.distributed._functional_collectives # noqa: F401 - import deepmd.pt.model.descriptor.sezm_nn.wignerd as wd if expected_registered: diff --git a/source/tests/pt/model/test_sezm_export.py b/source/tests/pt/model/test_sezm_export.py index 04a920d83f..315ccf6faf 100644 --- a/source/tests/pt/model/test_sezm_export.py +++ b/source/tests/pt/model/test_sezm_export.py @@ -405,6 +405,21 @@ def _assert_dict_allclose( err_msg=f"{context}: {key}", ) + def test_traced_graph_keeps_the_eager_wigner_assembly(self) -> None: + """The export seam must not capture the assemble operator. + + AOTInductor resolves custom ops by name in the C++ runtime, where + nothing registers ``sezm_nn::wigner_d_assemble``; the seam inside + ``trace_lower_exportable`` keeps the eager scatter, so every export + caller (freeze and direct ``torch.export``) inherits a loadable + graph. + """ + self.assertNotIn( + "wigner_d_assemble", + self.traced.code, + "the export trace must keep the eager Wigner assembly", + ) + def test_traced_matches_eager_same_shape(self) -> None: eager_out = _eager_forward(self.model, self.sample_inputs) traced_out = self.traced(*self.sample_inputs) @@ -551,6 +566,14 @@ def _assert_dict_allclose( err_msg=f"{context}: {key}", ) + def test_traced_graph_keeps_the_eager_wigner_assembly(self) -> None: + """The with-comm export seam keeps the eager assembly too.""" + self.assertNotIn( + "wigner_d_assemble", + self.traced.code, + "the with-comm export trace must keep the eager Wigner assembly", + ) + def test_traced_matches_eager(self) -> None: eager = _eager_parallel_forward(self.model, self.sample_inputs) traced_out = self.traced(*self.sample_inputs)