Fix Cuda Graph based MOE Tests Hang in CI

tests/pytorch/test_grouped_linear.py

@@ -1857,6 +1857,15 @@ def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch
         params_dtype=dtype,
         device=device,
     )
+    reference_grouped_linear = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=bias,
+        params_dtype=dtype,
+        device=device,
+    )
+    reference_grouped_linear.load_state_dict(grouped_linear.state_dict())
 
     static_x = torch.randn(total_tokens, in_features, dtype=dtype, device=device)
     static_x.requires_grad_(True)
@@ -1919,15 +1928,15 @@ def _train_step(x, dy, out_buf, *, use_graphed):
     expected_x = fresh_x.detach().clone().requires_grad_(True)
     expected_dy = fresh_dy.detach().clone()
     with autocast(enabled=use_fp8, recipe=fp8_recipe):
-        expected_out = grouped_linear(expected_x, static_m_splits)
+        expected_out = reference_grouped_linear(expected_x, static_m_splits)
     expected_out.backward(expected_dy)
 
     tols = dict(rtol=1e-2, atol=5e-3)
     if use_fp8:
         tols = dict(rtol=0.05, atol=0.05)
     torch.testing.assert_close(graph_out.float(), expected_out.float(), **tols)
     torch.testing.assert_close(graph_dx.float(), expected_x.grad.float(), **tols)
-    for graph_grad, param in zip(graph_param_grads, grouped_linear.parameters()):
+    for graph_grad, param in zip(graph_param_grads, reference_grouped_linear.parameters()):
         assert param.grad is not None
         torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols)
 

tests/pytorch/test_grouped_mlp.py

