[DO NOT REVIEW] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models

tests/unit_tests/cpu/test_quantization.py

@@ -385,6 +385,58 @@ def test_quantized_grouped_experts():
     assert hasattr(float8_cls.Config, "swiglu_limit")
 
 
+@pytest.mark.parametrize(
+    "module, recipe, experts_cls_name",
+    [
+        ("qwen3", "qwen3_moe_debug_mxfp8", "MXFP8GroupedExperts"),
+        ("gpt_oss", "gpt_oss_debugmodel_mxfp8", "MXFP8GptOssGroupedExperts"),
+    ],
+)
+def test_mxfp8_moe_configs_convert_experts_and_spare_router(
+    monkeypatch, module, recipe, experts_cls_name
+):
+    """The MoE MXFP8 recipes quantize experts and attention, nothing else.
+
+    Runs on CPU: the converters' SM100 gate is bypassed because the config-tree
+    transform under test does not depend on hardware.
+    """
+    pytest.importorskip("torchao")
+    if MXFP8Linear is None:
+        pytest.skip("torchao MXFP8Linear is unavailable")
+    import torchtitan.components.quantization.mxfp8.converter as converter_mod
+
+    monkeypatch.setattr(converter_mod, "has_cuda_capability", lambda *_: True)
+
+    config = ConfigManager().parse_args(["--module", module, "--config", recipe])
+    model_config = config.model_spec.model
+    assert has_quantization(model_config)
+
+    quantized, stock = [], []
+    for fqn, linear_config, _parent, _attr in model_config.traverse(Linear.Config):
+        target = quantized if isinstance(linear_config, MXFP8Linear.Config) else stock
+        target.append(fqn.split(".", 2)[-1])
+
+    # Attention projections are quantized; the router gate stays BF16 because
+    # its output drives routing, and lm_head because it is large and sensitive.
+    assert set(quantized) == {"attention.qkv_linear.wqkv", "attention.wo"}
+    assert set(stock) == {"moe.router.gate", "lm_head"}
+
+    # Only the fused QKV saves a quantized input: nothing else consumes its
+    # activation, whereas attention already retains wo's input.
+    for _fqn, linear_config, _parent, _attr in model_config.traverse(Linear.Config):
+        if not isinstance(linear_config, MXFP8Linear.Config):
+            continue
+        expected = "mxfp8" if _fqn.endswith("qkv_linear.wqkv") else "bf16"
+        assert linear_config.input_activation_format_for_backward == expected
+
+    # Experts subclass the model's own variant, so gpt_oss keeps its biases
+    # and SwiGLU clamp while only the grouped GEMM seam is replaced.
+    experts = [c for _fqn, c, _p, _a in model_config.traverse(GroupedExperts.Config)]
+    assert experts
+    for experts_config in experts:
+        assert type(experts_config)._owner.__name__ == experts_cls_name
+
+
 def test_mxfp8_grouped_experts_config_validation():
     """The grouped-expert config rejects unusable padding and activation formats."""
     experts_cls = get_mxfp8_grouped_experts_cls(GroupedExperts)

torchtitan/models/gpt_oss/config_registry.py

@@ -9,6 +9,10 @@
 from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss
 from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer
+from torchtitan.components.quantization import (
+    MXFP8GroupedExpertsConverter,
+    MXFP8LinearConverter,
+)
 from torchtitan.components.validate import Validator
 from torchtitan.config import ParallelismConfig, TrainingConfig
 from torchtitan.distributed.activation_checkpoint import FullAC
@@ -68,6 +72,56 @@ def gpt_oss_debugmodel_flex() -> Trainer.Config:
     return _gpt_oss_debugmodel(attn_backend="flex")
 
 
