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.