From ca0326d888bb1463475c9c84bec800bc0b888834 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 18:19:33 -0700 Subject: [PATCH 1/6] Add NVFP4 four-over-six grouped-experts converter and RL recipes Extends the four-over-six integration to the miles NVFP4 RL recipe shape (routed experts only, row-scaled activations, 1x16 weights, backward overrides): - NVFP4FourOverSixLinear / NVFP4FourOverSixLinearConverter gain backward_override ('quantized' | 'high_precision' | 'dequantized', mirroring TransformerEngine's NVTE_BACKWARD_OVERRIDE; None keeps the previous defaults) and weight_block ('1x16' mirrors NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1). Knob combinations are validated at config time. - NVFP4FourOverSixGroupedExpertsConverter: class-factory _grouped_mm override calling torchao four_over_six_grouped_mm (like the MXFP8 grouped-experts converter, with fqns include-list filtering). Four- over-six needs no RHT sign vector and no stochastic-rounding seed, so the stateless hook carries everything. Token dispatchers swap to padded variants (pad_multiple=128). - nvfp4_bf16_first_last_fqns: the miles --first-last-layers-bf16 analog. - deepseek_v3_debugmodel_nvfp4_four_over_six[_dequantized] recipes: the miles base point (row-scaled + MSE + bound 256 + 1x16 weights + high_precision backward, experts-only allow-list) and the GLM-5.2 analog (dequantized backward, first/last decoder layer bf16). - rl_grpo_qwen3_30b_a3b_varlen_nvfp4_four_over_six: the RL wiring. The trainer and vLLM generator share one model_spec, so both actors run the identical four-over-six forward -- the train/inference-consistency point of the miles recipe, with bf16 master weights syncing as usual. Unit tests: 35 passed (TE 2.19 devel container), covering converter targeting, knob plumbing, first/last-bf16 windows, and config-time validation. GPU smoke on GB200 verifies the hook's (E, K, N) -> (E, N, K) weight orientation, ragged-group padding, and both grads end to end. Depends on the TorchAO branch nvfp4-four-over-six-rowwise (wolfcomos/ao#7, through commit c47af4a7e). Co-Authored-By: Claude Fable 5 --- tests/unit_tests/test_quantization.py | 218 ++++++++++++++ .../components/quantization/__init__.py | 2 + torchtitan/components/quantization/nvfp4.py | 271 +++++++++++++++++- .../examples/alphabet_sort/config_registry.py | 32 +++ .../models/deepseek_v3/config_registry.py | 53 ++++ 5 files changed, 575 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_quantization.py b/tests/unit_tests/test_quantization.py index 74f6040d21..798ea475e8 100644 --- a/tests/unit_tests/test_quantization.py +++ b/tests/unit_tests/test_quantization.py @@ -422,3 +422,221 @@ def test_float8_grouped_experts_dcp_round_trip_needs_no_safe_globals(tmp_path): source.parameters(), target.parameters(), strict=True ): torch.testing.assert_close(target_parameter, source_parameter) + + +def test_nvfp4_bf16_first_last_fqns(): + from torchtitan.components.quantization.nvfp4 import nvfp4_bf16_first_last_fqns + + # 6 layers, one bf16 layer at each end -> convert layers 1..4. + fqns = nvfp4_bf16_first_last_fqns(6, 1, 1) + assert fqns == [f"layers.{i}." for i in range(1, 5)] + # Every fqn is trailing-dot anchored so "layers.1." matches layer 1 only, + # not "layers.10".."layers.19" (the converters substring-match). + assert all(f.startswith("layers.") and f.endswith(".") for f in fqns) + # No exclusions -> every layer converted. + assert nvfp4_bf16_first_last_fqns(4, 0, 0) == [ + "layers.0.", + "layers.1.", + "layers.2.", + "layers.3.", + ] + # A window that covers all layers leaves nothing to convert -> raise (an + # empty fqns list would instead convert *all* matching modules). + with pytest.raises(ValueError, match="nothing to convert"): + nvfp4_bf16_first_last_fqns(4, 2, 2) + + +def _four_over_six_grouped_cls(): + pytest.importorskip("torchao") + import torchtitan.components.quantization.nvfp4 as nvfp4_mod + + if nvfp4_mod.four_over_six_grouped_mm is None: + pytest.skip("torchao four-over-six grouped training prototype not available") + return nvfp4_mod._get_four_over_six_grouped_experts_cls + + +def test_four_over_six_grouped_experts_cls(): + """Four-over-six GroupedExperts factory: _owner, subclasses, knob fields.""" + factory = _four_over_six_grouped_cls() + + cls = factory(GroupedExperts) + assert cls.Config._owner is cls + assert issubclass(cls, GroupedExperts) + # Same class object on repeat calls (cached). + assert factory(GroupedExperts) is cls + + gpt_oss_cls = factory(GptOssGroupedExperts) + assert gpt_oss_cls.Config._owner is gpt_oss_cls + assert issubclass(gpt_oss_cls, GptOssGroupedExperts) + assert hasattr(gpt_oss_cls.Config, "swiglu_limit") + + config = cls.Config( + dim=256, + hidden_dim=128, + num_experts=2, + err_mode="mse", + e4m3_scale_bound=256, + row_scaled_activation=True, + backward_override="dequantized", + weight_block="1x16", + ) + assert config.backward_override == "dequantized" + + # The built module carries every knob into the grouped-GEMM call. + module = config.build() + assert module._four_over_six_kwargs == dict( + err_mode="mse", + e4m3_scale_bound=256, + row_scaled_activation=True, + weight_block="1x16", + backward_override="dequantized", + ) + + +def test_four_over_six_knob_validation(): + pytest.importorskip("torchao") + from torchtitan.components.quantization import NVFP4FourOverSixLinear + + if NVFP4FourOverSixLinear is None: + pytest.skip("torchao NVFP4 four-over-six training prototype not available") + + # Bad enum values are rejected at config time, not on the first forward. + with pytest.raises(ValueError, match="backward_override"): + NVFP4FourOverSixLinear.Config( + in_features=128, out_features=128, backward_override="bf16" + ) + with pytest.raises(ValueError, match="weight_block"): + NVFP4FourOverSixLinear.Config( + in_features=128, out_features=128, weight_block="32x32" + ) + # A row-scaled four-over-six tensor has no columnwise form for the + # quantized wgrad operand. + with pytest.raises(ValueError, match="no quantized backward"): + NVFP4FourOverSixLinear.Config( + in_features=128, + out_features=128, + row_scaled_activation=True, + backward_override="quantized", + ) + + factory = _four_over_six_grouped_cls() + grouped_config_cls = factory(GroupedExperts).Config + # Grouped four-over-six has no quantized backward at all. + with pytest.raises(ValueError, match="backward_override"): + grouped_config_cls( + dim=256, hidden_dim=128, num_experts=2, backward_override="quantized" + ) + # The grouped GEMM needs both projections' dims divisible by 128. + with pytest.raises(ValueError, match="divisible by 128"): + grouped_config_cls(dim=256, hidden_dim=96, num_experts=2) + + +def _parse_recipe(monkeypatch, module, recipe): + # Exercise convert() targeting independent of GPU: bypass the sm100 gate + # (hardware is irrelevant to the config-tree transform under test). + import torchtitan.components.quantization.nvfp4 as nvfp4_mod + + monkeypatch.setattr(nvfp4_mod, "has_cuda_capability", lambda *_: True) + config = ConfigManager().parse_args(["--module", module, "--config", recipe]) + return config.model_spec.model + + +def test_nvfp4_four_over_six_grouped_converter_targets_experts(monkeypatch): + factory = _four_over_six_grouped_cls() + from torchtitan.components.quantization import NVFP4FourOverSixLinear + from torchtitan.models.common.token_dispatcher import TorchAOTokenDispatcher + + model_config = _parse_recipe( + monkeypatch, "deepseek_v3", "deepseek_v3_debugmodel_nvfp4_four_over_six" + ) + + grouped_config_cls = factory(GroupedExperts).Config + converted = [ + (fqn, gc) + for fqn, gc, _parent, _attr in model_config.traverse(GroupedExperts.Config) + ] + # The debugmodel has 6 layers with layer 0 dense: every MoE layer converts. + assert len(converted) == 5 + for fqn, gc in converted: + assert isinstance(gc, grouped_config_cls), fqn + # The miles base-recipe point plumbs through to every expert config. + assert gc.row_scaled_activation is True + assert gc.err_mode == "mse" + assert gc.e4m3_scale_bound == 256 + assert gc.weight_block == "1x16" + assert gc.backward_override == "high_precision" + + # Every token dispatcher on a converted MoE layer pads to 128 rows. + dispatchers = [ + dc + for _fqn, dc, _parent, _attr in model_config.traverse( + TorchAOTokenDispatcher.Config + ) + ] + assert len(dispatchers) == 5 + assert all(dc.pad_multiple == 128 for dc in dispatchers) + + # Experts-only allow-list: no dense Linear is quantized (attention, dense + # FFN, shared experts, router gate, and lm_head all stay stock). + if NVFP4FourOverSixLinear is not None: + for fqn, lc, _parent, _attr in model_config.traverse(Linear.Config): + assert not isinstance(lc, NVFP4FourOverSixLinear.Config), fqn + + +def test_nvfp4_four_over_six_grouped_dequantized_keeps_first_last_bf16(monkeypatch): + factory = _four_over_six_grouped_cls() + + model_config = _parse_recipe( + monkeypatch, + "deepseek_v3", + "deepseek_v3_debugmodel_nvfp4_four_over_six_dequantized", + ) + + grouped_config_cls = factory(GroupedExperts).Config + converted_layers, stock_layers = set(), set() + for fqn, gc, _parent, _attr in model_config.traverse(GroupedExperts.Config): + layer = int(fqn.split(".")[1]) + if isinstance(gc, grouped_config_cls): + converted_layers.add(layer) + assert gc.backward_override == "dequantized" + else: + stock_layers.add(layer) + + # 6 layers: layer 0 is dense (no experts), the first/last-bf16 window keeps + # layer 5 stock, so layers 1..4 convert. + assert converted_layers == {1, 2, 3, 4} + assert stock_layers == {5} + + +def test_nvfp4_four_over_six_linear_converter_plumbs_new_knobs(monkeypatch): + pytest.importorskip("torchao") + import torchtitan.components.quantization.nvfp4 as nvfp4_mod + from torchtitan.components.quantization import ( + NVFP4FourOverSixLinear, + NVFP4FourOverSixLinearConverter, + ) + + if NVFP4FourOverSixLinear is None: + pytest.skip("torchao NVFP4 four-over-six training prototype not available") + monkeypatch.setattr(nvfp4_mod, "has_cuda_capability", lambda *_: True) + + config = ConfigManager().parse_args( + ["--module", "llama3", "--config", "llama3_debugmodel"] + ) + converter = NVFP4FourOverSixLinearConverter( + NVFP4FourOverSixLinearConverter.Config( + fqns=["layers"], + backward_override="dequantized", + weight_block="1x16", + ) + ) + model_config = converter.convert(config.model_spec.model) + + converted = [ + lc + for _fqn, lc, _parent, _attr in model_config.traverse(Linear.Config) + if isinstance(lc, NVFP4FourOverSixLinear.Config) + ] + assert converted + assert all(lc.backward_override == "dequantized" for lc in converted) + assert all(lc.weight_block == "1x16" for lc in converted) diff --git a/torchtitan/components/quantization/__init__.py b/torchtitan/components/quantization/__init__.py index 0c1f916ef4..6c5f2ef32a 100644 --- a/torchtitan/components/quantization/__init__.py +++ b/torchtitan/components/quantization/__init__.py @@ -43,6 +43,7 @@ class Config(ModelConfigConverter.Config): MXFP8LinearConverter, ) from .nvfp4 import ( # noqa: F401, E402 + NVFP4FourOverSixGroupedExpertsConverter, NVFP4FourOverSixLinear, NVFP4FourOverSixLinearConverter, NVFP4Linear, @@ -56,6 +57,7 @@ class Config(ModelConfigConverter.Config): "MXFP8GroupedExpertsConverter", "MXFP8Linear", "MXFP8LinearConverter", + "NVFP4FourOverSixGroupedExpertsConverter", "NVFP4FourOverSixLinear", "NVFP4FourOverSixLinearConverter", "NVFP4Linear", diff --git a/torchtitan/components/quantization/nvfp4.py b/torchtitan/components/quantization/nvfp4.py index 8e1540deb8..0324787c12 100644 --- a/torchtitan/components/quantization/nvfp4.py +++ b/torchtitan/components/quantization/nvfp4.py @@ -17,7 +17,7 @@ """ import math -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field, fields, replace from typing import cast import spmd_types as spmd @@ -27,11 +27,14 @@ from torchtitan.distributed.parallel_dims import MeshAxisName from torchtitan.models.common.decoder_sharding import dense_activation_placement from torchtitan.models.common.linear import Linear +from torchtitan.models.common.moe import GroupedExperts from torchtitan.protocols.module import Module from torchtitan.protocols.sharding import LocalMapConfig, SpmdLayout from torchtitan.tools.logging import logger from torchtitan.tools.utils import has_cuda_capability +from .utils import swap_token_dispatcher + TP = MeshAxisName.TP # TorchAO's NVFP4 Triton kernels require each local GEMM dimension to be a @@ -265,6 +268,66 @@ def nvfp4_bf16_tail_fqns(num_layers: int, bf16_tail_fraction: float) -> list[str return [f"layers.{i}." for i in range(convert_upto)] +def nvfp4_bf16_first_last_fqns( + num_layers: int, num_start_layers_bf16: int, num_end_layers_bf16: int +) -> list[str]: + """Converter ``fqns`` keeping the first ``num_start_layers_bf16`` and last + ``num_end_layers_bf16`` decoder layers in bf16 (the miles NVFP4 RL recipes' + ``--first-last-layers-bf16`` analog). + + Each fqn has a trailing '.' so 'layers.1.' matches layer 1 only, not + 'layers.10' (the converters substring-match). Raises if the window would + leave no layer to convert: an empty fqns list would instead convert *all* + matching modules (the ``not fqns`` branch in convert), the opposite of the + intent. + """ + convert_from = num_start_layers_bf16 + convert_upto = num_layers - num_end_layers_bf16 + if convert_from >= convert_upto: + raise ValueError( + f"num_start_layers_bf16={num_start_layers_bf16} and " + f"num_end_layers_bf16={num_end_layers_bf16} keep all {num_layers} " + "layers in bf16; nothing to convert to NVFP4." + ) + return [f"layers.{i}." for i in range(convert_from, convert_upto)] + + +def _validate_four_over_six_knobs( + backward_override: str | None, + weight_block: str, + *, + row_scaled_activation: bool = False, + grouped: bool = False, +) -> None: + """Reject invalid four-over-six knob combinations at config time. + + Mirrors the TorchAO ops' runtime checks so a bad recipe fails when the + config tree is built rather than on the first forward. Grouped GEMMs have + no quantized backward (TransformerEngine rejects four-over-six group + quantization), and a row-scaled four-over-six tensor has no columnwise + form for the quantized wgrad operand. + """ + allowed = ( + (None, "high_precision", "dequantized") + if grouped + else (None, "quantized", "high_precision", "dequantized") + ) + if backward_override not in allowed: + raise ValueError( + f"backward_override must be one of {allowed}; " + f"got {backward_override!r}" + ) + if backward_override == "quantized" and row_scaled_activation: + raise ValueError( + "row-scaled four-over-six has no quantized backward; use " + "'high_precision' or 'dequantized'" + ) + if weight_block not in ("1x16", "16x16"): + raise ValueError( + f"weight_block must be '1x16' or '16x16', got {weight_block!r}" + ) + + class NVFP4LinearConverter(QuantizationConverter): """Replace matching Linear.Config with NVFP4Linear.Config.""" @@ -358,6 +421,16 @@ class Config(Linear.Config): """One FP32 global scale per activation row instead of per tensor. Selects the bf16 backward.""" + backward_override: str | None = None + """'quantized', 'high_precision', or 'dequantized' (mirrors + TransformerEngine's NVTE_BACKWARD_OVERRIDE). None keeps the + recipe defaults: quantized backward per tensor, high-precision + backward when row-scaled.""" + + weight_block: str = "16x16" + """Weight tile granularity; '1x16' mirrors + NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1.""" + def __post_init__(self) -> None: for name in ("in_features", "out_features"): value = getattr(self, name) @@ -367,12 +440,19 @@ def __post_init__(self) -> None: f"got {name}={value}. NVFP4 cannot quantize this Linear; " "exclude it from the converter fqns." ) + _validate_four_over_six_knobs( + self.backward_override, + self.weight_block, + row_scaled_activation=self.row_scaled_activation, + ) def __init__(self, config: Config): Linear.__init__(self, config) self.err_mode = config.err_mode self.e4m3_scale_bound = config.e4m3_scale_bound self.row_scaled_activation = config.row_scaled_activation + self.backward_override = config.backward_override + self.weight_block = config.weight_block def forward(self, x: torch.Tensor) -> torch.Tensor: return four_over_six_linear( @@ -382,6 +462,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.err_mode, self.e4m3_scale_bound, self.row_scaled_activation, + self.backward_override, + self.weight_block, ) except ImportError: @@ -410,6 +492,15 @@ class Config(QuantizationConverter.Config): row_scaled_activation: bool = False """One FP32 global scale per activation row instead of per tensor.""" + backward_override: str | None = None + """'quantized', 'high_precision', or 'dequantized' (mirrors + TransformerEngine's NVTE_BACKWARD_OVERRIDE). None keeps the recipe + defaults.""" + + weight_block: str = "16x16" + """Weight tile granularity; '1x16' mirrors + NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1.""" + def __init__(self, config: Config): self.config = config @@ -442,6 +533,8 @@ def convert(self, model_config): err_mode=self.config.err_mode, e4m3_scale_bound=self.config.e4m3_scale_bound, row_scaled_activation=self.config.row_scaled_activation, + backward_override=self.config.backward_override, + weight_block=self.config.weight_block, ) if parent is None: model_config = new_config @@ -452,3 +545,179 @@ def convert(self, model_config): logger.info("Converted Linear layers to NVFP4FourOverSixLinear") return model_config + + +try: + from torchao.prototype.moe_training.nvfp4_training.four_over_six_grouped import ( + four_over_six_grouped_mm, + ) + +except ImportError: + four_over_six_grouped_mm = None + + +_four_over_six_experts_cache: dict[type, type] = {} + + +def _get_four_over_six_grouped_experts_cls(parent_cls: type) -> type: + """Get or create a four-over-six-quantized subclass of *parent_cls*. + + Works for any experts module exposing the ``_grouped_mm`` seam (the common + ``GroupedExperts`` and ``GptOssGroupedExperts``), like + ``_get_mxfp8_grouped_experts_cls``. The returned class has a proper + ``_owner`` set by ``__init_subclass__``. + + The subclass overrides ``_grouped_mm`` to call torchao's + ``four_over_six_grouped_mm``. Four-over-six needs no RHT sign vector and no + stochastic-rounding seed, so the stateless hook carries everything. + """ + if parent_cls in _four_over_six_experts_cache: + return _four_over_six_experts_cache[parent_cls] + + parent_config_cls = parent_cls.Config # type: ignore[attr-defined] + + class FourOverSixGroupedExperts(parent_cls): # type: ignore[valid-type, misc] + @dataclass(kw_only=True, slots=True) + class Config(parent_config_cls): # type: ignore[misc] + err_mode: str = "mae" + e4m3_scale_bound: int = 256 + row_scaled_activation: bool = False + backward_override: str | None = None + weight_block: str = "16x16" + + def __post_init__(self) -> None: + # The grouped GEMM requires K % 128 == 0 and N % 128 == 0; both + # projections (dim x hidden_dim and hidden_dim x dim) hit both + # dims, so reject the model-dim violations up front here. + for name in ("dim", "hidden_dim"): + value = getattr(self, name) + if value % _NVFP4_BLOCK: + raise ValueError( + f"NVFP4 requires {name} divisible by {_NVFP4_BLOCK}; " + f"got {name}={value}. NVFP4 cannot quantize these " + "grouped experts; exclude them from the converter fqns." + ) + _validate_four_over_six_knobs( + self.backward_override, self.weight_block, grouped=True + ) + + def __init__(self, config: Config): + super().__init__(config) + self._four_over_six_kwargs = dict( + err_mode=config.err_mode, + e4m3_scale_bound=config.e4m3_scale_bound, + row_scaled_activation=config.row_scaled_activation, + weight_block=config.weight_block, + backward_override=config.backward_override, + ) + + def _grouped_mm(self, *, A, B_t, offs): + # The hook receives B_t pre-transposed to (E, K, N); the torchao op + # takes expert weights in their stored (E, N, K) layout (it makes + # the operands contiguous itself). + return four_over_six_grouped_mm( + A, + B_t.transpose(-2, -1), + offs, + pad_token_groups_for_grouped_mm=True, + **self._four_over_six_kwargs, + ) + + FourOverSixGroupedExperts.__name__ = f"NVFP4FourOverSix{parent_cls.__name__}" + FourOverSixGroupedExperts.__qualname__ = f"NVFP4FourOverSix{parent_cls.__name__}" + _four_over_six_experts_cache[parent_cls] = FourOverSixGroupedExperts + return FourOverSixGroupedExperts + + +class NVFP4FourOverSixGroupedExpertsConverter(QuantizationConverter): + """Apply four-over-six NVFP4 quantization to MoE expert grouped GEMMs. + + The miles NVFP4 RL recipes quantize only the routed-expert projections; + with no Linear converter alongside, this converter alone reproduces that + allow-list (attention, dense MLP, shared experts, router, embeddings, and + the LM head all stay bf16). + """ + + @dataclass(kw_only=True, slots=True) + class Config(QuantizationConverter.Config): + fqns: list[str] = field(default_factory=list) + """ + List of fully qualified names of modules to apply four-over-six NVFP4 + quantization to. Only GroupedExperts.Config entries whose FQN contains + a match are converted. If empty, all GroupedExperts are converted. + """ + + err_mode: str = "mae" + """Candidate-selection error metric, 'mae' or 'mse'.""" + + e4m3_scale_bound: int = 256 + """Global E4M3 scale bound; 256 leaves map-to-4 headroom.""" + + row_scaled_activation: bool = False + """One FP32 global scale per activation row instead of per token + group.""" + + backward_override: str | None = None + """'high_precision' (the default when None) or 'dequantized'; grouped + four-over-six has no quantized backward.""" + + weight_block: str = "16x16" + """Weight tile granularity; '1x16' mirrors + NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1.""" + + pad_multiple: int = 128 + """ + Pad per-expert token groups to this multiple for NVFP4 grouped GEMM + alignment (the four-over-six grouped GEMM requires 128-row groups). + """ + + def __init__(self, config: Config): + self.config = config + + if four_over_six_grouped_mm is None: + raise ImportError( + "torchao is not installed or does not provide the NVFP4 " + "four-over-six grouped training prototype. Install a torchao " + "build with torchao.prototype.moe_training.nvfp4_training." + ) + + if not has_cuda_capability(10, 0): + raise ValueError("NVFP4 is only supported on SM100 or later architectures") + + if not self.config.model_compile_enabled: + logger.warning( + "torch.compile enablement is required for highest performance " + "of NVFP4 dynamic quantization." + ) + + def convert(self, model_config): + assert four_over_six_grouped_mm is not None + fqns = self.config.fqns + for fqn, config, parent, attr in model_config.traverse(GroupedExperts.Config): + if fqns and not any(target_fqn in fqn for target_fqn in fqns): + continue + # ``parent`` is the RoutedExperts.Config owning inner_experts + dispatcher. + swap_token_dispatcher(parent, self.config.pad_multiple) + base_module_cls = type(config)._owner + quantized_cls = _get_four_over_six_grouped_experts_cls(base_module_cls) + config_cls = quantized_cls.Config # type: ignore[attr-defined] + new_config = config_cls( + **{f.name: getattr(config, f.name) for f in fields(config)}, + err_mode=self.config.err_mode, + e4m3_scale_bound=self.config.e4m3_scale_bound, + row_scaled_activation=self.config.row_scaled_activation, + backward_override=self.config.backward_override, + weight_block=self.config.weight_block, + ) + if parent is None: + model_config = new_config + elif isinstance(parent, list): + parent[attr] = new_config + else: + setattr(parent, attr, new_config) + + logger.info( + "Converted GroupedExperts to use dynamic NVFP4 four-over-six " + "quantization for grouped_mm ops" + ) + return model_config diff --git a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py index dc2ea29c09..3648d256a0 100644 --- a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py +++ b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py @@ -16,6 +16,7 @@ from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer +from torchtitan.components.quantization import NVFP4FourOverSixGroupedExpertsConverter from torchtitan.config import ( CompileConfig, DebugConfig, @@ -788,6 +789,37 @@ def rl_grpo_qwen3_30b_a3b_varlen() -> Controller.Config: ) +def rl_grpo_qwen3_30b_a3b_varlen_nvfp4_four_over_six() -> Controller.Config: + """Qwen3-30B-A3B GRPO with four-over-six NVFP4 routed experts. + + The torchtitan analog of the miles NVFP4 RL reference recipe (which trains + the same model): only the routed-expert grouped GEMMs quantize, with + row-scaled activations, MSE candidate selection, E4M3 bound 256, 1x16 + weight blocks, and the high-precision backward. Because the trainer and + the vLLM generator share this one model_spec, both actors run the + identical four-over-six forward -- the policy trains on the same quantized + function the generator samples from, which is the recipe's + train/inference-consistency point (no quantized-weight sync needed; the + bf16 master weights sync as usual and each side re-quantizes dynamically). + """ + config = rl_grpo_qwen3_30b_a3b_varlen() + config.model_spec = model_registry( + "30B-A3B", + attn_backend="varlen", + converters=[ + NVFP4FourOverSixGroupedExpertsConverter.Config( + row_scaled_activation=True, + err_mode="mse", + e4m3_scale_bound=256, + weight_block="1x16", + backward_override="high_precision", + pad_multiple=128, + ), + ], + ) + return config + + def rl_grpo_qwen3_30b_a3b_varlen_perf() -> Controller.Config: """Qwen3-30B-A3B GRPO with throughput overrides (8 GPUs: 4 gen + 4 train). diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 892527a1bf..76309a3426 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -14,7 +14,9 @@ Float8LinearConverter, MXFP8GroupedExpertsConverter, MXFP8LinearConverter, + NVFP4FourOverSixGroupedExpertsConverter, ) +from torchtitan.components.quantization.nvfp4 import nvfp4_bf16_first_last_fqns from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed.activation_checkpoint import SelectiveAC from torchtitan.hf_datasets.text_datasets import DATASETS @@ -109,6 +111,57 @@ def deepseek_v3_debugmodel_mxfp8() -> Trainer.Config: return config +def _deepseek_v3_debugmodel_nvfp4_four_over_six( + backward_override: str, fqns: list[str] | None = None +) -> Trainer.Config: + # Quantize only the routed-expert grouped GEMMs with four-over-six NVFP4, + # mirroring the miles NVFP4 RL recipes: row-scaled activations, MSE + # candidate selection, E4M3 bound 256, and 1x16 weight blocks + # (NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1). With no Linear converter, + # attention, the dense-layer feed-forward, shared experts, the MoE router + # gate, embeddings, and the lm_head all stay bf16 -- the recipe's + # experts-only allow-list. + config = deepseek_v3_debugmodel() + model_compile_enabled = ( + config.compile.enable and "model" in config.compile.components + ) + config.model_spec = model_registry( + "debugmodel", + converters=[ + NVFP4FourOverSixGroupedExpertsConverter.Config( + model_compile_enabled=model_compile_enabled, + fqns=fqns or [], + row_scaled_activation=True, + err_mode="mse", + e4m3_scale_bound=256, + weight_block="1x16", + backward_override=backward_override, + pad_multiple=128, + ), + ], + ) + return config + + +def deepseek_v3_debugmodel_nvfp4_four_over_six() -> Trainer.Config: + # The miles NVFP4 RL base recipe point: high-precision backward + # (NVTE_BACKWARD_OVERRIDE=high_precision) on every routed-expert layer. + return _deepseek_v3_debugmodel_nvfp4_four_over_six("high_precision") + + +def deepseek_v3_debugmodel_nvfp4_four_over_six_dequantized() -> Trainer.Config: + # The advanced miles recipe variant (the GLM-5.2 NVFP4 e2e analog): the + # dequantized backward (NVTE_BACKWARD_OVERRIDE=dequantized) backpropagates + # through bf16 GEMMs on the dequantized fprop operands, and the first and + # last decoder layers stay bf16 (--first-last-layers-bf16 with one layer + # at each end). The debugmodel's layer 0 is dense (no routed experts), so + # the leading-layer exclusion is vacuous there but keeps the recipe shape. + # The debugmodel has 6 layers; layers 1..4 quantize. + return _deepseek_v3_debugmodel_nvfp4_four_over_six( + "dequantized", fqns=nvfp4_bf16_first_last_fqns(6, 1, 1) + ) + + def deepseek_v3_debugmodel_hybridep() -> Trainer.Config: config = deepseek_v3_debugmodel() config.model_spec = model_registry( From da8674a549a6e279cbc11354d677057816fbbc36 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 20:38:39 -0700 Subject: [PATCH 2/6] Slice grouped activations to offs[-1] for four-over-six expert GEMMs The TorchAOTokenDispatcher swapped in by the grouped-experts converter already 128-aligns every expert group, but it over-allocates the activation buffer past offs[-1], while torchao's four_over_six_grouped_mm requires offs[-1] == A.shape[0]. Passing the buffer through with pad_token_groups_for_grouped_mm=True violates that contract, and simply disabling padding would let the unwritten tail rows feed the per-group amaxes and silently poison the last expert's quantization in per-tensor mode. Slice the logical rows before the op, skip the op's own padding (the dispatcher's alignment already satisfies it), and zero-extend the output so downstream shapes match; pad routes zero gradients to the tail and the unpermute never reads those rows. --- torchtitan/components/quantization/nvfp4.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/torchtitan/components/quantization/nvfp4.py b/torchtitan/components/quantization/nvfp4.py index 0324787c12..2d12014b48 100644 --- a/torchtitan/components/quantization/nvfp4.py +++ b/torchtitan/components/quantization/nvfp4.py @@ -615,13 +615,23 @@ def _grouped_mm(self, *, A, B_t, offs): # The hook receives B_t pre-transposed to (E, K, N); the torchao op # takes expert weights in their stored (E, N, K) layout (it makes # the operands contiguous itself). - return four_over_six_grouped_mm( - A, + # + # The swapped TorchAOTokenDispatcher already 128-aligns every + # expert group but over-allocates the activation buffer past + # offs[-1], while the torchao op requires offs[-1] == A.shape[0] + # and the unwritten tail rows must not feed the per-group amaxes. + # Slice to the logical rows, skip the op's own padding, and + # zero-extend the output (pad routes zero grads to the tail; + # the unpermute never reads tail rows). + m_total = int(offs[-1]) + out = four_over_six_grouped_mm( + A[:m_total], B_t.transpose(-2, -1), offs, - pad_token_groups_for_grouped_mm=True, + pad_token_groups_for_grouped_mm=False, **self._four_over_six_kwargs, ) + return torch.nn.functional.pad(out, (0, 0, 0, A.shape[0] - m_total)) FourOverSixGroupedExperts.__name__ = f"NVFP4FourOverSix{parent_cls.__name__}" FourOverSixGroupedExperts.__qualname__ = f"NVFP4FourOverSix{parent_cls.__name__}" From 0532b02494e3e048855e6ea83fbdbeefb34d89a3 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 20:45:23 -0700 Subject: [PATCH 3/6] Disable CUDA graphs in the four-over-six DSv3 recipes The quantized grouped GEMM reads expert group offsets on the host and the row-scaled mode loops dense GEMMs per group, both of which CUDA-graph capture forbids (cudaErrorStreamCaptureUnsupported in the debugmodel's captured fwd+bwd). miles runs its quantized recipes with CUDA graphs off for the same class of reason. --- torchtitan/models/deepseek_v3/config_registry.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 76309a3426..7f0dfd7248 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -122,6 +122,11 @@ def _deepseek_v3_debugmodel_nvfp4_four_over_six( # gate, embeddings, and the lm_head all stay bf16 -- the recipe's # experts-only allow-list. config = deepseek_v3_debugmodel() + # The quantized grouped GEMM reads the expert group offsets on the host + # (and the row-scaled path loops dense GEMMs per group), which CUDA-graph + # capture forbids; miles likewise runs its quantized recipes with CUDA + # graphs off. + config.training.disable_cuda_graphs = True model_compile_enabled = ( config.compile.enable and "model" in config.compile.components ) From c259be88552c1ee9c90cac1f0571d492edbce9da Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 20:55:02 -0700 Subject: [PATCH 4/6] Reject compile with row-scaled four-over-six grouped experts The row-scaled grouped forward host-reads the group offsets and loops dense GEMMs per group; under torch.compile fullgraph (nonstrict_trace + capture_scalar_outputs) the offsets become unbacked SymInts and the loop bounds guard on data-dependent expressions with no graph-break escape. Fail at config time instead of at trace time. --- torchtitan/components/quantization/nvfp4.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/torchtitan/components/quantization/nvfp4.py b/torchtitan/components/quantization/nvfp4.py index 2d12014b48..34ae0879a0 100644 --- a/torchtitan/components/quantization/nvfp4.py +++ b/torchtitan/components/quantization/nvfp4.py @@ -646,6 +646,11 @@ class NVFP4FourOverSixGroupedExpertsConverter(QuantizationConverter): with no Linear converter alongside, this converter alone reproduces that allow-list (attention, dense MLP, shared experts, router, embeddings, and the LM head all stay bf16). + + The row-scaled grouped forward loops dense GEMMs per token group with + host-read offsets, so it cannot be captured by torch.compile; combining + row_scaled_activation with model compile is rejected at config time (a + fused single-GEMM row-scaled variant is future work). """ @dataclass(kw_only=True, slots=True) @@ -694,6 +699,14 @@ def __init__(self, config: Config): if not has_cuda_capability(10, 0): raise ValueError("NVFP4 is only supported on SM100 or later architectures") + if self.config.model_compile_enabled and self.config.row_scaled_activation: + raise ValueError( + "row-scaled four-over-six grouped GEMMs loop dense GEMMs per " + "token group with host-read offsets and cannot be captured by " + "torch.compile; run this converter eager (a fused single-GEMM " + "row-scaled variant is future work)." + ) + if not self.config.model_compile_enabled: logger.warning( "torch.compile enablement is required for highest performance " From 8906831c1c765101a2c3fdb71a5fecdf455f657c Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 20:55:02 -0700 Subject: [PATCH 5/6] Count four-over-six converters in has_quantization has_quantization enumerated only the Float8/MXFP8/NVFP4Linear configs and the float8/mxfp8 grouped-experts caches, so a model quantized by the four-over-six converters reported has_quantization=False and MFU was computed against the bf16 peak. --- torchtitan/components/quantization/utils.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/torchtitan/components/quantization/utils.py b/torchtitan/components/quantization/utils.py index 0de363dfa8..f7aefbfb4a 100644 --- a/torchtitan/components/quantization/utils.py +++ b/torchtitan/components/quantization/utils.py @@ -71,7 +71,11 @@ def has_quantization(model_config) -> bool: Float8Linear, ) from torchtitan.components.quantization.mx import _mxfp8_experts_cache, MXFP8Linear - from torchtitan.components.quantization.nvfp4 import NVFP4Linear + from torchtitan.components.quantization.nvfp4 import ( + _four_over_six_experts_cache, + NVFP4FourOverSixLinear, + NVFP4Linear, + ) quant_linear_types: list[type] = [] if Float8Linear is not None: @@ -80,6 +84,8 @@ def has_quantization(model_config) -> bool: quant_linear_types.append(MXFP8Linear.Config) if NVFP4Linear is not None: quant_linear_types.append(NVFP4Linear.Config) + if NVFP4FourOverSixLinear is not None: + quant_linear_types.append(NVFP4FourOverSixLinear.Config) has_quant_linear = bool(quant_linear_types) and any( isinstance(config, tuple(quant_linear_types)) @@ -87,7 +93,11 @@ def has_quantization(model_config) -> bool: ) quant_experts_types = tuple( cls.Config # type: ignore[attr-defined] - for cls in (*_float8_experts_cache.values(), *_mxfp8_experts_cache.values()) + for cls in ( + *_float8_experts_cache.values(), + *_mxfp8_experts_cache.values(), + *_four_over_six_experts_cache.values(), + ) ) has_quant_moe = bool(quant_experts_types) and any( isinstance(config, quant_experts_types) From 8717093f17da2aacc53c62bae86f70a10fadbab9 Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 21:37:23 -0700 Subject: [PATCH 6/6] Unit-test the compile rejection and has_quantization coverage Covers the config-time ValueError for model compile + row-scaled grouped four-over-six, and has_quantization returning True for models converted by either four-over-six converter (and False for the stock debugmodel). --- tests/unit_tests/test_quantization.py | 62 +++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/unit_tests/test_quantization.py b/tests/unit_tests/test_quantization.py index 798ea475e8..ad8d78bee2 100644 --- a/tests/unit_tests/test_quantization.py +++ b/tests/unit_tests/test_quantization.py @@ -640,3 +640,65 @@ def test_nvfp4_four_over_six_linear_converter_plumbs_new_knobs(monkeypatch): assert converted assert all(lc.backward_override == "dequantized" for lc in converted) assert all(lc.weight_block == "1x16" for lc in converted) + + +def test_four_over_six_grouped_converter_rejects_compile_with_row_scaled(monkeypatch): + _four_over_six_grouped_cls() + import torchtitan.components.quantization.nvfp4 as nvfp4_mod + from torchtitan.components.quantization import ( + NVFP4FourOverSixGroupedExpertsConverter, + ) + + monkeypatch.setattr(nvfp4_mod, "has_cuda_capability", lambda *_: True) + # The row-scaled grouped forward host-reads offsets and loops per group, + # which fullgraph compile cannot capture -> rejected at config time. + with pytest.raises(ValueError, match="torch.compile"): + NVFP4FourOverSixGroupedExpertsConverter( + NVFP4FourOverSixGroupedExpertsConverter.Config( + model_compile_enabled=True, + row_scaled_activation=True, + ) + ) + # Either knob alone stays accepted. + NVFP4FourOverSixGroupedExpertsConverter( + NVFP4FourOverSixGroupedExpertsConverter.Config( + model_compile_enabled=True, + ) + ) + NVFP4FourOverSixGroupedExpertsConverter( + NVFP4FourOverSixGroupedExpertsConverter.Config( + row_scaled_activation=True, + ) + ) + + +def test_has_quantization_counts_four_over_six(monkeypatch): + _four_over_six_grouped_cls() + from torchtitan.components.quantization import ( + NVFP4FourOverSixLinear, + NVFP4FourOverSixLinearConverter, + ) + + if NVFP4FourOverSixLinear is None: + pytest.skip("torchao NVFP4 four-over-six training prototype not available") + + stock = ConfigManager().parse_args( + ["--module", "deepseek_v3", "--config", "deepseek_v3_debugmodel"] + ) + assert not has_quantization(stock.model_spec.model) + + grouped = _parse_recipe( + monkeypatch, "deepseek_v3", "deepseek_v3_debugmodel_nvfp4_four_over_six" + ) + assert has_quantization(grouped) + + import torchtitan.components.quantization.nvfp4 as nvfp4_mod + + monkeypatch.setattr(nvfp4_mod, "has_cuda_capability", lambda *_: True) + dense = ConfigManager().parse_args( + ["--module", "llama3", "--config", "llama3_debugmodel"] + ) + converter = NVFP4FourOverSixLinearConverter( + NVFP4FourOverSixLinearConverter.Config(fqns=["layers"]) + ) + assert has_quantization(converter.convert(dense.model_spec.model))