diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index c1ac784f7a..4bd5447a85 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -3,6 +3,8 @@ # See LICENSE for license information. """Multi-process PyTorch EP tests, launched via torchrun (one process per GPU).""" +from contextlib import nullcontext +from dataclasses import replace import os import sys import unittest @@ -11,11 +13,16 @@ import torch import torch.distributed as dist +import transformer_engine.pytorch as te +from transformer_engine.pytorch import ops as te_ops from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, + EpConfig, ep_bootstrap, ep_finalize, + get_ep_drop_on_overflow, + get_ep_group, ep_prepare, ep_dispatch, ep_combine, @@ -25,6 +32,19 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ep_reference import ( + BlockScaledTensor, + MoeEpReference, + MoeFormat, + quantize_blockwise, +) +from transformer_engine.pytorch.ops.fused.moe_ep import ( + FusedMoeEp, + _cudnn_megamoe_supported, + _get_megamoe_combine_format, + _pack_cudnn_weights, +) +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" @@ -136,6 +156,31 @@ def _make_identity_inputs(rank, ep_size, device="cuda"): ) +def _make_moe_inputs(rank, ep_size, device="cuda"): + """Deterministic BF16 activations and FP32 top-k router weights.""" + generator = torch.Generator(device=device) + generator.manual_seed(2026 + rank) + tokens = ( + torch.randn( + TOKENS_PER_RANK, + HIDDEN_DIM, + generator=generator, + dtype=torch.float32, + device=device, + ) + * 0.25 + ).to(torch.bfloat16) + router_logits = torch.randn( + TOKENS_PER_RANK, + ep_size * NUM_LOCAL_EXPERTS, + generator=generator, + dtype=torch.float32, + device=device, + ) + topk_logits, topk_idx = torch.topk(router_logits, TOP_K, dim=-1) + return topk_idx, tokens, torch.softmax(topk_logits, dim=-1) + + def _degroup_mxfp8(recv_grouped, valid_counts=None): """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order. With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert @@ -146,6 +191,20 @@ def _degroup_mxfp8(recv_grouped, valid_counts=None): return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0) +def _reference_weights(op): + """Pack GroupedLinear weights into MoeEpReference ``(E, in, out)`` layout.""" + packed, _ = _pack_cudnn_weights(op, block_scaled_cls=BlockScaledTensor) + if isinstance(packed, torch.Tensor): + return packed.detach() + return BlockScaledTensor( + data=packed.data.detach(), + scale=packed.scale.detach(), + format=packed.format, + logical_shape=packed.logical_shape, + axis=packed.axis, + ) + + class _Cfg: rank: int world_size: int @@ -173,16 +232,24 @@ def _make_cfg() -> _Cfg: return cfg -class TestEP(unittest.TestCase): +class _EpTestCase(unittest.TestCase): + """Shared NCCL EP process state and pass selection.""" + cfg: _Cfg ep_group: dist.ProcessGroup @classmethod def setUpClass(cls): + if hasattr(_EpTestCase, "cfg"): + cls.cfg = _EpTestCase.cfg + cls.ep_group = _EpTestCase.ep_group + return if _device_sm() < 90: raise unittest.SkipTest(f"NCCL EP requires SM>=90 (got SM{_device_sm()})") - cls.cfg = _make_cfg() - cls.ep_group = _build_ep_group() + _EpTestCase.cfg = _make_cfg() + _EpTestCase.ep_group = _build_ep_group() + cls.cfg = _EpTestCase.cfg + cls.ep_group = _EpTestCase.ep_group ep_bootstrap( cls.ep_group, num_experts=cls.cfg.num_experts, @@ -220,24 +287,64 @@ def setUp(self): ): self.skipTest("not exercised in overflow mode") - def _make_buffer( + def _make_config( self, alignment=0, top_k=TOP_K, - dispatch_fwd_quant_recipe=None, - combine_bwd_quant_recipe=None, ): - return EpBuffer( + return EpConfig( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, + ep_group=self.ep_group, alignment=alignment, + zero_copy=ZERO_COPY, + drop_on_overflow=OVERFLOW, + ) + + def _make_buffer_from_config( + self, + config, + *, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): + return EpBuffer( + top_k=config.top_k, + max_tokens_per_rank=config.max_tokens_per_rank, + hidden_dim=config.hidden_dim, + num_local_experts=config.num_local_experts, + recv_capacity_per_rank=config.recv_capacity_per_rank, + alignment=config.alignment, + payload_dtype=config.payload_dtype, dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, combine_bwd_quant_recipe=combine_bwd_quant_recipe, ) + def _make_buffer( + self, + alignment=0, + top_k=TOP_K, + dispatch_fwd_quant_recipe=None, + combine_bwd_quant_recipe=None, + ): + config = self._make_config(alignment=alignment, top_k=top_k) + return self._make_buffer_from_config( + config, + dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe, + combine_bwd_quant_recipe=combine_bwd_quant_recipe, + ) + + +class TestEP(_EpTestCase): + """NCCL EP tests inherited from nvidia_origin/main.""" + + def test_bootstrap_accessors(self): + self.assertIs(get_ep_group(), self.ep_group) + self.assertEqual(get_ep_drop_on_overflow(), OVERFLOW) + def _expert_out(self, expert_out): """Stage the combine input into symm-mem under zero-copy (combine requires it).""" if not ZERO_COPY: @@ -829,6 +936,632 @@ def test_combine_autograd(self): torch.testing.assert_close(tokens_p.grad.float(), tokens.float(), atol=5e-2, rtol=5e-2) +class TestMoeEpSequential(_EpTestCase): + """Integration tests for Dispatch -> expert MLP -> Combine sequences.""" + + def _mxfp8_quantizer(self): + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + + def _require_mxfp8_shapes(self): + if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0: + self.skipTest( + "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 " + "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)" + ) + + def test_runtime_buffer_config_mismatch(self): + config = self._make_config() + buffer = self._make_buffer_from_config(config) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + dispatch = te_ops.MoeDispatch(config) + replacements = { + "top_k": config.top_k + 1, + "hidden_dim": config.hidden_dim + 1, + "num_local_experts": config.num_local_experts + 1, + "max_tokens_per_rank": config.max_tokens_per_rank + 1, + "recv_capacity_per_rank": config.recv_capacity_per_rank + 1, + "alignment": 1, + "payload_dtype": torch.float16, + "zero_copy": not config.zero_copy, + } + for field_name, wrong_value in replacements.items(): + with self.subTest(field_name=field_name): + original = getattr(buffer, field_name) + setattr(buffer, field_name, wrong_value) + try: + with self.assertRaisesRegex(ValueError, field_name): + dispatch( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + finally: + setattr(buffer, field_name, original) + + wrong_config = replace(config, drop_on_overflow=not config.drop_on_overflow) + with self.assertRaisesRegex(ValueError, "drop_on_overflow"): + te_ops.MoeDispatch(wrong_config)( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + with self.assertRaisesRegex(ValueError, "ep_group"): + te_ops.MoeDispatch(replace(config, ep_group=None))( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + expert_out = torch.empty( + self.cfg.recv_capacity_per_rank, + HIDDEN_DIM, + dtype=torch.bfloat16, + device=self.cfg.device, + ) + with self.assertRaisesRegex(ValueError, "runtime buffer config"): + original = buffer.hidden_dim + buffer.hidden_dim += 1 + try: + te_ops.MoeCombine(config)(expert_out, buffer=buffer) + finally: + buffer.hidden_dim = original + + def test_megamoe_combine_format_env(self): + old_value = os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE") + try: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "0" + self.assertEqual(_get_megamoe_combine_format(), "bf16") + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "1" + self.assertEqual(_get_megamoe_combine_format(), "mxfp8") + finally: + if old_value is None: + os.environ.pop("NVTE_MEGAMOE_MXFP8_COMBINE", None) + else: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = old_value + + @_mxfp8_align_test + def test_role_quantizer_requires_matching_buffer_recipe(self): + self._require_mxfp8_shapes() + config = self._make_config(alignment=128) + buffer = self._make_buffer_from_config(config) + dispatch = te_ops.MoeDispatch(config) + combine = te_ops.MoeCombine(config) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + recipe = MXFP8BlockScaling() + with te.autocast(enabled=True, recipe=recipe): + with self.assertRaisesRegex(ValueError, "does not have an MXFP8BlockScaling recipe"): + dispatch(tokens, topk_idx, topk_weights, buffer=buffer) + expert_out = torch.empty( + self.cfg.recv_capacity_per_rank, + HIDDEN_DIM, + dtype=torch.bfloat16, + device=self.cfg.device, + requires_grad=True, + ) + with self.assertRaisesRegex(ValueError, "does not have an MXFP8BlockScaling recipe"): + combine(expert_out, buffer=buffer) + + def _make_dispatch_combine_ops(self, *, mxfp8): + recipe = MXFP8BlockScaling() if mxfp8 else None + config = self._make_config(alignment=128 if mxfp8 else 0) + buffer = self._make_buffer_from_config( + config, + dispatch_fwd_quant_recipe=recipe, + combine_bwd_quant_recipe=recipe, + ) + return ( + buffer, + te_ops.MoeDispatch(config), + te_ops.MoeCombine(config), + ) + + def _run_dispatch_combine_identity(self, *, mxfp8): + """Route, apply top-k weights, and combine back to local token order.""" + if mxfp8: + self._require_mxfp8_shapes() + buffer, dispatch, combine = self._make_dispatch_combine_ops(mxfp8=mxfp8) + topk_idx, tokens, topk_weights = _make_identity_inputs( + self.cfg.rank, + self.cfg.ep_size, + ) + recipe = MXFP8BlockScaling() if mxfp8 else None + with te.autocast(enabled=mxfp8, recipe=recipe): + recv_tokens, tokens_per_expert, recv_weights = dispatch( + tokens, + topk_idx, + topk_weights, + buffer=buffer, + ) + if mxfp8: + recv_tokens = _degroup_mxfp8(recv_tokens) + recv_weights = recv_weights[: recv_tokens.shape[0]] + weighted_expert_output = (recv_tokens.float() * recv_weights.float().unsqueeze(-1)).to( + torch.bfloat16 + ) + output = combine( + weighted_expert_output, + buffer=buffer, + ) + torch.cuda.synchronize() + torch.testing.assert_close(output, tokens, atol=5e-2, rtol=5e-2) + self.assertEqual(tokens_per_expert.data_ptr(), buffer.tokens_per_expert.data_ptr()) + + @_eager_test_include + def test_dispatch_combine_identity_bf16(self): + """MoeDispatch and MoeCombine basic ops form an identity in BF16.""" + self._run_dispatch_combine_identity(mxfp8=False) + + @_eager_test_include + @_mxfp8_align_test + def test_dispatch_combine_identity_mxfp8(self): + """MoeDispatch and MoeCombine basic ops form an identity with MXFP8 transport.""" + self._run_dispatch_combine_identity(mxfp8=True) + + def _make_megamoe_model( + self, + *, + recipe, + accumulate_into_main_grad=False, + delay_wgrad_compute=False, + glu_interleave_size=None, + ): + """Build the exact five-op sequence recognized by MegaMoE fusion.""" + config = self._make_config(alignment=128 if recipe is not None else 0) + buffer = self._make_buffer_from_config( + config, + dispatch_fwd_quant_recipe=recipe, + combine_bwd_quant_recipe=recipe, + ) + dispatch = te_ops.MoeDispatch(config) + init_ctx = ( + te.quantized_model_init(enabled=True, recipe=recipe) + if recipe is not None + else nullcontext() + ) + previous_single_param = os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM") + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1" + try: + with init_ctx: + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * 256, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + activation = te_ops.ScaledSwiGLU( + glu_interleave_size=glu_interleave_size, + ) + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + 256, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + single_grouped_weight=True, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + finally: + if previous_single_param is None: + del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] + else: + os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = previous_single_param + combine = te_ops.MoeCombine(config) + + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=False) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=False) + fc1.set_extra_input_channel(0, "tokens_per_expert") + activation.set_extra_input_channel(0, "routing_weights") + fc2.set_extra_input_channel(0, "tokens_per_expert") + model = te_ops.Sequential(dispatch, fc1, activation, fc2, combine) + return model, fc1, fc2, buffer + + @_eager_test_include + def test_megamoe_bf16_numerics(self): + self._run_megamoe_vs_reference(quantization="bf16") + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_mxfp8_numerics(self): + self._run_megamoe_vs_reference(quantization="mxfp8") + + @_mxfp8_align_test + def test_megamoe_mxfp8_cuda_graph_matches_eager(self): + """Fused Sequential forward/backward graph replay matches eager execution.""" + if torch.cuda.get_device_capability() != (10, 7): + self.skipTest("FusedMoeEp CUDA graph test requires SM107") + if not _cudnn_megamoe_supported(): + self.skipTest("installed cuDNN frontend does not provide fixed training resources") + + recipe = MXFP8BlockScaling() + graph_model, graph_fc1, graph_fc2, graph_buffer = self._make_megamoe_model( + recipe=recipe, + glu_interleave_size=32, + ) + eager_model, eager_fc1, eager_fc2, eager_buffer = self._make_megamoe_model( + recipe=recipe, + glu_interleave_size=32, + ) + eager_model.load_state_dict(graph_model.state_dict()) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + static_tokens = tokens.detach().clone().requires_grad_(True) + static_topk_idx = topk_idx.detach().clone() + static_topk_weights = topk_weights.detach().clone().requires_grad_(True) + static_dy = torch.randn_like(static_tokens) + graph_op_kwargs = { + 0: {"buffer": graph_buffer}, + 4: {"buffer": graph_buffer}, + } + graphed_model = te.make_graphed_callables( + graph_model, + (static_tokens, static_topk_idx, static_topk_weights), + sample_kwargs={"op_kwargs": graph_op_kwargs}, + num_warmup_iters=3, + enabled=True, + recipe=recipe, + ) + + forward_ops = graph_model._module_groups[0]._forward_ops + backward_ops = graph_model._module_groups[0]._backward_ops + self.assertEqual(len(forward_ops), 1) + self.assertEqual(len(backward_ops), 1) + self.assertIsInstance(forward_ops[0][0], FusedMoeEp) + self.assertIs(backward_ops[0][0], forward_ops[0][0]) + + # Replace the capture-time contents while retaining captured addresses. + with torch.no_grad(): + static_tokens.copy_(torch.randn_like(static_tokens)) + static_topk_weights.copy_(torch.rand_like(static_topk_weights)) + static_dy.copy_(torch.randn_like(static_dy)) + + for parameter in graph_model.parameters(): + parameter.grad = torch.zeros_like(parameter) + if static_tokens.grad is not None: + static_tokens.grad.zero_() + if static_topk_weights.grad is not None: + static_topk_weights.grad.zero_() + with te.autocast(enabled=True, recipe=recipe): + graph_out = graphed_model( + static_tokens, + static_topk_idx, + static_topk_weights, + op_kwargs=graph_op_kwargs, + ) + graph_out.backward(static_dy) + torch.cuda.synchronize() + graph_results = ( + graph_out.detach().clone(), + static_tokens.grad.detach().clone(), + static_topk_weights.grad.detach().clone(), + graph_fc1.weight.grad.detach().clone(), + graph_fc2.weight.grad.detach().clone(), + ) + + eager_tokens = static_tokens.detach().clone().requires_grad_(True) + eager_topk_weights = static_topk_weights.detach().clone().requires_grad_(True) + for parameter in eager_model.parameters(): + parameter.grad = torch.zeros_like(parameter) + with te.autocast(enabled=True, recipe=recipe): + eager_out = eager_model( + eager_tokens, + static_topk_idx, + eager_topk_weights, + op_kwargs={ + 0: {"buffer": eager_buffer}, + 4: {"buffer": eager_buffer}, + }, + ) + eager_out.backward(static_dy) + torch.cuda.synchronize() + eager_results = ( + eager_out, + eager_tokens.grad, + eager_topk_weights.grad, + eager_fc1.weight.grad, + eager_fc2.weight.grad, + ) + + tolerances = {"rtol": 0.125, "atol": 0.25} + for graph_result, eager_result in zip(graph_results, eager_results): + torch.testing.assert_close(graph_result, eager_result, **tolerances) + + @_eager_test_include + def test_megamoe_main_grad_accumulation_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_main_grad_accumulation(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + ) + + @_eager_test_include + def test_megamoe_main_grad_overwrite_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + overwrite_main_grad=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_main_grad_overwrite(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + overwrite_main_grad=True, + ) + + @_eager_test_include + def test_megamoe_delayed_wgrad_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + delay_wgrad_compute=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_delayed_wgrad(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + delay_wgrad_compute=True, + ) + + @_eager_test_include + def test_megamoe_delayed_main_grad_bf16(self): + self._run_megamoe_vs_reference( + quantization="bf16", + accumulate_into_main_grad=True, + delay_wgrad_compute=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_delayed_main_grad(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + accumulate_into_main_grad=True, + delay_wgrad_compute=True, + ) + + @_eager_test_include + @_mxfp8_align_test + def test_megamoe_prequantized_input(self): + self._run_megamoe_vs_reference( + quantization="mxfp8", + prequantized_input=True, + ) + + def _run_megamoe_vs_reference( + self, + *, + quantization, + accumulate_into_main_grad=False, + overwrite_main_grad=False, + delay_wgrad_compute=False, + prequantized_input=False, + ): + """Compare the five-op MoE sequence with the PyTorch EP reference. + + The fuser selects MegaMoE when its runtime gates pass. Otherwise this + exercises the same sequence as separate NCCL EP and grouped-MLP ops. + """ + recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None + model, fc1, fc2, buffer = self._make_megamoe_model( + recipe=recipe, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for op in (fc1, fc2): + weights = op.weight.quantized_tensors + if weights is None: + weights = op.weight.split_into_quantized_tensors() + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + weights[expert].shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + weights[expert].copy_(weight) + if quantization == "mxfp8": + self.assertIsInstance(weights[expert], MXFP8Tensor) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + seq_tokens = tokens.detach().clone().requires_grad_(True) + seq_topk_weights = topk_weights.detach().clone().requires_grad_(True) + main_grad_sentinel = 0.5 + if accumulate_into_main_grad: + for op in (fc1, fc2): + op.weight.main_grad = torch.full( + op.weight.size(), + main_grad_sentinel, + dtype=torch.float32, + device=op.weight.device, + ) + op.weight.overwrite_main_grad = overwrite_main_grad + op.weight.zero_out_wgrad = False + op.weight.grad_added_to_main_grad = False + autocast_ctx = ( + te.autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() + ) + with autocast_ctx: + model_input = self._mxfp8_quantizer()(seq_tokens) if prequantized_input else seq_tokens + seq_out = model( + model_input, + topk_idx, + seq_topk_weights, + op_kwargs={ + 0: {"buffer": buffer}, + 4: {"buffer": buffer}, + }, + ) + + forward_ops = model._module_groups[0]._forward_ops + fused = len(forward_ops) == 1 and isinstance(forward_ops[0][0], FusedMoeEp) + if fused: + self.assertTrue(_cudnn_megamoe_supported()) + else: + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) + self.assertEqual(seq_out.dtype, torch.bfloat16) + + fc1_weight = _reference_weights(fc1) + fc2_weight = _reference_weights(fc2) + emulate_mxfp8 = fused or quantization == "mxfp8" + reference = MoeEpReference( + num_experts=self.cfg.num_experts, + hidden_size=HIDDEN_DIM, + intermediate_size=256, + top_k=TOP_K, + ep_group=self.ep_group, + max_tokens_per_rank=TOKENS_PER_RANK, + output_format=MoeFormat.BF16, + combine_format=( + MoeFormat.MXFP8 + if fused and os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0") == "1" + else MoeFormat.BF16 + ), + apply_topk_in_fc1=True, + generate_c=True, + intermediate_format=MoeFormat.MXFP8 if emulate_mxfp8 else None, + backward_operand_format=MoeFormat.MXFP8 if emulate_mxfp8 and not fused else None, + backward_wgrad_mode="operands" if fused else "none", + token_padding_size=256, + ) + if emulate_mxfp8 and not fused: + reference_activation = quantize_blockwise( + tokens.detach(), + MoeFormat.MXFP8, + axis=1, + ) + else: + reference_activation = tokens.detach() + if emulate_mxfp8 and not isinstance(fc1_weight, BlockScaledTensor): + fc1_weight = quantize_blockwise(fc1_weight, MoeFormat.MXFP8, axis=1) + if emulate_mxfp8 and not isinstance(fc2_weight, BlockScaledTensor): + fc2_weight = quantize_blockwise(fc2_weight, MoeFormat.MXFP8, axis=1) + reference_outputs = reference( + reference_activation, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + ) + if fused: + ref_out, fc1_c, route_metadata, wgrad_forward_stash = reference_outputs + else: + ref_out, fc1_c, route_metadata = reference_outputs + wgrad_forward_stash = None + + dy = ( + torch.randn( + seq_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + seq_out.backward(dy) + if delay_wgrad_compute: + fc1.backward_dw() + fc2.backward_dw() + reference_grads = reference.backward( + dy, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights.detach(), + fc1_c, + route_metadata, + wgrad_forward_stash=wgrad_forward_stash, + ) + if fused: + grad_tokens, grad_topk_weights, wgrad_operands = reference_grads + grad_fc1, grad_fc2 = wgrad_operands.dense_wgrads() + reference_wgrads = (grad_fc1, grad_fc2) + else: + grad_tokens, grad_topk_weights = reference_grads + reference_wgrads = (None, None) + + torch.cuda.synchronize() + tolerances = {"rtol": 0.125, "atol": 0.25} + torch.testing.assert_close(seq_out, ref_out, **tolerances) + torch.testing.assert_close( + seq_tokens.grad, + grad_tokens.to(dtype=seq_tokens.dtype), + **tolerances, + ) + torch.testing.assert_close(seq_topk_weights.grad, grad_topk_weights.float(), **tolerances) + for op, ref_grad in zip((fc1, fc2), reference_wgrads): + expected_grad = None if ref_grad is None else ref_grad.transpose(1, 2) + if accumulate_into_main_grad: + if expected_grad is not None and not overwrite_main_grad: + expected_grad = expected_grad + main_grad_sentinel + if expected_grad is not None: + torch.testing.assert_close( + op.weight.main_grad, + expected_grad.to(dtype=op.weight.main_grad.dtype), + **tolerances, + ) + else: + self.assertTrue(torch.isfinite(op.weight.main_grad).all()) + self.assertFalse(torch.all(op.weight.main_grad == main_grad_sentinel).item()) + self.assertTrue(op.weight.grad_added_to_main_grad) + self.assertIsNotNone(op.weight.grad) + continue + seq_grad = op.weight.grad + self.assertEqual(seq_grad.dtype, torch.bfloat16) + self.assertEqual( + tuple(seq_grad.shape), + (NUM_LOCAL_EXPERTS, op.out_features, op.in_features), + ) + self.assertTrue(seq_grad.is_contiguous()) + self.assertTrue(torch.isfinite(seq_grad).all()) + if expected_grad is not None: + torch.testing.assert_close( + seq_grad, + expected_grad.to(dtype=seq_grad.dtype), + **tolerances, + ) + + def _init_distributed(): dist.init_process_group(backend="nccl") torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) @@ -846,7 +1579,12 @@ def _init_distributed(): name_filter = os.environ.get("NVTE_EP_TEST_FILTER") if name_filter: loader.testMethodPrefix = name_filter - suite = loader.loadTestsFromTestCase(TestEP) + suite = unittest.TestSuite( + ( + loader.loadTestsFromTestCase(TestEP), + loader.loadTestsFromTestCase(TestMoeEpSequential), + ) + ) runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2) result = runner.run(suite) dist.barrier() diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 4e22d97ead..d78a5e3d18 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -286,6 +286,36 @@ def test_basic_construction_all_same_shape(self) -> None: assert grouped_tensor.get_common_last_dim() == 512 assert grouped_tensor.has_data() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_module_to_moves_grouped_parameter_storage(self) -> None: + """Module.to preserves GroupedTensor parameters initialized on CPU.""" + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=2, + shapes=[(4, 8), (4, 8)], + quantizer=None, + device="cpu", + dtype=torch.float32, + ) + values = torch.arange(grouped_tensor.numel(), dtype=torch.float32) + grouped_tensor.rowwise_data.copy_(values) + + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(grouped_tensor)) + original_parameter = module.weight + expected = values.to(device="cuda", dtype=torch.bfloat16) + + module.to(device="cuda", dtype=torch.bfloat16) + + assert module.weight is original_parameter + assert isinstance(module.weight, GroupedTensor) + assert module.weight.device.type == "cuda" + assert module.weight.dtype == torch.bfloat16 + assert module.weight.rowwise_data.device.type == "cuda" + assert module.weight.rowwise_data.dtype == torch.bfloat16 + torch.testing.assert_close(module.weight.rowwise_data, expected, rtol=0, atol=0) + members = module.weight.split_into_quantized_tensors() + assert all(member.device.type == "cuda" for member in members) + def test_basic_construction_varying_first_dim(self) -> None: """Test GroupedTensor construction with varying first dimension""" num_tensors = 3 diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 4464036fc7..34a3b0df68 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -17,16 +17,21 @@ from .cpu_offload import mark_not_offload from .distributed import symm_mem_alloc, release_symm_mem_pool -from .quantized_tensor import QuantizedTensor +from .quantized_tensor import QuantizedTensor, QuantizedTensorStorage +from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage # Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of # common.recipe (the concrete recipe classes are imported lazily where used). if TYPE_CHECKING: from ..common.recipe import Recipe + from .quantized_tensor import Quantizer __all__ = [ + "EpConfig", "EpBuffer", "ep_bootstrap", + "get_ep_drop_on_overflow", + "get_ep_group", "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", @@ -81,11 +86,28 @@ def _check_nccl_runtime_version() -> None: # omitted); ep_dispatch reads it to size the recv outputs from the per-step # recv-token total instead of a fixed recv_capacity_per_rank. _EAGER = False +_BOOTSTRAP_SETTINGS: Optional[dict[str, object]] = None + + +@dataclass(frozen=True, slots=True) +class EpConfig: + """Immutable configuration shared by EP MoE operations.""" + + top_k: int + hidden_dim: int + num_local_experts: int + max_tokens_per_rank: int + recv_capacity_per_rank: Optional[int] + ep_group: dist.ProcessGroup + alignment: int = 0 + payload_dtype: torch.dtype = torch.bfloat16 + zero_copy: bool = False + drop_on_overflow: bool = False def _atexit_finalize() -> None: """Best-effort teardown at interpreter shutdown; swallows errors.""" - global _BOOTSTRAPPED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS if _BOOTSTRAPPED: try: tex.ep_finalize() @@ -97,6 +119,7 @@ def _atexit_finalize() -> None: _BOOTSTRAPPED = False _EP_GROUP = None _EAGER = False + _BOOTSTRAP_SETTINGS = None def ep_bootstrap( @@ -132,7 +155,7 @@ def ep_bootstrap( ``drop_on_overflow`` drops tokens exceeding ``recv_capacity_per_rank`` instead of trapping. Requires ``recv_capacity_per_rank``. """ - global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS eager = recv_capacity_per_rank is None if _BOOTSTRAPPED: raise RuntimeError("ep_bootstrap was already called in this process") @@ -179,6 +202,18 @@ def ep_bootstrap( _BOOTSTRAPPED = True _EP_GROUP = ep_group _EAGER = bool(eager) + _BOOTSTRAP_SETTINGS = { + "top_k": int(num_topk), + "hidden_dim": int(hidden_dim), + "num_local_experts": int(num_experts) // ep_group.size(), + "max_tokens_per_rank": int(max_tokens_per_rank), + "recv_capacity_per_rank": ( + None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) + ), + "payload_dtype": max_token_dtype, + "zero_copy": bool(zero_copy), + "drop_on_overflow": bool(drop_on_overflow), + } if not _ATEXIT_REGISTERED: atexit.register(_atexit_finalize) _ATEXIT_REGISTERED = True @@ -189,6 +224,23 @@ def is_ep_bootstrapped() -> bool: return _BOOTSTRAPPED +def get_ep_group() -> Optional[dist.ProcessGroup]: + """Return the process group registered by :func:`ep_bootstrap`.""" + return _EP_GROUP + + +def get_ep_drop_on_overflow() -> Optional[bool]: + """Return the overflow policy registered by :func:`ep_bootstrap`.""" + if _BOOTSTRAP_SETTINGS is None: + return None + return bool(_BOOTSTRAP_SETTINGS["drop_on_overflow"]) + + +def _ep_is_eager() -> bool: + """Return whether the bootstrapped EP group uses variable-size eager buffers.""" + return _EAGER + + def ep_finalize() -> None: """Optional explicit EP teardown; idempotent. @@ -199,7 +251,7 @@ def ep_finalize() -> None: a caller that used ``symm_mem_alloc(use_pool=True)`` does not need a separate ``release_symm_mem_pool()`` before destroying the PG. """ - global _BOOTSTRAPPED, _EP_GROUP, _EAGER + global _BOOTSTRAPPED, _EP_GROUP, _EAGER, _BOOTSTRAP_SETTINGS if not _BOOTSTRAPPED: return try: @@ -210,6 +262,7 @@ def ep_finalize() -> None: _BOOTSTRAPPED = False _EP_GROUP = None _EAGER = False + _BOOTSTRAP_SETTINGS = None def is_symm_backed(t: torch.Tensor) -> bool: @@ -237,6 +290,7 @@ def is_symm_backed(t: torch.Tensor) -> bool: class EpBuffer: """Per-microbatch EP layer state: handle_mem, tokens_per_expert, and shape/dtype config. + Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). """ @@ -251,6 +305,7 @@ class EpBuffer: "payload_dtype", "device", "tokens_per_expert", + "num_local_tokens", "zero_copy", "eager", "total_recv_tokens", @@ -303,6 +358,7 @@ def __init__( self.tokens_per_expert = torch.empty( self.num_local_experts, dtype=torch.int64, device=device ) + self.num_local_tokens = self.max_tokens_per_rank # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) # Per-step recv-token total (int64 [1]), written by ep_prepare. Eager reads it @@ -545,7 +601,7 @@ class _DispatchState: def _ep_prepare_and_dispatch_fwd( - tokens: torch.Tensor, + tokens: torch.Tensor | MXFP8TensorStorage, topk_weights: torch.Tensor, topk_idx: torch.Tensor, buffer: "EpBuffer", @@ -566,7 +622,7 @@ def _ep_prepare_and_dispatch_fwd( num_recv_tokens = buffer.recv_capacity_per_rank payload_dtype = buffer.payload_dtype is_scaled = tokens_scale_inv is not None - tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens + tokens_data = tokens._rowwise_data if isinstance(tokens, MXFP8TensorStorage) else tokens assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]" hidden = tokens_data.shape[-1] if is_scaled and tokens._fp8_dtype != tex.DType.kFloat8E4M3: @@ -644,7 +700,7 @@ def _ep_prepare_and_dispatch_fwd( recv_scale_inv, tokens_per_expert, tokens._fp8_dtype, - tokens.dtype, + tokens._dtype, ) return recv_out, recv_topk_weights, state return recv_tokens, recv_topk_weights, state @@ -694,10 +750,10 @@ class _EpPrepareAndDispatch(torch.autograd.Function): recv-count, so no Python runs between the count read and the dispatch launch; caller-supplied buffers and zero-copy are then forbidden. Otherwise the recv outputs are allocated here to the static recv capacity (caller-supplied or symm-mem-backed under zero-copy) and passed in. When - ``tokens_scale_inv`` is set (MXFP8 for now), ``tokens`` is the quantized tensor kept as the - autograd operand so grad reaches the pre-quant input, and recv is returned as a per-expert - GroupedTensor. The compute lives in ``_ep_prepare_and_dispatch_fwd`` / ``_ep_dispatch_bwd``; - this wrapper only bridges autograd context handling.""" + When MXFP8 is configured, forward quantizes ``tokens`` to lightweight storage locally while + keeping the high-precision tensor as the autograd operand. The compute lives in + ``_ep_prepare_and_dispatch_fwd`` / ``_ep_dispatch_bwd``; this wrapper only bridges autograd + context handling.""" @staticmethod def forward( # type: ignore[override] @@ -708,12 +764,29 @@ def forward( # type: ignore[override] buffer: "EpBuffer", recv_tokens: Optional[torch.Tensor] = None, recv_topk_weights: Optional[torch.Tensor] = None, - tokens_scale_inv: Optional[torch.Tensor] = None, ): """Only tokens and topk_weights are differentiable, so the non-diff buffer tensors ride on the buffer object to keep the autograd operand list short.""" + tokens_scale_inv = None + if buffer.dispatch_fwd_quant_recipe is not None: + from .tensor.mxfp8_tensor import MXFP8Quantizer + + # Only MXFP8 Quantizer is supported for EP dispatch + quantizer = MXFP8Quantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantizer.internal = True + tokens, tokens_scale_inv = quantize_for_ep(tokens, quantizer) recv_out, recv_topk_weights, state = _ep_prepare_and_dispatch_fwd( - tokens, topk_weights, topk_idx, buffer, recv_tokens, recv_topk_weights, tokens_scale_inv + tokens, + topk_weights, + topk_idx, + buffer, + recv_tokens, + recv_topk_weights, + tokens_scale_inv, ) ctx.state = state # Detach so the long-lived buffers aren't tracked as differentiable outputs; autograd @@ -739,7 +812,6 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights): # type: ignore[override] None, # buffer None, # recv_tokens None, # recv_topk_weights - None, # tokens_scale_inv ) @@ -763,17 +835,21 @@ class _CombineState: def _ep_combine_fwd( expert_out: torch.Tensor, grad_out: Optional[torch.Tensor], - buffer: "EpBuffer", + *, + handle_mem: torch.Tensor, + token_counts: torch.Tensor, num_local_tokens: int, + hidden_dim: int, bwd_quant_recipe, + eager: bool, + zero_copy: bool, ): """Run combine and return ``(result, _CombineState)``. Eager mode is not graph-capturable, so it calls the backend op directly and skips the torch.library dispatch. No autograd; the caller owns context handling.""" - handle_mem = buffer.handle_mem - eager = buffer.eager + device = expert_out.device - result = torch.empty(num_local_tokens, buffer.hidden_dim, dtype=expert_out.dtype, device=device) + result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device) if eager: tex.ep_combine(handle_mem, expert_out, result) else: @@ -782,25 +858,30 @@ def _ep_combine_fwd( handle_mem=handle_mem, grad_out=grad_out, bwd_quant_recipe=bwd_quant_recipe, - token_counts=buffer.tokens_per_expert, + token_counts=token_counts, expert_out_shape=expert_out.shape, expert_out_dtype=expert_out.dtype, device=device, eager=eager, - zero_copy=buffer.zero_copy, + zero_copy=zero_copy, ) return result, state -def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): +def _ep_combine_bwd( + state: "_CombineState", + g_result: torch.Tensor, + quantized_grad: Optional[QuantizedTensorStorage] = None, + grad_scale_inv: Optional[torch.Tensor] = None, +): """Scatter the result-grad to expert positions and return the expert_out grad. High-precision sends the grad as-is; a quantized recipe (MXFP8 today) quantizes it and returns a per-expert - GroupedTensor. No autograd; the caller owns context handling.""" - if not g_result.is_contiguous(): + GroupedTensor. Optionally pre-mxfp8-quantized grad and scales can be provided.""" + if quantized_grad is None and not g_result.is_contiguous(): g_result = g_result.contiguous() handle_mem = state.handle_mem - if state.bwd_quant_recipe is None: + if state.bwd_quant_recipe is None and quantized_grad is None: grad_expert_out = state.grad_out if grad_expert_out is None: grad_expert_out = _alloc_io( @@ -811,16 +892,30 @@ def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): else: torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out) else: - mx, g_scale_inv = _quantize_mxfp8(g_result) + if quantized_grad is None: + from .tensor.mxfp8_tensor import MXFP8Quantizer + + # Only MXFP8 Quantizer is supported for EP combine bwd + quantizer = MXFP8Quantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + quantizer.internal = True + mx, grad_scale_inv = quantize_for_ep(g_result, quantizer) + else: + mx = quantized_grad + if grad_scale_inv is None: + raise ValueError("MXFP8 combine backward requires compact rowwise scales.") g_data = mx._rowwise_data recv_pr, hidden = state.expert_out_shape[0], state.expert_out_shape[-1] ge_data, ge_scale_inv = _scale_alloc_io( state.grad_out, recv_pr, hidden, - g_scale_inv.shape[-1], + grad_scale_inv.shape[-1], g_data.dtype, - g_scale_inv.dtype, + grad_scale_inv.dtype, state.device, state.zero_copy, ) @@ -828,10 +923,10 @@ def _ep_combine_bwd(state: "_CombineState", g_result: torch.Tensor): g_data_fp8 = g_data.view(torch.float8_e4m3fn) ge_data_fp8 = ge_data.view(torch.float8_e4m3fn) if state.eager: - tex.ep_combine_bwd(handle_mem, g_data_fp8, ge_data_fp8, g_scale_inv, ge_scale_inv) + tex.ep_combine_bwd(handle_mem, g_data_fp8, ge_data_fp8, grad_scale_inv, ge_scale_inv) else: torch.ops.transformer_engine_ep.combine_bwd( - handle_mem, g_data_fp8, ge_data_fp8, g_scale_inv, ge_scale_inv + handle_mem, g_data_fp8, ge_data_fp8, grad_scale_inv, ge_scale_inv ) grad_expert_out = _make_grouped_mxfp8( ge_data, ge_scale_inv, state.token_counts, mx._fp8_dtype, state.expert_out_dtype @@ -861,7 +956,15 @@ def forward( # type: ignore[override] """Combine fwd; stashes the backward state on ctx. When ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8.""" result, ctx.state = _ep_combine_fwd( - expert_out, grad_out, buffer, num_local_tokens, bwd_quant_recipe + expert_out, + grad_out, + handle_mem=buffer.handle_mem, + token_counts=buffer.tokens_per_expert, + num_local_tokens=num_local_tokens, + hidden_dim=buffer.hidden_dim, + bwd_quant_recipe=bwd_quant_recipe, + eager=buffer.eager, + zero_copy=buffer.zero_copy, ) return result @@ -910,31 +1013,54 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso return torch.empty(*shape, dtype=dtype, device=device) -def _quantize_mxfp8(x: torch.Tensor): - """Quantize a high-precision tensor to MXFP8 and return ``(quantized_tensor, scale_inv)`` where - ``scale_inv`` is the compact ``[T, H/block]`` scale the EP backend routes. The quantized tensor - is returned so callers can keep it as the autograd operand; its ``_rowwise_data`` is the fp8 - payload and ``scale_inv.shape[-1]`` the scale-column count. EP routes and returns E4M3 data in - both directions, so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to - the compact ``[T, H/block]`` layout; requires a 16-byte-aligned scale row.""" - from .constants import MXFP8_BLOCK_SCALING_SIZE +def quantize_for_ep( + input_: torch.Tensor | QuantizedTensorStorage, + quantizer: Optional["Quantizer"], +) -> tuple[MXFP8TensorStorage, torch.Tensor]: + """Return E4M3 MXFP8 storage and compact rowwise scales for EP transport. + + High-precision input is quantized with ``quantizer``; existing MXFP8 storage is accepted + as-is. The returned storage owns the FP8 payload in ``_rowwise_data``, while ``scale_inv`` has + the compact ``[T, H/block]`` layout routed by the EP backend. EP transports E4M3 in both + directions, so other FP8 formats are rejected. GEMM scale-row padding is stripped, and each + compact scale row must remain contiguous and 16-byte aligned. + """ + from .constants import DType, MXFP8_BLOCK_SCALING_SIZE from .tensor.mxfp8_tensor import MXFP8Quantizer - mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x) - if mx._with_gemm_swizzled_scales: - raise RuntimeError( - "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact." - ) - data = mx._rowwise_data - scale_inv = mx._rowwise_scale_inv + if quantizer is not None: + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + f"EP MXFP8 transport requires MXFP8Quantizer, got {type(quantizer).__name__}." + ) + if quantizer.dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") + + if isinstance(input_, MXFP8TensorStorage): + quantized = input_ + elif isinstance(input_, QuantizedTensorStorage): + raise TypeError(f"EP MXFP8 transport requires an MXFP8 input, got {type(input_).__name__}.") + else: + if quantizer is None: + raise ValueError("An MXFP8 quantizer is required for a non-quantized EP input.") + if not quantizer.internal: + quantizer = quantizer.copy() + quantizer.internal = True + quantized = quantizer(input_) + + if quantized._fp8_dtype != DType.kFloat8E4M3: + raise NotImplementedError("EP MXFP8 transport supports E4M3 only.") + if quantized._with_gemm_swizzled_scales: + raise ValueError("EP requires unswizzled MXFP8 scales.") + data = quantized._rowwise_data + scale_inv = quantized._rowwise_scale_inv if data is None or scale_inv is None: - raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.") - t_flat = x.shape[0] - hidden = x.shape[-1] - cols = hidden // MXFP8_BLOCK_SCALING_SIZE + raise ValueError("EP requires rowwise MXFP8 data and scales.") + t_flat, hidden = input_.shape + scale_cols = hidden // MXFP8_BLOCK_SCALING_SIZE # The backend forwards each token's scale row with a 16-byte-aligned store, so the row # (cols * dtype bytes) must be a multiple of 16. - scale_row_bytes = cols * scale_inv.element_size() + scale_row_bytes = scale_cols * scale_inv.element_size() if scale_row_bytes % 16 != 0: raise ValueError( f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives " @@ -944,13 +1070,10 @@ def _quantize_mxfp8(x: torch.Tensor): # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block] # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the # slice stays contiguous; assert rather than force a copy. - scale_inv = scale_inv[:t_flat, :cols] + scale_inv = scale_inv[:t_flat, :scale_cols] if not scale_inv.is_contiguous(): - raise ValueError( - "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a " - f"non-contiguous [{t_flat}, {cols}] slice." - ) - return mx, scale_inv + raise ValueError("EP requires compact contiguous MXFP8 scales.") + return quantized, scale_inv def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy): @@ -1049,9 +1172,6 @@ def ep_dispatch( "and cannot use caller-supplied recv_tokens / recv_topk_weights" ) - # Quantize up front (before prepare) so the quant kernels overlap the eager count sync and the - # quantized tensor stays the autograd operand; grad reaches the pre-quant input. - tokens_scale_inv = None if buffer.dispatch_fwd_quant_recipe is not None: from ..common.recipe import MXFP8BlockScaling @@ -1060,7 +1180,6 @@ def ep_dispatch( "EP block-scaled dispatch supports MXFP8BlockScaling only; got " f"{type(buffer.dispatch_fwd_quant_recipe).__name__}." ) - tokens, tokens_scale_inv = _quantize_mxfp8(tokens) # Fused prepare + dispatch in one C++ op. Eager sizes the recv outputs from the per-step host # recv-count (allocated in C++, no caller buffers); non-eager sizes them to the static recv @@ -1072,7 +1191,6 @@ def ep_dispatch( buffer, recv_tokens, recv_topk_weights, - tokens_scale_inv, ) return recv_out, recv_topk_weights, buffer.tokens_per_expert diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py new file mode 100644 index 0000000000..8614b762d2 --- /dev/null +++ b/transformer_engine/pytorch/ep_reference.py @@ -0,0 +1,1204 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism. + +The implementation deliberately favors readable semantics over performance. It +supports a one-rank execution path and a variable-size ``all_to_all_single`` EP +path, plus BF16, MXFP8, and NVFP4 block-scaled public outputs. + +The quantized tensor layouts are logical (unswizzled) layouts. A production +kernel may reorder scale factors internally without changing this API contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +class MoeFormat(str, Enum): + """Public and communication formats supported by the reference.""" + + BF16 = "bf16" + MXFP8 = "mxfp8" + NVFP4 = "nvfp4" + + +def _parse_format(value: Union[MoeFormat, str]) -> MoeFormat: + if isinstance(value, MoeFormat): + return value + try: + return MoeFormat(value.lower()) + except (AttributeError, ValueError) as exc: + choices = ", ".join(item.value for item in MoeFormat) + raise ValueError(f"unsupported format {value!r}; expected one of: {choices}") from exc + + +def _require_torch_dtype(name: str) -> torch.dtype: + dtype = getattr(torch, name, None) + if dtype is None: + raise RuntimeError(f"this PyTorch build does not provide torch.{name}") + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of range for a {ndim}-D tensor") + return normalized + + +def _shape_with_axis(shape: Sequence[int], axis: int, value: int) -> Tuple[int, ...]: + result = list(shape) + result[axis] = value + return tuple(result) + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator + + +@dataclass(frozen=True) +class BlockScaledTensor: + """Portable data-plus-scale representation for MXFP8 or NVFP4. + + ``logical_shape`` describes the dequantized tensor. For MXFP8, ``data`` + has that shape and uses E4M3. For NVFP4, ``data`` is a uint8 tensor with + two E2M1 values per byte along ``axis`` (low nibble first). ``scale`` + replaces that axis by one scale per block. + """ + + data: torch.Tensor + scale: torch.Tensor + format: Union[MoeFormat, str] + logical_shape: Tuple[int, ...] + axis: int = -1 + + def __post_init__(self) -> None: + fmt = _parse_format(self.format) + if fmt is MoeFormat.BF16: + raise ValueError("BlockScaledTensor only represents mxfp8 or nvfp4") + shape = tuple(int(dim) for dim in self.logical_shape) + if not shape or any(dim < 0 for dim in shape): + raise ValueError(f"logical_shape must contain non-negative dimensions, got {shape}") + axis = _normalize_axis(self.axis, len(shape)) + object.__setattr__(self, "format", fmt) + object.__setattr__(self, "logical_shape", shape) + object.__setattr__(self, "axis", axis) + self._validate_storage() + + @property + def block_size(self) -> int: + return 32 if self.format is MoeFormat.MXFP8 else 16 + + @property + def shape(self) -> Tuple[int, ...]: + return self.logical_shape + + @property + def device(self) -> torch.device: + return self.data.device + + def _validate_storage(self) -> None: + if self.data.device != self.scale.device: + raise ValueError("block-scaled data and scale must be on the same device") + + logical_extent = self.logical_shape[self.axis] + scale_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, self.block_size), + ) + if tuple(self.scale.shape) != scale_shape: + raise ValueError(f"scale shape must be {scale_shape}, got {tuple(self.scale.shape)}") + + if self.format is MoeFormat.MXFP8: + expected_dtype = _require_torch_dtype("float8_e4m3fn") + expected_scale_dtype = _require_torch_dtype("float8_e8m0fnu") + data_shape = self.logical_shape + if self.data.dtype != expected_dtype: + raise TypeError( + f"mxfp8 data must have dtype {expected_dtype}, got {self.data.dtype}" + ) + else: + expected_scale_dtype = _require_torch_dtype("float8_e4m3fn") + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if self.data.dtype != torch.uint8 and self.data.dtype != fp4_dtype: + raise TypeError("nvfp4 data must be packed uint8 or torch.float4_e2m1fn_x2") + data_shape = _shape_with_axis( + self.logical_shape, + self.axis, + _ceil_div(logical_extent, 2), + ) + + if tuple(self.data.shape) != data_shape: + raise ValueError(f"data shape must be {data_shape}, got {tuple(self.data.shape)}") + if self.scale.dtype != expected_scale_dtype: + raise TypeError(f"scale must have dtype {expected_scale_dtype}, got {self.scale.dtype}") + + def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """Return the logical tensor with block scales applied.""" + + logical_extent = self.logical_shape[self.axis] + scale = self.scale.movedim(self.axis, -1).float() + expanded_scale = scale.repeat_interleave(self.block_size, dim=-1)[..., :logical_extent] + + if self.format is MoeFormat.MXFP8: + values = self.data.movedim(self.axis, -1).float() + else: + packed = self.data + if packed.dtype != torch.uint8: + packed = packed.view(torch.uint8) + packed = packed.movedim(self.axis, -1) + low = packed & 0x0F + high = packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2)[..., :logical_extent] + table = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + values = table[codes.long()] + + return (values * expanded_scale).movedim(-1, self.axis).to(dtype) + + +def _nearest_e2m1_codes(values: torch.Tensor) -> torch.Tensor: + """Quantize to E2M1 nibble codes with round-to-nearest, ties-to-even.""" + + levels = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, + device=values.device, + ) + magnitudes = values.abs().unsqueeze(-1) + distances = (magnitudes - levels).abs() + minimum = distances.amin(dim=-1, keepdim=True) + candidates = distances == minimum + codes = torch.arange(8, dtype=torch.int64, device=values.device) + any_code = torch.where(candidates, codes, 8).amin(dim=-1) + even_code = torch.where(candidates & ((codes & 1) == 0), codes, 8).amin(dim=-1) + magnitude_code = torch.where(even_code < 8, even_code, any_code) + sign_code = torch.signbit(values).to(torch.int64) << 3 + return magnitude_code | sign_code + + +def quantize_blockwise( + tensor: torch.Tensor, + format: Union[MoeFormat, str], + *, + axis: int = -1, +) -> BlockScaledTensor: + """Quantize a floating tensor into logical MXFP8 or NVFP4 blocks. + + MXFP8 uses 32-value blocks, E4M3 payloads, and E8M0 scales rounded toward + positive infinity. NVFP4 uses 16-value blocks, packed E2M1 payloads, and + E4M3 scales rounded to nearest. + """ + + fmt = _parse_format(format) + if fmt is MoeFormat.BF16: + raise ValueError("quantize_blockwise requires mxfp8 or nvfp4") + if not tensor.is_floating_point(): + raise TypeError(f"tensor must be floating point, got {tensor.dtype}") + + axis = _normalize_axis(axis, tensor.ndim) + logical_shape = tuple(tensor.shape) + moved = tensor.float().movedim(axis, -1) + logical_extent = moved.shape[-1] + block_size = 32 if fmt is MoeFormat.MXFP8 else 16 + block_count = _ceil_div(logical_extent, block_size) + padded_extent = block_count * block_size + if padded_extent != logical_extent: + moved = F.pad(moved, (0, padded_extent - logical_extent)) + blocks = moved.reshape(*moved.shape[:-1], block_count, block_size) + + value_limit = 448.0 if fmt is MoeFormat.MXFP8 else 6.0 + scale_float = blocks.abs().amax(dim=-1) / value_limit + if fmt is MoeFormat.MXFP8: + safe_scale = torch.where(scale_float > 0, scale_float, 1.0) + scale_float = torch.where( + scale_float > 0, + torch.pow(2.0, torch.ceil(torch.log2(safe_scale))), + torch.zeros_like(scale_float), + ) + scale_dtype = _require_torch_dtype("float8_e8m0fnu") + else: + scale_dtype = _require_torch_dtype("float8_e4m3fn") + + scale = scale_float.to(scale_dtype) + scale_for_math = scale.float() + reciprocal = torch.where(scale_for_math > 0, scale_for_math.reciprocal(), 0.0) + normalized = (blocks * reciprocal.unsqueeze(-1)).clamp(-value_limit, value_limit) + + if fmt is MoeFormat.MXFP8: + data_dtype = _require_torch_dtype("float8_e4m3fn") + data = normalized.to(data_dtype).reshape(*moved.shape)[..., :logical_extent] + else: + codes = _nearest_e2m1_codes(normalized).reshape(*moved.shape) + low = codes[..., 0::2] + high = codes[..., 1::2] + data = (low | (high << 4)).to(torch.uint8)[..., : _ceil_div(logical_extent, 2)] + + return BlockScaledTensor( + data=data.movedim(-1, axis).contiguous(), + scale=scale.movedim(-1, axis).contiguous(), + format=fmt, + logical_shape=logical_shape, + axis=axis, + ) + + +MoeTensor = Union[torch.Tensor, BlockScaledTensor] + + +@dataclass(frozen=True) +class WgradForwardStashReference: + """Logical reference for the caller-owned forward wgrad stash. + + Unlike the production object, ``fc1_a`` bundles its logical E8M0 scales + with the E4M3 payload. It represents the padded, expert-concatenated + ``x.T`` operand after input MXFP8 staging and token-axis requantization. + """ + + fc1_a: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + +@dataclass(frozen=True) +class WgradOperandsReference: + """Logical MXFP8 operands and dense expert-weight-gradient oracle. + + The K dimension is a concatenation of local experts. Each expert's valid + routes come first, followed by zero rows up to its 256-route boundary. + Production scale tensors use a blocked physical layout; these reference + tensors keep ordinary logical scales so their represented values are easy + to inspect. + """ + + fc1_a: BlockScaledTensor + fc1_b: BlockScaledTensor + fc2_a: BlockScaledTensor + fc2_b: BlockScaledTensor + expert_offsets: torch.Tensor + valid_route_counts: torch.Tensor + route_metadata: torch.Tensor + + def dense_wgrads(self) -> Tuple[torch.Tensor, torch.Tensor]: + """Return dense ``dW1=x.T@dC`` and ``dW2=h.T@(p*dY)`` per expert.""" + + a1 = self.fc1_a.dequantize() + b1 = self.fc1_b.dequantize() + a2 = self.fc2_a.dequantize() + b2 = self.fc2_b.dequantize() + expert_count = int(self.expert_offsets.numel()) + dw1 = torch.zeros( + (expert_count, a1.shape[0], b1.shape[1]), + dtype=torch.float32, + device=a1.device, + ) + dw2 = torch.zeros( + (expert_count, a2.shape[0], b2.shape[1]), + dtype=torch.float32, + device=a2.device, + ) + begin = 0 + for expert, end_tensor in enumerate(self.expert_offsets): + end = int(end_tensor.item()) + if end > begin: + dw1[expert] = a1[:, begin:end] @ b1[begin:end] + dw2[expert] = a2[:, begin:end] @ b2[begin:end] + begin = end + return dw1, dw2 + + +@dataclass(frozen=True) +class _DispatchPlan: + """Send-side routing derived from ``topk_idx``; identical in fwd and bwd.""" + + send_expert: torch.Tensor # local expert id per sent route + send_weight: torch.Tensor # router weight per sent route + send_token_idx: torch.Tensor # source token per sent route + send_slot_idx: torch.Tensor # source top-k slot per sent route + send_counts: Tuple[int, ...] # routes sent to each rank + recv_counts: Tuple[int, ...] # routes received from each rank + + +def _tensor_device(tensor: MoeTensor) -> torch.device: + return tensor.device + + +def _decode_tensor( + tensor: MoeTensor, + *, + name: str, + expected_shape: Tuple[int, ...], + quantized_axis: int, +) -> torch.Tensor: + if isinstance(tensor, BlockScaledTensor): + if tensor.logical_shape != expected_shape: + raise ValueError( + f"{name} logical shape must be {expected_shape}, got {tensor.logical_shape}" + ) + if tensor.axis != _normalize_axis(quantized_axis, len(expected_shape)): + raise ValueError(f"{name} must be block-scaled along axis {quantized_axis}") + return tensor.dequantize() + + if tuple(tensor.shape) != expected_shape: + raise ValueError(f"{name} shape must be {expected_shape}, got {tuple(tensor.shape)}") + if not tensor.is_floating_point(): + raise TypeError(f"{name} must be floating point or BlockScaledTensor, got {tensor.dtype}") + return tensor.float() + + +def _format_round_trip_axis( + tensor: torch.Tensor, + format: MoeFormat, + *, + axis: int, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).float() + return quantize_blockwise(tensor, format, axis=axis).dequantize() + + +def _format_round_trip(tensor: torch.Tensor, format: MoeFormat) -> torch.Tensor: + return _format_round_trip_axis(tensor, format, axis=-1) + + +def forward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model GLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def backward_combine_round_trip( + tensor: torch.Tensor, + format: MoeFormat, +) -> torch.Tensor: + """Model dGLU combine conversion directly from its FP32 accumulator.""" + + return _format_round_trip(tensor, format) + + +def _padded_expert_rows( + rows: torch.Tensor, + expert_rows: torch.Tensor, + valid_counts: Sequence[int], + padded_ends: Sequence[int], +) -> torch.Tensor: + """Place compact expert-grouped rows at the start of padded ranges.""" + + padded_extent = int(padded_ends[-1]) if padded_ends else 0 + padded = torch.zeros( + (padded_extent, *rows.shape[1:]), + dtype=rows.dtype, + device=rows.device, + ) + begin = 0 + for expert, (count, end) in enumerate(zip(valid_counts, padded_ends)): + positions = torch.nonzero( + expert_rows == expert, + as_tuple=False, + ).flatten() + if int(positions.numel()) != int(count): + raise ValueError(f"expert {expert} has {positions.numel()} rows, expected {count}") + if count: + padded[begin : begin + count].copy_(rows.index_select(0, positions)) + begin = int(end) + return padded + + +class MoeEpReference: + """Reference implementation of routed SwiGLU experts plus EP dispatch. + + Global experts are assigned contiguously: rank ``r`` owns + ``[r * experts_per_rank, (r + 1) * experts_per_rank)``. Pass an explicit + initialized process group for multi-rank execution; ``None`` means a + one-rank reference even if the default distributed group is initialized. + + ``intermediate_format`` optionally applies a post-SwiGLU, pre-FC2 format + round trip to model fused kernels that materialize their FC2 input in low + precision. ``None`` preserves the raw mathematical reference semantics. + ``backward_operand_format`` additionally models dGLU staging of grad-output + and transposed weights along their backward reduction dimensions. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + ep_group: Optional[dist.ProcessGroup] = None, + max_tokens_per_rank: Optional[int] = None, + output_format: Union[MoeFormat, str] = MoeFormat.BF16, + combine_format: Union[MoeFormat, str] = MoeFormat.BF16, + intermediate_format: Optional[Union[MoeFormat, str]] = None, + backward_operand_format: Optional[Union[MoeFormat, str]] = None, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + backward_wgrad_mode: str = "none", + token_padding_size: int = 128, + ) -> None: + for name, value in ( + ("num_experts", num_experts), + ("hidden_size", hidden_size), + ("intermediate_size", intermediate_size), + ("top_k", top_k), + ): + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if top_k > num_experts: + raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})") + if max_tokens_per_rank is not None and max_tokens_per_rank < 0: + raise ValueError("max_tokens_per_rank must be non-negative") + if backward_wgrad_mode not in ("none", "operands"): + raise ValueError("backward_wgrad_mode must be 'none' or 'operands'") + if backward_wgrad_mode == "operands" and not generate_c: + raise ValueError("backward_wgrad_mode='operands' requires generate_c=True") + if not isinstance(token_padding_size, int) or token_padding_size <= 0: + raise ValueError("token_padding_size must be a positive integer") + if backward_wgrad_mode == "operands" and token_padding_size != 256: + raise ValueError("backward_wgrad_mode='operands' requires token_padding_size=256") + + if ep_group is None: + ep_size, ep_rank = 1, 0 + else: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError( + "ep_group requires an initialized torch.distributed process group" + ) + ep_size = dist.get_world_size(ep_group) + ep_rank = dist.get_rank(ep_group) + if num_experts % ep_size != 0: + raise ValueError( + f"num_experts ({num_experts}) must be divisible by EP size ({ep_size})" + ) + + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.ep_group = ep_group + self.ep_size = ep_size + self.ep_rank = ep_rank + self.experts_per_rank = num_experts // ep_size + self.max_tokens_per_rank = max_tokens_per_rank + self.output_format = _parse_format(output_format) + self.combine_format = _parse_format(combine_format) + self.intermediate_format = ( + None if intermediate_format is None else _parse_format(intermediate_format) + ) + self.backward_operand_format = ( + None if backward_operand_format is None else _parse_format(backward_operand_format) + ) + self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) + self.gate_up_clamp = None if gate_up_clamp is None else abs(float(gate_up_clamp)) + self.generate_c = bool(generate_c) + self.backward_wgrad_mode = backward_wgrad_mode + self.token_padding_size = token_padding_size + + for name, fmt in ( + ("output_format", self.output_format), + ("combine_format", self.combine_format), + ): + required_multiple = ( + 32 if fmt is MoeFormat.MXFP8 else 16 if fmt is MoeFormat.NVFP4 else 1 + ) + if hidden_size % required_multiple != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by {required_multiple} for" + f" {name}={fmt.value}" + ) + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(" + f"experts={self.num_experts}, local_experts={self.experts_per_rank}, " + f"hidden={self.hidden_size}, intermediate={self.intermediate_size}, " + f"top_k={self.top_k}, ep_rank={self.ep_rank}/{self.ep_size}, " + f"output={self.output_format.value}, combine={self.combine_format.value})" + ) + + def _collective_device(self, device: torch.device) -> torch.device: + """Device the process group can run ``all_to_all_single`` on. + + Gloo only implements all-to-all for CPU tensors, so CUDA tensors are + staged through host memory; NCCL groups communicate in place. + """ + if device.type != "cpu" and dist.get_backend(self.ep_group) == "gloo": + return torch.device("cpu") + return device + + def _exchange_counts(self, send_counts: torch.Tensor) -> torch.Tensor: + if self.ep_size == 1: + return send_counts.clone() + comm_device = self._collective_device(send_counts.device) + staged = send_counts.to(comm_device) + recv_counts = torch.empty_like(staged) + dist.all_to_all_single(recv_counts, staged, group=self.ep_group) + return recv_counts.to(send_counts.device) + + def _all_to_all( + self, + send: torch.Tensor, + send_counts: Sequence[int], + recv_counts: Sequence[int], + ) -> torch.Tensor: + if self.ep_size == 1: + return send.clone() + comm_device = self._collective_device(send.device) + staged = send.contiguous().to(comm_device) + output_shape = (sum(recv_counts), *send.shape[1:]) + recv = torch.empty(output_shape, dtype=send.dtype, device=comm_device) + dist.all_to_all_single( + recv, + staged, + output_split_sizes=list(recv_counts), + input_split_sizes=list(send_counts), + group=self.ep_group, + ) + return recv.to(send.device) + + def _dispatch_plan(self, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> _DispatchPlan: + """Route valid ``topk_idx`` entries to destination ranks, stably by rank. + + Backward reuses this so gradient re-dispatch reproduces the exact + forward route order. + """ + + device = topk_idx.device + token_count = topk_idx.shape[0] + flat_expert = topk_idx.reshape(-1).to(torch.int64) + flat_weight = topk_weights.reshape(-1).float() + valid = flat_expert != -1 + invalid_negative = flat_expert < -1 + invalid_high = flat_expert >= self.num_experts + if bool((invalid_negative | invalid_high).any().item()): + bad = flat_expert[invalid_negative | invalid_high][0].item() + raise ValueError(f"topk_idx contains out-of-range expert id {bad}") + + flat_token = torch.arange(token_count, device=device).repeat_interleave(self.top_k) + flat_slot = torch.arange(self.top_k, device=device).repeat(token_count) + expert = flat_expert[valid] + destination = torch.div(expert, self.experts_per_rank, rounding_mode="floor") + order = torch.argsort(destination, stable=True) + + send_counts_tensor = torch.bincount( + destination.index_select(0, order), minlength=self.ep_size + ).to(torch.int64) + recv_counts_tensor = self._exchange_counts(send_counts_tensor) + return _DispatchPlan( + send_expert=expert.index_select(0, order).remainder(self.experts_per_rank), + send_weight=flat_weight[valid].index_select(0, order), + send_token_idx=flat_token[valid].index_select(0, order), + send_slot_idx=flat_slot[valid].index_select(0, order), + send_counts=tuple(int(v) for v in send_counts_tensor.cpu().tolist()), + recv_counts=tuple(int(v) for v in recv_counts_tensor.cpu().tolist()), + ) + + def _run_local_experts( + self, + tokens: torch.Tensor, + local_expert_idx: torch.Tensor, + route_weight: torch.Tensor, + fc1_weight: torch.Tensor, + fc2_weight: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + output = torch.empty( + (tokens.shape[0], self.hidden_size), + dtype=torch.float32, + device=tokens.device, + ) + fc1_c_rows = [] if self.generate_c else None + for expert in range(self.experts_per_rank): + positions = torch.nonzero(local_expert_idx == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + expert_tokens = tokens.index_select(0, positions) + gate_up = expert_tokens @ fc1_weight[expert] + if fc1_c_rows is not None: + # Raw pre-SwiGLU accumulator: before clamp, no router weight. + fc1_c_rows.append(gate_up.to(torch.bfloat16)) + gate, up = gate_up.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + gate = gate.clamp(max=self.gate_up_clamp) + up = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + intermediate = F.silu(gate) * up + weights = route_weight.index_select(0, positions).unsqueeze(-1) + if self.apply_topk_in_fc1: + intermediate = intermediate * weights + if self.intermediate_format is not None: + intermediate = _format_round_trip( + intermediate, + self.intermediate_format, + ) + expert_output = intermediate @ fc2_weight[expert] + expert_output = forward_combine_round_trip( + expert_output, + self.combine_format, + ) + if not self.apply_topk_in_fc1: + # The upstream training kernel leaves scores out of dispatch + # and applies them in standalone TopkReduce after the combine + # wire-format round trip. + expert_output = expert_output * weights + output.index_copy_(0, positions, expert_output) + fc1_c = None + if fc1_c_rows is not None: + fc1_c = ( + torch.cat(fc1_c_rows) + if fc1_c_rows + else torch.empty( + (0, 2 * self.intermediate_size), dtype=torch.bfloat16, device=tokens.device + ) + ) + return output, fc1_c + + def __call__( + self, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Union[ + MoeTensor, + Tuple[MoeTensor, torch.Tensor, torch.Tensor], + Tuple[ + MoeTensor, + torch.Tensor, + torch.Tensor, + WgradForwardStashReference, + ], + ]: + """Run dispatch, local experts, return routing, top-k reduce, and encode. + + Shapes: + activation: ``(T, H)`` + fc1_weight: ``(E_local, H, 2 * I)`` + fc2_weight: ``(E_local, I, H)`` + topk_idx/topk_weights: ``(T, K)`` + + Returns the ``(T, H)`` result, or ``(result, fc1_c, route_metadata)`` + when constructed with ``generate_c=True``. In wgrad operand mode, a + fourth :class:`WgradForwardStashReference` item is returned. ``fc1_c`` + is the BF16 + pre-SwiGLU FC1 accumulator of every route this rank's experts + processed, ``(local_routes, 2 * I)``, grouped by local expert and + ordered within each expert by (source rank, source token-major route + order); captured before the gate/up clamp, without the router weight. + ``route_metadata`` is Int32 ``(local_routes, 4)`` with columns + ``(local_expert, src_rank, src_token, src_slot)``; row ``i`` identifies + the route behind ``fc1_c`` row ``i`` for the backward gradient + re-dispatch. + """ + + if topk_idx.ndim != 2: + raise ValueError(f"topk_idx must be 2-D, got shape {tuple(topk_idx.shape)}") + token_count = topk_idx.shape[0] + route_shape = (token_count, self.top_k) + if tuple(topk_idx.shape) != route_shape: + raise ValueError(f"topk_idx shape must be {route_shape}, got {tuple(topk_idx.shape)}") + if tuple(topk_weights.shape) != route_shape: + raise ValueError( + f"topk_weights shape must be {route_shape}, got {tuple(topk_weights.shape)}" + ) + if topk_idx.dtype not in (torch.int32, torch.int64): + raise TypeError(f"topk_idx must be int32 or int64, got {topk_idx.dtype}") + if not topk_weights.is_floating_point(): + raise TypeError(f"topk_weights must be floating point, got {topk_weights.dtype}") + if self.max_tokens_per_rank is not None and token_count > self.max_tokens_per_rank: + raise ValueError( + f"token count {token_count} exceeds max_tokens_per_rank={self.max_tokens_per_rank}" + ) + + device = _tensor_device(activation) + inputs = { + "fc1_weight": _tensor_device(fc1_weight), + "fc2_weight": _tensor_device(fc2_weight), + "topk_idx": topk_idx.device, + "topk_weights": topk_weights.device, + } + for name, input_device in inputs.items(): + if input_device != device: + raise ValueError(f"{name} must be on {device}, got {input_device}") + + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + ) + # The Rubin path first stages plain activation along H, then its + # forward column requantization forms x.T scales along routed K. + wgrad_activation_float = None + if self.backward_wgrad_mode == "operands": + wgrad_activation_float = _format_round_trip( + activation_float, + MoeFormat.MXFP8, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, 2 * self.intermediate_size), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + if self.backward_wgrad_mode == "operands": + # Plain forward operands are staged to the same public MXFP8 + # reduction-axis representation before the Rubin GEMMs. + fc1_float = _format_round_trip_axis( + fc1_float, + MoeFormat.MXFP8, + axis=1, + ) + fc2_float = _format_round_trip_axis( + fc2_float, + MoeFormat.MXFP8, + axis=1, + ) + + plan = self._dispatch_plan(topk_idx, topk_weights) + send_token_idx = plan.send_token_idx + send_slot_idx = plan.send_slot_idx + send_counts, recv_counts = plan.send_counts, plan.recv_counts + forward_activation_float = ( + wgrad_activation_float if wgrad_activation_float is not None else activation_float + ) + send_tokens = forward_activation_float.index_select( + 0, + send_token_idx, + ) + + recv_tokens = self._all_to_all(send_tokens, send_counts, recv_counts) + recv_wgrad_tokens = None + if wgrad_activation_float is not None: + recv_wgrad_tokens = self._all_to_all( + wgrad_activation_float.index_select(0, send_token_idx), + send_counts, + recv_counts, + ) + recv_expert = self._all_to_all(plan.send_expert, send_counts, recv_counts) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + + route_metadata = None + fc1_c_order = None + if self.generate_c: + recv_src_rank = torch.repeat_interleave( + torch.arange(self.ep_size, device=device), + torch.tensor(recv_counts, device=device), + ) + recv_token = self._all_to_all(send_token_idx, send_counts, recv_counts) + recv_slot = self._all_to_all(send_slot_idx, send_counts, recv_counts) + # Stable sort by local expert reproduces the fc1_c row order + # (grouped by expert; source order preserved within each group). + fc1_c_order = torch.argsort(recv_expert, stable=True) + route_metadata = ( + torch.stack((recv_expert, recv_src_rank, recv_token, recv_slot), dim=1) + .index_select(0, fc1_c_order) + .to(torch.int32) + ) + # recv rows are ordered by source rank, then that source's token-major + # route order, so the per-expert position grouping below realizes the + # documented fc1_c ordering. + recv_output, fc1_c = self._run_local_experts( + recv_tokens, + recv_expert, + recv_weight, + fc1_float, + fc2_float, + ) + + returned = self._all_to_all(recv_output, recv_counts, send_counts) + combine_plane = torch.zeros( + (token_count * self.top_k, self.hidden_size), + dtype=torch.float32, + device=device, + ) + send_flat_slot = send_token_idx * self.top_k + send_slot_idx + combine_plane.index_copy_(0, send_flat_slot, returned) + reduced = combine_plane.view(token_count, self.top_k, self.hidden_size).sum(dim=1) + + if self.output_format is MoeFormat.BF16: + output = reduced.to(torch.bfloat16) + else: + output = quantize_blockwise(reduced, self.output_format, axis=-1) + if self.generate_c: + if self.backward_wgrad_mode == "operands": + if recv_wgrad_tokens is None or fc1_c_order is None: + raise RuntimeError("wgrad forward staging was not built") + valid_counts = tuple( + int(value) + for value in torch.bincount( + recv_expert, + minlength=self.experts_per_rank, + ) + .cpu() + .tolist() + ) + padded_ends = [] + total = 0 + for count in valid_counts: + total += ( + _ceil_div( + count, + self.token_padding_size, + ) + * self.token_padding_size + ) + padded_ends.append(total) + ordered_tokens = recv_wgrad_tokens.index_select( + 0, + fc1_c_order, + ) + metadata_experts = route_metadata[:, 0].to(torch.int64) + padded_x = _padded_expert_rows( + ordered_tokens, + metadata_experts, + valid_counts, + padded_ends, + ) + wgrad_stash = WgradForwardStashReference( + fc1_a=quantize_blockwise( + padded_x.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + expert_offsets=torch.tensor( + padded_ends, + dtype=torch.int32, + device=device, + ), + valid_route_counts=torch.tensor( + valid_counts, + dtype=torch.int32, + device=device, + ), + route_metadata=route_metadata, + ) + return output, fc1_c, route_metadata, wgrad_stash + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + *, + wgrad_forward_stash: Optional[WgradForwardStashReference] = None, + ) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[ + torch.Tensor, + torch.Tensor, + WgradOperandsReference, + ], + ]: + """Backward pass consuming the ``generate_c=True`` stash. + + ``fc1_c`` is the recompute source: gate/up, the clamp masks, SwiGLU, + and the FC2 input are all rebuilt from it, so no post-SwiGLU forward + intermediate needs to be saved. ``route_metadata`` alone reconstructs + the mapping between re-dispatched rows and ``fc1_c`` rows and drives + the gradient return scatter. + + Quantization round-trips (input decode, ``combine_format``, + ``output_format``) are treated as straight-through identities; + ``grad_output`` is the ``(T, H)`` gradient of the dequantized output. + + Returns ``(grad_activation, grad_topk_weights)`` in float32. In wgrad + operand mode, a third :class:`WgradOperandsReference` item models the + caller-owned grouped-GEMM operands. + """ + + if not self.generate_c: + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) + if self.backward_wgrad_mode == "operands": + if not isinstance( + wgrad_forward_stash, + WgradForwardStashReference, + ): + raise TypeError("wgrad_forward_stash must be a WgradForwardStashReference") + if not torch.equal( + wgrad_forward_stash.route_metadata, + route_metadata, + ): + raise ValueError("wgrad_forward_stash route identity does not match route_metadata") + elif wgrad_forward_stash is not None: + raise ValueError("wgrad_forward_stash is only accepted in operands mode") + token_count = topk_idx.shape[0] + if tuple(grad_output.shape) != (token_count, self.hidden_size): + raise ValueError( + f"grad_output shape must be {(token_count, self.hidden_size)}, got" + f" {tuple(grad_output.shape)}" + ) + if not grad_output.is_floating_point(): + raise TypeError(f"grad_output must be floating point, got {grad_output.dtype}") + + device = _tensor_device(fc1_weight) + two_i = 2 * self.intermediate_size + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + ) + semantic_fc2_float = fc2_float + effective_backward_format = self.backward_operand_format + if effective_backward_format is None and self.backward_wgrad_mode == "operands": + effective_backward_format = MoeFormat.MXFP8 + if effective_backward_format is not None: + # The dGLU adapter requantizes both transposed weights along the + # backward GEMM reduction dimension. + fc1_float = _format_round_trip_axis( + fc1_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + fc2_float = _format_round_trip_axis( + fc2_float.transpose(1, 2), + effective_backward_format, + axis=1, + ).transpose(1, 2) + if fc1_c.shape != (int(route_metadata.shape[0]), two_i): + raise ValueError( + f"fc1_c shape must be {(int(route_metadata.shape[0]), two_i)}, got" + f" {tuple(fc1_c.shape)}" + ) + + # Re-dispatch router weights and output gradients along the identical + # forward routes. + plan = self._dispatch_plan(topk_idx, topk_weights) + send_counts, recv_counts = plan.send_counts, plan.recv_counts + semantic_grad_output = grad_output.float() + grad_output_float = semantic_grad_output + if effective_backward_format is not None: + grad_output_float = _format_round_trip( + grad_output_float, + effective_backward_format, + ) + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts) + recv_grad = self._all_to_all( + grad_output_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + recv_semantic_grad = self._all_to_all( + semantic_grad_output.index_select(0, plan.send_token_idx), + send_counts, + recv_counts, + ) + + # route_metadata rows are in fc1_c order; sorting them by + # (src_rank, src_token, src_slot) reproduces the receive order, giving + # the permutation between re-dispatched rows and fc1_c rows. + metadata = route_metadata.to(device=device, dtype=torch.int64) + local_routes = metadata.shape[0] + if local_routes > 0: + token_span = int(metadata[:, 2].max().item()) + 1 + recv_key = (metadata[:, 1] * token_span + metadata[:, 2]) * self.top_k + metadata[:, 3] + perm = torch.argsort(recv_key) # perm[j] = fc1_c row at receive position j + else: + perm = torch.empty((0,), dtype=torch.int64, device=device) + w_rows = torch.empty_like(recv_weight) + w_rows.index_copy_(0, perm, recv_weight) + dy_rows = torch.empty_like(recv_grad) + dy_rows.index_copy_(0, perm, recv_grad) + semantic_dy_rows = torch.empty_like(recv_semantic_grad) + semantic_dy_rows.index_copy_(0, perm, recv_semantic_grad) + + c_rows = fc1_c.float() + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros((local_routes, self.hidden_size), dtype=torch.float32, device=device) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + h_rows = torch.zeros( + (local_routes, self.intermediate_size), + dtype=torch.float32, + device=device, + ) + weighted_dy_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=torch.float32, + device=device, + ) + dc_rows = torch.zeros( + (local_routes, two_i), + dtype=torch.float32, + device=device, + ) + for expert in range(self.experts_per_rank): + positions = torch.nonzero(expert_rows == expert, as_tuple=False).flatten() + if positions.numel() == 0: + continue + c = c_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = dy_rows.index_select(0, positions) + semantic_d_y = semantic_dy_rows.index_select(0, positions) + + gate, up = c.split(self.intermediate_size, dim=-1) + if self.gate_up_clamp is not None: + g = gate.clamp(max=self.gate_up_clamp) + u = up.clamp(min=-self.gate_up_clamp, max=self.gate_up_clamp) + else: + g, u = gate, up + sig = torch.sigmoid(g) + s = g * sig + h = s * u + h_rows.index_copy_(0, positions, h) + weighted_dy_rows.index_copy_(0, positions, d_y * w) + + if self.apply_topk_in_fc1: + d_y_pre = d_y + else: + d_y_pre = d_y * w + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + semantic_d_h = semantic_d_y @ semantic_fc2_float[expert].transpose(0, 1) + d_w_rows[positions] = (semantic_d_h * h).sum(dim=-1) + else: + d_h = d_h_fc2 + d_w_rows[positions] = (semantic_d_y * (h @ semantic_fc2_float[expert])).sum(dim=-1) + + d_g = d_h * u * (sig * (1 + g * (1 - sig))) + d_u = d_h * s + if self.gate_up_clamp is not None: + d_gate = d_g * (gate <= self.gate_up_clamp) + d_up = d_u * ((up >= -self.gate_up_clamp) & (up <= self.gate_up_clamp)) + else: + d_gate, d_up = d_g, d_u + d_c = torch.cat((d_gate, d_up), dim=-1) + dc_rows.index_copy_(0, positions, d_c) + if self.intermediate_format is not None: + d_c = _format_round_trip(d_c, self.intermediate_format) + d_x = d_c @ fc1_float[expert].transpose(0, 1) + d_x_rows.index_copy_( + 0, + positions, + backward_combine_round_trip(d_x, self.combine_format), + ) + + # Return the route gradients to their source ranks and scatter-add. + returned_dx = self._all_to_all(d_x_rows.index_select(0, perm), recv_counts, send_counts) + returned_dw = self._all_to_all(d_w_rows.index_select(0, perm), recv_counts, send_counts) + grad_activation = torch.zeros( + (token_count, self.hidden_size), dtype=torch.float32, device=device + ) + grad_activation.index_add_(0, plan.send_token_idx, returned_dx) + grad_topk_weights = torch.zeros( + (token_count * self.top_k,), dtype=torch.float32, device=device + ) + grad_topk_weights.index_copy_( + 0, plan.send_token_idx * self.top_k + plan.send_slot_idx, returned_dw + ) + grad_topk_weights = grad_topk_weights.view( + token_count, + self.top_k, + ) + if self.backward_wgrad_mode == "operands": + stash = wgrad_forward_stash + padded_ends = tuple(int(value) for value in stash.expert_offsets.cpu().tolist()) + valid_counts = tuple(int(value) for value in stash.valid_route_counts.cpu().tolist()) + padded_dc = _padded_expert_rows( + dc_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_h = _padded_expert_rows( + h_rows, + expert_rows, + valid_counts, + padded_ends, + ) + padded_weighted_dy = _padded_expert_rows( + weighted_dy_rows, + expert_rows, + valid_counts, + padded_ends, + ) + operands = WgradOperandsReference( + fc1_a=stash.fc1_a, + fc1_b=quantize_blockwise( + padded_dc, + MoeFormat.MXFP8, + axis=0, + ), + fc2_a=quantize_blockwise( + padded_h.transpose(0, 1), + MoeFormat.MXFP8, + axis=1, + ), + fc2_b=quantize_blockwise( + padded_weighted_dy, + MoeFormat.MXFP8, + axis=0, + ), + expert_offsets=stash.expert_offsets, + valid_route_counts=stash.valid_route_counts, + route_metadata=stash.route_metadata, + ) + return grad_activation, grad_topk_weights, operands + return grad_activation, grad_topk_weights + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "WgradForwardStashReference", + "WgradOperandsReference", + "backward_combine_round_trip", + "forward_combine_round_trip", + "quantize_blockwise", +] diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index f39115c4c6..dea22d170c 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -11,6 +11,12 @@ import torch from transformer_engine_torch import FP8TensorMeta +from ..ep import ( + EpBuffer, + EpConfig, + get_ep_drop_on_overflow, + get_ep_group, +) from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..quantized_tensor import QuantizedTensorStorage, Quantizer @@ -43,6 +49,94 @@ def get_fused_normalization_quantizer( return None +def validate_buffer( + name: str, + buffer: Optional[torch.Tensor], + *, + shape: Optional[tuple[int, ...] | list[int]] = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + contiguous: Optional[bool] = None, +) -> Optional[torch.Tensor]: + """Validate requested buffer properties and return detached storage.""" + if buffer is None: + return None + if shape is not None and tuple(buffer.shape) != tuple(shape): + raise ValueError(f"{name} shape {tuple(buffer.shape)} does not match {tuple(shape)}.") + if dtype is not None and buffer.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {buffer.dtype}.") + if device is not None and buffer.device != device: + raise ValueError(f"{name} must be on {device}, got {buffer.device}.") + if contiguous is not None and buffer.is_contiguous() != contiguous: + requirement = "contiguous" if contiguous else "non-contiguous" + raise ValueError(f"{name} must be {requirement}.") + return buffer.detach() + + +def validate_ep_buffer( + op_name: str, + expected_config: EpConfig, + buffer: object, +) -> EpBuffer: + """Validate a runtime EP buffer against an operation's immutable config.""" + if not isinstance(buffer, EpBuffer): + raise TypeError(f"{op_name} requires buffer=EpBuffer(...), got {type(buffer).__name__}.") + + mismatches = { + name: (getattr(buffer, name), getattr(expected_config, name)) + for name in ( + "top_k", + "hidden_dim", + "num_local_experts", + "max_tokens_per_rank", + "recv_capacity_per_rank", + "alignment", + "payload_dtype", + "zero_copy", + ) + if getattr(buffer, name) != getattr(expected_config, name) + } + ep_group = get_ep_group() + if expected_config.ep_group is not ep_group: + mismatches["ep_group"] = (ep_group, expected_config.ep_group) + drop_on_overflow = get_ep_drop_on_overflow() + if expected_config.drop_on_overflow != drop_on_overflow: + mismatches["drop_on_overflow"] = ( + drop_on_overflow, + expected_config.drop_on_overflow, + ) + if mismatches: + details = ", ".join( + f"{name}={actual!r} (expected {expected!r})" + for name, (actual, expected) in mismatches.items() + ) + raise ValueError( + f"{op_name} runtime buffer config does not match its initialized config: {details}." + ) + return buffer + + +def validate_ep_comms_recipe( + op_name: str, + quantizer: Optional[Quantizer], + buffer_recipe: object, +) -> None: + """Require the buffer recipe to match the Op's quantizer role.""" + if isinstance(quantizer, MXFP8Quantizer): + from transformer_engine.common.recipe import MXFP8BlockScaling + + if not isinstance(buffer_recipe, MXFP8BlockScaling): + raise ValueError( + f"{op_name} selected MXFP8 Comms from its quantizer role, but the " + "runtime EpBuffer does not have an MXFP8BlockScaling recipe." + ) + elif buffer_recipe is not None: + raise ValueError( + f"{op_name} selected BF16 Comms from its quantizer role, but the " + f"runtime EpBuffer has recipe {type(buffer_recipe).__name__}." + ) + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 4eb2796b2c..aa327aabe5 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -22,7 +22,9 @@ from .all_reduce import AllReduce from .basic_linear import BasicLinear from .bias import Bias +from .combine import MoeCombine from .constant_scale import ConstantScale +from .dispatch import MoeDispatch from .dropout import Dropout from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported from .identity import Identity diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..ea8d29e0f7 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -418,7 +418,7 @@ def fuser_forward( dtype = extra_input.dtype x = maybe_dequantize(input_.contiguous(), dtype) - scales = maybe_dequantize(extra_input, dtype) + scales = extra_input y = self._scaled_unary_forward(x, scales) ctx = basic_op_ctxs[0] diff --git a/transformer_engine/pytorch/ops/basic/combine.py b/transformer_engine/pytorch/ops/basic/combine.py new file mode 100644 index 0000000000..1ad8602803 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -0,0 +1,173 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel combine operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import ( + EpBuffer, + EpConfig, + _ep_combine_bwd, + _ep_combine_fwd, + quantize_for_ep, +) +from ...quantization import QuantizerRole +from ...tensor import MXFP8Quantizer, Quantizer +from .._common import ( + maybe_dequantize, + validate_ep_buffer, + validate_ep_comms_recipe, +) +from ..op import BasicOperation, OperationContext + + +def _validate_combine_inputs( + input_: torch.Tensor, + buffer: EpBuffer, +) -> tuple[int, int]: + """Validate the expert output and routing metadata consumed by MoeCombine.""" + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") + if input_.ndim != 2: + raise ValueError(f"MoeCombine input must be 2D, got shape {tuple(input_.shape)}.") + if buffer.handle_mem.dtype is not torch.uint8 or buffer.handle_mem.device != input_.device: + raise ValueError("MoeCombine routing handle must be a uint8 tensor on the input device.") + if ( + buffer.tokens_per_expert.dtype is not torch.int64 + or buffer.tokens_per_expert.device != input_.device + ): + raise ValueError( + "MoeCombine tokens_per_expert must be an int64 tensor on the input device." + ) + return tuple(input_.shape) + + +class MoeCombine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation consumes routing state from a runtime :class:`EpBuffer`. + """ + + num_extra_inputs: int = 0 + + def __init__(self, config: EpConfig) -> None: + super().__init__() + if not isinstance(config, EpConfig): + raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + if config.zero_copy: + raise NotImplementedError("MoeCombine does not support zero-copy EP.") + self.config = config + + def num_quantizers(self, mode: str) -> int: + return 1 if mode == "backward" else 0 + + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + if mode == "backward": + # combine backward dispatches grad_output + name = getattr(self, "name", "") or "" + return [ + QuantizerRole( + module_type="combine", + tensor_type="dispatch_grad_output", + name=name, + ) + ] + return None + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + quantizer = self.get_quantizer("backward", 0) + if quantizer is not None: + quantizer.set_usage(rowwise=True, columnwise=False) + quantizer.optimize_for_gemm = False + quantizer.internal = True + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("MoeCombine uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("MoeCombine uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, list[tuple[()]]]: + # Combine's transport format is selected by Combine's own grad-output quantizer. + # If the preceding op expects a different gradient quantized format, + # it is requantized in that op's backward implementation (e.g., GroupedLinear backward). + del basic_op_extra_inputs, prev_op_grad_output_quantizer, next_op_input_quantizer + grad_output_quantizer = self.get_quantizer("backward", 0) + transport_quantizer = ( + grad_output_quantizer if isinstance(grad_output_quantizer, MXFP8Quantizer) else None + ) + kwargs = basic_op_kwargs[0] + buffer = validate_ep_buffer("MoeCombine", self.config, kwargs.get("buffer")) + validate_ep_comms_recipe( + "MoeCombine", + grad_output_quantizer, + buffer.combine_bwd_quant_recipe, + ) + # Only BF16 combine forward is supported for now. + input_ = maybe_dequantize(input_, torch.bfloat16) + _validate_combine_inputs(input_, buffer) + ctx = basic_op_ctxs[0] + result, combine_state = _ep_combine_fwd( + input_, + None, + handle_mem=buffer.handle_mem, + token_counts=buffer.tokens_per_expert, + num_local_tokens=buffer.num_local_tokens, + hidden_dim=input_.shape[-1], + bwd_quant_recipe=transport_quantizer, + eager=buffer.eager, + zero_copy=False, + ) + if ctx.requires_grad: + ctx.combine_state = combine_state + + return result, [()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + ctx = basic_op_ctxs[0] + grad_output_quantizer = self.get_quantizer("backward", 0) + grad_scale_inv = None + # Prepare grad_output (Quantize if necessary) + if isinstance(grad_output_quantizer, MXFP8Quantizer): + quantized_grad, grad_scale_inv = quantize_for_ep( + grad_output, + grad_output_quantizer, + ) + grad_output = quantized_grad + else: + grad_output = maybe_dequantize(grad_output, torch.bfloat16).contiguous() + quantized_grad = None + grad_input = _ep_combine_bwd( + ctx.combine_state, + grad_output, + quantized_grad, + grad_scale_inv, + ) + return grad_input, [()], [()] diff --git a/transformer_engine/pytorch/ops/basic/dispatch.py b/transformer_engine/pytorch/ops/basic/dispatch.py new file mode 100644 index 0000000000..f9e6db14e0 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -0,0 +1,202 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible NCCL expert-parallel dispatch operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +import torch + +from ...ep import ( + EpBuffer, + EpConfig, + _ep_dispatch_bwd, + _ep_prepare_and_dispatch_fwd, + quantize_for_ep, +) +from ...quantization import QuantizerRole +from ...tensor import MXFP8Quantizer, Quantizer +from .._common import ( + is_quantized_tensor, + maybe_dequantize, + validate_ep_buffer, + validate_ep_comms_recipe, +) +from ..op import BasicOperation, OperationContext + + +def _validate_dispatch_input( + input_: torch.Tensor, + buffer: EpBuffer, +) -> tuple[int, int]: + """Validate the local token matrix.""" + if ( + not isinstance(input_, torch.Tensor) + or is_quantized_tensor(input_) + or input_.dtype is not torch.bfloat16 + ): + raise TypeError( + f"MoeDispatch input must be a plain BF16 tensor, got {type(input_).__name__}." + ) + input_shape = tuple(input_.shape) + if len(input_shape) != 2 or input_shape[-1] != buffer.hidden_dim: + raise ValueError( + f"MoeDispatch input must have shape (T, {buffer.hidden_dim}), got {input_shape}." + ) + if input_.device != buffer.device: + raise ValueError(f"MoeDispatch input must be on {buffer.device}, got {input_.device}.") + return input_shape + + +def _validate_routing_inputs( + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + *, + device: torch.device, +) -> None: + """Validate routing properties not checked by the native binding.""" + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + for name, tensor in (("topk_idx", topk_idx), ("topk_weights", topk_weights)): + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + + +class MoeDispatch(BasicOperation): + """Quantize and dispatch BF16 tokens to local experts with NCCL EP. + + The extra inputs are routing indices and FP32 routing weights. The extra + outputs are local tokens-per-expert and received routing weights. + """ + + num_extra_inputs: int = 2 + # tokens-per-expert and received routing weights consumed by the expert MLP. + num_extra_outputs: int = 2 + + def __init__(self, config: EpConfig) -> None: + super().__init__() + if not isinstance(config, EpConfig): + raise TypeError(f"config must be an EpConfig, got {type(config).__name__}.") + if config.zero_copy: + raise NotImplementedError("MoeDispatch does not support zero-copy EP.") + self.config = config + + def num_quantizers(self, mode: str) -> int: + # quantized dispatch_bwd/combine is not supported. + return 1 if mode == "forward" else 0 + + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + if mode == "forward": + name = getattr(self, "name", "") or "" + return [ + QuantizerRole( + module_type="dispatch", + tensor_type="dispatch_input", + name=name, + ) + ] + return None + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + quantizer = self.get_quantizer("forward", 0) + if quantizer is not None: + # We just need data, scales for dispatch, and grouped tensor + # will be recreated after dispatch op. + quantizer.set_usage(rowwise=True, columnwise=False) + quantizer.optimize_for_gemm = False + quantizer.internal = True + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("MoeDispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("MoeDispatch uses fuser_backward") + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + del next_op_input_quantizer + # Dispatch uses unquantized transport without an input quantizer and + # MXFP8 transport with an MXFP8 input quantizer. + input_quantizer = self.get_quantizer("forward", 0) + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + buffer = validate_ep_buffer("MoeDispatch", self.config, kwargs.get("buffer")) + validate_ep_comms_recipe( + "MoeDispatch", + input_quantizer, + buffer.dispatch_fwd_quant_recipe, + ) + input_shape = _validate_dispatch_input(input_, buffer) + buffer.num_local_tokens = input_shape[0] + _validate_routing_inputs( + topk_idx, + topk_weights, + device=buffer.device, + ) + # Prepare the input + input_scale_inv = None + if isinstance(input_quantizer, MXFP8Quantizer): + input_, input_scale_inv = quantize_for_ep(input_, input_quantizer) + output, recv_topk_weights, dispatch_state = _ep_prepare_and_dispatch_fwd( + input_, + topk_weights, + topk_idx, + buffer, + None, + None, + input_scale_inv, + ) + tokens_per_expert = buffer.tokens_per_expert + # If next_op_input_quantizer is different from input_quantizer, + # we need to requantize the data, which is handled in grouped_linear anyway. + # We won't get any fusion benefit, so don't do it here. + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.dispatch_state = dispatch_state + ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer + + return output, [(tokens_per_expert, recv_topk_weights)] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + # Only BF16 Dispatch_bwd is supported for now. + grad_output = maybe_dequantize(grad_output, torch.bfloat16) + + grad_recv_weights = basic_op_grad_extra_outputs[0][1] + if grad_recv_weights is None: + grad_recv_weights = torch.zeros( + grad_output.shape[0], + dtype=torch.float32, + device=grad_output.device, + ) + else: + grad_recv_weights = grad_recv_weights.to(dtype=torch.float32) + + grad_input, grad_topk_weights = _ep_dispatch_bwd( + ctx.dispatch_state, + grad_output, + grad_recv_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index c19a09c2f6..71cb586079 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -527,18 +527,18 @@ def fuser_forward( ) extra_input = basic_op_extra_inputs[0][0] - + scale_dtype = None # Determine compute dtype if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") - elif isinstance(input_, torch.Tensor): - dtype = input_.dtype + scale_dtype = dtype else: - dtype = extra_input.dtype + dtype = input_.dtype + scale_dtype = extra_input.dtype - # Make sure inputs are in correct dtype + # Prepare the activation input in the compute dtype. input_ = maybe_dequantize(input_, dtype) - scales = maybe_dequantize(extra_input, dtype) + scales = maybe_dequantize(extra_input, scale_dtype) out = self._scaled_glu_forward(input_, scales) # Save state for backward pass diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index dc9dcd6dc3..7aa45a89f5 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -35,3 +35,4 @@ GroupedMLP_CuTeGEMMGLU, GroupedMLP_CuTeGEMMUnary, ) +from .moe_ep import FusedMoeEp # pylint: disable=wrong-import-position diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py new file mode 100644 index 0000000000..a90e2d4506 --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,615 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp).""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +import os +from typing import Any, Optional + +import torch +import transformer_engine_torch as tex + +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE +from ...quantization import Recipe +from ...tensor import GroupedTensor, Quantizer +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, + view_main_grad_as_grouped_buffer, +) +from ..basic import GroupedLinear, MoeCombine, MoeDispatch, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +def _cudnn_megamoe_supported() -> bool: + """Whether cuDNN FE provides the fixed-resource training API.""" + try: + from cudnn import grouped_gemm_wgrad_wrapper_sm100 # noqa: F401 + from cudnn.moe_ep import ( # noqa: F401 + MoeEp, + MoeEpTrainingWeights, + MoeEpTrainingWgradOperands, + ) + except (AttributeError, ImportError): + return False + return True + + +def _get_megamoe_combine_format() -> str: + """Return the MegaMoE combine wire format selected by the environment.""" + enabled = int(os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0")) + return "mxfp8" if enabled > 0 else "bf16" + + +def _get_megamoe_training_slot_count() -> int: + """Return the fixed number of concurrent MegaMoE training flights.""" + value = os.environ.get("NVTE_MEGAMOE_TRAINING_SLOT_COUNT", "8") + try: + slot_count = int(value) + except ValueError as exc: + raise ValueError( + f"NVTE_MEGAMOE_TRAINING_SLOT_COUNT must be a positive integer, got {value!r}" + ) from exc + if slot_count <= 0: + raise ValueError( + f"NVTE_MEGAMOE_TRAINING_SLOT_COUNT must be a positive integer, got {value!r}" + ) + return slot_count + + +def _pack_as_cudnn_moe_tensor( + data: torch.Tensor, + scale: Optional[torch.Tensor], + block_scaled_cls: type, +): + """Represent data and scales using the public cuDNN MoE tensor type.""" + if scale is None: + return data + return block_scaled_cls( + data=data, + scale=scale, + format="mxfp8", + logical_shape=tuple(data.shape), + axis=1, + ) + + +def _pack_cudnn_weights( + op: GroupedLinear, + *, + block_scaled_cls: Optional[type] = None, +): + """Pack TE rowwise/columnwise weight storage for cuDNN MoE. + + The rowwise binding is a zero-copy ``(E, in, out)`` K-major view over TE's + rowwise ``(E, out, in)`` storage. The columnwise binding is a zero-copy + ``(E, out, in)`` view backed by TE's columnwise storage. + """ + weight = op.weight + num_groups = op.num_groups + out_features = op.out_features + in_features = op.in_features + weight_quantizer = op.get_quantizer("forward", 1) + if weight_quantizer is None and weight.quantizer is None: + return ( + weight.rowwise_data.view( + num_groups, + out_features, + in_features, + ).permute(0, 2, 1), + None, + ) + + if weight_quantizer is not None and weight.quantizer is None: + weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = tex.group_quantize( + weight.rowwise_data.view(weight.logical_shape), + weight_quantizer, + op.num_groups, + None, + ) + data = ( + weight.rowwise_data.view(num_groups, out_features, in_features) + .view(torch.float8_e4m3fn) + .permute(0, 2, 1) + ) + scale = ( + weight.scale_inv.view( + num_groups, + out_features, + in_features // MXFP8_BLOCK_SCALING_SIZE, + ) + .view(torch.float8_e8m0fnu) + .permute(0, 2, 1) + ) + if block_scaled_cls is None: + from cudnn.moe_ep import BlockScaledTensor as block_scaled_cls + if weight.columnwise_data is None or weight.columnwise_scale_inv is None: + raise ValueError("FusedMoeEp training requires columnwise MXFP8 weight storage") + + columnwise_data = weight.columnwise_data.view( + num_groups, + out_features, + in_features, + ).view(torch.float8_e4m3fn) + columnwise_scale = weight.columnwise_scale_inv.view( + num_groups, + out_features // MXFP8_BLOCK_SCALING_SIZE, + in_features, + ).view(torch.float8_e8m0fnu) + return ( + _pack_as_cudnn_moe_tensor(data, scale, block_scaled_cls), + _pack_as_cudnn_moe_tensor( + columnwise_data, + columnwise_scale, + block_scaled_cls, + ), + ) + + +def _launch_grouped_wgrad_from_operands( + layer_operands: list[torch.Tensor], + unused: None, + output: torch.Tensor | GroupedTensor, + *, + offsets: torch.Tensor, + accumulate: bool, +) -> None: + """Compute one TE-layout grouped wgrad directly from MegaMoE's operands.""" + del unused + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + a_tensor, sfa_tensor, b_tensor, sfb_tensor = layer_operands + output_data = ( + output.rowwise_data.view(output.shape) if isinstance(output, GroupedTensor) else output + ) + # The fixed-resource producer ABI already represents A @ B in cuDNN's + # (in, out) layout. Write through a view of TE's contiguous (out, in) + # buffer; no operand transpose, permutation, or copy is needed. + # The public wrapper selects the Rubin specialization on SM107. + grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets, + output_mode="dense", + wgrad_tensor=output_data.transpose(1, 2), + wgrad_dtype=output_data.dtype, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, + accumulate_on_output=accumulate, + input_order="tensor2d", + ) + + +def _compute_grouped_weight_grad( + op: GroupedLinear, + operands, + prefix: str, +) -> list[Optional[torch.Tensor]]: + """Compute one dense weight gradient with cuDNN's grouped-WGrad API.""" + weight = op.weight + if not weight.requires_grad: + return [None] + + weight_shape = (op.out_features, op.in_features) + output_shape = (op.num_groups, *weight_shape) + accumulate = False + if op._accumulate_into_main_grad: + output_data = get_main_grad_from_param( + weight, + op_label=f"FusedMoeEp {prefix.upper()}", + ) + output_data = view_main_grad_as_grouped_buffer( + output_data, + op.num_groups, + weight_shape, + label=f"FusedMoeEp {prefix.upper()} weight", + ) + accumulate = get_accumulate_flag_in_param(weight) + else: + output_data = torch.empty( + output_shape, + dtype=weight.dtype, + device=weight.device, + ) + + layer_operands = [ + getattr(operands, f"{prefix}_a"), + getattr(operands, f"{prefix}_sfa"), + getattr(operands, f"{prefix}_b"), + getattr(operands, f"{prefix}_sfb"), + ] + _launch_grouped_wgrad_from_operands( + layer_operands, + None, + output_data, + offsets=operands.expert_offsets, + accumulate=accumulate, + ) + + if op._accumulate_into_main_grad: + return get_dummy_wgrads_for_params([weight]) + return [output_data] + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weight = op.weight if op.single_grouped_weight else None + weight_ok = ( + isinstance(weight, GroupedTensor) + and weight.dtype is torch.bfloat16 + and weight.rowwise_data is not None + ) + if weight_ok and weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + weight_ok = ( + recipe is not None + and recipe.mxfp8() + and not weight._with_gemm_swizzled_scales + and weight.rowwise_data is not None + and weight.scale_inv is not None + and weight.columnwise_data is not None + and weight.columnwise_scale_inv is not None + ) + else: + weight_ok = False + + return ( + not op.use_bias + and not op._scale_bias + and op.single_grouped_weight + and not op.single_grouped_bias + and not op._is_distributed_weight() + and not op.wgrad_store.delay_wgrad_compute() + and weight_ok + ) + + +def _import_cudnn_moe_ep(): + """Return ``cudnn.moe_ep.MoeEp`` or ``None`` if the package is missing.""" + try: + from cudnn.moe_ep import MoeEp + except ImportError: + return None + return MoeEp + + +def _routing_extras_internal( + dispatch: MoeDispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`cudnn.moe_ep.MoeEp`, so it can only replace the + sequence when those two outputs feed exactly these ops and are not + returned to the caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + +def _megamoe_supported(config, fc1: GroupedLinear, fc2: GroupedLinear) -> bool: + """Static MegaMoE capability gates that can be checked before first launch.""" + if not _cudnn_megamoe_supported(): + return False + if _import_cudnn_moe_ep() is None: + return False + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + return False + if config.max_tokens_per_rank <= 0: + return False + if config.recv_capacity_per_rank is None or config.recv_capacity_per_rank <= 0: + return False + if config.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: + return False + if config.top_k > 32: + return False + return True + + +def _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: + if len(window) != 5: + return False + if recipe is not None and not recipe.mxfp8(): + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, MoeDispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, MoeCombine) + ): + return False + config = dispatch.config + if combine.config != config or config.payload_dtype is not torch.bfloat16: + return False + if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): + return False + if activation.activation_recompute_in_mlp or activation.glu_interleave_size != 32: + return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False + if not _megamoe_supported(config, fc1, fc2): + return False + return ( + fc1.num_groups == config.num_local_experts + and fc2.num_groups == config.num_local_experts + and fc1.in_features == config.hidden_dim + and fc2.out_features == config.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint EP MoE fusion implemented with :class:`cudnn.moe_ep.MoeEp`.""" + + def __init__( + self, + *, + dispatch: MoeDispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: MoeCombine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + moe_ep_cls = _import_cudnn_moe_ep() + if moe_ep_cls is None: + raise ImportError( + "FusedMoeEp requires cudnn.moe_ep.MoeEp. Install the in-tree " + "cuDNN frontend with: pip install --force-reinstall " + "'./cudnn_frontend[moe_ep]'" + ) + from cudnn.moe_ep import BlockScaledTensor + + config = dispatch.config + ep_group = config.ep_group + ep_size = ep_group.size() + combine_format = _get_megamoe_combine_format() + self._block_scaled_cls = BlockScaledTensor + self._moe = moe_ep_cls( + num_experts=config.num_local_experts * ep_size, + hidden_size=config.hidden_dim, + intermediate_size=fc2.in_features, + top_k=config.top_k, + ep_group=ep_group, + max_tokens_per_rank=config.max_tokens_per_rank, + max_recv_size_per_rank=config.recv_capacity_per_rank, + drop_on_overflow=config.drop_on_overflow, + apply_topk_in_fc1=True, + token_padding_size=128, + sf_padding_size=128, + combine_format=combine_format, + output_format="bf16", + ) + self._training_resources = None + self._training_slot_count = _get_megamoe_training_slot_count() + self._free_training_slots = [] + self._active_training_slots = set() + + def _make_training_weights(self): + """Bind cuDNN weight views to the GroupedLinear parameters' current storage.""" + from cudnn.moe_ep import MoeEpTrainingWeights + + fc1_rowwise, fc1_columnwise = _pack_cudnn_weights( + self.fc1, + block_scaled_cls=self._block_scaled_cls, + ) + fc2_rowwise, fc2_columnwise = _pack_cudnn_weights( + self.fc2, + block_scaled_cls=self._block_scaled_cls, + ) + return MoeEpTrainingWeights( + forward_fc1=fc1_rowwise, + forward_fc2=fc2_rowwise, + backward_w2_transpose=fc2_columnwise, + backward_w1_transpose=fc1_columnwise, + ) + + def _begin_training_flight(self): + """Reserve one fixed training slot for a forward/backward flight.""" + if self._training_resources is None: + self._training_resources = self._moe.prepare_training_resources( + self._make_training_weights(), + slot_count=self._training_slot_count, + lane_count=1, + ) + self._free_training_slots.extend(self._training_resources.slots) + if not self._free_training_slots: + raise RuntimeError( + "FusedMoeEp has no free training slots; increase " + "NVTE_MEGAMOE_TRAINING_SLOT_COUNT " + f"(currently {self._training_slot_count}) " + "or complete backward for an outstanding microbatch" + ) + if not self._active_training_slots: + self._training_resources.refresh_weights() + slot = self._free_training_slots.pop(0) + self._active_training_slots.add(slot) + return slot + + def _release_training_flight(self, slot) -> None: + """Return a completed forward/backward flight's slot to the pool.""" + if slot not in self._active_training_slots: + raise RuntimeError("FusedMoeEp attempted to release an inactive training slot") + self._active_training_slots.remove(slot) + self._free_training_slots.append(slot) + + @property + def dispatch(self) -> MoeDispatch: + return self.basic_ops[0] + + @property + def fc1(self) -> GroupedLinear: + return self.basic_ops[1] + + @property + def fc2(self) -> GroupedLinear: + return self.basic_ops[3] + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: + if ( + not isinstance(input_, torch.Tensor) + or is_quantized_tensor(input_) + or input_.dtype is not torch.bfloat16 + ): + raise TypeError( + f"FusedMoeEp input must be a plain BF16 tensor, got {type(input_).__name__}." + ) + + topk_idx, topk_weights = basic_op_extra_inputs[0] + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + activation = input_ + if not activation.is_contiguous(): + activation = activation.contiguous() + if topk_idx.dtype is not torch.int32: + topk_idx = topk_idx.to(dtype=torch.int32) + if not topk_idx.is_contiguous(): + topk_idx = topk_idx.contiguous() + if not topk_weights.is_contiguous(): + topk_weights = topk_weights.contiguous() + + slot = self._begin_training_flight() + try: + output = self._training_resources.forward( + slot, + self._training_resources.lanes[0], + activation, + topk_idx, + topk_weights, + ) + except Exception: + self._release_training_flight(slot) + raise + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + basic_op_ctxs[0].moe_ep_training_slot = slot + basic_op_ctxs[0].prev_op_grad_output_quantizer = prev_op_grad_output_quantizer + else: + try: + self._training_resources.finalize_overflow( + (slot,), + self._training_resources.lanes[0], + ) + finally: + self._release_training_flight(slot) + + return output, [ + (None, None), + (), + (), + (), + (), + ] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + grad_output = maybe_dequantize( + grad_output, + torch.bfloat16, + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + slot = basic_op_ctxs[0].moe_ep_training_slot + try: + grad_input, grad_topk_weights, wgrad_operands = self._training_resources.backward( + slot, + self._training_resources.lanes[0], + grad_output, + ) + fc1_param_grads = _compute_grouped_weight_grad(self.fc1, wgrad_operands, "fc1") + fc2_param_grads = _compute_grouped_weight_grad(self.fc2, wgrad_operands, "fc2") + self._training_resources.finalize_overflow( + (slot,), + self._training_resources.lanes[0], + ) + finally: + self._release_training_flight(slot) + return ( + grad_input, + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), ()], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op EP MoE sequences into MegaMoE. + + Unfused Sequential (NCCL dispatch/combine + GroupedLinear + ScaledSwiGLU) + is the default. MegaMoE is claimed only when ``_matches`` succeeds. + """ + del unused + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + window = ops[idx : idx + 5] + if _matches(window, recipe): + dispatch, fc1, activation, fc2, combine = window + out.append( + FusedMoeEp( + dispatch=dispatch, + fc1=fc1, + activation=activation, + fc2=fc2, + combine=combine, + ) + ) + idx += 5 + else: + out.append(ops[idx]) + idx += 1 + return out + + +register_forward_backward_fusion(fuse_ops, prepend=True) + + +__all__ = ["FusedMoeEp"] diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 0cc03602a1..65615d2662 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -173,6 +173,69 @@ def __new__( ) return instance + def __tensor_flatten__(self): + """Expose grouped backing buffers to PyTorch's wrapper-subclass protocol.""" + tensor_attrs = ( + "rowwise_data", + "columnwise_data", + "scale_inv", + "columnwise_scale_inv", + "amax", + "columnwise_amax", + "scale", + "first_dims", + "last_dims", + "tensor_offsets", + ) + present = [name for name in tensor_attrs if getattr(self, name) is not None] + context = { + "cls": type(self), + "logical_shape": self.logical_shape, + "fake_dtype": self.fake_dtype, + "num_tensors": self.num_tensors, + "tensor_shapes": self.tensor_shapes, + "quantizer": self.quantizer, + "offsets": self.offsets, + "scale_inv_offsets": self.scale_inv_offsets, + "columnwise_scale_inv_offsets": self.columnwise_scale_inv_offsets, + "requires_grad": self.requires_grad, + "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + } + return present, context + + @staticmethod + def __tensor_unflatten__(inner_tensors, context, outer_size, outer_stride): + """Rebuild a GroupedTensor from PyTorch-managed backing buffers.""" + return context["cls"]( + shape=context["logical_shape"], + dtype=context["fake_dtype"], + num_tensors=context["num_tensors"], + shapes=context["tensor_shapes"], + quantizer=context["quantizer"], + data=inner_tensors.get("rowwise_data"), + columnwise_data=inner_tensors.get("columnwise_data"), + scale_inv=inner_tensors.get("scale_inv"), + columnwise_scale_inv=inner_tensors.get("columnwise_scale_inv"), + amax=inner_tensors.get("amax"), + columnwise_amax=inner_tensors.get("columnwise_amax"), + scale=inner_tensors.get("scale"), + first_dims=inner_tensors.get("first_dims"), + last_dims=inner_tensors.get("last_dims"), + tensor_offsets=inner_tensors.get("tensor_offsets"), + offsets=context["offsets"], + scale_inv_offsets=context["scale_inv_offsets"], + columnwise_scale_inv_offsets=context["columnwise_scale_inv_offsets"], + requires_grad=context["requires_grad"], + stride=list(outer_stride) if outer_stride is not None else None, + with_gemm_swizzled_scales=context["with_gemm_swizzled_scales"], + row_scaled_nvfp4=context["row_scaled_nvfp4"], + nvfp4_use_4over6=context["nvfp4_use_4over6"], + nvfp4_e4m3_max=context["nvfp4_e4m3_max"], + ) + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): """Dispatch by dequantizing grouped members, then requantizing writes.""" @@ -229,6 +292,74 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: return make_wrapper_like(src, requires_grad=False) return make_wrapper_like(src, requires_grad=src.requires_grad) + # Module.to()/cuda() use aten._to_copy on parameters. Rebuild the + # wrapper around migrated grouped storage instead of taking the default + # dequantize-and-stack path, which loses the GroupedTensor subclass and + # leaves its backing buffers on the original device. + if func == torch.ops.aten._to_copy.default: + src = args[0] + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") + + target_device = kwargs.get("device", src.device) + target_dtype = kwargs.get("dtype", src.dtype) + target_layout = kwargs.get("layout", src.layout) + pin_memory = kwargs.get("pin_memory", False) + non_blocking = kwargs.get("non_blocking", False) + + if target_layout != torch.strided: + raise NotImplementedError( + f"{cls.__name__} only supports strided layout, got {target_layout}" + ) + if pin_memory: + raise NotImplementedError(f"{cls.__name__} does not support pin_memory=True") + if src.quantizer is not None and target_dtype != src.dtype: + raise NotImplementedError( + f"{cls.__name__} cannot change the logical dtype of quantized storage " + f"from {src.dtype} to {target_dtype}" + ) + + def move_storage( + tensor: Optional[torch.Tensor], *, convert_dtype: bool = False + ) -> Optional[torch.Tensor]: + if tensor is None: + return None + dtype = target_dtype if convert_dtype else tensor.dtype + return tensor.to( + device=target_device, + dtype=dtype, + non_blocking=non_blocking, + copy=True, + ) + + convert_data_dtype = src.quantizer is None + return type(src)( + shape=src.logical_shape, + dtype=target_dtype, + num_tensors=src.num_tensors, + shapes=src.tensor_shapes, + quantizer=src.quantizer, + data=move_storage(src.rowwise_data, convert_dtype=convert_data_dtype), + columnwise_data=move_storage(src.columnwise_data, convert_dtype=convert_data_dtype), + scale_inv=move_storage(src.scale_inv), + columnwise_scale_inv=move_storage(src.columnwise_scale_inv), + amax=move_storage(src.amax), + columnwise_amax=move_storage(src.columnwise_amax), + scale=move_storage(src.scale), + first_dims=move_storage(src.first_dims), + last_dims=move_storage(src.last_dims), + tensor_offsets=move_storage(src.tensor_offsets), + offsets=src.offsets, + scale_inv_offsets=src.scale_inv_offsets, + columnwise_scale_inv_offsets=src.columnwise_scale_inv_offsets, + requires_grad=src.requires_grad, + stride=list(src.stride()), + with_gemm_swizzled_scales=src._with_gemm_swizzled_scales, + row_scaled_nvfp4=src.row_scaled_nvfp4, + nvfp4_use_4over6=src.nvfp4_use_4over6, + nvfp4_e4m3_max=src.nvfp4_e4m3_max, + ) + # Parameter construction may invoke aten.expand on tensor subclasses. # Handle this explicitly so grouped parameters can be created safely. if func == torch.ops.aten.expand.default: