diff --git a/tests/unit_tests/test_mxfp8_swiglu_override.py b/tests/unit_tests/test_mxfp8_fused_mlp.py similarity index 65% rename from tests/unit_tests/test_mxfp8_swiglu_override.py rename to tests/unit_tests/test_mxfp8_fused_mlp.py index db5c9e80b5..786eb4610f 100644 --- a/tests/unit_tests/test_mxfp8_swiglu_override.py +++ b/tests/unit_tests/test_mxfp8_fused_mlp.py @@ -4,15 +4,11 @@ # 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-SwiGLU overrides. - -The overrides must produce the MXFP8 fused modules with the right config -fields and (for the grouped path) the padded token dispatcher; the factories -must fail loud on non-stock configs. 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 (hardware is irrelevant to the transforms under -test). Numerics of the underlying composites are validated on SM100 hardware -in NVIDIA-internal CI. +"""Wiring 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. """ import unittest @@ -27,42 +23,37 @@ 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 -from torchtitan.overrides.fused_swiglu import FusedSwiGLU try: - from torchtitan.overrides.mxfp8_fused_swiglu import ( - mxfp8_fused_swiglu, - MXFP8FusedGroupedExperts, - MXFP8FusedSwiGLU, + from torchtitan.overrides.mxfp8_fused_mlp import ( + mxfp8_fused_mlp, + MXFP8FusedGroupedMLP, + MXFP8FusedMLP, ) except ImportError as e: # torchao (or a transitive dep) not installed raise unittest.SkipTest( - f"torchao is required for the MXFP8 SwiGLU overrides: {e}" + f"torchao is required for the MXFP8 fused-MLP overrides: {e}" ) from e -_DENSE_OVERRIDE = "torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_swiglu" -_GROUPED_OVERRIDE = ( - "torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_grouped_experts" -) +_DENSE_OVERRIDE = "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_mlp" +_GROUPED_OVERRIDE = "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_grouped_mlp" -class TestMXFP8FusedSwiGLUOverride(unittest.TestCase): +class TestMXFP8FusedMLPOverride(unittest.TestCase): def setUp(self): - # The factories gate on SM100 at config-application time; hardware is - # irrelevant to the config-tree transforms under test. patcher = mock.patch( - "torchtitan.overrides.mxfp8_fused_swiglu.has_cuda_capability", + "torchtitan.overrides.mxfp8_fused_mlp.has_cuda_capability", lambda *args: True, ) patcher.start() self.addCleanup(patcher.stop) - def test_dense_override_builds_mxfp8_fused_swiglu(self): + def test_dense_override_builds_mxfp8_fused_mlp(self): model_config = llama3_model_registry("debugmodel").model apply_overrides(OverrideConfig(imports=[_DENSE_OVERRIDE]), model_config) with torch.device("meta"): model = model_config.build() - fused = [m for m in model.modules() if isinstance(m, MXFP8FusedSwiGLU)] + fused = [m for m in model.modules() if isinstance(m, MXFP8FusedMLP)] self.assertTrue(fused) self.assertTrue(all(m.fuse_activation for m in fused)) @@ -74,7 +65,7 @@ def test_dense_override_kwargs_configure_the_composite(self): ) with torch.device("meta"): model = model_config.build() - fused = [m for m in model.modules() if isinstance(m, MXFP8FusedSwiGLU)] + fused = [m for m in model.modules() if isinstance(m, MXFP8FusedMLP)] self.assertTrue(fused) self.assertFalse(any(m.fuse_activation for m in fused)) @@ -84,7 +75,7 @@ def _grouped_model_config(self): return model_config def _grouped_experts_config(self, model_config): - nodes = list(model_config.traverse(MXFP8FusedGroupedExperts.Config)) + nodes = list(model_config.traverse(MXFP8FusedGroupedMLP.Config)) self.assertTrue(nodes) return nodes[0][1] @@ -100,9 +91,9 @@ def test_grouped_override_builds_experts_and_padded_dispatcher(self): self.assertTrue(all(pad == 128 for pad in pads)) with torch.device("meta"): model = model_config.build() - fused = [m for m in model.modules() if isinstance(m, MXFP8FusedGroupedExperts)] + fused = [m for m in model.modules() if isinstance(m, MXFP8FusedGroupedMLP)] self.assertTrue(fused) - self.assertTrue(all(type(m) is MXFP8FusedGroupedExperts for m in fused)) + self.assertTrue(all(type(m) is MXFP8FusedGroupedMLP for m in fused)) self.assertTrue(all(m.fuse_activation for m in fused)) def test_grouped_forward_validates_and_applies_the_function(self): @@ -113,9 +104,9 @@ def test_grouped_forward_validates_and_applies_the_function(self): x = torch.randn(2, cfg.dim) sentinel = torch.zeros(2, cfg.dim, dtype=torch.bfloat16) with mock.patch( - "torchtitan.overrides.mxfp8_fused_swiglu._validate_grouped_inputs" + "torchtitan.overrides.mxfp8_fused_mlp._validate_grouped_inputs" ) as validate, mock.patch( - "torchtitan.overrides.mxfp8_fused_swiglu._MXFP8SwiGLUGroupedMLP.apply", + "torchtitan.overrides.mxfp8_fused_mlp._MXFP8GroupedMLP.apply", return_value=sentinel, ) as function: out = module(x, num_tokens) @@ -134,7 +125,7 @@ def test_grouped_forward_validates_and_applies_the_function(self): validate.assert_called_once_with(args[0], args[1], args[2], args[3]) self.assertEqual(out.dtype, x.dtype) - def test_grouped_checkpoint_keys_and_param_shapes_unchanged(self): + def test_grouped_checkpoint_keys_unchanged(self): stock_nodes = list( deepseek_v3_model_registry("debugmodel").model.traverse( GroupedExperts.Config @@ -146,25 +137,51 @@ def test_grouped_checkpoint_keys_and_param_shapes_unchanged(self): stock = stock_nodes[0][1].build() fused = fused_cfg.build() self.assertEqual(set(fused.state_dict().keys()), set(stock.state_dict().keys())) - self.assertEqual( - tuple(fused.w13.shape), - (fused_cfg.num_experts, fused_cfg.hidden_dim, 2, fused_cfg.dim), - ) + + def test_fresh_init_matches_stock_bitwise(self): + # Stock parameters in stock registration order: fresh-init draws must + # be bitwise-identical to the corresponding stock module's. + def seeded_state_dict(cfg): + torch.manual_seed(42) + module = cfg.build() + module.init_states() + return module.state_dict() + + dense_stock = list( + llama3_model_registry("debugmodel").model.traverse(FeedForward.Config) + )[0][1] + dense_model = llama3_model_registry("debugmodel").model + apply_overrides(OverrideConfig(imports=[_DENSE_OVERRIDE]), dense_model) + dense_fused = list(dense_model.traverse(MXFP8FusedMLP.Config))[0][1] + grouped_stock = list( + deepseek_v3_model_registry("debugmodel").model.traverse( + GroupedExperts.Config + ) + )[0][1] + grouped_fused = self._grouped_experts_config(self._grouped_model_config()) + for stock_cfg, fused_cfg in ( + (dense_stock, dense_fused), + (grouped_stock, grouped_fused), + ): + stock_sd = seeded_state_dict(stock_cfg) + fused_sd = seeded_state_dict(fused_cfg) + self.assertEqual(set(fused_sd), set(stock_sd)) + for key, stock_tensor in stock_sd.items(): + self.assertTrue(torch.equal(fused_sd[key], stock_tensor), key) def test_dense_factory_raises_on_non_stock_ffn(self): - # A FeedForward.Config SUBCLASS (already fused) must raise, not no-op. + # A FeedForward.Config SUBCLASS (already overridden) must raise, not + # no-op. gate = Linear.Config(in_features=128, out_features=256) - cfg = FusedSwiGLU.Config( + cfg = MXFP8FusedMLP.Config( w1=gate, w2=Linear.Config(in_features=256, out_features=128), w3=gate, ) with self.assertRaisesRegex(ValueError, "stock FeedForward.Config"): - mxfp8_fused_swiglu(cfg) + mxfp8_fused_mlp(cfg) def test_dense_factory_raises_on_converted_projection(self): - # The composite quantizes every GEMM itself; combining with a linear - # quantization converter on the same module must raise. from torchtitan.components.quantization.mx import MXFP8Linear if MXFP8Linear is None: @@ -176,7 +193,7 @@ def test_dense_factory_raises_on_converted_projection(self): w3=gate, ) with self.assertRaisesRegex(ValueError, "quantization converter"): - mxfp8_fused_swiglu(cfg) + mxfp8_fused_mlp(cfg) if __name__ == "__main__": diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index eb26e320d8..8184dde4a4 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -109,22 +109,18 @@ def deepseek_v3_debugmodel_mxfp8() -> Trainer.Config: return config -def deepseek_v3_debugmodel_mxfp8_fused_swiglu() -> Trainer.Config: +def deepseek_v3_debugmodel_mxfp8_fused_mlp() -> Trainer.Config: config = deepseek_v3_debugmodel() - # Routed experts via the self-contained MXFP8 fused-SwiGLU override: one - # composite runs the whole expert MLP (both grouped GEMMs and the SwiGLU - # boundary) in MXFP8. The override swaps the token dispatcher for the - # padded variant its kernels require (128-row token groups), and that - # padded dispatch is only produced by the EP permute path, hence - # expert_parallel_degree=2. No quantization converter is needed: the - # composite quantizes every GEMM itself. + # Routed experts via the self-contained MXFP8 fused-MLP override: one + # composite runs the whole expert MLP in MXFP8, no quantization converter + # needed. The override swaps the token dispatcher for the 128-row-padded + # variant its kernels require, which only the EP permute path produces, + # hence expert_parallel_degree=2. config.compile = CompileConfig(enable=True, components=["model"]) config.override.imports.append( - "torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_grouped_experts" - ) - config.parallelism = ParallelismConfig( - expert_parallel_degree=2, + "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_grouped_mlp" ) + config.parallelism = ParallelismConfig(expert_parallel_degree=2) return config diff --git a/torchtitan/models/llama3/config_registry.py b/torchtitan/models/llama3/config_registry.py index 0d3d005ec5..31804ed2bf 100644 --- a/torchtitan/models/llama3/config_registry.py +++ b/torchtitan/models/llama3/config_registry.py @@ -264,15 +264,14 @@ def llama3_8b_mxfp8() -> Trainer.Config: return config -def llama3_debugmodel_mxfp8_fused_swiglu() -> Trainer.Config: +def llama3_debugmodel_mxfp8_fused_mlp() -> Trainer.Config: config = llama3_debugmodel() - # Dense FFN via the self-contained MXFP8 fused-SwiGLU override: one - # composite runs the whole MLP (both GEMMs and the SwiGLU boundary) in - # MXFP8. No quantization converter is needed: the composite quantizes - # every GEMM itself; attention and lm_head stay BF16. + # Dense FFN via the self-contained MXFP8 fused-MLP override: one composite + # runs the whole MLP in MXFP8, no quantization converter needed; attention + # and lm_head stay BF16. config.compile = CompileConfig(enable=True, components=["model"]) config.override.imports.append( - "torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_swiglu" + "torchtitan.overrides.mxfp8_fused_mlp.mxfp8_fused_mlp" ) return config diff --git a/torchtitan/overrides/mxfp8_fused_swiglu.py b/torchtitan/overrides/mxfp8_fused_mlp.py similarity index 66% rename from torchtitan/overrides/mxfp8_fused_swiglu.py rename to torchtitan/overrides/mxfp8_fused_mlp.py index 1939b347d6..29c04f790e 100644 --- a/torchtitan/overrides/mxfp8_fused_swiglu.py +++ b/torchtitan/overrides/mxfp8_fused_mlp.py @@ -6,53 +6,40 @@ # pyrefly: ignore-errors -"""Composite MXFP8 SwiGLU MLP for a fused w13 projection. +"""Composite MXFP8 MLP overrides. -One autograd function covers the full dense MLP +One autograd function covers the full SwiGLU MLP x -> MXFP8 w13 GEMM -> [gate | up] -> silu(gate) * up -> MXFP8 w2 GEMM with both directions of every quantization done by the CuTeDSL kernels. The +modules keep the stock parameters (``w1``/``w2``/``w3`` submodules dense, +``w1_EFD``/``w2_EDF``/``w3_EFD`` grouped) -- checkpoints, initialization, and +sharding are exactly the stock modules' -- and stack the gate and up weights +into the composite's fused ``w13`` operand at forward time. The ``fuse_activation`` flag selects how the activation boundary is quantized: - -* ``True``: the unified SwiGLU+MXFP8 kernel produces the rowwise and colwise - MXFP8 copies directly; the BF16 activation ``h`` is never written to global - memory. -* ``False``: ``h`` (forward) and ``[dGate | dUp]`` (backward) are materialized - in BF16 and quantized by the standalone 1x32 / 32x1 CuTeDSL kernels. - -Everything outside that boundary -- the w13/w2 GEMMs, their input, weight and -gradient casts -- is byte-for-byte identical between the two modes, so an A/B -comparison isolates the activation+quantization implementation. - -There is no silent fallback: configurations the kernels cannot execute -(missing CuTeDSL runtime, DTensor operands, non-BF16 dtypes, or dimensions -violating the kernels' 128-alignment contract) raise an actionable error so -the caller can change the config -- e.g. narrow the override's ``fqns`` or -drop it for the offending module -- rather than train silently on a -different numerical path. - -Two self-contained overrides wire the composites into a model: - -* ``mxfp8_fused_swiglu`` (dense ``FeedForward``) builds - :class:`MXFP8FusedSwiGLU`, a :class:`FusedSwiGLU` whose forward runs the - dense composite. -* ``mxfp8_fused_grouped_experts`` (``RoutedExperts``) builds - :class:`MXFP8FusedGroupedExperts` and swaps the token dispatcher for the - padded variant the grouped composite requires (``pad_multiple=128``). - -Activate by naming the factories, e.g. ``--override.imports -torchtitan.overrides.mxfp8_fused_swiglu.mxfp8_fused_swiglu``; both accept a -``fuse_activation`` kwarg via ``(target, kwargs)`` imports entries. The -composites quantize every GEMM themselves, so these overrides must not be -combined with the MXFP8 linear / grouped-experts converters on the same -modules (the factories raise if they are). +``True`` runs the unified SwiGLU+MXFP8 kernel (the BF16 activation never +reaches global memory); ``False`` materializes it in BF16 and quantizes with +the standalone 1x32 / 32x1 kernels. Everything outside that boundary is +identical between the two modes, so an A/B comparison isolates the fused +kernel. Configurations the kernels cannot execute (missing CuTeDSL runtime, +DTensor operands, non-BF16 dtypes, dimensions violating the 128-alignment +contract) raise an actionable error; there is no silent fallback. + +``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. """ from dataclasses import dataclass -import spmd_types as spmd - import torch import torch.nn.functional as F from torch.distributed.tensor import DTensor @@ -79,26 +66,17 @@ from torchtitan.components.quantization.utils import swap_token_dispatcher from torchtitan.config import derive, override -from torchtitan.models.common.decoder_sharding import dense_param_placement 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.overrides.fused_swiglu import ( - _fuse_w13_grouped_experts_param_init, - _fuse_w13_grouped_experts_sharding, - _make_fused_gate_up_init, - FusedGroupedExperts, - FusedSwiGLU, -) -from torchtitan.protocols.sharding import ShardingConfig from torchtitan.tools.utils import has_cuda_capability __all__ = [ - "MXFP8FusedGroupedExperts", - "MXFP8FusedSwiGLU", - "mxfp8_fused_grouped_experts", - "mxfp8_fused_swiglu", - "mxfp8_swiglu_mlp_w13", + "MXFP8FusedGroupedMLP", + "MXFP8FusedMLP", + "mxfp8_fused_grouped_mlp", + "mxfp8_fused_mlp", + "mxfp8_mlp_w13", ] _BLOCK_SIZE = 32 @@ -208,7 +186,7 @@ def _swiglu_backward_casts(grad_h, gated, fuse_activation): @torch._dynamo.allow_in_graph -class _MXFP8SwiGLUMLP(torch.autograd.Function): +class _MXFP8MLP(torch.autograd.Function): @staticmethod def forward(ctx, x, w13, w2, fuse_activation): x2d = x.reshape(-1, x.shape[-1]).contiguous() @@ -251,15 +229,15 @@ def _require_kernels(op_name): def _validate_dense_inputs(x, w13, w2): - _require_kernels("mxfp8_swiglu_mlp_w13") + _require_kernels("mxfp8_mlp_w13") if isinstance(x, DTensor) or isinstance(w13, DTensor) or isinstance(w2, DTensor): raise ValueError( - "mxfp8_swiglu_mlp_w13 takes plain local tensors, not DTensor; pass " + "mxfp8_mlp_w13 takes plain local tensors, not DTensor; pass " "local shards or exclude this module from the fused MXFP8 path." ) if not x.is_cuda: raise ValueError( - f"mxfp8_swiglu_mlp_w13 requires CUDA tensors, got device {x.device}" + f"mxfp8_mlp_w13 requires CUDA tensors, got device {x.device}" ) if ( x.dtype != torch.bfloat16 @@ -267,7 +245,7 @@ def _validate_dense_inputs(x, w13, w2): or w2.dtype != torch.bfloat16 ): raise ValueError( - "mxfp8_swiglu_mlp_w13 requires BF16 inputs and weights, got " + "mxfp8_mlp_w13 requires BF16 inputs and weights, got " f"x={x.dtype}, w13={w13.dtype}, down_weight={w2.dtype}" ) if w13.ndim != 3 or w13.shape[1] != 2 or w2.ndim != 2: @@ -290,11 +268,10 @@ def _validate_dense_inputs(x, w13, w2): "exclude this module from the fused MXFP8 path if its shapes cannot " "satisfy this." ) - # 32-bit index-math limit over BOTH A/B arms: the unified kernel's input - # layout reaches element 2*hidden*m - hidden - 1, but the unfused arm's - # standalone casts of the (m, 2*hidden) backward tensor reach - # 2*hidden*m - 1, and those kernels do not validate. Gate on the max so - # the two arms accept identical shapes. + # 32-bit index-math limit over BOTH A/B arms: the unfused arm's standalone + # casts of the (m, 2*hidden) backward tensor reach element 2*hidden*m - 1, + # past the unified kernel's own bound, and those kernels do not validate; + # gating on the max keeps the two arms' accepted shapes identical. if 2 * hidden * m - 1 > _INT32_MAX: raise ValueError( "tokens*hidden exceeds the kernels' 32-bit index math: " @@ -302,28 +279,17 @@ def _validate_dense_inputs(x, w13, w2): ) -def mxfp8_swiglu_mlp_w13(x, w13, down_weight, *, fuse_activation=True): +def mxfp8_mlp_w13(x, w13, down_weight, *, fuse_activation=True): """Dense MXFP8 SwiGLU MLP with a fused (H, 2, D) w13 weight. - Args: - x: BF16 input of shape (..., D). - w13: BF16 fused gate/up weight of shape (H, 2, D); w13[:, 0] is the - gate (w1) and w13[:, 1] the up (w3) projection. - down_weight: BF16 down-projection weight of shape (D_out, H). - fuse_activation: quantize the SwiGLU boundary with the unified - SwiGLU+MXFP8 kernel instead of standalone BF16 + cast kernels. - - Returns: - BF16 tensor of shape (..., D_out). - - Raises: - NotImplementedError: the MXFP8 CuTeDSL kernels are unavailable. - ValueError: DTensor operands, non-BF16 dtypes, or dimensions the - kernels cannot execute (every dim must be a multiple of 128). - There is no silent fallback; change the config instead. + ``x`` is BF16 of shape (..., D); ``w13[:, 0]`` is the gate (w1) and + ``w13[:, 1]`` the up (w3) projection; ``down_weight`` is (D_out, H). + Returns a BF16 tensor of shape (..., D_out). Raises instead of falling + back when the kernels are unavailable or the inputs violate their + contract (every dimension must be a multiple of 128). """ _validate_dense_inputs(x, w13, down_weight) - return _MXFP8SwiGLUMLP.apply(x, w13, down_weight, fuse_activation) + return _MXFP8MLP.apply(x, w13, down_weight, fuse_activation) def _pack_w13_grouped(w13): @@ -385,7 +351,7 @@ def _wgrad_k_groups(a_qdata, a_scales, b, offs, out_dtype): @torch._dynamo.allow_in_graph -class _MXFP8SwiGLUGroupedMLP(torch.autograd.Function): +class _MXFP8GroupedMLP(torch.autograd.Function): @staticmethod def forward(ctx, x, w13, w2_t, offs, fuse_activation): x = x.contiguous() @@ -432,11 +398,10 @@ def backward(ctx, grad_out): def _validate_grouped_inputs(x, w13, w2_t, offs): - # The only caller is MXFP8FusedGroupedExperts.forward, which guarantees - # plain local BF16 tensors in the module's own (M, D) / (E, F, 2, D) / - # (E, F, D_out) shapes; only environment, config dims, and the - # routing-dependent token count need checking. - _require_kernels("MXFP8FusedGroupedExperts") + # The only caller is MXFP8FusedGroupedMLP.forward, which guarantees plain + # local BF16 tensors in the module's own shapes; only environment, config + # dims, and the routing-dependent token count need checking. + _require_kernels("MXFP8FusedGroupedMLP") _, f, _, d = w13.shape m = x.shape[0] d_out = w2_t.shape[2] @@ -447,20 +412,17 @@ def _validate_grouped_inputs(x, w13, w2_t, offs): "this module from the fused MXFP8 path if its shapes cannot " "satisfy this." ) - # Group boundaries must additionally be 128-row aligned (the token - # dispatcher's pad_multiple guarantees it); checking offs here would sync. - # M is routing-dependent under compile (an unbacked SymInt, which type - # tests cannot tell apart from int inside traced code), so the M + # Group boundaries must additionally be 128-row aligned (the dispatcher's + # pad_multiple guarantees it); checking offs here would sync. M is + # routing-dependent under compile (an unbacked SymInt), so the M # conditions use identity tests: literal bools raise immediately, # symbolic ones become deferred runtime asserts. The m >= 128 and m % 32 - # forms are redundant with m % 128 (plus non-emptiness) but must be - # recorded separately: downstream cast-kernel wrappers and GEMM metas - # check exactly those forms, and the symbolic engine resolves them by - # expression match / value range, not by deriving them from mod-128. - # The last condition is the 32-bit index-math limit over BOTH A/B arms - # (the unfused arm's standalone casts of the (m, 2f) backward tensor - # reach element 2*f*m - 1, slightly past the unified kernel's own input - # bound, and those kernels do not validate). + # forms are redundant with m % 128 but recorded separately: downstream + # kernel wrappers and GEMM metas check exactly those forms, and the + # symbolic engine matches expressions rather than deriving them from + # mod-128. The last condition is the 32-bit index-math limit over BOTH + # A/B arms (the unfused arm's standalone casts reach element 2*f*m - 1, + # and those kernels do not validate). for cond, requirement in ( (m >= 128, "at least 128"), ( @@ -477,23 +439,20 @@ def _validate_grouped_inputs(x, w13, w2_t, offs): ): if cond is False: raise ValueError( - f"MXFP8FusedGroupedExperts: token count {m} (hidden={f}) " + f"MXFP8FusedGroupedMLP: token count {m} (hidden={f}) " f"must be {requirement}; there is no silent fallback." ) if cond is not True: torch._check(cond) -class MXFP8FusedSwiGLU(FusedSwiGLU): - """:class:`FusedSwiGLU` whose forward runs the composite MXFP8 SwiGLU MLP. - - Inherits the fused ``w13`` parameter, the stock-layout checkpoint hooks, - and the FSDP/TP sharding story from :class:`FusedSwiGLU`; only ``forward`` - changes. +class MXFP8FusedMLP(FeedForward): + """Stock :class:`FeedForward` whose forward runs the composite MXFP8 + SwiGLU MLP, stacking ``w1``/``w3`` into the fused ``w13`` operand. """ @dataclass(kw_only=True, slots=True) - class Config(FusedSwiGLU.Config): + class Config(FeedForward.Config): fuse_activation: bool = True """Quantize the SwiGLU boundary with the unified SwiGLU+MXFP8 kernel (False: standalone BF16 + cast kernels; identical GEMMs either way).""" @@ -505,13 +464,15 @@ def __init__(self, config: Config): def forward(self, x: torch.Tensor) -> torch.Tensor: if isinstance(x, DTensor): raise ValueError( - "MXFP8FusedSwiGLU does not support DTensor activations (dense " - "tensor parallelism); narrow the mxfp8_fused_swiglu override's " + "MXFP8FusedMLP does not support DTensor activations (dense " + "tensor parallelism); narrow the mxfp8_fused_mlp override's " "fqns or drop it for this module." ) - output = mxfp8_swiglu_mlp_w13( + # (H, 2, D) with [:, 0] = gate (w1) and [:, 1] = up (w3). + w13 = torch.stack([self.w1.weight, self.w3.weight], dim=1) + output = mxfp8_mlp_w13( x, - self.w13, + w13, self.w2.weight, fuse_activation=self.fuse_activation, ) @@ -520,18 +481,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return output -class MXFP8FusedGroupedExperts(FusedGroupedExperts): - """:class:`FusedGroupedExperts` whose forward runs the composite MXFP8 - SwiGLU grouped MLP. +class MXFP8FusedGroupedMLP(GroupedExperts): + """Routed experts whose forward runs the composite MXFP8 SwiGLU grouped + MLP. - Requires token groups padded to multiples of 128 rows (zero-filled) -- - the ``mxfp8_fused_grouped_experts`` factory swaps the token dispatcher - accordingly. Inherits ``w13``, checkpoint hooks, and sharding from - :class:`FusedGroupedExperts`. + 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. """ @dataclass(kw_only=True, slots=True) - class Config(FusedGroupedExperts.Config): + 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).""" @@ -545,20 +507,25 @@ def forward( x_RD: torch.Tensor, num_tokens_per_expert_E: torch.Tensor, ) -> torch.Tensor: - if isinstance(self.w13, DTensor): - w13 = self.w13.to_local() + if isinstance(self.w1_EFD, DTensor): + w1_EFD = self.w1_EFD.to_local() assert isinstance(self.w2_EDF, DTensor) w2_EDF = self.w2_EDF.to_local() + assert isinstance(self.w3_EFD, DTensor) + w3_EFD = self.w3_EFD.to_local() else: - w13 = self.w13 + w1_EFD = self.w1_EFD w2_EDF = self.w2_EDF + w3_EFD = self.w3_EFD offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) x = x_RD.bfloat16() - w13 = w13.bfloat16() + # (E, F, 2, D) with [:, :, 0] = gate (w1_EFD) and [:, :, 1] = up + # (w3_EFD). + w13 = torch.stack([w1_EFD, w3_EFD], dim=2).bfloat16() w2_t = w2_EDF.bfloat16().transpose(-2, -1) _validate_grouped_inputs(x, w13, w2_t, offsets_E) - return _MXFP8SwiGLUGroupedMLP.apply( + return _MXFP8GroupedMLP.apply( x, w13, w2_t, @@ -569,26 +536,24 @@ def forward( @override( target=FeedForward.Config, - description="Dense SwiGLU FFN via the composite MXFP8 SwiGLU MLP (fused w13).", + description="Dense SwiGLU FFN via the composite MXFP8 MLP.", ) -def mxfp8_fused_swiglu( +def mxfp8_fused_mlp( cfg: FeedForward.Config, *, fuse_activation: bool = True, -) -> "MXFP8FusedSwiGLU.Config": - # Config-application-time gate, matching the MXFP8 converters' UX; the - # composite re-validates at runtime. +) -> "MXFP8FusedMLP.Config": + # Config-application-time gate; the composite re-validates at runtime. if not has_cuda_capability(10, 0): raise ValueError( - "mxfp8_fused_swiglu requires SM100 or later; remove the override " + "mxfp8_fused_mlp requires SM100 or later; remove the override " "or run on supported hardware." ) - # Fail loud on anything but the stock config: this override owns the whole - # MLP's quantization, so composing it with another FFN variant or a linear - # quantization converter is a config error, not a silent no-op. + # Composing with another FFN variant or a linear quantization converter + # is a config error, not a silent no-op. if type(cfg) is not FeedForward.Config: raise ValueError( - "mxfp8_fused_swiglu targets the stock FeedForward.Config, got " + "mxfp8_fused_mlp targets the stock FeedForward.Config, got " f"{type(cfg).__qualname__}; narrow this override's fqns or remove " "the conflicting override/converter." ) @@ -596,65 +561,50 @@ def mxfp8_fused_swiglu( sub = getattr(cfg, name) if type(sub) is not Linear.Config: raise ValueError( - "mxfp8_fused_swiglu requires stock Linear.Config projections, " + "mxfp8_fused_mlp requires stock Linear.Config projections, " f"but {name} is {type(sub).__qualname__}. The composite " "quantizes every GEMM itself -- do not combine it with a " "linear quantization converter on the same module." ) + if cfg.w1.bias or cfg.w3.bias: + raise ValueError( + "mxfp8_fused_mlp supports a bias on w2 only; the composite has " + "no w1/w3 bias path." + ) - # Same param-init and sharding remaps as the fused_swiglu factory. - w1_init = (cfg.w1.param_init or {}).get("weight") - w3_init = (cfg.w3.param_init or {}).get("weight") - param_init = None - if w1_init is not None and w3_init is not None: - param_init = {"w13": _make_fused_gate_up_init(w1_init, w3_init, gate_up_axis=1)} - - fused = derive( - cfg, - MXFP8FusedSwiGLU.Config, - param_init=param_init, - fuse_activation=fuse_activation, - ) - base = cfg.sharding_config - fused.sharding_config = ShardingConfig( - state_shardings={"w13": dense_param_placement(tp=spmd.S(0))}, - in_src_shardings=base.in_src_shardings if base is not None else None, - in_dst_shardings=base.in_dst_shardings if base is not None else None, - ) - return fused + return derive(cfg, MXFP8FusedMLP.Config, fuse_activation=fuse_activation) @override( target=RoutedExperts.Config, - description="Routed experts via the composite MXFP8 SwiGLU grouped MLP " - "(fused w13, 128-row-padded token groups).", + description="Routed experts via the composite MXFP8 grouped MLP " + "(128-row-padded token groups).", ) -def mxfp8_fused_grouped_experts( +def mxfp8_fused_grouped_mlp( cfg: RoutedExperts.Config, *, fuse_activation: bool = True, ) -> RoutedExperts.Config: - # Config-application-time gate, matching the MXFP8 converters' UX; the - # composite re-validates at runtime. + # Config-application-time gate; the composite re-validates at runtime. if not has_cuda_capability(10, 0): raise ValueError( - "mxfp8_fused_grouped_experts requires SM100 or later; remove the " + "mxfp8_fused_grouped_mlp requires SM100 or later; remove the " "override or run on supported hardware." ) - # Targets RoutedExperts.Config (not GroupedExperts.Config) because the - # grouped composite constrains BOTH the experts and the token dispatcher: - # its kernels require every per-expert token group padded to a multiple of - # 128 rows (zero-filled), which only the padded dispatch path produces. + # 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 + # dispatch path produces. if type(cfg) is not RoutedExperts.Config: raise ValueError( - "mxfp8_fused_grouped_experts targets the stock " + "mxfp8_fused_grouped_mlp targets the stock " f"RoutedExperts.Config, got {type(cfg).__qualname__}; narrow this " "override's fqns or remove the conflicting override." ) inner = cfg.inner_experts if type(inner) is not GroupedExperts.Config: raise ValueError( - "mxfp8_fused_grouped_experts requires the stock " + "mxfp8_fused_grouped_mlp requires the stock " f"GroupedExperts.Config, but inner_experts is " f"{type(inner).__qualname__}. The composite quantizes every " "grouped GEMM itself -- do not combine it with the MXFP8 " @@ -663,16 +613,7 @@ def mxfp8_fused_grouped_experts( swap_token_dispatcher(cfg, pad_multiple=128) - # Same param-init and sharding remaps as the fused_grouped_experts factory. - param_init = _fuse_w13_grouped_experts_param_init(inner.param_init) - fused = derive( - inner, - MXFP8FusedGroupedExperts.Config, - param_init=param_init, - fuse_activation=fuse_activation, + cfg.inner_experts = derive( + inner, MXFP8FusedGroupedMLP.Config, fuse_activation=fuse_activation ) - base = inner.sharding_config - if base is not None: - fused.sharding_config = _fuse_w13_grouped_experts_sharding(base) - cfg.inner_experts = fused return cfg