@@ -515,7 +515,8 @@ def test_grouped_linear_cuda_graph_safe(
         split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
 
         # Pad input tokens to validate the sync-free flow
-        in_shape = (split_sizes.sum().item() + token_padding, in_features)
+        num_active_tokens = split_sizes.sum().item()
+        in_shape = (num_active_tokens + token_padding, in_features)
         out_shape = (in_shape[0], out_features)
 
         recipe = make_recipe(quantization)
@@ -531,33 +532,45 @@ def test_grouped_linear_cuda_graph_safe(
                 single_grouped_weight=single_grouped_weight,
                 single_grouped_bias=single_grouped_bias,
             )
+            reference_op = te.ops.GroupedLinear(
+                group_size,
+                in_features,
+                out_features,
+                bias=bias,
+                device=device,
+                dtype=dtype,
+                accumulate_into_main_grad=accumulate_into_main_grad,
+                single_grouped_weight=single_grouped_weight,
+                single_grouped_bias=single_grouped_bias,
+            )
+        reference_op.load_state_dict(op.state_dict())
 
-        def _weight_params() -> list[torch.nn.Parameter]:
+        def _weight_params(module: torch.nn.Module) -> list[torch.nn.Parameter]:
             if single_grouped_weight:
-                return [op.weight]
-            return [getattr(op, f"weight{i}") for i in range(group_size)]
+                return [module.weight]
+            return [getattr(module, f"weight{i}") for i in range(group_size)]
 
-        def _bias_params() -> list[torch.nn.Parameter]:
+        def _bias_params(module: torch.nn.Module) -> list[torch.nn.Parameter]:
             if not bias:
                 return []
             if single_grouped_bias:
-                return [op.bias]
-            return [getattr(op, f"bias{i}") for i in range(group_size)]
+                return [module.bias]
+            return [getattr(module, f"bias{i}") for i in range(group_size)]
 
-        def _init_main_grads(value: float = 0.0) -> None:
+        def _init_main_grads(module: torch.nn.Module, value: float = 0.0) -> None:
             if not accumulate_into_main_grad:
                 return
             with torch.no_grad():
-                for w in _weight_params():
+                for w in _weight_params(module):
                     if getattr(w, "main_grad", None) is None:
                         w.main_grad = torch.empty(w.size(), device=device, dtype=torch.float32)
                     w.main_grad.fill_(value)
 
-        def _collect_main_grads() -> list[torch.Tensor]:
-            return [w.main_grad.detach().clone() for w in _weight_params()]
+        def _collect_main_grads(module: torch.nn.Module) -> list[torch.Tensor]:
+            return [w.main_grad.detach().clone() for w in _weight_params(module)]
 
-        def _zero_param_grads() -> None:
-            for param in op.parameters():
+        def _zero_param_grads(module: torch.nn.Module) -> None:
+            for param in module.parameters():
                 if param.grad is None:
                     param.grad = torch.zeros_like(param)
                 else:
@@ -582,7 +595,7 @@ def train_step(
             out_buf.copy_(out)
             return out_buf
 
-        _init_main_grads(0.0)
+        _init_main_grads(op, 0.0)
 
         static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True)
         static_dy = torch.randn(out_shape, device=device, dtype=dtype)
@@ -605,8 +618,8 @@ def train_step(
             static_dy.copy_(fresh_dy)
 
         # Reset grads & main_grads so the captured iteration starts fresh.
-        _zero_param_grads()
-        _init_main_grads(0.5)
+        _zero_param_grads(op)
+        _init_main_grads(op, 0.5)
         if static_x.grad is not None:
             static_x.grad.zero_()
 
@@ -617,35 +630,36 @@ def train_step(
         torch.cuda.synchronize()
         graph_dx = static_x.grad.detach().clone()
         if accumulate_into_main_grad:
-            graph_main_grads = _collect_main_grads()
+            graph_main_grads = _collect_main_grads(op)
             graph_param_grads: list[torch.Tensor] = []
         else:
             graph_main_grads = []
             graph_param_grads = [param.grad.detach().clone() for param in op.parameters()]
 
-        # Reference: same op invoked eagerly with the same fresh inputs and
-        # the same starting grad/main_grad state.
-        _zero_param_grads()
-        _init_main_grads(0.5)
-        static_x.grad.zero_()
+        # Reference: an independent op invoked eagerly with the same fresh
+        # inputs and starting grad/main_grad state.
+        _zero_param_grads(reference_op)
+        _init_main_grads(reference_op, 0.5)
 
         expected_x = fresh_x.detach().clone().requires_grad_(True)
         expected_dy = fresh_dy.detach().clone()
         with te.autocast(enabled=quantization is not None, recipe=recipe):
-            expected_out = op(expected_x, static_split_sizes)
+            expected_out = reference_op(expected_x, static_split_sizes)
         expected_out.backward(expected_dy)
 
         tols = dtype_tols(dtype)
         if quantization is not None:
             tols = quantization_tols(quantization)
 
-        assert_close(graph_out, expected_out, **tols)
-        assert_close(graph_dx, expected_x.grad, **tols)
+        # Grouped GEMM only defines rows covered by split_sizes. The padded tail
+        # is intentionally outside every group and remains uninitialized.
+        assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols)
+        assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols)
         if accumulate_into_main_grad:
-            for g, w in zip(graph_main_grads, _weight_params()):
+            for g, w in zip(graph_main_grads, _weight_params(reference_op)):
                 assert_close(g, w.main_grad, **tols)
         else:
-            for g, param in zip(graph_param_grads, op.parameters()):
+            for g, param in zip(graph_param_grads, reference_op.parameters()):
                 assert_close(g, param.grad, **tols)
 
 
@@ -1491,7 +1505,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8(
         random.shuffle(split_sizes)
         split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device)
         # Pad the input tokens to validate the sync-free MOE
-        in_shape = (split_sizes.sum().item() + token_padding, hidden_size)
+        num_active_tokens = split_sizes.sum().item()
+        in_shape = (num_active_tokens + token_padding, hidden_size)
         recipe = make_recipe("mxfp8")
         with te.quantized_model_init(enabled=True, recipe=recipe):
             fc1 = te.ops.GroupedLinear(
@@ -1524,8 +1539,43 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8(
                 scaled_act,
                 fc2,
             )
+            reference_fc1 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                2 * hidden_size,
+                bias=False,
+                device=device,
+                dtype=dtype,
+                single_grouped_weight=single_grouped_weight,
+                accumulate_into_main_grad=accumulate_into_main_grad,
+            )
+            reference_fc2 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                hidden_size,
+                bias=False,
+                device=device,
+                dtype=dtype,
+                single_grouped_weight=single_grouped_weight,
+                accumulate_into_main_grad=accumulate_into_main_grad,
+            )
+            reference_scaled_act = (
+                te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size)
+                if activation == "scaled_swiglu"
+                else te.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size)
+            )
+            reference_module = te.ops.Sequential(
+                reference_fc1,
+                reference_scaled_act,
+                reference_fc2,
+            )
+        reference_module.load_state_dict(module.state_dict())
 