+def gpt_oss_mxfp8_linear_converter_config(
+    *, model_compile_enabled: bool
+) -> MXFP8LinearConverter.Config:
+    """Build the dense MXFP8 policy for the debug model.
+
+    ``fqns`` is an include-list, so the router gate and lm_head stay in BF16.
+
+    The fused QKV projection has a single-consumer input that nothing else
+    saves for backward, so its columnwise MXFP8 representation replaces BF16
+    storage. The output projection keeps the conservative BF16 format because
+    attention already retains its input.
+    """
+    return MXFP8LinearConverter.Config(
+        model_compile_enabled=model_compile_enabled,
+        fqns=["attention"],
+        linears_saving_inputs_for_backward_in_mxfp8=["attention.qkv_linear.wqkv"],
+    )
+
+
+def gpt_oss_debugmodel_mxfp8() -> Trainer.Config:
+    """Debug model with MXFP8 expert grouped GEMMs and dense linears.
+
+    The experts carry per-expert biases and a SwiGLU clamp, which the grouped
+    converter preserves by subclassing ``GptOssGroupedExperts``: only the
+    grouped GEMM seam is replaced, and the biases stay BF16 parameters.
+    """
+    config = _gpt_oss_debugmodel()
+    # The grouped converter swaps in a padding-capable token dispatcher, and
+    # TorchAOTokenDispatcher needs a CPU sync, which CUDA graphs reject. Run
+    # eager so this config works at any expert-parallel degree; the
+    # CUDA-graph path for MXFP8 experts needs HybridEP with a
+    # non_blocking_capacity_factor.
+    config.training.disable_cuda_graphs = True
+    model_compile_enabled = (
+        config.compile.enable and "model" in config.compile.components
+    )
+    config.model_spec = model_registry(
+        "debugmodel",
+        converters=[
+            gpt_oss_mxfp8_linear_converter_config(
+                model_compile_enabled=model_compile_enabled,
+            ),
+            MXFP8GroupedExpertsConverter.Config(
+                model_compile_enabled=model_compile_enabled,
+            ),
+        ],
+    )
+    return config
+
+
 def gpt_oss_20b() -> Trainer.Config:
     model_spec = model_registry("20b")
     return Trainer.Config(

torchtitan/models/qwen3/config_registry.py

@@ -22,7 +22,11 @@
     OptimizersContainer,
     ParamGroupConfig,
 )
-from torchtitan.components.quantization import NVFP4LinearConverter
+from torchtitan.components.quantization import (
+    MXFP8GroupedExpertsConverter,
+    MXFP8LinearConverter,
+    NVFP4LinearConverter,
+)
 from torchtitan.components.quantization.nvfp4 import nvfp4_bf16_tail_fqns
 from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig
 from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC
@@ -414,6 +418,53 @@ def qwen3_moe_debug() -> Trainer.Config:
     )
 
 
+def qwen3_mxfp8_linear_converter_config(
+    *, model_compile_enabled: bool
+) -> MXFP8LinearConverter.Config:
+    """Build the dense MXFP8 policy for the MoE debug model.
+
+    ``fqns`` is an include-list, so the router gate and lm_head stay in BF16;
+    the gate is tiny and its output drives routing, and the vocab projection
+    is both large and numerically sensitive.
+
+    The fused QKV projection has a single-consumer input that nothing else
+    saves for backward, so its columnwise MXFP8 representation replaces BF16
+    storage. The output projection keeps the conservative BF16 format because
+    flash attention already retains its input.
+    """
+    return MXFP8LinearConverter.Config(
+        model_compile_enabled=model_compile_enabled,
+        fqns=["attention"],
+        linears_saving_inputs_for_backward_in_mxfp8=["attention.qkv_linear.wqkv"],
+    )
+
+
+def qwen3_moe_debug_mxfp8() -> Trainer.Config:
+    """MoE debug model with MXFP8 expert grouped GEMMs and dense linears."""
+    config = qwen3_moe_debug()
+    # The grouped converter swaps in a padding-capable token dispatcher, and
+    # TorchAOTokenDispatcher needs a CPU sync, which CUDA graphs reject. Run
+    # eager so this config works at any expert-parallel degree; the
+    # CUDA-graph path for MXFP8 experts needs HybridEP with a
+    # non_blocking_capacity_factor.
+    config.training.disable_cuda_graphs = True
+    model_compile_enabled = (
+        config.compile.enable and "model" in config.compile.components
+    )
+    config.model_spec = model_registry(
+        "debugmodel_moe",
+        converters=[
+            qwen3_mxfp8_linear_converter_config(
+                model_compile_enabled=model_compile_enabled,
+            ),
+            MXFP8GroupedExpertsConverter.Config(
+                model_compile_enabled=model_compile_enabled,
+            ),
+        ],
+    )
+    return config
+
+
 def qwen3_moe_deepep() -> Trainer.Config:
     """Qwen3 debug MoE pretraining with the DeepEP v2 backend (compact training path), EP=4.