From 016009fe975342a9aeb57218a9b37b7adb9949e2 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Tue, 14 Jul 2026 21:40:06 +0000 Subject: [PATCH 1/2] fix grouped linear hang Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_linear.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index b561506b31..5dff390c9a 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/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,7 +1928,7 @@ 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) @@ -1927,7 +1936,7 @@ def _train_step(x, dy, out_buf, *, use_graphed): 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) From 1baace410ddcc56cb80f03043274aaa46157ee28 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Wed, 15 Jul 2026 00:39:22 +0000 Subject: [PATCH 2/2] make the same change in grouped mlp as well Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_mlp.py | 148 +++++++++++++++++++++--------- 1 file changed, 104 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e24fff9049..483f477f2f 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/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)