-        def _init_main_grads(value: float = 0.0) -> None:
+        def _init_main_grads(
+            fc1: torch.nn.Module,
+            fc2: torch.nn.Module,
+            value: float = 0.0,
+        ) -> None:
             if not accumulate_into_main_grad:
                 return
             with torch.no_grad():
@@ -1563,7 +1613,10 @@ def _init_main_grads(value: float = 0.0) -> None:
                         fc1_weight.main_grad.fill_(value)
                         fc2_weight.main_grad.fill_(value)
 
-        def _collect_main_grads() -> tuple[torch.Tensor, torch.Tensor]:
+        def _collect_main_grads(
+            fc1: torch.nn.Module,
+            fc2: torch.nn.Module,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
             if single_grouped_weight:
                 fc1_main_grad = fc1.weight.main_grad.detach().clone()
                 fc2_main_grad = fc2.weight.main_grad.detach().clone()
@@ -1604,7 +1657,7 @@ def train_step(
             out_buf.copy_(out)
             return out_buf
 
-        _init_main_grads(0.0)
+        _init_main_grads(fc1, fc2, 0.0)
 
         static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True)
         static_probs = torch.randn((in_shape[0],), device=device, dtype=dtype, requires_grad=True)
@@ -1643,7 +1696,7 @@ def train_step(
 
         for param in module.parameters():
             param.grad = torch.zeros_like(param)
-        _init_main_grads(0.5)
+        _init_main_grads(fc1, fc2, 0.5)
         if static_x.grad is not None:
             static_x.grad.zero_()
         if static_probs.grad is not None:
@@ -1658,21 +1711,19 @@ def train_step(
         graph_dx = static_x.grad.detach().clone()
         graph_dprobs = static_probs.grad.detach().clone()
         if accumulate_into_main_grad:
-            graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads()
+            graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads(fc1, fc2)
         else:
             graph_param_grads = [param.grad.detach().clone() for param in module.parameters()]
 
-        for param in module.parameters():
-            param.grad.zero_()
-        _init_main_grads(0.5)
-        static_x.grad.zero_()
-        static_probs.grad.zero_()
+        for param in reference_module.parameters():
+            param.grad = torch.zeros_like(param)
+        _init_main_grads(reference_fc1, reference_fc2, 0.5)
 
         expected_x = fresh_x.detach().clone().requires_grad_(True)
         expected_probs = fresh_probs.detach().clone().requires_grad_(True)
         expected_dy = fresh_dy.detach().clone()
         with te.autocast(enabled=True, recipe=recipe):
-            expected_out = module(
+            expected_out = reference_module(
                 expected_x,
                 static_split_sizes,
                 expected_probs,
@@ -1681,15 +1732,24 @@ def train_step(
         expected_out.backward(expected_dy)
 
         tols = dtype_tols(dtype)
-        assert_close(graph_out, expected_out, **tols)
-        assert_close(graph_dx, expected_x.grad, **tols)
-        assert_close(graph_dprobs, expected_probs.grad, **tols)
+        # The padded tail is outside every expert and its outputs/gradients are
+        # intentionally left uninitialized.
+        assert_close(graph_out[:num_active_tokens], expected_out[:num_active_tokens], **tols)
+        assert_close(graph_dx[:num_active_tokens], expected_x.grad[:num_active_tokens], **tols)
+        assert_close(
+            graph_dprobs[:num_active_tokens],
+            expected_probs.grad[:num_active_tokens],
+            **tols,
+        )
         if accumulate_into_main_grad:
-            expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads()
+            expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads(
+                reference_fc1,
+                reference_fc2,
+            )
             assert_close(graph_fc1_main_grad, expected_fc1_main_grad, **tols)
             assert_close(graph_fc2_main_grad, expected_fc2_main_grad, **tols)
         else:
-            for graph_grad, param in zip(graph_param_grads, module.parameters()):
+            for graph_grad, param in zip(graph_param_grads, reference_module.parameters()):
                 assert_close(graph_grad, param.grad, **tols)