Refactor the MXFP8 fused-MLP overrides onto stock parameters

tests/unit_tests/test_mxfp8_swiglu_override.py → 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__":

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
 
 

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