diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 7f7424e6c02..e4ab2fd2582 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -11,6 +11,8 @@ import re import warnings from contextlib import contextmanager, nullcontext +from itertools import accumulate +from numbers import Integral from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, cast import torch @@ -2358,6 +2360,7 @@ def __init__( name (str | None): module instance name passed top-down from its paranet module """ self.config = config + self.expert_gemm_backend = config.moe_expert_gemm_backend # TE returns a zero length Tensor when bias=False and # return_bias=True, but we prefer None. So in that case we @@ -2632,6 +2635,141 @@ def _split_grouped_checkpoint_tensor( f"into {self.num_gemms} GEMM shards." ) + def _torch_grouped_weight_layout_is_current(self) -> bool: + """Return whether each expert parameter still views the contiguous buffer.""" + grouped_weight = self._buffers.get('_torch_grouped_weight') + if grouped_weight is None or grouped_weight.shape[0] != self.num_gemms: + return False + + first_weight = self.weight0 + if ( + grouped_weight.device != first_weight.device + or grouped_weight.dtype != first_weight.dtype + or grouped_weight.shape[1:] != first_weight.shape + ): + return False + + storage_pointer = grouped_weight.untyped_storage().data_ptr() + expert_elements = first_weight.numel() + for index in range(self.num_gemms): + weight = getattr(self, f'weight{index}') + grouped_weight_view = grouped_weight[index] + if ( + weight.device != grouped_weight.device + or weight.dtype != grouped_weight.dtype + or weight.shape != grouped_weight_view.shape + or weight.untyped_storage().data_ptr() != storage_pointer + or weight.storage_offset() != index * expert_elements + or weight.stride() != grouped_weight_view.stride() + ): + return False + return True + + @torch.no_grad() + def prepare_torch_grouped_mm(self) -> int: + """Relocate frozen BF16 expert weights into one contiguous allocation. + + The existing per-expert ``Parameter`` objects become views of the new allocation, so + parameter names and checkpoint structure remain unchanged. The backing allocation is a + non-persistent buffer and therefore does not add a checkpoint entry. + + Returns: + Number of bytes in the contiguous backing allocation. + """ + if self.expert_gemm_backend != 'torch': + raise RuntimeError( + "prepare_torch_grouped_mm requires moe_expert_gemm_backend='torch'" + ) + self._resolve_torch_grouped_mm() + if self.use_bias: + raise RuntimeError("torch grouped expert GEMM does not support bias") + if getattr(self, 'single_grouped_weight', False): + raise RuntimeError( + "torch grouped expert GEMM expects per-GEMM weight parameters, but " + "moe_single_grouped_weight=True makes TE hold one grouped weight tensor" + ) + + weights = [getattr(self, f'weight{index}') for index in range(self.num_gemms)] + if any(weight.requires_grad for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires frozen base weights") + if any(weight.device.type != 'cuda' for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires CUDA weights") + if any(weight.dtype != torch.bfloat16 for weight in weights): + raise RuntimeError("torch grouped expert GEMM requires BF16 weights") + if any(weight.shape != weights[0].shape for weight in weights[1:]): + raise RuntimeError("torch grouped expert GEMM requires uniform expert shapes") + + if not self._torch_grouped_weight_layout_is_current(): + grouped_weight = torch.stack([weight.detach() for weight in weights]).contiguous() + for index, weight in enumerate(weights): + weight.data = grouped_weight[index] + if '_torch_grouped_weight' in self._buffers: + self._buffers['_torch_grouped_weight'] = grouped_weight + else: + self.register_buffer('_torch_grouped_weight', grouped_weight, persistent=False) + + grouped_weight = self._buffers['_torch_grouped_weight'] + return grouped_weight.numel() * grouped_weight.element_size() + + @property + def torch_grouped_mm_prepared(self) -> bool: + """Return whether the torch grouped GEMM backing allocation is ready.""" + return self._torch_grouped_weight_layout_is_current() + + @staticmethod + def _resolve_torch_grouped_mm(): + """Return the public torch grouped GEMM, with the legacy private op as fallback.""" + grouped_mm = getattr(F, 'grouped_mm', None) + if grouped_mm is None: + grouped_mm = getattr(torch, '_grouped_mm', None) + if grouped_mm is None: + raise RuntimeError( + "this PyTorch build does not provide torch.nn.functional.grouped_mm " + "or torch._grouped_mm" + ) + return grouped_mm + + def _validate_torch_grouped_mm_splits(self, m_splits) -> list[int]: + """Validate and normalize the token count for each local expert.""" + try: + split_count = len(m_splits) + except TypeError as error: + raise RuntimeError( + f"expected {self.num_gemms} expert splits, but expert splits are not a sequence" + ) from error + if split_count != self.num_gemms: + raise RuntimeError(f"expected {self.num_gemms} expert splits, got {split_count}") + if any( + isinstance(split, bool) or not isinstance(split, Integral) or split < 0 + for split in m_splits + ): + raise RuntimeError("expert splits must be non-negative integers") + return [int(split) for split in m_splits] + + def _torch_grouped_mm_forward(self, x, m_splits): + """Run the frozen expert base branch with torch grouped GEMM.""" + m_splits = self._validate_torch_grouped_mm_splits(m_splits) + if not self._torch_grouped_weight_layout_is_current(): + self.prepare_torch_grouped_mm() + + x_2d = x.reshape(-1, x.shape[-1]) + input_rows = sum(m_splits) + if input_rows != x_2d.shape[0]: + raise RuntimeError( + f"expert splits sum to {input_rows}, but input has {x_2d.shape[0]} rows" + ) + + grouped_weight = self._buffers['_torch_grouped_weight'] + if input_rows == 0: + output = x_2d.matmul(grouped_weight[0].transpose(0, 1)) + else: + offsets = torch.tensor( + list(accumulate(m_splits)), device=x.device, dtype=torch.int32 + ) + grouped_mm = self._resolve_torch_grouped_mm() + output = grouped_mm(x_2d, grouped_weight.transpose(1, 2), offs=offsets) + return output.reshape(*x.shape[:-1], output.shape[-1]), None + def finish_init(self, quantization_config: QuantizationConfig): """Post-init of quantization override""" if quantization_config is None: @@ -2647,6 +2785,9 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool: def forward(self, x, m_splits): """Forward.""" + if self.expert_gemm_backend == 'torch': + return self._torch_grouped_mm_forward(x, m_splits) + _is_first_microbatch = ( None if self.disable_parameter_transpose_cache else self.is_first_microbatch ) diff --git a/megatron/core/recompute.py b/megatron/core/recompute.py index bd0d1bcb3b2..fb68354781c 100644 --- a/megatron/core/recompute.py +++ b/megatron/core/recompute.py @@ -154,6 +154,13 @@ def chunk_runner(start: int, end: int, use_checkpoint: bool): if (start + layer_offset) in extract_layer_indices: intermediate_hidden_states.append(hidden_states) + if not hidden_states.requires_grad: + # Re-entrant checkpointing only attaches a grad_fn when some tensor input + # requires grad. With a frozen embedding (adapter-only training) every chunk + # output would otherwise carry no grad_fn and the adapters inside the chunks + # would receive no gradient. + hidden_states = hidden_states.detach().requires_grad_(True) + if self.config.recompute_method == 'uniform': # Uniformly divide the total number of layers and checkpoint # the input activation of each divided chunk. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..ee48f722c57 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -843,6 +843,14 @@ class TransformerConfig(ModelParallelConfig): parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` and ``add_bias_linear=True``.""" + moe_expert_gemm_backend: Literal['transformer_engine', 'torch'] = 'transformer_engine' + """Backend for grouped expert linear layers during training. + + ``torch`` uses ``torch._grouped_mm`` for frozen BF16, bias-free expert weights. It is intended + for parameter-efficient fine-tuning where only the adapter branch is trainable. The default + ``transformer_engine`` backend supports trainable expert weights and other precisions. + """ + moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. If a list of load balancing types is provided for `moe_router_load_balancing_type`, @@ -1551,6 +1559,60 @@ def __post_init__(self): if self.num_moe_experts is not None and self.num_moe_experts <= 0: raise ValueError("num_moe_experts must be non-negative.") + if self.moe_expert_gemm_backend not in ('transformer_engine', 'torch'): + raise ValueError( + "moe_expert_gemm_backend must be 'transformer_engine' or 'torch', " + f"got {self.moe_expert_gemm_backend!r}" + ) + if self.moe_expert_gemm_backend == 'torch': + if self.num_moe_experts is None: + raise ValueError("moe_expert_gemm_backend='torch' requires num_moe_experts") + if not self.moe_grouped_gemm: + raise ValueError("moe_expert_gemm_backend='torch' requires moe_grouped_gemm=True") + if not self.bf16 or self.params_dtype != torch.bfloat16: + raise ValueError("moe_expert_gemm_backend='torch' requires BF16 parameters") + if self.add_bias_linear: + raise ValueError("moe_expert_gemm_backend='torch' does not support expert bias") + if self.fp8 or self.fp4: + raise ValueError( + "moe_expert_gemm_backend='torch' does not support FP8 or FP4 experts" + ) + if self.moe_single_grouped_weight: + raise ValueError( + "moe_expert_gemm_backend='torch' requires per-GEMM expert weight parameters " + "and is incompatible with moe_single_grouped_weight=True" + ) + if self.use_transformer_engine_op_fuser: + raise ValueError( + "moe_expert_gemm_backend='torch' is incompatible with " + "use_transformer_engine_op_fuser=True" + ) + if self.delay_wgrad_compute: + raise ValueError( + "moe_expert_gemm_backend='torch' is incompatible with " + "delay_wgrad_compute=True" + ) + if self.overlap_dispatch_backward_with_experts_wgrad: + raise ValueError( + "moe_expert_gemm_backend='torch' is incompatible with " + "overlap_dispatch_backward_with_experts_wgrad=True" + ) + if self.transformer_impl != 'transformer_engine': + raise ValueError( + "moe_expert_gemm_backend='torch' requires " + "transformer_impl='transformer_engine' to construct TEGroupedMLP" + ) + + # Import lazily to avoid the transformer_engine -> TransformerConfig import cycle and + # prevent TEGroupedMLP from silently falling back to SequentialMLP. + from megatron.core.extensions.transformer_engine import TEColumnParallelGroupedLinear + + if TEColumnParallelGroupedLinear is None: + raise ValueError( + "moe_expert_gemm_backend='torch' requires Transformer Engine >= 1.9.0.dev0 " + "with GroupedLinear support" + ) + if self.num_moe_experts is not None and self.moe_ffn_hidden_size is None: self.moe_ffn_hidden_size = self.ffn_hidden_size warnings.warn("moe_ffn_hidden_size is not set, using ffn_hidden_size instead.") diff --git a/tests/unit_tests/ssm/test_hybrid_block.py b/tests/unit_tests/ssm/test_hybrid_block.py index 5d3c33264f4..d91a63e19e5 100644 --- a/tests/unit_tests/ssm/test_hybrid_block.py +++ b/tests/unit_tests/ssm/test_hybrid_block.py @@ -228,6 +228,58 @@ def run(block, hs, am): gb, gr = base_grads[name], rec_grads[name] assert torch.equal(gr, gb), f"Grad should be bitwise matched for {name}" + @pytest.mark.parametrize( + ("recompute_method", "num_layers", "recompute_num_layers"), + [pytest.param("uniform", 2, 1, id="uniform"), pytest.param("block", 3, 2, id="block")], + ) + def test_full_recompute_with_frozen_input_trains_adapter( + self, recompute_method: str, num_layers: int, recompute_num_layers: int + ): + """Full re-entrant recompute preserves gradients for adapter-only training.""" + block = self.get_hybrid_block( + Symbols.MLP * num_layers, + add_bias_linear=False, + recompute_granularity="full", + recompute_method=recompute_method, + recompute_num_layers=recompute_num_layers, + ).cuda() + block.requires_grad_(False) + block.train() + + def register_adapter(layer): + adapter = torch.nn.Sequential( + torch.nn.Linear(block.config.hidden_size, 4, bias=False), + torch.nn.Linear(4, block.config.hidden_size, bias=False), + ).cuda() + layer.add_module("adapter", adapter) + + def apply_adapter(_module, _args, kwargs, output): + hidden_states, context = output + return hidden_states + adapter(kwargs["hidden_states"]), context + + layer.register_forward_hook(apply_adapter, with_kwargs=True) + return adapter + + # Block recompute checkpoints layer 0 but not the final layer in this configuration. + adapters = [register_adapter(block.layers[index]) for index in (0, num_layers - 1)] + + sequence_length, micro_batch_size = 4, 1 + hidden_states = torch.randn( + sequence_length, micro_batch_size, block.config.hidden_size, device="cuda" + ) + attention_mask = torch.ones( + (micro_batch_size, 1, sequence_length, sequence_length), dtype=bool, device="cuda" + ) + + assert not hidden_states.requires_grad + output = block(hidden_states, attention_mask=attention_mask) + output.float().square().mean().backward() + + for adapter in adapters: + for parameter in adapter.parameters(): + assert parameter.grad is not None + assert torch.count_nonzero(parameter.grad) > 0 + def test_layer_types(self): """ Make sure that the layer types specified with layer_pattern diff --git a/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py new file mode 100644 index 00000000000..e8ed52c8752 --- /dev/null +++ b/tests/unit_tests/transformer/moe/test_frozen_torch_grouped_mlp.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch +import torch.nn.functional as F + +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.moe_layer import MoELayer +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import is_te_min_version +from tests.unit_tests.test_utilities import Utils + + +def _config(*, backend: str = "transformer_engine") -> TransformerConfig: + return TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + num_moe_experts=4, + moe_grouped_gemm=True, + moe_expert_gemm_backend=backend, + add_bias_linear=False, + gated_linear_unit=False, + activation_func=F.gelu, + bias_activation_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + ) + + +@pytest.mark.parametrize( + "override,error", + [ + ({"num_moe_experts": None}, "requires num_moe_experts"), + ({"moe_grouped_gemm": False}, "requires moe_grouped_gemm=True"), + ({"bf16": False, "params_dtype": torch.float32}, "requires BF16 parameters"), + ({"add_bias_linear": True}, "does not support expert bias"), + ( + {"use_transformer_engine_op_fuser": True}, + "is incompatible with use_transformer_engine_op_fuser=True", + ), + ], +) +def test_torch_grouped_expert_gemm_config_validation(override, error): + kwargs = { + "num_layers": 1, + "hidden_size": 16, + "num_attention_heads": 4, + "num_moe_experts": 2, + "moe_grouped_gemm": True, + "moe_expert_gemm_backend": "torch", + "add_bias_linear": False, + "bf16": True, + "params_dtype": torch.bfloat16, + } + kwargs.update(override) + + with pytest.raises(ValueError, match=error): + TransformerConfig(**kwargs) + + +@pytest.mark.skipif( + not is_te_min_version("1.9.0.dev0") + or not torch.cuda.is_available() + or not (hasattr(F, "grouped_mm") or hasattr(torch, "_grouped_mm")) + or torch.cuda.get_device_capability()[0] < 10, + reason="torch grouped expert GEMM requires TE and a compatible CUDA device", +) +class TestFrozenTorchGroupedMLP: + def setup_method(self): + Utils.initialize_model_parallel(1, 1) + + def teardown_method(self): + Utils.destroy_model_parallel() + + @staticmethod + def _model(config: TransformerConfig) -> MoELayer: + model = MoELayer( + config, + get_gpt_layer_with_transformer_engine_submodules( + config.num_moe_experts, moe_grouped_gemm=True + ).mlp.submodules, + ) + return Float16Module(config, model).module.cuda() + + def test_forward_input_gradient_and_checkpoint_parity(self): + baseline = self._model(_config()) + torch_model = self._model(_config(backend="torch")) + torch_model.load_state_dict(baseline.state_dict()) + torch_linears = (torch_model.experts.linear_fc1, torch_model.experts.linear_fc2) + for linear in torch_linears: + for parameter in linear.parameters(): + parameter.requires_grad = False + state_keys_before = set(torch_model.state_dict()) + + baseline_input = torch.rand( + (32, 2, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + torch_input = baseline_input.detach().clone().requires_grad_(True) + baseline_output, _ = baseline(baseline_input) + torch_output, _ = torch_model(torch_input) + grad_output = torch.randn_like(baseline_output) + baseline_output.backward(grad_output) + torch_output.backward(grad_output) + + torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02) + torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02) + assert set(torch_model.state_dict()) == state_keys_before + assert not any("_torch_grouped_weight" in key for key in state_keys_before) + for linear in torch_linears: + assert linear.torch_grouped_mm_prepared + assert ( + linear._buffers["_torch_grouped_weight"].untyped_storage().data_ptr() + == linear.weight0.untyped_storage().data_ptr() + ) + + def test_no_tokens(self): + model = self._model(_config(backend="torch")) + for parameter in model.experts.parameters(): + parameter.requires_grad = False + hidden_states = torch.empty( + (0, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + + output, _ = model.experts( + hidden_states, + tokens_per_expert=torch.zeros(4, dtype=torch.int32, device="cuda"), + permuted_probs=torch.empty(0, dtype=torch.float32, device="cuda"), + ) + output.sum().backward() + + assert output.shape == (0, 16) + assert hidden_states.grad is not None + + @pytest.mark.parametrize( + "m_splits,error", + [ + ([0, 0, 0], "expected 4 expert splits, got 3"), + ([0, 0, 0, 0, 0], "expected 4 expert splits, got 5"), + ([1, -1, 0, 0], "expert splits must be non-negative integers"), + ([0, 0, 0, 0.0], "expert splits must be non-negative integers"), + ], + ) + def test_no_tokens_rejects_malformed_splits(self, m_splits, error): + model = self._model(_config(backend="torch")) + linear = model.experts.linear_fc1 + hidden_states = torch.empty((0, 16), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(RuntimeError, match=error): + linear(hidden_states, m_splits) + + @pytest.mark.parametrize("available_api", ["public", "private"]) + def test_nonempty_forward_uses_available_grouped_mm(self, monkeypatch, available_api): + model = self._model(_config(backend="torch")) + linear = model.experts.linear_fc1 + for parameter in linear.parameters(): + parameter.requires_grad = False + hidden_states = torch.randn((4, 16), dtype=torch.bfloat16, device="cuda") + m_splits = [1, 1, 1, 1] + expected = torch.cat( + [ + hidden_states[index : index + 1] @ getattr(linear, f"weight{index}").transpose(0, 1) + for index in range(4) + ] + ) + seen_offsets = [] + + def grouped_mm(mat_a, mat_b, *, offs): + seen_offsets.append(offs.clone()) + outputs = [] + start = 0 + for expert_index, end in enumerate(offs.tolist()): + outputs.append(mat_a[start:end] @ mat_b[expert_index]) + start = end + return torch.cat(outputs) + + if available_api == "public": + monkeypatch.setattr(F, "grouped_mm", grouped_mm, raising=False) + monkeypatch.delattr(torch, "_grouped_mm", raising=False) + else: + monkeypatch.delattr(F, "grouped_mm", raising=False) + monkeypatch.setattr(torch, "_grouped_mm", grouped_mm, raising=False) + + output, output_bias = linear(hidden_states, m_splits) + + assert len(seen_offsets) == 1 + assert seen_offsets[0].tolist() == [1, 2, 3, 4] + torch.testing.assert_close(output, expected) + assert output_bias is None + + def test_reprepares_after_middle_expert_weight_replacement(self): + model = self._model(_config(backend="torch")) + linear = model.experts.linear_fc1 + for parameter in linear.parameters(): + parameter.requires_grad = False + linear.prepare_torch_grouped_mm() + replacement = torch.randn_like(linear.weight1) + + linear.weight1.data = replacement + + assert not linear.torch_grouped_mm_prepared + linear.prepare_torch_grouped_mm() + assert linear.torch_grouped_mm_prepared + torch.testing.assert_close(linear.weight1, replacement) + + def test_zero_token_groups(self): + baseline = self._model(_config()) + torch_model = self._model(_config(backend="torch")) + torch_model.load_state_dict(baseline.state_dict()) + for parameter in torch_model.experts.parameters(): + parameter.requires_grad = False + + baseline_input = torch.rand( + (5, 16), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + torch_input = baseline_input.detach().clone().requires_grad_(True) + tokens_per_expert = torch.tensor([2, 0, 3, 0], device="cuda") + permuted_probs = torch.rand(5, dtype=torch.float32, device="cuda") + baseline_output, _ = baseline.experts(baseline_input, tokens_per_expert, permuted_probs) + torch_output, _ = torch_model.experts(torch_input, tokens_per_expert, permuted_probs) + grad_output = torch.randn_like(baseline_output) + baseline_output.backward(grad_output) + torch_output.backward(grad_output) + + torch.testing.assert_close(torch_output, baseline_output, rtol=0.02, atol=0.02) + torch.testing.assert_close(torch_input.grad, baseline_input.grad, rtol=0.02, atol=0.02) diff --git a/tests/unit_tests/transformer/test_transformer_config.py b/tests/unit_tests/transformer/test_transformer_config.py index febb3842789..8a0637a4bea 100644 --- a/tests/unit_tests/transformer/test_transformer_config.py +++ b/tests/unit_tests/transformer/test_transformer_config.py @@ -1,7 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import pytest +import torch +from megatron.core.extensions import transformer_engine as transformer_engine_extension from megatron.core.transformer.transformer_config import TransformerConfig @@ -30,3 +32,53 @@ def test_ep_a2a_overlap_accepts_supported_mtp_layer_counts(mtp_num_layers: int | def test_ep_a2a_overlap_rejects_unsupported_mtp_layer_counts(mtp_num_layers: int): with pytest.raises(AssertionError, match="MTP supports at most one layer"): _make_overlap_config(mtp_num_layers) + + +def _make_torch_grouped_expert_config(**overrides) -> TransformerConfig: + kwargs = { + "num_layers": 1, + "hidden_size": 16, + "num_attention_heads": 4, + "num_moe_experts": 2, + "moe_grouped_gemm": True, + "moe_expert_gemm_backend": "torch", + "add_bias_linear": False, + "bf16": True, + "params_dtype": torch.bfloat16, + } + kwargs.update(overrides) + return TransformerConfig(**kwargs) + + +@pytest.mark.parametrize( + ("overrides", "error_match"), + [ + ({"delay_wgrad_compute": True}, "incompatible with delay_wgrad_compute=True"), + ( + {"overlap_dispatch_backward_with_experts_wgrad": True}, + "incompatible with overlap_dispatch_backward_with_experts_wgrad=True", + ), + ( + {"transformer_impl": "local"}, + "requires transformer_impl='transformer_engine' to construct TEGroupedMLP", + ), + ], +) +def test_torch_grouped_expert_backend_rejects_incompatible_modes(overrides, error_match): + with pytest.raises(ValueError, match=error_match): + _make_torch_grouped_expert_config(**overrides) + + +def test_torch_grouped_expert_backend_rejects_sequential_mlp_fallback(monkeypatch): + monkeypatch.setattr(transformer_engine_extension, "TEColumnParallelGroupedLinear", None) + + with pytest.raises(ValueError, match="requires Transformer Engine >= 1.9.0.dev0"): + _make_torch_grouped_expert_config() + + +def test_torch_grouped_expert_backend_accepts_te_grouped_linear(monkeypatch): + monkeypatch.setattr(transformer_engine_extension, "TEColumnParallelGroupedLinear", object) + + config = _make_torch_grouped_expert_config() + + assert config.moe_expert_gemm_backend == "torch"