diff --git a/tests/unit_tests/test_mxfp8_fused_mlp.py b/tests/unit_tests/test_mxfp8_fused_mlp.py index 786eb4610f..f31e732a90 100644 --- a/tests/unit_tests/test_mxfp8_fused_mlp.py +++ b/tests/unit_tests/test_mxfp8_fused_mlp.py @@ -4,39 +4,65 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Wiring tests for the self-contained MXFP8 fused-MLP overrides. +"""Tests for the self-contained MXFP8 fused-MLP overrides. -Everything here is a config-tree transform plus a meta-device build, so it -runs without a GPU: the factories' SM100 gate is patched out. Numerics of the -underlying composites are validated on SM100 hardware in NVIDIA-internal CI. +The wiring tests are config-tree transforms plus meta/CPU builds with the +composites mocked, so they run without a GPU: the factories' SM100 gate is +patched out and the grouped_gemm_swiglu plan's torchao-ops availability is +mocked. Numerics of the grouped_gemm_swiglu composite run in the class gated +on real SM100 hardware plus the torchao fused grouped-MLP ops (skipped +otherwise, e.g. in GPU-less stock-torchao CI); the swiglu composites' numerics +are validated on SM100 hardware in NVIDIA-internal CI. """ import unittest from unittest import mock import torch +import torch.nn.functional as F +from torchtitan.components.quantization import MXFP8GroupedExpertsConverter from torchtitan.config.override import apply_overrides, OverrideConfig from torchtitan.models.common.feed_forward import FeedForward from torchtitan.models.common.linear import Linear -from torchtitan.models.common.moe import GroupedExperts +from torchtitan.models.common.moe import GroupedExperts, RoutedExperts from torchtitan.models.common.token_dispatcher import TorchAOTokenDispatcher from torchtitan.models.deepseek_v3 import model_registry as deepseek_v3_model_registry from torchtitan.models.llama3 import model_registry as llama3_model_registry try: from torchtitan.overrides.mxfp8_fused_mlp import ( + _pack_w13_blocks, + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE, + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON, + mxfp8_fused_grouped_mlp, mxfp8_fused_mlp, MXFP8FusedGroupedMLP, MXFP8FusedMLP, ) + from torchtitan.overrides.mxfp8_fused_mlp import ( + _MXFP8GroupedGemmMLP, + ) except ImportError as e: # torchao (or a transitive dep) not installed raise unittest.SkipTest( f"torchao is required for the MXFP8 fused-MLP overrides: {e}" ) from e +# Reaching here means the override module (and thus torchao) imported; these +# torchao symbols exist in every torchao the override module accepts and are +# only exercised by the SM100-gated numerics class. +from torchao.prototype.mx_formats.config import ScaleCalculationMode # noqa: E402 +from torchao.prototype.mx_formats.mx_tensor import to_mx # noqa: E402 +from torchao.prototype.mx_formats.utils import to_blocked # noqa: E402 +from torchao.quantization.utils import compute_error # noqa: E402 + _DENSE_OVERRIDE = "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_mlp" _GROUPED_OVERRIDE = "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_grouped_mlp" +_GROUPED_PLAN_IMPORT = (_GROUPED_OVERRIDE, {"fusion_plan": "grouped_gemm_swiglu"}) + +_HAS_SM100_GPU = ( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0) +) class TestMXFP8FusedMLPOverride(unittest.TestCase): @@ -196,5 +222,615 @@ def test_dense_factory_raises_on_converted_projection(self): mxfp8_fused_mlp(cfg) +class TestGroupedGemmPlanWiring(unittest.TestCase): + """GPU-less wiring tests for the ``grouped_gemm_swiglu`` fusion plan: + config routing, the pad-256 dispatcher policy, the forward-time 32-block + pack, and the factory's fail-loud raises. The SM100 gate and the torchao + fused grouped-MLP ops' availability are mocked.""" + + def setUp(self): + for patcher in ( + mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp.has_cuda_capability", + lambda *args: True, + ), + mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp." + "_TORCHAO_GROUPED_MLP_OPS_AVAILABLE", + True, + ), + mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp.is_supported", + lambda dim, hidden_dim: dim % 128 == 0 and hidden_dim % 128 == 0, + ), + ): + patcher.start() + self.addCleanup(patcher.stop) + + def _plan_model_config(self): + model_config = deepseek_v3_model_registry("debugmodel").model + apply_overrides(OverrideConfig(imports=[_GROUPED_PLAN_IMPORT]), model_config) + return model_config + + def _plan_experts_config(self, model_config): + nodes = list(model_config.traverse(MXFP8FusedGroupedMLP.Config)) + self.assertTrue(nodes) + return nodes[0][1] + + def _stock_routed_experts_config(self): + model_config = deepseek_v3_model_registry("debugmodel").model + nodes = list(model_config.traverse(RoutedExperts.Config)) + self.assertTrue(nodes) + return nodes[0][1] + + def test_pack_w13_blocks_matches_kernel_32_block_mapping(self): + # Byte-verify the forward-time pack against an independent + # re-implementation of the kernels' 32-block GLU row order: block i + # holds gate (w1) rows [32i, 32i+32), then the SAME features' up (w3) + # rows. + e, f, d = 2, 96, 8 + w1 = torch.arange(e * f * d, dtype=torch.float32).reshape(e, f, d) + w3 = -1.0 - torch.arange(e * f * d, dtype=torch.float32).reshape(e, f, d) + packed = _pack_w13_blocks(w1, w3) + self.assertEqual(tuple(packed.shape), (e, 2 * f, d)) + expected = torch.empty(e, 2 * f, d) + for i in range(f // 32): + expected[:, 64 * i : 64 * i + 32] = w1[:, 32 * i : 32 * i + 32] + expected[:, 64 * i + 32 : 64 * i + 64] = w3[:, 32 * i : 32 * i + 32] + self.assertTrue(torch.equal(packed, expected)) + # Unpack round-trip through the (E, F//32, 2, 32, D) view is the + # identity. + v = packed.view(e, f // 32, 2, 32, d) + self.assertTrue(torch.equal(v[:, :, 0].reshape(e, f, d), w1)) + self.assertTrue(torch.equal(v[:, :, 1].reshape(e, f, d), w3)) + + def test_plan_routes_to_the_grouped_gemm_function(self): + cfg = self._plan_experts_config(self._plan_model_config()) + self.assertEqual(cfg.fusion_plan, "grouped_gemm_swiglu") + module = cfg.build() + torch.manual_seed(0) + with torch.no_grad(): + for param in module.parameters(): + param.normal_() + num_tokens = torch.zeros(cfg.num_experts, dtype=torch.int64) + num_tokens[0] = 2 + x = torch.randn(2, cfg.dim) + sentinel = torch.zeros(2, cfg.dim, dtype=torch.bfloat16) + with mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp._MXFP8GroupedGemmMLP.apply", + return_value=sentinel, + ) as fused, mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp._MXFP8GroupedMLP.apply" + ) as swiglu: + out = module(x, num_tokens) + fused.assert_called_once() + swiglu.assert_not_called() + x_arg, w13_arg, w2_arg, offs_arg = fused.call_args.args + self.assertEqual(x_arg.dtype, torch.bfloat16) + self.assertTrue( + torch.equal( + w13_arg, + _pack_w13_blocks(module.w1_EFD, module.w3_EFD).bfloat16(), + ) + ) + self.assertTrue(torch.equal(w2_arg, module.w2_EDF.bfloat16())) + self.assertEqual(offs_arg.dtype, torch.int32) + self.assertEqual( + offs_arg.tolist(), torch.cumsum(num_tokens, dim=0).tolist() + ) + self.assertEqual(out.dtype, x.dtype) + + def test_plan_pads_the_dispatcher_to_256(self): + model_config = self._plan_model_config() + pads = [ + dispatcher_cfg.pad_multiple + for _fqn, dispatcher_cfg, _parent, _attr in model_config.traverse( + TorchAOTokenDispatcher.Config + ) + ] + self.assertTrue(pads) + self.assertTrue(all(pad == 256 for pad in pads)) + nodes = list(model_config.traverse(MXFP8FusedGroupedMLP.Config)) + self.assertTrue(nodes) + for _fqn, cfg, _parent, _attr in nodes: + self.assertEqual(cfg.fusion_plan, "grouped_gemm_swiglu") + + def test_grouped_factory_raises_on_unknown_fusion_plan(self): + cfg = self._stock_routed_experts_config() + with self.assertRaisesRegex(ValueError, "unknown fusion_plan"): + mxfp8_fused_grouped_mlp(cfg, fusion_plan="fully_fused") + + def test_grouped_factory_raises_on_unfused_activation_under_the_plan(self): + cfg = self._stock_routed_experts_config() + with self.assertRaisesRegex(ValueError, "always fuses the activation"): + mxfp8_fused_grouped_mlp( + cfg, fuse_activation=False, fusion_plan="grouped_gemm_swiglu" + ) + + def test_grouped_factory_raises_when_the_torchao_ops_are_unavailable(self): + cfg = self._stock_routed_experts_config() + with mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp." + "_TORCHAO_GROUPED_MLP_OPS_AVAILABLE", + False, + ), mock.patch( + "torchtitan.overrides.mxfp8_fused_mlp." + "_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON", + "unavailable for the test", + ): + with self.assertRaisesRegex(ValueError, "unavailable for the test"): + mxfp8_fused_grouped_mlp(cfg, fusion_plan="grouped_gemm_swiglu") + + +# --------------------------------------------------------------------------- +# grouped_gemm_swiglu numerics (SM100 + torchao fused grouped-MLP ops only). +# Ported from the plan's gate-proven suite, adapted to stock parameters. +# --------------------------------------------------------------------------- + +_BLOCK = 32 +_E4M3 = torch.float8_e4m3fn +_RCEIL = ScaleCalculationMode.RCEIL + +# Tolerances derived from the measured variability of the unfused lane itself +# (calibrated 2026-08-18 on GB200, torch 2.14.0a0): the two unfused MXFP8 +# lanes score >= 112.3 dB against each other, fp32 scores 23.65-23.70 dB +# against the quantized reference, and the composite's measured band vs the +# reference is 35.30-35.70 dB (the fused GLU/dGLU kernels evaluate SwiGLU +# from their in-kernel FP32 accumulators while the reference rounds z and dh +# to BF16 first -- that one boundary dominates). The gate sits 5.3 dB below +# the measured band floor and 6.3 dB above the "independent-but-correct lane" +# level, so it still discriminates; a real dataflow/layout/offsets bug lands +# near 0 dB. Kernel-level exactness (60-160 dB) is enforced separately by the +# torchao op suite. +_SQNR_VS_REFERENCE_DB = 30.0 +# Secondary tracking gate vs the FP32 eager MLP (measured 23.65-23.70 dB): +# catches a blind spot shared by both MXFP8 lanes. +_SQNR_VS_FP32_DB = 21.0 + +# Fixture shapes: per-expert row counts are 256-multiples (the dispatcher's +# pad_multiple=256 ABI guarantee), zero-token experts are legal anywhere, and +# `tail` allocates inactive rows past offsets[-1] (A < R). D != F in the +# second case so a wrong-axis weight cast cannot cancel. +_CASES = { + "debugmodel_tail_zero_expert": dict( + d=256, f=256, sizes=[256, 0, 256, 512, 0, 256], tail=256, seed=0 + ), + "asym_d_ne_f_tail": dict(d=256, f=512, sizes=[256, 0, 512], tail=256, seed=1), +} + + +def _blk_view(w13): + """[E, 2F, D] 32-block order -> view [E, F//32, 2, 32, D] with the + gate/up axis at dim 2.""" + e, two_f, d = w13.shape + return w13.view(e, two_f // 64, 2, 32, d) + + +def _to_blk(w1, w3): + """Stock [E, F, D] pairs -> 32-block [E, 2F, D] (independent + re-implementation; the module's pack is byte-verified against this + mapping in the wiring tests).""" + e, f, d = w1.shape + return ( + torch.stack([w1, w3], dim=2) + .view(e, f // 32, 32, 2, d) + .permute(0, 1, 3, 2, 4) + .reshape(e, 2 * f, d) + ) + + +def _zsplit(z, f): + """[R, 2F] in 32-block order -> (gate [R, F], up [R, F]).""" + r = z.shape[0] + v = z.view(r, f // 32, 2, 32) + return v[:, :, 0, :].reshape(r, f), v[:, :, 1, :].reshape(r, f) + + +def _zmerge(gate, up): + """(gate [R, F], up [R, F]) -> [R, 2F] in 32-block order.""" + r, f = gate.shape + out = torch.empty(r, 2 * f, dtype=gate.dtype, device=gate.device) + v = out.view(r, f // 32, 2, 32) + v[:, :, 0, :] = gate.view(r, f // 32, 32) + v[:, :, 1, :] = up.view(r, f // 32, 32) + return out + + +def _make_case(*, d, f, sizes, tail, seed=0): + """Dispatcher-shaped fixture: expert-major x [R, D] with per-expert row + counts ``sizes`` (256-multiples), offsets = inclusive cumsum, and a + strict inactive tail of ``tail`` rows. Tail rows carry large deliberate + garbage plus NaN (a NaN-poisoning attack on the undefined inactive tail): + producers do not define them, kernels must never let them contaminate + active rows, and y/dx comparisons mask them because the mm op leaves + output tail rows unwritten.""" + g = len(sizes) + a = sum(sizes) + r = a + tail + torch.manual_seed(seed) + offsets = torch.tensor( + [sum(sizes[: i + 1]) for i in range(g)], device="cuda", dtype=torch.int32 + ) + x = torch.randn(r, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + dy = torch.randn(r, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w1 = torch.randn(g, f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w3 = torch.randn(g, f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + w13 = _to_blk(w1, w3) + w2 = torch.randn(g, d, f, device="cuda", dtype=torch.bfloat16) / d**0.5 + if tail: + x[a:] = 12345.0 + dy[a:] = -6789.0 + x[a : a + tail // 2] = float("nan") + dy[a : a + tail // 2] = float("nan") + return dict(x=x, dy=dy, w13=w13, w2=w2, offsets=offsets, sizes=sizes, a=a, r=r) + + +# Independent quantized-unfused reference: standalone torchao RCEIL casts +# (``to_mx``) + raw ``torch._scaled_grouped_mm`` + eager SwiGLU +# forward/backward (the BF16 round of z precedes SwiGLU; the BF16 rounds of +# h/dz precede their quantizers); wgrads are colwise quant-dequant + fp32 +# matmul per expert. Built from first principles, sharing no code with the +# module under test. Operates directly on the 32-block ``w13 [G, 2F, D]``: +# rowwise quantization is per-row (row order is irrelevant) and colwise +# 32-blocks along 2F are pure-gate or pure-up in this order, so quantization +# boundaries match the composite exactly. + + +def _rceil_rowwise(t): + scale, q = to_mx(t, _E4M3, _BLOCK, scaling_mode=_RCEIL) + return q, to_blocked(scale) + + +def _rceil_rowwise_3d(w): + qs, sfs = zip(*(_rceil_rowwise(w[g]) for g in range(w.shape[0]))) + return torch.stack(list(qs)), torch.stack(list(sfs)) + + +def _rceil_colwise_3d(w): + """[G, N, K] -> qdata [G, N, K] stride (N*K, 1, N) quantized along N + + per-group blocked scales (the ``mat2`` of a dgrad ``_scaled_grouped_mm``).""" + qs, sfs = zip(*(_rceil_rowwise(w[g].t().contiguous()) for g in range(w.shape[0]))) + return torch.stack(list(qs)).transpose(-2, -1), torch.stack(list(sfs)) + + +def _dequant(q, scale): + m, k = q.shape + return ( + q.float().view(m, k // _BLOCK, _BLOCK) + * scale.to(torch.float32).view(m, k // _BLOCK, 1) + ).view(m, k) + + +def _quant_dequant_colwise(t): + """[m, N] bf16 -> fp32 [N, m]: RCEIL-quantize along the row axis (32x1) + and dequantize. Per-expert slices quantize identically to the whole + matrix because 256-multiple group sizes keep every 32-value block inside + one group.""" + scale, q = to_mx(t.t().contiguous(), _E4M3, _BLOCK, scaling_mode=_RCEIL) + return _dequant(q, scale) + + +def _wgrad_expert(a, b): + """Normative wgrad for one expert: dequant(a_col).T @ dequant(b_col), + fp32 accumulation, one BF16 round. a [m, N], b [m, K] -> [N, K].""" + return (_quant_dequant_colwise(a) @ _quant_dequant_colwise(b).t()).to( + torch.bfloat16 + ) + + +def _reference_forward_backward(x, w13, w2, dy, offsets, sizes, a): + """Returns (y, dx, dw13 [G, 2F, D] 32-block order, dw2 [G, D, F]). y/dx + tail rows [a:] are defined as zero here (the real ops leave them + unwritten; callers mask them out of every comparison).""" + r, d = x.shape + g, two_f = w13.shape[0], w13.shape[1] + f = two_f // 2 + + # FC1 forward, then eager SwiGLU on the BF16-rounded z. The gate/up split + # follows the 32-block column order z inherits from the w13 row order. + x_q, x_sf = _rceil_rowwise(x) + w13_row_q, w13_row_sf = _rceil_rowwise_3d(w13) + z = torch._scaled_grouped_mm( + x_q, + w13_row_q.transpose(-2, -1), + x_sf.reshape(r, -1), + w13_row_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + z[a:] = 0 + gate_bf16, up_bf16 = _zsplit(z, f) + gate = gate_bf16.float() + up = up_bf16.float() + h = (F.silu(gate) * up).to(torch.bfloat16) + + # FC2 forward. + h_q, h_sf = _rceil_rowwise(h) + w2_row_q, w2_row_sf = _rceil_rowwise_3d(w2) + y = torch._scaled_grouped_mm( + h_q, + w2_row_q.transpose(-2, -1), + h_sf.reshape(r, -1), + w2_row_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + y[a:] = 0 + + # FC2 dgrad, then eager dSwiGLU on the BF16-rounded dh. + dy_q, dy_sf = _rceil_rowwise(dy) + w2_col_q, w2_col_sf = _rceil_colwise_3d(w2) + dh = torch._scaled_grouped_mm( + dy_q, + w2_col_q, + dy_sf.reshape(r, -1), + w2_col_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + dh[a:] = 0 + sig = torch.sigmoid(gate) + silu_g = gate * sig + dsilu = sig * (1.0 + gate * (1.0 - sig)) + dhf = dh.float() + dgate = (dhf * up * dsilu).to(torch.bfloat16) + dup = (dhf * silu_g).to(torch.bfloat16) + dz = _zmerge(dgate, dup) + + # FC1 dgrad. + dz_q, dz_sf = _rceil_rowwise(dz) + w13_col_q, w13_col_sf = _rceil_colwise_3d(w13) + dx = torch._scaled_grouped_mm( + dz_q, + w13_col_q, + dz_sf.reshape(r, -1), + w13_col_sf.reshape(g, -1), + offs=offsets, + out_dtype=torch.bfloat16, + ) + dx[a:] = 0 + + # Wgrads over active rows only; zero-token experts stay all-zero. + dw13 = torch.zeros(g, two_f, d, device=x.device, dtype=torch.bfloat16) + dw2 = torch.zeros(g, d, f, device=x.device, dtype=torch.bfloat16) + prev = 0 + for gi in range(g): + end = int(offsets[gi]) + if end > prev: + dw13[gi] = _wgrad_expert(dz[prev:end], x[prev:end]) + dw2[gi] = _wgrad_expert(dy[prev:end], h[prev:end]) + prev = end + return y, dx, dw13, dw2 + + +def _fp32_reference(x, w13, w2, dy, offsets, a): + """FP32 eager autograd MLP over the active rows. Returns + (y, dx, dw13 [G, 2F, D] 32-block order, dw2).""" + g, two_f = w13.shape[0], w13.shape[1] + f = two_f // 2 + x32 = x[:a].float().detach().requires_grad_(True) + w13_32 = w13.float().detach().requires_grad_(True) + w2_32 = w2.float().detach().requires_grad_(True) + v = _blk_view(w13_32) + outs, prev = [], 0 + for gi in range(g): + end = int(offsets[gi]) + w1_g = v[gi, :, 0].reshape(f, x.shape[1]) + w3_g = v[gi, :, 1].reshape(f, x.shape[1]) + gate = x32[prev:end] @ w1_g.t() + up = x32[prev:end] @ w3_g.t() + h = F.silu(gate) * up + outs.append(h @ w2_32[gi].t()) + prev = end + y_ref = torch.cat(outs, dim=0) + y_ref.backward(dy[:a].float()) + return y_ref, x32.grad, w13_32.grad, w2_32.grad + + +_MOD_D, _MOD_F, _MOD_E = 256, 256, 4 +_MOD_SIZES = [256, 512, 256, 256] + + +def _module_inputs(seed=0): + torch.manual_seed(seed) + r = sum(_MOD_SIZES) + x = torch.randn(r, _MOD_D, device="cuda", dtype=torch.bfloat16) / _MOD_D**0.5 + dy = torch.randn(r, _MOD_D, device="cuda", dtype=torch.bfloat16) / _MOD_D**0.5 + num_tokens = torch.tensor(_MOD_SIZES, device="cuda") + return x, dy, num_tokens + + +def _run_module(module, x, dy, num_tokens): + x = x.clone().detach().requires_grad_(True) + y = module(x, num_tokens) + y.backward(dy) + return y.detach(), x.grad + + +@unittest.skipUnless(_HAS_SM100_GPU, "Requires CUDA SM 10.0 (Blackwell)") +@unittest.skipUnless( + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE, + "torchao fused grouped-MLP ops unavailable: " + f"{_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON}", +) +class TestGroupedGemmPlanNumerics(unittest.TestCase): + """Numerics and fail-loud gating of the ``grouped_gemm_swiglu`` composite + on real SM100 hardware with the torchao fused grouped-MLP ops.""" + + def test_composite_matches_quantized_unfused_reference(self): + for case in sorted(_CASES): + with self.subTest(case=case): + self._check_composite_case(**_CASES[case]) + + def _check_composite_case(self, **case): + fx = _make_case(**case) + a, r, d = fx["a"], fx["r"], fx["x"].shape[1] + + x = fx["x"].clone().detach().requires_grad_(True) + w13 = fx["w13"].clone().detach().requires_grad_(True) + w2 = fx["w2"].clone().detach().requires_grad_(True) + y = _MXFP8GroupedGemmMLP.apply(x, w13, w2, fx["offsets"]) + self.assertEqual(tuple(y.shape), (r, d)) + self.assertEqual(y.dtype, torch.bfloat16) + y.backward(fx["dy"]) + + ref_y, ref_dx, ref_dw13, ref_dw2 = _reference_forward_backward( + fx["x"], fx["w13"], fx["w2"], fx["dy"], fx["offsets"], fx["sizes"], a + ) + fp32 = _fp32_reference( + fx["x"], fx["w13"], fx["w2"], fx["dy"], fx["offsets"], a + ) + + # Zero-token experts must produce exactly-zero weight gradients + # through the autograd path (the wgrad op writes empty-group outputs + # as zero). + for gi, m in enumerate(fx["sizes"]): + if m == 0: + self.assertEqual(w13.grad[gi].abs().max().item(), 0.0) + self.assertEqual(w2.grad[gi].abs().max().item(), 0.0) + + # y/dx are compared over active rows only: both lanes leave the + # inactive tail [A, R) unwritten, and the garbage+NaN planted in the + # x/dy tails (the negative control) must not move -- or NaN-poison -- + # any active output. + for name, got, ref, hp in [ + ("y", y[:a], ref_y[:a], fp32[0]), + ("dx", x.grad[:a], ref_dx[:a], fp32[1]), + ("dw13", w13.grad, ref_dw13, fp32[2]), + ("dw2", w2.grad, ref_dw2, fp32[3]), + ]: + self.assertTrue( + torch.isfinite(got).all(), f"{name} contains non-finite values" + ) + sqnr = compute_error(ref.float(), got.float()) + self.assertGreaterEqual( + sqnr, + _SQNR_VS_REFERENCE_DB, + f"{name} SQNR vs quantized-unfused reference", + ) + sqnr_hp = compute_error(hp.float(), got.float()) + self.assertGreaterEqual(sqnr_hp, _SQNR_VS_FP32_DB, f"{name} SQNR vs fp32") + + def test_composite_zero_routed_tokens(self): + # A local expert set receiving zero routed tokens (R == 0, every + # offset 0): the cast chain must produce empty quantized operands and + # the ops' documented R == 0 early-outs must return empty y/dx and + # exactly-zero weight grads without error. + d, f, g = 256, 256, 3 + torch.manual_seed(8) + x = torch.zeros(0, d, device="cuda", dtype=torch.bfloat16, requires_grad=True) + w13 = ( + torch.randn(g, 2 * f, d, device="cuda", dtype=torch.bfloat16) / d**0.5 + ).requires_grad_(True) + w2 = ( + torch.randn(g, d, f, device="cuda", dtype=torch.bfloat16) / d**0.5 + ).requires_grad_(True) + offsets = torch.zeros(g, device="cuda", dtype=torch.int32) + + y = _MXFP8GroupedGemmMLP.apply(x, w13, w2, offsets) + self.assertEqual(tuple(y.shape), (0, d)) + self.assertEqual(y.dtype, torch.bfloat16) + y.backward(torch.zeros_like(y)) + + self.assertIsNotNone(x.grad) + self.assertEqual(tuple(x.grad.shape), (0, d)) + self.assertIsNotNone(w13.grad) + self.assertEqual(w13.grad.abs().max().item(), 0.0) + self.assertIsNotNone(w2.grad) + self.assertEqual(w2.grad.abs().max().item(), 0.0) + + def _build_plan_module(self, seed): + torch.manual_seed(seed) + module = MXFP8FusedGroupedMLP.Config( + dim=_MOD_D, + hidden_dim=_MOD_F, + num_experts=_MOD_E, + fusion_plan="grouped_gemm_swiglu", + ).build() + module = module.to("cuda") + with torch.no_grad(): + # fp32 master weights: the .bfloat16() casts stay outside the + # Function so autograd routes bf16 grads back to fp32 params. + module.w1_EFD.normal_(0.0, _MOD_D**-0.5) + module.w3_EFD.normal_(0.0, _MOD_D**-0.5) + module.w2_EDF.normal_(0.0, _MOD_D**-0.5) + return module + + def test_torch_compile_matches_eager_bitwise(self): + x, dy, num_tokens = _module_inputs() + module = self._build_plan_module(seed=9) + + y_eager, dx_eager = _run_module(module, x, dy, num_tokens) + grads_eager = [ + module.w1_EFD.grad.clone(), + module.w3_EFD.grad.clone(), + module.w2_EDF.grad.clone(), + ] + module.zero_grad(set_to_none=True) + + compiled = torch.compile(module) + y_comp, dx_comp = _run_module(compiled, x, dy, num_tokens) + + # Deterministic kernels + compile-invariant surrounding data movement + # (cumsum, pack, dtype casts) => bitwise-identical results. + self.assertTrue(torch.equal(y_comp, y_eager)) + self.assertTrue(torch.equal(dx_comp, dx_eager)) + for got, ref in zip( + [module.w1_EFD.grad, module.w3_EFD.grad, module.w2_EDF.grad], + grads_eager, + ): + self.assertTrue(torch.equal(got, ref)) + + def test_plan_checkpoint_keys_match_stock(self): + stock_nodes = list( + deepseek_v3_model_registry("debugmodel").model.traverse( + GroupedExperts.Config + ) + ) + self.assertTrue(stock_nodes) + model_config = deepseek_v3_model_registry("debugmodel").model + apply_overrides(OverrideConfig(imports=[_GROUPED_PLAN_IMPORT]), model_config) + plan_nodes = list(model_config.traverse(MXFP8FusedGroupedMLP.Config)) + self.assertTrue(plan_nodes) + with torch.device("meta"): + stock = stock_nodes[0][1].build() + fused = plan_nodes[0][1].build() + self.assertEqual(set(fused.state_dict().keys()), set(stock.state_dict().keys())) + + def test_factory_raises_on_converter_quantized_experts(self): + # The composite quantizes every grouped GEMM itself; layering it on + # the MXFP8 grouped-experts converter's output is a config error that + # must raise -- never a silent fallback to the converter's unfused + # path. + model_config = deepseek_v3_model_registry( + "debugmodel", + converters=[MXFP8GroupedExpertsConverter.Config(pad_multiple=128)], + ).model + nodes = list(model_config.traverse(RoutedExperts.Config)) + self.assertTrue(nodes) + with self.assertRaisesRegex(ValueError, "grouped-experts converter"): + mxfp8_fused_grouped_mlp(nodes[0][1], fusion_plan="grouped_gemm_swiglu") + + def test_factory_raises_on_unsupported_dims(self): + model_config = deepseek_v3_model_registry("debugmodel").model + node = list(model_config.traverse(RoutedExperts.Config))[0][1] + node.inner_experts.hidden_dim = 100 # not a 128-multiple + with self.assertRaisesRegex(ValueError, "is_supported"): + mxfp8_fused_grouped_mlp(node, fusion_plan="grouped_gemm_swiglu") + + def test_factory_raises_on_non_alltoall_dispatcher(self): + # hybridep's padded dispatcher is not validated for the fused + # kernels' 256-row contract; the factory must refuse it rather than + # swap or accept it. + model_config = deepseek_v3_model_registry( + "debugmodel", + moe_comm_backend="hybridep", + non_blocking_capacity_factor=1.0, + ).model + node = list(model_config.traverse(RoutedExperts.Config))[0][1] + with self.assertRaisesRegex(ValueError, "TorchAO padded"): + mxfp8_fused_grouped_mlp(node, fusion_plan="grouped_gemm_swiglu") + + if __name__ == "__main__": unittest.main() diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 8184dde4a4..33c1bdec47 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -124,6 +124,29 @@ def deepseek_v3_debugmodel_mxfp8_fused_mlp() -> Trainer.Config: return config +def deepseek_v3_debugmodel_mxfp8_grouped_gemm_mlp() -> Trainer.Config: + config = deepseek_v3_debugmodel() + # Routed experts via the MXFP8 fused-MLP override's fully-fused plan + # (fusion_plan="grouped_gemm_swiglu"): one composite of four fused + # torchao grouped-GEMM ops owns the whole routed-expert path, and the + # factory installs the pad_multiple=256 TorchAO dispatcher those kernels + # require, produced by the EP permute path, hence + # expert_parallel_degree=2. disable_cuda_graphs is required by the + # TorchAOTokenDispatcher under EP>1; moe_force_load_balance keeps the + # per-rank routed row count fixed (a constant R avoids fused-kernel JIT + # churn per newly seen R). + config.override.imports.append( + ( + "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_grouped_mlp", + {"fusion_plan": "grouped_gemm_swiglu"}, + ) + ) + config.parallelism = ParallelismConfig(expert_parallel_degree=2) + config.training.disable_cuda_graphs = True + config.debug.moe_force_load_balance = True + return config + + def deepseek_v3_debugmodel_hybridep() -> Trainer.Config: config = deepseek_v3_debugmodel() config.model_spec = model_registry( diff --git a/torchtitan/overrides/mxfp8_fused_mlp.py b/torchtitan/overrides/mxfp8_fused_mlp.py index 29c04f790e..3ff50cdac9 100644 --- a/torchtitan/overrides/mxfp8_fused_mlp.py +++ b/torchtitan/overrides/mxfp8_fused_mlp.py @@ -29,16 +29,33 @@ ``mxfp8_fused_mlp`` (dense ``FeedForward``) builds :class:`MXFP8FusedMLP`; ``mxfp8_fused_grouped_mlp`` (``RoutedExperts``) builds :class:`MXFP8FusedGroupedMLP` and swaps the token dispatcher for the padded -variant the grouped composite requires. Activate by naming the factories in -``--override.imports``; both accept a ``fuse_activation`` kwarg via -``(target, kwargs)`` imports entries. The composites quantize every GEMM -themselves, so do not combine these overrides with the MXFP8 linear / -grouped-experts converters (the factories raise). Additional MXFP8 fusion -paths (e.g. fully fused grouped MLPs) extend this module with their own -composite and factory. +variant the selected grouped composite requires. Activate by naming the +factories in ``--override.imports``; both accept a ``fuse_activation`` kwarg +(and the grouped factory ``fusion_plan``) via ``(target, kwargs)`` imports +entries. The composites quantize every GEMM themselves, so do not combine +these overrides with the MXFP8 linear / grouped-experts converters (the +factories raise). + +The grouped override selects its composite with the ``fusion_plan`` kwarg: + +* ``"swiglu"`` (default): the composite above -- ``torch._scaled_grouped_mm`` + FC1/FC2 around the SwiGLU+MXFP8 boundary. Token groups are padded to + multiples of 128. +* ``"grouped_gemm_swiglu"``: the whole expert MLP runs on the four fused + torchao grouped-GEMM ops (``torchao::mxfp8_grouped_gemm_swiglu_fwd`` / + ``mxfp8_grouped_gemm`` / ``mxfp8_grouped_gemm_dswiglu_bwd`` / + ``mxfp8_grouped_gemm_wgrad``), so the FC1 GEMM, the SwiGLU boundary, and + every activation quantization are in-kernel. The stock ``w1_EFD``/``w3_EFD`` + weights are packed at forward time into the kernels' 32-block GLU + row-ordered ``[E, 2F, D]`` operand, and token groups must be padded to + multiples of 256 (128-only splits corrupt silently and + nondeterministically). ``fuse_activation`` has no unfused arm here -- the + factory raises on ``fuse_activation=False`` -- and the factory raises + actionably when the torchao ops are unavailable. """ from dataclasses import dataclass +from typing import Literal import torch import torch.nn.functional as F @@ -46,6 +63,7 @@ from torchao.prototype.moe_training.kernels.mxfp8 import ( triton_mx_block_rearrange_2d_K_groups, + triton_mx_block_rearrange_per_group_3d, ) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cutedsl_kernels_available, @@ -56,12 +74,47 @@ _compute_dgrad_sm100, _compute_fwd_sm100, ) + +# Importing the wrapper module registers the four torchao:: custom ops the +# grouped_gemm_swiglu plan runs on; the cudnn python package is only imported +# lazily inside the op bodies at first launch. The module is newer than +# several torchao releases, so its absence must surface as the factory's +# actionable config-time error (with this reason), not an ImportError at +# override-import time. Tests skip on this flag too. +try: + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import ( + _mxfp8_grouped_mlp_kernels_available, + is_supported, + ) +except ImportError as _exc: + is_supported = None + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE = False + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON = ( + "the installed torchao has no torchao.prototype.moe_training.kernels." + f"mxfp8.cutedsl_grouped_mlp module (a torchao build that ships the " + f"fused grouped-MLP custom ops is required): {_exc}" + ) +else: + _TORCHAO_GROUPED_MLP_OPS_AVAILABLE = bool(_mxfp8_grouped_mlp_kernels_available) + _TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON = ( + "" + if _TORCHAO_GROUPED_MLP_OPS_AVAILABLE + else ( + "torchao fused grouped-MLP ops are unavailable in this " + "environment (needs the cudnn python package >= 1.27 with the " + "grouped_gemm_*_wrapper_sm100 kernels)." + ) + ) from torchao.prototype.mx_formats.config import ( MXFP8Dim1CastKernelChoice, ScaleCalculationMode, ) +from torchao.prototype.mx_formats.kernels import triton_to_mxfp8_dim0 from torchao.prototype.mx_formats.mx_tensor import MXTensor -from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper +from torchao.prototype.mx_formats.utils import ( + _to_mxfp8_dim1_kernel_wrapper, + to_blocked, +) from torchao.quantization.quantize_.common.kernel_preference import KernelPreference from torchtitan.components.quantization.utils import swap_token_dispatcher @@ -69,6 +122,10 @@ from torchtitan.models.common.feed_forward import FeedForward from torchtitan.models.common.linear import Linear from torchtitan.models.common.moe import GroupedExperts, RoutedExperts +from torchtitan.models.common.token_dispatcher import ( + AllToAllTokenDispatcher, + TorchAOTokenDispatcher, +) from torchtitan.tools.utils import has_cuda_capability __all__ = [ @@ -84,6 +141,12 @@ _KERNEL_PREFERENCE = KernelPreference.AUTO _SCALE_MODE = ScaleCalculationMode.RCEIL _INT32_MAX = 2**31 - 1 +# grouped_gemm_swiglu plan: per-expert row groups must be 256-multiples +# (the fused kernels hard-code FIX_PAD_SIZE = 256); feature dims must be +# 128-multiples (blocked-scale tiles). The row guarantee is the dispatcher's +# pad_multiple; the ops re-validate R % 256 statically. +_ROW_ALIGNMENT = 256 +_DIM_ALIGNMENT = 128 def _wrap_rowwise(qdata, scales, orig_dtype): @@ -446,6 +509,245 @@ def _validate_grouped_inputs(x, w13, w2_t, offs): torch._check(cond) +def _cast_rowwise(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """1x32 rowwise RCEIL cast: row-major qdata + whole-matrix blocked scales + (identical to the per-group concatenation because every per-expert row + count is a 256-multiple, so 128-row scale tiles never straddle groups).""" + qdata, scales = triton_to_mxfp8_dim0(t, _BLOCK_SIZE, _SCALE_MODE.value) + return qdata, to_blocked(scales) + + +def _cast_weight_rowwise_3d(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[G, N, K]`` quantized along K: contiguous qdata + per-group blocked + scales for logical ``(N, K/32)`` per group — the rowwise ``b`` operand of + the fwd/mm ops.""" + qdata, scales = triton_to_mxfp8_dim0(w, _BLOCK_SIZE, _SCALE_MODE.value) + return qdata, triton_mx_block_rearrange_per_group_3d(scales) + + +def _cast_weight_colwise_3d(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[G, N, K]`` quantized along N: k-major per-group qdata + + per-group blocked scales for logical ``(K, N/32)``. + + Batched: ONE (32x1 RCEIL) cast of the flat ``[G*N, K]`` view along dim0 + + ONE ``K_groups`` swizzle with uniform scale-column offsets. Exact + because N is a 256-multiple, so 32-row quantization blocks never + straddle groups, and every group's N/32 scale columns are 4-multiples, + so the swizzle packs the same per-group ``to_blocked`` bytes densely + from the buffer start. qdata, scales, AND downstream op outputs are + BITWISE-equal to a naive per-group ``triton_to_mxfp8_dim1`` + + ``to_blocked`` loop (measured on GB200) at ~4x less time and ~6*G + fewer launches per weight. + + The cast's native ``[G, N, K]`` view carries an interleaved batch + stride ``(N, 1, G*N)``, which the fused-op wrappers reject (B must be + per-group-contiguous, k- or n-major); one fp8 repack to k-major — the + same major the rowwise casts pass — restores an accepted layout. + """ + g, n, k = w.shape + mx = _to_mxfp8_dim1_kernel_wrapper( + w.reshape(g * n, k), + _BLOCK_SIZE, + elem_dtype=_ELEM_DTYPE, + hp_dtype=w.dtype, + kernel_preference=_KERNEL_PREFERENCE, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=_SCALE_MODE, + ) + scale_offsets = ( + torch.arange(1, g + 1, device=w.device, dtype=torch.int32) + * (n // _BLOCK_SIZE) + ) + # Same pow2 quirk as _cast_colwise_grouped: the K_groups swizzle's + # tl.arange needs a power-of-2 group count; repeated end-offsets are + # zero-size groups the kernel skips. + g_pow2 = 1 << (g - 1).bit_length() + if g_pow2 != g: + scale_offsets = torch.cat( + [scale_offsets, scale_offsets[-1:].expand(g_pow2 - g)] + ) + col_scales = triton_mx_block_rearrange_2d_K_groups(mx.scale, scale_offsets) + k_pad = -(-k // 128) * 128 + flat = col_scales.reshape(-1)[: k_pad * (g * n // _BLOCK_SIZE)] + qdata = mx.qdata.view(k, g, n).permute(1, 2, 0).contiguous() + return qdata, flat.view(g, -1) + + +def _cast_colwise_grouped( + t: torch.Tensor, offsets: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Ragged colwise (32x1) RCEIL cast of ``[R, N]`` for the wgrad operands: + torchao-native qdata (``[R, N]``-logical, ``(1, R)`` strides — the fused + wgrad kernel accepts this major directly, verified on GB200) + PER-GROUP + blocked scales via ``triton_mx_block_rearrange_2d_K_groups``. + + Quantizing the whole ragged tensor in one launch is safe ONLY because + every per-expert row count is a 256-multiple, so 32-row quantization + blocks never straddle an expert boundary. + + The K_groups swizzle's ``tl.arange(0, num_groups)`` requires a power-of-2 + bound, so the scale-column offsets are padded to the next power of 2 by + repeating the final offset — repeated end-offsets are zero-size groups the + kernel skips. Its output also carries 4 trailing padding columns per + group slot (d2h-sync avoidance); with 256-multiple groups the real blocks + pack densely from the start of the buffer (total real content = + ``round_up(N, 128) * offsets[-1]/32`` elements <= ``... * R/32``), so the + flat buffer is statically sliced to ``round_up(N, 128) * R/32`` — the op's + documented maximum — without any device sync; the wgrad kernel never reads + past the offsets-bounded span. + """ + mx = _to_mxfp8_dim1_kernel_wrapper( + t, + _BLOCK_SIZE, + elem_dtype=_ELEM_DTYPE, + hp_dtype=t.dtype, + kernel_preference=_KERNEL_PREFERENCE, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=_SCALE_MODE, + ) + scale_offsets = offsets // _BLOCK_SIZE + g = scale_offsets.shape[0] + g_pow2 = 1 << (g - 1).bit_length() + if g_pow2 != g: + scale_offsets = torch.cat( + [scale_offsets, scale_offsets[-1:].expand(g_pow2 - g)] + ) + col_scales = triton_mx_block_rearrange_2d_K_groups(mx.scale, scale_offsets) + r, n = t.shape + n_pad = -(-n // 128) * 128 + # mx.qdata is [N, R]-shaped; .t() presents the op's [R, N]-logical view. + return mx.qdata.t(), col_scales.reshape(-1)[: n_pad * (r // _BLOCK_SIZE)] + + +@torch._dynamo.allow_in_graph +class _MXFP8GroupedGemmMLP(torch.autograd.Function): + """Fully-fused MXFP8 grouped SwiGLU MLP over the four torchao fused + grouped-GEMM ops: ``x [R, D] -> y [R, D]`` (BF16). + + All inputs are plain BF16 CUDA tensors (the module prologue casts and + un-DTensors them): expert-major padded rows ``x [R, D]`` with every + per-expert group a multiple of 256 rows; ``w13 [G, 2F, D]`` in 32-block + GLU row order (32 gate rows, then the same features' 32 up rows, ...); + ``w2 [G, D, F]``; int32 CUDA ``offsets [G]`` exclusive per-expert end + rows, ``offsets[-1] <= R``. Rows past ``offsets[-1]`` of ``y`` (and of + ``dx`` in backward) are left UNWRITTEN; ``dy`` arrives as contiguous + BF16 ``[R, D]``. All backward-only casts are lazy: forward quantizes + only what forward consumes (the rowwise views); backward requantizes the + colwise weight views and the colwise ``x`` from the saved BF16 + references — safe because the same-step backward always precedes the + optimizer update (an update in between trips the autograd version + counter), and cheaper under per-op SAC because the forward (and thus any + forward-side cast) re-runs in the recompute pass. + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + offsets: torch.Tensor, + ) -> torch.Tensor: + x_row_q, x_row_sf = _cast_rowwise(x) + w13_row_q, w13_row_sf = _cast_weight_rowwise_3d(w13) + z, h_row_q, h_row_sf, h_col_q, h_col_sf = ( + torch.ops.torchao.mxfp8_grouped_gemm_swiglu_fwd( + x_row_q, + x_row_sf, + w13_row_q, + w13_row_sf.reshape(-1), + offsets, + ) + ) + w2_row_q, w2_row_sf = _cast_weight_rowwise_3d(w2) + # FC2 forward: b [G, N=D, K=F] rowwise (quantized along F = the + # contraction), row-major as cast. + y = torch.ops.torchao.mxfp8_grouped_gemm( + h_row_q, + h_row_sf, + w2_row_q, + w2_row_sf.reshape(-1), + offsets, + ) + # x is saved BF16; its colwise cast is deferred to backward. Under + # per-op SAC the whole forward re-runs in the recompute pass, so a + # forward-side cast would execute twice per step for one consumer + # (the FC1 wgrad) — deferring makes it run exactly once. Safe for the + # same reason the weight casts are lazy: the same-step backward always + # precedes the optimizer update. + ctx.save_for_backward(z, h_col_q, h_col_sf, x, offsets, w13, w2) + return y + + @staticmethod + def backward(ctx, dy: torch.Tensor): + z, h_col_q, h_col_sf, x, offsets, w13, w2 = ctx.saved_tensors + if x.shape[0] == 0: + # A rank whose local experts received zero routed tokens: every + # grad is zero by construction, and torchao's CUDA colwise cast + # rejects 0-row inputs, so skip the cast/GEMM chain outright. + # (The forward needs no such guard: its rowwise casts accept 0 + # rows and the ops early-return at R == 0.) + return ( + torch.empty_like(x), + torch.zeros_like(w13), + torch.zeros_like(w2), + None, + ) + # The casts assert contiguity; dy is contiguous today (BF16 [R, D] + # stride (D, 1)) but that is a live invariant, not a given. + dy = dy.contiguous() + dy_row_q, dy_row_sf = _cast_rowwise(dy) + # w2 colwise (quantized along D = the dgrad contraction): the bwd op's + # ABI takes the [G, D, F]-logical cast output as-is. + w2_col_q, w2_col_sf = _cast_weight_colwise_3d(w2) + dz_row_q, dz_row_sf, dz_col_q, dz_col_sf = ( + torch.ops.torchao.mxfp8_grouped_gemm_dswiglu_bwd( + dy_row_q, + dy_row_sf, + w2_col_q, + w2_col_sf.reshape(-1), + z, + offsets, + ) + ) + # FC1 dgrad: b [G, N=D, K=2F] quantized along 2F. The colwise cast + # yields [G, 2F, D]; the mm op's b orientation is [G, N, K], so the + # call site transposes (unlike ``torch._scaled_grouped_mm``, whose + # [G, K, N] mat2 convention would take the cast output as-is). + w13_col_q, w13_col_sf = _cast_weight_colwise_3d(w13) + dx = torch.ops.torchao.mxfp8_grouped_gemm( + dz_row_q, + dz_row_sf, + w13_col_q.transpose(-2, -1), + w13_col_sf.reshape(-1), + offsets, + ) + dy_col_q, dy_col_sf = _cast_colwise_grouped(dy, offsets) + x_col_q, x_col_sf = _cast_colwise_grouped(x, offsets) + # dw2 [G, D, F] = dy^T @ h per expert; dw13 [G, 2F, D] = dz^T @ x per + # expert, landing directly in the 32-block w13 operand order. + dw2 = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dy_col_q, dy_col_sf, h_col_q, h_col_sf, offsets + ) + dw13 = torch.ops.torchao.mxfp8_grouped_gemm_wgrad( + dz_col_q, dz_col_sf, x_col_q, x_col_sf, offsets + ) + return dx, dw13, dw2, None + + +def _pack_w13_blocks(w1: torch.Tensor, w3: torch.Tensor) -> torch.Tensor: + # Stock (E, F, D) gate/up pair packed to the fused kernels' 32-block GLU + # row-ordered (E, 2F, D): 32 gate rows, then the same features' 32 up + # rows, ... The reshape of the permuted view copies. + e, f, d = w1.shape + return ( + torch.stack([w1, w3], dim=2) # (E, F, 2, D) + .view(e, f // 32, 32, 2, d) + .permute(0, 1, 3, 2, 4) + .reshape(e, 2 * f, d) + ) + + class MXFP8FusedMLP(FeedForward): """Stock :class:`FeedForward` whose forward runs the composite MXFP8 SwiGLU MLP, stacking ``w1``/``w3`` into the fused ``w13`` operand. @@ -482,25 +784,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class MXFP8FusedGroupedMLP(GroupedExperts): - """Routed experts whose forward runs the composite MXFP8 SwiGLU grouped - MLP. - - Keeps the stock ``w1_EFD``/``w2_EDF``/``w3_EFD`` parameters and stacks - the gate and up weights into the composite's ``(num_experts, hidden_dim, - 2, dim)`` ``w13`` operand at forward time. Requires token groups padded - to multiples of 128 rows (zero-filled) -- the ``mxfp8_fused_grouped_mlp`` - factory swaps the token dispatcher accordingly. + """Routed experts whose forward runs a fusion-plan-selected MXFP8 SwiGLU + grouped-MLP composite. + + Keeps the stock ``w1_EFD``/``w2_EDF``/``w3_EFD`` parameters and packs the + gate and up weights into the selected composite's ``w13`` operand at + forward time: ``(num_experts, hidden_dim, 2, dim)`` element-interleaved + for the ``swiglu`` plan, ``(num_experts, 2 * hidden_dim, dim)`` in the + fused kernels' 32-block GLU row order for the ``grouped_gemm_swiglu`` + plan. Requires token groups padded to multiples of 128 (``swiglu``) or + 256 (``grouped_gemm_swiglu``) rows (zero-filled) -- the + ``mxfp8_fused_grouped_mlp`` factory swaps the token dispatcher + accordingly. """ @dataclass(kw_only=True, slots=True) class Config(GroupedExperts.Config): fuse_activation: bool = True """Quantize the SwiGLU boundary with the unified SwiGLU+MXFP8 kernel - (False: standalone BF16 + cast kernels; identical GEMMs either way).""" + (False: standalone BF16 + cast kernels; identical GEMMs either way). + Only meaningful for the ``swiglu`` fusion plan.""" + + fusion_plan: Literal["swiglu", "grouped_gemm_swiglu"] = "swiglu" + """How much of the expert MLP one kernel covers: ``swiglu`` fuses the + activation+quantization boundary between two + ``torch._scaled_grouped_mm`` GEMMs; ``grouped_gemm_swiglu`` runs the + whole MLP on the four fused torchao grouped-GEMM ops.""" def __init__(self, config: Config): super().__init__(config) self.fuse_activation = config.fuse_activation + self.fusion_plan = config.fusion_plan def forward( self, @@ -519,6 +833,40 @@ def forward( w3_EFD = self.w3_EFD offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) + if self.fusion_plan == "grouped_gemm_swiglu": + # The factory gate can only validate the GLOBAL dims (the config + # carries sharding placements, not mesh degrees), so the local + # shard dims are validated here at first call: under dense tensor + # parallelism (expert_parallel_degree=1, + # tensor_parallel_degree>1) the weights are Shard-split on + # hidden_dim, and a TP degree with hidden_dim/tp not a + # 128-multiple would otherwise fail deep inside the first fused + # op launch. + local_f, local_d = w1_EFD.shape[1], w1_EFD.shape[2] + if local_f % _DIM_ALIGNMENT != 0 or local_d % _DIM_ALIGNMENT != 0: + raise ValueError( + f"MXFP8FusedGroupedMLP: the LOCAL expert shard dims " + f"(D={local_d}, F={local_f} from w1_EFD of local shape " + f"{tuple(w1_EFD.shape)}) must be positive multiples of " + f"{_DIM_ALIGNMENT} for the fused grouped-MLP kernels. " + "This typically means tensor parallelism split " + "hidden_dim into a non-128-multiple shard; choose a " + "tensor_parallel_degree such that hidden_dim / tp stays " + f"a multiple of {_DIM_ALIGNMENT}, or drop the " + "grouped_gemm_swiglu fusion plan for this module." + ) + # The .bfloat16() casts stay OUTSIDE the Function so autograd + # handles high-precision master-weight configs and dy reaches + # backward() BF16; dw13 flows back through the pack to the stock + # parameters. + y_RD = _MXFP8GroupedGemmMLP.apply( + x_RD.bfloat16(), + _pack_w13_blocks(w1_EFD, w3_EFD).bfloat16(), + w2_EDF.bfloat16(), + offsets_E, + ) + return y_RD.type_as(x_RD) + x = x_RD.bfloat16() # (E, F, 2, D) with [:, :, 0] = gate (w1_EFD) and [:, :, 1] = up # (w3_EFD). @@ -577,13 +925,14 @@ def mxfp8_fused_mlp( @override( target=RoutedExperts.Config, - description="Routed experts via the composite MXFP8 grouped MLP " - "(128-row-padded token groups).", + description="Routed experts via a composite MXFP8 grouped MLP " + "(fusion_plan selects the fused boundary).", ) def mxfp8_fused_grouped_mlp( cfg: RoutedExperts.Config, *, fuse_activation: bool = True, + fusion_plan: Literal["swiglu", "grouped_gemm_swiglu"] = "swiglu", ) -> RoutedExperts.Config: # Config-application-time gate; the composite re-validates at runtime. if not has_cuda_capability(10, 0): @@ -591,9 +940,28 @@ def mxfp8_fused_grouped_mlp( "mxfp8_fused_grouped_mlp requires SM100 or later; remove the " "override or run on supported hardware." ) - # Targets RoutedExperts.Config because the composite constrains BOTH the - # experts and the token dispatcher: its kernels need every token group - # padded to a 128-row multiple (zero-filled), which only the padded + if fusion_plan not in ("swiglu", "grouped_gemm_swiglu"): + raise ValueError( + f"mxfp8_fused_grouped_mlp: unknown fusion_plan {fusion_plan!r}; " + "expected 'swiglu' or 'grouped_gemm_swiglu'." + ) + if fusion_plan == "grouped_gemm_swiglu": + if not fuse_activation: + raise ValueError( + "mxfp8_fused_grouped_mlp: fuse_activation=False (the " + "standalone-cast test-reference arm) exists only for the " + "'swiglu' fusion plan; the grouped_gemm_swiglu plan always " + "fuses the activation. Drop fuse_activation=False or use " + "fusion_plan='swiglu'." + ) + if not _TORCHAO_GROUPED_MLP_OPS_AVAILABLE: + raise ValueError( + "mxfp8_fused_grouped_mlp: " + f"{_TORCHAO_GROUPED_MLP_UNAVAILABLE_REASON}" + ) + # Targets RoutedExperts.Config because the composites constrain BOTH the + # experts and the token dispatcher: their kernels need every token group + # padded (zero-filled) to the plan's row multiple, which only the padded # dispatch path produces. if type(cfg) is not RoutedExperts.Config: raise ValueError( @@ -611,9 +979,40 @@ def mxfp8_fused_grouped_mlp( "grouped-experts converter." ) - swap_token_dispatcher(cfg, pad_multiple=128) + if fusion_plan == "grouped_gemm_swiglu": + # GLOBAL dims only: the config carries sharding placements (e.g. a + # TP shard on hidden_dim) but not mesh degrees, so the per-rank + # shard dims cannot be computed here. MXFP8FusedGroupedMLP.forward + # re-validates the LOCAL dims at first call and raises with the + # config fix. + if not is_supported(inner.dim, inner.hidden_dim): + raise ValueError( + f"mxfp8_fused_grouped_mlp: is_supported(D={inner.dim}, " + f"F={inner.hidden_dim}) is False; both dims must be positive " + f"multiples of {_DIM_ALIGNMENT}." + ) + # The %256 padding contract is the factory's own work: the fused + # kernels require per-expert groups padded to 256 (FIX_PAD_SIZE) -- + # 128-multiple-only splits corrupt silently and nondeterministically. + dispatcher = cfg.token_dispatcher + if isinstance(dispatcher, TorchAOTokenDispatcher.Config): + dispatcher.pad_multiple = _ROW_ALIGNMENT + elif isinstance(dispatcher, AllToAllTokenDispatcher.Config): + swap_token_dispatcher(cfg, pad_multiple=_ROW_ALIGNMENT) + else: + raise ValueError( + f"mxfp8_fused_grouped_mlp: token_dispatcher is " + f"{type(dispatcher).__qualname__}; only the TorchAO padded " + "dispatcher (swapped in from the stock all-to-all) is " + "validated for the per-expert 256-row contract." + ) + else: + swap_token_dispatcher(cfg, pad_multiple=128) cfg.inner_experts = derive( - inner, MXFP8FusedGroupedMLP.Config, fuse_activation=fuse_activation + inner, + MXFP8FusedGroupedMLP.Config, + fuse_activation=fuse_activation, + fusion_plan=fusion_plan, ) return cfg