From 88e2e8928c83e6e941a5c15782de431ff68301b9 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Mon, 24 Aug 2026 02:09:43 -0700 Subject: [PATCH 1/2] Refactor the MXFP8 fused-MLP override to a self-owned structure MXFP8FusedGroupedMLP now subclasses stock GroupedExperts directly: the FusedGroupedExperts inheritance is replaced by self-owned w13 registration, stock-layout checkpoint hooks, and param-init/sharding remap helpers, so future MXFP8 fusion paths extend this module without that dependency (none are wired here). Everything renames under the "MXFP8 fused MLP" umbrella (module, classes, factories, debugmodel flavors); the dense class keeps its small FusedSwiGLU base. Bitwise-verified on both paths (5-step deterministic NGPU=2 runs, losses and grad norms identical to the previous revision). Co-Authored-By: Claude Fable 5 --- ...lu_override.py => test_mxfp8_fused_mlp.py} | 58 ++-- .../models/deepseek_v3/config_registry.py | 20 +- torchtitan/models/llama3/config_registry.py | 11 +- ...fp8_fused_swiglu.py => mxfp8_fused_mlp.py} | 301 ++++++++++-------- 4 files changed, 203 insertions(+), 187 deletions(-) rename tests/unit_tests/{test_mxfp8_swiglu_override.py => test_mxfp8_fused_mlp.py} (75%) rename torchtitan/overrides/{mxfp8_fused_swiglu.py => mxfp8_fused_mlp.py} (69%) diff --git a/tests/unit_tests/test_mxfp8_swiglu_override.py b/tests/unit_tests/test_mxfp8_fused_mlp.py similarity index 75% rename from tests/unit_tests/test_mxfp8_swiglu_override.py rename to tests/unit_tests/test_mxfp8_fused_mlp.py index db5c9e80b5..0766eda94d 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 @@ -30,39 +26,35 @@ 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 +66,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 +76,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 +92,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 +105,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) @@ -160,11 +152,9 @@ def test_dense_factory_raises_on_non_stock_ffn(self): 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 +166,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 69% rename from torchtitan/overrides/mxfp8_fused_swiglu.py rename to torchtitan/overrides/mxfp8_fused_mlp.py index 1939b347d6..aee3502772 100644 --- a/torchtitan/overrides/mxfp8_fused_swiglu.py +++ b/torchtitan/overrides/mxfp8_fused_mlp.py @@ -6,50 +6,36 @@ # pyrefly: ignore-errors -"""Composite MXFP8 SwiGLU MLP for a fused w13 projection. +"""Composite MXFP8 MLP overrides with a fused w13 projection. -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 ``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 +from collections.abc import Callable +from dataclasses import dataclass, replace import spmd_types as spmd @@ -83,22 +69,16 @@ 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.overrides.fused_swiglu import 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 +188,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 +231,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 +247,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 +270,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 +281,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 +353,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 +400,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 +414,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,14 +441,60 @@ 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): +def _make_fused_gate_up_init( + gate_init: Callable, + up_init: Callable, + *, + gate_up_axis: int, +) -> Callable: + """Build an initializer for a fused gate/up weight from per-half + initializers (index 0 on ``gate_up_axis`` = gate / stock w1, index 1 = up + / stock w3). Each half keeps its own initializer because the gate and up + inits differ (e.g. up shares w2's depth-scaled init).""" + + def _init(t: torch.Tensor) -> None: + gate_idx: list[int | slice] = [slice(None)] * t.ndim + up_idx: list[int | slice] = [slice(None)] * t.ndim + gate_idx[gate_up_axis] = 0 + up_idx[gate_up_axis] = 1 + gate_init(t[tuple(gate_idx)]) # gate (stock w1) + up_init(t[tuple(up_idx)]) # up (stock w3) + + return _init + + +def _fuse_w13_grouped_param_init(param_init: dict | None) -> dict | None: + """Remap ``w1_EFD`` / ``w3_EFD`` initializers onto the fused ``w13``; + other entries (e.g. ``w2_EDF``) are kept as-is.""" + if param_init is None: + return None + w1_init = param_init.get("w1_EFD") + w3_init = param_init.get("w3_EFD") + fused = {k: v for k, v in param_init.items() if k not in ("w1_EFD", "w3_EFD")} + if w1_init is not None and w3_init is not None: + fused["w13"] = _make_fused_gate_up_init(w1_init, w3_init, gate_up_axis=2) + return fused or None + + +def _fuse_w13_grouped_sharding(base: ShardingConfig) -> ShardingConfig: + """Replace the ``w1_EFD`` / ``w3_EFD`` state shardings with one for + ``w13``: the same axes as ``w1_EFD`` (EP on dim 0, TP on dim 1), the ``2`` + axis unsharded. Everything else is kept.""" + state = dict(base.state_shardings) + w1_layout = state.pop("w1_EFD") + state.pop("w3_EFD") + state["w13"] = w1_layout + return replace(base, state_shardings=state) + + +class MXFP8FusedMLP(FusedSwiGLU): """:class:`FusedSwiGLU` whose forward runs the composite MXFP8 SwiGLU MLP. Inherits the fused ``w13`` parameter, the stock-layout checkpoint hooks, @@ -505,11 +515,11 @@ 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( + output = mxfp8_mlp_w13( x, self.w13, self.w2.weight, @@ -520,24 +530,34 @@ 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. + The gate and up projections are fused into one ``w13`` parameter of shape + ``(num_experts, hidden_dim, 2, dim)`` (``w13[:, :, 0]`` = stock + ``w1_EFD``, ``w13[:, :, 1]`` = stock ``w3_EFD``); checkpoints save/load + the stock layout, staying interchangeable with the non-fused module. 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`. + 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).""" def __init__(self, config: Config): super().__init__(config) + del self.w1_EFD + del self.w3_EFD + self.w13 = torch.nn.Parameter( + torch.empty(config.num_experts, config.hidden_dim, 2, config.dim) + ) + self.register_state_dict_post_hook(self._split_w13_on_save) + self.register_load_state_dict_pre_hook(self._merge_w13_on_load) self.fuse_activation = config.fuse_activation def forward( @@ -558,7 +578,7 @@ def forward( w13 = w13.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, @@ -566,29 +586,43 @@ def forward( self.fuse_activation, ).type_as(x_RD) + @staticmethod + def _split_w13_on_save(module, state_dict, prefix, local_metadata) -> None: + """Save the fused ``w13`` as stock ``w1_EFD`` / ``w3_EFD``.""" + w13 = state_dict.pop(f"{prefix}w13") + state_dict[f"{prefix}w1_EFD"] = w13[:, :, 0, :].contiguous() + state_dict[f"{prefix}w3_EFD"] = w13[:, :, 1, :].contiguous() + + @staticmethod + def _merge_w13_on_load(module, state_dict, prefix, *args) -> None: + """Combine stock ``w1_EFD`` / ``w3_EFD`` back into the fused ``w13``.""" + w1_key, w3_key = f"{prefix}w1_EFD", f"{prefix}w3_EFD" + if w1_key in state_dict and w3_key in state_dict: + state_dict[f"{prefix}w13"] = torch.stack( + [state_dict.pop(w1_key), state_dict.pop(w3_key)], dim=2 + ) + @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 (fused w13).", ) -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,13 +630,12 @@ 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." ) - # 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 @@ -611,7 +644,7 @@ def mxfp8_fused_swiglu( fused = derive( cfg, - MXFP8FusedSwiGLU.Config, + MXFP8FusedMLP.Config, param_init=param_init, fuse_activation=fuse_activation, ) @@ -626,35 +659,34 @@ def mxfp8_fused_swiglu( @override( target=RoutedExperts.Config, - description="Routed experts via the composite MXFP8 SwiGLU grouped MLP " + description="Routed experts via the composite MXFP8 grouped MLP " "(fused w13, 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 +695,15 @@ 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) + param_init = _fuse_w13_grouped_param_init(inner.param_init) fused = derive( inner, - MXFP8FusedGroupedExperts.Config, + MXFP8FusedGroupedMLP.Config, param_init=param_init, fuse_activation=fuse_activation, ) base = inner.sharding_config if base is not None: - fused.sharding_config = _fuse_w13_grouped_experts_sharding(base) + fused.sharding_config = _fuse_w13_grouped_sharding(base) cfg.inner_experts = fused return cfg From 9f44233d338ba7e51d9df245ec310cac13453e6e Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Mon, 24 Aug 2026 03:22:31 -0700 Subject: [PATCH 2/2] Drop the fused w13 parameterization in favor of stock parameters Both override classes now keep the stock parameters (w1/w2/w3 dense, w1_EFD/w2_EDF/w3_EFD grouped) and stack the gate/up weights into the composite's w13 operand at forward time, byte-identical to the old merge-hook mapping; the module no longer depends on fused_swiglu. Ckpt-paired 2-GPU runs match the fused-param arm bitwise at print precision (grouped exactly; dense exactly once grad clipping is inert), and fresh-init state dicts match the stock modules bitwise. Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_mxfp8_fused_mlp.py | 43 ++++-- torchtitan/overrides/mxfp8_fused_mlp.py | 164 +++++------------------ 2 files changed, 72 insertions(+), 135 deletions(-) diff --git a/tests/unit_tests/test_mxfp8_fused_mlp.py b/tests/unit_tests/test_mxfp8_fused_mlp.py index 0766eda94d..786eb4610f 100644 --- a/tests/unit_tests/test_mxfp8_fused_mlp.py +++ b/tests/unit_tests/test_mxfp8_fused_mlp.py @@ -23,7 +23,6 @@ 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_mlp import ( @@ -126,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 @@ -138,15 +137,43 @@ 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, diff --git a/torchtitan/overrides/mxfp8_fused_mlp.py b/torchtitan/overrides/mxfp8_fused_mlp.py index aee3502772..29c04f790e 100644 --- a/torchtitan/overrides/mxfp8_fused_mlp.py +++ b/torchtitan/overrides/mxfp8_fused_mlp.py @@ -6,13 +6,17 @@ # pyrefly: ignore-errors -"""Composite MXFP8 MLP overrides with a fused w13 projection. +"""Composite MXFP8 MLP overrides. 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`` runs the unified SwiGLU+MXFP8 kernel (the BF16 activation never reaches global memory); ``False`` materializes it in BF16 and quantizes with @@ -34,10 +38,7 @@ composite and factory. """ -from collections.abc import Callable -from dataclasses import dataclass, replace - -import spmd_types as spmd +from dataclasses import dataclass import torch import torch.nn.functional as F @@ -65,12 +66,9 @@ 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 FusedSwiGLU -from torchtitan.protocols.sharding import ShardingConfig from torchtitan.tools.utils import has_cuda_capability __all__ = [ @@ -448,62 +446,13 @@ def _validate_grouped_inputs(x, w13, w2_t, offs): torch._check(cond) -def _make_fused_gate_up_init( - gate_init: Callable, - up_init: Callable, - *, - gate_up_axis: int, -) -> Callable: - """Build an initializer for a fused gate/up weight from per-half - initializers (index 0 on ``gate_up_axis`` = gate / stock w1, index 1 = up - / stock w3). Each half keeps its own initializer because the gate and up - inits differ (e.g. up shares w2's depth-scaled init).""" - - def _init(t: torch.Tensor) -> None: - gate_idx: list[int | slice] = [slice(None)] * t.ndim - up_idx: list[int | slice] = [slice(None)] * t.ndim - gate_idx[gate_up_axis] = 0 - up_idx[gate_up_axis] = 1 - gate_init(t[tuple(gate_idx)]) # gate (stock w1) - up_init(t[tuple(up_idx)]) # up (stock w3) - - return _init - - -def _fuse_w13_grouped_param_init(param_init: dict | None) -> dict | None: - """Remap ``w1_EFD`` / ``w3_EFD`` initializers onto the fused ``w13``; - other entries (e.g. ``w2_EDF``) are kept as-is.""" - if param_init is None: - return None - w1_init = param_init.get("w1_EFD") - w3_init = param_init.get("w3_EFD") - fused = {k: v for k, v in param_init.items() if k not in ("w1_EFD", "w3_EFD")} - if w1_init is not None and w3_init is not None: - fused["w13"] = _make_fused_gate_up_init(w1_init, w3_init, gate_up_axis=2) - return fused or None - - -def _fuse_w13_grouped_sharding(base: ShardingConfig) -> ShardingConfig: - """Replace the ``w1_EFD`` / ``w3_EFD`` state shardings with one for - ``w13``: the same axes as ``w1_EFD`` (EP on dim 0, TP on dim 1), the ``2`` - axis unsharded. Everything else is kept.""" - state = dict(base.state_shardings) - w1_layout = state.pop("w1_EFD") - state.pop("w3_EFD") - state["w13"] = w1_layout - return replace(base, state_shardings=state) - - -class MXFP8FusedMLP(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).""" @@ -519,9 +468,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: "tensor parallelism); narrow the mxfp8_fused_mlp override's " "fqns or drop it for this module." ) + # (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, ) @@ -534,13 +485,11 @@ class MXFP8FusedGroupedMLP(GroupedExperts): """Routed experts whose forward runs the composite MXFP8 SwiGLU grouped MLP. - The gate and up projections are fused into one ``w13`` parameter of shape - ``(num_experts, hidden_dim, 2, dim)`` (``w13[:, :, 0]`` = stock - ``w1_EFD``, ``w13[:, :, 1]`` = stock ``w3_EFD``); checkpoints save/load - the stock layout, staying interchangeable with the non-fused module. - Requires token groups padded to multiples of 128 rows (zero-filled) -- - the ``mxfp8_fused_grouped_mlp`` factory swaps the token dispatcher - accordingly. + 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) @@ -551,13 +500,6 @@ class Config(GroupedExperts.Config): def __init__(self, config: Config): super().__init__(config) - del self.w1_EFD - del self.w3_EFD - self.w13 = torch.nn.Parameter( - torch.empty(config.num_experts, config.hidden_dim, 2, config.dim) - ) - self.register_state_dict_post_hook(self._split_w13_on_save) - self.register_load_state_dict_pre_hook(self._merge_w13_on_load) self.fuse_activation = config.fuse_activation def forward( @@ -565,17 +507,22 @@ 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 _MXFP8GroupedMLP.apply( @@ -586,26 +533,10 @@ def forward( self.fuse_activation, ).type_as(x_RD) - @staticmethod - def _split_w13_on_save(module, state_dict, prefix, local_metadata) -> None: - """Save the fused ``w13`` as stock ``w1_EFD`` / ``w3_EFD``.""" - w13 = state_dict.pop(f"{prefix}w13") - state_dict[f"{prefix}w1_EFD"] = w13[:, :, 0, :].contiguous() - state_dict[f"{prefix}w3_EFD"] = w13[:, :, 1, :].contiguous() - - @staticmethod - def _merge_w13_on_load(module, state_dict, prefix, *args) -> None: - """Combine stock ``w1_EFD`` / ``w3_EFD`` back into the fused ``w13``.""" - w1_key, w3_key = f"{prefix}w1_EFD", f"{prefix}w3_EFD" - if w1_key in state_dict and w3_key in state_dict: - state_dict[f"{prefix}w13"] = torch.stack( - [state_dict.pop(w1_key), state_dict.pop(w3_key)], dim=2 - ) - @override( target=FeedForward.Config, - description="Dense SwiGLU FFN via the composite MXFP8 MLP (fused w13).", + description="Dense SwiGLU FFN via the composite MXFP8 MLP.", ) def mxfp8_fused_mlp( cfg: FeedForward.Config, @@ -635,32 +566,19 @@ def mxfp8_fused_mlp( "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." + ) - 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, - MXFP8FusedMLP.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 grouped MLP " - "(fused w13, 128-row-padded token groups).", + "(128-row-padded token groups).", ) def mxfp8_fused_grouped_mlp( cfg: RoutedExperts.Config, @@ -695,15 +613,7 @@ def mxfp8_fused_grouped_mlp( swap_token_dispatcher(cfg, pad_multiple=128) - param_init = _fuse_w13_grouped_param_init(inner.param_init) - fused = derive( - inner, - MXFP8FusedGroupedMLP.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_sharding(base) - cfg.inner_experts = fused return cfg