diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6e09ed316f..7a73a81223 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist +from transformer_engine.pytorch import ops as te_ops from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch.ep import ( EpBuffer, @@ -25,6 +26,7 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ops.fused.moe_ep import FusedMoeEp ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" @@ -39,6 +41,7 @@ # NVTE_EP_TOKENS_PER_RANK. HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512")) TOP_K = 2 +INTERMEDIATE_DIM = 16 TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32")) @@ -136,6 +139,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 @@ -277,6 +305,42 @@ def _moe_step(self, buffer, topk_idx, tokens, w): expert_out = self._weighted(recv_t, recv_w_out) return ep_combine(buffer, expert_out) + def _make_moe_model(self, *, fuse_ops=True): + """Build a BF16 EP MoE Sequential. + + With ``fuse_ops=True``, dispatch routing extras stay internal + (``output_to_caller=False``) so :class:`FusedMoeEp` can claim the + sequence. With ``fuse_ops=False``, those extras are returned to the + caller, which blocks fusion. + """ + buffer = self._make_buffer() + dispatch = te_ops.Dispatch(buffer) + fc1 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + HIDDEN_DIM, + 2 * INTERMEDIATE_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + activation = te_ops.ScaledSwiGLU() + fc2 = te_ops.GroupedLinear( + NUM_LOCAL_EXPERTS, + INTERMEDIATE_DIM, + HIDDEN_DIM, + bias=False, + device=self.cfg.device, + dtype=torch.bfloat16, + ) + combine = te_ops.Combine(buffer, num_local_tokens=TOKENS_PER_RANK) + + dispatch.set_extra_output_channel(0, "tokens_per_expert", output_to_caller=not fuse_ops) + dispatch.set_extra_output_channel(1, "routing_weights", output_to_caller=not fuse_ops) + 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") + return te_ops.Sequential(dispatch, fc1, activation, fc2, combine), fc1, fc2 + # Prepare @_eager_test_include @@ -528,6 +592,100 @@ def test_caller_provides_grad_expert_out(self): # the caller-owned buffer was used as the combine-bwd scatter target self.assertGreater(gbuf.abs().sum().item(), 0.0) + @_eager_test_include + def test_bf16_moe_sequential_fusion(self): + """Reference-backed fusion matches the unfused BF16 EP MoE sequence. + + ``fuse_ops=True`` keeps dispatch routing extras internal so fusion + fires; ``fuse_ops=False`` returns them to the caller and blocks it. + """ + if not EAGER: + self.skipTest("variable-size reference comparison requires eager EP mode") + + fused, fused_fc1, fused_fc2 = self._make_moe_model(fuse_ops=True) + unfused, unfused_fc1, unfused_fc2 = self._make_moe_model(fuse_ops=False) + generator = torch.Generator(device=self.cfg.device) + generator.manual_seed(3100 + self.cfg.rank) + with torch.no_grad(): + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + weight = ( + torch.randn( + getattr(fused_op, f"weight{expert}").shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + getattr(fused_op, f"weight{expert}").copy_(weight) + getattr(unfused_op, f"weight{expert}").copy_(weight) + + topk_idx, tokens, topk_weights = _make_moe_inputs( + self.cfg.rank, + self.cfg.ep_size, + self.cfg.device, + ) + fused_tokens = tokens.detach().clone().requires_grad_(True) + unfused_tokens = tokens.detach().clone().requires_grad_(True) + fused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + unfused_topk_weights = topk_weights.detach().clone().requires_grad_(True) + + fused_out = fused( + fused_tokens, + topk_idx, + fused_topk_weights, + ) + unfused_out, tokens_per_expert, recv_topk_weights = unfused( + unfused_tokens, + topk_idx, + unfused_topk_weights, + ) + + fused_forward_ops = fused._module_groups[0]._forward_ops + unfused_forward_ops = unfused._module_groups[0]._forward_ops + self.assertEqual(len(fused_forward_ops), 1) + self.assertIsInstance(fused_forward_ops[0][0], FusedMoeEp) + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in unfused_forward_ops)) + self.assertIsInstance(fused_out, torch.Tensor) + self.assertEqual(fused_out.dtype, torch.bfloat16) + self.assertEqual(unfused_out.dtype, torch.bfloat16) + self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) + self.assertEqual(tokens_per_expert.dtype, torch.int64) + self.assertEqual(recv_topk_weights.dtype, torch.float32) + + dy = ( + torch.randn( + fused_out.shape, + generator=generator, + dtype=torch.float32, + device=self.cfg.device, + ) + * 0.1 + ).to(torch.bfloat16) + fused_out.backward(dy) + unfused_out.backward(dy) + torch.cuda.synchronize() + + # The two BF16 paths use different grouped-GEMM and reduction orders. + # Observed forward abs error reaches ~3e-5 on near-zero values where + # rtol does not apply, so atol must stay above that (1e-5 is too tight). + tolerances = {"rtol": 1.6e-2, "atol": 5e-5} + torch.testing.assert_close(fused_out, unfused_out, **tolerances) + torch.testing.assert_close(fused_tokens.grad, unfused_tokens.grad, **tolerances) + torch.testing.assert_close( + fused_topk_weights.grad, + unfused_topk_weights.grad, + **tolerances, + ) + for fused_op, unfused_op in ((fused_fc1, unfused_fc1), (fused_fc2, unfused_fc2)): + for expert in range(NUM_LOCAL_EXPERTS): + fused_grad = getattr(fused_op, f"weight{expert}").grad + unfused_grad = getattr(unfused_op, f"weight{expert}").grad + self.assertEqual(fused_grad.dtype, torch.bfloat16) + self.assertEqual(unfused_grad.dtype, torch.bfloat16) + torch.testing.assert_close(fused_grad, unfused_grad, **tolerances) + @_zero_copy_test_include @_mxfp8_align_test def test_combine_bwd_mxfp8_caller_grad_out(self): diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 2799bfdf5d..08cac08ee1 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -26,6 +26,7 @@ __all__ = [ "EpBuffer", "ep_bootstrap", + "get_ep_group", "is_ep_bootstrapped", "ep_finalize", "ep_dispatch", @@ -183,6 +184,11 @@ 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 ep_finalize() -> None: """Optional explicit EP teardown; idempotent. diff --git a/transformer_engine/pytorch/ep_reference.py b/transformer_engine/pytorch/ep_reference.py new file mode 100644 index 0000000000..8a8d58ad0f --- /dev/null +++ b/transformer_engine/pytorch/ep_reference.py @@ -0,0 +1,881 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""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 _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, + dtype: torch.dtype, +) -> 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(dtype=dtype) + + 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.to(dtype=dtype) + + +def _format_round_trip( + tensor: torch.Tensor, + format: MoeFormat, + *, + dtype: torch.dtype, +) -> torch.Tensor: + if format is MoeFormat.BF16: + return tensor.to(torch.bfloat16).to(dtype) + return quantize_blockwise(tensor, format, axis=-1).dequantize(dtype=dtype) + + +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. + """ + + 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, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + generate_c: bool = False, + compute_dtype: torch.dtype = torch.float32, + ) -> 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 compute_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + f"compute_dtype must be torch.float32 or torch.bfloat16, got {compute_dtype}" + ) + + 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.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.compute_dtype = compute_dtype + + 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}, " + f"compute_dtype={self.compute_dtype})" + ) + + 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=self.compute_dtype, + 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 + expert_output = intermediate @ fc2_weight[expert] + if not self.apply_topk_in_fc1: + expert_output = expert_output * weights + expert_output = _format_round_trip( + expert_output, + self.combine_format, + dtype=self.compute_dtype, + ) + 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]]: + """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``. ``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, + dtype=self.compute_dtype, + ) + 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, + dtype=self.compute_dtype, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ) + + 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 + send_tokens = activation_float.index_select(0, send_token_idx) + + recv_tokens = self._all_to_all(send_tokens, 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).to( + dtype=self.compute_dtype + ) + + route_metadata = 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=self.compute_dtype, + 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: + return output, fc1_c, route_metadata + return output + + def backward( + self, + grad_output: torch.Tensor, + activation: MoeTensor, + fc1_weight: MoeTensor, + fc2_weight: MoeTensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + fc1_c: torch.Tensor, + route_metadata: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """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_fc1_weight, grad_fc2_weight, + grad_topk_weights)``. Activation and expert weight gradients use + ``compute_dtype``. Router weights are cast to ``compute_dtype`` for the + activation scaling path; the returned + ``grad_topk_weights`` remain float32. + """ + + if not self.generate_c: + raise RuntimeError( + "backward requires the operator to be constructed with generate_c=True" + ) + 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(activation) + two_i = 2 * self.intermediate_size + activation_float = _decode_tensor( + activation, + name="activation", + expected_shape=(token_count, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ) + fc1_float = _decode_tensor( + fc1_weight, + name="fc1_weight", + expected_shape=(self.experts_per_rank, self.hidden_size, two_i), + quantized_axis=1, + dtype=self.compute_dtype, + ) + fc2_float = _decode_tensor( + fc2_weight, + name="fc2_weight", + expected_shape=(self.experts_per_rank, self.intermediate_size, self.hidden_size), + quantized_axis=1, + dtype=self.compute_dtype, + ) + 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 the FC1 inputs, 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 + grad_output_float = grad_output.to(dtype=self.compute_dtype) + recv_tokens = self._all_to_all( + activation_float.index_select(0, plan.send_token_idx), send_counts, recv_counts + ) + # Same FP32-on-wire / compute_dtype-for-activation cast as forward. + recv_weight = self._all_to_all(plan.send_weight, send_counts, recv_counts).to( + dtype=self.compute_dtype + ) + recv_grad = self._all_to_all( + grad_output_float.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) + x_rows = torch.empty_like(recv_tokens) + x_rows.index_copy_(0, perm, recv_tokens) + 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) + + c_rows = fc1_c.to(dtype=self.compute_dtype) + expert_rows = metadata[:, 0] + d_x_rows = torch.zeros( + (local_routes, self.hidden_size), + dtype=self.compute_dtype, + device=device, + ) + d_w_rows = torch.zeros((local_routes,), dtype=torch.float32, device=device) + grad_fc1 = torch.zeros_like(fc1_float) + grad_fc2 = torch.zeros_like(fc2_float) + 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) + x = x_rows.index_select(0, positions) + w = w_rows.index_select(0, positions).unsqueeze(-1) + d_y = 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 + + if self.apply_topk_in_fc1: + h_fc2 = h * w + d_y_pre = d_y + else: + h_fc2 = h + d_y_pre = d_y * w + grad_fc2[expert] = h_fc2.transpose(0, 1) @ d_y_pre + d_h_fc2 = d_y_pre @ fc2_float[expert].transpose(0, 1) + if self.apply_topk_in_fc1: + d_h = d_h_fc2 * w + d_w_rows[positions] = (d_h_fc2 * h).sum(dim=-1).to(dtype=self.compute_dtype).float() + else: + d_h = d_h_fc2 + d_w_rows[positions] = ( + (d_y * (h @ fc2_float[expert])).sum(dim=-1).to(dtype=self.compute_dtype).float() + ) + + 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) + grad_fc1[expert] = x.transpose(0, 1) @ d_c + d_x_rows.index_copy_(0, positions, d_c @ fc1_float[expert].transpose(0, 1)) + + # 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=self.compute_dtype, + 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 + ) + return ( + grad_activation, + grad_fc1, + grad_fc2, + grad_topk_weights.view(token_count, self.top_k), + ) + + +__all__ = [ + "BlockScaledTensor", + "MoeEpReference", + "MoeFormat", + "MoeTensor", + "quantize_blockwise", +] diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index caab9c9c8d..29ebc1bf4b 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 Combine from .constant_scale import ConstantScale +from .dispatch import Dispatch 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/combine.py b/transformer_engine/pytorch/ops/basic/combine.py new file mode 100644 index 0000000000..4d4740c517 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/combine.py @@ -0,0 +1,138 @@ +# Copyright (c) 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, Optional + +import torch + +from ...ep import EpBuffer, _alloc_io, is_symm_backed +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_grad_buffer( + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"grad_out shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"grad_out must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"grad_out must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError("grad_out must be contiguous.") + if tensor.requires_grad: + raise ValueError("grad_out must not require gradients.") + return tensor + + +class Combine(BasicOperation): + """Combine pre-weighted local expert outputs with NCCL EP. + + The operation uses routing state produced by a :class:`Dispatch` with the + same :class:`EpBuffer`. + """ + + def __init__(self, buffer: EpBuffer, *, num_local_tokens: Optional[int] = None) -> None: + super().__init__() + self.buffer = buffer + self.num_local_tokens = ( + buffer.max_tokens_per_rank if num_local_tokens is None else int(num_local_tokens) + ) + if self.num_local_tokens < 0: + raise ValueError("num_local_tokens must be non-negative.") + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + *, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + **kwargs: Any, + ) -> torch.Tensor: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 combine input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Combine input must have shape (R, {self.buffer.hidden_dim}), " + f"got {tuple(input_.shape)}." + ) + + expert_out = input_ + if self.buffer.zero_copy: + expert_out = _alloc_io( + tuple(input_.shape), + input_.dtype, + input_.device, + True, + ) + expert_out.copy_(input_) + + result = torch.empty( + self.num_local_tokens, + self.buffer.hidden_dim, + dtype=input_.dtype, + device=input_.device, + ) + torch.ops.transformer_engine_ep.combine( + self.buffer.handle_mem, + expert_out, + result, + ) + + if ctx.requires_grad: + grad_out = kwargs.get("grad_out") + if self.buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes combine gradients per step and cannot use " + "a caller-supplied grad_out" + ) + grad_out = _validate_grad_buffer( + grad_out, + shape=tuple(input_.shape), + dtype=input_.dtype, + device=input_.device, + ) + if self.buffer.zero_copy and grad_out is not None and not is_symm_backed(grad_out): + raise ValueError("zero-copy Combine grad_out must be symmetric-memory-backed.") + ctx.grad_out = grad_out + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.save_for_backward(self.buffer.handle_mem) + + return result + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_input = ctx.grad_out + if grad_input is None: + grad_input = _alloc_io( + ctx.input_shape, + ctx.input_dtype, + grad_output.device, + self.buffer.zero_copy, + ) + torch.ops.transformer_engine_ep.combine_bwd( + handle_mem, + grad_output, + grad_input, + ) + 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..9762ed76f7 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/dispatch.py @@ -0,0 +1,207 @@ +# Copyright (c) 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, _alloc_io, ep_prepare +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext + + +def _validate_output_buffer( + name: str, + tensor: Optional[torch.Tensor], + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> Optional[torch.Tensor]: + if tensor is None: + return None + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} does not match {shape}.") + if tensor.dtype is not dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}.") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}.") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous.") + if tensor.requires_grad: + raise ValueError(f"{name} must not require gradients.") + return tensor + + +class Dispatch(BasicOperation): + """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 + num_extra_outputs: int = 2 + + def __init__(self, buffer: EpBuffer) -> None: + super().__init__() + self.buffer = buffer + + def op_forward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch uses fuser_forward") + + def op_backward(self, *args: Any, **kwargs: Any) -> None: + raise RuntimeError("Dispatch 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 prev_op_grad_output_quantizer, next_op_input_quantizer + topk_idx, topk_weights = basic_op_extra_inputs[0] + kwargs = basic_op_kwargs[0] + + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.") + if input_.ndim != 2 or input_.shape[-1] != self.buffer.hidden_dim: + raise ValueError( + f"Dispatch input must have shape (T, {self.buffer.hidden_dim}), " + f"got {tuple(input_.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}.") + expected_route_shape = (input_.shape[0], self.buffer.top_k) + if tuple(topk_idx.shape) != expected_route_shape: + raise ValueError( + f"topk_idx shape must be {expected_route_shape}, got {tuple(topk_idx.shape)}." + ) + if tuple(topk_weights.shape) != expected_route_shape: + raise ValueError( + f"topk_weights shape must be {expected_route_shape}, " + f"got {tuple(topk_weights.shape)}." + ) + 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 != input_.device: + raise ValueError(f"{name} must be on {input_.device}, got {tensor.device}.") + + recv_tokens = kwargs.get("recv_tokens") + recv_topk_weights = kwargs.get("recv_topk_weights") + if self.buffer.eager and (recv_tokens is not None or recv_topk_weights is not None): + raise ValueError( + "eager mode sizes dispatch outputs per step and cannot use " + "caller-supplied receive buffers" + ) + + tokens_per_expert = ep_prepare(self.buffer, topk_idx) + rows = ( + self.buffer._host_total_recv_tokens + if self.buffer.eager + else self.buffer.recv_capacity_per_rank + ) + if rows is None: + raise RuntimeError("NCCL EP dispatch receive size is unavailable.") + rows = int(rows) + recv_shape = (rows, self.buffer.hidden_dim) + recv_tokens = _validate_output_buffer( + "recv_tokens", + recv_tokens, + shape=recv_shape, + dtype=self.buffer.payload_dtype, + device=self.buffer.device, + ) + recv_topk_weights = _validate_output_buffer( + "recv_topk_weights", + recv_topk_weights, + shape=(rows,), + dtype=torch.float32, + device=self.buffer.device, + ) + if recv_tokens is None: + recv_tokens = _alloc_io( + recv_shape, + self.buffer.payload_dtype, + self.buffer.device, + self.buffer.zero_copy, + ) + if recv_topk_weights is None: + recv_topk_weights = _alloc_io( + (rows,), + torch.float32, + self.buffer.device, + self.buffer.zero_copy, + ) + + torch.ops.transformer_engine_ep.dispatch( + self.buffer.handle_mem, + topk_idx, + input_, + topk_weights, + recv_tokens, + recv_topk_weights, + ) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + ctx.input_shape = tuple(input_.shape) + ctx.input_dtype = input_.dtype + ctx.topk_weights_shape = tuple(topk_weights.shape) + ctx.save_for_backward(self.buffer.handle_mem) + + return recv_tokens, [(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] + (handle_mem,) = ctx.saved_tensors + grad_output = grad_output.contiguous() + + 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).contiguous() + + grad_input = torch.empty( + ctx.input_shape, + dtype=ctx.input_dtype, + device=grad_output.device, + ) + grad_topk_weights = torch.empty( + ctx.topk_weights_shape, + dtype=torch.float32, + device=grad_output.device, + ) + torch.ops.transformer_engine_ep.dispatch_bwd( + handle_mem, + grad_output, + grad_recv_weights, + grad_input, + grad_topk_weights, + ) + return grad_input, [()], [(None, grad_topk_weights)] 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..dd0f46c7ab --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,299 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Reference-backed BF16 expert-parallel MoE fusion.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import Any, Optional + +import torch + +from ...ep import get_ep_group +from ...ep_reference import MoeEpReference +from ...quantization import Recipe +from ...tensor import Quantizer +from ..basic import Combine, Dispatch, GroupedLinear, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +def _weight_list(op: GroupedLinear) -> list[torch.Tensor]: + """Return per-expert weights in their registered order.""" + return [getattr(op, f"weight{idx}") for idx in range(op.num_groups)] + + +def _reference_weights(op: GroupedLinear) -> torch.Tensor: + """Pack ``(out, in)`` expert weights into reference ``(E, in, out)`` layout.""" + return torch.stack([weight.transpose(0, 1) for weight in _weight_list(op)]) + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weights = _weight_list(op) if not op.single_grouped_weight else [] + return ( + not op.use_bias + and not op._scale_bias + and not op.single_grouped_weight + and not op.single_grouped_bias + and not op._accumulate_into_main_grad + and not op._is_distributed_weight() + and not op.wgrad_store.delay_wgrad_compute() + and bool(weights) + and all(weight.dtype is torch.bfloat16 for weight in weights) + ) + + +def _routing_extras_internal( + dispatch: Dispatch, + 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:`MoeEpReference`, 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 _routing_extras_internal( + dispatch: Dispatch, + 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:`MoeEpReference`, 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 _matches(window: Sequence[FusibleOperation], recipe: Optional[Recipe]) -> bool: + if recipe is not None or len(window) != 5: + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, Dispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, Combine) + ): + return False + if dispatch.buffer is not combine.buffer: + return False + buffer = dispatch.buffer + if not buffer.eager or buffer.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 is not None: + return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False + return ( + fc1.num_groups == buffer.num_local_experts + and fc2.num_groups == buffer.num_local_experts + and fc1.in_features == buffer.hidden_dim + and fc2.out_features == buffer.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint BF16 fusion implemented with :class:`MoeEpReference`.""" + + def __init__( + self, + *, + dispatch: Dispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: Combine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + ep_group = get_ep_group() + ep_size = 1 if ep_group is None else ep_group.size() + self._reference = MoeEpReference( + num_experts=dispatch.buffer.num_local_experts * ep_size, + hidden_size=dispatch.buffer.hidden_dim, + intermediate_size=fc2.in_features, + top_k=dispatch.buffer.top_k, + ep_group=ep_group, + max_tokens_per_rank=dispatch.buffer.max_tokens_per_rank, + compute_dtype=torch.bfloat16, + apply_topk_in_fc1=True, + generate_c=True, + ) + + @property + def dispatch(self) -> Dispatch: + 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]]]]: + del prev_op_grad_output_quantizer, next_op_input_quantizer + if input_.dtype is not torch.bfloat16: + raise NotImplementedError(f"FusedMoeEp requires BF16 input, got {input_.dtype}.") + if any(kwargs for kwargs in basic_op_kwargs): + raise NotImplementedError("FusedMoeEp does not support per-operation output buffers.") + + 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}.") + fc1_weight = _reference_weights(self.fc1) + fc2_weight = _reference_weights(self.fc2) + output, fc1_c, route_metadata = self._reference( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + ) + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + basic_op_ctxs[0].save_for_backward( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + # Dispatch extras are channel-bound with output_to_caller=False and are + # only consumed by ops inside this fusion, so they need not be materialized. + 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 + ( + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) = basic_op_ctxs[0].saved_tensors + grad_input, grad_fc1, grad_fc2, grad_topk_weights = self._reference.backward( + grad_output, + input_, + fc1_weight, + fc2_weight, + topk_idx, + topk_weights, + fc1_c, + route_metadata, + ) + + fc1_param_grads = [ + grad_fc1[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc1)) + ] + fc2_param_grads = [ + grad_fc2[idx].transpose(0, 1) if weight.requires_grad else None + for idx, weight in enumerate(_weight_list(self.fc2)) + ] + return ( + grad_input, + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights), (None,), (None,), (None,), ()], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op BF16 EP MoE sequences.""" + 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/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fd66529ba8..7009a2818c 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -432,6 +432,11 @@ def __init__( basic_ops.append(op) self._num_basic_ops: int = len(basic_ops) self._basic_ops: list[BasicOperation] = basic_ops + # Capture channel routing from each basic op. If any op later rebinds a + # channel it flips this flag, essentially invalidating this fuser's forward pass. + self._channels_stale: bool = False + for op in self._basic_ops: + op._capturing_op_fusers.add(self) # Number of extra tensor inputs self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops)