diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8f0117c68ec..46554765ee2 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -2103,6 +2103,10 @@ def sharded_state_dict( if HAVE_TE and is_te_min_version("1.9.0.dev0"): + _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = ( + "use_grouped_tensor" in inspect.signature(te.pytorch.GroupedLinear.__init__).parameters + ) + class TEGroupedLinear(te.pytorch.GroupedLinear): """ Wrapper for the Transformer-Engine's `GroupedLinear` layer. @@ -2205,6 +2209,14 @@ def __init__( config, "moe_single_grouped_bias", False ) + if _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR: + extra_kwargs["use_grouped_tensor"] = config.moe_use_grouped_tensor + elif config.moe_use_grouped_tensor and not config.use_transformer_engine_op_fuser: + raise RuntimeError( + "moe_use_grouped_tensor=True requires a Transformer Engine GroupedLinear " + "that exposes the use_grouped_tensor argument." + ) + self.te_quant_params: Optional[TEQuantizationParams] = None quant_config = get_quant_config_or_none(name, config.quant_recipe) self.finish_init(quant_config) @@ -2712,6 +2724,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): ) else: + _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = False TEGroupedLinear = None # type: ignore[assignment, misc] TEColumnParallelGroupedLinear = None # type: ignore[assignment, misc] TERowParallelGroupedLinear = None # type: ignore[assignment, misc] diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 59f8deffeca..deea7c223a9 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -301,25 +301,74 @@ def __init__( self.config.moe_mlp_glu_interleave_size, ) - if self.config.fp8 or self.config.fp4: - assert HAVE_TE, "FP8 and FP4 requires TE." - align_size = 256 if self._with_fused_impl else None + self._use_grouped_tensor = self.config.moe_use_grouped_tensor + if self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: + assert HAVE_TE, "Quantized or TE grouped-tensor GroupedMLP execution requires TE." + align_size = ( + get_align_size_for_quantization(self.config) if self._use_grouped_tensor else None + ) self.quantization_padding = Fp8Padding(self.num_local_experts, align_size=align_size) self.quantization_unpadding = Fp8Unpadding( self.num_local_experts, align_size=align_size ) + @staticmethod + def _apply_packed_bias(intermediate_parallel, packed_bias, tokens_per_expert, permuted_probs): + """Apply a packed expert bias without reading token counts on the host.""" + # TODO: get rid of the .float() by having fused kernel compute in FP32 + shape = intermediate_parallel.shape + hidden_size = shape[-1] + output_dtype = intermediate_parallel.dtype + flat_output = intermediate_parallel.view(-1, hidden_size).float() + flat_probs = permuted_probs.reshape(-1, 1).float() + + if tokens_per_expert.device != packed_bias.device: + raise ValueError("Packed MoE bias and tokens_per_expert must be on the same device.") + + # Permutation stores tokens contiguously by expert. Repeat bias row e by that expert's + # token count to create one bias row per permuted token: + # + # packed_bias = [bias_e0, bias_e1] + # tokens_per_expert = [ 2, 1] + # bias_per_token = [bias_e0, bias_e0, bias_e1] + # + # output_size avoids a stream synchronization to compute sum(tokens_per_expert). + # Cast before repeating so both forward arithmetic and repeat_interleave's backward + # reduction are computed in FP32. Autograd casts the final parameter gradient once. + bias_per_token = torch.repeat_interleave( + packed_bias.float(), tokens_per_expert, dim=0, output_size=flat_output.size(0) + ) + return (flat_output + bias_per_token * flat_probs).view(shape).to(output_dtype) + @staticmethod def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs): if bias_parallel is None: return intermediate_parallel + + # CUDA-graph-safe packed path. With single_grouped_bias=True, TE returns one packed + # GroupedTensor [num_experts, hidden_size]. The grouped-tensor backend also provides + # tokens_per_expert as a tensor on the same device. + if isinstance(bias_parallel, torch.Tensor) and isinstance(tokens_per_expert, torch.Tensor): + return TEGroupedMLP._apply_packed_bias( + intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs + ) + + # Eager-only CPU-metadata path. The legacy contract returns List[Tensor[hidden_size]], + # and torch.split plus the Python zip below require concrete host token counts. A packed + # bias paired with Python counts also uses this compatibility path. Converting a tensor + # with .tolist() synchronizes and copies device data to the host, so this path must never + # be included in a CUDA graph. + if isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = tokens_per_expert.tolist() + shape = intermediate_parallel.shape + flat_output = intermediate_parallel.view(-1, shape[-1]) return ( torch.cat( [ t + b * p for t, b, p in zip( - torch.split(intermediate_parallel.view(-1, shape[-1]), tokens_per_expert), + torch.split(flat_output, tokens_per_expert), bias_parallel, torch.split(permuted_probs, tokens_per_expert), ) @@ -634,9 +683,21 @@ def _fused_forward( # Apply padding if needed unpadded_tokens_per_expert = None + # Some dispatchers have already padded each expert's token segment before the tokens + # reach this module: + # * router padding changes the routing map before dispatch; + # * HybridEP pads as part of its fused dispatch/permute operation; + # * DeepEP can pad in the fused local permutation after communication. + # Padding those tensors again would insert a second set of dummy tokens and make + # tokens_per_expert disagree with the already-permuted token buffer, so skip the local + # Fp8Padding fallback in those cases. if skip_routed_expert_padding(self.config): pass - elif self.config.fp8 or self.config.fp4: + # Regular AllToAll normally reaches this branch because its permutation does not insert + # the padding needed by the fused grouped-MLP contract. FP8/FP4 require recipe-specific + # alignment, while the TE operation-fuser grouped-tensor path currently uses 256-token + # expert segments. + elif self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: tokens_per_expert = tokens_per_expert.tolist() unpadded_tokens_per_expert = tokens_per_expert permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( @@ -647,7 +708,22 @@ def _fused_forward( ) permuted_probs = permuted_probs.squeeze(-1) tokens_per_expert = torch.tensor( - tokens_per_expert, dtype=torch.int, device=permuted_probs.device + tokens_per_expert, dtype=torch.int64, device=permuted_probs.device + ) + + if self._use_grouped_tensor: + if not isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = torch.tensor( + tokens_per_expert, dtype=torch.int64, device=permuted_local_hidden_states.device + ) + else: + tokens_per_expert = tokens_per_expert.to( + device=permuted_local_hidden_states.device, dtype=torch.int64, non_blocking=True + ) + else: + raise RuntimeError( + "The Transformer Engine operation-fuser MoE path requires " + "moe_use_grouped_tensor=True." ) # if the number of tokens is 0, pad the hidden states to 256 @@ -766,11 +842,23 @@ def forward( # Apply padding if needed unpadded_tokens_per_expert = None - tokens_per_expert: list[int] = tokens_per_expert.tolist() permuted_probs = permuted_probs.unsqueeze(-1) + # The token buffer may already contain per-expert padding when padding was performed + # before expert compute: + # * router padding modified the routing map before dispatch; + # * HybridEP fused padding into dispatch/permute; + # * DeepEP fused padding into its post-communication local permutation. + # In those cases tokens_per_expert already describes the padded expert segments. Running + # Fp8Padding again would change the segment lengths without matching the existing token + # layout, so this module must leave both tensors unchanged. if skip_routed_expert_padding(self.config): pass - elif self.config.fp8 or self.config.fp4: + # Regular AllToAll normally supplies unpadded expert segments and therefore uses this + # explicit fallback. FP8/FP4 need their recipe-specific alignment. MCore currently also + # applies its common aligned-segment contract to the GroupedTensor backend so quantized + # grouped execution receives supported shapes + elif self.config.fp8 or self.config.fp4 or self._use_grouped_tensor: + tokens_per_expert = tokens_per_expert.tolist() unpadded_tokens_per_expert = tokens_per_expert permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( permuted_local_hidden_states, tokens_per_expert @@ -779,6 +867,18 @@ def forward( permuted_probs, unpadded_tokens_per_expert ) + if self._use_grouped_tensor: + if not isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = torch.tensor( + tokens_per_expert, dtype=torch.int64, device=permuted_local_hidden_states.device + ) + else: + tokens_per_expert = tokens_per_expert.to( + device=permuted_local_hidden_states.device, dtype=torch.int64, non_blocking=True + ) + elif isinstance(tokens_per_expert, torch.Tensor): + tokens_per_expert = tokens_per_expert.tolist() + if self.config.moe_apply_probs_on_input: assert ( self.config.moe_router_topk == 1 diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index dfdb9a14460..aa56cdf3016 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1378,8 +1378,11 @@ def forward( inp = inp.view(-1, inp_shape[-1]) if te_general_gemm is not None and router_dtype != torch.float64: - output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) - output = output[0] + # cuBLASLt's non-FP8 bias epilogue expects bias and output to have the same + # dtype. Router parameters may be BF16 while router logits are FP32, so cast the + # small bias vector before passing it to TE. + gemm_bias = bias.to(router_dtype) if bias is not None else None + output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=gemm_bias)[0] elif bias is None: output = torch.mm(inp.to(router_dtype), weight.to(router_dtype).t()) else: @@ -1457,22 +1460,33 @@ def get_align_size_for_quantization(config: TransformerConfig) -> int: Returns: int: The alignment size for quantization. """ - # CUTLASS kernel for grouped GEMM assumes 256 alignment. - if config.use_transformer_engine_op_fuser: + # TE's grouped-tensor and fused grouped-MLP kernels require 256-token alignment. + if config.use_transformer_engine_op_fuser or config.moe_use_grouped_tensor: return 256 if config.fp8: return get_fp8_align_size(config.fp8_recipe) if config.fp4: return get_fp4_align_size(config.fp4_recipe) - # Only FP8 or FP4 requires padding. Defaults to 0. + # Legacy high-precision grouped GEMM does not require padding. Defaults to 0. return 0 +def _deepep_permute_pads_grouped_tensor_input(config: TransformerConfig) -> bool: + """Whether DeepEP fused permutation pads input for TE grouped-tensor GEMM.""" + return ( + config.moe_use_grouped_tensor + and config.moe_token_dispatcher_type == "flex" + and config.moe_flex_dispatcher_backend == "deepep" + and config.moe_permute_fusion + and fused_permute_and_pad_with_probs is not None + ) + + def skip_routed_expert_padding(config: TransformerConfig) -> bool: """Whether the expert module should skip quantization padding. - Returns True when padding is already applied by the router or the - HybridEP / NCCL-EP dispatcher. + Returns True when padding is already applied by the router, the HybridEP / NCCL-EP + dispatcher, or DeepEP's fused permutation kernel. """ if config.moe_router_padding_for_quantization: return True @@ -1481,6 +1495,8 @@ def skip_routed_expert_padding(config: TransformerConfig) -> bool: "ncclep", ): return True + if _deepep_permute_pads_grouped_tensor_input(config): + return True return False diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py index 256f38cc6fb..cf5aeca9f6b 100644 --- a/megatron/core/transformer/moe/token_dispatcher.py +++ b/megatron/core/transformer/moe/token_dispatcher.py @@ -1136,8 +1136,9 @@ def dispatch( "HybridEP only supports float32 probs, please set --moe-router-dtype=fp32" ) self.token_probs = self.token_probs.float() # downcast or upcast - if self.config.fp8 or self.config.fp4: - self.pad_multiple = get_align_size_for_quantization(self.config) + align_size = get_align_size_for_quantization(self.config) + if align_size > 0: + self.pad_multiple = align_size if self._padded_num_tokens is not None and hidden_states.shape[0] < self._padded_num_tokens: pad_rows = self._padded_num_tokens - hidden_states.shape[0] hidden_states = torch.cat( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index af853276819..de08caef71d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -827,21 +827,29 @@ class TransformerConfig(ModelParallelConfig): """ moe_grouped_gemm: bool = False - """When there are multiple experts per rank, compress multiple local (potentially small) gemms - in a single kernel launch to improve the utilization and performance by leveraging the Grouped - GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). + """Use grouped GEMM to execute multiple local MoE experts together. + + The concrete implementation is selected by Transformer Engine. Set + ``moe_use_grouped_tensor=True`` to use its CUDA-graph-safe GroupedTensor path. + """ + + moe_use_grouped_tensor: bool = False + """Use Transformer Engine's native GroupedTensor path for grouped MoE GEMMs. + + This path uses padded expert segments and CUDA split metadata so it can be captured in CUDA + graphs. Enabling the Transformer Engine operation fuser also enables this option. """ moe_single_grouped_weight: bool = False """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and - ``use_transformer_engine_op_fuser=True``. + ``moe_use_grouped_tensor=True``. """ moe_single_grouped_bias: bool = False """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped - parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` - and ``add_bias_linear=True``.""" + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``, + ``moe_use_grouped_tensor=True``, and ``add_bias_linear=True``.""" moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. @@ -1298,6 +1306,12 @@ def __post_init__(self): """ super().__post_init__() + if self.use_transformer_engine_op_fuser and self.moe_grouped_gemm: + self.moe_use_grouped_tensor = True + + if self.moe_use_grouped_tensor and not self.moe_grouped_gemm: + raise ValueError("moe_use_grouped_tensor=True requires moe_grouped_gemm=True.") + # When fp32 residual connections are enabled, pipeline parallel communication must # use fp32 to match the dtype of the residual stream between pipeline stages. if self.fp32_residual_connection and self.pipeline_dtype is not None: @@ -1602,15 +1616,12 @@ def __post_init__(self): "(--fp4-param-gather). Without FP4 parameter gather, Transformer Engine " "uses a split-quantize fallback that is being deprecated." ) - if not self.use_transformer_engine_op_fuser: - raise ValueError( - "moe_single_grouped_weight requires " - "use_transformer_engine_op_fuser=True. The non-op-fuser TE GroupedLinear " - "path splits the grouped parameter into per-expert tensors and does not " - "support single-grouped-weight training." - ) + if not self.moe_use_grouped_tensor: + raise ValueError("moe_single_grouped_weight requires moe_use_grouped_tensor=True.") if self.moe_single_grouped_bias and not self.add_bias_linear: raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.") + if self.moe_single_grouped_bias and not self.moe_use_grouped_tensor: + raise ValueError("moe_single_grouped_bias requires moe_use_grouped_tensor=True.") if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index eb871568046..868c58b003c 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -185,6 +185,7 @@ "moe_flex_dispatcher_num_sms": None, "moe_grad_scale_func": None, "moe_grouped_gemm": True, + "moe_use_grouped_tensor": False, "moe_hybridep_num_sms": None, "moe_hybridep_num_sms_preprocessing": 108, "moe_hybridep_num_blocks_permute": None, diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b9e7fa346d2..f22c92efa43 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -31,11 +31,41 @@ def test_op_fuser_transformer_config_args_are_exposed(): _add_network_size_args(parser) args = parser.parse_args( - ["--use-transformer-engine-op-fuser", "--moe-mlp-glu-interleave-size", "16"] + [ + "--use-transformer-engine-op-fuser", + "--moe-mlp-glu-interleave-size", + "16", + "--moe-use-grouped-tensor", + ] ) assert args.use_transformer_engine_op_fuser is True assert args.moe_mlp_glu_interleave_size == 16 + assert args.moe_use_grouped_tensor is True + + +def test_op_fuser_enables_grouped_tensor(): + config = TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + num_moe_experts=2, + moe_grouped_gemm=True, + use_transformer_engine_op_fuser=True, + ) + + assert config.moe_use_grouped_tensor is True + + +def test_grouped_tensor_requires_grouped_gemm(): + with pytest.raises(ValueError, match="requires moe_grouped_gemm=True"): + TransformerConfig( + num_layers=1, + hidden_size=128, + num_attention_heads=4, + num_moe_experts=2, + moe_use_grouped_tensor=True, + ) def test_remove_glu_interleaving_restores_contiguous_gate_and_linear_halves(): @@ -171,8 +201,12 @@ def __call__(self, *args): moe_router_padding_for_quantization=False, moe_token_dispatcher_type=None, moe_flex_dispatcher_backend=None, + moe_use_grouped_tensor=True, moe_paged_stash=False, ) + module._use_grouped_tensor = True + module.quantization_padding = lambda tensor, token_counts: (tensor, token_counts) + module.quantization_unpadding = lambda tensor, token_counts: tensor module._fused_ops = None module.linear_fc2 = SimpleNamespace(use_bias=fc2_bias) fused_ops = FakeFusedOps() @@ -186,11 +220,11 @@ def __call__(self, *args): torch.testing.assert_close(output, torch.ones_like(hidden_states)) assert module._fused_ops[0] is fused_ops assert fused_ops.args[0] is hidden_states - assert fused_ops.args[1] is tokens_per_expert - assert fused_ops.args[2] is probs - assert fused_ops.args[3] is tokens_per_expert + torch.testing.assert_close(fused_ops.args[1], tokens_per_expert) + torch.testing.assert_close(fused_ops.args[2], probs) + torch.testing.assert_close(fused_ops.args[3], tokens_per_expert) if fc2_bias: - assert fused_ops.args[4] is probs + torch.testing.assert_close(fused_ops.args[4], probs) else: assert len(fused_ops.args) == 4 @@ -220,6 +254,20 @@ def test_apply_bias_combines_per_expert_bias_and_probs(): assert output.dtype == intermediate.dtype +def test_apply_bias_combines_packed_grouped_bias_and_accumulates_gradient(): + intermediate = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + packed_bias = torch.tensor([[10.0, 20.0], [100.0, 200.0]], requires_grad=True) + tokens_per_expert = torch.tensor([2, 1], dtype=torch.int64) + permuted_probs = torch.tensor([0.25, 0.5, 1.5]) + expected = torch.tensor([[3.5, 7.0], [8.0, 14.0], [155.0, 306.0]]) + + output = TEGroupedMLP._apply_bias(intermediate, packed_bias, tokens_per_expert, permuted_probs) + output.sum().backward() + + torch.testing.assert_close(output, expected) + torch.testing.assert_close(packed_bias.grad, torch.tensor([[0.75, 0.75], [1.5, 1.5]])) + + def test_make_fused_impl_pre_forward_hook_dispatches_submodule_hooks(): module = TEGroupedMLP.__new__(TEGroupedMLP) torch.nn.Module.__init__(module) @@ -1081,7 +1129,8 @@ def test_gpu_make_fused_ops_constructs_with_real_te(self): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.internal - def test_gpu_fused_path_scales_fc2_bias(self): + @pytest.mark.parametrize("single_grouped_bias", (False, True)) + def test_gpu_fused_path_scales_fc2_bias(self, monkeypatch, single_grouped_bias): """FC2 bias and its gradients must use the per-token router probability.""" try: from transformer_engine.pytorch.ops import GroupedLinear @@ -1091,6 +1140,8 @@ def test_gpu_fused_path_scales_fc2_bias(self): if "scale_bias" not in inspect.signature(GroupedLinear.__init__).parameters: pytest.skip("Installed TE op fuser GroupedLinear lacks `scale_bias` support") + if single_grouped_bias: + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") Utils.destroy_model_parallel() Utils.initialize_model_parallel(1, 1) @@ -1112,6 +1163,7 @@ def test_gpu_fused_path_scales_fc2_bias(self): moe_router_topk=1, moe_grouped_gemm=True, use_transformer_engine_op_fuser=True, + moe_single_grouped_bias=single_grouped_bias, ) _set_random_seed(seed_=123, data_parallel_random_init=False) submodules = get_submodules( @@ -1129,9 +1181,19 @@ def test_gpu_fused_path_scales_fc2_bias(self): for linear in (experts.linear_fc1, experts.linear_fc2): for expert_idx in range(self.num_experts): getattr(linear, f"weight{expert_idx}").zero_() - getattr(linear, f"bias{expert_idx}").zero_() - experts.linear_fc2.bias0.fill_(2.0) - experts.linear_fc2.bias1.fill_(4.0) + if not single_grouped_bias: + getattr(linear, f"bias{expert_idx}").zero_() + if single_grouped_bias: + linear.bias.rowwise_data.zero_() + if single_grouped_bias: + packed_fc2_bias = experts.linear_fc2.bias.rowwise_data.view( + self.num_experts, self.hidden_size + ) + packed_fc2_bias[0].fill_(2.0) + packed_fc2_bias[1].fill_(4.0) + else: + experts.linear_fc2.bias0.fill_(2.0) + experts.linear_fc2.bias1.fill_(4.0) hidden_states = torch.zeros( 3, self.hidden_size, dtype=torch.bfloat16, device="cuda", requires_grad=True @@ -1155,15 +1217,201 @@ def test_gpu_fused_path_scales_fc2_bias(self): torch.tensor([2.0, 2.0, 4.0], dtype=torch.bfloat16, device="cuda") * self.hidden_size ) torch.testing.assert_close(probs.grad, expected_prob_grad) - torch.testing.assert_close( - experts.linear_fc2.bias0.grad, - torch.ones_like(experts.linear_fc2.bias0) * probs[:2].detach().sum(), + if single_grouped_bias: + assert experts.linear_fc2.bias.grad is not None + packed_dbias = experts.linear_fc2.bias.grad.view(self.num_experts, self.hidden_size) + torch.testing.assert_close( + packed_dbias[0], torch.ones_like(packed_dbias[0]) * probs[:2].detach().sum() + ) + torch.testing.assert_close( + packed_dbias[1], torch.ones_like(packed_dbias[1]) * probs[2:].detach().sum() + ) + assert experts._fused_ops[0][2].bias is experts.linear_fc2.bias + else: + torch.testing.assert_close( + experts.linear_fc2.bias0.grad, + torch.ones_like(experts.linear_fc2.bias0) * probs[:2].detach().sum(), + ) + torch.testing.assert_close( + experts.linear_fc2.bias1.grad, + torch.ones_like(experts.linear_fc2.bias1) * probs[2:].detach().sum(), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + @pytest.mark.parametrize("use_op_fuser", (False, True), ids=("module", "op-fuser")) + @pytest.mark.parametrize( + "single_grouped_weight,single_grouped_bias", + ((True, False), (False, True), (True, True)), + ids=("single-weight", "single-bias", "single-weight-and-bias"), + ) + def test_gpu_single_grouped_parent_gradient_parity( + self, monkeypatch, use_op_fuser, single_grouped_weight, single_grouped_bias + ): + """Single grouped parents must receive the same gradients as discrete parameters. + + This is an integration test over the real MCore TEGroupedMLP wrapper. It covers both + gradient ownership mechanisms: + + * the module path returns the packed FC2 bias to ``_apply_packed_bias``, where normal + PyTorch autograd must update the registered grouped parent; + * the op-fuser path reattaches the same wrapper parameters to TE op shells, whose custom + backward must return packed wgrad/dbias in the corresponding parent slots. + + Comparing gradients directly is intentional. A forward or dgrad-only check would not + catch a disconnected grouped parent that the optimizer can never update. + """ + try: + from transformer_engine.pytorch.module import GroupedLinear as ModuleGroupedLinear + from transformer_engine.pytorch.ops import GroupedLinear as OpGroupedLinear + except ImportError: + pytest.skip("Required TE GroupedLinear APIs are not available") + import inspect + + module_parameters = inspect.signature(ModuleGroupedLinear.__init__).parameters + op_parameters = inspect.signature(OpGroupedLinear.__init__).parameters + if ( + "use_grouped_tensor" not in module_parameters + or "single_grouped_bias" not in module_parameters + or "single_grouped_bias" not in op_parameters + ): + pytest.skip("Installed TE lacks native single grouped bias support") + + monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1") + + def build_experts(single_weight, single_bias): + config = TransformerConfig( + num_layers=1, + hidden_size=self.hidden_size, + num_attention_heads=4, + num_moe_experts=self.num_experts, + use_cpu_initialization=False, + add_bias_linear=True, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bias_dropout_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + moe_grouped_gemm=True, + moe_use_grouped_tensor=True, + use_transformer_engine_op_fuser=use_op_fuser, + moe_single_grouped_weight=single_weight, + moe_single_grouped_bias=single_bias, + ) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + self.num_experts, moe_grouped_gemm=True + ).mlp + ) + assert isinstance(submodules, MoESubmodules) + layer = MoELayer(config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + assert isinstance(layer.experts, TEGroupedMLP) + return layer.experts + + def copy_linear_params(reference_linear, target_linear): + reference_weights = torch.stack( + [ + getattr(reference_linear, f"weight{idx}").detach() + for idx in range(self.num_experts) + ] + ) + reference_biases = torch.stack( + [ + getattr(reference_linear, f"bias{idx}").detach() + for idx in range(self.num_experts) + ] + ) + with torch.no_grad(): + if target_linear.single_grouped_weight: + target_linear.weight.rowwise_data.view_as(reference_weights).copy_( + reference_weights + ) + else: + for idx in range(self.num_experts): + getattr(target_linear, f"weight{idx}").copy_(reference_weights[idx]) + + if target_linear.single_grouped_bias: + target_linear.bias.rowwise_data.view_as(reference_biases).copy_( + reference_biases + ) + else: + for idx in range(self.num_experts): + getattr(target_linear, f"bias{idx}").copy_(reference_biases[idx]) + + def packed_grad(linear, name): + if getattr(linear, f"single_grouped_{name}"): + grad = getattr(linear, name).grad + assert grad is not None, f"Grouped {name} parent did not receive a gradient" + return grad.float() + grads = [getattr(linear, f"{name}{idx}").grad for idx in range(self.num_experts)] + assert all(grad is not None for grad in grads) + return torch.stack(grads).float() + + torch.manual_seed(1234) + reference = build_experts(False, False) + torch.manual_seed(5678) + target = build_experts(single_grouped_weight, single_grouped_bias) + copy_linear_params(reference.linear_fc1, target.linear_fc1) + copy_linear_params(reference.linear_fc2, target.linear_fc2) + + tokens_per_expert = torch.tensor([256, 256], dtype=torch.int64, device="cuda") + num_tokens = int(tokens_per_expert.sum().item()) + base_input = 0.1 * torch.randn( + num_tokens, self.hidden_size, dtype=torch.bfloat16, device="cuda" ) - torch.testing.assert_close( - experts.linear_fc2.bias1.grad, - torch.ones_like(experts.linear_fc2.bias1) * probs[2:].detach().sum(), + base_probs = torch.rand(num_tokens, dtype=torch.bfloat16, device="cuda") + grad_output = 0.1 * torch.randn( + num_tokens, self.hidden_size, dtype=torch.bfloat16, device="cuda" ) + reference_input = base_input.detach().clone().requires_grad_(True) + reference_probs = base_probs.detach().clone().requires_grad_(True) + reference_output, _ = reference(reference_input, tokens_per_expert, reference_probs) + reference_output.backward(grad_output) + + target_input = base_input.detach().clone().requires_grad_(True) + target_probs = base_probs.detach().clone().requires_grad_(True) + target_output, _ = target(target_input, tokens_per_expert, target_probs) + target_output.backward(grad_output) + + tolerances = {"rtol": 1e-2, "atol": 1e-2} + torch.testing.assert_close(target_output, reference_output, **tolerances) + torch.testing.assert_close(target_input.grad, reference_input.grad, **tolerances) + torch.testing.assert_close(target_probs.grad, reference_probs.grad, **tolerances) + + for target_linear, reference_linear in ( + (target.linear_fc1, reference.linear_fc1), + (target.linear_fc2, reference.linear_fc2), + ): + torch.testing.assert_close( + packed_grad(target_linear, "weight"), + packed_grad(reference_linear, "weight"), + **tolerances, + ) + torch.testing.assert_close( + packed_grad(target_linear, "bias"), + packed_grad(reference_linear, "bias"), + **tolerances, + ) + + if use_op_fuser: + ops = target._fused_ops[0] + fc1_op = ops[0] + fc2_op = ops[2] + # The fused-op shells must register MCore's original parent parameters, not + # detached copies or member views, so autograd and the optimizer update the same objects. + if single_grouped_weight: + assert fc1_op.weight is target.linear_fc1.weight + assert fc2_op.weight is target.linear_fc2.weight + if single_grouped_bias: + assert fc1_op.bias is target.linear_fc1.bias + assert fc2_op.bias is target.linear_fc2.bias + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.internal def test_gpu_fused_path_loss_decreases(self): diff --git a/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py b/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py new file mode 100644 index 00000000000..e7d34345052 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py @@ -0,0 +1,560 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Distributed MoE coverage for the TE grouped-tensor expert path. + +The three dispatchers place expert-token padding at different boundaries: + +* All-to-All returns unpadded expert segments and TEGroupedMLP pads them before FC1. +* DeepEP communicates first, then its local fused permutation pads expert segments. +* HybridEP fuses communication, permutation, and expert-segment padding. + +The numerical tests compare each grouped-tensor configuration with the old discrete-parameter, +CPU-split path on the same dispatcher. The lifecycle tests inspect the actual expert-compute +boundary to ensure padding rows are zero and the dispatcher-specific inverse removes them. +""" + +import inspect +import os +from typing import Dict + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import config as mcore_config +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_submodules, +) +from megatron.core.transformer.module import Float16Module +from megatron.core.transformer.moe.fused_a2a import ( + HAVE_DEEP_EP, + HAVE_HYBRIDEP, + reset_hybrid_ep_buffer, +) +from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.moe.moe_utils import fused_permute_and_pad_with_probs +from megatron.core.transformer.spec_utils import get_submodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.initialize import _set_random_seed +from tests.unit_tests.test_utilities import Utils + +pytestmark = [ + pytest.mark.internal, + pytest.mark.launch_on_gb200, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), +] + + +_ALIGN_SIZE = 256 +_HIDDEN_SIZE = 256 +_MOE_FFN_HIDDEN_SIZE = 256 +_NUM_LOCAL_EXPERTS = 2 +_NUM_LOCAL_TOKENS = 128 +_TOLERANCES = {"rtol": 1e-2, "atol": 1e-2} + +_PARAMETER_LAYOUTS = ( + pytest.param(False, False, False, id="discrete-weight-no-bias"), + pytest.param(True, False, False, id="single-weight-no-bias"), + pytest.param(False, True, False, id="discrete-weight-discrete-bias"), + pytest.param(True, True, False, id="single-weight-discrete-bias"), + pytest.param(False, True, True, id="discrete-weight-single-bias"), + pytest.param(True, True, True, id="single-weight-single-bias"), +) + + +def _require_test_environment(dispatcher: str) -> int: + """Validate runtime support and return the EP world size used by the test.""" + world_size = torch.distributed.get_world_size() + if world_size < 2: + pytest.skip("DeepEP/HybridEP parity requires at least two distributed ranks") + if dispatcher == "deepep" and not HAVE_DEEP_EP: + pytest.skip("DeepEP is not available") + if dispatcher == "hybridep" and not HAVE_HYBRIDEP: + pytest.skip("HybridEP is not available") + if dispatcher == "deepep" and fused_permute_and_pad_with_probs is None: + pytest.skip("DeepEP grouped-tensor padding requires TE fused permute-and-pad") + + try: + from transformer_engine.pytorch.module import GroupedLinear + except ImportError: + pytest.skip("Transformer Engine GroupedLinear is not available") + parameters = inspect.signature(GroupedLinear.__init__).parameters + if "use_grouped_tensor" not in parameters or "single_grouped_bias" not in parameters: + pytest.skip("Installed TE lacks native grouped-tensor parameter support") + return world_size + + +def _dispatcher_options(dispatcher: str) -> Dict[str, object]: + """Return the production padding configuration for one dispatcher.""" + if dispatcher == "alltoall": + return { + "moe_token_dispatcher_type": "alltoall", + "moe_flex_dispatcher_backend": None, + "moe_permute_fusion": True, + } + if dispatcher == "deepep": + # DeepEP communicates first. Its fused local permutation then groups and pads tokens. + return { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "deepep", + "moe_permute_fusion": True, + } + if dispatcher == "hybridep": + # HybridEP owns permutation and padding inside its fused dispatch/combine kernels. + return { + "moe_token_dispatcher_type": "flex", + "moe_flex_dispatcher_backend": "hybridep", + "moe_permute_fusion": True, + } + raise ValueError(f"Unknown dispatcher {dispatcher!r}") + + +def _build_moe_layer( + dispatcher: str, + *, + ep_size: int, + use_grouped_tensor: bool, + single_grouped_weight: bool, + use_bias: bool, + single_grouped_bias: bool, +) -> MoELayer: + """Build a small real TE MoE layer without using the TE operation fuser.""" + options = _dispatcher_options(dispatcher) + transformer_config = TransformerConfig( + num_layers=1, + hidden_size=_HIDDEN_SIZE, + num_attention_heads=8, + num_moe_experts=ep_size * _NUM_LOCAL_EXPERTS, + moe_ffn_hidden_size=_MOE_FFN_HIDDEN_SIZE, + use_cpu_initialization=False, + add_bias_linear=use_bias, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bias_dropout_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="none", + moe_router_topk=2, + moe_aux_loss_coeff=0.0, + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_use_grouped_tensor=use_grouped_tensor, + moe_single_grouped_weight=single_grouped_weight, + moe_single_grouped_bias=single_grouped_bias, + use_transformer_engine_op_fuser=False, + tensor_model_parallel_size=1, + expert_model_parallel_size=ep_size, + sequence_parallel=False, + **options, + ) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + num_experts=transformer_config.num_moe_experts, moe_grouped_gemm=True + ).mlp + ) + assert isinstance(submodules, MoESubmodules) + layer = MoELayer(transformer_config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + layer.set_layer_number(0) + return layer + + +def _copy_linear_parameters(reference, target) -> None: + """Copy discrete expert parameters into either a discrete or packed target layout.""" + for parameter_name in ("weight", "bias"): + if parameter_name == "bias" and not reference.use_bias: + continue + reference_parts = torch.stack( + [ + getattr(reference, f"{parameter_name}{idx}").detach() + for idx in range(reference.num_gemms) + ] + ) + target_is_grouped = getattr(target, f"single_grouped_{parameter_name}") + if target_is_grouped: + grouped_parameter = getattr(target, parameter_name) + grouped_parameter.rowwise_data.view_as(reference_parts).copy_(reference_parts) + else: + for idx, part in enumerate(reference_parts): + getattr(target, f"{parameter_name}{idx}").copy_(part) + + +@torch.no_grad() +def _copy_layer_parameters(reference: MoELayer, target: MoELayer) -> None: + """Give reference and target identical router and expert parameters.""" + target_parameters = dict(target.named_parameters()) + for name, parameter in reference.named_parameters(): + if not name.startswith("experts."): + target_parameters[name].copy_(parameter) + _copy_linear_parameters(reference.experts.linear_fc1, target.experts.linear_fc1) + _copy_linear_parameters(reference.experts.linear_fc2, target.experts.linear_fc2) + + +def _canonical_gradient(linear, parameter_name: str) -> torch.Tensor: + """Return expert gradients as one [experts, ...] FP32 tensor for either layout.""" + if getattr(linear, f"single_grouped_{parameter_name}"): + gradient = getattr(linear, parameter_name).grad + assert gradient is not None, f"Grouped {parameter_name} parent has no gradient" + return gradient.reshape(linear.num_gemms, -1).float() + + gradients = [getattr(linear, f"{parameter_name}{idx}").grad for idx in range(linear.num_gemms)] + assert all(gradient is not None for gradient in gradients) + return torch.stack([gradient.reshape(-1) for gradient in gradients]).float() + + +def _run_forward_backward( + layer: MoELayer, base_input: torch.Tensor, grad_output: torch.Tensor +) -> Dict[str, torch.Tensor]: + """Run one MoE step and collect values sensitive to dispatch and parameter layout.""" + layer.zero_grad(set_to_none=True) + hidden_states = base_input.detach().clone().requires_grad_(True) + output, _ = layer(hidden_states) + output.backward(grad_output) + + result = { + "output": output.detach(), + "input_grad": hidden_states.grad.detach(), + "router_grad": layer.router.weight.grad.detach(), + "fc1_weight_grad": _canonical_gradient(layer.experts.linear_fc1, "weight"), + "fc2_weight_grad": _canonical_gradient(layer.experts.linear_fc2, "weight"), + } + if layer.config.add_bias_linear: + result["fc1_bias_grad"] = _canonical_gradient(layer.experts.linear_fc1, "bias") + result["fc2_bias_grad"] = _canonical_gradient(layer.experts.linear_fc2, "bias") + return result + + +def _run_numerical_parity_case( + dispatcher: str, *, single_grouped_weight: bool, use_bias: bool, single_grouped_bias: bool +) -> None: + """Compare grouped-tensor execution with the old path on the same dispatcher.""" + ep_size = _require_test_environment(dispatcher) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, expert_model_parallel_size=ep_size + ) + mcore_config.ENABLE_EXPERIMENTAL = True + + _set_random_seed(seed_=1234, data_parallel_random_init=False) + reference = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=False, + single_grouped_weight=False, + use_bias=use_bias, + single_grouped_bias=False, + ) + target = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=True, + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + _copy_layer_parameters(reference, target) + + base_input = torch.randn( + _NUM_LOCAL_TOKENS, 1, _HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda" + ) + grad_output = torch.randn_like(base_input) + + reference_result = _run_forward_backward(reference, base_input, grad_output) + target_result = _run_forward_backward(target, base_input, grad_output) + + assert target.experts._use_grouped_tensor + assert not reference.experts._use_grouped_tensor + assert reference_result.keys() == target_result.keys() + for name in reference_result: + torch.testing.assert_close( + target_result[name], + reference_result[name], + msg=lambda message, value=name: f"{dispatcher} {value} mismatch: {message}", + **_TOLERANCES, + ) + + +def _make_padding_mask( + real_tokens_per_expert: torch.Tensor, padded_tokens_per_expert: torch.Tensor +) -> torch.Tensor: + """Build a mask for the padding suffix in every expert-major segment. + + For example, using an alignment of four for readability: + + ``real_tokens_per_expert = tensor([3, 2])`` + ``padded_tokens_per_expert = tensor([4, 4])`` + + The packed expert-major rows are ``[e0 real x3][e0 pad x1][e1 real x2][e1 pad x2]``, + so this function returns ``[F, F, F, T, F, F, T, T]``. Indexing the packed hidden states or + probabilities with that mask selects only the synthetic rows that must contain exact zeros. + """ + masks = [] + for real_count, padded_count in zip( + real_tokens_per_expert.cpu().tolist(), padded_tokens_per_expert.cpu().tolist() + ): + # Each expert contributes a False prefix for real rows followed by a True padding suffix. + masks.append(torch.zeros(real_count, dtype=torch.bool, device="cuda")) + masks.append(torch.ones(padded_count - real_count, dtype=torch.bool, device="cuda")) + return torch.cat(masks) + + +def _infer_real_tokens_per_expert( + padded_probs: torch.Tensor, padded_tokens_per_expert: torch.Tensor +) -> torch.Tensor: + """Infer real expert counts from the zero suffix in every padded probability segment. + + This dropless test uses FP32 softmax router probabilities, so every routed token has a + nonzero probability. All padding implementations use exact zeros and append them after the + real rows in each expert-major segment. This gives one dispatcher-independent source of truth + without inspecting backend-specific routing metadata. + """ + real_counts = [] + offset = 0 + for padded_count in padded_tokens_per_expert.cpu().tolist(): + segment = padded_probs[offset : offset + padded_count] + nonzero_rows = segment != 0 + real_count = int(nonzero_rows.sum().item()) + + # Padding must be one contiguous suffix. Interspersed zero rows would preserve the total + # nonzero count while still violating the expert-major layout expected by grouped GEMM. + assert torch.all(nonzero_rows[:real_count]) + assert not torch.any(nonzero_rows[real_count:]) + real_counts.append(real_count) + offset += padded_count + + assert offset == padded_probs.numel() + return torch.tensor(real_counts, dtype=torch.int64, device=padded_probs.device) + + +def _install_padding_probes(layer: MoELayer, monkeypatch): + """Observe the tensors crossing each padding boundary without changing execution. + + ``register_forward_pre_hook`` runs immediately before the selected module's ``forward``. + The hook receives the module and the tuple of positional arguments that forward is about to + consume. Returning ``None`` leaves those arguments unchanged, so these hooks are read-only + probes rather than replacements for any production operation. + + There are two relevant boundaries. ``layer.experts`` sees what the token dispatcher hands to + TEGroupedMLP, while ``linear_fc1`` sees the final padded tensors and CUDA expert counts that + TE's grouped-tensor GEMM actually consumes. They are different for AllToAll, where TEGroupedMLP + owns padding, but identical for DeepEP and HybridEP, whose fused dispatch paths already pad. + """ + captured = {"padding_calls": 0, "unpadding_calls": 0} + + def capture_dispatcher_input(_module, args): + # TEGroupedMLP.forward(hidden, tokens_per_expert, permuted_probs) is about to run. + # DeepEP and HybridEP have already padded at this boundary, so preserve their router + # probabilities for the generic real-row inference below. Detach so the test does not + # retain the graph, and clone in case downstream computation reuses the input storage. + captured["dispatcher_probs"] = args[2].detach().clone() + + def capture_fc1_input(_module, args): + # GroupedLinear.forward(hidden, m_splits, ...) is about to run. This is the authoritative + # view of the rows and device-side split tensor presented to the grouped GEMM. + captured["padded_hidden"] = args[0].detach().clone() + captured["padded_counts"] = args[1].detach().clone() + + # Pre-hooks observe module inputs before either module can transform them. The returned hook + # handles need not be retained because each test owns this layer and executes one forward. + layer.experts.register_forward_pre_hook(capture_dispatcher_input) + layer.experts.linear_fc1.register_forward_pre_hook(capture_fc1_input) + + def capture_padding(_module, args, output): + # quantization_padding is called once for hidden states and once for router probabilities + # in the AllToAll path. A forward hook is used here because the padded tensor is its output. + captured["padding_calls"] += 1 + padded_tensor = output[0] + # Probability padding receives [tokens, 1], whereas hidden-state padding receives + # [tokens, hidden_size]. Preserve the padded probabilities for an exact-zero assertion. + if args[0].shape[-1] == 1: + captured["padded_probs"] = padded_tensor.detach().clone().reshape(-1) + + def capture_unpadding(_module, _args, output): + # AllToAll uses TEGroupedMLP's explicit unpadding module after expert compute. Recording + # the call proves that its locally inserted padding reaches the matching removal path. + captured["unpadding_calls"] += 1 + + layer.experts.quantization_padding.register_forward_hook(capture_padding) + layer.experts.quantization_unpadding.register_forward_hook(capture_unpadding) + + # All dispatchers share this final restoration API. + original_combine_postprocess = layer.token_dispatcher.combine_postprocess + + def capture_combine_postprocess(hidden_states, *args, **kwargs): + output = original_combine_postprocess(hidden_states, *args, **kwargs) + captured["restored_shape"] = output.shape + return output + + monkeypatch.setattr(layer.token_dispatcher, "combine_postprocess", capture_combine_postprocess) + + return captured + + +def _run_padding_lifecycle_case(dispatcher: str, monkeypatch) -> None: + """Verify exact zero padding, 256 alignment, no double-padding, and unpadding. + + This test follows one real MoE forward through dispatch, expert compute, and restoration. It + independently reconstructs the number of real tokens assigned to each local expert, then + compares that metadata with the padded CUDA ``m_splits`` observed directly at FC1. Numerical + parity with the legacy backend is tested separately; this case focuses on padding ownership + and the structural contract required by TE's grouped-tensor kernels. + """ + # EP spans the whole torchrun world so every tested dispatcher performs real communication. + # TP remains one to keep the independently reconstructed local-expert counts unambiguous. + ep_size = _require_test_environment(dispatcher) + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, expert_model_parallel_size=ep_size + ) + mcore_config.ENABLE_EXPERIMENTAL = True + + # Use the strictest parameter layout. If packed weight and packed bias reach the native + # grouped-tensor path correctly, discrete parameter layouts use the same padding lifecycle. + _set_random_seed(seed_=1357, data_parallel_random_init=False) + layer = _build_moe_layer( + dispatcher, + ep_size=ep_size, + use_grouped_tensor=True, + single_grouped_weight=True, + use_bias=True, + single_grouped_bias=True, + ) + # Install observers before the forward so they capture dispatcher output, FC1 input, and the + # common dispatcher boundary that returns to the original token layout. + captured = _install_padding_probes(layer, monkeypatch) + + torch.manual_seed(9753) + hidden_states = torch.randn( + _NUM_LOCAL_TOKENS, 1, _HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + output, _ = layer(hidden_states) + + padded_counts = captured["padded_counts"] + + # TE's grouped-tensor API requires device-resident int64 splits. Every expert segment must be + # represented by its m_split, and the physical FC1 input must contain their total. + assert padded_counts.device.type == "cuda" + assert padded_counts.dtype == torch.int64 + assert torch.all(padded_counts % _ALIGN_SIZE == 0) + assert captured["padded_hidden"].shape[0] == padded_counts.sum().item() + + if dispatcher == "alltoall": + # All-to-All returns real expert rows; TEGroupedMLP pads hidden states and probs itself. + assert captured["padding_calls"] == 2 + assert captured["unpadding_calls"] == 1 + else: + # DeepEP and HybridEP already return padded rows. TEGroupedMLP must not pad them again. + assert captured["padding_calls"] == 0 + assert captured["unpadding_calls"] == 0 + # For fused dispatchers, probabilities are already padded when they enter TEGroupedMLP, + # so the experts pre-hook is the correct observation point for the zero check below. + captured["padded_probs"] = captured["dispatcher_probs"].reshape(-1) + + # Router probabilities provide a common representation across all dispatchers: real routed + # rows are nonzero and padded rows are an exact-zero suffix. + real_counts = _infer_real_tokens_per_expert(captured["padded_probs"], padded_counts) + expected_padded_counts = ((real_counts + _ALIGN_SIZE - 1) // _ALIGN_SIZE) * _ALIGN_SIZE + torch.testing.assert_close(padded_counts, expected_padded_counts, rtol=0, atol=0) + + # Expert-major layout is [expert 0 real][expert 0 pad][expert 1 real][expert 1 pad]... + # Build that exact mask and require both hidden states and routing probabilities to use + # numerical zero for every synthetic row. Merely allocating the right shape is insufficient. + padding_mask = _make_padding_mask(real_counts, padded_counts) + # An expert may already have a 256-aligned token count and legitimately need no padding. + # Boolean indexing with an empty mask is valid; otherwise these checks inspect every pad row. + assert not torch.any(captured["padded_hidden"][padding_mask]) + assert not torch.any(captured["padded_probs"][padding_mask]) + + # Regardless of where a backend removes padding, the common dispatcher postprocess contract + # must restore exactly the shape that entered this MoE layer. + assert captured["restored_shape"] == hidden_states.shape + + # The public MoE contract is unchanged by internal alignment. Run backward as a final check + # that padding/unpadding preserved a connected, finite autograd path to the original input. + assert output.shape == hidden_states.shape + output.float().square().mean().backward() + assert hidden_states.grad is not None + assert torch.isfinite(hidden_states.grad).all() + + +class TestGroupedTensorDispatcherNumerics: + """Distributed numerical and padding coverage for grouped-tensor dispatchers.""" + + def setup_method(self, method): + if not torch.distributed.is_available() or Utils.world_size < 2: + pytest.skip("Distributed dispatcher tests must be launched with torchrun") + self._old_single_param_env = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + self._previous_experimental = mcore_config.ENABLE_EXPERIMENTAL + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + Utils.initialize_distributed() + + def teardown_method(self, method): + try: + mcore_config.ENABLE_EXPERIMENTAL = self._previous_experimental + reset_hybrid_ep_buffer() + Utils.destroy_model_parallel() + finally: + if self._old_single_param_env is None: + os.environ.pop("NVTE_GROUPED_LINEAR_SINGLE_PARAM", None) + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = self._old_single_param_env + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_alltoall_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """All-to-All grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "alltoall", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_deepep_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """DeepEP grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "deepep", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.parametrize( + "single_grouped_weight,use_bias,single_grouped_bias", _PARAMETER_LAYOUTS + ) + @pytest.mark.timeout(180) + def test_hybridep_grouped_tensor_moe_parity( + self, single_grouped_weight, use_bias, single_grouped_bias + ): + """HybridEP grouped-tensor MoE forward/backward matches its legacy expert path.""" + _run_numerical_parity_case( + "hybridep", + single_grouped_weight=single_grouped_weight, + use_bias=use_bias, + single_grouped_bias=single_grouped_bias, + ) + + @pytest.mark.timeout(180) + def test_alltoall_grouped_tensor_padding_lifecycle(self, monkeypatch): + """All-to-All explicitly pads in TEGroupedMLP and removes it before combine.""" + _run_padding_lifecycle_case("alltoall", monkeypatch) + + @pytest.mark.timeout(180) + def test_deepep_grouped_tensor_padding_lifecycle(self, monkeypatch): + """DeepEP fused local permutation pads, and local unpermute removes those rows.""" + _run_padding_lifecycle_case("deepep", monkeypatch) + + @pytest.mark.timeout(180) + def test_hybridep_grouped_tensor_padding_lifecycle(self, monkeypatch): + """HybridEP fused dispatch pads, and fused combine returns the original token shape.""" + _run_padding_lifecycle_case("hybridep", monkeypatch) diff --git a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py index 3e1bdb244d4..c6e2780ad97 100644 --- a/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py +++ b/tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py @@ -54,8 +54,12 @@ _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM = ( "single_grouped_weight" in inspect.signature(TEGroupedLinear.__init__).parameters ) + _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR = ( + "use_grouped_tensor" in inspect.signature(TEGroupedLinear.__init__).parameters + ) except (ImportError, AttributeError): _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM = False + _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR = False pytestmark = [ pytest.mark.internal, @@ -68,6 +72,10 @@ not _TE_GROUPED_LINEAR_SUPPORTS_SINGLE_PARAM, reason="Installed TE GroupedLinear does not expose single_grouped_weight", ), + pytest.mark.skipif( + not _TE_GROUPED_LINEAR_SUPPORTS_USE_GROUPED_TENSOR, + reason="Installed TE GroupedLinear does not expose use_grouped_tensor", + ), ] @@ -194,6 +202,7 @@ def create_test_args( args.num_experts = 2 args.moe_layer_freq = 1 args.moe_grouped_gemm = True + args.moe_use_grouped_tensor = True args.moe_single_grouped_weight = single_weight args.moe_token_dispatcher_type = "alltoall" args.moe_router_topk = 1 @@ -683,7 +692,22 @@ def test_mxfp8_single_weight_torch_dist_checkpoint_matches_discrete_baseline( self.assert_all_ranks_passed(local_passed, local_error) - @pytest.mark.parametrize("precision", ["bf16", "mxfp8", "nvfp4"]) + @pytest.mark.parametrize( + "precision", + [ + "bf16", + "mxfp8", + pytest.param( + "nvfp4", + marks=pytest.mark.skip( + reason=( + "NVFP4 single grouped weights are not supported by the " + "TransformerEngine native grouped-tensor path yet." + ) + ), + ), + ], + ) @pytest.mark.parametrize("gradient_accumulation_fusion", [False, True]) def test_single_grouped_weight_parity_with_primary_param_gather( self, precision, gradient_accumulation_fusion @@ -726,17 +750,18 @@ def test_single_grouped_weight_parity_without_primary_param_gather( use_transformer_engine_op_fuser=True, ) - def test_single_grouped_weight_parity_module_grouped_linear(self): - """Single grouped weights require the TE op-fuser execution path.""" - args = self.create_test_args( - precision="bf16", - primary_param_gather=False, - single_weight=True, - gradient_accumulation_fusion=False, + @pytest.mark.parametrize( + "precision,primary_param_gather", [("bf16", False), ("mxfp8", False), ("mxfp8", True)] + ) + @pytest.mark.parametrize("gradient_accumulation_fusion", [False, True]) + def test_single_grouped_weight_parity_module_grouped_linear( + self, precision, primary_param_gather, gradient_accumulation_fusion + ): + """Compare native TE GroupedLinear single and discrete parameter layouts.""" + _skip_if_unsupported(precision) + self.run_parity_case( + precision=precision, + primary_param_gather=primary_param_gather, + gradient_accumulation_fusion=gradient_accumulation_fusion, use_transformer_engine_op_fuser=False, ) - with pytest.raises( - ValueError, - match="moe_single_grouped_weight requires use_transformer_engine_op_fuser=True", - ): - core_transformer_config_from_args(args)