[Pytorch] MOE Sequential Block

tests/pytorch/distributed/run_ep.py

@@ -11,6 +11,7 @@
 import torch
 import torch.distributed as dist
 
+from transformer_engine.pytorch import ops as te_ops
 from transformer_engine.common.recipe import MXFP8BlockScaling
 from transformer_engine.pytorch.ep import (
     EpBuffer,
@@ -25,6 +26,7 @@
     _ep_combine_raw,
     _ep_dispatch_raw,
 )
+from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp
 
 ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1"
 EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1"
@@ -39,6 +41,7 @@
 # NVTE_EP_TOKENS_PER_RANK.
 HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512"))
 TOP_K = 2
+INTERMEDIATE_DIM = 16
 TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32"))
 
 
@@ -136,6 +139,31 @@ def _make_identity_inputs(rank, ep_size, device="cuda"):
     )
 
 
+def _make_moe_inputs(rank, ep_size, device="cuda"):
+    """Deterministic BF16 activations and FP32 top-k router weights."""
+    generator = torch.Generator(device=device)
+    generator.manual_seed(2026 + rank)
+    tokens = (
+        torch.randn(
+            TOKENS_PER_RANK,
+            HIDDEN_DIM,
+            generator=generator,
+            dtype=torch.float32,
+            device=device,
+        )
+        * 0.25
+    ).to(torch.bfloat16)
+    router_logits = torch.randn(
+        TOKENS_PER_RANK,
+        ep_size * NUM_LOCAL_EXPERTS,
+        generator=generator,
+        dtype=torch.float32,
+        device=device,
+    )
+    topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1)
+    return topk_idx, tokens, torch.softmax(topk_logits, dim=-1)
+
+
 def _degroup_mxfp8(recv_grouped, valid_counts=None):
     """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order.
     With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert
@@ -277,6 +305,42 @@ def _moe_step(self, buffer, topk_idx, tokens, w):
         expert_out = self._weighted(recv_t, recv_w_out)
         return ep_combine(buffer, expert_out)
 
+    def _make_moe_model(self, *, fuse_ops=True):
+        """Build a BF16 EP MoE Sequential.
+
+        With ``fuse_ops=True``, dispatch routing extras stay internal
+        (``output_to_caller=False``) so :class:`FusedMoeEp` can claim the
+        sequence. With ``fuse_ops=False``, those extras are returned to the
+        caller, which blocks fusion.
+        """
+        buffer = self._make_buffer()
+        dispatch = te_ops.Dispatch(buffer)
+        fc1 = te_ops.GroupedLinear(
+            NUM_LOCAL_EXPERTS,
+            HIDDEN_DIM,
+            2 * INTERMEDIATE_DIM,
+            bias=False,
+            device=self.cfg.device,
+            dtype=torch.bfloat16,
+        )
+        activation = te_ops.ScaledSwiGLU()
+        fc2 = te_ops.GroupedLinear(
+            NUM_LOCAL_EXPERTS,
+            INTERMEDIATE_DIM,
+            HIDDEN_DIM,
+            bias=False,
+            device=self.cfg.device,
+            dtype=torch.bfloat16,
+        )
+        combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK)
+
+        dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops)
+        dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops)
+        fc1.set_extra_input_channel(0, "tokens_per_expert")
+        activation.set_extra_input_channel(0, "routing_weights")
+        fc2.set_extra_input_channel(0, "tokens_per_expert")
+        return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2
+
     # Prepare
 
     @_eager_test_include
@@ -528,6 +592,100 @@ def test_caller_provides_grad_expert_out(self):
         # the caller-owned buffer was used as the combine-bwd scatter target
         self.assertGreater(gbuf.abs().sum().item(), 0.0)
 
+    @_eager_test_include
+    def test_bf16_moe_sequential_fusion(self):
+        """Reference-backed fusion matches the unfused BF16 EP MoE sequence.
+
+        ``fuse_ops=True`` keeps dispatch routing extras internal so fusion
+        fires; ``fuse_ops=False`` returns them to the caller and blocks it.
+        """
+        if not EAGER:
+            self.skipTest("variable-size reference comparison requires eager EP mode")
+
+        fused, fused_fc1, fused_fc2 = self._make_moe_model(fuse_ops=True)
+        unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fuse_ops=False)
+        generator = torch.Generator(device=self.cfg.device)
+        generator.manual_seed(3100 + self.cfg.rank)
+        with torch.no_grad():
+            for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)):
+                for expert in range(NUM_LOCAL_EXPERTS):
+                    weight = (
+                        torch.randn(
+                            getattr(fused_op, f"weight{expert}").shape,
+                            generator=generator,
+                            dtype=torch.float32,
+                            device=self.cfg.device,
+                        )
+                        * 0.1
+                    ).to(torch.bfloat16)
+                    getattr(fused_op, f"weight{expert}").copy_(weight)
+                    getattr(unfused_op, f"weight{expert}").copy_(weight)
+
+        topk_idx, tokens, topk_weights = _make_moe_inputs(
+            self.cfg.rank,
+            self.cfg.ep_size,
+            self.cfg.device,
+        )
+        fused_tokens = tokens.detach().clone().requires_grad_(True)
+        unfused_tokens = tokens.detach().clone().requires_grad_(True)
+        fused_topk_weights = topk_weights.detach().clone().requires_grad_(True)
+        unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True)
+
+        fused_out = fused(
+            fused_tokens,
+            topk_idx,
+            fused_topk_weights,
+        )
+        unfused_out, tokens_per_expert, recv_topk_weights = unfused(
+            unfused_tokens,
+            topk_idx,
+            unfused_topk_weights,
+        )
+
+        fused_forward_ops = fused._module_groups[0]._forward_ops
+        unfused_forward_ops = unfused._module_groups[0]._forward_ops
+        self.assertEqual(len(fused_forward_ops), 1)
+        self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp)
+        self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops))
+        self.assertIsInstance(fused_out, torch.Tensor)
+        self.assertEqual(fused_out.dtype, torch.bfloat16)
+        self.assertEqual(unfused_out.dtype, torch.bfloat16)
+        self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,))
+        self.assertEqual(tokens_per_expert.dtype, torch.int64)
+        self.assertEqual(recv_topk_weights.dtype, torch.float32)
+
+        dy = (
+            torch.randn(
+                fused_out.shape,
+                generator=generator,
+                dtype=torch.float32,
+                device=self.cfg.device,
+            )
+            * 0.1
+        ).to(torch.bfloat16)
+        fused_out.backward(dy)
+        unfused_out.backward(dy)
+        torch.cuda.synchronize()
+
+        # The two BF16 paths use different grouped-GEMM and reduction orders.
+        # Observed forward abs error reaches ~3e-5 on near-zero values where
+        # rtol does not apply, so atol must stay above that (1e-5 is too tight).
+        tolerances = {"rtol": 1.6e-2, "atol": 5e-5}
+        torch.testing.assert_close(fused_out, unfused_out, **tolerances)
+        torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances)
+        torch.testing.assert_close(
+            fused_topk_weights.grad,
+            unfused_topk_weights.grad,
+            **tolerances,
+        )
+        for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)):
+            for expert in range(NUM_LOCAL_EXPERTS):
+                fused_grad = getattr(fused_op, f"weight{expert}").grad
+                unfused_grad = getattr(unfused_op, f"weight{expert}").grad
+                self.assertEqual(fused_grad.dtype, torch.bfloat16)
+                self.assertEqual(unfused_grad.dtype, torch.bfloat16)
+                torch.testing.assert_close(fused_grad, unfused_grad, **tolerances)
+
     @_zero_copy_test_include
     @_mxfp8_align_test
     def test_combine_bwd_mxfp8_caller_grad_out(self):

transformer_engine/pytorch/ep.py

@@ -26,6 +26,7 @@
 __all__ = [
     "EpBuffer",
     "ep_bootstrap",
+    "get_ep_group",
     "is_ep_bootstrapped",
     "ep_finalize",
     "ep_dispatch",
@@ -183,6 +184,11 @@ def is_ep_bootstrapped() -> bool:
     return _BOOTSTRAPPED
 
 
+def get_ep_group() -> Optional[dist.ProcessGroup]:
+    """Return the process group registered by :func:`ep_bootstrap`."""
+    return _EP_GROUP
+
+
 def ep_finalize() -> None:
     """Optional explicit EP teardown; idempotent.