From 45c85116d60de5140352dc02eb3a5c0807984092 Mon Sep 17 00:00:00 2001 From: Animesh Jain Date: Fri, 28 Aug 2026 14:14:26 -0700 Subject: [PATCH] Add MXFP8 recipes for the qwen3 and gpt_oss MoE models The grouped-expert support landed without a way to turn it on outside DeepSeek-V3, so the only model exercising it was the one it was developed against. Add debug recipes for the two other MoE models whose experts the converter already handles. Both quantize the attention projections and the expert grouped GEMMs, and leave the router gate and lm_head in BF16: the gate is tiny and its output drives routing, and the vocab projection is large and numerically sensitive. Only the fused QKV saves a quantized input activation for WGRAD, since nothing else consumes it, while attention already retains the output projection's input. gpt_oss exercises a path DeepSeek-V3 does not: its experts carry per-expert biases and a SwiGLU clamp, and the converter preserves both by subclassing GptOssGroupedExperts rather than replacing it -- only the grouped GEMM seam changes, and the 2D biases are skipped when installing compute weights. Add CPU coverage for both recipes. MXFP8 needs SM100, which CI does not have, but the config-tree transform does not, so the test bypasses the hardware gate and asserts which modules are converted, which input-activation format each linear saves, and that the experts subclass the model's own variant. Verified on 2xGB200: qwen3_moe_debug_mxfp8 trains with loss tracking BF16 (6.673 vs 6.603 at step 3) and 0.6 GiB lower peak memory. --- tests/unit_tests/cpu/test_quantization.py | 52 +++++++++++++++++++ torchtitan/models/gpt_oss/config_registry.py | 54 ++++++++++++++++++++ torchtitan/models/qwen3/config_registry.py | 53 ++++++++++++++++++- 3 files changed, 158 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/cpu/test_quantization.py b/tests/unit_tests/cpu/test_quantization.py index 0198dfbf33..71930449f9 100644 --- a/tests/unit_tests/cpu/test_quantization.py +++ b/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) diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index 29368d36f8..d98e61ed40 100644 --- a/torchtitan/models/gpt_oss/config_registry.py +++ b/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( diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py index 4c79519d56..6dba65685d 100644 --- a/torchtitan/models/qwen3/config_registry.py +++ b/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.