From 64708a2472b0bac19b9a300ad959c9e67cd062b7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 13 Aug 2026 12:54:31 -0700 Subject: [PATCH 01/68] finegrained kernels --- src/transformers/__init__.py | 2 + src/transformers/integrations/__init__.py | 2 + src/transformers/integrations/finegrained.py | 1480 +++++++++++++++++ .../integrations/finegrained_fp8.py | 10 + src/transformers/integrations/hub_kernels.py | 5 +- src/transformers/modeling_utils.py | 6 +- src/transformers/quantizers/auto.py | 22 +- .../quantizers/quantizer_finegrained.py | 394 +++++ .../quantizers/quantizer_finegrained_fp8.py | 11 + src/transformers/utils/quantization_config.py | 57 +- 10 files changed, 1974 insertions(+), 15 deletions(-) create mode 100644 src/transformers/integrations/finegrained.py create mode 100644 src/transformers/quantizers/quantizer_finegrained.py diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 47e5160d9521..bbbc023f641a 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -259,6 +259,7 @@ "CompressedTensorsConfig", "EetqConfig", "FbgemmFp8Config", + "FineGrainedConfig", "FineGrainedFP8Config", "FourOverSixConfig", "FPQuantConfig", @@ -787,6 +788,7 @@ from .utils.quantization_config import CompressedTensorsConfig as CompressedTensorsConfig from .utils.quantization_config import EetqConfig as EetqConfig from .utils.quantization_config import FbgemmFp8Config as FbgemmFp8Config + from .utils.quantization_config import FineGrainedConfig as FineGrainedConfig from .utils.quantization_config import FineGrainedFP8Config as FineGrainedFP8Config from .utils.quantization_config import FourOverSixConfig as FourOverSixConfig from .utils.quantization_config import FPQuantConfig as FPQuantConfig diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 4dc3aec9f71b..b22274b31f4f 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,6 +50,7 @@ ], "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], + "finegrained": ["FineGrainedExperts", "FineGrainedLinear", "replace_with_finegrained_layer"], "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "gemma_quant": [ @@ -212,6 +213,7 @@ ) from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear + from .finegrained import FineGrainedExperts, FineGrainedLinear, replace_with_finegrained_layer from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .gemma_quant import ( diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py new file mode 100644 index 000000000000..e00fa94cf20a --- /dev/null +++ b/src/transformers/integrations/finegrained.py @@ -0,0 +1,1480 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import functools +import importlib +import os +import sys +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn as nn +from torch.nn import functional as F + +from ..activations import ACT2FN +from ..core_model_loading import ConversionOps +from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module +from ..utils import logging +from ..utils.deprecation import deprecate_kwarg +from ..utils.import_utils import is_kernels_available +from .deepgemm import ( + deepgemm_fp8_fp4_experts_forward, + deepgemm_fp8_fp4_linear, + deepgemm_fp8_fp4_megamoe_experts_forward, + is_deepgemm_loadable, + is_sm100, +) +from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel +from .moe import ExpertsInterface, use_experts_implementation +from .tensor_parallel import to_local + + +logger = logging.get_logger(__name__) + + +_FP8_DTYPE = torch.float8_e4m3fn +_FP8_MIN = torch.finfo(_FP8_DTYPE).min +_FP8_MAX = torch.finfo(_FP8_DTYPE).max + + +@functools.cache +def _get_ue8m0_dtype() -> torch.dtype: + """Return ``torch.float8_e8m0fnu`` or raise a clear error on torch without FP8 support. + + UE8M0 scales are always stored/consumed as this single dtype — the kernels (Triton + finegrained + DeepGEMM) read it natively, and supporting the same scales in mixed + container dtypes would be a mess — so fail loudly rather than fall back.""" + if not hasattr(torch, "float8_e8m0fnu"): + raise RuntimeError( + "scale_fmt='ue8m0' requires torch.float8_e8m0fnu, which is only available in " + f"PyTorch >= 2.7 (found {torch.__version__}). Upgrade torch to use UE8M0 FP8 checkpoints." + ) + return torch.float8_e8m0fnu + + +def _first_attr(obj, *names): + for name in names: + if hasattr(obj, name): + return getattr(obj, name) + raise AttributeError(f"{type(obj).__name__} has none of: {names}") + + +@dataclass(frozen=True) +class FineGrained: + """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. + + Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the three + matmuls, with the recipe resolved off the weight/scale dtypes inside the kernel; the + per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary classes ship + in the same build. MoE blocks are COMPOSED from the matmuls rather than wrapped: the + experts forwards pass ``epilogue`` to the gate_up GEMM when the activation is fusable + (no bias, supported act_fn) and fall back to the unfused two-GEMM form otherwise. All + symbols are required — a build missing any raises at load with the full list. + """ + + matmul: Callable + batched_matmul: Callable + grouped_matmul: Callable + nvfp4_quantize_two_level: Callable + swizzle_mx_scales: Callable + compute_grouped_scheduling: Callable + weighted_reduce: Callable + Quantization: Callable + Epilogue: Callable + + +# Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env +# can change between attempts. A module global (not `@functools.cache`) avoids Dynamo warning about +# tracing a cache-wrapped function on every compile. +_FINEGRAINED: FineGrained | None = None + + +def _import_local_finegrained(): + """A locally importable `finegrained_kernels` package takes precedence over the hub build: + `FINEGRAINED_KERNELS_PATH` (a checkout's `torch-ext` directory, or any directory containing + the package) is prepended to `sys.path`, then a plain import is attempted either way — an + installed / already-on-path package also wins. Returns the module, or `None` to fall back + to the `kernels` hub load.""" + path = os.environ.get("FINEGRAINED_KERNELS_PATH") + if path: + if not os.path.isdir(path): + raise ImportError(f"FINEGRAINED_KERNELS_PATH does not exist: {path}") + if path not in sys.path: + sys.path.insert(0, path) + try: + return importlib.import_module("finegrained_kernels") + except ImportError: + if path: + raise # an explicit local path that fails to import is a setup error, not a fallback + return None + + +@torch._dynamo.allow_in_graph +def _load_finegrained_kernel() -> None: + """ + Load the finegrained-fp8 Triton kernel once into the `_FINEGRAINED` module global. + + `@allow_in_graph` makes `torch.compile` treat the untraceable hub download + dynamic import as a + single opaque node instead of tracing into it; it returns `None` (proxyable) and populates the + global, which `load_finegrained_kernel` then returns. + + Under NO circumstances may this function return a value: an `@allow_in_graph` fx node's + return must be proxyable, and returning the bundle (e.g. from the warm-cache + short-circuit) breaks torch.compile with `Unsupported: torch.* op returned non-Tensor`. + + Raises `ImportError` if the `kernels` package is missing, or the kernel or required + symbols cannot be found. + """ + global _FINEGRAINED + if _FINEGRAINED is not None: + return + + kernel = _import_local_finegrained() + if kernel is None: + if not is_kernels_available(): + raise ImportError(f"finegrained-fp8 kernel unavailable: {_MISSING_KERNELS_MESSAGE}") + kernel = lazy_load_kernel("finegrained-kernels") + if kernel is None: + raise ImportError( + "Failed to load the finegrained-kernels kernel — check that `kernels-community/finegrained-kernels` " + "has a build matching the current torch/CUDA." + ) + + required = ( + "matmul_2d", + "matmul_batched", + "matmul_grouped", + "nvfp4_quantize_two_level", + "swizzle_mx_scales", + "compute_grouped_scheduling", + "weighted_reduce", + "Quantization", + "Epilogue", + ) + symbols = {name: getattr(kernel, name, None) for name in required} + missing = [name for name, attr in symbols.items() if attr is None] + if missing: + raise ImportError( + f"finegrained-kernels build is missing required symbols: {', '.join(missing)}. {_MISSING_KERNELS_MESSAGE}" + ) + + _FINEGRAINED = FineGrained( + matmul=symbols["matmul_2d"], + batched_matmul=symbols["matmul_batched"], + grouped_matmul=symbols["matmul_grouped"], + nvfp4_quantize_two_level=symbols["nvfp4_quantize_two_level"], + compute_grouped_scheduling=symbols["compute_grouped_scheduling"], + weighted_reduce=symbols["weighted_reduce"], + swizzle_mx_scales=symbols["swizzle_mx_scales"], + Quantization=symbols["Quantization"], + Epilogue=symbols["Epilogue"], + ) + + +def load_finegrained_kernel() -> FineGrained: + _load_finegrained_kernel() + return _FINEGRAINED + + +def _cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return (a + b - 1) // b + + +@dataclass(frozen=True) +class _WeightFormat: + """Storage layout of one quantized weight format — the single source both module classes + derive their parameter shapes from. ``scale_dtype`` ``None`` defers to the config's + ``scale_fmt`` (block-FP8 ships fp32 or UE8M0 containers; the group formats pin theirs). + ``scale_group`` ``None`` means the block comes from the quant config's ``weight_block_size``.""" + + weight_dtype: torch.dtype + values_per_byte: int = 1 + scale_dtype: torch.dtype | None = None + scale_group: tuple[int, int] | None = None + has_global_scale: bool = False + + +def _weight_formats() -> dict[str, _WeightFormat]: + return { + # block-scaled E4M3, block from the quant config, fp32/UE8M0 scale container + "fp8": _WeightFormat(weight_dtype=_FP8_DTYPE), + # E4M3 values, UE8M0 group-32 scales + "mxfp8": _WeightFormat(weight_dtype=_FP8_DTYPE, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32)), + # packed E2M1 values (2/byte), UE8M0 group-32 scales (dsv4 ships scale_fmt's container) + "mxfp4": _WeightFormat(weight_dtype=torch.int8, values_per_byte=2, scale_group=(1, 32)), + # packed E2M1 values, E4M3 group-16 block scales, per-tensor/per-expert fp32 global + "nvfp4": _WeightFormat( + weight_dtype=torch.int8, + values_per_byte=2, + scale_dtype=_FP8_DTYPE, + scale_group=(1, 16), + has_global_scale=True, + ), + } + + +def resolve_weight_format( + weight_format: str, + scale_fmt: str = "float", + block_size: tuple[int, int] | None = None, +) -> tuple[_WeightFormat, torch.dtype, tuple[int, int] | None]: + """``(format, scale_dtype, (sf_gran_n, sf_gran_k))`` for one format name, with the config's + ``scale_fmt``/``weight_block_size`` filling the slots the format leaves open.""" + formats = _weight_formats() + if weight_format not in formats: + raise ValueError(f"unknown weight_format {weight_format!r}; expected one of {sorted(formats)}") + fmt = formats[weight_format] + scale_dtype = fmt.scale_dtype + if scale_dtype is None: + scale_dtype = _get_ue8m0_dtype() if scale_fmt == "ue8m0" else torch.float32 + return fmt, scale_dtype, fmt.scale_group if fmt.scale_group is not None else block_size + + +def _alloc_expert_proj( + num_experts: int, + proj_out: int, + proj_in: int, + weight_dtype: torch.dtype, + sf_dtype: torch.dtype, + weight_k_div: int = 1, + sf_gran_n: int | None = None, + sf_gran_k: int | None = None, + min_sf_out: int = 1, +) -> tuple[nn.Parameter, nn.Parameter]: + """Allocate `(weight, weight_scale_inv)` parameters for one expert projection. + + `weight_k_div` halves the K dim for FP4-packed storage (2 e2m1 values per byte). + `sf_gran_n` / `sf_gran_k` set per-block (None → per-row/per-tensor) SF granularity. + `min_sf_out` floors the SF tensor's output dim — used by the fused gate_up + projection to keep room for both halves (pass `2`) even when `proj_out < sf_gran_n` + would otherwise collapse the SF dim to 1. + """ + weight_t = torch.empty(num_experts, proj_out, proj_in // weight_k_div, dtype=weight_dtype) + weight = nn.Parameter(weight_t, requires_grad=weight_t.is_floating_point()) + sf_out = max(_cdiv(proj_out, sf_gran_n) if sf_gran_n is not None else 1, min_sf_out) + sf_in = _cdiv(proj_in, sf_gran_k) if sf_gran_k is not None else 1 + sf_t = torch.empty(num_experts, sf_out, sf_in, dtype=sf_dtype) + sf = nn.Parameter(sf_t, requires_grad=sf_t.is_floating_point()) + return weight, sf + + +@deprecate_kwarg("output_dtype", version="v5.16") +def finegrained_triton_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + block_size: list[int] | None = None, + bias: torch.Tensor | None = None, + activation_scale: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + weight_global_scale: torch.Tensor | None = None, + activation_format: str | None = None, +) -> torch.Tensor: + """Triton fine-grained linear: fused act-quant + matmul, then optional bias add. + + Serves every weight recipe the kernel resolves off the tensors themselves — block-FP8 + (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4 (``int8``-packed values; the two-level + per-tensor ``weight_global_scale`` recovers on the accumulator). ``block_size`` is accepted + for back-compat and ignored: the quantization block is derived from the scale tensor's + shape. ``activation_scale=None`` → dynamic activation quant (inline); a per-tensor scalar → + static quant against it. ``activation_format`` picks the activation recipe where the weights + leave it open (``None`` = the weight-native choice; ``"bf16"`` = weight-only, no activation + quant — W4A16). ``output_dtype`` defaults to ``input.dtype`` when left ``None``. + """ + kernel = load_finegrained_kernel() + quantization = None + if activation_format is not None: + recipe = None if activation_format == "bf16" else activation_format + quantization = kernel.Quantization(input_recipe=recipe) + original_shape = input.shape + output = kernel.matmul( + input.reshape(-1, original_shape[-1]), + weight, + activation_scale, + weight_scale_inv, + quantization=quantization, + output_dtype=output_dtype if output_dtype is not None else input.dtype, + b_global_scale=weight_global_scale, + ) + output = output.reshape(*original_shape[:-1], output.shape[-1]) + if bias is not None: + output.add_(bias) + return output + + +@deprecate_kwarg("output_dtype", version="v5.16") +def finegrained_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + block_size: list[int] | None = None, + bias: torch.Tensor | None = None, + activation_scale: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + allow_deepgemm: bool = True, + weight_global_scale: torch.Tensor | None = None, + activation_format: str | None = None, +) -> torch.Tensor: + """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. + + Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales: + 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — 3-6× faster on the shapes it supports. + Preferred for FP4, UE8M0 SFs, and 128×128 block FP8. + 2. Triton finegrained-fp8 fallback — used when DeepGEMM is unavailable, when the + caller passes ``activation_scale`` (DeepGEMM is dynamic-only), or for any + shape DeepGEMM declined. + + Args: + input: (..., K) bf16/fp16 activations. + weight: (N, K) `float8_e4m3fn` or (N, K // 2) `int8` (FP4-packed). + weight_scale_inv: per-block weight scales — `float32` (V3-style) or `float8_e8m0fnu` + (V4-style; reinterpreted as int32 at the DeepGEMM kernel boundary). + block_size: [block_n, block_k] for FP8 block-wise quant, or None/[N, K] for per-tensor. + Ignored for FP4 weights (the kernel infers SF granularity from the dtype). + bias: optional bias added to the matmul output. + activation_scale: pass a per-tensor scalar to use static activation quant; leave `None` + for dynamic (per-token) quant. + allow_deepgemm: set ``False`` to force the Triton fallback for this call. Used when the + model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound + to a single CUDA context and produce garbage across devices (see the multi-device guard + in ``quantizer_finegrained.py``). + """ + # DeepGEMM is CUDA-only, dynamic-only, SM90+ only, FP4/FP8-block-128-only. On SM100 its FP8 GEMM only + # consumes UE8M0 scales; float32 scales would be ceil-rounded to UE8M0 without requantizing and + # silently corrupt the output (#47030), so we skip DeepGEMM up front for that combo rather than + # attempt-then-fall-back every call — `_assert_sm100_requirements` still guards it as a backstop. + # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces the Triton fallback for this single dispatcher (the + # experts ``"deepgemm"`` impl is unaffected — use ``set_experts_implementation`` for that). Used by the + # FP8 MoE batched_mm / grouped_mm paths to avoid a still-unexplained DeepGEMM-vs-Triton interaction that + # degrades end-to-end generation on B200 (per-row kernel outputs still measure bit-perfect, but final + # tokens drift; not reproducible with the DeepGEMM linear off). + deepgemm_preferred = ( + allow_deepgemm + and is_deepgemm_loadable() + and activation_scale is None + # DeepGEMM serves neither the NVFP4 two-level global nor an explicit activation format + and weight_global_scale is None + and activation_format is None + and not (is_sm100() and weight_scale_inv.dtype == torch.float32) + and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) + and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" + ) + + if deepgemm_preferred: + try: + return deepgemm_fp8_fp4_linear( + input, + weight, + weight_scale_inv, + block_size=block_size, + activation_scale=activation_scale, + bias=bias, + ) + except (ImportError, NotImplementedError, ValueError) as e: + # DeepGEMM is loadable but declined this specific input; fall back to Triton, which is more + # permissive (handles FP8/FP4 with float32 or UE8M0 scales on any arch, plus input dtypes + # DeepGEMM rejects). If Triton can't serve it either, it raises its own error. + # - NotImplementedError: an arch/input combo DeepGEMM has no kernel for (FP4 on Hopper — + # `is_deepgemm_loadable` is dtype-agnostic and passes there — or float32 scales on Blackwell); + # - ValueError: an input DeepGEMM rejects but Triton supports (e.g. activations it won't quantize); + # - ImportError: a symbol/build gap. + logger.warning_once( + f"DeepGEMM declined this call, falling back to Triton. Reason: {e} " + "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely." + ) + + return finegrained_triton_linear( + input, + weight, + weight_scale_inv, + block_size, + bias, + activation_scale, + weight_global_scale=weight_global_scale, + activation_format=activation_format, + ) + + +class FineGrainedLinear(nn.Linear): + # Internal, temporary flag — not public API, don't set it directly. `_disable_deepgemm_on_multi_device` + # flips it True at load when the model spans >1 CUDA device in one process (DeepGEMM's context-bound + # kernels corrupt across devices); removable once the kernel ships a context-free loader. + _deepgemm_disabled = False + + def __init__( + self, + in_features: int, + out_features: int, + block_size: tuple[int, int] | None = None, + activation_scheme: str = "dynamic", + scale_fmt: str = "float", + has_bias: bool = False, + weight_format: str = "fp8", + activation_format: str | None = None, + has_global_scale: bool | None = None, + ): + super().__init__(in_features, out_features) + + self.has_bias = has_bias + self.weight_format = weight_format + self.block_size = block_size + self.activation_scheme = activation_scheme + self.activation_format = activation_format + # the format table decides storage: value dtype (packed E2M1 = 2 values per int8 byte + # along K), scale dtype/granularity, and whether a two-level global exists + fmt, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + if has_global_scale is None: + has_global_scale = fmt.has_global_scale + in_storage = in_features // fmt.values_per_byte + self.weight = torch.nn.Parameter( + torch.empty(out_features, in_storage, dtype=fmt.weight_dtype), + requires_grad=fmt.weight_dtype.is_floating_point, + ) + if has_global_scale: + # NVFP4 two-level: the per-tensor fp32 global the kernel recovers on the accumulator + self.weight_global_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) + else: + self.register_parameter("weight_global_scale", None) + + if scale_group is None: + # no group and no block: one per-tensor scale + self.weight_scale_inv = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) + else: + scale_out_features = _cdiv(out_features, scale_group[0]) + scale_in_features = _cdiv(in_features, scale_group[1]) + self.weight_scale_inv = nn.Parameter( + torch.empty(scale_out_features, scale_in_features, dtype=sf_dtype), + requires_grad=sf_dtype.is_floating_point, + ) + + if self.activation_scheme == "static": + self.activation_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) + else: + self.register_parameter("activation_scale", None) + + if self.has_bias: + self.bias = nn.Parameter(torch.empty(self.out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.weight.element_size() > 1: + return F.linear(input, self.weight, self.bias) + + weight = to_local(self.weight) + scale_inv = to_local(self.weight_scale_inv) + + return finegrained_linear( + input, + weight, + scale_inv, + block_size=self.block_size, + activation_scale=self.activation_scale, + bias=self.bias, + allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, + activation_format=self.activation_format, + ) + + +class FineGrainedGroupedLinear(FineGrainedLinear): + """FP8 drop-in for block-diagonal grouped linears. + + The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)` + weight; logically that's `n_groups` independent `(out_per_group, in_per_group)` + sub-matrices, each consuming a disjoint slice of the input's last-but-one dim. + Forward expects input of shape `(..., n_groups, in_per_group)` and returns + `(..., n_groups, out_per_group)` — same contract as the vanilla bf16 grouped + linear it replaces. + + """ + + def __init__( + self, + in_features_per_group: int, + out_features: int, + n_groups: int, + block_size: tuple[int, int] | None = None, + activation_scheme: str = "dynamic", + scale_fmt: str = "float", + has_bias: bool = False, + ): + super().__init__( + in_features=in_features_per_group, + out_features=out_features, + block_size=block_size, + activation_scheme=activation_scheme, + scale_fmt=scale_fmt, + has_bias=has_bias, + ) + self.n_groups = n_groups + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_shape = x.shape[:-2] + hidden_dim = x.shape[-1] + + if self.weight.element_size() > 1: + w = self.weight.view(self.n_groups, -1, hidden_dim).transpose(1, 2) + x = x.reshape(-1, self.n_groups, hidden_dim).transpose(0, 1) + y = torch.bmm(x, w).transpose(0, 1) + y = y.reshape(*input_shape, self.n_groups, -1) + if self.has_bias: + y.add_(self.bias.view(self.n_groups, -1)) + return y + + w = to_local(self.weight) + scale_inv = to_local(self.weight_scale_inv) + + w = w.view(self.n_groups, -1, hidden_dim) + x = x.movedim(-2, 0).reshape(-1, hidden_dim) + scale_inv = scale_inv.view(self.n_groups, scale_inv.size(0) // self.n_groups, scale_inv.size(1)) + + tokens_per_group = x.size(0) // self.n_groups + # (E+1,) row boundaries — the kernels' expert_start schedule + expert_start = torch.arange(0, self.n_groups + 1, device=x.device, dtype=torch.int32) * tokens_per_group + + kernel = load_finegrained_kernel() + y = kernel.grouped_matmul( + x, + w, + None, + scale_inv, + expert_start=expert_start, + quantization=_kernel_quantization(kernel, self.activation_format), + b_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, + ) + y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) + if self.has_bias: + y.add_(self.bias.view(self.n_groups, -1)) + return y + + +_FUSABLE_ACT_FNS = ("silu", "gelu", "relu") + + +def _kernel_epilogue(kernel, module) -> object | None: + """The fused gate|up + GLU epilogue for the gate_up GEMM, when the kernel can apply it: + gated experts whose activation the epilogue implements and whose gate_up has NO bias (a + fused GLU leaves nowhere to add one — biased models take the unfused two-GEMM form).""" + act_name = getattr(module, "act_fn_name", None) + if not module.has_gate or module.has_bias or act_name not in _FUSABLE_ACT_FNS: + return None + return kernel.Epilogue( + gate=True, + act_fn=act_name, + swiglu_alpha=module.swiglu_alpha, + swiglu_limit=module.swiglu_limit, + ) + + +def _kernel_quantization(kernel, activation_format: str | None): + """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` + keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" + if activation_format is None: + return None + recipe = None if activation_format == "bf16" else activation_format + return kernel.Quantization(input_recipe=recipe) + + +@functools.lru_cache(maxsize=64) +def _batched_gather_idx(num_tokens: int, num_top_k: int, device: torch.device) -> torch.Tensor: + """The batched routed-row gather map (row ``s`` reads token ``s // num_top_k``): depends + only on the SHAPE, so cache it — its build is pure eager launch overhead at decode.""" + return torch.arange(num_tokens * num_top_k, device=device, dtype=torch.int32) // num_top_k + + +def _activation_recipe(module) -> str | None: + """The activation quantization recipe this module's chain runs: the explicit + ``activation_format`` when set (``"bf16"`` = weight-only -> None), else the + weight-native default (fp8 weights quantize activations to block-FP8, MX/NV weights + to their own family — the same resolution the kernels apply).""" + fmt = module.activation_format + if fmt is not None: + return None if fmt == "bf16" else fmt + return module.weight_format if module.weight_format is not None else "fp8" + + +def _gate_up_quantization(kernel, module, epilogue): + """The gate_up GEMM's ``Quantization`` — WITH the fused intermediate requant whenever the + GLU is fused: ``output_recipe`` makes the epilogue emit the quantized intermediate + ``(C, Cs)`` that the down GEMM consumes directly (``As=Cs``), matching the kernels' own + ``moe_fused_*`` chains. Without the fused epilogue the intermediate is produced host-side + in bf16 and the down GEMM quantizes it internally (the unfused form).""" + recipe = _activation_recipe(module) + if epilogue is None or recipe is None: + return _kernel_quantization(kernel, module.activation_format) + return kernel.Quantization(input_recipe=recipe, output_recipe=recipe) + + +def _proj_scale(module, proj: str) -> torch.Tensor: + """A projection's weight scales — the pre-swizzled cache when the post-load hook built + one (the tcgen05 fast path), else the affine Parameter.""" + swizzled = getattr(module, f"{proj}_scale_inv_swizzled", None) + return swizzled if swizzled is not None else to_local(getattr(module, f"{proj}_scale_inv")) + + +def _cached_kernel_objects(kernel, module): + """The module's static kernel-call objects, built once (per-call dataclass construction + is measurable at decode): ``(epilogue, gate_up_quantization, down_quantization)``.""" + if not hasattr(module, "_cached_epilogue"): + module._cached_epilogue = _kernel_epilogue(kernel, module) + module._cached_gate_up_quantization = _gate_up_quantization(kernel, module, module._cached_epilogue) + module._cached_down_quantization = _kernel_quantization(kernel, module.activation_format) + return module._cached_epilogue, module._cached_gate_up_quantization, module._cached_down_quantization + + +def _apply_unfused_gate_up(self, proj_out, up_bias_ids): + """The unfused gate_up tail (fused-epilogue models never reach this): optional + per-expert bias (rows indexed by ``up_bias_ids``; sentinel rows clamp to a valid + expert — the reduce skips them), then the gating/activation in torch.""" + if self.has_bias: + up_bias = to_local(self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) + proj_out = proj_out + up_bias[up_bias_ids.clamp(max=self.num_experts - 1)] + return self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) + + +def _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, out_dtype): + """The shared bookend after the down GEMM (rows in ROUTED order): optional down bias + (BEFORE the routing-weight multiply — each expert output is ``x @ W_d^T + b_d``), then + the kernels' routing-weighted top-k reduce (skips EP-sentinel rows, whose GEMM output + is uninitialized by contract).""" + if self.has_bias: + down_bias = to_local(self.down_proj_bias) + routed_ids = top_k_index.reshape(-1) + proj_out = proj_out + down_bias[routed_ids.clamp(max=self.num_experts - 1)] + return kernel.weighted_reduce(proj_out, top_k_index, top_k_weights, self.num_experts).to(out_dtype) + + +def finegrained_batched_mm_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Batched (decode) experts forward — the kernels' fused chain: gate_up per routed row + (fused GLU epilogue + intermediate requant where supported) -> down consuming the + requant handoff -> ``weighted_reduce``.""" + if self.activation_scheme == "static": + raise NotImplementedError( + "batched_mm experts dispatch does not support activation_scheme='static'. " + "Use the default eager dispatch or switch to activation_scheme='dynamic'." + ) + + kernel = load_finegrained_kernel() + num_top_k = top_k_index.size(-1) + + # per-routed-row dispatch: the kernel gathers row s's token (s // num_top_k) itself — + # no (S, H) replica copy; the ops take the router's int64 ids natively + gather_idx = _batched_gather_idx(hidden_states.size(0), num_top_k, hidden_states.device) + expert_ids = top_k_index.reshape(-1) # (S,) + + up_name = "gate_up_proj" if self.has_gate else "up_proj" + epilogue, gate_up_quantization, down_quantization = _cached_kernel_objects(kernel, self) + + proj_out = kernel.batched_matmul( + hidden_states, + to_local(getattr(self, up_name)), + None, + _proj_scale(self, up_name), + expert_ids=expert_ids, + gather_idx=gather_idx, + epilogue=epilogue, + quantization=gate_up_quantization, + b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, + ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) or (S, I) + + inter_scale = None + if isinstance(proj_out, (tuple, list)): + proj_out, inter_scale = proj_out # the fused requant's quantized intermediate + + if epilogue is None: + # batched rows are in ROUTED order, so the bias indexes by the flat routed ids + proj_out = _apply_unfused_gate_up(self, proj_out, expert_ids) + + proj_out = kernel.batched_matmul( + proj_out, + to_local(self.down_proj), + inter_scale, # fused requant handoff: pre-quantized intermediate scales (None = raw) + _proj_scale(self, "down_proj"), + expert_ids=expert_ids, + quantization=None if inter_scale is not None else down_quantization, + b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, + ) # (S, hidden_dim), routed order + + return _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, hidden_states.dtype) + + +def finegrained_grouped_mm_experts_forward( + self: torch.nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + """Grouped (prefill) experts forward — the kernels' fused chain: one on-device routing + pass -> gate_up over the expert-sorted schedule (fused GLU epilogue + intermediate + requant where supported) -> down consuming the requant handoff and scattering back to + routed rows -> ``weighted_reduce``.""" + if self.activation_scheme == "static": + raise NotImplementedError( + "grouped_mm experts dispatch does not support activation_scheme='static'. " + "Use the default eager dispatch or switch to activation_scheme='dynamic'." + ) + + kernel = load_finegrained_kernel() + num_top_k = top_k_index.size(-1) + + # one on-device routing pass (counting sort, no host sync): expert-sorted row starts + + # the gather/scatter maps. EP sentinels (id >= num_experts) fall past expert_start[-1]; + # their output rows stay uninitialized and the weighted_reduce bookend skips them. + expert_start, gather_idx, scatter_idx = kernel.compute_grouped_scheduling(top_k_index, self.num_experts, num_top_k) + + up_name = "gate_up_proj" if self.has_gate else "up_proj" + epilogue, gate_up_quantization, down_quantization = _cached_kernel_objects(kernel, self) + + proj_out = kernel.grouped_matmul( + hidden_states, + to_local(getattr(self, up_name)), + None, + _proj_scale(self, up_name), + expert_start=expert_start, + gather_idx=gather_idx, + epilogue=epilogue, + quantization=gate_up_quantization, + b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, + ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) + + inter_scale = None + if isinstance(proj_out, (tuple, list)): + proj_out, inter_scale = proj_out # the fused requant's quantized intermediate + + if epilogue is None: + # grouped gate_up rows are expert-SORTED; only biased unfused models pay the sort + up_bias_ids = torch.sort(top_k_index.reshape(-1)).values if self.has_bias else top_k_index.reshape(-1) + proj_out = _apply_unfused_gate_up(self, proj_out, up_bias_ids) + + proj_out = kernel.grouped_matmul( + proj_out, + to_local(self.down_proj), + inter_scale, # fused requant handoff: pre-quantized intermediate scales (None = raw) + _proj_scale(self, "down_proj"), + expert_start=expert_start, + scatter_idx=scatter_idx, # scatter straight to routed rows — weighted_reduce's layout + quantization=None if inter_scale is not None else down_quantization, + b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, + ) # (S, hidden_dim), ROUTED order after the scatter + + return _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, hidden_states.dtype) + + +class FineGrainedExperts(nn.Module): + # Internal, temporary flag — not public API, don't set it directly. `_disable_deepgemm_on_multi_device` + # flips it True at load when the model spans >1 CUDA device in one process (DeepGEMM's context-bound + # kernels corrupt across devices); removable once the kernel ships a context-free loader. + _deepgemm_disabled = False + + # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. + # The plan dicts store `{module-path-pattern: parallel-layer-kind}`; this maps an + # old kind to a new kind, and the quantizer rewrites every plan VALUE that matches. + # The default `MoeTensorParalellExperts` kind is impl-agnostic; some impls need a + # distinct TP layer (e.g. megamoe needs no gradient-sync hooks and an EP + # `process_group` injection). Declared here so the quantizer doesn't have to know + # about impl-specific TP needs — extend this dict when adding new impls. + _impl_tp_layer_overrides: dict[str, dict[str, str]] = { + "deepgemm_megamoe": { + "moe_tp_experts": "megamoe_experts", + "ep_router": "megamoe_router", + }, + } + + def __init__( + self, + config, + block_size: tuple[int, int] | None = None, + activation_scheme: str = "dynamic", + scale_fmt: str = "float", + has_bias: bool = False, + has_gate: bool = True, + weight_format: str | None = None, + activation_format: str | None = None, + has_global_scale: bool | None = None, + ): + super().__init__() + + self.config = config + self.has_bias = has_bias + self.has_gate = has_gate + self.block_size = block_size + self.hidden_dim = config.hidden_size + self.has_global_scale = has_global_scale + self.activation_format = activation_format + self.activation_scheme = activation_scheme + self.swiglu_alpha = getattr(config, "swiglu_alpha", None) + self.swiglu_limit = getattr(config, "swiglu_limit", None) + self.num_experts = _first_attr(config, "num_local_experts", "num_experts") + self.intermediate_dim = _first_attr(config, "moe_intermediate_size", "intermediate_size") + self.act_fn_name = _first_attr(config, "hidden_activation", "hidden_act") + self.limit = getattr(config, "swiglu_limit", None) + self.act_fn = ACT2FN[self.act_fn_name] + + # Expert weight storage is declared by the QUANT config's format key, not the model + # config: `weight_format` arrives from `replace_with_finegrained_layer` as the + # checkpoint's quant_method ("fp8", "mxfp8", "mxfp4", "nvfp4"). The DeepSeek V4-style + # `config.expert_dtype = "fp4"` model-config side-channel predates it and is kept as + # the legacy fallback when no format is passed (dsv4 fp4 experts under the "fp8" key). + if weight_format is None: + weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" + self.weight_format = weight_format + fmt, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + if has_global_scale is None: + has_global_scale = fmt.has_global_scale + # the forwards gate on the ATTRIBUTE — it must carry the format-resolved value, not + # the raw ctor arg (None would silently drop the nvfp4 global at every forward) + self.has_global_scale = has_global_scale + alloc_kwargs = { + "weight_dtype": fmt.weight_dtype, + "sf_dtype": sf_dtype, + "weight_k_div": fmt.values_per_byte, + "sf_gran_n": scale_group[0] if scale_group is not None else None, + "sf_gran_k": scale_group[1] if scale_group is not None else None, + } + + if self.has_gate: + self.gate_up_proj, self.gate_up_proj_scale_inv = _alloc_expert_proj( + self.num_experts, 2 * self.intermediate_dim, self.hidden_dim, min_sf_out=2, **alloc_kwargs + ) + if self.has_bias: + self.gate_up_proj_bias = nn.Parameter( + torch.empty(self.num_experts, 2 * self.intermediate_dim, dtype=torch.float32) + ) + else: + self.register_parameter("gate_up_proj_bias", None) + else: + self.up_proj, self.up_proj_scale_inv = _alloc_expert_proj( + self.num_experts, self.intermediate_dim, self.hidden_dim, **alloc_kwargs + ) + if self.has_bias: + self.up_proj_bias = nn.Parameter( + torch.empty(self.num_experts, self.intermediate_dim, dtype=torch.float32) + ) + else: + self.register_parameter("up_proj_bias", None) + + self.down_proj, self.down_proj_scale_inv = _alloc_expert_proj( + self.num_experts, self.hidden_dim, self.intermediate_dim, **alloc_kwargs + ) + # NVFP4 two-level: per-expert fp32 globals the kernels recover on the accumulator + for proj in (("gate_up_proj" if self.has_gate else "up_proj"), "down_proj"): + if self.has_global_scale: + setattr( + self, + f"{proj}_global_scale", + nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)), + ) + else: + self.register_parameter(f"{proj}_global_scale", None) + if self.has_bias: + self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, dtype=torch.float32)) + else: + self.register_parameter("down_proj_bias", None) + + if self.activation_scheme == "static": + self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) + self.down_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) + + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: + gate, up = gate_up.chunk(2, dim=-1) + if self.swiglu_alpha is not None: + # Clamped SwiGLU-OAI gate (same math as the model's non-quantized experts). + gate = gate.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + glu = gate * torch.sigmoid(gate * self.swiglu_alpha) + return (up + 1.0) * glu + elif self.limit is not None: + gate = gate.clamp(max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + return self.act_fn(gate) * up + + def forward( + self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor + ) -> torch.Tensor: + # index_add_ will accumulate using the dtype of the tensor we write into + # so we use float32 for the accumulation to avoid numerical issues in bf16/fp16 + final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32) + + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1) + + for expert_idx in expert_hit: + if expert_idx == self.num_experts: + continue + + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate_up_act_scale = ( + self.gate_up_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None + ) + proj_out = self.linear( + current_state, + self.gate_up_proj[expert_idx] if self.has_gate else self.up_proj[expert_idx], + self.gate_up_proj_scale_inv[expert_idx] if self.has_gate else self.up_proj_scale_inv[expert_idx], + activation_scale=gate_up_act_scale, + ) + proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) + down_act_scale = ( + self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None + ) + proj_out = self.linear( + proj_out, + self.down_proj[expert_idx], + self.down_proj_scale_inv[expert_idx], + activation_scale=down_act_scale, + ) + routing_weights = top_k_weights[token_idx, top_k_pos, None] + weighted_out = proj_out * routing_weights.to(proj_out.dtype) + final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) + return final_hidden_states.to(hidden_states.dtype) + + def linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + activation_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + if weight.element_size() > 1: + return F.linear(input, weight, None) + + return finegrained_linear( + input, + weight, + weight_scale_inv, + self.block_size, + activation_scale=activation_scale, + allow_deepgemm=not self._deepgemm_disabled, + ) + + +class FineGrainedExpertsInterface(ExpertsInterface): + """Interface for registering custom FP8 experts forward functions.""" + + _global_mapping = { + "batched_mm": finegrained_batched_mm_experts_forward, + "grouped_mm": finegrained_grouped_mm_experts_forward, + "deepgemm": deepgemm_fp8_fp4_experts_forward, + "deepgemm_megamoe": deepgemm_fp8_fp4_megamoe_experts_forward, + } + + +ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() + + +def swizzle_scales_after_loading(model: nn.Module) -> None: + """Pre-swizzle MX/NVFP4 weight block scales into the ``SWIZZLE_32_4_4`` layout the Blackwell + tcgen05 scaled-MMA reads directly (one triton launch per matrix, values unchanged) — the + kernels accept both layouts, but plain row-major forces a per-tile gather that caps the + scaled dot below the fp8/fp4 peak. Runs once post-load; the swizzled artifact is cached as + a BUFFER next to the affine scale Parameter (which stays the state_dict source of truth, so + save/dequantize round-trip untouched; scale tensors are small). gate_up scales take the + gate-interleaved artifact the fused GATE arm reads natively. + ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` skips the pass (debug / pre-Blackwell A/B).""" + if os.environ.get("TRANSFORMERS_FINEGRAINED_NO_SWIZZLE", "0") == "1": + return + try: + kernel = load_finegrained_kernel() + except ImportError: + return # dequantize fallback path: nothing to swizzle for + for module in model.modules(): + if not isinstance(module, FineGrainedExperts): + continue + for proj in (("gate_up_proj" if module.has_gate else "up_proj"), "down_proj"): + scale = getattr(module, f"{proj}_scale_inv", None) + # only the group-scaled families have a swizzled layout (E8M0 / E4M3 group scales); + # block-FP8's fp32 (N/128, K/128) grid has no tcgen05 layout to pre-arrange + if scale is None or scale.dtype == torch.float32 or scale.device.type == "meta": + continue + gate = proj == "gate_up_proj" + n_rows = to_local(scale).shape[-2] + # the swizzled layout needs whole 128-row blocks per projection (under gate the + # split sits at N, mid-block unless N % 128 == 0, e.g. GPT-OSS N=2880); the + # kernels read the affine layout directly, so just skip the pre-swizzle + if (n_rows // 2 if gate else n_rows) % 128: + continue + swizzled = kernel.swizzle_mx_scales(to_local(scale).data, gate=gate) + module.register_buffer(f"{proj}_scale_inv_swizzled", swizzled, persistent=False) + + +def _disable_deepgemm_on_multi_device(model: nn.Module) -> None: + """Internal, temporary helper (not public API): flag every FP8 module to skip DeepGEMM when the + model spans >1 CUDA device in one process. + + DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` handle to the + CUDA context live at load time; driving that cached handle from another device launches it against + the wrong context and produces garbage. (Build-time fix: compile DeepGEMM with + `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until our wheel picks that up + we avoid single-process multi-device.) Setting `_deepgemm_disabled` routes both the linear and + experts paths through Triton/grouped_mm. A model that fits on one device keeps DeepGEMM even with + other GPUs visible; TP/EP put one device per process, so this is a no-op there. + """ + quantized_modules = [m for m in model.modules() if isinstance(m, (FineGrainedLinear, FineGrainedExperts))] + cuda_devices = set() + for m in quantized_modules: + param = next(m.parameters(), None) + if param is not None and param.device.type == "cuda": + cuda_devices.add(param.device.index) + if len(cuda_devices) <= 1: + return + for m in quantized_modules: + m._deepgemm_disabled = True + logger.warning_once( + "This FP8 model spans multiple CUDA devices in one process; routing its FP8 linear and experts " + "layers through Triton/grouped_mm instead of DeepGEMM (DeepGEMM's cached kernels are bound to a " + "single CUDA context and corrupt across devices). Run tensor/expert parallel (one device per " + "process) to use the faster DeepGEMM path." + ) + + +def replace_with_finegrained_layer( + model, modules_to_not_convert: list[str] | None = None, quantization_config=None, pre_quantized=False +): + """ + A helper function to replace all `torch.nn.Linear` modules by `FineGrainedLinear` modules. + + Parameters: + model (`torch.nn.Module`): + Input model or `torch.nn.Module` as the function is run recursively. + modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`): + Names of the modules to not convert. In practice we keep the `lm_head` in full precision for numerical stability reasons. + quantization_config (`FineGrainedConfig`): + The quantization config object that contains the quantization parameters. + pre_quantized (`bool`, defaults to `False`): + Whether the model is pre-quantized or not + """ + + if quantization_config.dequantize: + return model + + # The checkpoint's quant_method IS the weight format ("fp8", "mxfp8", "mxfp4", "nvfp4"). + # Under the "fp8" key pass None: dsv4-style checkpoints declare fp4 EXPERTS through the + # legacy `config.expert_dtype` model-config side-channel, which the ctor falls back to. + quant_method = getattr(quantization_config, "quant_method", None) + quant_method = getattr(quant_method, "value", quant_method) + # modelopt exports are NVFP4-only (validated by the config translation) + if quant_method == "modelopt": + quant_method = "nvfp4" + weight_format = quant_method if quant_method in ("mxfp8", "mxfp4", "nvfp4") else None + + has_been_replaced = False + for module_name, module in model.named_modules(): + if not should_convert_module(module_name, modules_to_not_convert): + continue + + new_module = None + with torch.device("meta"): + if module_name.endswith(".experts"): + has_gate = getattr(module, "has_gate", True) + has_bias = getattr(module, "has_bias", False) + config = getattr(module, "config", model.config.get_text_config()) + new_class = use_experts_implementation( + experts_class=FineGrainedExperts, + experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + has_bias=has_bias, + has_gate=has_gate, + ) + new_module = new_class( + config=config, + block_size=quantization_config.weight_block_size, + activation_scheme=quantization_config.activation_scheme, + scale_fmt=quantization_config.scale_fmt, + has_bias=has_bias, + has_gate=has_gate, + weight_format=weight_format, + activation_format=getattr(quantization_config, "activation_format", None), + ) + # GPT-OSS hardcodes its clamped-SwiGLU parameters as MODULE attributes + # (alpha=1.702, limit=7.0) rather than config fields — carry them over so + # `_apply_gate` (and the fused epilogue) keep the reference numerics. + if getattr(module, "alpha", None) is not None and new_module.swiglu_alpha is None: + new_module.swiglu_alpha = module.alpha + new_module.swiglu_limit = getattr(module, "limit", None) + elif type(module) is nn.Linear: + # Vanilla `nn.Linear` → standard FineGrainedLinear swap. + new_module = FineGrainedLinear( + in_features=module.in_features, + out_features=module.out_features, + block_size=quantization_config.weight_block_size, + activation_scheme=quantization_config.activation_scheme, + scale_fmt=quantization_config.scale_fmt, + has_bias=module.bias is not None, + weight_format=weight_format or "fp8", + activation_format=getattr(quantization_config, "activation_format", None), + ) + elif isinstance(module, nn.Linear) and "GroupedLinear" in type(module).__name__: + # Block-diagonal grouped linear (e.g. DSv4's `DeepseekV4GroupedLinear`): + # one underlying weight conceptually split into `n_groups` independent + # sub-matmuls fed by disjoint input slices. Vanilla `FineGrainedLinear` would + # collapse those groups into one giant linear and yield the wrong + # output dim, so swap to `FineGrainedGroupedLinear` which keeps the per-group + # bmm contract and runs each block as its own FP8 matmul. + new_module = FineGrainedGroupedLinear( + in_features_per_group=module.in_features, + out_features=module.out_features, + n_groups=module.n_groups, + block_size=quantization_config.weight_block_size, + activation_scheme=quantization_config.activation_scheme, + scale_fmt=quantization_config.scale_fmt, + has_bias=module.bias is not None, + ) + if new_module is not None: + model.set_submodule(module_name, new_module) + has_been_replaced = True + + if not has_been_replaced: + logger.warning( + "You are loading your model using fp8 but no linear modules were found in your model." + " Please double check your model architecture." + ) + return model + + +def _deinterleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor: + """GPT-OSS fuses gate|up INTERLEAVED along the output dim ([g0, u0, g1, u1, ...]); the + finegrained kernels (and `FineGrainedExperts._apply_gate`'s ``chunk``) read the STACKED + form ([gate; up]). One permutation at load, values untouched.""" + return torch.cat([t[:, 0::2], t[:, 1::2]], dim=1) + + +class FineGrainedMxfp4Deserialize(ConversionOps): + """Convert GPT-OSS MXFP4 checkpoint tensors to the finegrained layout: ``{proj}_blocks`` + ``(E, N, K/32, 16)`` uint8 (two low-nibble-first E2M1 values per byte) reshapes to the + packed ``(E, N, K/2)`` int8 the kernels read directly — same nibble order — and + ``{proj}_scales`` (biased-127 exponent bytes) bitcast to ``float8_e8m0fnu``. No swizzle + (the kernels take affine scales; the post-load hook builds the swizzled cache). gate_up + rows de-interleave from GPT-OSS's [g0, u0, ...] to the stacked [gate; up] form.""" + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def convert( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module | None = None, + full_layer_name: str | None = None, + missing_keys: list[str] | None = None, + **kwargs, + ) -> dict[str, torch.Tensor]: + proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj" + blocks = input_dict[f"{proj}_blocks"] + scales = input_dict[f"{proj}_scales"] + blocks = blocks[0] if isinstance(blocks, list) else blocks + scales = scales[0] if isinstance(scales, list) else scales + + num_experts, rows = blocks.shape[0], blocks.shape[1] + weight = blocks.reshape(num_experts, rows, -1).view(torch.int8) + # deinterleave BEFORE the e8m0 bitcast: cat/index ops aren't implemented for + # float8_e8m0fnu on CUDA, and the reorder is byte-identical either way + if proj == "gate_up_proj": + weight = _deinterleave_gate_up_rows(weight) + scales = _deinterleave_gate_up_rows(scales) + scale_inv = scales.contiguous().view(_get_ue8m0_dtype()) + prefix = full_layer_name.rsplit(".", 1)[0] + return {f"{prefix}.{proj}": weight.contiguous(), f"{prefix}.{proj}_scale_inv": scale_inv} + + +class FineGrainedGateUpBiasDeinterleave(ConversionOps): + """GPT-OSS ``gate_up_proj_bias`` follows the interleaved output layout; de-interleave to + match the stacked weight rows produced by `FineGrainedMxfp4Deserialize`.""" + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): + # single-source converter: input_dict is keyed by the SOURCE PATTERN string (with + # any regex anchors), so don't look up by a hardcoded name + bias = next(iter(input_dict.values())) + bias = bias[0] if isinstance(bias, list) else bias + return {full_layer_name: _deinterleave_gate_up_rows(bias)} + + +class FineGrainedViewPackedInt8(ConversionOps): + """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules + store (same bytes — ``copy_`` into an int8 param would numerically CONVERT and corrupt + values >= 128). Passes non-uint8 tensors through untouched, so it can ride converters + whose pattern also matches unquantized modules.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): + v = next(iter(input_dict.values())) + v = v[0] if isinstance(v, list) else v + if torch.is_tensor(v) and v.dtype == torch.uint8: + v = v.view(torch.int8) + return {full_layer_name: v} + + +class FineGrainedFuseEqualGlobals(ConversionOps): + """Reduce modelopt's per-half second-level globals to the single per-expert global the + stacked gate_up GEMM applies: the halves are calibrated together and ship bit-identical + (verified across the GLM-5.2-NVFP4 checkpoint), so this asserts equality and keeps one. + Also handles the single-source (down_proj) case, where it just flattens to ``(E,)``.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): + stacks = [] + for v in input_dict.values(): + v = torch.stack(v, dim=0) if isinstance(v, list) else v + stacks.append(v.reshape(-1).float()) + first = stacks[0] + for other in stacks[1:]: + if not torch.equal(first, other): + raise ValueError( + f"modelopt per-half globals differ for {full_layer_name} — the stacked " + "gate_up GEMM applies one global per expert; this checkpoint needs a " + "block-scale refold, which is not implemented." + ) + return {full_layer_name: first} + + +class FineGrainedQuantize(ConversionOps): + """ + A quantization operation that creates two tensors, weight and scale out of a weight. + """ + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def _resolve_block_size(self, value: torch.Tensor) -> tuple[int, int]: + block_size = None + if self.hf_quantizer.quantization_config is not None: + if isinstance(self.hf_quantizer.quantization_config, dict): + block_size = self.hf_quantizer.quantization_config.get("weight_block_size") + else: + block_size = getattr(self.hf_quantizer.quantization_config, "weight_block_size", None) + if block_size is None: + block_size = (value.shape[-2], value.shape[-1]) + return tuple(block_size) + + def _quantize_one(self, key: str, value: torch.Tensor) -> dict[str, torch.Tensor]: + # Pass through tensors that aren't tileable (1D norms / biases, or shapes + # that don't divide cleanly by the configured block) — they were never + # FP8-quantized on the load side, so the reverse op shouldn't touch them. + if value.ndim < 2: + return {key: value} + block_m, block_n = self._resolve_block_size(value) + rows, cols = value.shape[-2], value.shape[-1] + if rows % block_m != 0 or cols % block_n != 0: + return {key: value} + + # Leading dims can be empty (2D) or include num_experts/... (3D+) + leading_shape = value.shape[:-2] + rows_tiles = rows // block_m + cols_tiles = cols // block_n + original_shape = value.shape + value_fp32 = value.to(torch.float32) + # Reshape to (..., rows_tiles, block_m, cols_tiles, block_n) + reshaped = value_fp32.reshape(*leading_shape, rows_tiles, block_m, cols_tiles, block_n) + # Per-tile max-abs over the block dims (block_m at -3, block_n at -1) + max_abs = reshaped.abs().amax(dim=(-3, -1)) + safe_max_abs = torch.where(max_abs > 0, max_abs, torch.ones_like(max_abs)) + # We store inverse scale to match the upstream ``weight_scale_inv`` convention + scales = _FP8_MAX / safe_max_abs + scales = torch.where(max_abs > 0, scales, torch.ones_like(scales)) # keep zeros stable + inv_scales = (1.0 / scales).to(torch.float32) + # ue8m0 stores weight_scale_inv as a power of two. Round it before quantizing and derive the + # forward scale from it, so dequant multiplies by the exact scale the weight was divided by. + if self.hf_quantizer.quantization_config.scale_fmt == "ue8m0": + inv_scales = torch.pow(2.0, torch.ceil(torch.log2(inv_scales.clamp(min=torch.finfo(torch.float32).tiny)))) + inv_scales = inv_scales.to(_get_ue8m0_dtype()) + scales = 1.0 / inv_scales.to(torch.float32) # forward scale = exact reciprocal of the stored inverse + # Broadcast scales over the block dims and quantize + scales_broadcast = scales.unsqueeze(-1).unsqueeze(-3) # (..., rows_tiles, 1, cols_tiles, 1) + scaled = reshaped * scales_broadcast + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) + quantized = quantized.reshape(original_shape) + scale_key = key.rsplit(".", 1)[0] + ".weight_scale_inv" if key.endswith(".weight") else key + "_scale_inv" + return {key: quantized, scale_key: inv_scales} + + def convert(self, input_dict: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]: + # Quantize every (key, tensor) entry in the dict. Single-tensor case (legacy + # callers that pass one key) and multi-tensor case (reverse of an expert + # ``MergeModulelist`` that emits one key per expert) are handled the same way. + result: dict[str, torch.Tensor] = {} + for key, value in input_dict.items(): + tensor = value[0] if isinstance(value, list) else value + result.update(self._quantize_one(key, tensor)) + return result + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedDequantize(self.hf_quantizer) + + +class FineGrainedDequantize(ConversionOps): + """Dequantize FP8 weights using their per-block ``weight_scale_inv``. + + Designed to run as the *first* op in any :class:`WeightConverter` chain when + loading with ``dequantize=True`` — :meth:`update_weight_conversions` on the + FP8 quantizer attaches it to each existing model-specific converter so that + per-expert (weight, scale) pairs are folded into full-precision tensors before + the chain's merge / concat ops collapse the per-expert structure. + + Pattern semantics + Input ``input_dict`` carries one entry per source pattern; each value is a + list of tensors (one per ``*`` match). For every weight pattern that has a + sibling ``*.weight_scale_inv`` pattern in the dict, this op pairs them up by + index, dequantizes per-pair, and emits the dequantized list under the + original *weight* key. Scale entries are dropped from the output so the + remaining ops only see weights. + """ + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def _scale_pattern_for(self, weight_pattern: str) -> str: + # Strip the optional ``$`` regex anchor so we can match the underlying name. + anchored = weight_pattern.endswith("$") + base = weight_pattern[:-1] if anchored else weight_pattern + if base.endswith(".weight"): + scale = base[: -len(".weight")] + ".weight_scale_inv" + elif base == "weight": + scale = "weight_scale_inv" + else: + scale = base + "_scale_inv" + return scale + "$" if anchored else scale + + # E2M1 (FP4) value table — checkpoints sometimes ship MoE experts as packed FP4 + # (two e2m1 nibbles per int8 byte), so the "weight" dtype lands as ``int8`` / + # ``float4_e2m1fn_x2`` and we have to unpack before applying the scale grid. + _FP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: + """Two ``e2m1`` FP4 values per byte → float32 tensor twice as wide on the last dim.""" + lut = torch.tensor(self._FP4_E2M1_LUT, dtype=torch.float32, device=packed.device) + u8 = packed.contiguous().view(torch.uint8) + low = (u8 & 0xF).long() + high = ((u8 >> 4) & 0xF).long() + unpacked = torch.stack([lut[low], lut[high]], dim=-1) + return unpacked.reshape(*packed.shape[:-1], 2 * packed.shape[-1]) + + def _dequantize_one( + self, quantized: torch.Tensor, scales: torch.Tensor, output_dtype: torch.dtype | None = None + ) -> torch.Tensor: + # FP4 path: int8 / float4_e2m1fn_x2 stores two nibbles per byte. Unpack to fp32 + # first so the rest of the routine sees a normal (rows, cols) float matrix. + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if quantized.dtype == torch.int8 or (fp4_dtype is not None and quantized.dtype == fp4_dtype): + quantized_fp32 = self._unpack_fp4(quantized) + else: + quantized_fp32 = quantized.to(torch.float32) + rows, cols = quantized_fp32.shape[-2:] + # Derive block size from the scale grid rather than the global config: MoE experts + # ship MXFP4 with a ``[1, 32]`` block, dense linears ship FP8 with ``[128, 128]``, + # and the same dequant has to handle both within one checkpoint. + try: + scale_rows, scale_cols = scales.shape[-2:] + except Exception: + # scale can be a single tensor in extreme cases where it was not wrapped properly but is [1,0]. + scale_rows, scale_cols = 1, 1 + if rows % scale_rows or cols % scale_cols: + raise ValueError( + f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." + ) + block_m = rows // scale_rows + block_n = cols // scale_cols + # ``ue8m0`` (``float8_e8m0fnu``) scales have no CUDA ``mul`` kernel, and casting + # the FP8 weight to that dtype loses precision. Promote both sides to fp32 for + # the math; prefer the destination parameter's dtype when known so eager modules + # (e.g. plain ``nn.Linear``) keep the model's compute dtype after load. + if output_dtype is None: + output_dtype = ( + scales.dtype if scales.dtype.is_floating_point and scales.element_size() >= 2 else torch.bfloat16 + ) + # MXFP8 checkpoints ship E8M0 exponents stored as ``torch.uint8`` (one byte per + # block) — the actual scale is `2 ** (byte - 127)`. Interpreting the raw bytes + # as scalar multipliers would be silently wrong, so unpack to fp32 here. + if scales.dtype == torch.uint8: + s_fp32 = (scales.to(torch.float32) - 127.0).exp2() + else: + s_fp32 = scales.to(torch.float32) + original_shape = quantized_fp32.shape + q = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) + s = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) + return (q * s).to(output_dtype).reshape(original_shape) + + def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: + if model is None or full_layer_name is None: + return None + module, tensor_name = get_module_from_name(model, full_layer_name) + param = getattr(module, tensor_name, None) + return getattr(param, "dtype", None) + + def convert( + self, + input_dict: dict[str, list[torch.Tensor] | torch.Tensor], + full_layer_name: str | None = None, + model: torch.nn.Module | None = None, + **kwargs, + ) -> dict[str, list[torch.Tensor] | torch.Tensor]: + output_dtype = self._get_target_dtype(model, full_layer_name) + # Backward-compatible single-tensor path (the legacy fallback converter declares + # ``["weight$", "weight_scale_inv", "activation_scale"]`` and produces a single + # ``weight`` target). Also handles the no-scale case (e.g. RMSNorm weights that + # match ``weight$`` but ship no ``weight_scale_inv`` alongside). + if "weight$" in input_dict: + # The downstream renamer in `core_model_loading._convert_one_module` uses the + # output dict's *key*, not its content, to derive prefix/suffix; if `full_layer_name` + # is unset (direct invocation / tests) fall back to the legacy converter's target. + target_key = full_layer_name if full_layer_name is not None else "weight" + quantized = input_dict["weight$"] + quantized = quantized[0] if isinstance(quantized, list) else quantized + if "weight_scale_inv" in input_dict: + scales = input_dict["weight_scale_inv"] + scales = scales[0] if isinstance(scales, list) else scales + return {target_key: self._dequantize_one(quantized, scales, output_dtype=output_dtype)} + return {target_key: quantized} + + # Generic chain path: dequantize every weight pattern that has a sibling scale. + result: dict[str, list[torch.Tensor] | torch.Tensor] = {} + for key, value in input_dict.items(): + if "activation_scale" in key or "weight_scale_inv" in key: + continue # consumed by the dequant; drop from the chain + scale_key = self._scale_pattern_for(key) + if scale_key not in input_dict: + # No scale to apply (e.g. unrelated entry) — pass through untouched. + result[key] = value + continue + weights = value if isinstance(value, list) else [value] + scales = input_dict[scale_key] + scales = scales if isinstance(scales, list) else [scales] + if len(weights) != len(scales): + raise ValueError( + f"FineGrainedDequantize: weight/scale count mismatch for {key} " + f"({len(weights)} weights vs {len(scales)} scales)." + ) + result[key] = [self._dequantize_one(w, s, output_dtype=output_dtype) for w, s in zip(weights, scales)] + return result + + @property + def reverse_op(self) -> ConversionOps: + # Round-trip: dequantize on load -> re-quantize on save, so the saved + # checkpoint preserves the FP8 format (weight + per-block ``weight_scale_inv``) + # whether the in-memory state stayed quantized or was dequantized for compute. + return FineGrainedQuantize(self.hf_quantizer) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 8f82c9b15196..a4b6c5fc2e1a 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -15,6 +15,7 @@ import functools import os +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -40,6 +41,15 @@ from .tensor_parallel import to_local +warnings.warn( + "finegrained_fp8 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 030cb44969ed..152e18e72593 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -577,9 +577,12 @@ def register_kernel_mapping_transformers(*args, **kwargs): _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { - "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2}, "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, + "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, + # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); + # exports the same three matmuls plus the fused-MoE forwards and per-recipe quant helpers + "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 1}, } _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {} diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 44fdec436bb1..92450042d773 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -72,7 +72,7 @@ ) from .integrations.deepspeed import _load_state_dict_into_zero3_model from .integrations.eager_paged import eager_paged_attention_forward -from .integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS +from .integrations.finegrained import ALL_FINEGRAINED_EXPERTS_FUNCTIONS from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward @@ -1964,7 +1964,9 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in def get_correct_experts_implementation(self, requested_experts: str | None) -> str: applicable_experts = "grouped_mm" if requested_experts is None else requested_experts - base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys())) + base_experts_fns = ["eager"] + list( + set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FINEGRAINED_EXPERTS_FUNCTIONS.keys()) + ) valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns] valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1] valid_experts_str = ", ".join(valid_experts_str_list) diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 3e9c1e1257f3..1ca9f80d0c78 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -25,6 +25,7 @@ CompressedTensorsConfig, EetqConfig, FbgemmFp8Config, + FineGrainedConfig, FineGrainedFP8Config, FourOverSixConfig, FPQuantConfig, @@ -53,7 +54,7 @@ from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer from .quantizer_eetq import EetqHfQuantizer from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer -from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer +from .quantizer_finegrained import FineGrainedHfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer from .quantizer_gemma import GemmaQuantizer @@ -61,7 +62,6 @@ from .quantizer_higgs import HiggsHfQuantizer from .quantizer_hqq import HqqHfQuantizer from .quantizer_metal import MetalHfQuantizer -from .quantizer_mxfp4 import Mxfp4HfQuantizer from .quantizer_quanto import QuantoHfQuantizer from .quantizer_quark import QuarkHfQuantizer from .quantizer_sinq import SinqHfQuantizer @@ -89,13 +89,17 @@ "bitnet": BitNetHfQuantizer, "vptq": VptqHfQuantizer, "spqr": SpQRHfQuantizer, - "fp8": FineGrainedFP8HfQuantizer, + "fp8": FineGrainedHfQuantizer, + # NVFP4 (two-level E2M1: E4M3 group-16 block scales x fp32 per-tensor global) — served by the + # same finegrained quantizer; the format is resolved off the checkpoint tensors. + "nvfp4": FineGrainedHfQuantizer, + "modelopt": FineGrainedHfQuantizer, # MXFP8 = FP8 (E4M3 weights) with per-block ``[1, 32]`` E8M0 (uint8) scales — # reuses the FineGrainedFP8 dequant path, with the E8M0 byte→exponent # unpacking handled inside ``Fp8Dequantize._dequantize_one``. - "mxfp8": FineGrainedFP8HfQuantizer, + "mxfp8": FineGrainedHfQuantizer, "auto-round": AutoRoundQuantizer, - "mxfp4": Mxfp4HfQuantizer, + "mxfp4": FineGrainedHfQuantizer, "metal": MetalHfQuantizer, "sinq": SinqHfQuantizer, "gemma": GemmaQuantizer, @@ -120,10 +124,12 @@ "bitnet": BitNetQuantConfig, "vptq": VptqConfig, "spqr": SpQRConfig, - "fp8": FineGrainedFP8Config, - "mxfp8": FineGrainedFP8Config, + "fp8": FineGrainedConfig, + "nvfp4": FineGrainedConfig, + "modelopt": FineGrainedConfig, + "mxfp8": FineGrainedConfig, "auto-round": AutoRoundConfig, - "mxfp4": Mxfp4Config, + "mxfp4": FineGrainedConfig, "metal": MetalConfig, "sinq": SinqConfig, "gemma": GemmaQuantizationConfig, diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py new file mode 100644 index 000000000000..aabbd53bdedb --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -0,0 +1,394 @@ +from typing import TYPE_CHECKING + +from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..utils.quantization_config import FineGrainedFP8Config + +logger = logging.get_logger(__name__) + + +class FineGrainedHfQuantizer(HfQuantizer): + """ + FP8 quantization implementation supporting both standard and MoE models. + Supports both e4m3fn formats based on platform. + """ + + requires_calibration = False + quantization_config: "FineGrainedFP8Config" + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") + + if self.quantization_config.dequantize: + return + + if not torch.cuda.is_available() and not is_torch_xpu_available(): + if self.pre_quantized: + logger.warning_once( + "Using FP8 quantized models requires a GPU or XPU, we will default to dequantizing the model to bf16 since no GPU or XPU is available" + ) + self.quantization_config.dequantize = True + return + else: + raise RuntimeError("No GPU or XPU found. A GPU or XPU is needed for FP8 quantization.") + + if torch.cuda.is_available(): + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + if (major < 8) or (major == 8 and minor < 9): + logger.warning_once( + "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)" + f", actual = `{major}.{minor}`. We will default to dequantizing the model to bf16. Feel free " + f"to use a different quantization method like bitsandbytes or torchao" + ) + self.quantization_config.dequantize = True + return + + device_map = kwargs.get("device_map") + if device_map is None: + logger.warning_once( + "You have loaded an FP8 model on CPU and have a CUDA or XPU device available, make sure to set " + "your model on a GPU or XPU device in order to run your model. To remove this warning, " + "pass device_map = 'cuda' or 'xpu'. " + ) + elif isinstance(device_map, dict): + if ( + not self.pre_quantized + and len(device_map) > 1 + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + raise ValueError( + "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device." + "This is not supported when the model is quantized on the fly. " + "Please use a quantized checkpoint or remove the cpu/disk device from the device_map." + ) + + def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: + # `FP8GroupedLinear` is a subclass of `FineGrainedLinear`, so the tuple covers it implicitly. + from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear + + module, tensor_name = get_module_from_name(model, param_name) + if isinstance(module, (FineGrainedLinear, FineGrainedExperts)): + if self.pre_quantized or tensor_name == "bias": + return False + else: + return True + return False + + def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: + "Return the element size (in bytes) for `param_name`." + if self.param_needs_quantization(model, param_name): + # 8 bit, this is needed as when `pre_quantized`` is False, we don't set the dtype of the FineGrainedLinear in order to correctly load the weights + return 1 + return super().param_element_size(model, param_name, param) + + def _normalize_modules_to_not_convert(self, model: "PreTrainedModel"): + """Rewrite the skip-list to the model's own module tree. + For models that were already released, if they have a list of modules to not quantize + we need to apply the weight renaming / weight conversion opérations to get the actual + layer name of the model in `transformers`. + """ + skip = self.quantization_config.modules_to_not_convert + if not skip: + return + + from ..conversion_mapping import get_model_conversion_mapping + + renamings = get_model_conversion_mapping(model) + remapped = [] + for name in skip: + renamed = name + for rename in renamings: + renamed, _ = rename.rename_source_key(renamed) + remapped.append(renamed) + self.quantization_config.modules_to_not_convert = remapped + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + **kwargs, + ): + from ..integrations.finegrained import replace_with_finegrained_layer + + self._normalize_modules_to_not_convert(model) + if self._quant_method() == "mxfp4" and getattr(self.quantization_config, "activation_format", None) is None: + # GPT-OSS MXFP4 runs weight-only (W4A16): raw bf16 activations against packed + # weights. The kernels' weight-native default would quantize activations to mxfp4. + self.quantization_config.activation_format = "bf16" + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules + ) + + model = replace_with_finegrained_layer( + model, + modules_to_not_convert=self.modules_to_not_convert, + quantization_config=self.quantization_config, + pre_quantized=self.pre_quantized, + ) + + def _process_model_after_weight_loading(self, model, **kwargs): + # dsv4-flash-base stores its (power-of-two) ue8m0 scales in a float32 container under + # `.scale`; those renamed keys keep the on-disk float32 dtype, so cast them to the UE8M0 + # dtype the kernels expect (exact, since the values are powers of two). Checkpoints that + # already ship the native float8 E8M0 dtype (e.g. dsv4-flash) are left untouched. + if self.quantization_config.scale_fmt == "ue8m0": + from ..integrations.finegrained import _get_ue8m0_dtype + + ue8m0 = _get_ue8m0_dtype() + float32_scales = [ + name + for name, param in model.named_parameters() + if name.endswith("_scale_inv") and param.dtype == torch.float32 + ] + for name in float32_scales: + module_name, _, attr = name.rpartition(".") + module = model.get_submodule(module_name) + scale = getattr(module, attr) + setattr(module, attr, torch.nn.Parameter(scale.data.to(ue8m0), requires_grad=False)) + + # Single-process multi-device is unsafe for DeepGEMM (its kernels are bound to one CUDA + # context); route those models through Triton/grouped_mm instead. + from ..integrations.finegrained import _disable_deepgemm_on_multi_device + + _disable_deepgemm_on_multi_device(model) + + from ..integrations.finegrained import swizzle_scales_after_loading + + swizzle_scales_after_loading(model) + return model + + def update_tp_plan(self, config): + if "Qwen3" in config.__class__.__name__: + text_plan = { + "layers.*.self_attn.q_proj.weight": "colwise", + "layers.*.self_attn.q_proj.weight_scale_inv": "colwise", + "layers.*.self_attn.k_proj.weight": "colwise", + "layers.*.self_attn.k_proj.weight_scale_inv": "colwise", + "layers.*.self_attn.v_proj.weight": "colwise", + "layers.*.self_attn.v_proj.weight_scale_inv": "colwise", + "layers.*.self_attn.o_proj.weight": "rowwise", + "layers.*.self_attn.o_proj.weight_scale_inv": "rowwise", + "layers.*.mlp.gate_proj.weight": "colwise", + "layers.*.mlp.gate_proj.weight_scale_inv": "colwise", + "layers.*.mlp.up_proj.weight": "colwise", + "layers.*.mlp.up_proj.weight_scale_inv": "colwise", + "layers.*.mlp.down_proj.weight": "rowwise", + "layers.*.mlp.down_proj.weight_scale_inv": "rowwise", + } + + config.base_model_tp_plan = text_plan + + # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes + # on top of any plan written above (e.g. the Qwen3 dense plan). Models carry the + # experts mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite + # both. See `FineGrainedExperts._impl_tp_layer_overrides`. + from ..integrations.finegrained import FineGrainedExperts + + impl = getattr(config, "_experts_implementation", None) + layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl) + if layer_overrides: + for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): + base_plan = getattr(config, plan_attr, None) or {} + updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} + if updated_plan != base_plan: + setattr(config, plan_attr, updated_plan) + + return config + + def is_serializable(self): + return True + + @property + def is_trainable(self) -> bool: + return False + + @property + def is_compileable(self) -> bool: + return True + + def get_quantize_ops(self): + from ..integrations.finegrained import FineGrainedQuantize + + return FineGrainedQuantize(self) + + def _quant_method(self) -> str: + method = getattr(self.quantization_config, "quant_method", None) + return getattr(method, "value", method) + + def get_weight_conversions(self): + from ..core_model_loading import WeightConverter + from ..integrations.finegrained import ( + FineGrainedDequantize, + FineGrainedGateUpBiasDeinterleave, + FineGrainedMxfp4Deserialize, + ) + + # GPT-OSS-style MXFP4 checkpoints ship {proj}_blocks + {proj}_scales (interleaved + # gate|up rows); deserialize them into the packed weight + e8m0 scale pair the + # finegrained kernels read, de-interleaving gate_up rows (and the bias to match). + if self.pre_quantized and self._quant_method() == "mxfp4" and not self.quantization_config.dequantize: + return [ + WeightConverter( + source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], + target_patterns=r"gate_up_proj$", + operations=[FineGrainedMxfp4Deserialize(self)], + ), + WeightConverter( + source_patterns=["down_proj_blocks", "down_proj_scales"], + target_patterns=r"down_proj$", + operations=[FineGrainedMxfp4Deserialize(self)], + ), + WeightConverter( + source_patterns=[r"gate_up_proj_bias$"], + target_patterns=r"gate_up_proj_bias$", + operations=[FineGrainedGateUpBiasDeinterleave(self)], + ), + ] + + # NVIDIA modelopt NVFP4 checkpoints: per-expert split gate/up projections with + # packed-uint8 weights, E4M3 group-16 `weight_scale`, and fp32 `weight_scale_2` + # second-level globals (bit-identical across the gate/up halves — asserted by the + # fuse op). The arch conversion plan already stacks the expert WEIGHTS; these add + # the same stacking for scales and globals, plus the uint8 -> int8 packed bitcast. + # Calibrated `input_scale` entries are left unconsumed (activations run the + # kernels' dynamic quant); they surface as unexpected keys, which is intended. + if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: + from ..core_model_loading import Concatenate, MergeModulelist + from ..integrations.finegrained import FineGrainedFuseEqualGlobals + + return [ + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.weight_scale$", + "mlp.experts.*.up_proj.weight_scale$", + ], + target_patterns="mlp.experts.gate_up_proj_scale_inv", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight_scale$", + target_patterns="mlp.experts.down_proj_scale_inv", + operations=[MergeModulelist(dim=0)], + ), + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.weight_scale_2", + "mlp.experts.*.up_proj.weight_scale_2", + ], + target_patterns="mlp.experts.gate_up_proj_global_scale", + operations=[FineGrainedFuseEqualGlobals(self)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight_scale_2", + target_patterns="mlp.experts.down_proj_global_scale", + operations=[FineGrainedFuseEqualGlobals(self)], + ), + ] + # No converters for non-expert modules: this checkpoint family quantizes ONLY + # the routed experts (every attention / shared-expert / dense-MLP module sits + # in the ignore list and ships bf16). Generic weight_scale* renames here would + # run BEFORE converter collection and mangle the expert keys out from under + # the converters above (per-expert `weight_global_scale` orphans and + # `_scale_inv_inv` double-suffix targets — the saturation bug of round 8). + # Quantized-dense modelopt checkpoints need scoped per-module converters when + # one shows up. + + if self.pre_quantized and self.quantization_config.dequantize: + return [ + # either use the dollar sign, or permute the source patterns to start matching against the scales first + # We also collect the activation scales, they will not be used + WeightConverter( + source_patterns=["weight$", "weight_scale_inv", "activation_scale"], + target_patterns="weight", + operations=[FineGrainedDequantize(self)], + ) + ] + return [] + + def update_weight_conversions(self, weight_conversions): + """When loading with ``dequantize=True``, attach an :class:`FineGrainedDequantize` op to + every existing :class:`WeightConverter` so that per-block scales are folded into + the weight *before* any later merge/concat ops collapse the per-expert structure. + + For each model-supplied converter that has a ``.weight`` source, we: + 1. anchor the existing weight patterns with ``$`` so they don't accidentally + also match the ``.weight_scale_inv`` keys (the regex is searched, so the + unanchored prefix would match both, sending scales to the wrong bucket); + 2. add anchored ``*.weight_scale_inv`` sources next to each weight pattern so + the loader collects scale tensors alongside the weight tensors into the + *same* converter bucket (both keys rewrite to the same target); + 3. prepend a fresh :class:`FineGrainedDequantize` op so dequant runs first, before + any merge/concat collapses the per-expert structure. + + The generic ``weight$ + weight_scale_inv → weight`` converter from + :meth:`get_weight_conversions` is still appended at the end as a fallback for + plain ``nn.Linear`` weights with no model-specific converter. + """ + from ..core_model_loading import WeightConverter, WeightRenaming + from ..integrations.finegrained import FineGrainedDequantize + + # `*.scale` → `*.weight_scale_inv`. Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) + # ship per-block scales under `.scale`; the model expects `.weight_scale_inv`. + # Lives here (not in each model's `conversion_mapping`) so non-FP8 round-trips + # don't see a stray rule. Needed in both dequantize modes — `dequantize=False` + # loads scales as parameters, `dequantize=True` feeds them into `FineGrainedDequantize`. + scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") + weight_conversions = [scale_rename] + list(weight_conversions) + + if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: + # Anchor the arch plan's `.weight` sources so they don't also swallow the + # modelopt `weight_scale`/`weight_scale_2` keys into the weight merge chain + # (unanchored patterns are searched — same trap the dequantize path guards). + from ..integrations.finegrained import FineGrainedViewPackedInt8 + + updated = [] + for conv in weight_conversions: + if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): + conv = WeightConverter( + source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], + target_patterns=conv.target_patterns, + # the stacked packed-fp4 expert weights need the same uint8 -> int8 + # bitcast the standalone-linear converter applies (pass-through on + # any unquantized weights the plan also routes here) + operations=list(conv.operations) + [FineGrainedViewPackedInt8(self)], + ) + updated.append(conv) + return updated + self.get_weight_conversions() + if not (self.pre_quantized and self.quantization_config.dequantize): + return weight_conversions + self.get_weight_conversions() + + updated: list = [] + for conv in weight_conversions: + # Only WeightConverter has ``.operations`` to extend with the dequant op; + # WeightRenaming (e.g. the ``scale_rename`` we prepended) just passes through. + if not isinstance(conv, WeightConverter): + updated.append(conv) + continue + weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] + if weight_sources: + anchored_weight = [p + "$" for p in weight_sources] + scale_sources = [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] + other = [p for p in conv.source_patterns if not p.endswith(".weight")] + new_sources = anchored_weight + scale_sources + other + new_ops = [FineGrainedDequantize(self)] + list(conv.operations) + conv = WeightConverter( + source_patterns=new_sources, + target_patterns=conv._original_target_patterns, + operations=new_ops, + ) + updated.append(conv) + # Generic fallback for plain ``nn.Linear`` weights with no model-specific converter. + updated.extend(self.get_weight_conversions()) + return updated diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index f9f49b8ec64b..6ef4c2a9c75b 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -1,3 +1,4 @@ +import warnings from typing import TYPE_CHECKING from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging @@ -8,6 +9,16 @@ if is_torch_available(): import torch + +warnings.warn( + "quantizer_finegrained_fp8 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, + stacklevel=2, +) + + if TYPE_CHECKING: from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import FineGrainedFP8Config diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 9e9816122d3d..13554f4dc19a 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -18,6 +18,7 @@ import importlib.metadata import json import os +import re from dataclasses import dataclass from enum import Enum from typing import Any, Optional, Union @@ -1691,15 +1692,22 @@ def post_init(self): @dataclass -class FineGrainedFP8Config(QuantizationConfigMixin): +class FineGrainedConfig(QuantizationConfigMixin): """ - FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. + Configuration for the fine-grained quantization family served by the + `kernels-community/finegrained-kernels` package: block-FP8 (DeepSeek-style, fp32 or UE8M0 + scales), MXFP8, MXFP4 and NVFP4 weights, with dynamic, static or weight-matched activation + quantization. The weight FORMAT is never declared here — it is resolved from the checkpoint + tensors themselves (value dtype, scale dtype/shape, presence of a global scale), exactly the + way the kernels resolve it, so the config cannot disagree with the weights. Args: activation_scheme (`str`, *optional*, defaults to `"dynamic"`): - The scheme used for activation, the defaults and only support scheme for now is "dynamic". + The scheme used for activation quantization: "dynamic" (inline per-token/per-block) + or "static" (calibrated per-tensor scale stored in the checkpoint). weight_block_size (`typing.tuple[int, int]`, *optional*, defaults to `(128, 128)`): - The size of the weight blocks for quantization, default is (128, 128). + The size of the weight blocks for block-FP8 quantization, default is (128, 128). + Group-scaled formats (MX/NV) carry their granularity in the scale tensors instead. dequantize (`bool`, *optional*, defaults to `False`): Whether to dequantize the model during loading. modules_to_not_convert (`list`, *optional*): @@ -1707,6 +1715,10 @@ class FineGrainedFP8Config(QuantizationConfigMixin): scale_fmt (`str`, *optional*, defaults to `"float"`): Storage dtype of the per-block weight scales: `"float"` (fp32, V3-style) or `"ue8m0"` (1-byte `torch.float8_e8m0fnu`, V4-style). + activation_format (`str`, *optional*): + Activation quantization format, needed only where the weights leave it ambiguous + (MXFP4 weights run as W4A16 with `"bf16"`, W4A8 with `"mxfp8"`, W4A4 with + `"mxfp4"`). `None` (default) matches the kernels' weight-native choice. """ def __init__( @@ -1716,12 +1728,40 @@ def __init__( dequantize: bool = False, modules_to_not_convert: list | None = None, scale_fmt: str = "float", + activation_format: str | None = None, **kwargs, ): self.quant_method = kwargs.pop("quant_method", QuantizationMethod.FP8) + self.activation_format = activation_format # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") + # NVIDIA modelopt exports (quant_method "modelopt"): translate the payload — + # quant_algo names the format (NVFP4 is the one the finegrained kernels serve), + # ``ignore`` is a glob-style skip list, and the calibrated per-tensor + # ``input_scale``/``kv_cache_scheme`` entries are dropped for now (activations run + # the kernels' dynamic quant, KV cache stays in the compute dtype). + if str(self.quant_method) == "modelopt" or kwargs.get("quant_algo") is not None: + quant_algo = kwargs.pop("quant_algo", None) + kwargs.pop("config_groups", None) + kwargs.pop("kv_cache_scheme", None) + kwargs.pop("producer", None) + if quant_algo != "NVFP4": + raise ValueError(f"modelopt checkpoints are supported for quant_algo='NVFP4' only; got {quant_algo!r}") + self.activation_format = activation_format or "nvfp4" + ignore = kwargs.pop("ignore", None) + if modules_to_not_convert is None and ignore is not None: + # modelopt ships glob-style subtree entries ("model.layers.0*", + # "model.layers.1.*"); translate each to a BOUNDED regex — a naive + # "*" -> ".*" leaves the dots unescaped and "layers.1..*" swallows + # layers 10-19 (and ``should_convert_module``'s bare-prefix clause + # would even match "layers.16" from "layers.1") + + def _subtree_regex(glob): + prefix = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") + return re.escape(prefix) + r"(\..*)?$" + + modules_to_not_convert = [_subtree_regex(g) for g in ignore] self.modules_to_not_convert = modules_to_not_convert self.activation_scheme = activation_scheme self.weight_block_size = weight_block_size @@ -1742,11 +1782,20 @@ def post_init(self): raise ValueError("weight_block_size must be a tuple of two positive integers") if self.scale_fmt not in ("float", "ue8m0"): raise ValueError(f"scale_fmt must be 'float' or 'ue8m0'; got {self.scale_fmt!r}") + if self.activation_format not in (None, "bf16", "fp8", "mxfp8", "mxfp4", "nvfp4"): + raise ValueError( + "activation_format must be one of None, 'bf16', 'fp8', 'mxfp8', 'mxfp4', " + f"'nvfp4'; got {self.activation_format!r}" + ) def get_loading_attributes(self): return {"dequantize": self.dequantize, "modules_to_not_convert": self.modules_to_not_convert} +# Back-compat alias: serialized checkpoints and released code reference the FP8-era name. +FineGrainedFP8Config = FineGrainedConfig + + class QuarkConfig(QuantizationConfigMixin): def __init__( self, From 7b21bdb7f14176ac3b80a81a4e8b6dcd16a790a7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 15 Aug 2026 05:39:23 -0700 Subject: [PATCH 02/68] deepgemm and tp/ep fixes --- src/transformers/integrations/finegrained.py | 49 ++++++++++++++----- .../integrations/tensor_parallel.py | 16 ++++++ .../quantizers/quantizer_finegrained.py | 8 ++- tests/tensor_parallel/test_tensor_parallel.py | 47 ++++++++++++++++++ 4 files changed, 107 insertions(+), 13 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index e00fa94cf20a..327dfb543ec0 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -332,11 +332,12 @@ def finegrained_linear( """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales: - 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — 3-6× faster on the shapes it supports. - Preferred for FP4, UE8M0 SFs, and 128×128 block FP8. - 2. Triton finegrained-fp8 fallback — used when DeepGEMM is unavailable, when the - caller passes ``activation_scale`` (DeepGEMM is dynamic-only), or for any - shape DeepGEMM declined. + 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — on the pre-SM100 shapes it supports (FP4, + UE8M0 SFs, 128×128 block FP8). Never preferred on SM100: the Triton path is tuned + per shape there and is the only one that reads the pre-swizzled SWIZZLE_32_4_4 + scales, and DeepGEMM's context-bound kernels are the multi-device hazard below. + 2. Triton finegrained fallback — everywhere else: SM100, an ``activation_scale`` + (DeepGEMM is dynamic-only), or any shape DeepGEMM declined. Args: input: (..., K) bf16/fp16 activations. @@ -363,13 +364,23 @@ def finegrained_linear( # degrades end-to-end generation on B200 (per-row kernel outputs still measure bit-perfect, but final # tokens drift; not reproducible with the DeepGEMM linear off). deepgemm_preferred = ( - allow_deepgemm + # SM100 perf: Triton measures 2.9x (decode) / 1.45x (prefill) over DeepGEMM on the DSV4 + # qkv linear — the 128x128 block-FP8 shape DeepGEMM is otherwise preferred for, which the + # gate above never catches (block-FP8's (N/128, K/128) grid has no swizzled layout). + not is_sm100() + # If the model is on multiple devices, DeepGEMM's context-bound kernels corrupt across devices. The + # multi-device guard in `quantizer_finegrained.py` flips `_deepgemm_disabled` True at load, which + # this dispatcher sees and respects. The Triton fallback is context-free and safe. + and allow_deepgemm + # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major: DeepGEMM would + # consume the permuted buffer as affine and silently return garbage. Correctness gate, + # true on any arch, for whatever `swizzle_scales_after_loading` has already swizzled. + and weight_scale_inv.ndim <= 2 and is_deepgemm_loadable() and activation_scale is None # DeepGEMM serves neither the NVFP4 two-level global nor an explicit activation format and weight_global_scale is None and activation_format is None - and not (is_sm100() and weight_scale_inv.dtype == torch.float32) and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" ) @@ -992,21 +1003,37 @@ def swizzle_scales_after_loading(model: nn.Module) -> None: a BUFFER next to the affine scale Parameter (which stays the state_dict source of truth, so save/dequantize round-trip untouched; scale tensors are small). gate_up scales take the gate-interleaved artifact the fused GATE arm reads natively. - ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` skips the pass (debug / pre-Blackwell A/B).""" + + SM100 only. The layout exists for the Blackwell tcgen05 scaled-MMA; elsewhere it buys + nothing, costs the swizzled arm's tile constraints (``swizzled_scales_bm_pruner`` pins + BM=BN=128), and collides with DeepGEMM — which stays the preferred linear backend below + SM100 and reads plain row-major scales. + ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` skips the pass (debug / A/B).""" if os.environ.get("TRANSFORMERS_FINEGRAINED_NO_SWIZZLE", "0") == "1": return + if not is_sm100(): + return try: kernel = load_finegrained_kernel() except ImportError: return # dequantize fallback path: nothing to swizzle for + formats = _weight_formats() for module in model.modules(): if not isinstance(module, FineGrainedExperts): continue + # Only the GROUP-scaled families have a layout to pre-arrange: SWIZZLE_32_4_4 is the + # scale OPERAND layout the tcgen05 scaled-MMA reads (MX group-32, NVFP4 group-16). + # Block-FP8's (N/128, K/128) grid never reaches a scaled-MMA — the block-dynamic loop + # applies it on the fp32 accumulator around a plain dot — so it has no swizzled form. + # Keyed off the declared recipe, not the scale dtype: V4-style block-FP8 ships UE8M0 + # scales, so a `dtype == float32` test lets it through (it then survives only by + # accident, because N/128 is rarely a multiple of 128). + fmt = formats.get(module.weight_format) + if fmt is None or fmt.scale_group is None: + continue for proj in (("gate_up_proj" if module.has_gate else "up_proj"), "down_proj"): scale = getattr(module, f"{proj}_scale_inv", None) - # only the group-scaled families have a swizzled layout (E8M0 / E4M3 group scales); - # block-FP8's fp32 (N/128, K/128) grid has no tcgen05 layout to pre-arrange - if scale is None or scale.dtype == torch.float32 or scale.device.type == "meta": + if scale is None or scale.device.type == "meta": continue gate = proj == "gate_up_proj" n_rows = to_local(scale).shape[-2] diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 8eca2f441d4d..aa1cd993de97 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -355,6 +355,13 @@ def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int world_size = reduce(operator.mul, mesh_shape) # Get param shape: works for both torch.Tensor and safetensors TensorInfo param_shape = list(param.shape) if isinstance(param, torch.Tensor) else param.get_shape() + # A 0-dim parameter has no axis to shard — every rank needs the whole value. Guarded here + # rather than at the call sites (there are nine) so the invariant holds for all of them: + # below, `dim` normalizes to -1 and indexes `param_shape[-1]` on an empty list. Per-tensor + # scalar scales make this reachable in practice — a ModelOpt NVFP4 checkpoint ships one + # `weight_scale_2` / `input_scale` per projection, thousands of 0-dim tensors per model. + if not param_shape: + return param[...] if not isinstance(param, torch.Tensor) else param if dim < 0: dim = param_dim + dim if empty_param.dim() == 3 and dim == 1 and len(param_shape) == 2: @@ -1077,6 +1084,15 @@ def shard_tensor( end = (self.rank + 1) * shard_size # special case we don't "shard" just send this entire tensor to the correct rank. shape = param.get_shape() if not isinstance(param, torch.Tensor) else param.shape + # A 0-dim scalar has no expert axis to split on: ModelOpt NVFP4 checkpoints ship one + # `weight_scale_2` / `input_scale` per expert projection. Every path below indexes, and + # both `param[:]` and `param[start:end]` raise on a scalar (`[...]` is the one form a + # lazy safetensors slice accepts), so route scalars before reaching them — taking the + # value whole for this rank's experts and dropping the rest, exactly as for >=1-dim. + if len(shape) == 0: + if tensor_idx is None or start <= tensor_idx < end: + return param[...].to(device=device, dtype=dtype) + return None if tensor_idx is not None and start <= tensor_idx < end: # this tensor does need to be materialized on this device: return param[:].to(device=device) diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index aabbd53bdedb..c809146b9841 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -288,12 +288,16 @@ def get_weight_conversions(self): "mlp.experts.*.up_proj.weight_scale_2", ], target_patterns="mlp.experts.gate_up_proj_global_scale", - operations=[FineGrainedFuseEqualGlobals(self)], + # MergeModulelist is what stamps each source with its expert index, which + # is what lets expert parallelism keep only this rank's experts (see + # `tensor_idx` in core_model_loading). Without it every rank collects all + # E globals and the forward asserts on the per-expert count. + operations=[MergeModulelist(dim=0), FineGrainedFuseEqualGlobals(self)], ), WeightConverter( source_patterns="mlp.experts.*.down_proj.weight_scale_2", target_patterns="mlp.experts.down_proj_global_scale", - operations=[FineGrainedFuseEqualGlobals(self)], + operations=[MergeModulelist(dim=0), FineGrainedFuseEqualGlobals(self)], ), ] # No converters for non-expert modules: this checkpoint family quantizes ONLY diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 3d0e644f69cd..1c46aca7dcdc 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -27,6 +27,7 @@ RowwiseParallel, add_tensor_parallel_hooks_to_module, get_packed_weights, + get_tensor_shard, repack_weights, ) from transformers.testing_utils import TestCasePlus, is_tensor_parallel_test @@ -445,3 +446,49 @@ def test_packed_rowwise_shard_tensor(self): sharded = layer.shard_tensor(full_unpacked) expected_shape = (16, 16) # last dim is not packed, so just sharded self.assertEqual(sharded.shape, expected_shape) + + +@is_tensor_parallel_test +class TestZeroDimParameterSharding(TestCasePlus): + """A 0-dim parameter has no axis to shard, so every rank needs the whole value. + + This is reachable in practice: quantized checkpoints carry per-tensor scalar scales, and a + ModelOpt NVFP4 model ships one `weight_scale_2` / `input_scale` per projection — thousands of + 0-dim tensors. Both sharding paths used to index a shape that isn't there, and neither failure + is visible single-process: `get_tensor_shard` normalizes `dim` to -1 and reads + `param_shape[-1]` on an empty list, while `GroupedGemmParallel` slices `param[start:end]`. + """ + + class _Mesh: + shape = (2,) + + def size(self): + return 2 + + def test_get_tensor_shard_replicates_a_scalar(self): + """Root guard: nine call sites reach this, so it is fixed here rather than at each.""" + scalar = torch.tensor(3.5) + for rank in (0, 1): + out = get_tensor_shard(scalar, scalar, self._Mesh(), rank, -1) + assert out.shape == torch.Size([]), f"rank {rank} got {out.shape}, expected 0-dim" + assert out.item() == 3.5, "a replicated scalar must keep its value on every rank" + + def test_get_tensor_shard_still_shards_a_matrix(self): + """The guard must not swallow the normal path.""" + param = torch.randn(8, 4) + shard = get_tensor_shard(param, param, self._Mesh(), 0, 0) + assert shard.shape == (4, 4), f"expected a half-shard, got {shard.shape}" + + def test_grouped_gemm_replicates_a_scalar_for_local_experts(self): + """Expert parallelism: a scalar belonging to this rank's expert range comes back whole, + and one belonging to another rank is dropped — the same contract as for >=1-dim.""" + layer = GroupedGemmParallel() + layer.empty_param = torch.empty(4, 8, 8) # 4 experts over 2 ranks -> 2 each + layer.device_mesh = self._Mesh() + layer.rank = 0 + + mine = layer.shard_tensor(torch.tensor(1.25), tensor_idx=1, device="cpu", dtype=torch.float32) + assert mine is not None and mine.item() == 1.25 + + theirs = layer.shard_tensor(torch.tensor(1.25), tensor_idx=3, device="cpu", dtype=torch.float32) + assert theirs is None, "an expert owned by another rank must be dropped, not materialized" From 05bd0f3587030916f38647f95da2d7fb96dd20a2 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 16 Aug 2026 10:43:00 -0700 Subject: [PATCH 03/68] finegrained: adopt the interleaved gate|up layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernels now store gate|up row-interleaved (gate j at row 2j, up j at 2j+1) rather than as two N-apart halves, so the integration follows: - GPT-OSS already ships this order on disk, so its de-interleave converters go away entirely: _deinterleave_gate_up_rows, its use in FineGrainedMxfp4Deserialize, and the whole FineGrainedGateUpBiasDeinterleave op plus its registration. - Checkpoints that ship gate_up stacked are interleaved once post-load by interleave_gate_up_after_loading, ordered BEFORE swizzle_scales_after_loading (the swizzle packs whatever row order it finds) and skipped for mxfp4. This cannot live in conversion_mapping's MergeModulelist+Concatenate chain — that is shared with non-finegrained MoE models. - _apply_gate de-interleaves with a stride-2 split, mirroring the fused epilogue's split_gate_up; the two must agree or fused and unfused stop being comparable. - swizzle_scales_after_loading drops gate= and gates on the doubled extent (n_rows % 128, i.e. N % 64), so GPT-OSS N=2880 now pre-swizzles instead of falling back to an affine scale read. Adds tests/kernels/test_finegrained.py (24 tests), covering the post-load interleave, that mxfp4 is left untouched, dispatch/swizzle gating, and the MergeModulelist path. tests/kernels/test_finegrained.py: 24 passed. --- src/transformers/integrations/finegrained.py | 80 +-- .../quantizers/quantizer_finegrained.py | 19 +- tests/kernels/test_finegrained.py | 600 ++++++++++++++++++ 3 files changed, 651 insertions(+), 48 deletions(-) create mode 100644 tests/kernels/test_finegrained.py diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 327dfb543ec0..b8c3fb44d1f7 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -907,7 +907,9 @@ def __init__( self.down_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: - gate, up = gate_up.chunk(2, dim=-1) + # the GEMM's output columns alternate gate/up; mirrors the fused epilogue's + # `split_gate_up`, which the fused/unfused parity tests compare against + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] if self.swiglu_alpha is not None: # Clamped SwiGLU-OAI gate (same math as the model's non-quantized experts). gate = gate.clamp(max=self.swiglu_limit) @@ -995,14 +997,44 @@ class FineGrainedExpertsInterface(ExpertsInterface): ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() +def interleave_gate_up_after_loading(model: nn.Module, *, already_interleaved: bool = False) -> None: + """Put gate_up into the ``[g0, u0, g1, u1, ...]`` row order the kernels read. + + Checkpoints that ship the two projections separately arrive stacked ``[gate; up]`` (the shared + ``MergeModulelist`` + ``Concatenate`` chain in ``conversion_mapping``, which serves unquantized + MoE models too and so cannot carry a finegrained-specific reorder). GPT-OSS-style MXFP4 already + ships interleaved — ``already_interleaved`` skips those. + + Weight, block-scale grid and bias all key off the same output-row axis, so all three move + together. Runs BEFORE ``swizzle_scales_after_loading``, which then swizzles the final order. + """ + if already_interleaved: + return + for module in model.modules(): + if not isinstance(module, FineGrainedExperts) or not module.has_gate: + continue + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias"): + t = getattr(module, name, None) + if t is None or t.device.type == "meta": + continue + local = to_local(t) + axis = -1 if local.ndim == 1 or name.endswith("bias") else -2 + n = local.shape[axis] // 2 + if axis == -1: + rows = torch.stack([local[..., :n], local[..., n:]], dim=-1) + else: + rows = torch.stack([local[..., :n, :], local[..., n:, :]], dim=-2) + local.data.copy_(rows.reshape(local.shape)) + + def swizzle_scales_after_loading(model: nn.Module) -> None: """Pre-swizzle MX/NVFP4 weight block scales into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05 scaled-MMA reads directly (one triton launch per matrix, values unchanged) — the kernels accept both layouts, but plain row-major forces a per-tile gather that caps the scaled dot below the fp8/fp4 peak. Runs once post-load; the swizzled artifact is cached as a BUFFER next to the affine scale Parameter (which stays the state_dict source of truth, so - save/dequantize round-trip untouched; scale tensors are small). gate_up scales take the - gate-interleaved artifact the fused GATE arm reads natively. + save/dequantize round-trip untouched; scale tensors are small). A gate_up grid swizzles as one + matrix over its 2N rows. SM100 only. The layout exists for the Blackwell tcgen05 scaled-MMA; elsewhere it buys nothing, costs the swizzled arm's tile constraints (``swizzled_scales_bm_pruner`` pins @@ -1035,14 +1067,13 @@ def swizzle_scales_after_loading(model: nn.Module) -> None: scale = getattr(module, f"{proj}_scale_inv", None) if scale is None or scale.device.type == "meta": continue - gate = proj == "gate_up_proj" n_rows = to_local(scale).shape[-2] - # the swizzled layout needs whole 128-row blocks per projection (under gate the - # split sits at N, mid-block unless N % 128 == 0, e.g. GPT-OSS N=2880); the - # kernels read the affine layout directly, so just skip the pre-swizzle - if (n_rows // 2 if gate else n_rows) % 128: + # the swizzled layout needs whole 128-row blocks; for gate_up that is the 2N extent, + # i.e. N % 64 (GPT-OSS N=2880 qualifies). The kernels read the affine layout directly, + # so a non-conforming projection just skips the pre-swizzle. + if n_rows % 128: continue - swizzled = kernel.swizzle_mx_scales(to_local(scale).data, gate=gate) + swizzled = kernel.swizzle_mx_scales(to_local(scale).data) module.register_buffer(f"{proj}_scale_inv_swizzled", swizzled, persistent=False) @@ -1179,20 +1210,13 @@ def replace_with_finegrained_layer( return model -def _deinterleave_gate_up_rows(t: torch.Tensor) -> torch.Tensor: - """GPT-OSS fuses gate|up INTERLEAVED along the output dim ([g0, u0, g1, u1, ...]); the - finegrained kernels (and `FineGrainedExperts._apply_gate`'s ``chunk``) read the STACKED - form ([gate; up]). One permutation at load, values untouched.""" - return torch.cat([t[:, 0::2], t[:, 1::2]], dim=1) - - class FineGrainedMxfp4Deserialize(ConversionOps): """Convert GPT-OSS MXFP4 checkpoint tensors to the finegrained layout: ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 (two low-nibble-first E2M1 values per byte) reshapes to the packed ``(E, N, K/2)`` int8 the kernels read directly — same nibble order — and ``{proj}_scales`` (biased-127 exponent bytes) bitcast to ``float8_e8m0fnu``. No swizzle - (the kernels take affine scales; the post-load hook builds the swizzled cache). gate_up - rows de-interleave from GPT-OSS's [g0, u0, ...] to the stacked [gate; up] form.""" + (the kernels take affine scales; the post-load hook builds the swizzled cache). GPT-OSS ships + gate_up in the interleaved [g0, u0, ...] row order the kernels read, so rows pass through.""" def __init__(self, hf_quantizer): self.hf_quantizer = hf_quantizer @@ -1213,31 +1237,11 @@ def convert( num_experts, rows = blocks.shape[0], blocks.shape[1] weight = blocks.reshape(num_experts, rows, -1).view(torch.int8) - # deinterleave BEFORE the e8m0 bitcast: cat/index ops aren't implemented for - # float8_e8m0fnu on CUDA, and the reorder is byte-identical either way - if proj == "gate_up_proj": - weight = _deinterleave_gate_up_rows(weight) - scales = _deinterleave_gate_up_rows(scales) scale_inv = scales.contiguous().view(_get_ue8m0_dtype()) prefix = full_layer_name.rsplit(".", 1)[0] return {f"{prefix}.{proj}": weight.contiguous(), f"{prefix}.{proj}_scale_inv": scale_inv} -class FineGrainedGateUpBiasDeinterleave(ConversionOps): - """GPT-OSS ``gate_up_proj_bias`` follows the interleaved output layout; de-interleave to - match the stacked weight rows produced by `FineGrainedMxfp4Deserialize`.""" - - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer - - def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): - # single-source converter: input_dict is keyed by the SOURCE PATTERN string (with - # any regex anchors), so don't look up by a hardcoded name - bias = next(iter(input_dict.values())) - bias = bias[0] if isinstance(bias, list) else bias - return {full_layer_name: _deinterleave_gate_up_rows(bias)} - - class FineGrainedViewPackedInt8(ConversionOps): """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store (same bytes — ``copy_`` into an int8 param would numerically CONVERT and corrupt diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index c809146b9841..84e73f35927c 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -164,8 +164,13 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) - from ..integrations.finegrained import swizzle_scales_after_loading + from ..integrations.finegrained import ( + interleave_gate_up_after_loading, + swizzle_scales_after_loading, + ) + # order matters: the swizzle below packs whatever row order gate_up ends up in + interleave_gate_up_after_loading(model, already_interleaved=self._quant_method() == "mxfp4") swizzle_scales_after_loading(model) return model @@ -231,13 +236,12 @@ def get_weight_conversions(self): from ..core_model_loading import WeightConverter from ..integrations.finegrained import ( FineGrainedDequantize, - FineGrainedGateUpBiasDeinterleave, FineGrainedMxfp4Deserialize, ) - # GPT-OSS-style MXFP4 checkpoints ship {proj}_blocks + {proj}_scales (interleaved - # gate|up rows); deserialize them into the packed weight + e8m0 scale pair the - # finegrained kernels read, de-interleaving gate_up rows (and the bias to match). + # GPT-OSS-style MXFP4 checkpoints ship {proj}_blocks + {proj}_scales; deserialize them + # into the packed weight + e8m0 scale pair the finegrained kernels read. The gate_up bias + # needs no converter — it already follows the interleaved row order. if self.pre_quantized and self._quant_method() == "mxfp4" and not self.quantization_config.dequantize: return [ WeightConverter( @@ -250,11 +254,6 @@ def get_weight_conversions(self): target_patterns=r"down_proj$", operations=[FineGrainedMxfp4Deserialize(self)], ), - WeightConverter( - source_patterns=[r"gate_up_proj_bias$"], - target_patterns=r"gate_up_proj_bias$", - operations=[FineGrainedGateUpBiasDeinterleave(self)], - ), ] # NVIDIA modelopt NVFP4 checkpoints: per-expert split gate/up projections with diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py new file mode 100644 index 000000000000..8b7631e8a616 --- /dev/null +++ b/tests/kernels/test_finegrained.py @@ -0,0 +1,600 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Marshalling tests for the multi-recipe ``finegrained`` integration: the kernel bundle is +mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` +(the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" + +import unittest +from dataclasses import dataclass +from unittest import mock + +import torch + +import transformers.integrations.finegrained as fg +from transformers.integrations.finegrained import ( + FineGrainedExperts, + FineGrainedLinear, + finegrained_linear, + load_finegrained_kernel, +) +from transformers.testing_utils import require_torch, require_torch_gpu + + +@dataclass +class _Call: + args: tuple + kwargs: dict + + +class _Recorder: + """Records every call and returns a correctly-shaped zero tensor for matmuls.""" + + def __init__(self): + self.calls = {} + + def _op(self, name): + def op(*args, **kwargs): + self.calls.setdefault(name, []).append(_Call(args, kwargs)) + if name == "swizzle_mx_scales": + return args[0] + a, b = args[0], args[1] + n = b.shape[-2] + gather = kwargs.get("gather_idx") + rows = gather.shape[0] if gather is not None else a.shape[0] + return a.new_zeros(rows, n, dtype=torch.bfloat16) + + return op + + +def _fake_bundle(): + rec = _Recorder() + kernel = mock.Mock() + for name in ("matmul_2d", "matmul_batched", "matmul_grouped", "nvfp4_quantize_two_level", "swizzle_mx_scales"): + setattr(kernel, name, rec._op(name)) + + def compute_grouped_scheduling(top_k_index, num_experts, num_top_k): + rec.calls.setdefault("compute_grouped_scheduling", []).append(_Call((top_k_index, num_experts, num_top_k), {})) + flat = top_k_index.reshape(-1) + sorted_ids, perm = torch.sort(flat) + counts = torch.bincount(sorted_ids.clamp(max=num_experts - 1), minlength=num_experts) + expert_start = torch.cat([counts.new_zeros(1), counts.cumsum(0)]).to(torch.int32) + gather_idx = (perm // num_top_k).to(torch.int32) + scatter_idx = torch.empty_like(perm) + scatter_idx[torch.arange(perm.numel())] = perm + return expert_start, gather_idx, scatter_idx.to(torch.int32) + + def weighted_reduce(rows, top_k_index, top_k_weights, num_experts, simulate_unfused=False): + rec.calls.setdefault("weighted_reduce", []).append(_Call((rows, top_k_index, top_k_weights, num_experts), {})) + num_tokens, num_top_k = top_k_index.shape + w = top_k_weights.reshape(-1, 1).to(rows.dtype) + w = w * (top_k_index.reshape(-1, 1) < num_experts) + return (rows * w).view(num_tokens, num_top_k, -1).sum(1) + + kernel.compute_grouped_scheduling = compute_grouped_scheduling + kernel.weighted_reduce = weighted_reduce + + @dataclass + class Quantization: + input_recipe: str | None = "weights" + output_recipe: str | None = None + + @dataclass + class Epilogue: + gate: bool = False + act_fn: str = "silu" + swiglu_alpha: float | None = None + swiglu_limit: float | None = None + + kernel.Quantization = Quantization + kernel.Epilogue = Epilogue + return kernel, rec + + +class _Cfg: + hidden_size = 64 + num_local_experts = 4 + intermediate_size = 32 + hidden_act = "silu" + + +def _loaded(kernel): + return ( + mock.patch.object(fg, "_FINEGRAINED", None), + mock.patch.object(fg, "is_kernels_available", return_value=True), + mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), + # a locally importable checkout (FINEGRAINED_KERNELS_PATH / installed package) + # takes precedence over the hub loader; tests must stay hermetic to the fake bundle + mock.patch.object(fg, "_import_local_finegrained", return_value=None), + ) + + +@require_torch +class FineGrainedLoaderTest(unittest.TestCase): + def test_loader_requires_every_symbol(self): + kernel, _ = _fake_bundle() + del kernel.matmul_grouped + # Mock auto-creates attributes; force the miss + kernel.matmul_grouped = None + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, self.assertRaises(ImportError) as ctx: + load_finegrained_kernel() + self.assertIn("matmul_grouped", str(ctx.exception)) + + def test_loader_binds_all_symbols(self): + kernel, _ = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + bundle = load_finegrained_kernel() + self.assertIs(bundle.Quantization, kernel.Quantization) + self.assertIs(bundle.Epilogue, kernel.Epilogue) + + +@require_torch +class FineGrainedLinearMarshallingTest(unittest.TestCase): + def _run(self, **linear_kwargs): + kernel, rec = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + x = torch.randn(3, 5, 64, dtype=torch.bfloat16) + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + ws = torch.randn(1, 1, dtype=torch.float32) + out = finegrained_linear(x, w, ws, **linear_kwargs) + return out, rec.calls["matmul_2d"][-1] + + def test_scale_rides_positionally_and_block_size_dies(self): + out, call = self._run(block_size=[128, 128]) + # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here + self.assertEqual(call.args[0].shape, (15, 64)) + self.assertIsNone(call.args[2]) + self.assertEqual(call.args[3].dtype, torch.float32) + self.assertNotIn("block_size", call.kwargs) + self.assertEqual(out.shape, (3, 5, 32)) + + def test_static_activation_scale_is_As(self): + scale = torch.tensor(0.5) + _, call = self._run(activation_scale=scale) + self.assertIs(call.args[2], scale) + + def test_nvfp4_global_and_weight_only_format(self): + g = torch.tensor(2.0) + _, call = self._run(weight_global_scale=g, activation_format="bf16") + self.assertIs(call.kwargs["b_global_scale"], g) + self.assertIsNone(call.kwargs["quantization"].input_recipe) + + def test_module_forward_threads_everything(self): + kernel, rec = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + m = FineGrainedLinear( + 64, + 32, + block_size=None, + weight_format="nvfp4", + activation_format="bf16", + ) + m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) + out = m(torch.randn(2, 64, dtype=torch.bfloat16)) + call = rec.calls["matmul_2d"][-1] + self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) + self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertEqual(out.shape, (2, 32)) + + +@require_torch +class FineGrainedExpertsMarshallingTest(unittest.TestCase): + def _experts(self, **kw): + cfg = _Cfg() + weight_format = kw.pop("weight_format", "fp8") + m = FineGrainedExperts(cfg, block_size=(4, 4), weight_format=weight_format, **kw) + for name, p in m.named_parameters(): + if p is None: + continue + if p.dtype == torch.float8_e4m3fn: + p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) + elif p.dtype.is_floating_point: + p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) + return m + + def _route(self, tokens=6): + hs = torch.randn(tokens, 64, dtype=torch.bfloat16) + idx = torch.randint(0, 4, (tokens, 2)) + wts = torch.rand(tokens, 2) + return hs, idx, wts + + def test_batched_marshalling(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + up, down = rec.calls["matmul_batched"] + self.assertIsNone(up.args[2]) # As positional None + self.assertEqual(up.args[0].shape[0], 6) # UNEXPANDED tokens — the kernel gathers + self.assertEqual(up.kwargs["gather_idx"].shape, (12,)) + self.assertEqual(up.kwargs["expert_ids"].shape, (12,)) + self.assertNotIn("block_size", up.kwargs) + self.assertIsNone(up.kwargs["b_global_scale"]) + self.assertEqual(out.shape, (6, 64)) + + def test_grouped_marshalling_expert_start_boundaries(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + up = rec.calls["matmul_grouped"][0] + es = up.kwargs["expert_start"] + self.assertEqual(es.shape, (5,)) # (E+1,) boundaries + self.assertEqual(es[0].item(), 0) + self.assertEqual(es[-1].item(), 12) # S = tokens * top_k + + def test_nvfp4_experts_thread_per_expert_globals(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, weight_format="nvfp4") + # the format table must resolve the ATTRIBUTE the forwards gate on, not just the + # param allocation — a None here silently drops the global at every forward + self.assertIs(m.has_global_scale, True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + up, down = rec.calls["matmul_batched"] + self.assertIs(up.kwargs["b_global_scale"], m.gate_up_proj_global_scale) + self.assertIs(down.kwargs["b_global_scale"], m.down_proj_global_scale) + + def test_bias_lands_before_gate_and_before_routing(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, has_bias=True) + hs, idx, wts = self._route() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = fg.finegrained_batched_mm_experts_forward(m, hs, idx, wts) + # zero matmul outputs + bias: result = sum_k w_k * (gate(bias_up) @ ... ) — just assert + # the bias params were consumed (non-None) and output is finite + self.assertTrue(torch.isfinite(out).all()) + self.assertEqual(len(rec.calls["matmul_batched"]), 2) + + +@require_torch +class FrozenFp8ShimTest(unittest.TestCase): + def test_frozen_module_warns_and_is_self_contained(self): + import importlib + import warnings + + import transformers.integrations.finegrained_fp8 as frozen + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + importlib.reload(frozen) + self.assertTrue(any(issubclass(x.category, DeprecationWarning) for x in w)) + # distinct machinery: the frozen classes are not the live ones + self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) + + +@require_torch +class FineGrainedMxfp4ConverterTest(unittest.TestCase): + """The gpt-oss checkpoint layout through the real ConversionOps: blocks (E, N, K/32, 16) + uint8 low-nibble-first E2M1 + biased-127 exponent scales, gate|up rows INTERLEAVED — + deserialized output must dequantize identically to the reference LUT unpack.""" + + FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + + def _reference_dequant(self, blocks, scales): + # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) + lut = torch.tensor(self.FP4_VALUES) + lo = lut[(blocks & 0xF).long()] + hi = lut[(blocks >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) + exp = (scales.long() - 127).unsqueeze(-1) + return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) + + def test_deserialize_matches_reference(self): + from transformers.integrations.finegrained import FineGrainedMxfp4Deserialize, _get_ue8m0_dtype + + torch.manual_seed(0) + E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows + blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) + scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) + + op = FineGrainedMxfp4Deserialize(hf_quantizer=None) + out = op.convert( + {"gate_up_proj_blocks": blocks, "gate_up_proj_scales": scales}, + full_layer_name="model.layers.0.mlp.experts.gate_up_proj", + ) + weight = out["model.layers.0.mlp.experts.gate_up_proj"] + scale_inv = out["model.layers.0.mlp.experts.gate_up_proj_scale_inv"] + self.assertEqual(weight.dtype, torch.int8) + self.assertEqual(weight.shape, (E, N, K // 2)) + self.assertEqual(scale_inv.dtype, _get_ue8m0_dtype()) + + ref = self._reference_dequant(blocks, scales) # rows pass through interleaved + + # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) + lut = torch.tensor(self.FP4_VALUES) + w_u8 = weight.view(torch.uint8) + lo = lut[(w_u8 & 0xF).long()] + hi = lut[(w_u8 >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) + exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) + got = vals * torch.pow(torch.tensor(2.0), exp) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + def test_interleave_gate_up_after_loading(self): + """Stacked-checkpoint families get gate_up interleaved post-load; GPT-OSS-style MXFP4 + already ships that order and is skipped.""" + from transformers.integrations.finegrained import ( + FineGrainedExperts, + interleave_gate_up_after_loading, + ) + + module = FineGrainedExperts.__new__(FineGrainedExperts) + torch.nn.Module.__init__(module) + module.has_gate = True + # rows [g0,g1,g2,u0,u1,u2]; bias follows the same output axis + module.gate_up_proj = torch.nn.Parameter( + torch.arange(24, dtype=torch.float32).reshape(1, 6, 4), requires_grad=False + ) + module.gate_up_proj_bias = torch.nn.Parameter( + torch.arange(6, dtype=torch.float32), requires_grad=False + ) + stacked_w = module.gate_up_proj.detach().clone() + stacked_b = module.gate_up_proj_bias.detach().clone() + + model = torch.nn.Module() + model.experts = module + interleave_gate_up_after_loading(model) + + for j in range(3): + torch.testing.assert_close(module.gate_up_proj[0, 2 * j], stacked_w[0, j]) + torch.testing.assert_close(module.gate_up_proj[0, 2 * j + 1], stacked_w[0, 3 + j]) + torch.testing.assert_close(module.gate_up_proj_bias[2 * j], stacked_b[j]) + torch.testing.assert_close(module.gate_up_proj_bias[2 * j + 1], stacked_b[3 + j]) + + def test_interleave_skipped_when_already_interleaved(self): + """GPT-OSS ships [g0,u0,...] already — the pass must leave it alone.""" + from transformers.integrations.finegrained import ( + FineGrainedExperts, + interleave_gate_up_after_loading, + ) + + module = FineGrainedExperts.__new__(FineGrainedExperts) + torch.nn.Module.__init__(module) + module.has_gate = True + module.gate_up_proj = torch.nn.Parameter( + torch.arange(24, dtype=torch.float32).reshape(1, 6, 4), requires_grad=False + ) + before = module.gate_up_proj.detach().clone() + + model = torch.nn.Module() + model.experts = module + interleave_gate_up_after_loading(model, already_interleaved=True) + torch.testing.assert_close(module.gate_up_proj.data, before) + + +@require_torch +class FineGrainedDeepGemmDispatchTest(unittest.TestCase): + """`deepgemm_preferred` carries two independent gates: a correctness one (a pre-swizzled + scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine + and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 + scales have no swizzled layout, so the first never fires for the shape the second catches.""" + + def _routed_to(self, *, sm100, scale_ndim): + """Which backend a block-FP8 linear reaches, given arch and scale layout.""" + kernel, _ = _fake_bundle() + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) + p1, p2, p3, p4 = _loaded(kernel) + with ( + p1, + p2, + p3, + p4, + mock.patch.object(fg, "is_deepgemm_loadable", return_value=True), + mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, + ): + fg.finegrained_linear(torch.randn(4, 64, dtype=torch.bfloat16), w, s, block_size=[128, 128]) + return "deepgemm" if dg.called else "triton" + + def test_sm100_never_prefers_deepgemm(self): + self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") + + def test_pre_sm100_still_uses_deepgemm(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") + + def test_swizzled_scales_never_reach_deepgemm(self): + """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + + +@require_torch +class FineGrainedSwizzlePassTest(unittest.TestCase): + """`swizzle_scales_after_loading` builds the Blackwell SWIZZLE_32_4_4 artifact. It is gated + on arch AND on the recipe having group scales — keyed off the declared weight format, not + the scale dtype, because V4-style block-FP8 ships UE8M0 scales and a dtype test lets it + through (it then survives only by accident, since N/128 is rarely a multiple of 128).""" + + def _experts(self, weight_format, scale_dtype): + cfg = _Cfg() + experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format=weight_format) + for proj in ("gate_up_proj", "down_proj"): + rows = 256 + setattr( + experts, + f"{proj}_scale_inv", + torch.nn.Parameter( + torch.zeros(cfg.num_local_experts, rows, 8, dtype=scale_dtype), + requires_grad=False, + ), + ) + model = torch.nn.Module() + model.experts = experts + return model, experts + + def _run_pass(self, model, *, sm100): + kernel, _ = _fake_bundle() + kernel.swizzle_mx_scales = mock.Mock(side_effect=lambda s, gate=False: s) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, mock.patch.object(fg, "is_sm100", return_value=sm100): + fg.swizzle_scales_after_loading(model) + return kernel.swizzle_mx_scales + + def test_skipped_off_sm100(self): + """The layout is a Blackwell tcgen05 artifact; elsewhere it buys nothing, pins the tile + space, and collides with DeepGEMM — which stays the preferred backend below SM100.""" + model, _ = self._experts("mxfp8", torch.float8_e8m0fnu) + self.assertEqual(self._run_pass(model, sm100=False).call_count, 0) + + def test_applied_for_group_scaled_recipes_on_sm100(self): + model, _ = self._experts("mxfp8", torch.float8_e8m0fnu) + self.assertGreater(self._run_pass(model, sm100=True).call_count, 0) + + def test_block_fp8_skipped_even_with_ue8m0_scales(self): + """The latent case: block-FP8's (N/128, K/128) grid never reaches a scaled-MMA, so it has + no swizzled form — but it ships UE8M0 scales, so only a recipe-keyed guard excludes it.""" + model, _ = self._experts("fp8", torch.float8_e8m0fnu) + self.assertEqual(self._run_pass(model, sm100=True).call_count, 0) + + +@require_torch +class FineGrainedModeloptConverterTest(unittest.TestCase): + """The ModelOpt (NVFP4) weight conversions. `MergeModulelist` is load-bearing beyond the + merge itself: `core_model_loading` stamps a source with its expert index ONLY when one is + present in the chain, and expert parallelism selects experts by that index. Without it every + rank collects all E globals and the forward asserts on the per-expert count.""" + + def _modelopt_conversions(self): + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4") + quantizer = FineGrainedHfQuantizer(cfg) + return quantizer.get_weight_conversions() + + def test_global_scale_converters_carry_an_expert_index(self): + from transformers.core_model_loading import MergeModulelist + + globals_converters = [c for c in self._modelopt_conversions() if any("global_scale" in t for t in _targets(c))] + self.assertTrue(globals_converters, "no global-scale converter found") + for conv in globals_converters: + self.assertTrue( + any(isinstance(op, MergeModulelist) for op in conv.operations), + f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " + "experts and every rank would collect all of them", + ) + + +def _targets(conv): + t = conv.target_patterns + return t if isinstance(t, (list, tuple)) else [t] + + +def _sm100(): + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 + + +@require_torch_gpu +@unittest.skipUnless(_sm100(), "the fine-grained kernels target SM100+") +class FineGrainedRealKernelTest(unittest.TestCase): + """Runs the REAL kernels. The mocked tests above pin the calling contract; these pin that the + contract produces correct numbers, which no amount of argument-matching can show.""" + + @classmethod + def setUpClass(cls): + try: + load_finegrained_kernel() + except ImportError as e: + raise unittest.SkipTest(f"finegrained kernels unavailable: {e}") + + def _block_fp8_weight(self, N, K, *, ue8m0): + """(N, K) E4M3 + its (N/128, K/128) inv-scale grid, and the exact dequantized values the + kernel will see — so the reference is the quantization floor, not the pre-quant weight.""" + w = torch.randn(N, K, device="cuda", dtype=torch.float32) + blocks = w.reshape(N // 128, 128, K // 128, 128) + amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12) + inv = amax / 448.0 + if ue8m0: # power-of-two scales: the tcgen05 dot_scaled recipe + inv = torch.pow(2.0, torch.ceil(torch.log2(inv))) + q = (blocks / inv).to(torch.float8_e4m3fn) + deq = (q.float() * inv).reshape(N, K) + return q.reshape(N, K), inv.reshape(N // 128, K // 128).contiguous(), deq + + def test_linear_forward_matches_the_quantization_floor(self): + """A real quantized linear against its own dequantized weight. Loose vs bf16 would pass + even if the kernel silently dropped a K-block; comparing to the floor does not.""" + M, N, K = 256, 512, 384 + for ue8m0 in (False, True): + with self.subTest(ue8m0=ue8m0): + w, s, deq = self._block_fp8_weight(N, K, ue8m0=ue8m0) + x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + out = finegrained_linear(x, w, s, block_size=[128, 128]) + ref = x.float() @ deq.t() + rel = ((out.float() - ref).norm() / ref.norm()).item() + self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") + + def test_swizzled_scales_do_not_change_the_result(self): + """The SWIZZLE_32_4_4 pass reorders scale BYTES; it must not move a single value. Run the + same MXFP8 GEMM against the affine grid and the swizzled artifact and demand agreement — + a wrong reorder shows up here and nowhere in the mocked tests.""" + kernel = load_finegrained_kernel() + N, K, M = 512, 256, 128 + w = torch.randn(N, K, device="cuda", dtype=torch.float32) + groups = w.reshape(N, K // 32, 32) + inv = torch.pow(2.0, torch.ceil(torch.log2(groups.abs().amax(-1, keepdim=True) / 448.0).clamp(min=-127))) + q = (groups / inv).to(torch.float8_e4m3fn).reshape(N, K) + scales = inv.reshape(N, K // 32).to(torch.float8_e8m0fnu) + x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + + affine = kernel.matmul(x, q, None, scales, output_dtype=torch.bfloat16) + swizzled = kernel.matmul( + x, q, None, kernel.swizzle_mx_scales(scales), output_dtype=torch.bfloat16 + ) + torch.cuda.synchronize() + self.assertTrue( + torch.equal(affine, swizzled), + "swizzling changed the result; the artifact is a byte reorder, values must be identical", + ) + + def test_experts_swizzle_pass_keeps_the_forward_correct(self): + """End-to-end over the integration: build real MXFP8 experts, run the post-load swizzle, + and check the fused forward still matches what it produced beforehand.""" + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() + for proj, rows in (("gate_up_proj", 2 * cfg.intermediate_size), ("down_proj", cfg.hidden_size)): + cols = cfg.hidden_size if proj == "gate_up_proj" else cfg.intermediate_size + setattr(experts, proj, torch.nn.Parameter( + torch.randn(cfg.num_local_experts, rows, cols, device="cuda").to(torch.float8_e4m3fn), + requires_grad=False)) + setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter( + torch.full((cfg.num_local_experts, rows, cols // 32), 127, dtype=torch.uint8, device="cuda") + .view(torch.float8_e8m0fnu), requires_grad=False)) + + x = torch.randn(8, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) + idx = torch.randint(0, cfg.num_local_experts, (8, 2), device="cuda", dtype=torch.long) + wts = torch.rand(8, 2, device="cuda", dtype=torch.bfloat16) + + before = experts(x, idx, wts) + model = torch.nn.Module() + model.experts = experts + fg.swizzle_scales_after_loading(model) + after = experts(x, idx, wts) + torch.cuda.synchronize() + + self.assertTrue( + any(getattr(experts, f"{p}_scale_inv_swizzled", None) is not None + for p in ("gate_up_proj", "down_proj")), + "the swizzle pass produced no artifact, so this asserts nothing", + ) + rel = ((after.float() - before.float()).norm() / before.float().norm().clamp(min=1e-9)).item() + self.assertLess(rel, 1e-6, f"swizzled forward diverged from the affine one: {rel:.2e}") From 686238ed59a73cce7aedb8623b91eceebaa17dda Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 18 Aug 2026 02:01:24 -0700 Subject: [PATCH 04/68] finegrained: keep DeepGEMM Mega MoE on the stacked gate|up layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernels store gate|up row-interleaved, but ``transform_weights_for_mega_moe`` does its OWN gate/up interleave and so takes the stacked form. Interleaving those modules at load and undoing it at the boundary is a pointless round trip — and interleaving twice is not the identity, it is a different permutation, i.e. silent garbage rather than a crash. So modules bound for megamoe are simply never interleaved: the post-load pass skips them and records what the module actually holds in ``_gate_up_interleaved`` (the bytes cannot be inspected to tell, and ``set_experts_implementation`` can switch backends after load). ``setup_megamoe_weights`` reads that flag and raises if it is ever handed an interleaved module, rather than re-permuting on an assumption. The other DeepGEMM expert paths need no change: their grouped GEMM is layout-agnostic and the gate split happens in ``_apply_gate``, which already de-interleaves stride-2. Verified: triton backend interleaves and flags; megamoe backend is left untouched with no round trip. The raise itself is unverified here — DeepGEMM's JIT needs a CUDA toolkit >= 12.9 and CUDA_HOME is unset on this box, so no megamoe path executes. --- src/transformers/integrations/deepgemm.py | 10 ++++++++++ src/transformers/integrations/finegrained.py | 10 ++++++++++ 2 files changed, 20 insertions(+) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 20a482aa6523..442f91921834 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -800,6 +800,16 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() + # `transform_weights_for_mega_moe` does its own gate/up interleave, so it takes the stacked + # [gate; up] form. Modules loaded for this backend are never interleaved (see + # `interleave_gate_up_after_loading`); reaching here interleaved means the backend was switched + # after load, and silently re-permuting would be a wrong answer rather than a failure. + if module.has_gate and getattr(module, "_gate_up_interleaved", False): + raise RuntimeError( + "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " + "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — set " + "`_experts_implementation` before loading the model." + ) intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index b8c3fb44d1f7..2a8611505cb2 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1010,9 +1010,19 @@ def interleave_gate_up_after_loading(model: nn.Module, *, already_interleaved: b """ if already_interleaved: return + # Backends that pack gate|up themselves (DeepGEMM Mega MoE) take the stacked form, so skip + # them rather than interleaving and having them undo it. `_gate_up_interleaved` records what + # the module actually holds, since `set_experts_implementation` can switch backends after + # load — a consumer must read the flag, not re-derive the backend. + impl = getattr(getattr(model, "config", None), "_experts_implementation", None) + packs_gate_up_itself = impl in ("deepgemm_megamoe",) for module in model.modules(): if not isinstance(module, FineGrainedExperts) or not module.has_gate: continue + if packs_gate_up_itself: + module._gate_up_interleaved = False + continue + module._gate_up_interleaved = True for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias"): t = getattr(module, name, None) if t is None or t.device.type == "meta": From 641fcb2cfcbb043fd03d67714b96591efb3e9ebf Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 18 Aug 2026 03:36:51 -0700 Subject: [PATCH 05/68] fix --- src/transformers/integrations/finegrained.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 2a8611505cb2..f1ec55784cbb 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1030,11 +1030,16 @@ def interleave_gate_up_after_loading(model: nn.Module, *, already_interleaved: b local = to_local(t) axis = -1 if local.ndim == 1 or name.endswith("bias") else -2 n = local.shape[axis] // 2 + # The reorder is byte-level, so run it on a uint8 view for the 1-byte float8 + # scale/weight dtypes: `torch.stack` lowers to `cat`, and "cat_cuda" is not + # implemented for float8_e8m0fnu (nor e4m3) — it raises mid-load under TP/EP. + byte_view = local.element_size() == 1 and local.dtype.is_floating_point + src = local.view(torch.uint8) if byte_view else local if axis == -1: - rows = torch.stack([local[..., :n], local[..., n:]], dim=-1) + rows = torch.stack([src[..., :n], src[..., n:]], dim=-1) else: - rows = torch.stack([local[..., :n, :], local[..., n:, :]], dim=-2) - local.data.copy_(rows.reshape(local.shape)) + rows = torch.stack([src[..., :n, :], src[..., n:, :]], dim=-2) + src.copy_(rows.reshape(src.shape)) def swizzle_scales_after_loading(model: nn.Module) -> None: From cb04cbe0a5b1cbdd76ae2837e1c13fbfdd4ee222 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 21 Aug 2026 00:10:03 -0700 Subject: [PATCH 06/68] finegrained: build the per-call kernel objects instead of caching them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The memoized `(epilogue, gate_up_quantization, down_quantization)` triple bought nothing where it matters. Constructing the three dataclasses measured ~1.4us per layer (~85us per step at 61 layers), and only in EAGER decode — under cudagraphs the host path does not run on replay, which is the mode this integration deploys in. Against that it held hidden mutable state on the module, and cached a gate_up quantization carrying an `output_recipe` that silently mismatches any caller asking for the unfused form. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 27 ++++++++++++-------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index f1ec55784cbb..55cb8c667bc9 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -638,14 +638,21 @@ def _proj_scale(module, proj: str) -> torch.Tensor: return swizzled if swizzled is not None else to_local(getattr(module, f"{proj}_scale_inv")) -def _cached_kernel_objects(kernel, module): - """The module's static kernel-call objects, built once (per-call dataclass construction - is measurable at decode): ``(epilogue, gate_up_quantization, down_quantization)``.""" - if not hasattr(module, "_cached_epilogue"): - module._cached_epilogue = _kernel_epilogue(kernel, module) - module._cached_gate_up_quantization = _gate_up_quantization(kernel, module, module._cached_epilogue) - module._cached_down_quantization = _kernel_quantization(kernel, module.activation_format) - return module._cached_epilogue, module._cached_gate_up_quantization, module._cached_down_quantization +def _kernel_objects(kernel, module): + """The three per-call kernel objects: ``(epilogue, gate_up_quantization, down_quantization)``. + + Built per call, not memoized on the module. Caching them measured ~1.4us per layer (~85us per + step on a 61-layer model) and only in EAGER decode — under cudagraphs the host path does not + run on replay, so it bought nothing in the mode this integration actually deploys in. Against + that: hidden mutable state on the module, and a cached gate_up quantization carrying an + ``output_recipe`` that silently mismatches any caller asking for the unfused form. + """ + epilogue = _kernel_epilogue(kernel, module) + return ( + epilogue, + _gate_up_quantization(kernel, module, epilogue), + _kernel_quantization(kernel, module.activation_format), + ) def _apply_unfused_gate_up(self, proj_out, up_bias_ids): @@ -694,7 +701,7 @@ def finegrained_batched_mm_experts_forward( expert_ids = top_k_index.reshape(-1) # (S,) up_name = "gate_up_proj" if self.has_gate else "up_proj" - epilogue, gate_up_quantization, down_quantization = _cached_kernel_objects(kernel, self) + epilogue, gate_up_quantization, down_quantization = _kernel_objects(kernel, self) proj_out = kernel.batched_matmul( hidden_states, @@ -754,7 +761,7 @@ def finegrained_grouped_mm_experts_forward( expert_start, gather_idx, scatter_idx = kernel.compute_grouped_scheduling(top_k_index, self.num_experts, num_top_k) up_name = "gate_up_proj" if self.has_gate else "up_proj" - epilogue, gate_up_quantization, down_quantization = _cached_kernel_objects(kernel, self) + epilogue, gate_up_quantization, down_quantization = _kernel_objects(kernel, self) proj_out = kernel.grouped_matmul( hidden_states, From e1c7a921d82cfe3fc803fad84d13eb6be176fe08 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 21 Aug 2026 02:35:41 -0700 Subject: [PATCH 07/68] finegrained: fuse both MoE biases into the GEMMs The kernels now take a per-expert bias as a `bias=` operand, so a biased model no longer falls back to the unfused two-GEMM form just for having a bias. `_kernel_epilogue` stops bailing on `has_bias`, and both host-side adds go away: `_apply_unfused_gate_up` keeps only the activation, `_finish_down` only the routing-weighted reduce. That also drops the `torch.sort` the grouped path paid to index the gate_up bias by expert-sorted row -- the kernel indexes by the tile's own expert id instead. An unsupported act_fn no longer costs the bias its fusion either: the bias is an operand beside the epilogue rather than part of it, so it fuses whether or not the GLU does, and `epilogue is None` goes back to meaning exactly one thing. The bias rides the same output-row axis as the weight and scale grid, so `interleave_gate_up_after_loading` already delivers it in the kernels' interleaved order -- no reordering here. Tests assert the operand reaches both GEMMs, plus a new case pinning that an unfusable activation still fuses its bias. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 59 +++++++++++--------- tests/kernels/test_finegrained.py | 24 +++++++- 2 files changed, 53 insertions(+), 30 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 55cb8c667bc9..30f62d5f3b72 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -80,9 +80,10 @@ class FineGrained: matmuls, with the recipe resolved off the weight/scale dtypes inside the kernel; the per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary classes ship in the same build. MoE blocks are COMPOSED from the matmuls rather than wrapped: the - experts forwards pass ``epilogue`` to the gate_up GEMM when the activation is fusable - (no bias, supported act_fn) and fall back to the unfused two-GEMM form otherwise. All - symbols are required — a build missing any raises at load with the full list. + experts forwards pass ``epilogue`` to both GEMMs — it carries the per-expert biases + always, and the fused GLU when the activation is one the kernels implement (an unsupported + act_fn still fuses its bias and runs only the activation in the host tail). All symbols + are required — a build missing any raises at load with the full list. """ matmul: Callable @@ -578,11 +579,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _kernel_epilogue(kernel, module) -> object | None: - """The fused gate|up + GLU epilogue for the gate_up GEMM, when the kernel can apply it: - gated experts whose activation the epilogue implements and whose gate_up has NO bias (a - fused GLU leaves nowhere to add one — biased models take the unfused two-GEMM form).""" + """The fused gate|up + GLU epilogue for the gate_up GEMM, when the kernels implement the + activation. ``None`` leaves the activation to the host tail. A bias does NOT gate this — + it is a separate ``bias=`` operand on the GEMM (see ``_proj_bias``), fused either way.""" act_name = getattr(module, "act_fn_name", None) - if not module.has_gate or module.has_bias or act_name not in _FUSABLE_ACT_FNS: + if not module.has_gate or act_name not in _FUSABLE_ACT_FNS: return None return kernel.Epilogue( gate=True, @@ -592,6 +593,16 @@ def _kernel_epilogue(kernel, module) -> object | None: ) +def _proj_bias(module, proj: str) -> torch.Tensor | None: + """A projection's per-expert bias, or ``None``. The kernels add it to the accumulator: for + gate_up before the gated split (so an ungated GEMM emitting the raw ``2*I`` pre-activation + takes the same add at doubled width), for down before the routing-weight multiply — each + expert output is ``x @ W_d^T + b_d``. gate_up's is already in the kernels' interleaved row + order (``interleave_gate_up_after_loading`` moves weight, scale grid and bias together off + the same output-row axis).""" + return to_local(getattr(module, f"{proj}_bias")) if module.has_bias else None + + def _kernel_quantization(kernel, activation_format: str | None): """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" @@ -655,25 +666,18 @@ def _kernel_objects(kernel, module): ) -def _apply_unfused_gate_up(self, proj_out, up_bias_ids): - """The unfused gate_up tail (fused-epilogue models never reach this): optional - per-expert bias (rows indexed by ``up_bias_ids``; sentinel rows clamp to a valid - expert — the reduce skips them), then the gating/activation in torch.""" - if self.has_bias: - up_bias = to_local(self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) - proj_out = proj_out + up_bias[up_bias_ids.clamp(max=self.num_experts - 1)] +def _apply_unfused_gate_up(self, proj_out): + """The gate_up tail for models whose activation the kernels do not implement: the + gating/activation in torch. The bias is NOT here — it rides the GEMM's epilogue, which + applies whether or not the GLU itself fused.""" return self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) def _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, out_dtype): - """The shared bookend after the down GEMM (rows in ROUTED order): optional down bias - (BEFORE the routing-weight multiply — each expert output is ``x @ W_d^T + b_d``), then - the kernels' routing-weighted top-k reduce (skips EP-sentinel rows, whose GEMM output - is uninitialized by contract).""" - if self.has_bias: - down_bias = to_local(self.down_proj_bias) - routed_ids = top_k_index.reshape(-1) - proj_out = proj_out + down_bias[routed_ids.clamp(max=self.num_experts - 1)] + """The shared bookend after the down GEMM (rows in ROUTED order): the kernels' + routing-weighted top-k reduce (skips EP-sentinel rows, whose GEMM output is uninitialized + by contract). The down bias is already in ``proj_out`` — it is a ``bias=`` operand on the + down GEMM, which puts it before the routing-weight multiply as required.""" return kernel.weighted_reduce(proj_out, top_k_index, top_k_weights, self.num_experts).to(out_dtype) @@ -711,6 +715,7 @@ def finegrained_batched_mm_experts_forward( expert_ids=expert_ids, gather_idx=gather_idx, epilogue=epilogue, + bias=_proj_bias(self, up_name), quantization=gate_up_quantization, b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) or (S, I) @@ -720,8 +725,7 @@ def finegrained_batched_mm_experts_forward( proj_out, inter_scale = proj_out # the fused requant's quantized intermediate if epilogue is None: - # batched rows are in ROUTED order, so the bias indexes by the flat routed ids - proj_out = _apply_unfused_gate_up(self, proj_out, expert_ids) + proj_out = _apply_unfused_gate_up(self, proj_out) proj_out = kernel.batched_matmul( proj_out, @@ -729,6 +733,7 @@ def finegrained_batched_mm_experts_forward( inter_scale, # fused requant handoff: pre-quantized intermediate scales (None = raw) _proj_scale(self, "down_proj"), expert_ids=expert_ids, + bias=_proj_bias(self, "down_proj"), quantization=None if inter_scale is not None else down_quantization, b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, ) # (S, hidden_dim), routed order @@ -771,6 +776,7 @@ def finegrained_grouped_mm_experts_forward( expert_start=expert_start, gather_idx=gather_idx, epilogue=epilogue, + bias=_proj_bias(self, up_name), quantization=gate_up_quantization, b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) @@ -780,9 +786,7 @@ def finegrained_grouped_mm_experts_forward( proj_out, inter_scale = proj_out # the fused requant's quantized intermediate if epilogue is None: - # grouped gate_up rows are expert-SORTED; only biased unfused models pay the sort - up_bias_ids = torch.sort(top_k_index.reshape(-1)).values if self.has_bias else top_k_index.reshape(-1) - proj_out = _apply_unfused_gate_up(self, proj_out, up_bias_ids) + proj_out = _apply_unfused_gate_up(self, proj_out) proj_out = kernel.grouped_matmul( proj_out, @@ -791,6 +795,7 @@ def finegrained_grouped_mm_experts_forward( _proj_scale(self, "down_proj"), expert_start=expert_start, scatter_idx=scatter_idx, # scatter straight to routed rows — weighted_reduce's layout + bias=_proj_bias(self, "down_proj"), quantization=None if inter_scale is not None else down_quantization, b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, ) # (S, hidden_dim), ROUTED order after the scatter diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 8b7631e8a616..ec3c677ba9ac 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -259,10 +259,28 @@ def test_bias_lands_before_gate_and_before_routing(self): p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: out = fg.finegrained_batched_mm_experts_forward(m, hs, idx, wts) - # zero matmul outputs + bias: result = sum_k w_k * (gate(bias_up) @ ... ) — just assert - # the bias params were consumed (non-None) and output is finite self.assertTrue(torch.isfinite(out).all()) - self.assertEqual(len(rec.calls["matmul_batched"]), 2) + up, down = rec.calls["matmul_batched"] + # both biases ride their GEMM's epilogue: the kernel adds gate_up's before the gated + # split and down's before the routing-weight multiply. Nothing is added host-side, so + # a biased model keeps the fused form instead of falling back to two GEMMs. + self.assertIs(up.kwargs["bias"], m.gate_up_proj_bias) + self.assertTrue(up.kwargs["epilogue"].gate) + self.assertIs(down.kwargs["bias"], m.down_proj_bias) + + def test_unfusable_act_fn_still_fuses_the_bias(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, has_bias=True) + m.act_fn_name = "quick_gelu" # not in _FUSABLE_ACT_FNS + hs, idx, wts = self._route() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, hs, idx, wts) + up, _ = rec.calls["matmul_batched"] + # the GLU cannot fuse, but the ungated GEMM emitting the raw 2*I pre-activation takes + # the same bias add at doubled width — only the activation falls back to the host tail + self.assertIsNone(up.kwargs["epilogue"]) + self.assertIs(up.kwargs["bias"], m.gate_up_proj_bias) @require_torch From eac0ee87e929d2b6e72df5445be522bc21b73eef Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 11 Sep 2026 02:09:29 -0700 Subject: [PATCH 08/68] finegrained: expert layouts as conversion ops; experts forwards on the kernels' MoE chain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The experts forwards are adapters over the kernels' `moe_fused_batched/grouped`: `_moe_operands` hands them the module's tensors and the activation — a `get_supported_act_fns()` name the gate_up epilogue fuses, else the module's own GLU as a callable the kernels run on the host, so any activation works without a kernel release. The kernel bundle is loaded by symbol name (`matmul_2d/batched/grouped`, the two forwards, swizzle/unswizzle, `get_supported_act_fns`). Every layout difference between a checkpoint and what the modules hold is a `ConversionOps` with a reverse, attached by the quantizer, so `save_pretrained` restores the checkpoint bitwise: - `FineGrainedInterleaveGateUp`: stacked [gate; up] rows -> the kernels' [g0, u0, ...] order, skipped when the model declares `is_concatenated=False` (GPT-OSS; the flag travels through `use_experts_implementation`, which stamps its defaults after `__init__`) or the backend packs gate|up itself (DeepGEMM Mega MoE); - `FineGrainedScaleContainer`: a scale into the dtype its module holds — the same bytes for a uint8 UE8M0 container (MiniMax), an exact cast for float32 values (dsv4-flash-base) — and the module records the shipped container so the reverse restores it; - `FineGrainedSwizzleScales`: the SWIZZLE_32_4_4 artifact for modules that hold 5-D scales (`FineGrainedExperts.__init__` allocates that shape for every non-DeepGEMM backend on SM100 with a group-scaled format and activation quant; whole 128-row/4-col blocks only); - `FineGrainedPackedBlocks` regroups GPT-OSS `{proj}_blocks` into the packed weight, so the blocks/scales format is two one-to-one converters and round-trips too. Catch-all converters cover keys that already arrive under the fused names and dense scale keys. The post-load hooks, `_save_to_state_dict` override and post-load dtype cast are gone. Also: the eager per-expert loop reads a swizzled stack's `[e:e+1]` slice and passes bias, NVFP4 global and activation format (it dropped all three); expert biases take the model dtype; the `mxfp4` format pins UE8M0 scales; `FineGrainedEmbedding` (FP8 table, per-tensor scale) for `modules_to_convert`; DeepGEMM backends refuse modules holding swizzled scales or interleaved rows before loading the kernel; `disable_deepgemm_on_multi_device` is public. Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/__init__.py | 16 +- src/transformers/integrations/deepgemm.py | 72 +- src/transformers/integrations/finegrained.py | 985 ++++++++---------- .../quantizers/quantizer_finegrained.py | 268 ++--- tests/kernels/test_finegrained.py | 707 +++++++++---- 5 files changed, 1147 insertions(+), 901 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index f3231ea81efc..5e21b782694f 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,7 +50,13 @@ ], "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], - "finegrained": ["FineGrainedExperts", "FineGrainedLinear", "replace_with_finegrained_layer"], + "finegrained": [ + "FineGrainedEmbedding", + "FineGrainedExperts", + "FineGrainedLinear", + "replace_with_finegrained_embedding", + "replace_with_finegrained_layer", + ], "finegrained_fp8": ["FP8Embedding", "FP8Linear", "replace_with_fp8_embedding", "replace_with_fp8_linear"], "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "gemma_quant": [ @@ -211,7 +217,13 @@ ) from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear - from .finegrained import FineGrainedExperts, FineGrainedLinear, replace_with_finegrained_layer + from .finegrained import ( + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedLinear, + replace_with_finegrained_embedding, + replace_with_finegrained_layer, + ) from .finegrained_fp8 import FP8Embedding, FP8Linear, replace_with_fp8_embedding, replace_with_fp8_linear from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .gemma_quant import ( diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index f29c312acf87..4ac05fde7e12 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -684,6 +684,32 @@ def deepgemm_bf16_experts_forward( ) +def _assert_affine_scales(module: torch.nn.Module) -> None: + """DeepGEMM reads row-major block scales. A module loaded for a triton backend holds them in the + SWIZZLE_32_4_4 layout (5-D); consuming that as affine returns garbage, so a backend switched to + DeepGEMM after such a load fails here instead.""" + for name in ("gate_up_proj_scale_inv", "up_proj_scale_inv", "down_proj_scale_inv"): + scale = getattr(module, name, None) + if scale is not None and scale.ndim == 5: + raise RuntimeError( + f"DeepGEMM experts need row-major block scales, but `{name}` is held swizzled for the " + "triton experts backends. Switching to a DeepGEMM backend after load is not supported — " + "pass `experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + +def _assert_stacked_gate_up(module: torch.nn.Module) -> None: + """Mega MoE's ``transform_weights_for_mega_moe`` does its own gate/up interleave, so it takes the + stacked ``[gate; up]`` rows. A module loaded for a triton backend holds them interleaved; switched + to this backend after load, silently re-permuting would be a wrong answer rather than a failure.""" + if module.has_gate and getattr(module, "_gate_up_interleaved", False): + raise RuntimeError( + "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " + "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — pass " + "`experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -692,15 +718,20 @@ def deepgemm_fp8_fp4_experts_forward( ) -> torch.Tensor: if self._deepgemm_disabled: # Set at load when the model spans >1 CUDA device in this process, where DeepGEMM's - # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`). + # context-bound kernels corrupt across devices (see `disable_deepgemm_on_multi_device`). raise RuntimeError( "DeepGEMM experts selected on a model spanning multiple CUDA devices in one process; " "its kernels are bound to a single CUDA context and corrupt across devices. Use " "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. - _assert_sm100_requirements(self.down_proj, self.down_proj_scale_inv) + # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as + # affine and turn into garbage — refuse before any kernel work. + _assert_affine_scales(self) + # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is + # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the + # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. + _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) deepgemm = load_deepgemm_kernel() @@ -794,33 +825,29 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: Unwraps any ``DTensor`` wrappers FSDP2/EP may have placed around the loader- side Parameters — the kernel takes raw pointers. """ - deepgemm = load_deepgemm_kernel() - gate_up_sf_raw = module.gate_up_proj_scale_inv.data - down_sf_raw = module.down_proj_scale_inv.data - # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. - gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() - down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() - # `transform_weights_for_mega_moe` does its own gate/up interleave, so it takes the stacked - # [gate; up] form. Modules loaded for this backend are never interleaved (see - # `interleave_gate_up_after_loading`); reaching here interleaved means the backend was switched - # after load, and silently re-permuting would be a wrong answer rather than a failure. - if module.has_gate and getattr(module, "_gate_up_interleaved", False): - raise RuntimeError( - "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " - "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — set " - "`_experts_implementation` before loading the model." - ) - + # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as + # affine and turn into garbage — refuse before the (hub-download + JIT) kernel load. + _assert_affine_scales(module) + # Same for its interleaved gate|up rows: `transform_weights_for_mega_moe` interleaves stacked rows + # itself, and re-permuting already-interleaved ones would be a silent wrong answer. + _assert_stacked_gate_up(module) + # The UE8M0 scale grid is group-32 along both dims, so the kernel's SF layout needs both divisible. intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts hidden_dim = module.hidden_dim - if hidden_dim % 32 != 0 or intermediate_hidden % 32 != 0: raise ValueError( f"DeepGEMM Mega MoE requires `hidden_dim` and `intermediate_hidden` divisible by 32 " f"(FP8 SF granularity); got hidden_dim={hidden_dim}, intermediate_hidden={intermediate_hidden}." ) + deepgemm = load_deepgemm_kernel() + gate_up_sf_raw = module.gate_up_proj_scale_inv.data + down_sf_raw = module.down_proj_scale_inv.data + # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. + gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() + down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() + gate_up_sf = deepgemm.transform_sf_into_required_layout( gate_up_sf_raw.float(), 2 * intermediate_hidden, @@ -870,6 +897,9 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( `transform_weights_for_mega_moe((gate_up, gate_up_sf), (down, down_sf))`. - `config.swiglu_limit` (optional): SwiGLU clamp; absent → unclamped. """ + # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as + # affine and turn into garbage — refuse before any kernel work. + _assert_affine_scales(self) # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 30f62d5f3b72..ae115490db7f 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -25,10 +25,9 @@ from torch.nn import functional as F from ..activations import ACT2FN -from ..core_model_loading import ConversionOps +from ..core_model_loading import ConversionOps, Interleave from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module from ..utils import logging -from ..utils.deprecation import deprecate_kwarg from ..utils.import_utils import is_kernels_available from .deepgemm import ( deepgemm_fp8_fp4_experts_forward, @@ -76,23 +75,24 @@ def _first_attr(obj, *names): class FineGrained: """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. - Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the three - matmuls, with the recipe resolved off the weight/scale dtypes inside the kernel; the - per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary classes ship - in the same build. MoE blocks are COMPOSED from the matmuls rather than wrapped: the - experts forwards pass ``epilogue`` to both GEMMs — it carries the per-expert biases - always, and the fused GLU when the activation is one the kernels implement (an unsupported - act_fn still fuses its bias and runs only the activation in the host tail). All symbols - are required — a build missing any raises at load with the full list. + Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and + the two MoE forwards, with the recipe resolved off the weight/scale dtypes inside the + kernel; the per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary + classes ship in the same build. The experts forwards hand the kernels' MoE forwards the + module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, + biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as + the module's own callable and runs on the host between the two GEMMs. All symbols are + required — a build missing any raises at load with the full list. """ - matmul: Callable - batched_matmul: Callable - grouped_matmul: Callable - nvfp4_quantize_two_level: Callable + matmul_2d: Callable + matmul_batched: Callable + matmul_grouped: Callable + moe_fused_batched: Callable + moe_fused_grouped: Callable swizzle_mx_scales: Callable - compute_grouped_scheduling: Callable - weighted_reduce: Callable + unswizzle_mx_scales: Callable + get_supported_act_fns: Callable Quantization: Callable Epilogue: Callable @@ -154,35 +154,14 @@ def _load_finegrained_kernel() -> None: "has a build matching the current torch/CUDA." ) - required = ( - "matmul_2d", - "matmul_batched", - "matmul_grouped", - "nvfp4_quantize_two_level", - "swizzle_mx_scales", - "compute_grouped_scheduling", - "weighted_reduce", - "Quantization", - "Epilogue", - ) - symbols = {name: getattr(kernel, name, None) for name in required} + # the bundle's fields ARE the required symbols, under the kernel's own names + symbols = {name: getattr(kernel, name, None) for name in FineGrained.__dataclass_fields__} missing = [name for name, attr in symbols.items() if attr is None] if missing: raise ImportError( f"finegrained-kernels build is missing required symbols: {', '.join(missing)}. {_MISSING_KERNELS_MESSAGE}" ) - - _FINEGRAINED = FineGrained( - matmul=symbols["matmul_2d"], - batched_matmul=symbols["matmul_batched"], - grouped_matmul=symbols["matmul_grouped"], - nvfp4_quantize_two_level=symbols["nvfp4_quantize_two_level"], - compute_grouped_scheduling=symbols["compute_grouped_scheduling"], - weighted_reduce=symbols["weighted_reduce"], - swizzle_mx_scales=symbols["swizzle_mx_scales"], - Quantization=symbols["Quantization"], - Epilogue=symbols["Epilogue"], - ) + _FINEGRAINED = FineGrained(**symbols) def load_finegrained_kernel() -> FineGrained: @@ -215,8 +194,10 @@ def _weight_formats() -> dict[str, _WeightFormat]: "fp8": _WeightFormat(weight_dtype=_FP8_DTYPE), # E4M3 values, UE8M0 group-32 scales "mxfp8": _WeightFormat(weight_dtype=_FP8_DTYPE, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32)), - # packed E2M1 values (2/byte), UE8M0 group-32 scales (dsv4 ships scale_fmt's container) - "mxfp4": _WeightFormat(weight_dtype=torch.int8, values_per_byte=2, scale_group=(1, 32)), + # packed E2M1 values (2/byte), UE8M0 group-32 scales + "mxfp4": _WeightFormat( + weight_dtype=torch.int8, values_per_byte=2, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32) + ), # packed E2M1 values, E4M3 group-16 block scales, per-tensor/per-expert fp32 global "nvfp4": _WeightFormat( weight_dtype=torch.int8, @@ -249,39 +230,60 @@ def _alloc_expert_proj( num_experts: int, proj_out: int, proj_in: int, - weight_dtype: torch.dtype, - sf_dtype: torch.dtype, - weight_k_div: int = 1, - sf_gran_n: int | None = None, - sf_gran_k: int | None = None, - min_sf_out: int = 1, + fmt: _WeightFormat, + scale_dtype: torch.dtype, + scale_group: tuple[int, int] | None, + min_scale_out: int = 1, + swizzled: bool = False, ) -> tuple[nn.Parameter, nn.Parameter]: - """Allocate `(weight, weight_scale_inv)` parameters for one expert projection. + """``(weight, weight_scale_inv)`` Parameters for one expert projection: the weight in the + format's storage (FP4 packs two values per byte along K), the scale grid at the format's + granularity (``None`` = one scale per expert). ``min_scale_out`` floors the grid's output dim — + the fused gate_up projection keeps room for both halves (``2``) even when ``proj_out`` is + smaller than the block. ``swizzled`` allocates the grid in the ``SWIZZLE_32_4_4`` artifact + layout instead — 128-row blocks x 4-column groups of ``(2, 256)`` byte tiles, the expert axis + leading so EP shards and gathers it on dim 0 — which needs whole blocks; a projection that + does not tile stays affine, which the kernels read directly.""" + weight = torch.empty(num_experts, proj_out, proj_in // fmt.values_per_byte, dtype=fmt.weight_dtype) + if scale_group is None: + scale_shape = (num_experts, max(1, min_scale_out), 1) + else: + rows, cols = max(_cdiv(proj_out, scale_group[0]), min_scale_out), _cdiv(proj_in, scale_group[1]) + scale_shape = (num_experts, rows, cols) + if swizzled and rows % 128 == 0 and cols % 4 == 0: + scale_shape = (num_experts, rows // 128, cols // 4, 2, 256) + scale = torch.empty(scale_shape, dtype=scale_dtype) + return ( + nn.Parameter(weight, requires_grad=weight.is_floating_point()), + nn.Parameter(scale, requires_grad=scale.is_floating_point() and not swizzled), + ) - `weight_k_div` halves the K dim for FP4-packed storage (2 e2m1 values per byte). - `sf_gran_n` / `sf_gran_k` set per-block (None → per-row/per-tensor) SF granularity. - `min_sf_out` floors the SF tensor's output dim — used by the fused gate_up - projection to keep room for both halves (pass `2`) even when `proj_out < sf_gran_n` - would otherwise collapse the SF dim to 1. - """ - weight_t = torch.empty(num_experts, proj_out, proj_in // weight_k_div, dtype=weight_dtype) - weight = nn.Parameter(weight_t, requires_grad=weight_t.is_floating_point()) - sf_out = max(_cdiv(proj_out, sf_gran_n) if sf_gran_n is not None else 1, min_sf_out) - sf_in = _cdiv(proj_in, sf_gran_k) if sf_gran_k is not None else 1 - sf_t = torch.empty(num_experts, sf_out, sf_in, dtype=sf_dtype) - sf = nn.Parameter(sf_t, requires_grad=sf_t.is_floating_point()) - return weight, sf + +def _swizzles_scales(config, fmt: _WeightFormat, activation_format: str | None) -> bool: + """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the Blackwell + tcgen05 scaled-MMA reads directly (plain row-major forces a per-tile gather that caps the scaled + dot below the fp8/fp4 peak): SM100, a triton dispatch (the fused forwards and the eager loop both + read it; the DeepGEMM backends read affine scales), a group-scaled format (MX group-32, NVFP4 + group-16 — block-FP8's grid never reaches a scaled-MMA, even with UE8M0 scales) and a chain that + quantizes activations (weight-only reads scales per group affinely). + ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` keeps every module affine (debug / A/B).""" + if os.environ.get("TRANSFORMERS_FINEGRAINED_NO_SWIZZLE", "0") == "1": + return False + return ( + getattr(config, "_experts_implementation", None) not in ("deepgemm", "deepgemm_megamoe") + and fmt.scale_group is not None + and _block_recipe(activation_format) is not None + and torch.cuda.is_available() + and is_sm100() + ) -@deprecate_kwarg("output_dtype", version="v5.16") def finegrained_triton_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor, - block_size: list[int] | None = None, bias: torch.Tensor | None = None, activation_scale: torch.Tensor | None = None, - output_dtype: torch.dtype | None = None, weight_global_scale: torch.Tensor | None = None, activation_format: str | None = None, ) -> torch.Tensor: @@ -289,26 +291,21 @@ def finegrained_triton_linear( Serves every weight recipe the kernel resolves off the tensors themselves — block-FP8 (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4 (``int8``-packed values; the two-level - per-tensor ``weight_global_scale`` recovers on the accumulator). ``block_size`` is accepted - for back-compat and ignored: the quantization block is derived from the scale tensor's - shape. ``activation_scale=None`` → dynamic activation quant (inline); a per-tensor scalar → - static quant against it. ``activation_format`` picks the activation recipe where the weights - leave it open (``None`` = the weight-native choice; ``"bf16"`` = weight-only, no activation - quant — W4A16). ``output_dtype`` defaults to ``input.dtype`` when left ``None``. + per-tensor ``weight_global_scale`` recovers on the accumulator). ``activation_scale=None`` → + dynamic activation quant (inline); a per-tensor scalar → static quant against it. + ``activation_format`` picks the activation recipe where the weights leave it open (``None`` = + the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). """ kernel = load_finegrained_kernel() - quantization = None - if activation_format is not None: - recipe = None if activation_format == "bf16" else activation_format - quantization = kernel.Quantization(input_recipe=recipe) + quantization = _kernel_quantization(kernel, activation_format) original_shape = input.shape - output = kernel.matmul( + output = kernel.matmul_2d( input.reshape(-1, original_shape[-1]), weight, activation_scale, weight_scale_inv, quantization=quantization, - output_dtype=output_dtype if output_dtype is not None else input.dtype, + output_dtype=input.dtype, b_global_scale=weight_global_scale, ) output = output.reshape(*original_shape[:-1], output.shape[-1]) @@ -317,7 +314,6 @@ def finegrained_triton_linear( return output -@deprecate_kwarg("output_dtype", version="v5.16") def finegrained_linear( input: torch.Tensor, weight: torch.Tensor, @@ -325,7 +321,6 @@ def finegrained_linear( block_size: list[int] | None = None, bias: torch.Tensor | None = None, activation_scale: torch.Tensor | None = None, - output_dtype: torch.dtype | None = None, allow_deepgemm: bool = True, weight_global_scale: torch.Tensor | None = None, activation_format: str | None = None, @@ -352,30 +347,27 @@ def finegrained_linear( for dynamic (per-token) quant. allow_deepgemm: set ``False`` to force the Triton fallback for this call. Used when the model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound - to a single CUDA context and produce garbage across devices (see the multi-device guard - in ``quantizer_finegrained.py``). + to a single CUDA context and produce garbage across devices (see + ``disable_deepgemm_on_multi_device``). + weight_global_scale: the NVFP4 two-level per-tensor global (None for other recipes). + activation_format: the activation recipe where the weights leave it open (see + ``finegrained_triton_linear``). """ - # DeepGEMM is CUDA-only, dynamic-only, SM90+ only, FP4/FP8-block-128-only. On SM100 its FP8 GEMM only - # consumes UE8M0 scales; float32 scales would be ceil-rounded to UE8M0 without requantizing and - # silently corrupt the output (#47030), so we skip DeepGEMM up front for that combo rather than - # attempt-then-fall-back every call — `_assert_sm100_requirements` still guards it as a backstop. - # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces the Triton fallback for this single dispatcher (the - # experts ``"deepgemm"`` impl is unaffected — use ``set_experts_implementation`` for that). Used by the - # FP8 MoE batched_mm / grouped_mm paths to avoid a still-unexplained DeepGEMM-vs-Triton interaction that - # degrades end-to-end generation on B200 (per-row kernel outputs still measure bit-perfect, but final - # tokens drift; not reproducible with the DeepGEMM linear off). + # DeepGEMM is CUDA-only, dynamic-only, SM90+, FP4 or 128x128-block FP8; the combos it would + # silently corrupt (float32 scales on SM100, #47030) are skipped up front rather than attempted. + # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces Triton for this dispatcher only. deepgemm_preferred = ( # SM100 perf: Triton measures 2.9x (decode) / 1.45x (prefill) over DeepGEMM on the DSV4 # qkv linear — the 128x128 block-FP8 shape DeepGEMM is otherwise preferred for, which the # gate above never catches (block-FP8's (N/128, K/128) grid has no swizzled layout). not is_sm100() # If the model is on multiple devices, DeepGEMM's context-bound kernels corrupt across devices. The - # multi-device guard in `quantizer_finegrained.py` flips `_deepgemm_disabled` True at load, which + # multi-device guard `disable_deepgemm_on_multi_device` flips `_deepgemm_disabled` True at load, which # this dispatcher sees and respects. The Triton fallback is context-free and safe. and allow_deepgemm # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major: DeepGEMM would # consume the permuted buffer as affine and silently return garbage. Correctness gate, - # true on any arch, for whatever `swizzle_scales_after_loading` has already swizzled. + # true on any arch, for whatever the loader's `FineGrainedSwizzleScales` swizzled. and weight_scale_inv.ndim <= 2 and is_deepgemm_loadable() and activation_scale is None @@ -413,7 +405,6 @@ def finegrained_linear( input, weight, weight_scale_inv, - block_size, bias, activation_scale, weight_global_scale=weight_global_scale, @@ -421,12 +412,21 @@ def finegrained_linear( ) -class FineGrainedLinear(nn.Linear): - # Internal, temporary flag — not public API, don't set it directly. `_disable_deepgemm_on_multi_device` +class _FineGrainedModule: + """What every finegrained module (dense linear or experts) carries besides its parameters.""" + + # the dtype the checkpoint shipped this module's block scales in, when it differs from the held + # one (uint8 exponent bytes, float32 values) — set by the loader's `FineGrainedScaleContainer`, + # read back by its reverse so a save restores the checkpoint's container + scale_container_dtype: torch.dtype | None = None + + # Internal, temporary flag — not public API, don't set it directly. `disable_deepgemm_on_multi_device` # flips it True at load when the model spans >1 CUDA device in one process (DeepGEMM's context-bound # kernels corrupt across devices); removable once the kernel ships a context-free loader. _deepgemm_disabled = False + +class FineGrainedLinear(_FineGrainedModule, nn.Linear): def __init__( self, in_features: int, @@ -503,6 +503,23 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: ) +class FineGrainedEmbedding(nn.Embedding): + """``nn.Embedding`` whose table is stored in FP8 with one per-tensor ``weight_scale``. A table + worth quantizing is one too big to dequantize at load (Qwen4-Exp's hashed n-gram table is + ~48 GiB in FP8), so the rescale runs on the rows a lookup gathers. Built on the ``meta`` + device: ``nn.Embedding`` allocates a full-precision table first.""" + + def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int | None = None): + super().__init__(num_embeddings, embedding_dim, padding_idx=padding_idx) + self.weight = nn.Parameter(torch.empty(num_embeddings, embedding_dim, dtype=_FP8_DTYPE), requires_grad=False) + # `(1,)` rather than a scalar: the shape checkpoints serialize per-tensor scales with + self.weight_scale = nn.Parameter(torch.ones(1), requires_grad=False) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + rows = super().forward(input) + return rows.to(self.weight_scale.dtype) * self.weight_scale + + class FineGrainedGroupedLinear(FineGrainedLinear): """FP8 drop-in for block-diagonal grouped linears. @@ -560,7 +577,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: expert_start = torch.arange(0, self.n_groups + 1, device=x.device, dtype=torch.int32) * tokens_per_group kernel = load_finegrained_kernel() - y = kernel.grouped_matmul( + y = kernel.matmul_grouped( x, w, None, @@ -575,34 +592,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -_FUSABLE_ACT_FNS = ("silu", "gelu", "relu") - - -def _kernel_epilogue(kernel, module) -> object | None: - """The fused gate|up + GLU epilogue for the gate_up GEMM, when the kernels implement the - activation. ``None`` leaves the activation to the host tail. A bias does NOT gate this — - it is a separate ``bias=`` operand on the GEMM (see ``_proj_bias``), fused either way.""" - act_name = getattr(module, "act_fn_name", None) - if not module.has_gate or act_name not in _FUSABLE_ACT_FNS: - return None - return kernel.Epilogue( - gate=True, - act_fn=act_name, - swiglu_alpha=module.swiglu_alpha, - swiglu_limit=module.swiglu_limit, - ) - - -def _proj_bias(module, proj: str) -> torch.Tensor | None: - """A projection's per-expert bias, or ``None``. The kernels add it to the accumulator: for - gate_up before the gated split (so an ungated GEMM emitting the raw ``2*I`` pre-activation - takes the same add at doubled width), for down before the routing-weight multiply — each - expert output is ``x @ W_d^T + b_d``. gate_up's is already in the kernels' interleaved row - order (``interleave_gate_up_after_loading`` moves weight, scale grid and bias together off - the same output-row axis).""" - return to_local(getattr(module, f"{proj}_bias")) if module.has_bias else None - - def _kernel_quantization(kernel, activation_format: str | None): """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" @@ -612,203 +601,74 @@ def _kernel_quantization(kernel, activation_format: str | None): return kernel.Quantization(input_recipe=recipe) -@functools.lru_cache(maxsize=64) -def _batched_gather_idx(num_tokens: int, num_top_k: int, device: torch.device) -> torch.Tensor: - """The batched routed-row gather map (row ``s`` reads token ``s // num_top_k``): depends - only on the SHAPE, so cache it — its build is pure eager launch overhead at decode.""" - return torch.arange(num_tokens * num_top_k, device=device, dtype=torch.int32) // num_top_k - - -def _activation_recipe(module) -> str | None: - """The activation quantization recipe this module's chain runs: the explicit - ``activation_format`` when set (``"bf16"`` = weight-only -> None), else the - weight-native default (fp8 weights quantize activations to block-FP8, MX/NV weights - to their own family — the same resolution the kernels apply).""" - fmt = module.activation_format - if fmt is not None: - return None if fmt == "bf16" else fmt - return module.weight_format if module.weight_format is not None else "fp8" - - -def _gate_up_quantization(kernel, module, epilogue): - """The gate_up GEMM's ``Quantization`` — WITH the fused intermediate requant whenever the - GLU is fused: ``output_recipe`` makes the epilogue emit the quantized intermediate - ``(C, Cs)`` that the down GEMM consumes directly (``As=Cs``), matching the kernels' own - ``moe_fused_*`` chains. Without the fused epilogue the intermediate is produced host-side - in bf16 and the down GEMM quantizes it internally (the unfused form).""" - recipe = _activation_recipe(module) - if epilogue is None or recipe is None: - return _kernel_quantization(kernel, module.activation_format) - return kernel.Quantization(input_recipe=recipe, output_recipe=recipe) - - -def _proj_scale(module, proj: str) -> torch.Tensor: - """A projection's weight scales — the pre-swizzled cache when the post-load hook built - one (the tcgen05 fast path), else the affine Parameter.""" - swizzled = getattr(module, f"{proj}_scale_inv_swizzled", None) - return swizzled if swizzled is not None else to_local(getattr(module, f"{proj}_scale_inv")) - - -def _kernel_objects(kernel, module): - """The three per-call kernel objects: ``(epilogue, gate_up_quantization, down_quantization)``. - - Built per call, not memoized on the module. Caching them measured ~1.4us per layer (~85us per - step on a 61-layer model) and only in EAGER decode — under cudagraphs the host path does not - run on replay, so it bought nothing in the mode this integration actually deploys in. Against - that: hidden mutable state on the module, and a cached gate_up quantization carrying an - ``output_recipe`` that silently mismatches any caller asking for the unfused form. - """ - epilogue = _kernel_epilogue(kernel, module) - return ( - epilogue, - _gate_up_quantization(kernel, module, epilogue), - _kernel_quantization(kernel, module.activation_format), - ) - +def _block_recipe(activation_format: str | None) -> str | None: + """The module-level ``activation_format`` as the kernels' block ``recipe``: ``None`` keeps + the weight family's own format, ``"bf16"`` is weight-only (no activation quant), an explicit + name is itself.""" + if activation_format is None: + return "weights" + return None if activation_format == "bf16" else activation_format + + +def _moe_operands(kernel, module) -> dict: + """Everything the kernels' MoE forwards take from the module: the two projections with their + scales (held swizzled for dot_scaled chains), the per-expert globals and biases, + and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the module's + own GLU as a callable the kernels run on the host, so any activation works without a kernel + change.""" + up = "gate_up_proj" if module.has_gate else "up_proj" + + def local(name): + t = getattr(module, name, None) + return to_local(t) if t is not None else None + + act_name = module.act_fn_name + fused = act_name in kernel.get_supported_act_fns() and (module.swiglu_alpha is None or act_name == "silu") + if fused: + act_fn = act_name + else: + act_fn = module._apply_gate if module.has_gate else module.act_fn + return { + "gate_up_proj": local(up), + "down_proj": local("down_proj"), + "gate_up_proj_scale_inv": local(f"{up}_scale_inv"), + "down_proj_scale_inv": local("down_proj_scale_inv"), + "gate_up_proj_global_scale": local(f"{up}_global_scale"), + "down_proj_global_scale": local("down_proj_global_scale"), + "gate_up_proj_bias": local(f"{up}_bias"), + "down_proj_bias": local("down_proj_bias"), + "act_fn": act_fn, + "swiglu_alpha": module.swiglu_alpha, + "swiglu_limit": module.swiglu_limit, + "recipe": _block_recipe(module.activation_format), + "gate": module.has_gate, + } -def _apply_unfused_gate_up(self, proj_out): - """The gate_up tail for models whose activation the kernels do not implement: the - gating/activation in torch. The bias is NOT here — it rides the GEMM's epilogue, which - applies whether or not the GLU itself fused.""" - return self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) +def _fused_experts_forward(module, kernel_forward: str, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: + """The kernels' fused MoE chain (gate_up with the fused GLU epilogue and intermediate requant + where supported -> down -> the routing-weighted top-k reduce) over the module's tensors.""" + if module.activation_scheme == "static": + raise NotImplementedError( + f"the {kernel_forward} experts dispatch does not support activation_scheme='static'; use " + "experts_implementation='eager' or activation_scheme='dynamic'." + ) + kernel = load_finegrained_kernel() + return getattr(kernel, kernel_forward)(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, module)) -def _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, out_dtype): - """The shared bookend after the down GEMM (rows in ROUTED order): the kernels' - routing-weighted top-k reduce (skips EP-sentinel rows, whose GEMM output is uninitialized - by contract). The down bias is already in ``proj_out`` — it is a ``bias=`` operand on the - down GEMM, which puts it before the routing-weight multiply as required.""" - return kernel.weighted_reduce(proj_out, top_k_index, top_k_weights, self.num_experts).to(out_dtype) +def finegrained_batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: + """Batched (decode) experts forward: one program per routed row.""" + return _fused_experts_forward(self, "moe_fused_batched", hidden_states, top_k_index, top_k_weights) -def finegrained_batched_mm_experts_forward( - self: torch.nn.Module, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Batched (decode) experts forward — the kernels' fused chain: gate_up per routed row - (fused GLU epilogue + intermediate requant where supported) -> down consuming the - requant handoff -> ``weighted_reduce``.""" - if self.activation_scheme == "static": - raise NotImplementedError( - "batched_mm experts dispatch does not support activation_scheme='static'. " - "Use the default eager dispatch or switch to activation_scheme='dynamic'." - ) - kernel = load_finegrained_kernel() - num_top_k = top_k_index.size(-1) - - # per-routed-row dispatch: the kernel gathers row s's token (s // num_top_k) itself — - # no (S, H) replica copy; the ops take the router's int64 ids natively - gather_idx = _batched_gather_idx(hidden_states.size(0), num_top_k, hidden_states.device) - expert_ids = top_k_index.reshape(-1) # (S,) - - up_name = "gate_up_proj" if self.has_gate else "up_proj" - epilogue, gate_up_quantization, down_quantization = _kernel_objects(kernel, self) - - proj_out = kernel.batched_matmul( - hidden_states, - to_local(getattr(self, up_name)), - None, - _proj_scale(self, up_name), - expert_ids=expert_ids, - gather_idx=gather_idx, - epilogue=epilogue, - bias=_proj_bias(self, up_name), - quantization=gate_up_quantization, - b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, - ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) or (S, I) - - inter_scale = None - if isinstance(proj_out, (tuple, list)): - proj_out, inter_scale = proj_out # the fused requant's quantized intermediate - - if epilogue is None: - proj_out = _apply_unfused_gate_up(self, proj_out) - - proj_out = kernel.batched_matmul( - proj_out, - to_local(self.down_proj), - inter_scale, # fused requant handoff: pre-quantized intermediate scales (None = raw) - _proj_scale(self, "down_proj"), - expert_ids=expert_ids, - bias=_proj_bias(self, "down_proj"), - quantization=None if inter_scale is not None else down_quantization, - b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, - ) # (S, hidden_dim), routed order - - return _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, hidden_states.dtype) - - -def finegrained_grouped_mm_experts_forward( - self: torch.nn.Module, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, -) -> torch.Tensor: - """Grouped (prefill) experts forward — the kernels' fused chain: one on-device routing - pass -> gate_up over the expert-sorted schedule (fused GLU epilogue + intermediate - requant where supported) -> down consuming the requant handoff and scattering back to - routed rows -> ``weighted_reduce``.""" - if self.activation_scheme == "static": - raise NotImplementedError( - "grouped_mm experts dispatch does not support activation_scheme='static'. " - "Use the default eager dispatch or switch to activation_scheme='dynamic'." - ) +def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: + """Grouped (prefill) experts forward: one on-device routing pass, then the chain over the + expert-sorted schedule, scattering back to routed rows (EP-sentinel rows skipped).""" + return _fused_experts_forward(self, "moe_fused_grouped", hidden_states, top_k_index, top_k_weights) - kernel = load_finegrained_kernel() - num_top_k = top_k_index.size(-1) - - # one on-device routing pass (counting sort, no host sync): expert-sorted row starts + - # the gather/scatter maps. EP sentinels (id >= num_experts) fall past expert_start[-1]; - # their output rows stay uninitialized and the weighted_reduce bookend skips them. - expert_start, gather_idx, scatter_idx = kernel.compute_grouped_scheduling(top_k_index, self.num_experts, num_top_k) - - up_name = "gate_up_proj" if self.has_gate else "up_proj" - epilogue, gate_up_quantization, down_quantization = _kernel_objects(kernel, self) - - proj_out = kernel.grouped_matmul( - hidden_states, - to_local(getattr(self, up_name)), - None, - _proj_scale(self, up_name), - expert_start=expert_start, - gather_idx=gather_idx, - epilogue=epilogue, - bias=_proj_bias(self, up_name), - quantization=gate_up_quantization, - b_global_scale=to_local(getattr(self, f"{up_name}_global_scale")) if self.has_global_scale else None, - ) # fused+requant: (C, Cs); fused: (S, intermediate_dim) GLU intermediate; unfused: (S, 2*I) - - inter_scale = None - if isinstance(proj_out, (tuple, list)): - proj_out, inter_scale = proj_out # the fused requant's quantized intermediate - - if epilogue is None: - proj_out = _apply_unfused_gate_up(self, proj_out) - - proj_out = kernel.grouped_matmul( - proj_out, - to_local(self.down_proj), - inter_scale, # fused requant handoff: pre-quantized intermediate scales (None = raw) - _proj_scale(self, "down_proj"), - expert_start=expert_start, - scatter_idx=scatter_idx, # scatter straight to routed rows — weighted_reduce's layout - bias=_proj_bias(self, "down_proj"), - quantization=None if inter_scale is not None else down_quantization, - b_global_scale=to_local(self.down_proj_global_scale) if self.has_global_scale else None, - ) # (S, hidden_dim), ROUTED order after the scatter - - return _finish_down(self, kernel, proj_out, top_k_index, top_k_weights, hidden_states.dtype) - - -class FineGrainedExperts(nn.Module): - # Internal, temporary flag — not public API, don't set it directly. `_disable_deepgemm_on_multi_device` - # flips it True at load when the model spans >1 CUDA device in one process (DeepGEMM's context-bound - # kernels corrupt across devices); removable once the kernel ships a context-free loader. - _deepgemm_disabled = False +class FineGrainedExperts(_FineGrainedModule, nn.Module): # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. # The plan dicts store `{module-path-pattern: parallel-layer-kind}`; this maps an # old kind to a new kind, and the quantizer rewrites every plan VALUE that matches. @@ -834,15 +694,19 @@ def __init__( weight_format: str | None = None, activation_format: str | None = None, has_global_scale: bool | None = None, + is_concatenated: bool = True, ): super().__init__() self.config = config self.has_bias = has_bias self.has_gate = has_gate + # the model's own gate|up row order (transformers' experts flag): stacked ``[gate; up]`` or, + # like GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]`` — decides whether the + # loader interleaves + self.is_concatenated = is_concatenated self.block_size = block_size self.hidden_dim = config.hidden_size - self.has_global_scale = has_global_scale self.activation_format = activation_format self.activation_scheme = activation_scheme self.swiglu_alpha = getattr(config, "swiglu_alpha", None) @@ -850,7 +714,6 @@ def __init__( self.num_experts = _first_attr(config, "num_local_experts", "num_experts") self.intermediate_dim = _first_attr(config, "moe_intermediate_size", "intermediate_size") self.act_fn_name = _first_attr(config, "hidden_activation", "hidden_act") - self.limit = getattr(config, "swiglu_limit", None) self.act_fn = ACT2FN[self.act_fn_name] # Expert weight storage is declared by the QUANT config's format key, not the model @@ -861,58 +724,41 @@ def __init__( if weight_format is None: weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" self.weight_format = weight_format - fmt, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) - if has_global_scale is None: - has_global_scale = fmt.has_global_scale - # the forwards gate on the ATTRIBUTE — it must carry the format-resolved value, not - # the raw ctor arg (None would silently drop the nvfp4 global at every forward) - self.has_global_scale = has_global_scale - alloc_kwargs = { - "weight_dtype": fmt.weight_dtype, - "sf_dtype": sf_dtype, - "weight_k_div": fmt.values_per_byte, - "sf_gran_n": scale_group[0] if scale_group is not None else None, - "sf_gran_k": scale_group[1] if scale_group is not None else None, - } - - if self.has_gate: - self.gate_up_proj, self.gate_up_proj_scale_inv = _alloc_expert_proj( - self.num_experts, 2 * self.intermediate_dim, self.hidden_dim, min_sf_out=2, **alloc_kwargs - ) - if self.has_bias: - self.gate_up_proj_bias = nn.Parameter( - torch.empty(self.num_experts, 2 * self.intermediate_dim, dtype=torch.float32) - ) - else: - self.register_parameter("gate_up_proj_bias", None) - else: - self.up_proj, self.up_proj_scale_inv = _alloc_expert_proj( - self.num_experts, self.intermediate_dim, self.hidden_dim, **alloc_kwargs + fmt, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + # the forwards gate on this attribute, so it carries the format-resolved value (a None ctor + # arg would silently drop the nvfp4 global at every forward) + self.has_global_scale = fmt.has_global_scale if has_global_scale is None else has_global_scale + + # what the module HOLDS, as the loader's conversion ops deliver it (and their reverses + # restore for the checkpoint): gate|up rows in the kernels' interleaved order unless the + # backend packs gate|up itself (DeepGEMM Mega MoE), block scales swizzled where the + # scaled-MMA reads them + impl = getattr(config, "_experts_implementation", None) + self._gate_up_interleaved = self.has_gate and impl != "deepgemm_megamoe" + swizzled = _swizzles_scales(config, fmt, activation_format) + + up_name = "gate_up_proj" if self.has_gate else "up_proj" + up_rows = (2 if self.has_gate else 1) * self.intermediate_dim + for proj, rows, in_dim, min_scale_out in ( + (up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1), + ("down_proj", self.hidden_dim, self.intermediate_dim, 1), + ): + weight, scale = _alloc_expert_proj( + self.num_experts, rows, in_dim, fmt, scale_dtype, scale_group, min_scale_out, swizzled ) + setattr(self, proj, weight) + setattr(self, f"{proj}_scale_inv", scale) if self.has_bias: - self.up_proj_bias = nn.Parameter( - torch.empty(self.num_experts, self.intermediate_dim, dtype=torch.float32) - ) + # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it + # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype + setattr(self, f"{proj}_bias", nn.Parameter(torch.empty(self.num_experts, rows))) else: - self.register_parameter("up_proj_bias", None) - - self.down_proj, self.down_proj_scale_inv = _alloc_expert_proj( - self.num_experts, self.hidden_dim, self.intermediate_dim, **alloc_kwargs - ) - # NVFP4 two-level: per-expert fp32 globals the kernels recover on the accumulator - for proj in (("gate_up_proj" if self.has_gate else "up_proj"), "down_proj"): + self.register_parameter(f"{proj}_bias", None) + # NVFP4 two-level: per-expert fp32 globals the kernels recover on the accumulator if self.has_global_scale: - setattr( - self, - f"{proj}_global_scale", - nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)), - ) + setattr(self, f"{proj}_global_scale", nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32))) else: self.register_parameter(f"{proj}_global_scale", None) - if self.has_bias: - self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, dtype=torch.float32)) - else: - self.register_parameter("down_proj_bias", None) if self.activation_scheme == "static": self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) @@ -928,9 +774,9 @@ def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) glu = gate * torch.sigmoid(gate * self.swiglu_alpha) return (up + 1.0) * glu - elif self.limit is not None: - gate = gate.clamp(max=self.limit) - up = up.clamp(min=-self.limit, max=self.limit) + elif self.swiglu_limit is not None: + gate = gate.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) return self.act_fn(gate) * up def forward( @@ -956,42 +802,43 @@ def forward( ) proj_out = self.linear( current_state, - self.gate_up_proj[expert_idx] if self.has_gate else self.up_proj[expert_idx], - self.gate_up_proj_scale_inv[expert_idx] if self.has_gate else self.up_proj_scale_inv[expert_idx], + "gate_up_proj" if self.has_gate else "up_proj", + expert_idx, activation_scale=gate_up_act_scale, ) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) down_act_scale = ( self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None ) - proj_out = self.linear( - proj_out, - self.down_proj[expert_idx], - self.down_proj_scale_inv[expert_idx], - activation_scale=down_act_scale, - ) + proj_out = self.linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) routing_weights = top_k_weights[token_idx, top_k_pos, None] weighted_out = proj_out * routing_weights.to(proj_out.dtype) final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) return final_hidden_states.to(hidden_states.dtype) def linear( - self, - input: torch.Tensor, - weight: torch.Tensor, - weight_scale_inv: torch.Tensor, - activation_scale: torch.Tensor | None = None, + self, input: torch.Tensor, proj: str, expert_idx, activation_scale: torch.Tensor | None = None ) -> torch.Tensor: + """One expert's ``proj`` as a dense linear: the same weight, scale (one expert's slice of a + swizzled stack is the ``(1, ...)`` artifact ``matmul_2d`` reads), bias, NVFP4 global and + activation format the fused forwards hand the kernels.""" + weight = getattr(self, proj)[expert_idx] + bias = getattr(self, f"{proj}_bias") + bias = bias[expert_idx] if bias is not None else None if weight.element_size() > 1: - return F.linear(input, weight, None) - + return F.linear(input, weight, bias) + scale = getattr(self, f"{proj}_scale_inv") + global_scale = getattr(self, f"{proj}_global_scale") return finegrained_linear( input, weight, - weight_scale_inv, + scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx], self.block_size, + bias=bias, activation_scale=activation_scale, allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=global_scale[expert_idx] if global_scale is not None else None, + activation_format=self.activation_format, ) @@ -999,9 +846,9 @@ class FineGrainedExpertsInterface(ExpertsInterface): """Interface for registering custom FP8 experts forward functions.""" _global_mapping = { + "deepgemm": deepgemm_fp8_fp4_experts_forward, "batched_mm": finegrained_batched_mm_experts_forward, "grouped_mm": finegrained_grouped_mm_experts_forward, - "deepgemm": deepgemm_fp8_fp4_experts_forward, "deepgemm_megamoe": deepgemm_fp8_fp4_megamoe_experts_forward, } @@ -1009,114 +856,17 @@ class FineGrainedExpertsInterface(ExpertsInterface): ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() -def interleave_gate_up_after_loading(model: nn.Module, *, already_interleaved: bool = False) -> None: - """Put gate_up into the ``[g0, u0, g1, u1, ...]`` row order the kernels read. - - Checkpoints that ship the two projections separately arrive stacked ``[gate; up]`` (the shared - ``MergeModulelist`` + ``Concatenate`` chain in ``conversion_mapping``, which serves unquantized - MoE models too and so cannot carry a finegrained-specific reorder). GPT-OSS-style MXFP4 already - ships interleaved — ``already_interleaved`` skips those. - - Weight, block-scale grid and bias all key off the same output-row axis, so all three move - together. Runs BEFORE ``swizzle_scales_after_loading``, which then swizzles the final order. - """ - if already_interleaved: - return - # Backends that pack gate|up themselves (DeepGEMM Mega MoE) take the stacked form, so skip - # them rather than interleaving and having them undo it. `_gate_up_interleaved` records what - # the module actually holds, since `set_experts_implementation` can switch backends after - # load — a consumer must read the flag, not re-derive the backend. - impl = getattr(getattr(model, "config", None), "_experts_implementation", None) - packs_gate_up_itself = impl in ("deepgemm_megamoe",) - for module in model.modules(): - if not isinstance(module, FineGrainedExperts) or not module.has_gate: - continue - if packs_gate_up_itself: - module._gate_up_interleaved = False - continue - module._gate_up_interleaved = True - for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias"): - t = getattr(module, name, None) - if t is None or t.device.type == "meta": - continue - local = to_local(t) - axis = -1 if local.ndim == 1 or name.endswith("bias") else -2 - n = local.shape[axis] // 2 - # The reorder is byte-level, so run it on a uint8 view for the 1-byte float8 - # scale/weight dtypes: `torch.stack` lowers to `cat`, and "cat_cuda" is not - # implemented for float8_e8m0fnu (nor e4m3) — it raises mid-load under TP/EP. - byte_view = local.element_size() == 1 and local.dtype.is_floating_point - src = local.view(torch.uint8) if byte_view else local - if axis == -1: - rows = torch.stack([src[..., :n], src[..., n:]], dim=-1) - else: - rows = torch.stack([src[..., :n, :], src[..., n:, :]], dim=-2) - src.copy_(rows.reshape(src.shape)) - - -def swizzle_scales_after_loading(model: nn.Module) -> None: - """Pre-swizzle MX/NVFP4 weight block scales into the ``SWIZZLE_32_4_4`` layout the Blackwell - tcgen05 scaled-MMA reads directly (one triton launch per matrix, values unchanged) — the - kernels accept both layouts, but plain row-major forces a per-tile gather that caps the - scaled dot below the fp8/fp4 peak. Runs once post-load; the swizzled artifact is cached as - a BUFFER next to the affine scale Parameter (which stays the state_dict source of truth, so - save/dequantize round-trip untouched; scale tensors are small). A gate_up grid swizzles as one - matrix over its 2N rows. - - SM100 only. The layout exists for the Blackwell tcgen05 scaled-MMA; elsewhere it buys - nothing, costs the swizzled arm's tile constraints (``swizzled_scales_bm_pruner`` pins - BM=BN=128), and collides with DeepGEMM — which stays the preferred linear backend below - SM100 and reads plain row-major scales. - ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` skips the pass (debug / A/B).""" - if os.environ.get("TRANSFORMERS_FINEGRAINED_NO_SWIZZLE", "0") == "1": - return - if not is_sm100(): - return - try: - kernel = load_finegrained_kernel() - except ImportError: - return # dequantize fallback path: nothing to swizzle for - formats = _weight_formats() - for module in model.modules(): - if not isinstance(module, FineGrainedExperts): - continue - # Only the GROUP-scaled families have a layout to pre-arrange: SWIZZLE_32_4_4 is the - # scale OPERAND layout the tcgen05 scaled-MMA reads (MX group-32, NVFP4 group-16). - # Block-FP8's (N/128, K/128) grid never reaches a scaled-MMA — the block-dynamic loop - # applies it on the fp32 accumulator around a plain dot — so it has no swizzled form. - # Keyed off the declared recipe, not the scale dtype: V4-style block-FP8 ships UE8M0 - # scales, so a `dtype == float32` test lets it through (it then survives only by - # accident, because N/128 is rarely a multiple of 128). - fmt = formats.get(module.weight_format) - if fmt is None or fmt.scale_group is None: - continue - for proj in (("gate_up_proj" if module.has_gate else "up_proj"), "down_proj"): - scale = getattr(module, f"{proj}_scale_inv", None) - if scale is None or scale.device.type == "meta": - continue - n_rows = to_local(scale).shape[-2] - # the swizzled layout needs whole 128-row blocks; for gate_up that is the 2N extent, - # i.e. N % 64 (GPT-OSS N=2880 qualifies). The kernels read the affine layout directly, - # so a non-conforming projection just skips the pre-swizzle. - if n_rows % 128: - continue - swizzled = kernel.swizzle_mx_scales(to_local(scale).data) - module.register_buffer(f"{proj}_scale_inv_swizzled", swizzled, persistent=False) - - -def _disable_deepgemm_on_multi_device(model: nn.Module) -> None: - """Internal, temporary helper (not public API): flag every FP8 module to skip DeepGEMM when the - model spans >1 CUDA device in one process. - - DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` handle to the - CUDA context live at load time; driving that cached handle from another device launches it against - the wrong context and produces garbage. (Build-time fix: compile DeepGEMM with - `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until our wheel picks that up - we avoid single-process multi-device.) Setting `_deepgemm_disabled` routes both the linear and - experts paths through Triton/grouped_mm. A model that fits on one device keeps DeepGEMM even with - other GPUs visible; TP/EP put one device per process, so this is a no-op there. - """ - quantized_modules = [m for m in model.modules() if isinstance(m, (FineGrainedLinear, FineGrainedExperts))] +def disable_deepgemm_on_multi_device(model: nn.Module) -> None: + """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one + process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` + handle to the CUDA context live at load time; driving that cached handle from another device + launches it against the wrong context and produces garbage. (Build-time fix: compile DeepGEMM + with `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until our wheel picks + that up we avoid single-process multi-device.) Setting `_deepgemm_disabled` routes both the + linear and experts paths through Triton/grouped_mm. A model that fits on one device keeps + DeepGEMM even with other GPUs visible; TP/EP put one device per process, so this is a no-op + there.""" + quantized_modules = [m for m in model.modules() if isinstance(m, _FineGrainedModule)] cuda_devices = set() for m in quantized_modules: param = next(m.parameters(), None) @@ -1134,22 +884,26 @@ def _disable_deepgemm_on_multi_device(model: nn.Module) -> None: ) -def replace_with_finegrained_layer( - model, modules_to_not_convert: list[str] | None = None, quantization_config=None, pre_quantized=False -): - """ - A helper function to replace all `torch.nn.Linear` modules by `FineGrainedLinear` modules. - - Parameters: - model (`torch.nn.Module`): - Input model or `torch.nn.Module` as the function is run recursively. - modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`): - Names of the modules to not convert. In practice we keep the `lm_head` in full precision for numerical stability reasons. - quantization_config (`FineGrainedConfig`): - The quantization config object that contains the quantization parameters. - pre_quantized (`bool`, defaults to `False`): - Whether the model is pre-quantized or not - """ +def replace_with_finegrained_embedding(model, patterns: list[str], modules_to_not_convert: list[str] | None = None): + """Swap every ``nn.Embedding`` whose name ends with one of ``patterns`` (the quant config's + ``modules_to_convert``) for ``FineGrainedEmbedding``. Runs under ``dequantize=True`` too: a table + worth quantizing is one we cannot afford to expand.""" + for name, module in list(model.named_modules()): + if type(module) is nn.Embedding and any(name.endswith(p) for p in patterns): + if not should_convert_module(name, modules_to_not_convert): + continue + with torch.device("meta"): + new_module = FineGrainedEmbedding(module.num_embeddings, module.embedding_dim, module.padding_idx) + new_module._hf_quantized_needs_local_tp = True + model.set_submodule(name, new_module) + return model + + +def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | None = None, quantization_config=None): + """Swap the model's ``nn.Linear`` (and grouped-linear) modules for ``FineGrainedLinear`` and its + ``.experts`` modules for ``FineGrainedExperts``, both allocated in the quantized storage the + checkpoint's format declares. ``modules_to_not_convert`` names the modules kept as they are + (typically ``lm_head``); nothing is replaced under ``dequantize=True``.""" if quantization_config.dequantize: return model @@ -1174,12 +928,16 @@ def replace_with_finegrained_layer( if module_name.endswith(".experts"): has_gate = getattr(module, "has_gate", True) has_bias = getattr(module, "has_bias", False) + is_concatenated = getattr(module, "is_concatenated", True) config = getattr(module, "config", model.config.get_text_config()) + # the decorator stamps these flags on the instance after __init__, so the model's + # own gate|up convention has to travel through it too new_class = use_experts_implementation( experts_class=FineGrainedExperts, experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, has_bias=has_bias, has_gate=has_gate, + is_concatenated=is_concatenated, ) new_module = new_class( config=config, @@ -1190,6 +948,7 @@ def replace_with_finegrained_layer( has_gate=has_gate, weight_format=weight_format, activation_format=getattr(quantization_config, "activation_format", None), + is_concatenated=is_concatenated, ) # GPT-OSS hardcodes its clamped-SwiGLU parameters as MODULE attributes # (alpha=1.702, limit=7.0) rather than config fields — carry them over so @@ -1237,36 +996,170 @@ def replace_with_finegrained_layer( return model -class FineGrainedMxfp4Deserialize(ConversionOps): - """Convert GPT-OSS MXFP4 checkpoint tensors to the finegrained layout: ``{proj}_blocks`` - ``(E, N, K/32, 16)`` uint8 (two low-nibble-first E2M1 values per byte) reshapes to the - packed ``(E, N, K/2)`` int8 the kernels read directly — same nibble order — and - ``{proj}_scales`` (biased-127 exponent bytes) bitcast to ``float8_e8m0fnu``. No swizzle - (the kernels take affine scales; the post-load hook builds the swizzled cache). GPT-OSS ships - gate_up in the interleaved [g0, u0, ...] row order the kernels read, so rows pass through.""" +# --------------------------------------------------------------------------------------------------- +# Conversion ops. The loader runs a converter's ops on the rank-local tensors it collected for one +# target, and `save_pretrained` runs their reverses on the gathered tensors, so everything below +# maps a checkpoint layout onto the layout the modules above hold — and back. +# --------------------------------------------------------------------------------------------------- - def __init__(self, hf_quantizer): + +def _keyed_by_target(input_dict: dict, target_patterns) -> dict: + """A converter's ops see their single tensor under the SOURCE pattern until an op re-keys it to the + target (the core ops do; the loader then expands the target into the full name) — so a + layout op that is first in its converter re-keys the same way. Multi-tensor dicts (a + deserializer's fully named outputs) pass through.""" + if target_patterns and len(input_dict) == 1: + value = next(iter(input_dict.values())) + return {target_patterns[0]: value[0] if isinstance(value, list) else value} + return input_dict + + +def _held_scale(model, full_layer_name, key): + """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else + ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's + name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" + param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[2] + if not param_name.endswith("_scale_inv") or model is None: + return None, None + module = model.get_submodule(full_layer_name.rpartition(".")[0]) + if not isinstance(module, _FineGrainedModule): + return None, None + return module, getattr(module, param_name, None) + + +def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" + return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) + + +def _checkpoint_scale_container(model) -> torch.dtype | None: + """The scale container the checkpoint shipped, as recorded on its finegrained modules at load + (one checkpoint, one container); None when it was already the held dtype.""" + if model is None: + return None + return next((m.scale_container_dtype for m in model.modules() if getattr(m, "scale_container_dtype", None)), None) + + +class FineGrainedInterleaveGateUp(ConversionOps): + """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order (weight, + scale grid and bias share the row axis) — the core ``Interleave`` along dim 1. Decided off the + model's experts modules, on load and on save alike: skipped when the model's own rows already + come interleaved (``is_concatenated=False``, GPT-OSS) or the backend holds them stacked + (``_gate_up_interleaved`` False, DeepGEMM Mega MoE).""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer + self.inverse = inverse - def convert( - self, - input_dict: dict[str, torch.Tensor], - model: torch.nn.Module | None = None, - full_layer_name: str | None = None, - missing_keys: list[str] | None = None, - **kwargs, - ) -> dict[str, torch.Tensor]: - proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj" - blocks = input_dict[f"{proj}_blocks"] - scales = input_dict[f"{proj}_scales"] - blocks = blocks[0] if isinstance(blocks, list) else blocks - scales = scales[0] if isinstance(scales, list) else scales + def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + input_dict = _keyed_by_target(input_dict, target_patterns) + experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None + if experts is None or not experts.is_concatenated or not experts._gate_up_interleaved: + return input_dict + interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved + return {key: interleave.convert({key: value}, None, [key])[key] for key, value in input_dict.items()} + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInterleaveGateUp(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedScaleContainer(ConversionOps): + """A block scale into the dtype its module holds. Two containers exist for UE8M0: the same + exponent byte under ``uint8`` (MiniMax-style mxfp8 checkpoints) — a bit reinterpretation into + ``float8_e8m0fnu`` — and the power-of-two value in a float32 tensor (dsv4-flash-base) — an exact + numeric cast. Other tensors, and scales already in the held dtype, pass through. The module + records the checkpoint's container (``scale_container_dtype``), so the reverse restores it + exactly on save.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + if self.inverse: + container = _checkpoint_scale_container(model) + # e8m0 is only ever a scale, so the dtype alone picks the tensors to restore + if container is not None and value.dtype == _get_ue8m0_dtype(): + value = _recontain(value, container) + else: + module, held = _held_scale(model, full_layer_name, key) + if held is not None and value.dtype != held.dtype: + module.scale_container_dtype = value.dtype + value = _recontain(value, held.dtype) + out[key] = value + return out + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedScaleContainer(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedSwizzleScales(ConversionOps): + """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D scale + Parameter, see ``_swizzles_scales``): one triton launch per rank-local shard, values unchanged. + Every other tensor the converter carries (weights, biases, the scales of a module kept affine) + passes through. The reverse restores the checkpoint's affine grid from a 5-D tensor's own shape + (whole 128-row blocks and 4-column groups, so no padding to trim), so saving is exact.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) + return out + + @staticmethod + def _unswizzle(value): + if value.ndim != 5: + return value + experts, row_blocks, col_groups = value.shape[:3] + kernel = load_finegrained_kernel() + return kernel.unswizzle_mx_scales(value, row_blocks * 128, col_groups * 4, num_experts=experts) + + @staticmethod + def _swizzle(key, value, model, full_layer_name): + _, held = _held_scale(model, full_layer_name, key) + if held is None or held.ndim != 5 or value.ndim == 5: + return value + return load_finegrained_kernel().swizzle_mx_scales(value) - num_experts, rows = blocks.shape[0], blocks.shape[1] - weight = blocks.reshape(num_experts, rows, -1).view(torch.int8) - scale_inv = scales.contiguous().view(_get_ue8m0_dtype()) - prefix = full_layer_name.rsplit(".", 1)[0] - return {f"{prefix}.{proj}": weight.contiguous(), f"{prefix}.{proj}_scale_inv": scale_inv} + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedSwizzleScales(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedPackedBlocks(ConversionOps): + """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 — two low-nibble-first E2M1 values per + byte, 16 bytes per group of 32 — as the packed ``(E, N, K/2)`` int8 the kernels read: the same + bytes, grouped per row. The reverse folds the rows back into 16-byte groups under uint8. The + ``{proj}_scales`` companion goes through ``FineGrainedScaleContainer`` (exponent bytes -> + ``float8_e8m0fnu``) in its own converter.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + if self.inverse: + out[key] = value.view(torch.uint8).reshape(*value.shape[:-1], value.shape[-1] // 16, 16) + else: + out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) + return out + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedPackedBlocks(self.hf_quantizer, inverse=not self.inverse) class FineGrainedViewPackedInt8(ConversionOps): @@ -1278,12 +1171,12 @@ class FineGrainedViewPackedInt8(ConversionOps): def __init__(self, hf_quantizer=None): self.hf_quantizer = hf_quantizer - def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): - v = next(iter(input_dict.values())) - v = v[0] if isinstance(v, list) else v - if torch.is_tensor(v) and v.dtype == torch.uint8: - v = v.view(torch.int8) - return {full_layer_name: v} + def convert(self, input_dict, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + out[key] = value.view(torch.int8) if torch.is_tensor(value) and value.dtype == torch.uint8 else value + return out class FineGrainedFuseEqualGlobals(ConversionOps): @@ -1370,10 +1263,9 @@ def _quantize_one(self, key: str, value: torch.Tensor) -> dict[str, torch.Tensor scale_key = key.rsplit(".", 1)[0] + ".weight_scale_inv" if key.endswith(".weight") else key + "_scale_inv" return {key: quantized, scale_key: inv_scales} - def convert(self, input_dict: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]: - # Quantize every (key, tensor) entry in the dict. Single-tensor case (legacy - # callers that pass one key) and multi-tensor case (reverse of an expert - # ``MergeModulelist`` that emits one key per expert) are handled the same way. + def convert(self, input_dict: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]: + # every (key, tensor) entry — one key, or one per expert from the reverse of a + # ``MergeModulelist`` — quantizes the same way result: dict[str, torch.Tensor] = {} for key, value in input_dict.items(): tensor = value[0] if isinstance(value, list) else value @@ -1492,14 +1384,11 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor] | torch.Tensor]: output_dtype = self._get_target_dtype(model, full_layer_name) - # Backward-compatible single-tensor path (the legacy fallback converter declares - # ``["weight$", "weight_scale_inv", "activation_scale"]`` and produces a single - # ``weight`` target). Also handles the no-scale case (e.g. RMSNorm weights that - # match ``weight$`` but ship no ``weight_scale_inv`` alongside). + # The dense converter (``["weight$", "weight_scale_inv", "activation_scale"] -> weight``) + # hands one weight, with or without a scale (RMSNorm weights match ``weight$`` too). if "weight$" in input_dict: - # The downstream renamer in `core_model_loading._convert_one_module` uses the - # output dict's *key*, not its content, to derive prefix/suffix; if `full_layer_name` - # is unset (direct invocation / tests) fall back to the legacy converter's target. + # the loader derives prefix/suffix from the output KEY; without a `full_layer_name` + # (direct invocation) key it by the converter's target target_key = full_layer_name if full_layer_name is not None else "weight" quantized = input_dict["weight$"] quantized = quantized[0] if isinstance(quantized, list) else quantized diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 84e73f35927c..94667bad4344 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -1,3 +1,4 @@ +import re from typing import TYPE_CHECKING from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging @@ -24,9 +25,6 @@ class FineGrainedHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedFP8Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") @@ -120,7 +118,7 @@ def _process_model_before_weight_loading( model: "PreTrainedModel", **kwargs, ): - from ..integrations.finegrained import replace_with_finegrained_layer + from ..integrations.finegrained import replace_with_finegrained_embedding, replace_with_finegrained_layer self._normalize_modules_to_not_convert(model) if self._quant_method() == "mxfp4" and getattr(self.quantization_config, "activation_format", None) is None: @@ -131,47 +129,18 @@ def _process_model_before_weight_loading( model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules ) - model = replace_with_finegrained_layer( - model, - modules_to_not_convert=self.modules_to_not_convert, - quantization_config=self.quantization_config, - pre_quantized=self.pre_quantized, + if self.pre_quantized and self.quantization_config.modules_to_convert: + replace_with_finegrained_embedding( + model, self.quantization_config.modules_to_convert, self.modules_to_not_convert + ) + replace_with_finegrained_layer( + model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config ) def _process_model_after_weight_loading(self, model, **kwargs): - # dsv4-flash-base stores its (power-of-two) ue8m0 scales in a float32 container under - # `.scale`; those renamed keys keep the on-disk float32 dtype, so cast them to the UE8M0 - # dtype the kernels expect (exact, since the values are powers of two). Checkpoints that - # already ship the native float8 E8M0 dtype (e.g. dsv4-flash) are left untouched. - if self.quantization_config.scale_fmt == "ue8m0": - from ..integrations.finegrained import _get_ue8m0_dtype - - ue8m0 = _get_ue8m0_dtype() - float32_scales = [ - name - for name, param in model.named_parameters() - if name.endswith("_scale_inv") and param.dtype == torch.float32 - ] - for name in float32_scales: - module_name, _, attr = name.rpartition(".") - module = model.get_submodule(module_name) - scale = getattr(module, attr) - setattr(module, attr, torch.nn.Parameter(scale.data.to(ue8m0), requires_grad=False)) - - # Single-process multi-device is unsafe for DeepGEMM (its kernels are bound to one CUDA - # context); route those models through Triton/grouped_mm instead. - from ..integrations.finegrained import _disable_deepgemm_on_multi_device - - _disable_deepgemm_on_multi_device(model) + from ..integrations.finegrained import disable_deepgemm_on_multi_device - from ..integrations.finegrained import ( - interleave_gate_up_after_loading, - swizzle_scales_after_loading, - ) - - # order matters: the swizzle below packs whatever row order gate_up ends up in - interleave_gate_up_after_loading(model, already_interleaved=self._quant_method() == "mxfp4") - swizzle_scales_after_loading(model) + disable_deepgemm_on_multi_device(model) return model def update_tp_plan(self, config): @@ -236,33 +205,38 @@ def get_weight_conversions(self): from ..core_model_loading import WeightConverter from ..integrations.finegrained import ( FineGrainedDequantize, - FineGrainedMxfp4Deserialize, + FineGrainedInterleaveGateUp, + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, ) - # GPT-OSS-style MXFP4 checkpoints ship {proj}_blocks + {proj}_scales; deserialize them - # into the packed weight + e8m0 scale pair the finegrained kernels read. The gate_up bias - # needs no converter — it already follows the interleaved row order. + # {proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed + # weight, the exponent bytes take the scale container op; then the usual layout ops. Bare + # patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`. if self.pre_quantized and self._quant_method() == "mxfp4" and not self.quantization_config.dequantize: - return [ - WeightConverter( - source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], - target_patterns=r"gate_up_proj$", - operations=[FineGrainedMxfp4Deserialize(self)], - ), - WeightConverter( - source_patterns=["down_proj_blocks", "down_proj_scales"], - target_patterns=r"down_proj$", - operations=[FineGrainedMxfp4Deserialize(self)], - ), - ] - - # NVIDIA modelopt NVFP4 checkpoints: per-expert split gate/up projections with - # packed-uint8 weights, E4M3 group-16 `weight_scale`, and fp32 `weight_scale_2` - # second-level globals (bit-identical across the gate/up halves — asserted by the - # fuse op). The arch conversion plan already stacks the expert WEIGHTS; these add - # the same stacking for scales and globals, plus the uint8 -> int8 packed bitcast. - # Calibrated `input_scale` entries are left unconsumed (activations run the - # kernels' dynamic quant); they surface as unexpected keys, which is intended. + converters = [] + for proj in ("gate_up_proj", "down_proj"): + interleave = [FineGrainedInterleaveGateUp(self)] if proj == "gate_up_proj" else [] + converters += [ + WeightConverter( + source_patterns=rf"{proj}_blocks$", + target_patterns=proj, + operations=[FineGrainedPackedBlocks(self), *interleave, FineGrainedSwizzleScales(self)], + ), + WeightConverter( + source_patterns=rf"{proj}_scales$", + target_patterns=f"{proj}_scale_inv", + operations=[*interleave, FineGrainedScaleContainer(self), FineGrainedSwizzleScales(self)], + ), + ] + return converters + + # modelopt NVFP4 (per-expert gate/up, E4M3 group-16 `weight_scale`, fp32 `weight_scale_2` + # globals): the arch plan stacks the expert weights; these stack the scales and globals the + # same way. Only the routed experts are quantized, so no dense converters — a generic + # `weight_scale*` rename would run before converter collection and mangle these keys. + # Calibrated `input_scale` keys stay unconsumed (dynamic act quant) and surface as unexpected. if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: from ..core_model_loading import Concatenate, MergeModulelist from ..integrations.finegrained import FineGrainedFuseEqualGlobals @@ -299,19 +273,11 @@ def get_weight_conversions(self): operations=[MergeModulelist(dim=0), FineGrainedFuseEqualGlobals(self)], ), ] - # No converters for non-expert modules: this checkpoint family quantizes ONLY - # the routed experts (every attention / shared-expert / dense-MLP module sits - # in the ignore list and ships bf16). Generic weight_scale* renames here would - # run BEFORE converter collection and mangle the expert keys out from under - # the converters above (per-expert `weight_global_scale` orphans and - # `_scale_inv_inv` double-suffix targets — the saturation bug of round 8). - # Quantized-dense modelopt checkpoints need scoped per-module converters when - # one shows up. if self.pre_quantized and self.quantization_config.dequantize: return [ - # either use the dollar sign, or permute the source patterns to start matching against the scales first - # We also collect the activation scales, they will not be used + # anchored `weight$` so the scale keys land in the scale slots; the activation + # scales are collected only to be dropped WeightConverter( source_patterns=["weight$", "weight_scale_inv", "activation_scale"], target_patterns="weight", @@ -321,77 +287,119 @@ def get_weight_conversions(self): return [] def update_weight_conversions(self, weight_conversions): - """When loading with ``dequantize=True``, attach an :class:`FineGrainedDequantize` op to - every existing :class:`WeightConverter` so that per-block scales are folded into - the weight *before* any later merge/concat ops collapse the per-expert structure. - - For each model-supplied converter that has a ``.weight`` source, we: - 1. anchor the existing weight patterns with ``$`` so they don't accidentally - also match the ``.weight_scale_inv`` keys (the regex is searched, so the - unanchored prefix would match both, sending scales to the wrong bucket); - 2. add anchored ``*.weight_scale_inv`` sources next to each weight pattern so - the loader collects scale tensors alongside the weight tensors into the - *same* converter bucket (both keys rewrite to the same target); - 3. prepend a fresh :class:`FineGrainedDequantize` op so dequant runs first, before - any merge/concat collapses the per-expert structure. - - The generic ``weight$ + weight_scale_inv → weight`` converter from - :meth:`get_weight_conversions` is still appended at the end as a fallback for - plain ``nn.Linear`` weights with no model-specific converter. - """ + """Rewrite the model's conversion plan for this checkpoint: + + - every mode: ``*.scale`` -> ``*.weight_scale_inv`` (DeepSeek-V4-Flash ships per-block + scales under ``.scale``; kept here rather than in each model's mapping so non-FP8 + round-trips never see the rule); + - ``dequantize=True``: a :class:`FineGrainedDequantize` runs first in every converter with a + ``.weight`` source (anchored, with the sibling ``.weight_scale_inv`` collected alongside), + so per-expert (weight, scale) pairs fold into full-precision tensors before merge/concat + ops collapse the per-expert structure; + - every mode: sharded ``.weight`` converter targets are anchored, since they become source + patterns on save and would otherwise match their scale keys; + - modelopt: the arch plan's ``.weight`` sources are anchored so they don't swallow the + ``weight_scale`` / ``weight_scale_2`` keys, and the packed weights get the uint8 -> int8 view; + - otherwise the expert converters get the module layout ops (``_with_expert_layout_ops``). + + :meth:`get_weight_conversions` is appended in every mode.""" from ..core_model_loading import WeightConverter, WeightRenaming - from ..integrations.finegrained import FineGrainedDequantize + from ..integrations.finegrained import FineGrainedDequantize, FineGrainedViewPackedInt8 - # `*.scale` → `*.weight_scale_inv`. Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) - # ship per-block scales under `.scale`; the model expects `.weight_scale_inv`. - # Lives here (not in each model's `conversion_mapping`) so non-FP8 round-trips - # don't see a stray rule. Needed in both dequantize modes — `dequantize=False` - # loads scales as parameters, `dequantize=True` feeds them into `FineGrainedDequantize`. scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") - weight_conversions = [scale_rename] + list(weight_conversions) + weight_conversions = [scale_rename, *weight_conversions] + # a converter's targets become source patterns on save: anchor sharded `.weight` targets + # (Qwen4-Exp's `ngram_embedding`) so they don't also match their scale keys then + for conv in weight_conversions: + if isinstance(conv, WeightConverter): + conv._original_target_patterns = [ + f"{p}$" if p.endswith(".weight") else p for p in conv._original_target_patterns + ] - if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: - # Anchor the arch plan's `.weight` sources so they don't also swallow the - # modelopt `weight_scale`/`weight_scale_2` keys into the weight merge chain - # (unanchored patterns are searched — same trap the dequantize path guards). - from ..integrations.finegrained import FineGrainedViewPackedInt8 + if self.pre_quantized and self.quantization_config.dequantize: + updated = [] + for conv in weight_conversions: + weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] + if isinstance(conv, WeightConverter) and weight_sources: + conv = WeightConverter( + source_patterns=[p + "$" for p in weight_sources] + + [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] + + [p for p in conv.source_patterns if not p.endswith(".weight")], + target_patterns=conv._original_target_patterns, + operations=[FineGrainedDequantize(self), *conv.operations], + ) + updated.append(conv) + return updated + self.get_weight_conversions() + if self.pre_quantized and self._quant_method() == "modelopt": updated = [] for conv in weight_conversions: if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): conv = WeightConverter( source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], target_patterns=conv.target_patterns, - # the stacked packed-fp4 expert weights need the same uint8 -> int8 - # bitcast the standalone-linear converter applies (pass-through on - # any unquantized weights the plan also routes here) - operations=list(conv.operations) + [FineGrainedViewPackedInt8(self)], + operations=[*conv.operations, FineGrainedViewPackedInt8(self)], ) updated.append(conv) - return updated + self.get_weight_conversions() - if not (self.pre_quantized and self.quantization_config.dequantize): - return weight_conversions + self.get_weight_conversions() + weight_conversions = updated + + return self._with_expert_layout_ops(weight_conversions + self.get_weight_conversions()) + + def _with_expert_layout_ops(self, weight_conversions): + """Give every converter that produces expert tensors the layout ``FineGrainedExperts`` holds + (see its ``__init__``), in chain order: gate|up rows interleaved (weight, scale grid and bias + share the row axis), the scale in the held container dtype, then the swizzled scale layout. + Each op decides from the model it runs against and each has a reverse, so saving restores + the checkpoint layout. Catch-all converters at the end take keys that arrive already under + the fused names (a plain rename, a transformers-format checkpoint) and dense linears' scale + keys, so they get the same treatment.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained import ( + FineGrainedInterleaveGateUp, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) - updated: list = [] + def layout_ops(target: str) -> list: + ops = [] + if "gate_up_proj" in target: + ops.append(FineGrainedInterleaveGateUp(self)) + if "_bias" not in target: + ops += [ + FineGrainedScaleContainer(self), + FineGrainedSwizzleScales(self), + ] # the swizzle reads 1-byte scales + return ops + + updated = [] for conv in weight_conversions: - # Only WeightConverter has ``.operations`` to extend with the dequant op; - # WeightRenaming (e.g. the ``scale_rename`` we prepended) just passes through. - if not isinstance(conv, WeightConverter): - updated.append(conv) - continue - weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] - if weight_sources: - anchored_weight = [p + "$" for p in weight_sources] - scale_sources = [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] - other = [p for p in conv.source_patterns if not p.endswith(".weight")] - new_sources = anchored_weight + scale_sources + other - new_ops = [FineGrainedDequantize(self)] + list(conv.operations) + targets = conv.target_patterns if isinstance(conv, WeightConverter) else [] + expert_targets = [ + t for t in targets if re.search(r"experts\.(gate_up_proj|down_proj)(_scale_inv|_bias)?\$?$", t) + ] + if expert_targets: conv = WeightConverter( - source_patterns=new_sources, - target_patterns=conv._original_target_patterns, - operations=new_ops, + source_patterns=conv.source_patterns, + target_patterns=conv.target_patterns, + operations=list(conv.operations) + layout_ops(expert_targets[0]), ) updated.append(conv) - # Generic fallback for plain ``nn.Linear`` weights with no model-specific converter. - updated.extend(self.get_weight_conversions()) + # dense linears have no converter: their scale keys take the container cast alone + updated.append( + WeightConverter( + source_patterns=r"weight_scale_inv$", + target_patterns="weight_scale_inv", + operations=[FineGrainedScaleContainer(self)], + ) + ) + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias", "down_proj_scale_inv"): + if layout_ops(name): + # the target is replacement text (op outputs are keyed by it), so no regex escapes + updated.append( + WeightConverter( + source_patterns=rf"experts\.{name}$", + target_patterns=f"experts.{name}", + operations=layout_ops(name), + ) + ) return updated diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index ec3c677ba9ac..fa6749c2b077 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -48,6 +48,9 @@ def op(*args, **kwargs): self.calls.setdefault(name, []).append(_Call(args, kwargs)) if name == "swizzle_mx_scales": return args[0] + if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) + hidden, down = args[0], kwargs["down_proj"] + return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) a, b = args[0], args[1] n = b.shape[-2] gather = kwargs.get("gather_idx") @@ -60,29 +63,17 @@ def op(*args, **kwargs): def _fake_bundle(): rec = _Recorder() kernel = mock.Mock() - for name in ("matmul_2d", "matmul_batched", "matmul_grouped", "nvfp4_quantize_two_level", "swizzle_mx_scales"): + for name in ( + "matmul_2d", + "matmul_batched", + "matmul_grouped", + "moe_fused_batched", + "moe_fused_grouped", + "swizzle_mx_scales", + "unswizzle_mx_scales", + ): setattr(kernel, name, rec._op(name)) - - def compute_grouped_scheduling(top_k_index, num_experts, num_top_k): - rec.calls.setdefault("compute_grouped_scheduling", []).append(_Call((top_k_index, num_experts, num_top_k), {})) - flat = top_k_index.reshape(-1) - sorted_ids, perm = torch.sort(flat) - counts = torch.bincount(sorted_ids.clamp(max=num_experts - 1), minlength=num_experts) - expert_start = torch.cat([counts.new_zeros(1), counts.cumsum(0)]).to(torch.int32) - gather_idx = (perm // num_top_k).to(torch.int32) - scatter_idx = torch.empty_like(perm) - scatter_idx[torch.arange(perm.numel())] = perm - return expert_start, gather_idx, scatter_idx.to(torch.int32) - - def weighted_reduce(rows, top_k_index, top_k_weights, num_experts, simulate_unfused=False): - rec.calls.setdefault("weighted_reduce", []).append(_Call((rows, top_k_index, top_k_weights, num_experts), {})) - num_tokens, num_top_k = top_k_index.shape - w = top_k_weights.reshape(-1, 1).to(rows.dtype) - w = w * (top_k_index.reshape(-1, 1) < num_experts) - return (rows * w).view(num_tokens, num_top_k, -1).sum(1) - - kernel.compute_grouped_scheduling = compute_grouped_scheduling - kernel.weighted_reduce = weighted_reduce + kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") @dataclass class Quantization: @@ -191,6 +182,36 @@ def test_module_forward_threads_everything(self): self.assertEqual(out.shape, (2, 32)) +@require_torch +class FineGrainedEmbeddingTest(unittest.TestCase): + """A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, + rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`.""" + + def test_lookup_rescales_the_gathered_rows(self): + torch.manual_seed(0) + table = torch.randn(64, 16) + scale = table.abs().max() / 448.0 + embedding = fg.FineGrainedEmbedding(64, 16) + embedding.weight = torch.nn.Parameter((table / scale).to(torch.float8_e4m3fn), requires_grad=False) + embedding.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16).reshape(1), requires_grad=False) + ids = torch.tensor([[3, 7], [60, 0]]) + out = embedding(ids) + self.assertEqual(out.dtype, torch.bfloat16) + expected = (table / scale).to(torch.float8_e4m3fn)[ids].to(torch.bfloat16) * scale.to(torch.bfloat16) + torch.testing.assert_close(out, expected) + + def test_replacement_follows_the_patterns_and_the_skip_list(self): + model = torch.nn.Module() + model.ngram_embedding = torch.nn.Embedding(8, 4) + model.embed_tokens = torch.nn.Embedding(8, 4) + model.other = torch.nn.Embedding(8, 4) + fg.replace_with_finegrained_embedding(model, ["ngram_embedding", "other"], modules_to_not_convert=["other"]) + self.assertIsInstance(model.ngram_embedding, fg.FineGrainedEmbedding) + self.assertTrue(model.ngram_embedding._hf_quantized_needs_local_tp) + self.assertIs(type(model.embed_tokens), torch.nn.Embedding) + self.assertIs(type(model.other), torch.nn.Embedding) + + @require_torch class FineGrainedExpertsMarshallingTest(unittest.TestCase): def _experts(self, **kw): @@ -212,32 +233,70 @@ def _route(self, tokens=6): wts = torch.rand(tokens, 2) return hs, idx, wts + def test_eager_linear_threads_bias_global_and_activation_format(self): + """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: + the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" + kernel, rec = _fake_bundle() + m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") + p1, p2, p3, p4 = _loaded(kernel) + linear = mock.Mock(wraps=fg.finegrained_linear) + with ( + p1, + p2, + p3, + p4, + mock.patch.object(fg, "is_deepgemm_loadable", return_value=False), + mock.patch.object(fg, "finegrained_linear", linear), + ): + m(*self._route()) + self.assertTrue(linear.call_args_list) + for call in linear.call_args_list: + self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) + self.assertEqual(call.kwargs["weight_global_scale"].ndim, 0) # this expert's NVFP4 global + self.assertEqual(call.kwargs["activation_format"], "bf16") + for call in rec.calls["matmul_2d"]: + self.assertIsNotNone(call.kwargs["b_global_scale"]) + self.assertIsNone(call.kwargs["quantization"].input_recipe) + def test_batched_marshalling(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) - up, down = rec.calls["matmul_batched"] - self.assertIsNone(up.args[2]) # As positional None - self.assertEqual(up.args[0].shape[0], 6) # UNEXPANDED tokens — the kernel gathers - self.assertEqual(up.kwargs["gather_idx"].shape, (12,)) - self.assertEqual(up.kwargs["expert_ids"].shape, (12,)) - self.assertNotIn("block_size", up.kwargs) - self.assertIsNone(up.kwargs["b_global_scale"]) + (call,) = rec.calls["moe_fused_batched"] + hidden, top_k_index, top_k_weights = call.args + self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers + self.assertEqual(top_k_index.shape, (6, 2)) + self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) + # one scale tensor per projection: the affine Parameter, or the swizzled cache when the + # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) + self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) + self.assertIsNone(call.kwargs["gate_up_proj_global_scale"]) + self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name + self.assertIs(call.kwargs["gate"], True) + self.assertEqual(call.kwargs["recipe"], "weights") # activation_format None = weight family self.assertEqual(out.shape, (6, 64)) - def test_grouped_marshalling_expert_start_boundaries(self): + def test_grouped_marshalling(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - up = rec.calls["matmul_grouped"][0] - es = up.kwargs["expert_start"] - self.assertEqual(es.shape, (5,)) # (E+1,) boundaries - self.assertEqual(es[0].item(), 0) - self.assertEqual(es[-1].item(), 12) # S = tokens * top_k + (call,) = rec.calls["moe_fused_grouped"] + self.assertIs(call.kwargs["down_proj"], m.down_proj) + self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) + self.assertNotIn("moe_fused_batched", rec.calls) + + def test_activation_format_maps_to_the_block_recipe(self): + kernel, rec = _fake_bundle() + for activation_format, recipe in ((None, "weights"), ("bf16", None), ("mxfp8", "mxfp8")): + m = self._experts(has_gate=True, activation_format=activation_format) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["recipe"], recipe) def test_nvfp4_experts_thread_per_expert_globals(self): kernel, rec = _fake_bundle() @@ -248,39 +307,35 @@ def test_nvfp4_experts_thread_per_expert_globals(self): p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_batched_mm_experts_forward(m, *self._route()) - up, down = rec.calls["matmul_batched"] - self.assertIs(up.kwargs["b_global_scale"], m.gate_up_proj_global_scale) - self.assertIs(down.kwargs["b_global_scale"], m.down_proj_global_scale) + (call,) = rec.calls["moe_fused_batched"] + self.assertIs(call.kwargs["gate_up_proj_global_scale"], m.gate_up_proj_global_scale) + self.assertIs(call.kwargs["down_proj_global_scale"], m.down_proj_global_scale) - def test_bias_lands_before_gate_and_before_routing(self): + def test_biases_ride_the_kernel_chain(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True, has_bias=True) - hs, idx, wts = self._route() p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: - out = fg.finegrained_batched_mm_experts_forward(m, hs, idx, wts) + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) self.assertTrue(torch.isfinite(out).all()) - up, down = rec.calls["matmul_batched"] - # both biases ride their GEMM's epilogue: the kernel adds gate_up's before the gated - # split and down's before the routing-weight multiply. Nothing is added host-side, so - # a biased model keeps the fused form instead of falling back to two GEMMs. - self.assertIs(up.kwargs["bias"], m.gate_up_proj_bias) - self.assertTrue(up.kwargs["epilogue"].gate) - self.assertIs(down.kwargs["bias"], m.down_proj_bias) - - def test_unfusable_act_fn_still_fuses_the_bias(self): + (call,) = rec.calls["moe_fused_batched"] + # the kernels add gate_up's bias before the gated split and down's before the + # routing-weight multiply; nothing is added host-side + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) + + def test_unfusable_act_fn_is_passed_as_the_module_glu(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True, has_bias=True) - m.act_fn_name = "quick_gelu" # not in _FUSABLE_ACT_FNS - hs, idx, wts = self._route() + m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: - fg.finegrained_batched_mm_experts_forward(m, hs, idx, wts) - up, _ = rec.calls["matmul_batched"] - # the GLU cannot fuse, but the ungated GEMM emitting the raw 2*I pre-activation takes - # the same bias add at doubled width — only the activation falls back to the host tail - self.assertIsNone(up.kwargs["epilogue"]) - self.assertIs(up.kwargs["bias"], m.gate_up_proj_bias) + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels run the module's own GLU on the host between the two GEMMs — a new + # activation never waits for a kernel release; the bias still rides the GEMM + self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) @require_torch @@ -316,27 +371,38 @@ def _reference_dequant(self, blocks, scales): exp = (scales.long() - 127).unsqueeze(-1) return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) - def test_deserialize_matches_reference(self): - from transformers.integrations.finegrained import FineGrainedMxfp4Deserialize, _get_ue8m0_dtype + def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): + from transformers.integrations.finegrained import FineGrainedPackedBlocks, FineGrainedScaleContainer torch.manual_seed(0) E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) - op = FineGrainedMxfp4Deserialize(hf_quantizer=None) - out = op.convert( - {"gate_up_proj_blocks": blocks, "gate_up_proj_scales": scales}, - full_layer_name="model.layers.0.mlp.experts.gate_up_proj", + weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ + "gate_up_proj" + ] + self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) + back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] + self.assertTrue(torch.equal(back, blocks)) + + # the scales are the container op's job: a module holding e8m0 receives the exponent bytes as is + module = fg.FineGrainedExperts.__new__(fg.FineGrainedExperts) + torch.nn.Module.__init__(module) + module.gate_up_proj_scale_inv = torch.nn.Parameter( + torch.zeros(E, N, K // 32, dtype=torch.float8_e8m0fnu), requires_grad=False ) - weight = out["model.layers.0.mlp.experts.gate_up_proj"] - scale_inv = out["model.layers.0.mlp.experts.gate_up_proj_scale_inv"] - self.assertEqual(weight.dtype, torch.int8) - self.assertEqual(weight.shape, (E, N, K // 2)) - self.assertEqual(scale_inv.dtype, _get_ue8m0_dtype()) + model = torch.nn.Module() + model.experts = module + scale_inv = FineGrainedScaleContainer().convert( + {"gate_up_proj_scales$": scales}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + target_patterns=["gate_up_proj_scale_inv"], + )["gate_up_proj_scale_inv"] + self.assertEqual(scale_inv.dtype, torch.float8_e8m0fnu) ref = self._reference_dequant(blocks, scales) # rows pass through interleaved - # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) lut = torch.tensor(self.FP4_VALUES) w_u8 = weight.view(torch.uint8) @@ -347,57 +413,6 @@ def test_deserialize_matches_reference(self): got = vals * torch.pow(torch.tensor(2.0), exp) torch.testing.assert_close(got, ref, rtol=0, atol=0) - def test_interleave_gate_up_after_loading(self): - """Stacked-checkpoint families get gate_up interleaved post-load; GPT-OSS-style MXFP4 - already ships that order and is skipped.""" - from transformers.integrations.finegrained import ( - FineGrainedExperts, - interleave_gate_up_after_loading, - ) - - module = FineGrainedExperts.__new__(FineGrainedExperts) - torch.nn.Module.__init__(module) - module.has_gate = True - # rows [g0,g1,g2,u0,u1,u2]; bias follows the same output axis - module.gate_up_proj = torch.nn.Parameter( - torch.arange(24, dtype=torch.float32).reshape(1, 6, 4), requires_grad=False - ) - module.gate_up_proj_bias = torch.nn.Parameter( - torch.arange(6, dtype=torch.float32), requires_grad=False - ) - stacked_w = module.gate_up_proj.detach().clone() - stacked_b = module.gate_up_proj_bias.detach().clone() - - model = torch.nn.Module() - model.experts = module - interleave_gate_up_after_loading(model) - - for j in range(3): - torch.testing.assert_close(module.gate_up_proj[0, 2 * j], stacked_w[0, j]) - torch.testing.assert_close(module.gate_up_proj[0, 2 * j + 1], stacked_w[0, 3 + j]) - torch.testing.assert_close(module.gate_up_proj_bias[2 * j], stacked_b[j]) - torch.testing.assert_close(module.gate_up_proj_bias[2 * j + 1], stacked_b[3 + j]) - - def test_interleave_skipped_when_already_interleaved(self): - """GPT-OSS ships [g0,u0,...] already — the pass must leave it alone.""" - from transformers.integrations.finegrained import ( - FineGrainedExperts, - interleave_gate_up_after_loading, - ) - - module = FineGrainedExperts.__new__(FineGrainedExperts) - torch.nn.Module.__init__(module) - module.has_gate = True - module.gate_up_proj = torch.nn.Parameter( - torch.arange(24, dtype=torch.float32).reshape(1, 6, 4), requires_grad=False - ) - before = module.gate_up_proj.detach().clone() - - model = torch.nn.Module() - model.experts = module - interleave_gate_up_after_loading(model, already_interleaved=True) - torch.testing.assert_close(module.gate_up_proj.data, before) - @require_torch class FineGrainedDeepGemmDispatchTest(unittest.TestCase): @@ -434,54 +449,343 @@ def test_swizzled_scales_never_reach_deepgemm(self): """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + def test_deepgemm_experts_refuse_swizzled_scales(self): + """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM + experts backend afterwards must fail loudly rather than read the permuted buffer as affine.""" + from transformers.integrations.deepgemm import ( + deepgemm_fp8_fp4_experts_forward, + deepgemm_fp8_fp4_megamoe_experts_forward, + ) + + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" + with ( + mock.patch.object(fg, "is_sm100", return_value=True), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts(cfg, weight_format="mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) + hs = torch.randn(2, 256, dtype=torch.bfloat16) + idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.assertRaisesRegex(RuntimeError, "held swizzled"): + forward(experts, hs, idx, wts) + @require_torch -class FineGrainedSwizzlePassTest(unittest.TestCase): - """`swizzle_scales_after_loading` builds the Blackwell SWIZZLE_32_4_4 artifact. It is gated - on arch AND on the recipe having group scales — keyed off the declared weight format, not - the scale dtype, because V4-style block-FP8 ships UE8M0 scales and a dtype test lets it - through (it then survives only by accident, since N/128 is rarely a multiple of 128).""" +class FineGrainedScaleLayoutTest(unittest.TestCase): + """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 + scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the + scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's + `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" + + E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks - def _experts(self, weight_format, scale_dtype): + def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): cfg = _Cfg() - experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format=weight_format) - for proj in ("gate_up_proj", "down_proj"): - rows = 256 - setattr( - experts, - f"{proj}_scale_inv", - torch.nn.Parameter( - torch.zeros(cfg.num_local_experts, rows, 8, dtype=scale_dtype), - requires_grad=False, - ), + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E + cfg._experts_implementation = impl + with ( + mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts( + cfg, block_size=(128, 128), weight_format=weight_format, activation_format=activation_format ) model = torch.nn.Module() + model.config = cfg model.experts = experts return model, experts - def _run_pass(self, model, *, sm100): + def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) + self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertTrue(experts._gate_up_interleaved) + # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a + # float32 grid there; GPT-OSS's config has no scale_fmt at all) + _, experts = self._experts("mxfp4", activation_format="bf16") + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + + def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): + for kw in ( + {"weight_format": "mxfp8", "sm100": False}, + {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely + {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales + {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, + ): + with self.subTest(**kw): + _, experts = self._experts(**kw) + self.assertEqual( + experts.gate_up_proj_scale_inv.shape, + (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), + ) + # block-FP8's (N/128, K/128) grid never reaches a scaled-MMA, whatever its scale dtype + _, experts = self._experts("fp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) + with mock.patch.dict("os.environ", {"TRANSFORMERS_FINEGRAINED_NO_SWIZZLE": "1"}): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) + + def test_megamoe_holds_gate_up_stacked(self): + _, experts = self._experts("fp8", impl="deepgemm_megamoe") + self.assertFalse(experts._gate_up_interleaved) + + def _op_kernel(self): kernel, _ = _fake_bundle() - kernel.swizzle_mx_scales = mock.Mock(side_effect=lambda s, gate=False: s) + kernel.swizzle_mx_scales = mock.Mock( + side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) + ) + kernel.unswizzle_mx_scales = mock.Mock( + side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) + ) + return kernel + + def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): + from transformers.integrations.finegrained import FineGrainedSwizzleScales + + model, experts = self._experts("mxfp8") + kernel = self._op_kernel() + op = FineGrainedSwizzleScales(hf_quantizer=None) + grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(fg, "is_sm100", return_value=sm100): - fg.swizzle_scales_after_loading(model) - return kernel.swizzle_mx_scales + with p1, p2, p3, p4: + # the loader hands a single target under its pattern, the full name (with suffix) alongside + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" + ) + self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) + # the same converter's weight passes through + out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") + self.assertIs(out["mlp.experts.gate_up_proj"], weight) + # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully + out = op.convert( + { + "experts.down_proj": weight, + "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), + }, + model=model, + full_layer_name="experts.down_proj", + ) + self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) + self.assertIs(out["experts.down_proj"], weight) + # a module that holds affine scales keeps them + affine_model, _ = self._experts("mxfp4", activation_format="bf16") + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, + model=affine_model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + self.assertIs(out["mlp.experts.gate_up_proj"], grid) + self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) + + def test_scale_container_op_round_trips_every_container(self): + """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: + the op brings both into the held e8m0 (exact cast / same bytes), records the container on + the module, and its reverse restores that container on save; weights and native scales + pass both ways.""" + from transformers.integrations.finegrained import FineGrainedScaleContainer + + op = FineGrainedScaleContainer(hf_quantizer=None) + native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) + for container, shipped in ((torch.float32, native.float()), (torch.uint8, native.view(torch.uint8))): + with self.subTest(container=container): + model, experts = self._experts("mxfp8", sm100=False) # affine e8m0 scale Parameters + out = op.convert( + {"mlp.experts.gate_up_proj": shipped}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + held = out["mlp.experts.gate_up_proj"] + self.assertEqual(held.dtype, torch.float8_e8m0fnu) + self.assertTrue(torch.equal(held.view(torch.uint8), native.view(torch.uint8))) + self.assertIs(experts.scale_container_dtype, container) + back = op.reverse_op.convert({"x": held}, model=model)["x"] + self.assertEqual(back.dtype, container) + self.assertTrue(torch.equal(back, shipped)) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + self.assertIs(op.reverse_op.convert({"w": weight}, model=model)["w"], weight) + model, experts = self._experts("mxfp8", sm100=False) + self.assertIs( + op.convert({"mlp.experts.down_proj": native}, model=model, full_layer_name="experts.down_proj_scale_inv")[ + "mlp.experts.down_proj" + ], + native, + ) + self.assertIsNone(experts.scale_container_dtype) + self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays + + def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): + from transformers.integrations.finegrained import FineGrainedSwizzleScales + + kernel = self._op_kernel() + reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op + artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) + affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") + self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) + kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) + out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") + self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) - def test_skipped_off_sm100(self): - """The layout is a Blackwell tcgen05 artifact; elsewhere it buys nothing, pins the tile - space, and collides with DeepGEMM — which stays the preferred backend below SM100.""" - model, _ = self._experts("mxfp8", torch.float8_e8m0fnu) - self.assertEqual(self._run_pass(model, sm100=False).call_count, 0) + def _quantizer(self, quant_method): + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + return FineGrainedHfQuantizer(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions + + def _arch_converters(self): + from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter + + return [ + WeightConverter( + source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"], + target_patterns="mlp.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight", + target_patterns="mlp.experts.down_proj", + operations=[MergeModulelist(dim=0)], + ), + WeightConverter( + source_patterns=["mlp.fc1.weight", "mlp.fc2.weight"], + target_patterns="mlp.fc.weight", + operations=[Concatenate(dim=0)], + ), + ] + + def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): + from transformers.core_model_loading import WeightConverter + from transformers.integrations.finegrained import ( + FineGrainedInterleaveGateUp, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) - def test_applied_for_group_scaled_recipes_on_sm100(self): - model, _ = self._experts("mxfp8", torch.float8_e8m0fnu) - self.assertGreater(self._run_pass(model, sm100=True).call_count, 0) + converters = self._quantizer("mxfp8")(self._arch_converters()) + by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} + gate_up_ops = by_target["mlp.experts.gate_up_proj"].operations + # stacked [gate; up] -> [g0, u0, ...], the scale into the held dtype, then the swizzle packs + # the final row order (it reads 1-byte scales) + self.assertIsInstance(gate_up_ops[-3], FineGrainedInterleaveGateUp) + self.assertFalse(gate_up_ops[-3].inverse) + self.assertIsInstance(gate_up_ops[-2], FineGrainedScaleContainer) + self.assertIsInstance(gate_up_ops[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["mlp.experts.down_proj"].operations[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["weight_scale_inv"].operations[0], FineGrainedScaleContainer) # dense linears + self.assertEqual(len(by_target["mlp.fc.weight"].operations), 1) # dense converters untouched + # keys arriving under the fused names get the same layout + self.assertIn("experts.gate_up_proj_bias", by_target) # patterns are stored anchor-stripped + self.assertIsInstance(by_target["experts.down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) + # ...and saving reverses in the opposite order: unswizzle first, then de-interleave + reverse_ops = by_target["mlp.experts.gate_up_proj"].reverse_transform().operations + self.assertIsInstance(reverse_ops[0], FineGrainedSwizzleScales) + self.assertTrue(reverse_ops[0].inverse) + self.assertIsInstance(reverse_ops[2], FineGrainedInterleaveGateUp) + self.assertTrue(reverse_ops[2].inverse) + + def test_blocks_scales_converters_carry_the_layout_ops(self): + from transformers.core_model_loading import WeightConverter + from transformers.integrations.finegrained import FineGrainedInterleaveGateUp, FineGrainedSwizzleScales + + converters = self._quantizer("mxfp4")([]) + by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} + ops = by_target["gate_up_proj"].operations + self.assertIsInstance(ops[-2], FineGrainedInterleaveGateUp) + self.assertIsInstance(ops[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["gate_up_proj_scale_inv"].operations[0], FineGrainedInterleaveGateUp) + self.assertIsInstance(by_target["down_proj"].operations[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) + + def test_catch_all_converters_deliver_the_full_parameter_name(self): + """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax + down_proj_scale_inv) goes through the real converter plumbing: the ops receive the tensor + under the source PATTERN and must hand back the target, which the loader expands to the + full name.""" + from transformers.core_model_loading import WeightConverter, rename_source_key + + model, _ = self._experts("mxfp8") + converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] + for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): + renamed, source_pattern = rename_source_key(key, [], converters) + self.assertEqual(renamed, key) + converter = next(c for c in converters if source_pattern in c.source_patterns) + tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) + converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) + kernel = self._op_kernel() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = converter.convert(renamed, model=model) + self.assertEqual(list(out), [key]) + + def test_replacement_keeps_the_models_gate_up_convention(self): + """`replace_with_finegrained_layer` goes through `use_experts_implementation`, which stamps + the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must + survive it (it decides whether the loader interleaves).""" + from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM + from transformers.utils.quantization_config import FineGrainedConfig - def test_block_fp8_skipped_even_with_ue8m0_scales(self): - """The latent case: block-FP8's (N/128, K/128) grid never reaches a scaled-MMA, so it has - no swizzled form — but it ships UE8M0 scales, so only a recipe-keyed guard excludes it.""" - model, _ = self._experts("fp8", torch.float8_e8m0fnu) - self.assertEqual(self._run_pass(model, sm100=True).call_count, 0) + gpt_oss = GptOssConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=32, + num_local_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + sliding_window=16, + ) + qwen = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + ) + for cls, cfg, quant_method, expected in ( + (GptOssForCausalLM, gpt_oss, "mxfp4", False), + (Qwen3MoeForCausalLM, qwen, "mxfp8", True), + ): + with self.subTest(model=cls.__name__), torch.device("meta"): + model = cls(cfg) + fg.replace_with_finegrained_layer( + model, modules_to_not_convert=[], quantization_config=FineGrainedConfig(quant_method=quant_method) + ) + experts = model.model.layers[0].mlp.experts + self.assertIsInstance(experts, fg.FineGrainedExperts) + self.assertIs(experts.is_concatenated, expected) + + def test_interleave_op_follows_the_model_and_the_experts_backend(self): + """The op reads the model's experts modules: a model whose own rows already come + interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself + and keeps the stack; every other case gets the interleaved order, and the reverse restores + the stack.""" + from transformers.integrations.finegrained import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + model, experts = self._experts("mxfp8") + out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) + back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(back, stacked) + experts.is_concatenated = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) + experts.is_concatenated = True + experts._gate_up_interleaved = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) @require_torch @@ -560,59 +864,62 @@ def test_linear_forward_matches_the_quantization_floor(self): rel = ((out.float() - ref).norm() / ref.norm()).item() self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") - def test_swizzled_scales_do_not_change_the_result(self): - """The SWIZZLE_32_4_4 pass reorders scale BYTES; it must not move a single value. Run the - same MXFP8 GEMM against the affine grid and the swizzled artifact and demand agreement — - a wrong reorder shows up here and nowhere in the mocked tests.""" - kernel = load_finegrained_kernel() - N, K, M = 512, 256, 128 - w = torch.randn(N, K, device="cuda", dtype=torch.float32) - groups = w.reshape(N, K // 32, 32) - inv = torch.pow(2.0, torch.ceil(torch.log2(groups.abs().amax(-1, keepdim=True) / 448.0).clamp(min=-127))) - q = (groups / inv).to(torch.float8_e4m3fn).reshape(N, K) - scales = inv.reshape(N, K // 32).to(torch.float8_e8m0fnu) - x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) - - affine = kernel.matmul(x, q, None, scales, output_dtype=torch.bfloat16) - swizzled = kernel.matmul( - x, q, None, kernel.swizzle_mx_scales(scales), output_dtype=torch.bfloat16 - ) - torch.cuda.synchronize() - self.assertTrue( - torch.equal(affine, swizzled), - "swizzling changed the result; the artifact is a byte reorder, values must be identical", - ) + def test_experts_scale_layout_op_keeps_the_forward_correct(self): + """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op + fills them from the affine grid, the fused forward matches the affine module's, and the + reverse op hands the affine grid back bitwise.""" + from transformers.integrations.finegrained import FineGrainedSwizzleScales - def test_experts_swizzle_pass_keeps_the_forward_correct(self): - """End-to-end over the integration: build real MXFP8 experts, run the post-load swizzle, - and check the fused forward still matches what it produced beforehand.""" cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = ( + 512, + 256, + 4, + ) # K >= 256: the swizzled arm has tiles + cfg._experts_implementation = "grouped_mm" experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() + with mock.patch.dict("os.environ", {"TRANSFORMERS_FINEGRAINED_NO_SWIZZLE": "1"}): + affine_experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() + model = torch.nn.Module() + model.experts = experts + op = FineGrainedSwizzleScales(hf_quantizer=None) + + affine = {} for proj, rows in (("gate_up_proj", 2 * cfg.intermediate_size), ("down_proj", cfg.hidden_size)): cols = cfg.hidden_size if proj == "gate_up_proj" else cfg.intermediate_size - setattr(experts, proj, torch.nn.Parameter( - torch.randn(cfg.num_local_experts, rows, cols, device="cuda").to(torch.float8_e4m3fn), - requires_grad=False)) - setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter( - torch.full((cfg.num_local_experts, rows, cols // 32), 127, dtype=torch.uint8, device="cuda") - .view(torch.float8_e8m0fnu), requires_grad=False)) + weight = torch.randn(cfg.num_local_experts, rows, cols, device="cuda").to(torch.float8_e4m3fn) + grid = torch.randint(120, 134, (cfg.num_local_experts, rows, cols // 32), dtype=torch.uint8, device="cuda") + affine[proj] = grid.view(torch.float8_e8m0fnu) + for module in (experts, affine_experts): + setattr(module, proj, torch.nn.Parameter(weight.clone(), requires_grad=False)) + setattr(affine_experts, f"{proj}_scale_inv", torch.nn.Parameter(affine[proj].clone(), requires_grad=False)) + self.assertEqual( + getattr(experts, f"{proj}_scale_inv").dim(), 5, "the module should hold this projection swizzled" + ) + out = op.convert( + {f"mlp.experts.{proj}": affine[proj]}, model=model, full_layer_name=f"experts.{proj}_scale_inv" + ) + swizzled = out[f"mlp.experts.{proj}"] + self.assertEqual(swizzled.shape, getattr(experts, f"{proj}_scale_inv").shape) + setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(swizzled, requires_grad=False)) x = torch.randn(8, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) idx = torch.randint(0, cfg.num_local_experts, (8, 2), device="cuda", dtype=torch.long) wts = torch.rand(8, 2, device="cuda", dtype=torch.bfloat16) - - before = experts(x, idx, wts) - model = torch.nn.Module() - model.experts = experts - fg.swizzle_scales_after_loading(model) - after = experts(x, idx, wts) + reference = fg.finegrained_grouped_mm_experts_forward(affine_experts, x, idx, wts) + out = fg.finegrained_grouped_mm_experts_forward(experts, x, idx, wts) torch.cuda.synchronize() - - self.assertTrue( - any(getattr(experts, f"{p}_scale_inv_swizzled", None) is not None - for p in ("gate_up_proj", "down_proj")), - "the swizzle pass produced no artifact, so this asserts nothing", - ) - rel = ((after.float() - before.float()).norm() / before.float().norm().clamp(min=1e-9)).item() + rel = ((out.float() - reference.float()).norm() / reference.float().norm().clamp(min=1e-9)).item() self.assertLess(rel, 1e-6, f"swizzled forward diverged from the affine one: {rel:.2e}") + # the eager per-expert loop reads one expert's slice of the swizzled stack directly + eager_reference = affine_experts(x, idx, wts) + eager_out = experts(x, idx, wts) + rel = ( + (eager_out.float() - eager_reference.float()).norm() / eager_reference.float().norm().clamp(min=1e-9) + ).item() + self.assertLess(rel, 1e-6, f"eager forward on swizzled scales diverged from affine: {rel:.2e}") + + reverse = op.reverse_op + for proj, grid in affine.items(): + restored = reverse.convert({proj: getattr(experts, f"{proj}_scale_inv").data})[proj] + self.assertTrue(torch.equal(restored.view(torch.uint8), grid.view(torch.uint8)), proj) From 1d94e5572e3f119b989a9eaec658185c9279bbe1 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 11 Sep 2026 02:20:10 -0700 Subject: [PATCH 09/68] style --- src/transformers/distributed/sharding_utils.py | 4 +--- src/transformers/integrations/tensor_parallel.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 536588f0469b..009001db7945 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -117,9 +117,7 @@ def shard_tensor( # here: replicated on the dense path, per-expert-ownership-filtered on the MoE path. if not source_shape: if tensor_idx is not None: - has_axis0_shard = any( - self._normalize_param_dim(placement.dim) == 0 for _, placement in dim_placements - ) + has_axis0_shard = any(self._normalize_param_dim(placement.dim) == 0 for _, placement in dim_placements) if has_axis0_shard and not ( self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size ): diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 8cb485ba64fe..896782d7e30c 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -78,6 +78,7 @@ def to_local(t): return t.to_local() return t + __all__ = [ "to_local", "ALL_PARALLEL_STYLES", From fb16472d09fb5794dd9edb7a6bdf781428e4861f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 11 Sep 2026 05:36:07 -0700 Subject: [PATCH 10/68] finegrained: on-the-fly quantization for every format; loader keeps checkpoint precision for quantized params MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `FineGrainedQuantize` reads the module's weight format: block-FP8 in torch (per-block E4M3, fp32 or power-of-two UE8M0 inverse scales, from the module's block and held scale dtype), MXFP8 / MXFP4 / NVFP4 through the kernels' row-wise quantizers in one launch per tensor (NVFP4 normalized by the per-tensor / per-expert global `amax / (6 * 448)`), and emits the scale (and global) in the layout the module holds — the container dtype, the swizzled artifact — since the loader runs the quantization op after the converter ops. Core loader: a parameter the quantizer's op is about to quantize is materialized in the checkpoint's precision instead of the empty parameter's storage dtype (int8 storage zeroed FP4 weights; float8 double-rounded the existing FP8 path). Also: `auto.py` groups the finegrained keys under one comment and drops unused NVFP4 imports; the quantizer's config type hint names `FineGrainedConfig`; DeepGEMM guards (`_assert_affine_scales`, `_assert_stacked_gate_up`, group-32 divisibility) run before the kernel load; expert biases take the model dtype; comment trims for the noisy-comments check; tests for every path. Co-Authored-By: Claude Fable 5.1 --- src/transformers/core_model_loading.py | 7 +- src/transformers/integrations/finegrained.py | 167 +++++++++++------- src/transformers/quantizers/auto.py | 35 ++-- .../quantizers/quantizer_finegrained.py | 4 +- tests/kernels/test_finegrained.py | 76 +++++++- 5 files changed, 201 insertions(+), 88 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 580f46849c3c..a2df840ac33c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1699,8 +1699,11 @@ def convert_and_load_state_dict_in_model( matched_dtype_pattern = dtype_policy_alt.search(renamed_key) if matched_dtype_pattern is not None: _dtype = dtype_plan[dtype_policy_by_group_name[matched_dtype_pattern.lastgroup]] - elif empty_param is not None and empty_param.dtype != _dtype: - _dtype = empty_param.dtype # usually correct when initializing + elif empty_param is not None and empty_param.dtype != _dtype and not needs_quantization: + # usually correct when initializing — but a parameter the quantizer's op is about to + # quantize must arrive in the checkpoint's precision, not pre-cast to its quantized + # storage dtype (int8 storage would zero it, float8 would double-round it) + _dtype = empty_param.dtype # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the # distributed op selects whole experts. The signal is a `MergeModulelist` diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index ae115490db7f..bd6f969f4d49 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -81,8 +81,9 @@ class FineGrained: classes ship in the same build. The experts forwards hand the kernels' MoE forwards the module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as - the module's own callable and runs on the host between the two GEMMs. All symbols are - required — a build missing any raises at load with the full list. + the module's own callable and runs on the host between the two GEMMs. The quantizers serve + on-the-fly quantization into the group formats. All symbols are required — a build missing any + raises at load with the full list. """ matmul_2d: Callable @@ -92,6 +93,9 @@ class FineGrained: moe_fused_grouped: Callable swizzle_mx_scales: Callable unswizzle_mx_scales: Callable + mxfp8_act_quant: Callable + mxfp4_act_quant: Callable + nvfp4_act_quant: Callable get_supported_act_fns: Callable Quantization: Callable Epilogue: Callable @@ -1205,78 +1209,121 @@ def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=Non class FineGrainedQuantize(ConversionOps): - """ - A quantization operation that creates two tensors, weight and scale out of a weight. - """ + """Quantize a full-precision weight on load into the format its module holds, emitting the + scale (and the NVFP4 global) in the module's layout — the held container dtype, the swizzled + artifact where the module holds one. Block-FP8 is computed in torch (per-block E4M3, fp32 or + power-of-two UE8M0 inverse scales); the group formats run the kernels' row-wise quantizers + (``mxfp8_act_quant`` / ``mxfp4_act_quant`` / ``nvfp4_act_quant``), NVFP4 after normalizing by the + canonical per-tensor / per-expert global ``amax / (6 * 448)``. + Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) + pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" def __init__(self, hf_quantizer): self.hf_quantizer = hf_quantizer - def _resolve_block_size(self, value: torch.Tensor) -> tuple[int, int]: - block_size = None - if self.hf_quantizer.quantization_config is not None: - if isinstance(self.hf_quantizer.quantization_config, dict): - block_size = self.hf_quantizer.quantization_config.get("weight_block_size") - else: - block_size = getattr(self.hf_quantizer.quantization_config, "weight_block_size", None) - if block_size is None: - block_size = (value.shape[-2], value.shape[-1]) - return tuple(block_size) - - def _quantize_one(self, key: str, value: torch.Tensor) -> dict[str, torch.Tensor]: - # Pass through tensors that aren't tileable (1D norms / biases, or shapes - # that don't divide cleanly by the configured block) — they were never - # FP8-quantized on the load side, so the reverse op shouldn't touch them. - if value.ndim < 2: - return {key: value} - block_m, block_n = self._resolve_block_size(value) - rows, cols = value.shape[-2], value.shape[-1] - if rows % block_m != 0 or cols % block_n != 0: - return {key: value} - - # Leading dims can be empty (2D) or include num_experts/... (3D+) - leading_shape = value.shape[:-2] - rows_tiles = rows // block_m - cols_tiles = cols // block_n - original_shape = value.shape - value_fp32 = value.to(torch.float32) - # Reshape to (..., rows_tiles, block_m, cols_tiles, block_n) - reshaped = value_fp32.reshape(*leading_shape, rows_tiles, block_m, cols_tiles, block_n) - # Per-tile max-abs over the block dims (block_m at -3, block_n at -1) - max_abs = reshaped.abs().amax(dim=(-3, -1)) - safe_max_abs = torch.where(max_abs > 0, max_abs, torch.ones_like(max_abs)) - # We store inverse scale to match the upstream ``weight_scale_inv`` convention - scales = _FP8_MAX / safe_max_abs - scales = torch.where(max_abs > 0, scales, torch.ones_like(scales)) # keep zeros stable - inv_scales = (1.0 / scales).to(torch.float32) - # ue8m0 stores weight_scale_inv as a power of two. Round it before quantizing and derive the - # forward scale from it, so dequant multiplies by the exact scale the weight was divided by. - if self.hf_quantizer.quantization_config.scale_fmt == "ue8m0": - inv_scales = torch.pow(2.0, torch.ceil(torch.log2(inv_scales.clamp(min=torch.finfo(torch.float32).tiny)))) - inv_scales = inv_scales.to(_get_ue8m0_dtype()) - scales = 1.0 / inv_scales.to(torch.float32) # forward scale = exact reciprocal of the stored inverse - # Broadcast scales over the block dims and quantize - scales_broadcast = scales.unsqueeze(-1).unsqueeze(-3) # (..., rows_tiles, 1, cols_tiles, 1) - scaled = reshaped * scales_broadcast - quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) - quantized = quantized.reshape(original_shape) - scale_key = key.rsplit(".", 1)[0] + ".weight_scale_inv" if key.endswith(".weight") else key + "_scale_inv" - return {key: quantized, scale_key: inv_scales} - - def convert(self, input_dict: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]: - # every (key, tensor) entry — one key, or one per expert from the reverse of a - # ``MergeModulelist`` — quantizes the same way + def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: result: dict[str, torch.Tensor] = {} for key, value in input_dict.items(): tensor = value[0] if isinstance(value, list) else value - result.update(self._quantize_one(key, tensor)) + result.update(self._quantize_one(key, tensor, self._weight_holder(model, key))) return result + @staticmethod + def _weight_holder(model, key: str): + """``(module, scale_param_name)`` when ``key`` is a finegrained module's weight, else ``None``.""" + if model is None: + return None + module, name = get_module_from_name(model, key) + if not isinstance(module, _FineGrainedModule) or getattr(module, name, None) is None: + return None + if name == "weight": + return module, "weight_scale_inv" + if name in ("gate_up_proj", "up_proj", "down_proj"): + return module, f"{name}_scale_inv" + return None + + def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torch.Tensor]: + if value.ndim < 2: + return {key: value} + prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key + module = holder[0] if holder is not None else None + held_scale = getattr(module, holder[1]) if holder is not None else None + fmt = _weight_formats()[module.weight_format] if module is not None else None + if fmt is None or fmt.scale_group is None: + block = tuple(module.block_size) if module is not None and module.block_size else self._config_block(value) + ue8m0 = ( + held_scale.dtype == _get_ue8m0_dtype() + if held_scale is not None + else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" + ) + quantized = _quantize_block_fp8(value, block, ue8m0) + return {key: value} if quantized is None else {key: quantized[0], f"{prefix}_scale_inv": quantized[1]} + + if value.device.type != "cuda": + raise ValueError(f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on CUDA") + weight, scale, global_scale = _quantize_group(value, fmt) + scale = _recontain(scale, held_scale.dtype) + if held_scale.ndim == 5: + scale = load_finegrained_kernel().swizzle_mx_scales(scale) + out = {key: weight, f"{prefix}_scale_inv": scale} + if global_scale is not None: + out[f"{prefix}_global_scale"] = global_scale + return out + + def _config_block(self, value: torch.Tensor) -> tuple[int, int]: + config = self.hf_quantizer.quantization_config + block = ( + config.get("weight_block_size") if isinstance(config, dict) else getattr(config, "weight_block_size", None) + ) + return tuple(block) if block is not None else (value.shape[-2], value.shape[-1]) + @property def reverse_op(self) -> ConversionOps: return FineGrainedDequantize(self.hf_quantizer) +def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): + """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or + ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two + before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" + block_m, block_n = block + rows, cols = value.shape[-2], value.shape[-1] + if rows % block_m or cols % block_n: + return None + tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) + max_abs = tiles.abs().amax(dim=(-3, -1)) + inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) + if ue8m0: + inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) + scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) + return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) + + +def _quantize_group(value: torch.Tensor, fmt: _WeightFormat): + """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the + kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has + one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the + smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" + kernel = load_finegrained_kernel() + rows = value.reshape(-1, value.shape[-1]) + global_scale = None + if fmt.has_global_scale: + per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) + global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) + rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) + packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) + global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale + else: + quantize = kernel.mxfp4_act_quant if fmt.values_per_byte == 2 else kernel.mxfp8_act_quant + packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) + return ( + packed.view(fmt.weight_dtype).reshape(*value.shape[:-1], -1), + scale.reshape(*value.shape[:-1], -1), + global_scale, + ) + + class FineGrainedDequantize(ConversionOps): """Dequantize FP8 weights using their per-block ``weight_scale_inv``. diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index bccb17f47d1e..43c2bb60528a 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -36,7 +36,6 @@ HqqConfig, MetalConfig, Mxfp4Config, - NVFP4Config, # noqa: F401 (dense on-the-fly path; not mapped — see below) QuantizationConfigMixin, QuantizationMethod, QuantoConfig, @@ -65,7 +64,6 @@ from .quantizer_higgs import HiggsHfQuantizer from .quantizer_hqq import HqqHfQuantizer from .quantizer_metal import MetalHfQuantizer -from .quantizer_nvfp4 import NVFP4HfQuantizer # noqa: F401 (dense on-the-fly path, see below) from .quantizer_quanto import QuantoHfQuantizer from .quantizer_quark import QuarkHfQuantizer from .quantizer_sinq import SinqHfQuantizer @@ -93,25 +91,17 @@ "bitnet": BitNetHfQuantizer, "vptq": VptqHfQuantizer, "spqr": SpQRHfQuantizer, - "fp8": FineGrainedHfQuantizer, "gguf": GgufHfQuantizer, - # NVFP4 (two-level E2M1: E4M3 group-16 block scales x fp32 per-tensor global) — served by the - # finegrained quantizer; the format is resolved off the checkpoint tensors. - # - # This key is contested: `NVFP4HfQuantizer` (added upstream in #47883) also claims it, but it - # quantizes a bf16 checkpoint ON THE FLY and only replaces `nn.Linear`, so it cannot serve a - # PRE-QUANTIZED checkpoint and never reaches MoE experts (fused 3-D parameters, not Linear). - # The finegrained path covers both, so it owns the key; `NVFP4Linear` remains reachable - # directly for the dense on-the-fly case. - "nvfp4": FineGrainedHfQuantizer, - "modelopt": FineGrainedHfQuantizer, - # MXFP8 = FP8 (E4M3 weights) with per-block ``[1, 32]`` E8M0 (uint8) scales — - # reuses the FineGrainedFP8 dequant path, with the E8M0 byte→exponent - # unpacking handled inside ``Fp8Dequantize._dequantize_one``. - "mxfp8": FineGrainedHfQuantizer, + "metal": MetalHfQuantizer, "auto-round": AutoRoundQuantizer, + # the finegrained quantizer serves every block/group-scaled format — block-FP8, MXFP8, MXFP4, + # NVFP4 and modelopt's NVFP4 export — for dense linears and MoE experts alike; the format is + # resolved off the checkpoint tensors + "fp8": FineGrainedHfQuantizer, + "mxfp8": FineGrainedHfQuantizer, "mxfp4": FineGrainedHfQuantizer, - "metal": MetalHfQuantizer, + "nvfp4": FineGrainedHfQuantizer, + "modelopt": FineGrainedHfQuantizer, "sinq": SinqHfQuantizer, "gemma": GemmaQuantizer, } @@ -136,13 +126,14 @@ "bitnet": BitNetQuantConfig, "vptq": VptqConfig, "spqr": SpQRConfig, + "metal": MetalConfig, + "auto-round": AutoRoundConfig, + # the finegrained formats (see AUTO_QUANTIZER_MAPPING) "fp8": FineGrainedConfig, - "nvfp4": FineGrainedConfig, # see AUTO_QUANTIZER_MAPPING above for why this key is ours - "modelopt": FineGrainedConfig, "mxfp8": FineGrainedConfig, - "auto-round": AutoRoundConfig, "mxfp4": FineGrainedConfig, - "metal": MetalConfig, + "nvfp4": FineGrainedConfig, + "modelopt": FineGrainedConfig, "sinq": SinqConfig, "gemma": GemmaQuantizationConfig, } diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 94667bad4344..c55fd4e0b604 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from ..modeling_utils import PreTrainedModel - from ..utils.quantization_config import FineGrainedFP8Config + from ..utils.quantization_config import FineGrainedConfig logger = logging.get_logger(__name__) @@ -23,7 +23,7 @@ class FineGrainedHfQuantizer(HfQuantizer): """ requires_calibration = False - quantization_config: "FineGrainedFP8Config" + quantization_config: "FineGrainedConfig" def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index fa6749c2b077..f87c3814b810 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -71,6 +71,9 @@ def _fake_bundle(): "moe_fused_grouped", "swizzle_mx_scales", "unswizzle_mx_scales", + "mxfp8_act_quant", + "mxfp4_act_quant", + "nvfp4_act_quant", ): setattr(kernel, name, rec._op(name)) kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") @@ -788,6 +791,37 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) +@require_torch +class FineGrainedOnTheFlyQuantizeTest(unittest.TestCase): + """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run the + kernels' quantizers — covered with real kernels below), emitted in the module's layout.""" + + def test_block_fp8_round_trips_within_its_floor(self): + from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + for scale_fmt in ("float", "ue8m0"): + with self.subTest(scale_fmt=scale_fmt), mock.patch.object(fg, "is_sm100", return_value=False): + experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", scale_fmt=scale_fmt) + model = torch.nn.Module() + model.experts = experts + original = torch.randn(4, 256, 256) * 0.02 + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] + self.assertEqual( + (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) + ) + self.assertEqual( + (scale.dtype, scale.shape), + (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), + ) + deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) + rel = ((deq - original).norm() / original.norm()).item() + self.assertLess(rel, 4e-2, f"{scale_fmt}: {rel:.3f}") + + @require_torch class FineGrainedModeloptConverterTest(unittest.TestCase): """The ModelOpt (NVFP4) weight conversions. `MergeModulelist` is load-bearing beyond the @@ -864,6 +898,43 @@ def test_linear_forward_matches_the_quantization_floor(self): rel = ((out.float() - ref).norm() / ref.norm()).item() self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") + def test_on_the_fly_group_formats_round_trip_within_their_floor(self): + """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in + the module's layout (swizzled where held so); dequantizing them back lands within each + format's floor.""" + from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + for fmt, floor in (("mxfp8", 4e-2), ("mxfp4", 2e-1), ("nvfp4", 2e-1)): + with self.subTest(fmt=fmt): + experts = FineGrainedExperts(cfg, weight_format=fmt).cuda() + model = torch.nn.Module() + model.experts = experts + original = (torch.randn(4, 512, 512, device="cuda") * 0.02).to(torch.bfloat16) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] + self.assertEqual( + (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) + ) + self.assertEqual( + (scale.dtype, scale.shape), + (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), + ) + if scale.ndim == 5: + kernel = load_finegrained_kernel() + scale = kernel.unswizzle_mx_scales(scale, 512, scale.shape[2] * 4, num_experts=4) + deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) + if fmt == "nvfp4": + deq = deq * out["experts.gate_up_proj_global_scale"].reshape(-1, 1, 1) + else: + self.assertNotIn("experts.gate_up_proj_global_scale", out) + rel = ((deq - original.float()).norm() / original.float().norm()).item() + self.assertLess(rel, floor, f"{fmt}: {rel:.3f}") + self.assertGreater(rel, 0.0) + def test_experts_scale_layout_op_keeps_the_forward_correct(self): """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op fills them from the affine grid, the fused forward matches the affine module's, and the @@ -910,14 +981,15 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): out = fg.finegrained_grouped_mm_experts_forward(experts, x, idx, wts) torch.cuda.synchronize() rel = ((out.float() - reference.float()).norm() / reference.float().norm().clamp(min=1e-9)).item() - self.assertLess(rel, 1e-6, f"swizzled forward diverged from the affine one: {rel:.2e}") + # different tuned arms reorder the fp32 reduction (~1e-5); a layout bug is O(1) + self.assertLess(rel, 1e-3, f"swizzled forward diverged from the affine one: {rel:.2e}") # the eager per-expert loop reads one expert's slice of the swizzled stack directly eager_reference = affine_experts(x, idx, wts) eager_out = experts(x, idx, wts) rel = ( (eager_out.float() - eager_reference.float()).norm() / eager_reference.float().norm().clamp(min=1e-9) ).item() - self.assertLess(rel, 1e-6, f"eager forward on swizzled scales diverged from affine: {rel:.2e}") + self.assertLess(rel, 1e-3, f"eager forward on swizzled scales diverged from affine: {rel:.2e}") reverse = op.reverse_op for proj, grid in affine.items(): From 07dae5009358894309fb122028c17abc976ca8fa Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 11 Sep 2026 07:45:38 -0700 Subject: [PATCH 11/68] finegrained: pass input_recipe to the dispatchers, drop the Quantization/Epilogue bundle fields The kernels' dispatchers now take input_recipe/output_recipe strings, so the bundle no longer needs the Quantization and Epilogue config classes; the linears hand the module-level activation_format straight through as input_recipe and the _kernel_quantization adapter is gone. Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/finegrained.py | 35 +++++++------------- tests/kernels/test_finegrained.py | 24 +++----------- 2 files changed, 17 insertions(+), 42 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index bd6f969f4d49..c2a64242645a 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -76,12 +76,12 @@ class FineGrained: """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and - the two MoE forwards, with the recipe resolved off the weight/scale dtypes inside the - kernel; the per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary - classes ship in the same build. The experts forwards hand the kernels' MoE forwards the - module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, - biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as - the module's own callable and runs on the host between the two GEMMs. The quantizers serve + the two MoE forwards, with the weight recipe resolved off the weight/scale dtypes inside the + kernel and ``input_recipe`` naming the activation format at the op boundary. The experts + forwards hand the kernels' MoE forwards the module's tensors and let them run the chain + (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an + activation outside ``get_supported_act_fns()`` is passed as the module's own callable and + runs on the host between the two GEMMs. The quantizers serve on-the-fly quantization into the group formats. All symbols are required — a build missing any raises at load with the full list. """ @@ -97,8 +97,6 @@ class FineGrained: mxfp4_act_quant: Callable nvfp4_act_quant: Callable get_supported_act_fns: Callable - Quantization: Callable - Epilogue: Callable # Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env @@ -301,14 +299,13 @@ def finegrained_triton_linear( the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). """ kernel = load_finegrained_kernel() - quantization = _kernel_quantization(kernel, activation_format) original_shape = input.shape output = kernel.matmul_2d( input.reshape(-1, original_shape[-1]), weight, activation_scale, weight_scale_inv, - quantization=quantization, + input_recipe=_block_recipe(activation_format), output_dtype=input.dtype, b_global_scale=weight_global_scale, ) @@ -587,7 +584,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: None, scale_inv, expert_start=expert_start, - quantization=_kernel_quantization(kernel, self.activation_format), + input_recipe=_block_recipe(self.activation_format), b_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) @@ -596,19 +593,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -def _kernel_quantization(kernel, activation_format: str | None): - """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` - keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" - if activation_format is None: - return None - recipe = None if activation_format == "bf16" else activation_format - return kernel.Quantization(input_recipe=recipe) - - def _block_recipe(activation_format: str | None) -> str | None: - """The module-level ``activation_format`` as the kernels' block ``recipe``: ``None`` keeps - the weight family's own format, ``"bf16"`` is weight-only (no activation quant), an explicit - name is itself.""" + """The module-level ``activation_format`` as the kernels' activation recipe (``input_recipe`` + of the dispatchers, ``recipe`` of the MoE forwards): ``None`` keeps the weight family's own + format (``"weights"``), ``"bf16"`` is weight-only (no activation quant), an explicit name is + itself.""" if activation_format is None: return "weights" return None if activation_format == "bf16" else activation_format diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index f87c3814b810..436a18411ba1 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -78,20 +78,6 @@ def _fake_bundle(): setattr(kernel, name, rec._op(name)) kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") - @dataclass - class Quantization: - input_recipe: str | None = "weights" - output_recipe: str | None = None - - @dataclass - class Epilogue: - gate: bool = False - act_fn: str = "silu" - swiglu_alpha: float | None = None - swiglu_limit: float | None = None - - kernel.Quantization = Quantization - kernel.Epilogue = Epilogue return kernel, rec @@ -130,8 +116,8 @@ def test_loader_binds_all_symbols(self): p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: bundle = load_finegrained_kernel() - self.assertIs(bundle.Quantization, kernel.Quantization) - self.assertIs(bundle.Epilogue, kernel.Epilogue) + self.assertIs(bundle.matmul_2d, kernel.matmul_2d) + self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) @require_torch @@ -164,7 +150,7 @@ def test_nvfp4_global_and_weight_only_format(self): g = torch.tensor(2.0) _, call = self._run(weight_global_scale=g, activation_format="bf16") self.assertIs(call.kwargs["b_global_scale"], g) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertIsNone(call.kwargs["input_recipe"]) def test_module_forward_threads_everything(self): kernel, rec = _fake_bundle() @@ -181,7 +167,7 @@ def test_module_forward_threads_everything(self): out = m(torch.randn(2, 64, dtype=torch.bfloat16)) call = rec.calls["matmul_2d"][-1] self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertIsNone(call.kwargs["input_recipe"]) self.assertEqual(out.shape, (2, 32)) @@ -259,7 +245,7 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): self.assertEqual(call.kwargs["activation_format"], "bf16") for call in rec.calls["matmul_2d"]: self.assertIsNotNone(call.kwargs["b_global_scale"]) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertIsNone(call.kwargs["input_recipe"]) def test_batched_marshalling(self): kernel, rec = _fake_bundle() From 0b6ca4c9936e92c4f1f66abcb3a79d36e181972b Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sun, 13 Sep 2026 03:48:17 -0700 Subject: [PATCH 12/68] Revert "finegrained: pass input_recipe to the dispatchers, drop the Quantization/Epilogue bundle fields" This reverts commit 07dae5009358894309fb122028c17abc976ca8fa. The kernels keep the Quantization/Epilogue op-boundary configs as the dispatcher API, so the bundle carries them and the linears build kernel.Quantization again. Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/finegrained.py | 35 +++++++++++++------- tests/kernels/test_finegrained.py | 24 +++++++++++--- 2 files changed, 42 insertions(+), 17 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index c2a64242645a..bd6f969f4d49 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -76,12 +76,12 @@ class FineGrained: """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and - the two MoE forwards, with the weight recipe resolved off the weight/scale dtypes inside the - kernel and ``input_recipe`` naming the activation format at the op boundary. The experts - forwards hand the kernels' MoE forwards the module's tensors and let them run the chain - (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an - activation outside ``get_supported_act_fns()`` is passed as the module's own callable and - runs on the host between the two GEMMs. The quantizers serve + the two MoE forwards, with the recipe resolved off the weight/scale dtypes inside the + kernel; the per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary + classes ship in the same build. The experts forwards hand the kernels' MoE forwards the + module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, + biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as + the module's own callable and runs on the host between the two GEMMs. The quantizers serve on-the-fly quantization into the group formats. All symbols are required — a build missing any raises at load with the full list. """ @@ -97,6 +97,8 @@ class FineGrained: mxfp4_act_quant: Callable nvfp4_act_quant: Callable get_supported_act_fns: Callable + Quantization: Callable + Epilogue: Callable # Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env @@ -299,13 +301,14 @@ def finegrained_triton_linear( the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). """ kernel = load_finegrained_kernel() + quantization = _kernel_quantization(kernel, activation_format) original_shape = input.shape output = kernel.matmul_2d( input.reshape(-1, original_shape[-1]), weight, activation_scale, weight_scale_inv, - input_recipe=_block_recipe(activation_format), + quantization=quantization, output_dtype=input.dtype, b_global_scale=weight_global_scale, ) @@ -584,7 +587,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: None, scale_inv, expert_start=expert_start, - input_recipe=_block_recipe(self.activation_format), + quantization=_kernel_quantization(kernel, self.activation_format), b_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) @@ -593,11 +596,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y +def _kernel_quantization(kernel, activation_format: str | None): + """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` + keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" + if activation_format is None: + return None + recipe = None if activation_format == "bf16" else activation_format + return kernel.Quantization(input_recipe=recipe) + + def _block_recipe(activation_format: str | None) -> str | None: - """The module-level ``activation_format`` as the kernels' activation recipe (``input_recipe`` - of the dispatchers, ``recipe`` of the MoE forwards): ``None`` keeps the weight family's own - format (``"weights"``), ``"bf16"`` is weight-only (no activation quant), an explicit name is - itself.""" + """The module-level ``activation_format`` as the kernels' block ``recipe``: ``None`` keeps + the weight family's own format, ``"bf16"`` is weight-only (no activation quant), an explicit + name is itself.""" if activation_format is None: return "weights" return None if activation_format == "bf16" else activation_format diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 436a18411ba1..f87c3814b810 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -78,6 +78,20 @@ def _fake_bundle(): setattr(kernel, name, rec._op(name)) kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") + @dataclass + class Quantization: + input_recipe: str | None = "weights" + output_recipe: str | None = None + + @dataclass + class Epilogue: + gate: bool = False + act_fn: str = "silu" + swiglu_alpha: float | None = None + swiglu_limit: float | None = None + + kernel.Quantization = Quantization + kernel.Epilogue = Epilogue return kernel, rec @@ -116,8 +130,8 @@ def test_loader_binds_all_symbols(self): p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: bundle = load_finegrained_kernel() - self.assertIs(bundle.matmul_2d, kernel.matmul_2d) - self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) + self.assertIs(bundle.Quantization, kernel.Quantization) + self.assertIs(bundle.Epilogue, kernel.Epilogue) @require_torch @@ -150,7 +164,7 @@ def test_nvfp4_global_and_weight_only_format(self): g = torch.tensor(2.0) _, call = self._run(weight_global_scale=g, activation_format="bf16") self.assertIs(call.kwargs["b_global_scale"], g) - self.assertIsNone(call.kwargs["input_recipe"]) + self.assertIsNone(call.kwargs["quantization"].input_recipe) def test_module_forward_threads_everything(self): kernel, rec = _fake_bundle() @@ -167,7 +181,7 @@ def test_module_forward_threads_everything(self): out = m(torch.randn(2, 64, dtype=torch.bfloat16)) call = rec.calls["matmul_2d"][-1] self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) - self.assertIsNone(call.kwargs["input_recipe"]) + self.assertIsNone(call.kwargs["quantization"].input_recipe) self.assertEqual(out.shape, (2, 32)) @@ -245,7 +259,7 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): self.assertEqual(call.kwargs["activation_format"], "bf16") for call in rec.calls["matmul_2d"]: self.assertIsNotNone(call.kwargs["b_global_scale"]) - self.assertIsNone(call.kwargs["input_recipe"]) + self.assertIsNone(call.kwargs["quantization"].input_recipe) def test_batched_marshalling(self): kernel, rec = _fake_bundle() From 513e231457467cb4eb31b5c5d28114ae2240260a Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 14 Sep 2026 00:03:43 -0700 Subject: [PATCH 13/68] finegrained: pass activation_format straight through; let deep-gemm skip the kernels arch check The kernels now speak the module's vocabulary: activation_format (None = the weights' format, "bf16" = weight-only) goes straight to matmul_2d / matmul_grouped / the MoE forwards, so the Quantization/Epilogue bundle fields and both adapter helpers are gone. hub_kernels: kernels >= 0.16 checks a build's declared archs against the device; the deep-gemm build declares only 9.0a although it JIT-compiles per device, so the mapping opts out via a new check_arch passthrough (only forwarded when False, so older kernels releases without the keyword keep working). Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/finegrained.py | 52 ++++++-------------- src/transformers/integrations/hub_kernels.py | 23 +++++++-- tests/kernels/test_finegrained.py | 37 +++++--------- 3 files changed, 47 insertions(+), 65 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index bd6f969f4d49..2fec11dc9cb7 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -75,13 +75,14 @@ def _first_attr(obj, *names): class FineGrained: """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. - Every recipe (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and - the two MoE forwards, with the recipe resolved off the weight/scale dtypes inside the - kernel; the per-recipe quant helpers and the ``Epilogue``/``Quantization`` op-boundary - classes ship in the same build. The experts forwards hand the kernels' MoE forwards the - module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, - biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as - the module's own callable and runs on the host between the two GEMMs. The quantizers serve + Every format (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and + the two MoE forwards, with the weight format resolved off the weight/scale dtypes inside the + kernel and the module's ``activation_format`` passed through unchanged — the kernels speak the + same vocabulary (``None`` = the weights' format, ``"bf16"`` = weight-only). The experts + forwards hand the kernels' MoE forwards the module's tensors and let them run the chain + (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an + activation outside ``get_supported_act_fns()`` is passed as the module's own callable and runs + on the host between the two GEMMs. The quantizers serve on-the-fly quantization into the group formats. All symbols are required — a build missing any raises at load with the full list. """ @@ -97,8 +98,6 @@ class FineGrained: mxfp4_act_quant: Callable nvfp4_act_quant: Callable get_supported_act_fns: Callable - Quantization: Callable - Epilogue: Callable # Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env @@ -276,7 +275,7 @@ def _swizzles_scales(config, fmt: _WeightFormat, activation_format: str | None) return ( getattr(config, "_experts_implementation", None) not in ("deepgemm", "deepgemm_megamoe") and fmt.scale_group is not None - and _block_recipe(activation_format) is not None + and activation_format != "bf16" and torch.cuda.is_available() and is_sm100() ) @@ -293,22 +292,21 @@ def finegrained_triton_linear( ) -> torch.Tensor: """Triton fine-grained linear: fused act-quant + matmul, then optional bias add. - Serves every weight recipe the kernel resolves off the tensors themselves — block-FP8 + Serves every weight format the kernel resolves off the tensors themselves — block-FP8 (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4 (``int8``-packed values; the two-level per-tensor ``weight_global_scale`` recovers on the accumulator). ``activation_scale=None`` → dynamic activation quant (inline); a per-tensor scalar → static quant against it. - ``activation_format`` picks the activation recipe where the weights leave it open (``None`` = + ``activation_format`` picks the activation format where the weights leave it open (``None`` = the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). """ kernel = load_finegrained_kernel() - quantization = _kernel_quantization(kernel, activation_format) original_shape = input.shape output = kernel.matmul_2d( input.reshape(-1, original_shape[-1]), weight, activation_scale, weight_scale_inv, - quantization=quantization, + activation_format=activation_format, output_dtype=input.dtype, b_global_scale=weight_global_scale, ) @@ -353,8 +351,8 @@ def finegrained_linear( model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound to a single CUDA context and produce garbage across devices (see ``disable_deepgemm_on_multi_device``). - weight_global_scale: the NVFP4 two-level per-tensor global (None for other recipes). - activation_format: the activation recipe where the weights leave it open (see + weight_global_scale: the NVFP4 two-level per-tensor global (None for other formats). + activation_format: the activation format where the weights leave it open (see ``finegrained_triton_linear``). """ # DeepGEMM is CUDA-only, dynamic-only, SM90+, FP4 or 128x128-block FP8; the combos it would @@ -587,7 +585,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: None, scale_inv, expert_start=expert_start, - quantization=_kernel_quantization(kernel, self.activation_format), + activation_format=self.activation_format, b_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) @@ -596,24 +594,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -def _kernel_quantization(kernel, activation_format: str | None): - """Map the module-level ``activation_format`` onto the kernel's ``Quantization``. ``None`` - keeps the kernel's weight-native default; ``"bf16"`` = weight-only (no activation quant).""" - if activation_format is None: - return None - recipe = None if activation_format == "bf16" else activation_format - return kernel.Quantization(input_recipe=recipe) - - -def _block_recipe(activation_format: str | None) -> str | None: - """The module-level ``activation_format`` as the kernels' block ``recipe``: ``None`` keeps - the weight family's own format, ``"bf16"`` is weight-only (no activation quant), an explicit - name is itself.""" - if activation_format is None: - return "weights" - return None if activation_format == "bf16" else activation_format - - def _moe_operands(kernel, module) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their scales (held swizzled for dot_scaled chains), the per-expert globals and biases, @@ -644,7 +624,7 @@ def local(name): "act_fn": act_fn, "swiglu_alpha": module.swiglu_alpha, "swiglu_limit": module.swiglu_limit, - "recipe": _block_recipe(module.activation_format), + "activation_format": module.activation_format, "gate": module.has_gate, } diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 61dfd673c26e..d1b3483115e8 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -622,7 +622,9 @@ def register_kernel_mapping_transformers(*args, **kwargs): _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { - "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2}, + # DeepGEMM JIT-compiles its CUDA kernels for the running device; the build metadata declares only + # the arch it was packaged on (9.0a), so the `kernels` arch check would wrongly reject sm_100. + "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2, "check_arch": False}, "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); @@ -756,7 +758,13 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ if version is None and revision is None: version = 1 - kernel = get_kernel(repo_id, revision=revision, version=version, allow_all_kernels=ALLOW_ALL_KERNELS) + kernel = get_kernel( + repo_id, + revision=revision, + version=version, + allow_all_kernels=ALLOW_ALL_KERNELS, + check_arch=_HUB_KERNEL_MAPPING[kernel_name].get("check_arch", True), + ) mapping[kernel_name] = kernel except FileNotFoundError as e: mapping[kernel_name] = None @@ -818,6 +826,7 @@ def get_kernel( revision: str | None = None, version: int | str | None = None, allow_all_kernels: bool = False, + check_arch: bool = True, ) -> ModuleType: from .. import __version__ @@ -825,8 +834,16 @@ def get_kernel( raise ImportError(_MISSING_KERNELS_MESSAGE) user_agent = {"framework": "transformers", "version": __version__, "repo_id": kernel_name} + # `check_arch` (kernels >= 0.16) is only passed when a caller opts out, so older `kernels` + # releases without the keyword keep working. + arch_kwargs = {} if check_arch else {"check_arch": False} return get_kernel_hub( - kernel_name, revision=revision, version=version, user_agent=user_agent, trust_remote_code=allow_all_kernels + kernel_name, + revision=revision, + version=version, + user_agent=user_agent, + trust_remote_code=allow_all_kernels, + **arch_kwargs, ) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index f87c3814b810..a50537e6447c 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Marshalling tests for the multi-recipe ``finegrained`` integration: the kernel bundle is +"""Marshalling tests for the multi-format ``finegrained`` integration: the kernel bundle is mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` (the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" @@ -77,21 +77,6 @@ def _fake_bundle(): ): setattr(kernel, name, rec._op(name)) kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") - - @dataclass - class Quantization: - input_recipe: str | None = "weights" - output_recipe: str | None = None - - @dataclass - class Epilogue: - gate: bool = False - act_fn: str = "silu" - swiglu_alpha: float | None = None - swiglu_limit: float | None = None - - kernel.Quantization = Quantization - kernel.Epilogue = Epilogue return kernel, rec @@ -130,8 +115,8 @@ def test_loader_binds_all_symbols(self): p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: bundle = load_finegrained_kernel() - self.assertIs(bundle.Quantization, kernel.Quantization) - self.assertIs(bundle.Epilogue, kernel.Epilogue) + self.assertIs(bundle.matmul_2d, kernel.matmul_2d) + self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) @require_torch @@ -164,7 +149,7 @@ def test_nvfp4_global_and_weight_only_format(self): g = torch.tensor(2.0) _, call = self._run(weight_global_scale=g, activation_format="bf16") self.assertIs(call.kwargs["b_global_scale"], g) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertEqual(call.kwargs["activation_format"], "bf16") def test_module_forward_threads_everything(self): kernel, rec = _fake_bundle() @@ -181,7 +166,7 @@ def test_module_forward_threads_everything(self): out = m(torch.randn(2, 64, dtype=torch.bfloat16)) call = rec.calls["matmul_2d"][-1] self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertEqual(call.kwargs["activation_format"], "bf16") self.assertEqual(out.shape, (2, 32)) @@ -259,7 +244,7 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): self.assertEqual(call.kwargs["activation_format"], "bf16") for call in rec.calls["matmul_2d"]: self.assertIsNotNone(call.kwargs["b_global_scale"]) - self.assertIsNone(call.kwargs["quantization"].input_recipe) + self.assertEqual(call.kwargs["activation_format"], "bf16") def test_batched_marshalling(self): kernel, rec = _fake_bundle() @@ -278,7 +263,7 @@ def test_batched_marshalling(self): self.assertIsNone(call.kwargs["gate_up_proj_global_scale"]) self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name self.assertIs(call.kwargs["gate"], True) - self.assertEqual(call.kwargs["recipe"], "weights") # activation_format None = weight family + self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format self.assertEqual(out.shape, (6, 64)) def test_grouped_marshalling(self): @@ -292,14 +277,14 @@ def test_grouped_marshalling(self): self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) self.assertNotIn("moe_fused_batched", rec.calls) - def test_activation_format_maps_to_the_block_recipe(self): + def test_activation_format_passes_through(self): kernel, rec = _fake_bundle() - for activation_format, recipe in ((None, "weights"), ("bf16", None), ("mxfp8", "mxfp8")): + for activation_format in (None, "bf16", "mxfp8"): m = self._experts(has_gate=True, activation_format=activation_format) p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_batched_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["recipe"], recipe) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) def test_nvfp4_experts_thread_per_expert_globals(self): kernel, rec = _fake_bundle() @@ -879,7 +864,7 @@ def _block_fp8_weight(self, N, K, *, ue8m0): blocks = w.reshape(N // 128, 128, K // 128, 128) amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12) inv = amax / 448.0 - if ue8m0: # power-of-two scales: the tcgen05 dot_scaled recipe + if ue8m0: # power-of-two scales: the tcgen05 dot_scaled format inv = torch.pow(2.0, torch.ceil(torch.log2(inv))) q = (blocks / inv).to(torch.float8_e4m3fn) deq = (q.float() * inv).reshape(N, K) From 240abc9bebf469430c35aa4259d74a88642fe7d0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 14 Sep 2026 02:03:09 -0700 Subject: [PATCH 14/68] hub_kernels: finegrained-kernels ships as version 0 Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/hub_kernels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index d1b3483115e8..0dde1d24721d 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -629,7 +629,7 @@ def register_kernel_mapping_transformers(*args, **kwargs): "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); # exports the same three matmuls plus the fused-MoE forwards and per-recipe quant helpers - "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 1}, + "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 0}, # dense NVFP4 GEMM behind `NVFP4Linear`; MoE experts route to finegrained-kernels instead "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, } From 30ab2c21d8c3d60e2c6d7bab6aefdeaa76520b79 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 14 Sep 2026 04:05:40 -0700 Subject: [PATCH 15/68] hub_kernels: load finegrained-kernels from the kernels-staging build (pr-1018) Co-Authored-By: Claude Fable 5.1 --- src/transformers/integrations/hub_kernels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 0dde1d24721d..083388a0a65a 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -629,7 +629,7 @@ def register_kernel_mapping_transformers(*args, **kwargs): "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); # exports the same three matmuls plus the fused-MoE forwards and per-recipe quant helpers - "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 0}, + "finegrained-kernels": {"repo_id": "kernels-staging/finegrained-kernels", "revision": "pr-1018"}, # dense NVFP4 GEMM behind `NVFP4Linear`; MoE experts route to finegrained-kernels instead "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, } From a562fb8c41ff59e4023f3a415958106d394a9b42 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 02:31:06 -0700 Subject: [PATCH 16/68] finegrained: consume modelopt's calibrated globals, and fuse a model's per-expert output norm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit modelopt NVFP4 checkpoints ship a second-level global per projection and a calibrated `input_scale`; both now reach the kernels, in either of the two layouts such a checkpoint uses — one tensor per expert per projection (GLM-5.2) or one stacked tensor per layer (the fused vLLM layout). A gate|up stack arrives with two weight globals per expert, calibrated separately. One conversion op owns every global of a layer and merges them to the one-per-expert the kernels take: the stack keeps the gate's, and the up half's folds onto the down projection, whose weight global scales the expert output back up and whose calibrated input global moves the other way, keeping the requantized intermediate on the range the checkpoint calibrated. That is exact arithmetic on fp32 scalars, unlike rescaling the up half's e4m3 block scales. Many-to-many converters are now opt-in per op (`ConversionOps.supports_many_to_many`) rather than a private whitelist. A model states its per-expert output norm the way it states its activation: `use_experts_implementation(post_expert_norm="")` plus an `_apply_post_norm` the class must define. A name the kernels implement rides as that name and is folded into their reduce; anything else is the module, called on the routed rows. The DeepGEMM experts paths, which have no slot for such a norm, refuse instead of dropping it. Weight-only modules hold no activation global, since nothing there quantizes activations: the parameters are not allocated, the converters for them are not emitted, and a calibrated checkpoint's `input_scale` keys are ignored for that run. Co-Authored-By: Claude Opus 5 --- src/transformers/core_model_loading.py | 20 +- src/transformers/integrations/deepgemm.py | 23 +- src/transformers/integrations/finegrained.py | 301 ++++++++++--- src/transformers/integrations/moe.py | 23 + .../quantizers/quantizer_finegrained.py | 116 ++++- src/transformers/utils/quantization_config.py | 13 +- tests/kernels/test_finegrained.py | 421 +++++++++++++++++- 7 files changed, 814 insertions(+), 103 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index a2df840ac33c..3918c3800431 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,6 +81,11 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" + # Whether this op maps SEVERAL sources onto SEVERAL targets in one call — a converter is + # otherwise restricted to one-to-many, one-to-one or many-to-one, since an m:n mapping only + # makes sense when a single operation owns the whole relation (see `WeightConverter`). + supports_many_to_many: bool = False + def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -613,6 +618,8 @@ class ErnieFuseAndSplitTextVisionExperts(ConversionOps): The final fusions are defined by the amount of original module lists. """ + supports_many_to_many = True + def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -675,6 +682,8 @@ class ErnieSplitAndDecoupleTextVisionExperts(ConversionOps): The final decoupled module lists are defined by the amount of keys. """ + supports_many_to_many = True + def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -1142,13 +1151,6 @@ def reverse_transform(self) -> WeightTransform: return result -# List of classes that are known to be able to use m:n -_INTERNAL_MANY_TO_MANY_CONVERSIONS = ( - ErnieFuseAndSplitTextVisionExperts, - ErnieSplitAndDecoupleTextVisionExperts, -) - - class WeightConverter(WeightTransform): __slots__ = ("operations", "force_cpu") @@ -1164,8 +1166,8 @@ def __init__( self.force_cpu = force_cpu if bool(len(self.source_patterns) - 1) + bool(len(self.target_patterns) - 1) >= 2: - # We allow many-to-many only if we use an internal operation that can handle it - if not any(isinstance(op, _INTERNAL_MANY_TO_MANY_CONVERSIONS) for op in self.operations): + # We allow many-to-many only if an operation declares that it handles the whole relation + if not any(op.supports_many_to_many for op in self.operations): raise ValueError( f"source keys={self.source_patterns}, target_patterns={self.target_patterns} but you can only have one to many, one to one or many to one." ) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 4ac05fde7e12..b815fd0c0d30 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -619,12 +619,24 @@ def deepgemm_fp8_fp4_linear( return output +def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: + """DeepGEMM's experts chain ends at the routing-weighted reduce, with no slot for a model's + per-expert output norm — which it would drop silently, so refuse instead.""" + if getattr(self, "post_expert_norm", None) is not None: + raise NotImplementedError( + "DeepGEMM experts dispatch cannot apply this model's per-expert output norm; use " + "`experts_implementation='grouped_mm'` (or 'batched_mm')." + ) + + def deepgemm_bf16_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: + _assert_no_post_expert_norm(self) + if hidden_states.dtype != torch.bfloat16: raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") @@ -702,7 +714,7 @@ def _assert_stacked_gate_up(module: torch.nn.Module) -> None: """Mega MoE's ``transform_weights_for_mega_moe`` does its own gate/up interleave, so it takes the stacked ``[gate; up]`` rows. A module loaded for a triton backend holds them interleaved; switched to this backend after load, silently re-permuting would be a wrong answer rather than a failure.""" - if module.has_gate and getattr(module, "_gate_up_interleaved", False): + if module.has_gate and getattr(module, "holds_interleaved_gate_up", False): raise RuntimeError( "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — pass " @@ -725,6 +737,7 @@ def deepgemm_fp8_fp4_experts_forward( "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) + _assert_no_post_expert_norm(self) # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as # affine and turn into garbage — refuse before any kernel work. _assert_affine_scales(self) @@ -733,12 +746,13 @@ def deepgemm_fp8_fp4_experts_forward( # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) - deepgemm = load_deepgemm_kernel() - if self.activation_scheme == "static": raise NotImplementedError("DeepGEMM experts dispatch does not support static activation quantization.") if hidden_states.dtype != torch.bfloat16: raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") + + deepgemm = load_deepgemm_kernel() + grouped_fp8_fp4_matmul = ( deepgemm.grouped_fp8_fp4_matmul_nn if self.is_transposed else deepgemm.grouped_fp8_fp4_matmul_nt ) @@ -897,6 +911,8 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( `transform_weights_for_mega_moe((gate_up, gate_up_sf), (down, down_sf))`. - `config.swiglu_limit` (optional): SwiGLU clamp; absent → unclamped. """ + + _assert_no_post_expert_norm(self) # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as # affine and turn into garbage — refuse before any kernel work. _assert_affine_scales(self) @@ -910,7 +926,6 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( f"DeepGEMM Mega MoE requires FP4-packed expert weights (dtype=`int8`), got " f"`{self.gate_up_proj.dtype}`. Use the 'deepgemm' dispatch for FP8 experts." ) - if process_group is None: raise ValueError( "DeepGEMM Mega MoE requires a `process_group` for the EP group. The TP wrapping " diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 2fec11dc9cb7..2dced4f365c9 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -17,8 +17,10 @@ import importlib import os import sys +from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass +from types import MethodType import torch import torch.nn as nn @@ -82,7 +84,8 @@ class FineGrained: forwards hand the kernels' MoE forwards the module's tensors and let them run the chain (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an activation outside ``get_supported_act_fns()`` is passed as the module's own callable and runs - on the host between the two GEMMs. The quantizers serve + on the host between the two GEMMs; a per-expert output norm the kernels implement + (``get_supported_norms()``) rides the same way, by name, and is fused into the chain's reduce. The quantizers serve on-the-fly quantization into the group formats. All symbols are required — a build missing any raises at load with the full list. """ @@ -98,6 +101,8 @@ class FineGrained: mxfp4_act_quant: Callable nvfp4_act_quant: Callable get_supported_act_fns: Callable + get_supported_norms: Callable + rms_norm_rows: Callable # Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env @@ -288,13 +293,16 @@ def finegrained_triton_linear( bias: torch.Tensor | None = None, activation_scale: torch.Tensor | None = None, weight_global_scale: torch.Tensor | None = None, + input_global_scale: torch.Tensor | None = None, activation_format: str | None = None, ) -> torch.Tensor: """Triton fine-grained linear: fused act-quant + matmul, then optional bias add. Serves every weight format the kernel resolves off the tensors themselves — block-FP8 (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4 (``int8``-packed values; the two-level - per-tensor ``weight_global_scale`` recovers on the accumulator). ``activation_scale=None`` → + per-tensor ``weight_global_scale`` recovers on the accumulator; ``input_global_scale`` is + NVFP4's CALIBRATED activation global, the checkpoint's ``input_scale`` — the quant divides by it + and the accumulator multiplies it back). ``activation_scale=None`` → dynamic activation quant (inline); a per-tensor scalar → static quant against it. ``activation_format`` picks the activation format where the weights leave it open (``None`` = the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). @@ -308,6 +316,7 @@ def finegrained_triton_linear( weight_scale_inv, activation_format=activation_format, output_dtype=input.dtype, + a_global_scale=input_global_scale, b_global_scale=weight_global_scale, ) output = output.reshape(*original_shape[:-1], output.shape[-1]) @@ -325,6 +334,7 @@ def finegrained_linear( activation_scale: torch.Tensor | None = None, allow_deepgemm: bool = True, weight_global_scale: torch.Tensor | None = None, + input_global_scale: torch.Tensor | None = None, activation_format: str | None = None, ) -> torch.Tensor: """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. @@ -351,7 +361,10 @@ def finegrained_linear( model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound to a single CUDA context and produce garbage across devices (see ``disable_deepgemm_on_multi_device``). - weight_global_scale: the NVFP4 two-level per-tensor global (None for other formats). + weight_global_scale: the NVFP4 two-level weight global, one value for this matrix + (None for other formats). + input_global_scale: the NVFP4 calibrated activation global, the checkpoint's + ``input_scale`` (None = quantize against the block scales alone). activation_format: the activation format where the weights leave it open (see ``finegrained_triton_linear``). """ @@ -373,8 +386,9 @@ def finegrained_linear( and weight_scale_inv.ndim <= 2 and is_deepgemm_loadable() and activation_scale is None - # DeepGEMM serves neither the NVFP4 two-level global nor an explicit activation format + # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format and weight_global_scale is None + and input_global_scale is None and activation_format is None and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" @@ -410,6 +424,7 @@ def finegrained_linear( bias, activation_scale, weight_global_scale=weight_global_scale, + input_global_scale=input_global_scale, activation_format=activation_format, ) @@ -594,6 +609,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y +def _norm_eps(norm: nn.Module) -> float: + """A norm module's epsilon, under either name transformers modules give it. ``nn.RMSNorm`` + leaves it ``None`` and resolves it to the dtype's own, so this does the same.""" + eps = getattr(norm, "eps", None) + if eps is None: + eps = getattr(norm, "variance_epsilon", None) + return float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps + + def _moe_operands(kernel, module) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their scales (held swizzled for dot_scaled chains), the per-expert globals and biases, @@ -602,30 +626,53 @@ def _moe_operands(kernel, module) -> dict: change.""" up = "gate_up_proj" if module.has_gate else "up_proj" - def local(name): - t = getattr(module, name, None) - return to_local(t) if t is not None else None - act_name = module.act_fn_name fused = act_name in kernel.get_supported_act_fns() and (module.swiglu_alpha is None or act_name == "silu") if fused: act_fn = act_name else: act_fn = module._apply_gate if module.has_gate else module.act_fn + + norm, norm_eps, norm_weight = module.post_expert_norm, 1e-6, None + if norm is None: + post_expert_norm = None + elif module.post_expert_norm_name in kernel.get_supported_norms(): + post_expert_norm = module.post_expert_norm_name + norm_weight, norm_eps = norm.weight, _norm_eps(norm) + else: + post_expert_norm = module._apply_post_norm + + def local(name): + t = getattr(module, name, None) + return to_local(t) if t is not None else None + + def both(suffix: str) -> dict: + """One kernel argument per projection. The kernels always say ``gate_up_proj``; the module + names that weight ``up_proj`` when the model has no gate. A tensor the module does not + hold passes as ``None``, which is the argument the kernels take for it.""" + return { + f"gate_up_proj{suffix}": local(f"{up}{suffix}"), + f"down_proj{suffix}": local(f"down_proj{suffix}"), + } + return { - "gate_up_proj": local(up), - "down_proj": local("down_proj"), - "gate_up_proj_scale_inv": local(f"{up}_scale_inv"), - "down_proj_scale_inv": local("down_proj_scale_inv"), - "gate_up_proj_global_scale": local(f"{up}_global_scale"), - "down_proj_global_scale": local("down_proj_global_scale"), - "gate_up_proj_bias": local(f"{up}_bias"), - "down_proj_bias": local("down_proj_bias"), + **both(""), + **both("_scale_inv"), + **both("_weight_global_scale"), + **both("_input_global_scale"), + **both("_bias"), + # a supported activation NAME is fused into the gate_up epilogue; any other callable + # leaves that GEMM plain and runs on the host between the two "act_fn": act_fn, + "gate": module.has_gate, "swiglu_alpha": module.swiglu_alpha, "swiglu_limit": module.swiglu_limit, "activation_format": module.activation_format, - "gate": module.has_gate, + # a supported norm NAME is folded into the routing-weighted reduce, one pass ahead of it + # computing the rows' rsqrt; any other callable is applied to the routed rows before it + "post_expert_norm": post_expert_norm, + "post_expert_norm_weight": norm_weight, + "post_expert_norm_eps": norm_eps, } @@ -685,19 +732,19 @@ def __init__( self.config = config self.has_bias = has_bias self.has_gate = has_gate - # the model's own gate|up row order (transformers' experts flag): stacked ``[gate; up]`` or, - # like GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]`` — decides whether the - # loader interleaves - self.is_concatenated = is_concatenated self.block_size = block_size - self.hidden_dim = config.hidden_size + # The model's own gate|up row order (transformers' experts flag): stacked ``[gate; up]`` + # or, like GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``. It describes the + # CHECKPOINT only — what this module holds is `holds_interleaved_gate_up` below. + self.is_concatenated = is_concatenated self.activation_format = activation_format self.activation_scheme = activation_scheme - self.swiglu_alpha = getattr(config, "swiglu_alpha", None) - self.swiglu_limit = getattr(config, "swiglu_limit", None) + self.hidden_dim = _first_attr(config, "moe_hidden_size", "hidden_size") self.num_experts = _first_attr(config, "num_local_experts", "num_experts") self.intermediate_dim = _first_attr(config, "moe_intermediate_size", "intermediate_size") self.act_fn_name = _first_attr(config, "hidden_activation", "hidden_act") + self.swiglu_alpha = getattr(config, "swiglu_alpha", None) + self.swiglu_limit = getattr(config, "swiglu_limit", None) self.act_fn = ACT2FN[self.act_fn_name] # Expert weight storage is declared by the QUANT config's format key, not the model @@ -713,12 +760,13 @@ def __init__( # arg would silently drop the nvfp4 global at every forward) self.has_global_scale = fmt.has_global_scale if has_global_scale is None else has_global_scale - # what the module HOLDS, as the loader's conversion ops deliver it (and their reverses + # What the module HOLDS, as the loader's conversion ops deliver it (and their reverses # restore for the checkpoint): gate|up rows in the kernels' interleaved order unless the # backend packs gate|up itself (DeepGEMM Mega MoE), block scales swizzled where the - # scaled-MMA reads them + # scaled-MMA reads them. `is_concatenated` above is the checkpoint's own order: the + # loader interleaves where the two disagree, and everything at runtime reads this one. impl = getattr(config, "_experts_implementation", None) - self._gate_up_interleaved = self.has_gate and impl != "deepgemm_megamoe" + self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" swizzled = _swizzles_scales(config, fmt, activation_format) up_name = "gate_up_proj" if self.has_gate else "up_proj" @@ -738,19 +786,47 @@ def __init__( setattr(self, f"{proj}_bias", nn.Parameter(torch.empty(self.num_experts, rows))) else: self.register_parameter(f"{proj}_bias", None) - # NVFP4 two-level: per-expert fp32 globals the kernels recover on the accumulator + # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator. The + # weight's is one per expert — a gate|up stack whose halves modelopt calibrated + # separately is merged to that at load (`FineGrainedWeightGlobals`), SwiGLU being + # linear in the up half. The activation's is the checkpoint's calibrated + # `input_scale`: the quant divides by it and the accumulator multiplies it back. The + # gate_up's is ONE value (its rows are the hidden states, quantized once, before + # routing); the down's is per expert, because its rows ARE per expert — the gate_up + # epilogue requantizes each against its own expert's value. if self.has_global_scale: - setattr(self, f"{proj}_global_scale", nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32))) + setattr( + self, + f"{proj}_weight_global_scale", + nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)), + ) else: - self.register_parameter(f"{proj}_global_scale", None) + self.register_parameter(f"{proj}_weight_global_scale", None) + if self.has_global_scale and self.activation_format != "bf16": + setattr( + self, + f"{proj}_input_global_scale", + nn.Parameter(torch.ones(1 if proj == up_name else self.num_experts, dtype=torch.float32)), + ) + else: + self.register_parameter(f"{proj}_input_global_scale", None) if self.activation_scheme == "static": self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) self.down_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) + # The model's per-expert output norm and the name it gave the form (its + # `use_experts_implementation(post_expert_norm=...)`), both carried over by the swap. + # A submodule slot, not a class attribute: a class attribute of the same name shadows + # what `nn.Module.__setattr__` registers, so the norm would never be read back. + self.post_expert_norm = None + self.post_expert_norm_name = None + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: - # the GEMM's output columns alternate gate/up; mirrors the fused epilogue's - # `split_gate_up`, which the fused/unfused parity tests compare against + # the interleaved rows give alternating output columns, mirroring the fused epilogue's + # `split_gate_up` that the fused/unfused parity tests compare against. Only the chains + # that hold that layout come through here — DeepGEMM Mega MoE, the one backend keeping + # the stack, runs its own forward end to end. gate, up = gate_up[..., 0::2], gate_up[..., 1::2] if self.swiglu_alpha is not None: # Clamped SwiGLU-OAI gate (same math as the model's non-quantized experts). @@ -795,6 +871,8 @@ def forward( self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None ) proj_out = self.linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) + if self.post_expert_norm is not None: + proj_out = self._apply_post_norm(proj_out) routing_weights = top_k_weights[token_idx, top_k_pos, None] weighted_out = proj_out * routing_weights.to(proj_out.dtype) final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) @@ -812,7 +890,9 @@ def linear( if weight.element_size() > 1: return F.linear(input, weight, bias) scale = getattr(self, f"{proj}_scale_inv") - global_scale = getattr(self, f"{proj}_global_scale") + weight_globals = getattr(self, f"{proj}_weight_global_scale") + input_globals = getattr(self, f"{proj}_input_global_scale") + weight_global = None if weight_globals is None else weight_globals[expert_idx] return finegrained_linear( input, weight, @@ -821,7 +901,12 @@ def linear( bias=bias, activation_scale=activation_scale, allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=global_scale[expert_idx] if global_scale is not None else None, + weight_global_scale=weight_global, + input_global_scale=( + None + if input_globals is None or self.activation_format == "bf16" + else input_globals.reshape(-1)[expert_idx if input_globals.numel() > 1 else 0].reshape(1) + ), activation_format=self.activation_format, ) @@ -940,6 +1025,21 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No if getattr(module, "alpha", None) is not None and new_module.swiglu_alpha is None: new_module.swiglu_alpha = module.alpha new_module.swiglu_limit = getattr(module, "limit", None) + # A per-expert output norm belongs to the model, not the quantization: carry the + # submodule, the name, and the model's own `_apply_post_norm` over, or the swap + # would drop the norm or quietly substitute different math for applying it. The + # method is rebound to THIS module, which holds the same submodule — binding the + # original's would pin the unquantized weights alive. + if getattr(module, "post_expert_norm", None) is not None: + apply_post_norm = getattr(type(module), "_apply_post_norm", None) + if apply_post_norm is None: + raise TypeError( + f"{type(module).__name__} holds a post-expert norm but does not define " + "`_apply_post_norm(self, expert_out)`; see `use_experts_implementation`." + ) + new_module.post_expert_norm = module.post_expert_norm + new_module.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) + new_module._apply_post_norm = MethodType(apply_post_norm, new_module) elif type(module) is nn.Linear: # Vanilla `nn.Linear` → standard FineGrainedLinear swap. new_module = FineGrainedLinear( @@ -1026,10 +1126,9 @@ def _checkpoint_scale_container(model) -> torch.dtype | None: class FineGrainedInterleaveGateUp(ConversionOps): """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order (weight, - scale grid and bias share the row axis) — the core ``Interleave`` along dim 1. Decided off the - model's experts modules, on load and on save alike: skipped when the model's own rows already - come interleaved (``is_concatenated=False``, GPT-OSS) or the backend holds them stacked - (``_gate_up_interleaved`` False, DeepGEMM Mega MoE).""" + scale grid and bias share the row axis) — the core ``Interleave`` along dim 1. Runs on load and + on save alike wherever the checkpoint's row order (``is_concatenated``) differs from the one + the experts hold (``holds_interleaved_gate_up``).""" def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer @@ -1038,7 +1137,7 @@ def __init__(self, hf_quantizer=None, inverse: bool = False): def convert(self, input_dict, model=None, target_patterns=None, **kwargs): input_dict = _keyed_by_target(input_dict, target_patterns) experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None - if experts is None or not experts.is_concatenated or not experts._gate_up_interleaved: + if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): return input_dict interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved return {key: interleave.convert({key: value}, None, [key])[key] for key, value in input_dict.items()} @@ -1163,29 +1262,108 @@ def convert(self, input_dict, target_patterns=None, **kwargs): return out -class FineGrainedFuseEqualGlobals(ConversionOps): - """Reduce modelopt's per-half second-level globals to the single per-expert global the - stacked gate_up GEMM applies: the halves are calibrated together and ship bit-identical - (verified across the GLM-5.2-NVFP4 checkpoint), so this asserts equality and keeps one. - Also handles the single-source (down_proj) case, where it just flattens to ``(E,)``.""" +class FineGrainedWeightGlobals(ConversionOps): + """modelopt's second-level globals, and the calibrated ``input_scale`` they have to travel + with, as the kernels index them: one fp32 global per expert per projection. + + A gate|up stack arrives with TWO per expert, because modelopt calibrates the two projections + separately — they differ by up to 1.6x per expert on Muse-Spark, and agree on GLM-5.2. The + stack keeps the gate's and the up half's folds onto the DOWN projection, which is why this one + op owns every global of a layer: ``inter = silu(gate) * up`` is linear in the up half, so + running the stack on the gate's global alone scales the expert's output by ``1 / ratio``, and + the down's weight global scales it back. The down's calibrated input global moves the other + way, because it is what the fused epilogue normalizes the intermediate by before requantizing + it: without that the requant drifts off the calibrated range and a large enough ratio pushes + the e4m3 block scales past their 448 ceiling. + + Rescaling the up half's e4m3 block scales instead — the other way to merge — would re-round + them against codes already chosen for the old global, moving every value in a block by up to + 6%. + + Sources are either the per-expert keys (``experts.*.gate_proj.weight_scale_2``, merged into a + list per expert) or one stacked ``(E, 2)`` tensor per layer (the fused vLLM layout).""" + + supports_many_to_many = True def __init__(self, hf_quantizer=None): self.hf_quantizer = hf_quantizer - def convert(self, input_dict, model=None, full_layer_name=None, missing_keys=None, **kwargs): - stacks = [] - for v in input_dict.values(): - v = torch.stack(v, dim=0) if isinstance(v, list) else v - stacks.append(v.reshape(-1).float()) - first = stacks[0] - for other in stacks[1:]: - if not torch.equal(first, other): - raise ValueError( - f"modelopt per-half globals differ for {full_layer_name} — the stacked " - "gate_up GEMM applies one global per expert; this checkpoint needs a " - "block-scale refold, which is not implemented." - ) - return {full_layer_name: first} + def convert(self, input_dict, target_patterns=None, model=None, **kwargs): + sources = defaultdict(dict) + for key, value in input_dict.items(): + sources[_global_role(key)][key] = value + globals_ = {} + for target in target_patterns: + if role_sources := sources.get(_global_role(target)): + globals_[target] = _stack_globals(role_sources) + targets = {_global_role(target): target for target in globals_} + stack = targets.get(("gate_up", "weight")) + + # a stack calibrated per half keeps the gate's global; the up half's leaves as a ratio, + # which the down's weight global takes on (scaling the expert output back up) and its + # input global gives back (keeping the requantized intermediate on the calibrated range) + if stack is not None and globals_[stack].shape[1] == 2: + self._assert_no_gate_up_bias(model) + gate, up = globals_[stack][:, 0], globals_[stack][:, 1] + globals_[stack] = gate + ratio = (up / gate)[:, None] + for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): + if target := targets.get(role): + globals_[target] = globals_[target] * factor + return {target: value.reshape(-1).contiguous() for target, value in globals_.items()} + + @staticmethod + def _assert_no_gate_up_bias(model) -> None: + """A gate|up bias is added AFTER the global, so merging the halves would have to divide + its up rows by the same ratio — in the bias converter, which the layout ops own. No + checkpoint pairs the two, so refuse rather than carry a silent half-correction.""" + experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None + if experts is not None and experts.has_bias: + raise NotImplementedError( + "this checkpoint calibrates the gate|up halves separately AND carries an expert " + "bias; the up half's bias would have to be folded with the globals." + ) + + +def _global_role(key: str) -> tuple[str, str]: + """What a global-scale key holds, for a checkpoint key and a module target alike: which + projection it belongs to, and which of modelopt's two levels it is — the weight's second-level + global (``weight_scale_2``) or the calibrated activation one (``input_scale``).""" + projection = "down" if "down_proj" in key or "w2" in key else "gate_up" + level = "input" if "input_scale" in key or "input_global" in key else "weight" + return projection, level + + +def _stack_globals(sources: dict) -> torch.Tensor: + """One ``(E, calibrated projections)`` fp32 tensor from a layer's globals, gate column first: + two sources are the stack's halves, one source is the fused ``(E, 2)`` pair or a single + projection's ``(E,)``.""" + columns = [] + for key, value in sorted(sources.items(), key=lambda kv: 0 if "gate_proj" in kv[0] else 1): + value = torch.stack(value, dim=0) if isinstance(value, list) else value + columns.append(value.float().reshape(value.shape[0], -1)) + return torch.cat(columns, dim=1) + + +class FineGrainedInputGlobals(ConversionOps): + """modelopt's calibrated ``input_scale`` in the layout the module holds: one value for the + gate_up (its rows are the hidden states, quantized once BEFORE routing, so the halves and + experts reduce to their max — the global is a split of the block scale, exact for any value + the accumulator multiplies back), and one per expert for the down, whose rows are per expert + by construction.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, full_layer_name=None, **kwargs): + values = [] + for value in input_dict.values(): + value = torch.stack(value, dim=0) if isinstance(value, list) else value + values.append(value.float()) + stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) + per_expert = stacked.reshape(stacked.shape[0], -1).amax(dim=1) + one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") + return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} class FineGrainedQuantize(ConversionOps): @@ -1247,7 +1425,12 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc scale = load_finegrained_kernel().swizzle_mx_scales(scale) out = {key: weight, f"{prefix}_scale_inv": scale} if global_scale is not None: - out[f"{prefix}_global_scale"] = global_scale + out[f"{prefix}_weight_global_scale"] = global_scale + # activations are quantized dynamically against the block scales alone here — there is no + # calibration pass — so the module's activation globals are the identity + held_input = getattr(module, holder[1].replace("_scale_inv", "_input_global_scale"), None) + if held_input is not None: + out[f"{prefix}_input_global_scale"] = torch.ones_like(held_input) return out def _config_block(self, value: torch.Tensor) -> tuple[int, int]: diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index fb3b736c0301..504e9a0ce775 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -161,6 +161,10 @@ def batched_mm_experts_forward( proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed ) # (S, hidden_dim) + # Normalize each expert application, where such a norm is defined, before it is weighted + if self.post_expert_norm_name is not None: + proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) + # Apply routing weights weighted_out = proj_out * sample_weights.unsqueeze(-1) # (S, hidden_dim) @@ -464,6 +468,10 @@ def grouped_mm_experts_forward( # Same: zero the uninitialized sentinel-tail rows. proj_out = proj_out.masked_fill(sentinel_mask, 0.0) + # Normalize each expert application, where such a norm is defined, before it is weighted + if self.post_expert_norm_name is not None: + proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) + # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) @@ -531,6 +539,7 @@ def use_experts_implementation( is_transposed: bool = False, has_bias: bool = False, has_gate: bool = True, + post_expert_norm: str | None = None, ) -> type[torch.nn.Module]: """Decorator to modify experts class to support different experts implementations. @@ -549,6 +558,12 @@ def use_experts_implementation( has_gate (`bool`, *optional*, defaults to `True`): Whether the experts use a gating mechanism or not. Whether it has gate_up_proj weights or just up_proj weights. + post_expert_norm (`str`, *optional*): + Names the normalization the experts apply to the down output before the routing + weights, when they apply one — the way `hidden_act` names the activation + `_apply_gate` applies. Naming one requires the class to define + `_apply_post_norm(self, expert_out)`, which is where its own math lives; a backend + that implements the named form can fuse it instead of calling that. Returns: `type[torch.nn.Module]`: The modified experts class. @@ -566,6 +581,7 @@ def __init__(self, config, *args, **kwargs): self.has_bias = has_bias self.is_transposed = is_transposed self.is_concatenated = is_concatenated + self.post_expert_norm_name = post_expert_norm @wraps(original_forward) def forward(self, *args, **kwargs): @@ -575,6 +591,13 @@ def forward(self, *args, **kwargs): if not hasattr(experts_class, "_apply_gate"): experts_class._apply_gate = _default_apply_gate + if post_expert_norm is not None and not hasattr(experts_class, "_apply_post_norm"): + raise TypeError( + f"{experts_class.__name__} names a post-expert norm ({post_expert_norm!r}) but does not " + "define `_apply_post_norm(self, expert_out)`, which applies it to one expert " + "application's rows before the routing weights." + ) + experts_class.__init__ = __init__ experts_class.forward = forward return experts_class diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index c55fd4e0b604..21cd0862664a 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -136,6 +136,13 @@ def _process_model_before_weight_loading( replace_with_finegrained_layer( model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config ) + if getattr(self.quantization_config, "activation_format", None) == "bf16": + # Weight-only: the experts hold no activation global, so a calibrated checkpoint's + # `input_scale` has no slot to load into. It is not an unexpected key, it is one this + # run has no use for — a later W4A4 run of the same checkpoint reads it. + model._keys_to_ignore_on_load_unexpected = set(model._keys_to_ignore_on_load_unexpected or []) | { + r"input_scale$" + } def _process_model_after_weight_loading(self, model, **kwargs): from ..integrations.finegrained import disable_deepgemm_on_multi_device @@ -232,16 +239,34 @@ def get_weight_conversions(self): ] return converters - # modelopt NVFP4 (per-expert gate/up, E4M3 group-16 `weight_scale`, fp32 `weight_scale_2` - # globals): the arch plan stacks the expert weights; these stack the scales and globals the - # same way. Only the routed experts are quantized, so no dense converters — a generic - # `weight_scale*` rename would run before converter collection and mangle these keys. - # Calibrated `input_scale` keys stay unconsumed (dynamic act quant) and surface as unexpected. + # modelopt NVFP4: E4M3 group-16 `weight_scale`, fp32 `weight_scale_2` weight globals and + # calibrated fp32 `input_scale` activation globals, in EITHER of the two layouts a modelopt + # checkpoint ships them in — one tensor per expert per projection (GLM-5.2-NVFP4) or one + # already-stacked tensor per layer (the fused vLLM layout, Muse-Spark). Both sets are + # returned: a checkpoint matches one of them and the other never fires. The arch plan + # stacks the expert WEIGHTS; these bring the scales and globals to the same layout, and + # `_with_expert_layout_ops` adds the gate|up interleave + swizzle to the block scales. + # Only the routed experts are quantized, so no dense converters — a generic `weight_scale*` + # rename would run before converter collection and mangle these keys. if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: + # A weight-only run keeps activations bf16, so its experts hold no activation global + # and the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys + # to the unexpected-key filter below rather than converting them onto a module slot + # that does not exist. + calibrated = getattr(self.quantization_config, "activation_format", None) != "bf16" from ..core_model_loading import Concatenate, MergeModulelist - from ..integrations.finegrained import FineGrainedFuseEqualGlobals + from ..integrations.finegrained import ( + FineGrainedInputGlobals, + FineGrainedScaleContainer, + FineGrainedViewPackedInt8, + FineGrainedWeightGlobals, + ) - return [ + # MergeModulelist is what stamps each source with its expert index, which is what lets + # expert parallelism keep only this rank's experts (see `tensor_idx` in + # core_model_loading). Without it every rank collects all E values and the forward + # asserts on the per-expert count. + per_expert = [ WeightConverter( source_patterns=[ "mlp.experts.*.gate_proj.weight_scale$", @@ -255,24 +280,83 @@ def get_weight_conversions(self): target_patterns="mlp.experts.down_proj_scale_inv", operations=[MergeModulelist(dim=0)], ), + # every second-level global of a layer in ONE converter: the gate|up stack's two + # calibrated halves merge into one global per expert by folding the up half's onto + # the down projection, so the weight globals and the down's input scale have to be + # decided together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ "mlp.experts.*.gate_proj.weight_scale_2", "mlp.experts.*.up_proj.weight_scale_2", + "mlp.experts.*.down_proj.weight_scale_2", + *(["mlp.experts.*.down_proj.input_scale"] if calibrated else []), ], - target_patterns="mlp.experts.gate_up_proj_global_scale", - # MergeModulelist is what stamps each source with its expert index, which - # is what lets expert parallelism keep only this rank's experts (see - # `tensor_idx` in core_model_loading). Without it every rank collects all - # E globals and the forward asserts on the per-expert count. - operations=[MergeModulelist(dim=0), FineGrainedFuseEqualGlobals(self)], + target_patterns=[ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[MergeModulelist(dim=0), FineGrainedWeightGlobals(self)], ), + ] + if calibrated: + per_expert.append( + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.input_scale", + "mlp.experts.*.up_proj.input_scale", + ], + target_patterns="mlp.experts.gate_up_proj_input_global_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputGlobals(self)], + ) + ) + fused = [ + # the container cast the module's scale dtype needs; `_with_expert_layout_ops` + # adds the interleave and the swizzle on top (it re-adds the cast, which is + # idempotent) + WeightConverter( + source_patterns=rf"experts\.{proj}_weight_scale$", + target_patterns=f"experts.{proj}_scale_inv", + operations=[FineGrainedScaleContainer(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + fused += [ + # the same one-converter rule as the per-expert layout above, over the keys the + # fused layout ships WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight_scale_2", - target_patterns="mlp.experts.down_proj_global_scale", - operations=[MergeModulelist(dim=0), FineGrainedFuseEqualGlobals(self)], + source_patterns=[ + r"experts\.gate_up_proj_weight_scale_2$", + r"experts\.down_proj_weight_scale_2$", + *([r"experts\.down_proj_input_scale$"] if calibrated else []), + ], + target_patterns=[ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[FineGrainedWeightGlobals(self)], ), ] + if calibrated: + fused.append( + WeightConverter( + source_patterns=r"experts\.gate_up_proj_input_scale$", + target_patterns="experts.gate_up_proj_input_global_scale", + operations=[FineGrainedInputGlobals(self)], + ) + ) + # the fused expert WEIGHTS take the packed uint8 -> int8 view the `.weight`-anchored + # rule below gives the per-expert layout; the layout ops ride on top of both + fused += [ + WeightConverter( + source_patterns=rf"experts\.{proj}$", + target_patterns=f"experts.{proj}", + operations=[FineGrainedViewPackedInt8(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + return per_expert + fused if self.pre_quantized and self.quantization_config.dequantize: return [ diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 193da851d2a4..ca97f0d46e54 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1736,10 +1736,11 @@ def __init__( if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") # NVIDIA modelopt exports (quant_method "modelopt"): translate the payload — - # quant_algo names the format (NVFP4 is the one the finegrained kernels serve), - # ``ignore`` is a glob-style skip list, and the calibrated per-tensor - # ``input_scale``/``kv_cache_scheme`` entries are dropped for now (activations run - # the kernels' dynamic quant, KV cache stays in the compute dtype). + # quant_algo names the format (NVFP4 is the one the finegrained kernels serve), the + # skip list is a glob-style ``ignore`` (or ``exclude_modules``, the name modelopt's + # own ``hf_quant_config.json`` uses), and ``kv_cache_scheme`` is dropped (the KV cache + # stays in the compute dtype). The calibrated ``input_scale`` TENSORS are consumed by + # the loader's converters — the kernels quantize activations against them. if str(self.quant_method) == "modelopt" or kwargs.get("quant_algo") is not None: quant_algo = kwargs.pop("quant_algo", None) kwargs.pop("config_groups", None) @@ -1748,7 +1749,7 @@ def __init__( if quant_algo != "NVFP4": raise ValueError(f"modelopt checkpoints are supported for quant_algo='NVFP4' only; got {quant_algo!r}") self.activation_format = activation_format or "nvfp4" - ignore = kwargs.pop("ignore", None) + ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) if modules_to_not_convert is None and ignore is not None: # modelopt ships glob-style subtree entries ("model.layers.0*", # "model.layers.1.*"); translate each to a BOUNDED regex — a naive @@ -1756,6 +1757,8 @@ def __init__( # layers 10-19 (and ``should_convert_module``'s bare-prefix clause # would even match "layers.16" from "layers.1") + # (``exclude_modules`` entries are bare subtrees — "self_attn", "layers.0." — + # which land on the same bounded form) def _subtree_regex(glob): prefix = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") return re.escape(prefix) + r"(\..*)?$" diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index a50537e6447c..d2b7665d3b2d 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -77,6 +77,7 @@ def _fake_bundle(): ): setattr(kernel, name, rec._op(name)) kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") + kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") return kernel, rec @@ -240,8 +241,12 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): self.assertTrue(linear.call_args_list) for call in linear.call_args_list: self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) - self.assertEqual(call.kwargs["weight_global_scale"].ndim, 0) # this expert's NVFP4 global + # this expert's NVFP4 weight global: one value for the down, and (1, 2) for the + # gate|up stack, whose halves modelopt calibrates separately + self.assertIn(tuple(call.kwargs["weight_global_scale"].shape), {(), (1, 2)}) self.assertEqual(call.kwargs["activation_format"], "bf16") + # weight-only: the activations stay bf16, so no activation global is quantized against + self.assertIsNone(call.kwargs["input_global_scale"]) for call in rec.calls["matmul_2d"]: self.assertIsNotNone(call.kwargs["b_global_scale"]) self.assertEqual(call.kwargs["activation_format"], "bf16") @@ -260,7 +265,7 @@ def test_batched_marshalling(self): # one scale tensor per projection: the affine Parameter, or the swizzled cache when the # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) - self.assertIsNone(call.kwargs["gate_up_proj_global_scale"]) + self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name self.assertIs(call.kwargs["gate"], True) self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format @@ -296,8 +301,98 @@ def test_nvfp4_experts_thread_per_expert_globals(self): with p1, p2, p3, p4: fg.finegrained_batched_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_batched"] - self.assertIs(call.kwargs["gate_up_proj_global_scale"], m.gate_up_proj_global_scale) - self.assertIs(call.kwargs["down_proj_global_scale"], m.down_proj_global_scale) + self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) + self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) + # one weight global per expert on both projections: a gate|up stack calibrated per half + # is merged to that at load, so nothing downstream carries the pair + self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) + # the calibrated activation globals (the checkpoint's `input_scale`): one value for the + # gate_up, whose rows are quantized once before routing, and one per expert for the down, + # whose rows the gate_up epilogue requantizes per expert + self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) + self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) + self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) + self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) + + def test_weight_only_experts_hold_no_activation_global(self): + """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation + quant for a calibrated global to normalize: the module never allocates one, and the ops + would refuse a two-level path without it.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIsNone(m.gate_up_proj_input_global_scale) + self.assertIsNone(m.down_proj_input_global_scale) + self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) + self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) + # the WEIGHT globals are the format's own second level and stay + self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) + + def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): + """A model whose experts norm the down output before the routing weights (Muse-Spark): + the swap carries the norm over, the kernel chain runs it on the routed rows, and the + eager loop applies it per expert application — the same place the reference forwards do. + A form the kernels do not implement rides as the module's own ``_apply_post_norm``.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + norm = torch.nn.LayerNorm(m.hidden_dim) + m.post_expert_norm, m._apply_post_norm = norm, norm + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m._apply_post_norm) + + class _Recording(torch.nn.Module): + def __init__(self): + super().__init__() + self.rows = 0 + + def forward(self, x): + self.rows += x.shape[0] + return x + + m.post_expert_norm = m._apply_post_norm = _Recording() + with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + m(*self._route()) + self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") + + def test_naming_a_post_expert_norm_requires_the_hook_that_applies_it(self): + """The shared experts interface: a class that names a post-expert norm has to define + `_apply_post_norm`, where its own math lives — there is no default passthrough to fall + back on, the way `act_fn` has no default activation.""" + from transformers.integrations.moe import use_experts_implementation + + with self.assertRaises(TypeError): + + @use_experts_implementation(post_expert_norm="rms_norm") + class _NoHook(torch.nn.Module): + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states + + def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): + """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels + when they implement it, with the weight and epsilon the fused form needs, so the chain + folds it into its reduce instead of calling the module.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + m.post_expert_norm = m._apply_post_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4) + p1, p2, p3, p4 = _loaded(kernel) + for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): + m.post_expert_norm_name = name + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + call = rec.calls["moe_fused_grouped"][-1] + if fused: + self.assertEqual(call.kwargs["post_expert_norm"], name) + self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) + self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) + else: + self.assertEqual(call.kwargs["post_expert_norm"], m._apply_post_norm) + self.assertIsNone(call.kwargs["post_expert_norm_weight"]) def test_biases_ride_the_kernel_chain(self): kernel, rec = _fake_bundle() @@ -490,7 +585,7 @@ def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) - self.assertTrue(experts._gate_up_interleaved) + self.assertTrue(experts.holds_interleaved_gate_up) # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a # float32 grid there; GPT-OSS's config has no scale_fmt at all) _, experts = self._experts("mxfp4", activation_format="bf16") @@ -518,7 +613,7 @@ def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): def test_megamoe_holds_gate_up_stacked(self): _, experts = self._experts("fp8", impl="deepgemm_megamoe") - self.assertFalse(experts._gate_up_interleaved) + self.assertFalse(experts.holds_interleaved_gate_up) def _op_kernel(self): kernel, _ = _fake_bundle() @@ -772,7 +867,7 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): experts.is_concatenated = False self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) experts.is_concatenated = True - experts._gate_up_interleaved = False + experts.holds_interleaved_gate_up = False self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) @@ -825,7 +920,14 @@ def _modelopt_conversions(self): def test_global_scale_converters_carry_an_expert_index(self): from transformers.core_model_loading import MergeModulelist - globals_converters = [c for c in self._modelopt_conversions() if any("global_scale" in t for t in _targets(c))] + # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked + # tensor per layer has no per-expert sources to stamp — it is sharded like the fused + # expert weight next to it + globals_converters = [ + c + for c in self._modelopt_conversions() + if any("global_scale" in t for t in _targets(c)) and any("experts.*." in p for p in c.source_patterns) + ] self.assertTrue(globals_converters, "no global-scale converter found") for conv in globals_converters: self.assertTrue( @@ -834,6 +936,162 @@ def test_global_scale_converters_carry_an_expert_index(self): "experts and every rank would collect all of them", ) + def test_both_modelopt_scale_layouts_have_converters(self): + """modelopt ships the scales either as one tensor per expert per projection (GLM-5.2) or + as one already-stacked tensor per layer (the fused vLLM layout, Muse-Spark). Both are + converted; the block scale takes the module's layout ops either way.""" + sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} + for suffix in ("weight_scale$", "weight_scale_2", "input_scale"): + self.assertTrue( + any(p.startswith("mlp.experts.*.gate_proj.") and p.endswith(suffix) for p in sources), + f"no per-expert converter for {suffix}", + ) + self.assertTrue( + any("gate_up_proj_" in p and p.endswith(suffix.rstrip("$") + "$") for p in sources), + f"no fused converter for {suffix}", + ) + + def test_one_converter_merges_a_layer_s_globals(self): + """modelopt calibrates the gate|up halves separately. One converter owns every global of + a layer, because merging them to the one-per-expert the kernels take moves the up half's + onto the down projection: SwiGLU is linear in the up half, so the stack keeps the gate's + global and the down's weight global scales the expert output back. The down's calibrated + input scale moves the other way, keeping the requantized intermediate on the range the + checkpoint calibrated.""" + from transformers.integrations.finegrained import FineGrainedWeightGlobals + + gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) + down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) + targets = [ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + "mlp.experts.down_proj_input_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "mlp.experts.*.up_proj.weight_scale_2": up, + "mlp.experts.*.gate_proj.weight_scale_2": gate, + "mlp.experts.*.down_proj.weight_scale_2": down, + "mlp.experts.*.down_proj.input_scale": down_input, + }, + target_patterns=targets, + ) + ratio = up / gate + torch.testing.assert_close(out[targets[0]], gate) + torch.testing.assert_close(out[targets[1]], down * ratio) + torch.testing.assert_close(out[targets[2]], down_input / ratio) + + def test_globals_with_one_calibrated_projection_pass_through(self): + """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a + single global) has nothing to merge, so every global reaches its module unchanged.""" + from transformers.integrations.finegrained import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0], [2.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([3.0, 4.0]), + "experts.down_proj_input_scale": torch.tensor([0.5, 0.25]), + }, + target_patterns=targets, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([3.0, 4.0])) + torch.testing.assert_close(out[targets[2]], torch.tensor([0.5, 0.25])) + + def test_the_fused_layout_ships_both_halves_in_one_tensor(self): + """Muse-Spark's `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two + keys, and merges the same way.""" + from transformers.integrations.finegrained import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 3.0], [2.0, 4.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([5.0 * 3.0, 6.0 * 2.0])) + + def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_down(self): + """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated + `input_scale` reduces to one value; the down's rows are per expert, so its stays per + expert — the requant epilogue normalizes each row by its own expert's value.""" + from transformers.integrations.finegrained import FineGrainedInputGlobals + + up = FineGrainedInputGlobals().convert( + { + "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), + "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), + }, + full_layer_name="mlp.experts.gate_up_proj_input_global_scale", + )["mlp.experts.gate_up_proj_input_global_scale"] + torch.testing.assert_close(up, torch.tensor([3.0])) + down = FineGrainedInputGlobals().convert( + {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, + full_layer_name="mlp.experts.down_proj_input_global_scale", + )["mlp.experts.down_proj_input_global_scale"] + torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) + + +class _InputScaledRMSNorm(torch.nn.Module): + """Muse-Spark's post-expert norm: scale by ``1 + weight`` BEFORE normalizing, so the row's + mean square is taken on the scaled values (not the same function as normalizing first).""" + + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16)) + + def forward(self, x): + out = x.float() * (1.0 + self.weight.float()) + return (out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) + + +class _UnfusableNorm(torch.nn.Module): + """A norm whose math is none of the fused forms (a bias the kernels have no slot for): the + binder has to leave it a module call rather than pick the nearest form.""" + + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim, dtype=torch.bfloat16)) + self.bias = torch.nn.Parameter(torch.full((dim,), 0.5, dtype=torch.bfloat16)) + + def forward(self, x): + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return (normed * self.weight.float() + self.bias.float()).type_as(x) + + +def _nvfp4_two_level(kernel, weight): + """Canonical two-level NVFP4 quant of one matrix: the global is ``amax / (6 * 448)`` — the + smallest one that keeps every e4m3 block scale in range — and the block quant runs on the + normalized values. Returns ``(packed_e2m1, e4m3 block scales, fp32 global)``.""" + global_scale = (weight.abs().amax() / (6.0 * 448.0)).clamp(min=torch.finfo(torch.float32).tiny) + packed, block = kernel.nvfp4_act_quant((weight.float() / global_scale).contiguous()) + return packed.view(torch.int8), block, global_scale + + +def _unpack_e2m1(packed): + """Packed-E2M1 bytes (two codes per byte along the last axis) back to fp32 values.""" + lut = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + device=packed.device, + ) + b = packed.view(torch.uint8) + pairs = torch.stack([lut[(b & 0x0F).long()], lut[(b >> 4).long()]], dim=-1) + return pairs.reshape(*packed.shape[:-1], packed.shape[-1] * 2) + def _targets(conv): t = conv.target_patterns @@ -913,13 +1171,156 @@ def test_on_the_fly_group_formats_round_trip_within_their_floor(self): scale = kernel.unswizzle_mx_scales(scale, 512, scale.shape[2] * 4, num_experts=4) deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) if fmt == "nvfp4": - deq = deq * out["experts.gate_up_proj_global_scale"].reshape(-1, 1, 1) + global_scale = out["experts.gate_up_proj_weight_global_scale"] + deq = deq * global_scale.reshape(-1, 1, 1) else: - self.assertNotIn("experts.gate_up_proj_global_scale", out) + self.assertNotIn("experts.gate_up_proj_weight_global_scale", out) rel = ((deq - original.float()).norm() / original.float().norm()).item() self.assertLess(rel, floor, f"{fmt}: {rel:.3f}") self.assertGreater(rel, 0.0) + def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes=(8.0, 1.0)): + """NVFP4 experts in the layout a modelopt checkpoint delivers: every (expert, half) of the + gate|up stack quantized against its OWN global (``weight_scale_2`` is per projection, and + the two differ per expert on Muse-Spark), the down per expert, plus calibrated activation + globals. Returns the module, the exact dequantized weights (so the reference is the + quantization floor rather than the pre-quant weight), and the dequantization a reader that + applied ONE of the two globals to both halves would get.""" + from transformers.integrations.finegrained import FineGrainedWeightGlobals + + kernel = load_finegrained_kernel() + experts = FineGrainedExperts(cfg, weight_format="nvfp4", activation_format=activation_format).cuda() + dequantized, single_global, checkpoint_globals = {}, {}, {} + for proj, rows, cols in ( + ("gate_up_proj", 2 * cfg.intermediate_size, cfg.hidden_size), + ("down_proj", cfg.hidden_size, cfg.intermediate_size), + ): + reference = torch.randn(cfg.num_local_experts, rows, cols, device="cuda") + if proj == "gate_up_proj": # the halves interleave by row; draw them far apart + reference[:, 0::2] *= magnitudes[0] + reference[:, 1::2] *= magnitudes[1] + packed = torch.empty(cfg.num_local_experts, rows, cols // 2, dtype=torch.int8, device="cuda") + block = torch.empty(cfg.num_local_experts, rows, cols // 16, dtype=torch.float8_e4m3fn, device="cuda") + halves = 2 if proj == "gate_up_proj" else 1 + globals_ = torch.empty(cfg.num_local_experts, halves, device="cuda") + deq, deq_one = torch.empty_like(reference), torch.empty_like(reference) + for e in range(cfg.num_local_experts): + for h in range(halves): + rows_h = slice(h, None, halves) + q, b, g = _nvfp4_two_level(kernel, reference[e, rows_h]) + packed[e, rows_h], block[e, rows_h], globals_[e, h] = q, b, g + values = _unpack_e2m1(q) * b.float().repeat_interleave(16, dim=-1) + deq[e, rows_h] = values * g + if halves == 2: + # the control: BOTH halves read the up half's global, which is what a kernel + # that ignored the layout does. It has to be the up one — putting the gate's + # global on both scales the expert's whole output uniformly, and the + # post-expert RMSNorm divides exactly that out; a wrong GATE scale goes + # through the SiLU and survives. + for h in (0, 1): + rows_h = slice(h, None, 2) + deq_one[e, rows_h] = deq[e, rows_h] * (globals_[e, 1] / globals_[e, h]) + scale = block if getattr(experts, f"{proj}_scale_inv").dim() != 5 else kernel.swizzle_mx_scales(block) + setattr(experts, proj, torch.nn.Parameter(packed, requires_grad=False)) + setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(scale, requires_grad=False)) + checkpoint_globals[f"experts.{proj}_weight_scale_2"] = globals_ + dequantized[proj], single_global[proj] = deq, deq_one + # calibrated input_scale: one value for the gate_up (quantized before routing), one per + # expert for the down (its rows are per expert). A weight-only module holds neither — + # nothing there quantizes activations — so the checkpoint's keys stay unused. + if activation_format != "bf16": + experts.gate_up_proj_input_global_scale = torch.nn.Parameter( + torch.full((1,), 0.05, device="cuda"), requires_grad=False + ) + checkpoint_globals["experts.down_proj_input_scale"] = torch.linspace( + 0.02, 0.08, cfg.num_local_experts, device="cuda" + ) + single_global["down_proj"] = dequantized["down_proj"] # the down has one global either way + # The globals go in through the loader's own converter, so the module holds exactly what a + # real checkpoint puts there: one per expert, with the gate|up pair merged onto the down. + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + for target, value in FineGrainedWeightGlobals().convert(checkpoint_globals, target_patterns=targets).items(): + setattr(experts, target.split(".")[-1], torch.nn.Parameter(value, requires_grad=False)) + return experts, dequantized, single_global + + def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): + """The experts chain in torch over the dequantized weights: gate|up (interleaved columns), + SwiGLU, down, the optional per-expert output norm, then the routing-weighted sum.""" + out = torch.zeros_like(x, dtype=torch.float32) + for token in range(x.shape[0]): + for slot in range(idx.shape[1]): + e = int(idx[token, slot]) + pre = x[token].float() @ dequantized["gate_up_proj"][e].t() + inter = torch.nn.functional.silu(pre[0::2]) * pre[1::2] + row = inter @ dequantized["down_proj"][e].t() + if post_norm is not None: + row = post_norm(row.to(x.dtype)).float() + out[token] += row * float(wts[token, slot]) + return out + + def test_modelopt_nvfp4_experts_merge_their_globals_and_fuse_the_post_norm(self): + """The Muse-Spark checkpoint shape end to end through the integration: a gate|up stack + calibrated per half (merged at load), calibrated activation globals, and a per-expert + output norm the binder fuses, on all three forwards. Weight-only pins the numbers against the dequantized weights (nothing else + rounds); the W4A4 chain rides the 4-bit activation floor, so it is pinned by being far + closer to the right globals than to one global folded over both halves — what a kernel + that ignored the layout computes.""" + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + x = torch.randn(8, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) * 0.1 + idx = torch.randint(0, cfg.num_local_experts, (8, 2), device="cuda", dtype=torch.long) + wts = torch.rand(8, 2, device="cuda", dtype=torch.bfloat16) + forwards = ( + ("grouped", fg.finegrained_grouped_mm_experts_forward), + ("batched", fg.finegrained_batched_mm_experts_forward), + ("eager", lambda module, *args: module(*args)), + ) + + def relative(out, reference): + return ((out.float() - reference).norm() / reference.norm()).item() + + for activation_format, floor in (("bf16", 0.02), (None, 0.35)): + experts, dequantized, one_global = self._nvfp4_experts_modelopt_layout(cfg, activation_format) + reference = self._moe_reference(dequantized, x, idx, wts) + wrong = self._moe_reference(one_global, x, idx, wts) + for name, forward in forwards: + with self.subTest(forward=name, activation_format=activation_format): + out = forward(experts, x, idx, wts) + rel, rel_wrong = relative(out, reference), relative(out, wrong) + self.assertLess(rel, floor, f"{name} diverged from the dequantized reference: {rel:.3f}") + self.assertLess( + 3 * rel, + rel_wrong, + f"{name} is as close to ONE global over both gate|up halves ({rel_wrong:.3f}) " + f"as to the per-half ones ({rel:.3f}) — the fold is not reading the halves", + ) + + # the same chain with the model's per-expert output norm, which runs on the routed + # rows: every forward has to apply it where the reference forwards do. A name the + # kernels implement is folded into the chain's reduce and anything else is the module + # itself, so both have to land on the same numbers as the reference. + for norm, norm_name in ( + (torch.nn.RMSNorm(cfg.hidden_size, device="cuda", dtype=torch.bfloat16), "rms_norm"), + (_InputScaledRMSNorm(cfg.hidden_size).cuda(), "input_scaled_rms_norm"), + # a form the kernels do not implement: the module is called on the rows + (_UnfusableNorm(cfg.hidden_size).cuda(), "a_models_own_norm"), + ): + norm.weight.data.uniform_(0.5, 1.5) + experts.post_expert_norm, experts.post_expert_norm_name = norm, norm_name + experts._apply_post_norm = norm # what the swap carries from the model's class + normed = self._moe_reference(dequantized, x, idx, wts, norm) + self.assertGreater(relative(normed, reference), 0.05, "the norm has to change the output") + for name, forward in forwards: + with self.subTest(forward=name, activation_format=activation_format, norm=norm_name): + rel = relative(forward(experts, x, idx, wts), normed) + self.assertLess(rel, floor, f"{name} diverged with the post-expert norm: {rel:.3f}") + def test_experts_scale_layout_op_keeps_the_forward_correct(self): """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op fills them from the affine grid, the fused forward matches the affine module's, and the From 89267bf3cd9c47c08773f54d47c196cd214157e4 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 09:58:07 -0700 Subject: [PATCH 17/68] finegrained: address the integration review One guard decorator on the DeepGEMM experts forwards instead of four statements per body, and that file's comments cut to what is not already in the code. `prefers_deepgemm_linear` moves to `deepgemm.py`, where its measured verdicts belong, so the dispatcher just asks. `_FineGrainedModule.local(name)` is the one accessor every forward reads operands through, so the DTensor unwrap appears once rather than at seven call sites. `_WeightFormat` is public (it is in `resolve_weight_format`'s signature), the `TRANSFORMERS_FINEGRAINED_NO_SWIZZLE` debug hatch is gone (tests patch the predicate), and `_set_optional_parameter` replaces six register-or-assign blocks. `check_arch` is passed straight through: `kernels` is pinned >= 0.16, where it landed, so the compatibility shim was dead. The scalar-shard branch becomes `_owns_expert`, stated as the general fact it is rather than as an NVFP4 anecdote. `FineGrainedConfig`'s docstring opens with a format table. The frozen fp8 modules' `stacklevel=2` is dropped: at module scope it named importlib rather than the module being deprecated. Co-Authored-By: Claude Opus 5 --- src/transformers/core_model_loading.py | 5 +- .../distributed/sharding_utils.py | 22 +- src/transformers/integrations/deepgemm.py | 120 +++++++---- src/transformers/integrations/finegrained.py | 201 ++++++++---------- .../integrations/finegrained_fp8.py | 1 - src/transformers/integrations/hub_kernels.py | 8 +- .../quantizers/quantizer_finegrained_fp8.py | 1 - src/transformers/utils/quantization_config.py | 28 ++- 8 files changed, 193 insertions(+), 193 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 3918c3800431..0c48db3bfffc 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1702,9 +1702,8 @@ def convert_and_load_state_dict_in_model( if matched_dtype_pattern is not None: _dtype = dtype_plan[dtype_policy_by_group_name[matched_dtype_pattern.lastgroup]] elif empty_param is not None and empty_param.dtype != _dtype and not needs_quantization: - # usually correct when initializing — but a parameter the quantizer's op is about to - # quantize must arrive in the checkpoint's precision, not pre-cast to its quantized - # storage dtype (int8 storage would zero it, float8 would double-round it) + # usually correct when initializing; only exception can be quants (int8 storage + # would zero it, float8 would double-round it) _dtype = empty_param.dtype # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 009001db7945..ca3cb2c2f75a 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -110,18 +110,11 @@ def shard_tensor( (mesh_dim, placement) for mesh_dim, placement in enumerate(self.placements) if hasattr(placement, "dim") ] - # A 0-dim tensor has no axis to shard — every owning rank takes the whole value. Reachable - # in practice: ModelOpt NVFP4 checkpoints ship one `weight_scale_2` / `input_scale` scalar - # per expert projection, thousands per model. Both paths below index the (empty) shape — - # `source[()]` is not accepted by a lazy safetensors slice, `[...]` is — so route scalars - # here: replicated on the dense path, per-expert-ownership-filtered on the MoE path. + # A 0-dim tensor has no axis to shard, so every rank that owns it takes the whole value + # (`source[...]`, since a lazy safetensors slice rejects `source[()]`). if not source_shape: - if tensor_idx is not None: - has_axis0_shard = any(self._normalize_param_dim(placement.dim) == 0 for _, placement in dim_placements) - if has_axis0_shard and not ( - self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size - ): - return None + if tensor_idx is not None and not self._owns_expert(tensor_idx, dim_placements): + return None return source[...].to(device=device, dtype=dtype) # Dense path @@ -331,6 +324,13 @@ def _normalize_param_dim(self, dim: int) -> int: # if dim is negative, it should be normalized to the last axis return dim if dim >= 0 else self.param_ndim + dim + def _owns_expert(self, tensor_idx: int, dim_placements: list) -> bool: + """Whether this rank holds expert `tensor_idx`. True when nothing shards axis 0 — the + experts are replicated, so every rank owns every one of them.""" + if not any(self._normalize_param_dim(placement.dim) == 0 for _, placement in dim_placements): + return True + return self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size + def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTensor: """Wrap `local_tensor` as a DTensor that mirrors `ref`'s mesh, placements, diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index b815fd0c0d30..609b02bd8eec 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -66,12 +66,9 @@ class DeepGEMM: transform_weights_for_mega_moe: Callable get_symm_buffer_for_mega_moe: Callable fp8_fp4_mega_moe: Callable - # M/K-dimension alignment for TMA-based contiguous grouped GEMM. Sourced from - # `get_mk_alignment_for_contiguous_layout()` at load time. The kernel exposes a - # `set_mk_alignment_for_contiguous_layout` setter, but we don't call it: the - # build-time default (128) was empirically the best across MoE workloads - # (bench showed kernel-recommended 240 is slower and 256 doesn't even compile). - # Same stance as vLLM, which caches and never sets it. + # M/K alignment for the TMA contiguous grouped GEMM, read at load time. The kernel's setter + # stays unused: the build-time default (128) measured best across MoE workloads (its own + # recommendation, 240, is slower; 256 does not compile). vLLM takes the same stance. m_alignment: int @@ -222,9 +219,8 @@ def is_deepgemm_loadable(raise_error: bool = False) -> bool: return True -# Cache the loaded bundle but not failures: re-checking each call is cheap and intended, since the env -# can change between attempts. A module global (not `@functools.cache`) avoids Dynamo warning about -# tracing a cache-wrapped function on every compile. +# Successes are cached, failures are not: the env can change between attempts. A module global +# rather than `@functools.cache`, which Dynamo warns about tracing on every compile. _DEEPGEMM: DeepGEMM | None = None @@ -517,23 +513,19 @@ def _dispatch_routed_input( expert_ids = top_k_index.reshape(-1) # (S,) sample_weights = top_k_weights.reshape(-1) # (S,) - # Sort by expert for grouped processing expert_ids_g, perm = torch.sort(expert_ids) sorted_hidden_states_g = hidden_states[perm // num_top_k] sample_weights_g = sample_weights[perm] - # Build the M-grouped padded layout (DeepGEMM contract: each expert's rows - # start on the kernel's M-alignment boundary, sentinels routed past valid - # expert blocks). + # DeepGEMM's contract: each expert's rows start on the kernel's M-alignment boundary, and + # sentinels route past every valid expert block. sorted_to_padded, grouped_layout, total_padded_rows = _build_deepgemm_contiguous_layout( expert_ids_g, num_experts, m_alignment, use_psum_layout ) - # EP sentinel mask is captured before the in-place clamp; used by the post-mask in - # `_combine_routed_output` to zero sentinel rows before the per-token reduction. The clamp - # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls - # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the - # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself. + # Captured before the in-place clamp, which keeps a per-row gather (bias) in bounds — those + # rows the kernel skips anyway. `_combine_routed_output` zeroes them before the per-token + # reduction; the layout above was built from the unclamped ids, so mutating now is safe. sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1) expert_ids_g.clamp_(max=num_experts - 1) return ( @@ -575,6 +567,43 @@ def _combine_routed_output( @deprecate_kwarg("output_dtype", version="v5.16") +def prefers_deepgemm_linear( + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + *, + block_size: list[int] | None, + activation_scale: torch.Tensor | None, + weight_global_scale: torch.Tensor | None, + input_global_scale: torch.Tensor | None, + activation_format: str | None, + allow_deepgemm: bool, +) -> bool: + """Whether a dense fine-grained linear should go to DeepGEMM rather than Triton. + + DeepGEMM is CUDA-only, dynamic-only, SM90+, FP4 or 128x128-block FP8; the combos it would + silently corrupt are refused here rather than attempted. + ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces Triton for this dispatcher only.""" + return ( + # SM100 perf: Triton measures 2.9x (decode) / 1.45x (prefill) over DeepGEMM on the DSV4 + # qkv linear — the 128x128 block-FP8 shape DeepGEMM is otherwise preferred for. + not is_sm100() + # False when the model spans devices: DeepGEMM's kernels are bound to one CUDA context and + # corrupt across them, and the Triton fallback is context-free. + and allow_deepgemm + # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would + # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. + and weight_scale_inv.ndim <= 2 + and is_deepgemm_loadable() + and activation_scale is None + # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format + and weight_global_scale is None + and input_global_scale is None + and activation_format is None + and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) + and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" + ) + + def deepgemm_fp8_fp4_linear( input: torch.Tensor, weight: torch.Tensor, @@ -629,14 +658,35 @@ def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: ) +def deepgemm_experts_guards(*, affine_scales: bool = False, sm100: bool = False): + """State an experts forward's requirements on the module, checked before any kernel work. + + Every arm refuses a per-expert output norm. ``affine_scales`` additionally refuses the + SWIZZLE_32_4_4 scales a module loaded for a triton backend holds, and ``sm100`` fails before + the hub download + JIT when the device cannot serve these dtypes.""" + + def decorate(forward): + @functools.wraps(forward) + def guarded(self, *args, **kwargs): + _assert_no_post_expert_norm(self) + if affine_scales: + _assert_affine_scales(self) + if sm100: + _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) + return forward(self, *args, **kwargs) + + return guarded + + return decorate + + +@deepgemm_experts_guards() def deepgemm_bf16_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: - _assert_no_post_expert_norm(self) - if hidden_states.dtype != torch.bfloat16: raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") @@ -667,7 +717,6 @@ def deepgemm_bf16_experts_forward( up_bias = (self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) if self.has_bias else None down_bias = self.down_proj_bias if self.has_bias else None - # Up projection. up_out_dim = weight_up.shape[-1] if self.is_transposed else weight_up.shape[1] act = _pad_for_deepgemm(sorted_hidden, sorted_to_padded, total_padded_rows) proj_out = torch.empty(total_padded_rows, up_out_dim, device=device, dtype=hidden_states.dtype) @@ -677,7 +726,6 @@ def deepgemm_bf16_experts_forward( proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - # Down projection. out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=hidden_states.dtype) grouped_bf16_matmul(proj_out, weight_down, out, grouped_layout, use_psum_layout=is_sm100()) if self.has_bias: @@ -722,6 +770,7 @@ def _assert_stacked_gate_up(module: torch.nn.Module) -> None: ) +@deepgemm_experts_guards(affine_scales=True, sm100=True) def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -737,15 +786,6 @@ def deepgemm_fp8_fp4_experts_forward( "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) - _assert_no_post_expert_norm(self) - # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as - # affine and turn into garbage — refuse before any kernel work. - _assert_affine_scales(self) - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is - # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the - # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. - _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) - if self.activation_scheme == "static": raise NotImplementedError("DeepGEMM experts dispatch does not support static activation quantization.") if hidden_states.dtype != torch.bfloat16: @@ -782,7 +822,6 @@ def deepgemm_fp8_fp4_experts_forward( ) sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None - # Up projection. act_fp8, act_scales = deepgemm.per_token_cast_to_fp8(sorted_hidden, **cast_kwargs) act_fp8 = _pad_for_deepgemm(act_fp8, sorted_to_padded, total_padded_rows) act_scales = _pad_for_deepgemm(act_scales, sorted_to_padded, total_padded_rows) @@ -797,7 +836,6 @@ def deepgemm_fp8_fp4_experts_forward( ) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - # Down projection. proj_fp8, proj_scales = deepgemm.per_token_cast_to_fp8(proj_out, **cast_kwargs) out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=torch.bfloat16) grouped_fp8_fp4_matmul( @@ -886,6 +924,7 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) +@deepgemm_experts_guards(affine_scales=True, sm100=True) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -912,15 +951,6 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( - `config.swiglu_limit` (optional): SwiGLU clamp; absent → unclamped. """ - _assert_no_post_expert_norm(self) - # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as - # affine and turn into garbage — refuse before any kernel work. - _assert_affine_scales(self) - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is - # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the - # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. - _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) - if self.gate_up_proj.dtype != torch.int8: raise NotImplementedError( f"DeepGEMM Mega MoE requires FP4-packed expert weights (dtype=`int8`), got " @@ -934,10 +964,8 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( deepgemm = load_deepgemm_kernel() - # First-forward one-shot: pack UE8M0 SFs and interleave the L1/L2 weights for UTCCP. - # Kept lazy here (instead of in a quantizer load-time hook) so the megamoe-specific - # setup lives alongside the megamoe forward — `set_experts_implementation` refuses - # to flip in/out of `deepgemm_megamoe` at runtime, so the flag won't go stale. + # One-shot on the first forward, kept here rather than in a load-time hook so the setup sits + # with the forward it serves; the flag cannot go stale because the backend cannot be flipped. if not getattr(self, "_megamoe_transformed", False): setup_megamoe_weights(self) self._megamoe_transformed = True diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 2dced4f365c9..9be49026f2c9 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1,4 +1,4 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,8 +35,8 @@ deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, deepgemm_fp8_fp4_megamoe_experts_forward, - is_deepgemm_loadable, is_sm100, + prefers_deepgemm_linear, ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel from .moe import ExpertsInterface, use_experts_implementation @@ -183,7 +183,7 @@ def _cdiv(a: int, b: int) -> int: @dataclass(frozen=True) -class _WeightFormat: +class WeightFormat: """Storage layout of one quantized weight format — the single source both module classes derive their parameter shapes from. ``scale_dtype`` ``None`` defers to the config's ``scale_fmt`` (block-FP8 ships fp32 or UE8M0 containers; the group formats pin theirs). @@ -196,18 +196,18 @@ class _WeightFormat: has_global_scale: bool = False -def _weight_formats() -> dict[str, _WeightFormat]: +def _weight_formats() -> dict[str, WeightFormat]: return { # block-scaled E4M3, block from the quant config, fp32/UE8M0 scale container - "fp8": _WeightFormat(weight_dtype=_FP8_DTYPE), + "fp8": WeightFormat(weight_dtype=_FP8_DTYPE), # E4M3 values, UE8M0 group-32 scales - "mxfp8": _WeightFormat(weight_dtype=_FP8_DTYPE, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32)), + "mxfp8": WeightFormat(weight_dtype=_FP8_DTYPE, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32)), # packed E2M1 values (2/byte), UE8M0 group-32 scales - "mxfp4": _WeightFormat( + "mxfp4": WeightFormat( weight_dtype=torch.int8, values_per_byte=2, scale_dtype=_get_ue8m0_dtype(), scale_group=(1, 32) ), # packed E2M1 values, E4M3 group-16 block scales, per-tensor/per-expert fp32 global - "nvfp4": _WeightFormat( + "nvfp4": WeightFormat( weight_dtype=torch.int8, values_per_byte=2, scale_dtype=_FP8_DTYPE, @@ -221,7 +221,7 @@ def resolve_weight_format( weight_format: str, scale_fmt: str = "float", block_size: tuple[int, int] | None = None, -) -> tuple[_WeightFormat, torch.dtype, tuple[int, int] | None]: +) -> tuple[WeightFormat, torch.dtype, tuple[int, int] | None]: """``(format, scale_dtype, (sf_gran_n, sf_gran_k))`` for one format name, with the config's ``scale_fmt``/``weight_block_size`` filling the slots the format leaves open.""" formats = _weight_formats() @@ -234,11 +234,17 @@ def resolve_weight_format( return fmt, scale_dtype, fmt.scale_group if fmt.scale_group is not None else block_size +def _set_optional_parameter(module: nn.Module, name: str, tensor: torch.Tensor | None) -> None: + """Register `name` as a Parameter over `tensor`, or as an absent one when `tensor` is None — + a slot a checkpoint may or may not fill (a bias, an NVFP4 global, a static activation scale).""" + module.register_parameter(name, None if tensor is None else nn.Parameter(tensor)) + + def _alloc_expert_proj( num_experts: int, proj_out: int, proj_in: int, - fmt: _WeightFormat, + fmt: WeightFormat, scale_dtype: torch.dtype, scale_group: tuple[int, int] | None, min_scale_out: int = 1, @@ -267,16 +273,13 @@ def _alloc_expert_proj( ) -def _swizzles_scales(config, fmt: _WeightFormat, activation_format: str | None) -> bool: +def _swizzles_scales(config, fmt: WeightFormat, activation_format: str | None) -> bool: """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05 scaled-MMA reads directly (plain row-major forces a per-tile gather that caps the scaled dot below the fp8/fp4 peak): SM100, a triton dispatch (the fused forwards and the eager loop both read it; the DeepGEMM backends read affine scales), a group-scaled format (MX group-32, NVFP4 group-16 — block-FP8's grid never reaches a scaled-MMA, even with UE8M0 scales) and a chain that - quantizes activations (weight-only reads scales per group affinely). - ``TRANSFORMERS_FINEGRAINED_NO_SWIZZLE=1`` keeps every module affine (debug / A/B).""" - if os.environ.get("TRANSFORMERS_FINEGRAINED_NO_SWIZZLE", "0") == "1": - return False + quantizes activations (weight-only reads scales per group affinely).""" return ( getattr(config, "_experts_implementation", None) not in ("deepgemm", "deepgemm_megamoe") and fmt.scale_group is not None @@ -298,14 +301,11 @@ def finegrained_triton_linear( ) -> torch.Tensor: """Triton fine-grained linear: fused act-quant + matmul, then optional bias add. - Serves every weight format the kernel resolves off the tensors themselves — block-FP8 - (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4 (``int8``-packed values; the two-level - per-tensor ``weight_global_scale`` recovers on the accumulator; ``input_global_scale`` is - NVFP4's CALIBRATED activation global, the checkpoint's ``input_scale`` — the quant divides by it - and the accumulator multiplies it back). ``activation_scale=None`` → - dynamic activation quant (inline); a per-tensor scalar → static quant against it. - ``activation_format`` picks the activation format where the weights leave it open (``None`` = - the weight-native choice; ``"bf16"`` = weight-only, no activation quant — W4A16). + The weight format is read off the tensors (see ``WeightFormat``); what the caller decides is + the activation side. ``activation_scale`` ``None`` quantizes inline per token, a scalar + quantizes statically against it, and ``activation_format`` picks the format where the weights + leave it open (``"bf16"`` = weight-only). The two NVFP4 globals are the checkpoint's + ``weight_scale_2`` and ``input_scale``; both come back on the accumulator. """ kernel = load_finegrained_kernel() original_shape = input.shape @@ -368,33 +368,16 @@ def finegrained_linear( activation_format: the activation format where the weights leave it open (see ``finegrained_triton_linear``). """ - # DeepGEMM is CUDA-only, dynamic-only, SM90+, FP4 or 128x128-block FP8; the combos it would - # silently corrupt (float32 scales on SM100, #47030) are skipped up front rather than attempted. - # ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces Triton for this dispatcher only. - deepgemm_preferred = ( - # SM100 perf: Triton measures 2.9x (decode) / 1.45x (prefill) over DeepGEMM on the DSV4 - # qkv linear — the 128x128 block-FP8 shape DeepGEMM is otherwise preferred for, which the - # gate above never catches (block-FP8's (N/128, K/128) grid has no swizzled layout). - not is_sm100() - # If the model is on multiple devices, DeepGEMM's context-bound kernels corrupt across devices. The - # multi-device guard `disable_deepgemm_on_multi_device` flips `_deepgemm_disabled` True at load, which - # this dispatcher sees and respects. The Triton fallback is context-free and safe. - and allow_deepgemm - # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major: DeepGEMM would - # consume the permuted buffer as affine and silently return garbage. Correctness gate, - # true on any arch, for whatever the loader's `FineGrainedSwizzleScales` swizzled. - and weight_scale_inv.ndim <= 2 - and is_deepgemm_loadable() - and activation_scale is None - # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format - and weight_global_scale is None - and input_global_scale is None - and activation_format is None - and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) - and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" - ) - - if deepgemm_preferred: + if prefers_deepgemm_linear( + weight, + weight_scale_inv, + block_size=block_size, + activation_scale=activation_scale, + weight_global_scale=weight_global_scale, + input_global_scale=input_global_scale, + activation_format=activation_format, + allow_deepgemm=allow_deepgemm, + ): try: return deepgemm_fp8_fp4_linear( input, @@ -442,6 +425,13 @@ class _FineGrainedModule: # kernels corrupt across devices); removable once the kernel ships a context-free loader. _deepgemm_disabled = False + def local(self, name: str) -> torch.Tensor | None: + """One of this module's parameters as the kernels take it: the local shard of whatever + FSDP2/EP wrapped it in, and ``None`` for a slot this module does not hold. Every operand + goes through here so no forward hands a `DTensor` to a kernel.""" + tensor = getattr(self, name, None) + return to_local(tensor) if tensor is not None else None + class FineGrainedLinear(_FineGrainedModule, nn.Linear): def __init__( @@ -473,11 +463,10 @@ def __init__( torch.empty(out_features, in_storage, dtype=fmt.weight_dtype), requires_grad=fmt.weight_dtype.is_floating_point, ) - if has_global_scale: - # NVFP4 two-level: the per-tensor fp32 global the kernel recovers on the accumulator - self.weight_global_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) - else: - self.register_parameter("weight_global_scale", None) + # NVFP4 two-level: the per-tensor fp32 global the kernel recovers on the accumulator + _set_optional_parameter( + self, "weight_global_scale", torch.tensor(1.0, dtype=torch.float32) if has_global_scale else None + ) if scale_group is None: # no group and no block: one per-tensor scale @@ -490,32 +479,23 @@ def __init__( requires_grad=sf_dtype.is_floating_point, ) - if self.activation_scheme == "static": - self.activation_scale = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) - else: - self.register_parameter("activation_scale", None) - - if self.has_bias: - self.bias = nn.Parameter(torch.empty(self.out_features)) - else: - self.register_parameter("bias", None) + static = self.activation_scheme == "static" + _set_optional_parameter(self, "activation_scale", torch.tensor(1.0, dtype=torch.float32) if static else None) + _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) def forward(self, input: torch.Tensor) -> torch.Tensor: if self.weight.element_size() > 1: return F.linear(input, self.weight, self.bias) - weight = to_local(self.weight) - scale_inv = to_local(self.weight_scale_inv) - return finegrained_linear( input, - weight, - scale_inv, + self.local("weight"), + self.local("weight_scale_inv"), block_size=self.block_size, activation_scale=self.activation_scale, bias=self.bias, allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, + weight_global_scale=self.local("weight_global_scale"), activation_format=self.activation_format, ) @@ -582,10 +562,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y.add_(self.bias.view(self.n_groups, -1)) return y - w = to_local(self.weight) - scale_inv = to_local(self.weight_scale_inv) - - w = w.view(self.n_groups, -1, hidden_dim) + w = self.local("weight").view(self.n_groups, -1, hidden_dim) + scale_inv = self.local("weight_scale_inv") x = x.movedim(-2, 0).reshape(-1, hidden_dim) scale_inv = scale_inv.view(self.n_groups, scale_inv.size(0) // self.n_groups, scale_inv.size(1)) @@ -601,7 +579,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: scale_inv, expert_start=expert_start, activation_format=self.activation_format, - b_global_scale=to_local(self.weight_global_scale) if self.weight_global_scale is not None else None, + b_global_scale=self.local("weight_global_scale"), ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) if self.has_bias: @@ -642,17 +620,13 @@ def _moe_operands(kernel, module) -> dict: else: post_expert_norm = module._apply_post_norm - def local(name): - t = getattr(module, name, None) - return to_local(t) if t is not None else None - def both(suffix: str) -> dict: """One kernel argument per projection. The kernels always say ``gate_up_proj``; the module names that weight ``up_proj`` when the model has no gate. A tensor the module does not hold passes as ``None``, which is the argument the kernels take for it.""" return { - f"gate_up_proj{suffix}": local(f"{up}{suffix}"), - f"down_proj{suffix}": local(f"down_proj{suffix}"), + f"gate_up_proj{suffix}": module.local(f"{up}{suffix}"), + f"down_proj{suffix}": module.local(f"down_proj{suffix}"), } return { @@ -780,36 +754,27 @@ def __init__( ) setattr(self, proj, weight) setattr(self, f"{proj}_scale_inv", scale) - if self.has_bias: - # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it - # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype - setattr(self, f"{proj}_bias", nn.Parameter(torch.empty(self.num_experts, rows))) - else: - self.register_parameter(f"{proj}_bias", None) - # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator. The - # weight's is one per expert — a gate|up stack whose halves modelopt calibrated - # separately is merged to that at load (`FineGrainedWeightGlobals`), SwiGLU being - # linear in the up half. The activation's is the checkpoint's calibrated - # `input_scale`: the quant divides by it and the accumulator multiplies it back. The - # gate_up's is ONE value (its rows are the hidden states, quantized once, before - # routing); the down's is per expert, because its rows ARE per expert — the gate_up - # epilogue requantizes each against its own expert's value. - if self.has_global_scale: - setattr( - self, - f"{proj}_weight_global_scale", - nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)), - ) - else: - self.register_parameter(f"{proj}_weight_global_scale", None) - if self.has_global_scale and self.activation_format != "bf16": - setattr( - self, - f"{proj}_input_global_scale", - nn.Parameter(torch.ones(1 if proj == up_name else self.num_experts, dtype=torch.float32)), - ) - else: - self.register_parameter(f"{proj}_input_global_scale", None) + # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it + # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype + _set_optional_parameter( + self, f"{proj}_bias", torch.empty(self.num_experts, rows) if self.has_bias else None + ) + # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator, one per + # expert for the weight (`FineGrainedWeightGlobals` merges a separately calibrated + # gate|up stack down to that). The activation's is the checkpoint's `input_scale` — + # one value for gate_up, whose rows are the pre-routing hidden states, and one per + # expert for down, whose rows are per expert. + _set_optional_parameter( + self, + f"{proj}_weight_global_scale", + torch.ones(self.num_experts, dtype=torch.float32) if self.has_global_scale else None, + ) + calibrated = self.has_global_scale and self.activation_format != "bf16" + _set_optional_parameter( + self, + f"{proj}_input_global_scale", + torch.ones(1 if proj == up_name else self.num_experts, dtype=torch.float32) if calibrated else None, + ) if self.activation_scheme == "static": self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) @@ -1234,7 +1199,9 @@ def convert(self, input_dict, target_patterns=None, **kwargs): out = {} for key, value in _keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value - if self.inverse: + if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) + out[key] = value + elif self.inverse: out[key] = value.view(torch.uint8).reshape(*value.shape[:-1], value.shape[-1] // 16, 16) else: out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) @@ -1463,7 +1430,7 @@ def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) -def _quantize_group(value: torch.Tensor, fmt: _WeightFormat): +def _quantize_group(value: torch.Tensor, fmt: WeightFormat): """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the @@ -1512,7 +1479,10 @@ def _scale_pattern_for(self, weight_pattern: str) -> str: # Strip the optional ``$`` regex anchor so we can match the underlying name. anchored = weight_pattern.endswith("$") base = weight_pattern[:-1] if anchored else weight_pattern - if base.endswith(".weight"): + if base.endswith("_blocks"): + # GPT-OSS packs its experts as `{proj}_blocks` + `{proj}_scales` + scale = base[: -len("_blocks")] + "_scales" + elif base.endswith(".weight"): scale = base[: -len(".weight")] + ".weight_scale_inv" elif base == "weight": scale = "weight_scale_inv" @@ -1609,9 +1579,10 @@ def convert( return {target_key: quantized} # Generic chain path: dequantize every weight pattern that has a sibling scale. + consumed = {self._scale_pattern_for(key) for key in input_dict} result: dict[str, list[torch.Tensor] | torch.Tensor] = {} for key, value in input_dict.items(): - if "activation_scale" in key or "weight_scale_inv" in key: + if "activation_scale" in key or key in consumed: continue # consumed by the dequant; drop from the chain scale_key = self._scale_pattern_for(key) if scale_key not in input_dict: diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index ba2b7268bfad..4cab0f6a64be 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -44,7 +44,6 @@ "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", DeprecationWarning, - stacklevel=2, ) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 083388a0a65a..fc4ead023d73 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -754,6 +754,7 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"] revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None) version = _HUB_KERNEL_MAPPING[kernel_name].get("version", None) + check_arch = _HUB_KERNEL_MAPPING[kernel_name].get("check_arch", True) # Default version as it's mandatory if version is None and revision is None: version = 1 @@ -763,7 +764,7 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ revision=revision, version=version, allow_all_kernels=ALLOW_ALL_KERNELS, - check_arch=_HUB_KERNEL_MAPPING[kernel_name].get("check_arch", True), + check_arch=check_arch, ) mapping[kernel_name] = kernel except FileNotFoundError as e: @@ -834,16 +835,13 @@ def get_kernel( raise ImportError(_MISSING_KERNELS_MESSAGE) user_agent = {"framework": "transformers", "version": __version__, "repo_id": kernel_name} - # `check_arch` (kernels >= 0.16) is only passed when a caller opts out, so older `kernels` - # releases without the keyword keep working. - arch_kwargs = {} if check_arch else {"check_arch": False} return get_kernel_hub( kernel_name, revision=revision, version=version, user_agent=user_agent, trust_remote_code=allow_all_kernels, - **arch_kwargs, + check_arch=check_arch, ) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 4ae56e89ccba..40d98c730f2e 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -15,7 +15,6 @@ "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", DeprecationWarning, - stacklevel=2, ) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index ca97f0d46e54..7ee9dc8d77c6 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1691,11 +1691,19 @@ def post_init(self): class FineGrainedConfig(QuantizationConfigMixin): """ Configuration for the fine-grained quantization family served by the - `kernels-community/finegrained-kernels` package: block-FP8 (DeepSeek-style, fp32 or UE8M0 - scales), MXFP8, MXFP4 and NVFP4 weights, with dynamic, static or weight-matched activation - quantization. The weight FORMAT is never declared here — it is resolved from the checkpoint - tensors themselves (value dtype, scale dtype/shape, presence of a global scale), exactly the - way the kernels resolve it, so the config cannot disagree with the weights. + `kernels-community/finegrained-kernels` package. + + | format | weights | block scales | a checkpoint that ships it | + |-----------|--------------------------------|-------------------------|--------------------------------| + | block-FP8 | `float8_e4m3fn` | fp32 or UE8M0, 128x128 | `deepseek-ai/DeepSeek-V3` | + | MXFP8 | `float8_e4m3fn` | UE8M0, group-32 | `MiniMaxAI/MiniMax-M3` | + | MXFP4 | E2M1, two values per `int8` | UE8M0, group-32 | `openai/gpt-oss-20b` | + | NVFP4 | E2M1, two values per `int8` | E4M3 group-16 + fp32 global | `nvidia/GLM-5.2-NVFP4` | + + The weight FORMAT is never declared here — it is resolved from the checkpoint tensors + themselves (value dtype, scale dtype/shape, presence of a global scale), exactly the way the + kernels resolve it, so the config cannot disagree with the weights. What the config does carry + is everything the tensors leave open: how activations are quantized, and which modules to skip. Args: activation_scheme (`str`, *optional*, defaults to `"dynamic"`): @@ -1735,12 +1743,10 @@ def __init__( # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") - # NVIDIA modelopt exports (quant_method "modelopt"): translate the payload — - # quant_algo names the format (NVFP4 is the one the finegrained kernels serve), the - # skip list is a glob-style ``ignore`` (or ``exclude_modules``, the name modelopt's - # own ``hf_quant_config.json`` uses), and ``kv_cache_scheme`` is dropped (the KV cache - # stays in the compute dtype). The calibrated ``input_scale`` TENSORS are consumed by - # the loader's converters — the kernels quantize activations against them. + # NVIDIA modelopt exports: `quant_algo` names the format, the skip list is a glob-style + # `ignore` (or `exclude_modules`, modelopt's own spelling), and `kv_cache_scheme` is + # dropped — the KV cache stays in the compute dtype. The calibrated `input_scale` TENSORS + # are the loader's business, not this config's. if str(self.quant_method) == "modelopt" or kwargs.get("quant_algo") is not None: quant_algo = kwargs.pop("quant_algo", None) kwargs.pop("config_groups", None) From 09ea21c0445f79091bb257fbcd6183e857aea0bc Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 09:58:19 -0700 Subject: [PATCH 18/68] finegrained: freeze the mxfp4 and nvfp4 integrations, and serve what they served MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every checkpoint `quant_method` the two answer to — `fp8`, `mxfp8`, `mxfp4`, `nvfp4`, `modelopt` — already routes to `FineGrainedHfQuantizer`, so both carry the same DeprecationWarning the fp8 pair does. Auditing that supersession found one real gap. `dequantize=True` on a GPT-OSS MXFP4 checkpoint produced no converter at all: the `_blocks`/`_scales` keys were wired only for the quantized path, so they went unmatched, where `Mxfp4Config(dequantize=True)` had handled them. The finegrained chain now dequantizes them, bit-equal to `mxfp4.convert_moe_packed_tensors` — the implementation it replaces — which the new test compares against directly. Two shared ops needed repair to get there: `FineGrainedPackedBlocks` regrouped a sibling `_scales` entry as if it were blocks, and `FineGrainedDequantize` left the scale it had consumed in the chain. The same path on NVFP4 or modelopt would have folded one block scale and dropped the second-level global, handing back a weight scaled by `1 / global`. Neither frozen integration supported that either, but it failed quietly — including where `dequantize` turns itself on (no GPU, compute capability < 8.9). It raises now. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/mxfp4.py | 11 +++ src/transformers/integrations/nvfp4.py | 9 +++ .../quantizers/quantizer_finegrained.py | 67 +++++++++++++++---- .../quantizers/quantizer_mxfp4.py | 9 +++ .../quantizers/quantizer_nvfp4.py | 9 +++ tests/kernels/test_finegrained.py | 40 ++++++++--- 6 files changed, 124 insertions(+), 21 deletions(-) diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index c15a8b844353..b5efea30f582 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings + from ..utils import is_torch_available, logging @@ -26,6 +28,15 @@ logger = logging.get_logger(__name__) + +warnings.warn( + "mxfp4 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, +) + + FP4_VALUES = [ +0.0, +0.5, diff --git a/src/transformers/integrations/nvfp4.py b/src/transformers/integrations/nvfp4.py index 476e459df405..51570d197768 100644 --- a/src/transformers/integrations/nvfp4.py +++ b/src/transformers/integrations/nvfp4.py @@ -14,6 +14,7 @@ from __future__ import annotations import functools +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -33,6 +34,14 @@ import torch.nn as nn +warnings.warn( + "nvfp4 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, +) + + # NVFP4 uses one FP8 scale factor for every 16 consecutive weight values. SF_VEC_SIZE = 16 diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 21cd0862664a..72925356ee7d 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -1,3 +1,16 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import re from typing import TYPE_CHECKING @@ -17,23 +30,41 @@ class FineGrainedHfQuantizer(HfQuantizer): - """ - FP8 quantization implementation supporting both standard and MoE models. - Supports both e4m3fn formats based on platform. + """Quantizer for the fine-grained family the `kernels-community/finegrained-kernels` package + serves: block-FP8 (fp32 or UE8M0 scales), MXFP8, MXFP4 and NVFP4, for dense linears, embeddings + and MoE experts alike. + + It never declares the weight format — that is resolved from the checkpoint tensors themselves, + the way the kernels resolve it. What it does own is the vocabulary each producer ships its + scales under (modelopt's `weight_scale_2` / `input_scale`, GPT-OSS's `_blocks` / `_scales`) and + the conversion ops that bring them to the layout `FineGrainedExperts` holds. """ requires_calibration = False quantization_config: "FineGrainedConfig" + def _assert_dequantize_supported(self) -> None: + """The dequantize chain folds a weight into full precision with its per-block scale alone. + NVFP4 carries a second level (modelopt's `weight_scale_2`) that the chain has no slot for, + so refuse rather than hand back a weight scaled by `1 / global`.""" + if self._quant_method() in ("nvfp4", "modelopt"): + raise NotImplementedError( + f"`dequantize=True` is not supported for {self._quant_method()!r} checkpoints: their " + "two-level scales cannot be folded into a full-precision weight by this path. Load " + "them quantized on a GPU that serves NVFP4, or start from a bf16 checkpoint." + ) + def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") if self.quantization_config.dequantize: + self._assert_dequantize_supported() return if not torch.cuda.is_available() and not is_torch_xpu_available(): if self.pre_quantized: + self._assert_dequantize_supported() logger.warning_once( "Using FP8 quantized models requires a GPU or XPU, we will default to dequantizing the model to bf16 since no GPU or XPU is available" ) @@ -46,6 +77,7 @@ def validate_environment(self, *args, **kwargs): compute_capability = torch.cuda.get_device_capability() major, minor = compute_capability if (major < 8) or (major == 8 and minor < 9): + self._assert_dequantize_supported() logger.warning_once( "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)" f", actual = `{major}.{minor}`. We will default to dequantizing the model to bf16. Feel free " @@ -239,15 +271,26 @@ def get_weight_conversions(self): ] return converters - # modelopt NVFP4: E4M3 group-16 `weight_scale`, fp32 `weight_scale_2` weight globals and - # calibrated fp32 `input_scale` activation globals, in EITHER of the two layouts a modelopt - # checkpoint ships them in — one tensor per expert per projection (GLM-5.2-NVFP4) or one - # already-stacked tensor per layer (the fused vLLM layout, Muse-Spark). Both sets are - # returned: a checkpoint matches one of them and the other never fires. The arch plan - # stacks the expert WEIGHTS; these bring the scales and globals to the same layout, and - # `_with_expert_layout_ops` adds the gate|up interleave + swizzle to the block scales. - # Only the routed experts are quantized, so no dense converters — a generic `weight_scale*` - # rename would run before converter collection and mangle these keys. + # The same GPT-OSS keys under `dequantize=True`: the blocks regroup into the packed rows + # and the exponent-byte scales come back out of them as bf16, in the (E, hidden, 2I) + # orientation the unquantized experts hold. + if self.pre_quantized and self._quant_method() == "mxfp4" and self.quantization_config.dequantize: + from ..core_model_loading import Transpose + + return [ + WeightConverter( + source_patterns=[rf"{proj}_blocks$", rf"{proj}_scales$"], + target_patterns=proj, + operations=[FineGrainedPackedBlocks(self), FineGrainedDequantize(self), Transpose(1, 2)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + + # modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts + # a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already + # stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and + # the other never fires. Routed experts only — a generic `weight_scale*` rename would run + # before converter collection and mangle these keys. if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: # A weight-only run keeps activations bf16, so its experts hold no activation global # and the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index ec6f62ba3fee..fee0b8b6af68 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import warnings from typing import TYPE_CHECKING from .base import HfQuantizer @@ -36,6 +37,14 @@ from ..core_model_loading import WeightConverter +warnings.warn( + "quantizer_mxfp4 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, +) + + logger = logging.get_logger(__name__) triton_kernels_hub = None diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index cb58ca213578..1906afe944b2 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import warnings from typing import TYPE_CHECKING from ..utils import is_kernels_available, is_torch_available @@ -29,6 +30,14 @@ from ..utils.quantization_config import NVFP4Config +warnings.warn( + "quantizer_nvfp4 is frozen for backward compatibility and no longer " + "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", + DeprecationWarning, +) + + def _as_cuda_device(device) -> torch.device | None: if isinstance(device, int): return torch.device("cuda", device) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index d2b7665d3b2d..1fcbfcd16098 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -21,6 +21,7 @@ import torch +import transformers.integrations.deepgemm as deepgemm import transformers.integrations.finegrained as fg from transformers.integrations.finegrained import ( FineGrainedExperts, @@ -125,7 +126,7 @@ class FineGrainedLinearMarshallingTest(unittest.TestCase): def _run(self, **linear_kwargs): kernel, rec = _fake_bundle() p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): x = torch.randn(3, 5, 64, dtype=torch.bfloat16) w = torch.randn(32, 64).to(torch.float8_e4m3fn) ws = torch.randn(1, 1, dtype=torch.float32) @@ -155,7 +156,7 @@ def test_nvfp4_global_and_weight_only_format(self): def test_module_forward_threads_everything(self): kernel, rec = _fake_bundle() p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): m = FineGrainedLinear( 64, 32, @@ -234,7 +235,7 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): p2, p3, p4, - mock.patch.object(fg, "is_deepgemm_loadable", return_value=False), + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), mock.patch.object(fg, "finegrained_linear", linear), ): m(*self._route()) @@ -356,7 +357,7 @@ def forward(self, x): return x m.post_expert_norm = m._apply_post_norm = _Recording() - with p1, p2, p3, p4, mock.patch.object(fg, "is_deepgemm_loadable", return_value=False): + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): m(*self._route()) self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") @@ -499,7 +500,7 @@ def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): @require_torch class FineGrainedDeepGemmDispatchTest(unittest.TestCase): - """`deepgemm_preferred` carries two independent gates: a correctness one (a pre-swizzled + """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 scales have no swizzled layout, so the first never fires for the shape the second catches.""" @@ -515,8 +516,8 @@ def _routed_to(self, *, sm100, scale_ndim): p2, p3, p4, - mock.patch.object(fg, "is_deepgemm_loadable", return_value=True), - mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), + mock.patch.object(deepgemm, "is_sm100", return_value=sm100), mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, ): fg.finegrained_linear(torch.randn(4, 64, dtype=torch.bfloat16), w, s, block_size=[128, 128]) @@ -607,7 +608,7 @@ def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): # block-FP8's (N/128, K/128) grid never reaches a scaled-MMA, whatever its scale dtype _, experts = self._experts("fp8") self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) - with mock.patch.dict("os.environ", {"TRANSFORMERS_FINEGRAINED_NO_SWIZZLE": "1"}): + with mock.patch.object(fg, "_swizzles_scales", return_value=False): _, experts = self._experts("mxfp8") self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) @@ -876,6 +877,27 @@ class FineGrainedOnTheFlyQuantizeTest(unittest.TestCase): """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run the kernels' quantizers — covered with real kernels below), emitted in the module's layout.""" + def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): + """`dequantize=True` on a GPT-OSS MXFP4 checkpoint used to be the mxfp4 integration's job. + The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation + the unquantized experts hold — so it is compared against that integration directly.""" + from transformers.core_model_loading import Transpose + from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedPackedBlocks + from transformers.integrations.mxfp4 import convert_moe_packed_tensors + + torch.manual_seed(0) + num_experts, rows, hidden = 2, 8, 64 + blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) + scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) + reference = convert_moe_packed_tensors(blocks.clone(), scales.clone(), dtype=torch.bfloat16) + + sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] + tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} + for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): + tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) + self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain + torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) + def test_block_fp8_round_trips_within_its_floor(self): from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedQuantize @@ -1335,7 +1357,7 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): ) # K >= 256: the swizzled arm has tiles cfg._experts_implementation = "grouped_mm" experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() - with mock.patch.dict("os.environ", {"TRANSFORMERS_FINEGRAINED_NO_SWIZZLE": "1"}): + with mock.patch.object(fg, "_swizzles_scales", return_value=False): affine_experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() model = torch.nn.Module() model.experts = experts From e139d2a099c7edfa409b33826ea5b43206ee2111 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 10:35:26 -0700 Subject: [PATCH 19/68] fix style --- docs/source/en/main_classes/quantization.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index ea2080b00e3f..2a2a8c27b7aa 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -98,6 +98,10 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] SpQRConfig +## FineGrainedConfig + +[[autodoc]] FineGrainedConfig + ## FineGrainedFP8Config [[autodoc]] FineGrainedFP8Config From c66564416ec4e0ef5df165f69f3473323e34a786 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 10:54:36 -0700 Subject: [PATCH 20/68] FineGrainedFP8Config is its own class again, not an alias of the new one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It had become `FineGrainedFP8Config = FineGrainedConfig`, so the name still imported but the class was gone. That makes "frozen for backward compatibility" a fiction — the old name silently gained the new behaviour (`activation_format`, the modelopt payload parsing, the wider `post_init`) — and it rendered the same class under two `[[autodoc]]` headings. Restored verbatim, with a docstring line pointing at the config that serves the whole family. Splitting them means the two `isinstance` sites in `auto.py` no longer cover the new config by accident. `LOADING_ATTRIBUTES_CONFIG_TYPES` is the load-bearing one: it is what carries `dequantize` / `modules_to_not_convert` from a config passed to `from_pretrained` onto the checkpoint's own, so leaving it alone would have quietly stopped `FineGrainedConfig(dequantize=True)` from taking effect on every fine-grained checkpoint. Both configs are in it now. `quant_method="fp8"` still builds `FineGrainedConfig`: the frozen class is what you get by naming it, not what a checkpoint resolves to. Co-Authored-By: Claude Opus 5 --- src/transformers/quantizers/auto.py | 5 +- src/transformers/utils/quantization_config.py | 65 ++++++++++++++++++- 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 43c2bb60528a..f00f4865bb6d 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -146,6 +146,7 @@ CompressedTensorsConfig, Mxfp4Config, MetalConfig, + FineGrainedConfig, FineGrainedFP8Config, ) @@ -276,7 +277,9 @@ def merge_quantization_configs( if loading_attr_dict: warning_msg += f"However, loading attributes (e.g. {list(loading_attr_dict.keys())}) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored." - if warning_msg != "" and not isinstance(quantization_config, (Mxfp4Config, MetalConfig, FineGrainedFP8Config)): + if warning_msg != "" and not isinstance( + quantization_config, (Mxfp4Config, MetalConfig, FineGrainedConfig, FineGrainedFP8Config) + ): warnings.warn(warning_msg) else: # in the case of mxfp4, we don't want to print the warning message, bit confusing for users diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 7ee9dc8d77c6..1ac3a460166b 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1802,8 +1802,69 @@ def get_loading_attributes(self): return {"dequantize": self.dequantize, "modules_to_not_convert": self.modules_to_not_convert} -# Back-compat alias: serialized checkpoints and released code reference the FP8-era name. -FineGrainedFP8Config = FineGrainedConfig +class FineGrainedFP8Config(QuantizationConfigMixin): + """ + FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. + + Frozen for backward compatibility: it configures the frozen `finegrained_fp8` integration and no + longer receives new recipes. [`FineGrainedConfig`] is the one that serves the whole family + (block-FP8, MXFP8, MXFP4, NVFP4), and a checkpoint declaring `quant_method="fp8"` builds that one. + + Args: + activation_scheme (`str`, *optional*, defaults to `"dynamic"`): + The scheme used for activation, the defaults and only support scheme for now is "dynamic". + weight_block_size (`typing.tuple[int, int]`, *optional*, defaults to `(128, 128)`): + The size of the weight blocks for quantization, default is (128, 128). + dequantize (`bool`, *optional*, defaults to `False`): + Whether to dequantize the model during loading. + modules_to_not_convert (`list`, *optional*): + A list of module names that should not be converted during quantization. + modules_to_convert (`list`, *optional*): + A list of additional module names, such as embedding tables, that should be converted during quantization. + scale_fmt (`str`, *optional*, defaults to `"float"`): + Storage dtype of the per-block weight scales: `"float"` (fp32, V3-style) or + `"ue8m0"` (1-byte `torch.float8_e8m0fnu`, V4-style). + """ + + def __init__( + self, + activation_scheme: str = "dynamic", + weight_block_size: tuple[int, int] = (128, 128), + dequantize: bool = False, + modules_to_not_convert: list | None = None, + modules_to_convert: list | None = None, + scale_fmt: str = "float", + **kwargs, + ): + self.quant_method = kwargs.pop("quant_method", QuantizationMethod.FP8) + # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. + if modules_to_not_convert is None and "ignored_layers" in kwargs: + modules_to_not_convert = kwargs.pop("ignored_layers") + self.modules_to_not_convert = modules_to_not_convert + # TODO: check overlap with not to convert + self.modules_to_convert = modules_to_convert + self.activation_scheme = activation_scheme + self.weight_block_size = weight_block_size + self.dequantize = dequantize + self.scale_fmt = scale_fmt + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + self.activation_scheme = self.activation_scheme.lower() + if self.activation_scheme not in ["dynamic", "static"]: + raise ValueError(f"Activation scheme {self.activation_scheme} not supported") + if self.weight_block_size is not None and len(self.weight_block_size) != 2: + raise ValueError("weight_block_size must be a tuple of two integers") + if self.weight_block_size is not None and (self.weight_block_size[0] <= 0 or self.weight_block_size[1] <= 0): + raise ValueError("weight_block_size must be a tuple of two positive integers") + if self.scale_fmt not in ("float", "ue8m0"): + raise ValueError(f"scale_fmt must be 'float' or 'ue8m0'; got {self.scale_fmt!r}") + + def get_loading_attributes(self): + return {"dequantize": self.dequantize, "modules_to_not_convert": self.modules_to_not_convert} class QuarkConfig(QuantizationConfigMixin): From 29cca042b787b6d9034d22a37d96c6893418be66 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 15 Sep 2026 10:58:01 -0700 Subject: [PATCH 21/68] FineGrainedFP8Config: drop the frozen-notice docstring line MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `c66564416e` restored the class and added a line saying it "configures the frozen `finegrained_fp8` integration". That is not true: it sets `quant_method="fp8"`, which `AUTO_QUANTIZER_MAPPING` routes to `FineGrainedHfQuantizer` — the frozen quantizer is not reachable through `from_pretrained` at all. Removing it also leaves the class byte-identical to main, which is what a class frozen for backward compatibility should look like in a diff. Which config serves which family is already stated where a reader meets it, in `FineGrainedConfig`'s own docstring on the same docs page. Co-Authored-By: Claude Opus 5 --- src/transformers/utils/quantization_config.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 1ac3a460166b..8868190034ba 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1806,10 +1806,6 @@ class FineGrainedFP8Config(QuantizationConfigMixin): """ FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. - Frozen for backward compatibility: it configures the frozen `finegrained_fp8` integration and no - longer receives new recipes. [`FineGrainedConfig`] is the one that serves the whole family - (block-FP8, MXFP8, MXFP4, NVFP4), and a checkpoint declaring `quant_method="fp8"` builds that one. - Args: activation_scheme (`str`, *optional*, defaults to `"dynamic"`): The scheme used for activation, the defaults and only support scheme for now is "dynamic". From f63ad97d3f8793f7ba847536f21211b3020b79ab Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 01:10:10 -0700 Subject: [PATCH 22/68] Tolerate an experts module without the norm, and a machine without CUDA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three CI failures on a CPU-only runner, three causes: `post_expert_norm_name` is set by `use_experts_implementation`, but the shared forwards are the ones every backend adapts onto — including duck-typed stand-ins built without the decorator. The norm is optional, so read it as one. `_assert_stacked_gate_up` already read its second attribute defensively; `has_gate` now matches. `prefers_deepgemm_linear` calls `is_sm100()`, and `get_device_capability()` asserts rather than answering False on a CPU-only build. Caught there rather than gated on `is_available()`, so that a test faking the capability alone still answers — which is the convention the existing deepgemm tests are written to. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/deepgemm.py | 11 +++++++++-- src/transformers/integrations/moe.py | 12 ++++++++---- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 381ba9524b6b..b0e52c6d8d08 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -145,8 +145,15 @@ def is_sm100() -> bool: `torch.cuda.get_device_capability()` pybind out of the graph — so the branches on it (SF layout, cast kwargs, psum layout, and the arch guards) stay compile-safe. Re-queries each eager call, so a faked capability in tests is honoured. + + False rather than raising where there is no device to ask: `get_device_capability()` asserts on a + CPU-only build. Caught rather than gated on `is_available()` so that a test faking the capability + alone still answers — which is the convention the callers' tests are written to. """ - return torch.cuda.get_device_capability()[0] >= 10 + try: + return torch.cuda.get_device_capability()[0] >= 10 + except (AssertionError, RuntimeError): + return False @torch._dynamo.assume_constant_result @@ -772,7 +779,7 @@ def _assert_stacked_gate_up(module: torch.nn.Module) -> None: """Mega MoE's ``transform_weights_for_mega_moe`` does its own gate/up interleave, so it takes the stacked ``[gate; up]`` rows. A module loaded for a triton backend holds them interleaved; switched to this backend after load, silently re-permuting would be a wrong answer rather than a failure.""" - if module.has_gate and getattr(module, "holds_interleaved_gate_up", False): + if getattr(module, "has_gate", True) and getattr(module, "holds_interleaved_gate_up", False): raise RuntimeError( "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — pass " diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 004f6387c565..fcd3d8d39ae5 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -164,8 +164,10 @@ def batched_mm_experts_forward( proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed ) # (S, hidden_dim) - # Normalize each expert application, where such a norm is defined, before it is weighted - if self.post_expert_norm_name is not None: + # Normalize each expert application, where such a norm is defined, before it is weighted. + # `getattr`: an experts module built without the decorator (or predating the norm) has no such + # attribute, and this forward is the shared one every backend adapts onto. + if getattr(self, "post_expert_norm_name", None) is not None: proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) # Apply routing weights @@ -479,8 +481,10 @@ def grouped_mm_experts_forward( if sentinel_mask is not None: proj_out = proj_out.masked_fill(sentinel_mask, 0.0) - # Normalize each expert application, where such a norm is defined, before it is weighted - if self.post_expert_norm_name is not None: + # Normalize each expert application, where such a norm is defined, before it is weighted. + # `getattr`: an experts module built without the decorator (or predating the norm) has no such + # attribute, and this forward is the shared one every backend adapts onto. + if getattr(self, "post_expert_norm_name", None) is not None: proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) # Apply routing weights From 1265f0cce5c2004c1c1e998b8fa6e7d2f9c6abe5 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 01:10:10 -0700 Subject: [PATCH 23/68] Pin that every modelopt converter target is a slot the module holds A converter whose target no module holds is a load failure, and the activation global is exactly that under a weight-only run: the experts allocate one only when they quantize activations, so the declarations are the only thing standing between the two runs. Checked against a real module's parameter names in both. Verified it catches the failure by removing the conditional: the weight-only run then converts `gate_up_proj_input_global_scale` onto a module without that slot. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 5 ++-- tests/kernels/test_finegrained.py | 28 ++++++++++++++++++-- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 9be49026f2c9..29697b54ced2 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1284,8 +1284,9 @@ def _assert_no_gate_up_bias(model) -> None: """A gate|up bias is added AFTER the global, so merging the halves would have to divide its up rows by the same ratio — in the bias converter, which the layout ops own. No checkpoint pairs the two, so refuse rather than carry a silent half-correction.""" - experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None - if experts is not None and experts.has_bias: + if model is not None and any( + isinstance(module, FineGrainedExperts) and module.has_bias for module in model.modules() + ): raise NotImplementedError( "this checkpoint calibrates the gate|up halves separately AND carries an expert " "bias; the up half's bias would have to be folded with the globals." diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 1fcbfcd16098..c142dbfd1d25 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -931,14 +931,38 @@ class FineGrainedModeloptConverterTest(unittest.TestCase): present in the chain, and expert parallelism selects experts by that index. Without it every rank collects all E globals and the forward asserts on the per-expert count.""" - def _modelopt_conversions(self): + def _modelopt_conversions(self, **cfg_kwargs): from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig - cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4") + cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) quantizer = FineGrainedHfQuantizer(cfg) return quantizer.get_weight_conversions() + def test_every_declared_target_is_a_slot_the_module_holds(self): + """A converter target that no module holds is a load failure, and the activation global is + exactly such a target under a weight-only run: the experts allocate one only when they + quantize activations. Both runs are checked against a real module's parameter names, since + the declarations are the only thing standing between the two.""" + cfg = _Cfg() + for activation_format, holds_global in ((None, True), ("bf16", False)): + experts = FineGrainedExperts( + cfg, block_size=(4, 4), weight_format="nvfp4", activation_format=activation_format + ) + slots = {name for name, _ in experts.named_parameters()} | set(experts._buffers) + self.assertEqual( + any("input_global_scale" in slot for slot in slots), + holds_global, + f"module with activation_format={activation_format!r} disagrees about the slot", + ) + targets = { + t.split("experts.")[-1].rstrip("$") + for c in self._modelopt_conversions(activation_format=activation_format) + for t in _targets(c) + } + for target in targets: + self.assertIn(target, slots, f"activation_format={activation_format!r} converts onto a missing slot") + def test_global_scale_converters_carry_an_expert_index(self): from transformers.core_model_loading import MergeModulelist From 97dfa1063eb2904e4d96db8039ae9adbc6f60aa6 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 01:12:23 -0700 Subject: [PATCH 24/68] is_sm100: answer False without CUDA, and fake both halves in the tests The integration reads better as `is_available() and capability >= 10`; the reason it wasn't was that three test sites faked the capability alone, which that form ignores. That is the tests constraining the integration, so the tests move instead: they now patch availability alongside the capability, and the module docstring says so. Replaces the try/except that was catching what the availability check now prevents. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/deepgemm.py | 12 +++--------- tests/kernels/test_deepgemm.py | 19 +++++++++++-------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index b0e52c6d8d08..04da8e891cd8 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -144,16 +144,10 @@ def is_sm100() -> bool: evaluate it once at trace time and inline the bool — keeping the untraceable `torch.cuda.get_device_capability()` pybind out of the graph — so the branches on it (SF layout, cast kwargs, psum layout, and the arch guards) stay compile-safe. Re-queries each eager call, so a faked - capability in tests is honoured. - - False rather than raising where there is no device to ask: `get_device_capability()` asserts on a - CPU-only build. Caught rather than gated on `is_available()` so that a test faking the capability - alone still answers — which is the convention the callers' tests are written to. + capability in tests is honoured — which means faking BOTH, since `get_device_capability()` asserts + rather than answering on a CPU-only build and the availability check is what keeps it unreached. """ - try: - return torch.cuda.get_device_capability()[0] >= 10 - except (AssertionError, RuntimeError): - return False + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 @torch._dynamo.assume_constant_result diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 3542cd4c20c7..778ca7e2c90b 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -28,9 +28,9 @@ (packed int32 UE8M0 SFs, int32 grouped layout, `(qtensor, sf)` operand tuples, recipes, transformed Mega MoE weights, ...). -Arch-gated paths are exercised by faking the device capability (`is_sm100()` reads it): the -SF-packing / TMA-alignment / psum-layout code it selects is pure tensor arithmetic, so mocking the -capability to SM100 drives the Blackwell paths on any device. +Arch-gated paths are exercised by faking the device (`is_sm100()` reads availability, then the +capability): the SF-packing / TMA-alignment / psum-layout code it selects is pure tensor arithmetic, +so mocking both to a Blackwell device drives those paths anywhere, a CPU-only runner included. """ import contextlib @@ -361,12 +361,14 @@ def setUp(self): def _bundle(self, *, is_sm100): captured = {} bundle = _make_bundle(captured) - # The forwards read the arch via `is_sm100()` (which queries `get_device_capability`), so fake the - # device to the requested arch: lets SM100 dispatch/packing run on this SM80 box and drives the - # Hopper-rejection guards. `[0]` is all `is_sm100()` reads. + # The forwards read the arch via `is_sm100()`, so fake the device to the requested arch: lets + # SM100 dispatch/packing run on this SM80 box and drives the Hopper-rejection guards. `[0]` is + # all `is_sm100()` reads of the capability, and it reaches that only when CUDA is available — + # so a CPU-only runner needs both faked, or every arch-gated branch below goes untested. capability = (10, 0) if is_sm100 else (9, 0) with ( mock.patch.object(dg, "load_deepgemm_kernel", return_value=bundle), + mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object(torch.cuda, "get_device_capability", return_value=capability), ): yield captured @@ -431,11 +433,12 @@ def test_assert_sm100_requirements(self): int8_w = torch.zeros(1, 1, dtype=torch.int8, device=torch_device) f32_sf = torch.ones(1, 1, dtype=torch.float32, device=torch_device) ue8m0_sf = f32_sf.to(torch.float8_e8m0fnu) - with mock.patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)): # SM90 + available = mock.patch.object(torch.cuda, "is_available", return_value=True) + with available, mock.patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)): # SM90 with self.assertRaisesRegex(NotImplementedError, "Blackwell"): dg._assert_sm100_requirements(int8_w, ue8m0_sf) # FP4 has no Hopper kernel dg._assert_sm100_requirements(fp8_w, f32_sf) # float32 SF is fine on SM90 -> no raise - with mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)): # SM100 + with available, mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)): # SM100 with self.assertRaisesRegex(NotImplementedError, "float32 scale-factor path"): dg._assert_sm100_requirements(fp8_w, f32_sf) # no float32 SF path on Blackwell dg._assert_sm100_requirements(int8_w, ue8m0_sf) # UE8M0 on SM100 -> no raise From 788d02dd9b2dc236c72e237353d448244e63a179 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 01:34:44 -0700 Subject: [PATCH 25/68] fix --- tests/kernels/test_finegrained.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index c142dbfd1d25..997a034650a2 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -87,6 +87,7 @@ class _Cfg: num_local_experts = 4 intermediate_size = 32 hidden_act = "silu" + _experts_implementation = None def _loaded(kernel): From 53a8970aa16680d0d3a968d88aab683c2cbd39f3 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 08:04:29 -0700 Subject: [PATCH 26/68] finegrained: drop the DTensor unwrap, and pass the kernel operands by name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A sharded parameter is already local by the time a forward runs, so the unwrap was doing nothing: FSDP2 unshards in its pre-forward hook, and a TP/EP plan swaps each DTensor for its local shard around the call (`MoeExpertsParallel` -> `_use_local_dtensor_params`). Verified both: under `fully_shard` alone a param reads as a plain Parameter inside forward, not a DTensor. So the invariant belongs to the distributed layer that wrapped the parameter, not to every backend that reads one. That removes `to_local` and the `local()` accessor. Absent slots still read back as None because `_set_optional_parameter` registers them as None parameters, which is what the kernels take for a missing bias or global. `_moe_operands` also stops building its dict through a `both(suffix)` helper: the kernels' arguments are now written out, so `down_proj_scale_inv=` can be found by grep. The five `getattr`s that remain are load-bearing — the kernels always say `gate_up_proj`, the module says `up_proj` when the model has no gate. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 47 ++++++++------------ 1 file changed, 19 insertions(+), 28 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 29697b54ced2..ada31c7da7ef 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -40,7 +40,6 @@ ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel from .moe import ExpertsInterface, use_experts_implementation -from .tensor_parallel import to_local logger = logging.get_logger(__name__) @@ -425,13 +424,6 @@ class _FineGrainedModule: # kernels corrupt across devices); removable once the kernel ships a context-free loader. _deepgemm_disabled = False - def local(self, name: str) -> torch.Tensor | None: - """One of this module's parameters as the kernels take it: the local shard of whatever - FSDP2/EP wrapped it in, and ``None`` for a slot this module does not hold. Every operand - goes through here so no forward hands a `DTensor` to a kernel.""" - tensor = getattr(self, name, None) - return to_local(tensor) if tensor is not None else None - class FineGrainedLinear(_FineGrainedModule, nn.Linear): def __init__( @@ -489,13 +481,13 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return finegrained_linear( input, - self.local("weight"), - self.local("weight_scale_inv"), + self.weight, + self.weight_scale_inv, block_size=self.block_size, activation_scale=self.activation_scale, bias=self.bias, allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=self.local("weight_global_scale"), + weight_global_scale=self.weight_global_scale, activation_format=self.activation_format, ) @@ -562,8 +554,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y.add_(self.bias.view(self.n_groups, -1)) return y - w = self.local("weight").view(self.n_groups, -1, hidden_dim) - scale_inv = self.local("weight_scale_inv") + w = self.weight.view(self.n_groups, -1, hidden_dim) + scale_inv = self.weight_scale_inv x = x.movedim(-2, 0).reshape(-1, hidden_dim) scale_inv = scale_inv.view(self.n_groups, scale_inv.size(0) // self.n_groups, scale_inv.size(1)) @@ -579,7 +571,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: scale_inv, expert_start=expert_start, activation_format=self.activation_format, - b_global_scale=self.local("weight_global_scale"), + b_global_scale=self.weight_global_scale, ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) if self.has_bias: @@ -620,21 +612,20 @@ def _moe_operands(kernel, module) -> dict: else: post_expert_norm = module._apply_post_norm - def both(suffix: str) -> dict: - """One kernel argument per projection. The kernels always say ``gate_up_proj``; the module - names that weight ``up_proj`` when the model has no gate. A tensor the module does not - hold passes as ``None``, which is the argument the kernels take for it.""" - return { - f"gate_up_proj{suffix}": module.local(f"{up}{suffix}"), - f"down_proj{suffix}": module.local(f"down_proj{suffix}"), - } - + # The kernels always say `gate_up_proj`; the module names that weight `up_proj` when the model + # has no gate, which is the only reason these are not plain attribute reads. A slot the module + # does not hold reads back as None — the argument the kernels take for an absent one. return { - **both(""), - **both("_scale_inv"), - **both("_weight_global_scale"), - **both("_input_global_scale"), - **both("_bias"), + "gate_up_proj": getattr(module, up), + "gate_up_proj_scale_inv": getattr(module, f"{up}_scale_inv"), + "gate_up_proj_weight_global_scale": getattr(module, f"{up}_weight_global_scale"), + "gate_up_proj_input_global_scale": getattr(module, f"{up}_input_global_scale"), + "gate_up_proj_bias": getattr(module, f"{up}_bias"), + "down_proj": module.down_proj, + "down_proj_scale_inv": module.down_proj_scale_inv, + "down_proj_weight_global_scale": module.down_proj_weight_global_scale, + "down_proj_input_global_scale": module.down_proj_input_global_scale, + "down_proj_bias": module.down_proj_bias, # a supported activation NAME is fused into the gate_up epilogue; any other callable # leaves that GEMM plain and runs on the host between the two "act_fn": act_fn, From e8e6b879079df96e9fb591b8c4f394c95323d1cc Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 16 Sep 2026 08:04:29 -0700 Subject: [PATCH 27/68] deepgemm: only Mega MoE cannot apply a post-expert norm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guard refused a per-expert output norm on every DeepGEMM arm, but only Mega MoE has to: it fuses the routing-weighted reduce into its kernel, so there is no seam for the norm. The other two arms reduce in `_combine_routed_output`, so the norm rides the rows just before it — the same place the reference forwards apply it. `deepgemm_experts_guards` gains `post_expert_norm`, which only the Mega MoE arm sets False. A model naming a norm — Muse-Spark names `input_scaled_rms_norm` — was locked out of every DeepGEMM path and now loses only Mega MoE. The old behaviour had no test, which is why relaxing it broke nothing; the new one covers both halves so the Mega MoE restriction cannot silently disappear either. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/deepgemm.py | 31 +++++++++++++++-------- tests/kernels/test_deepgemm.py | 30 ++++++++++++++++++++++ 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 04da8e891cd8..82d93e6db906 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -654,26 +654,37 @@ def deepgemm_fp8_fp4_linear( def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: - """DeepGEMM's experts chain ends at the routing-weighted reduce, with no slot for a model's - per-expert output norm — which it would drop silently, so refuse instead.""" + """Mega MoE fuses the routing-weighted reduce into its kernel, so a model's per-expert output + norm has nowhere to go — it would be dropped silently, so refuse instead. The other DeepGEMM + arms reduce in `_combine_routed_output` and apply the norm on the rows just before it.""" if getattr(self, "post_expert_norm", None) is not None: raise NotImplementedError( - "DeepGEMM experts dispatch cannot apply this model's per-expert output norm; use " - "`experts_implementation='grouped_mm'` (or 'batched_mm')." + "DeepGEMM Mega MoE cannot apply this model's per-expert output norm; use " + "`experts_implementation='deepgemm'`, 'grouped_mm' or 'batched_mm'." ) -def deepgemm_experts_guards(*, affine_scales: bool = False, sm100: bool = False): +def _apply_post_expert_norm(self: torch.nn.Module, rows: torch.Tensor) -> torch.Tensor: + """The model's per-expert output norm, on the routed rows before the routing weights — the + same place the reference forwards apply it. A no-op for a model that names none.""" + if getattr(self, "post_expert_norm_name", None) is None: + return rows + return self._apply_post_norm(rows) + + +def deepgemm_experts_guards(*, affine_scales: bool = False, sm100: bool = False, post_expert_norm: bool = True): """State an experts forward's requirements on the module, checked before any kernel work. - Every arm refuses a per-expert output norm. ``affine_scales`` additionally refuses the + ``post_expert_norm=False`` refuses a model that names a per-expert output norm — Mega MoE + only, whose fused reduce has no seam for it. ``affine_scales`` additionally refuses the SWIZZLE_32_4_4 scales a module loaded for a triton backend holds, and ``sm100`` fails before the hub download + JIT when the device cannot serve these dtypes.""" def decorate(forward): @functools.wraps(forward) def guarded(self, *args, **kwargs): - _assert_no_post_expert_norm(self) + if not post_expert_norm: + _assert_no_post_expert_norm(self) if affine_scales: _assert_affine_scales(self) if sm100: @@ -743,7 +754,7 @@ def deepgemm_bf16_experts_forward( out.index_add_(0, sorted_to_padded, down_bias[expert_ids_g]) return _combine_routed_output( - out, + _apply_post_expert_norm(self, out), sorted_weights, sentinel_mask, perm, @@ -865,7 +876,7 @@ def deepgemm_fp8_fp4_experts_forward( ) return _combine_routed_output( - out, + _apply_post_expert_norm(self, out), sorted_weights, sentinel_mask, perm, @@ -941,7 +952,7 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) -@deepgemm_experts_guards(affine_scales=True, sm100=True) +@deepgemm_experts_guards(affine_scales=True, sm100=True, post_expert_norm=False) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 778ca7e2c90b..391133ef7a4c 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -375,6 +375,36 @@ def _bundle(self, *, is_sm100): # ── deepgemm_fp8_fp4_linear ──────────────────────────────────────────────── + def test_post_expert_norm_runs_on_the_reducing_arms_and_is_refused_by_megamoe(self): + """Mega MoE fuses the routing-weighted reduce, so a per-expert output norm has nowhere to + go and is refused. The other arms reduce in `_combine_routed_output`, so the norm rides the + rows just before it — the same place the reference forwards apply it.""" + import transformers.integrations.deepgemm as dg + + applied = [] + + class _Experts(torch.nn.Module): + post_expert_norm_name = "input_scaled_rms_norm" + post_expert_norm = object() + + def _apply_post_norm(self, rows): + applied.append(tuple(rows.shape)) + return rows + + module = _Experts() + rows = torch.randn(4, 8) + # the reducing arms: the norm is applied and the rows pass through + self.assertIs(dg._apply_post_expert_norm(module, rows), rows) + self.assertEqual(applied, [(4, 8)]) + + # a model naming no norm is untouched + plain = torch.nn.Module() + self.assertIs(dg._apply_post_expert_norm(plain, rows), rows) + + # mega moe still refuses, and only mega moe + with self.assertRaises(NotImplementedError): + dg._assert_no_post_expert_norm(module) + def test_linear_fp8_sm90_kernel_inputs(self): # FP8 weights + float32 block SF on SM90: recipe stays None, SFs are handed over row-major # float32 (SM90 dispatch transforms SFA itself and only checks SFB — see `_coerce_sf_for_kernel`). From 43b9a748d397a69181023dbb621c6d986cd0856f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 17 Sep 2026 06:47:55 -0700 Subject: [PATCH 28/68] Shard quantized expert companions under EP and TP A stacked experts projection is a Parameter, not a Module, so its scale grid, bias and NVFP4 globals are SIBLINGS rather than children: the plan matcher's one module-level fallback never reaches them, and the module entry shards nothing. Each companion has to be named in the plan or the weight becomes this rank's expert slice while its scales stay whole. EP had a hand-written rule for exactly one of them (`grouped_gemm` -> `_scale_inv`) that this branch had dropped; restore it and generalise to `_bias`, both NVFP4 globals and a static activation scale. The gate_up input global stays replicated: one per-tensor value for the pre-routing hidden states. TP never had the equivalent rule at all, since it uses packed_colwise / rowwise and that loop only matched `grouped_gemm`. Add it, plus a second fix: packed_colwise splits the output axis as two packed halves, which is the [gate; up] STACK, while a quantized module holds the kernels' INTERLEAVED rows. Register param-only styles for the axis actually wanted -- the built-in styles count from the END of the shape, landing inside a swizzled scale grid's inner tile and on a bias's expert axis. Verified on 2 GPUs against a pre-quantized DeepSeek-V3 fixture: EP and TP logits match an unsharded reference, and removing the TP half makes that test fail with the scale left whole against a sharded weight. 8-rank smoke runs pass on DeepSeek-V4-Flash (EP) and DeepSeek-V3 (TP); device_map smoke passes on GLM-5.2-NVFP4, DeepSeek-V3, MiniMax-M3-MXFP8 and OLMoE. Also in this pass, from the second review: - `to_local` deprecated rather than deleted; every call site is gone now that `_hf_quantized_needs_local_tp` hands modules local tensors. - `validate_environment` gains its first tests -- the only ones that existed built the frozen FP8 quantizer, which no `from_pretrained` can reach since AUTO_QUANTIZER_MAPPING["fp8"] became FineGrainedHfQuantizer. - `_swizzles_scales` no longer claims block-FP8 never reaches a scaled MMA; it does, via UE8M0, but broadcasts one 128-block scalar in register so there is no per-group grid to lay out. - modelopt remapped to nvfp4 at construction; glob skip-lists moved to the shared `subtree_pattern_to_regex`; GPT-OSS carries swiglu_alpha / swiglu_limit on its config. Co-Authored-By: Claude Opus 5 --- src/transformers/core_model_loading.py | 30 +- .../distributed/sharding_utils.py | 11 +- src/transformers/integrations/deepgemm.py | 138 +++--- src/transformers/integrations/finegrained.py | 450 ++++++++---------- .../integrations/finegrained_fp8.py | 15 +- src/transformers/integrations/hub_kernels.py | 2 +- src/transformers/integrations/moe.py | 30 +- .../integrations/tensor_parallel.py | 18 +- .../models/gpt_oss/configuration_gpt_oss.py | 9 + .../models/gpt_oss/modeling_gpt_oss.py | 4 +- .../models/gpt_oss/modular_gpt_oss.py | 4 +- .../configuration_openai_privacy_filter.py | 9 + .../modeling_openai_privacy_filter.py | 4 +- .../quantizers/quantizer_finegrained.py | 449 ++++++++++------- .../quantizers/quantizers_utils.py | 16 + src/transformers/utils/quantization_config.py | 50 +- tests/kernels/test_deepgemm.py | 41 +- tests/kernels/test_finegrained.py | 391 ++++++++++++++- tests/kernels/test_kernels.py | 4 +- .../quantization/finegrained_fp8/test_fp8.py | 3 +- tests/utils/test_import_utils.py | 8 +- 21 files changed, 1074 insertions(+), 612 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 0c48db3bfffc..1526b825ea64 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,11 +81,6 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" - # Whether this op maps SEVERAL sources onto SEVERAL targets in one call — a converter is - # otherwise restricted to one-to-many, one-to-one or many-to-one, since an m:n mapping only - # makes sense when a single operation owns the whole relation (see `WeightConverter`). - supports_many_to_many: bool = False - def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -618,8 +613,6 @@ class ErnieFuseAndSplitTextVisionExperts(ConversionOps): The final fusions are defined by the amount of original module lists. """ - supports_many_to_many = True - def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -682,8 +675,6 @@ class ErnieSplitAndDecoupleTextVisionExperts(ConversionOps): The final decoupled module lists are defined by the amount of keys. """ - supports_many_to_many = True - def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -1151,6 +1142,23 @@ def reverse_transform(self) -> WeightTransform: return result +def _internal_many_to_many_conversions() -> tuple[type[ConversionOps], ...]: + """The classes known to be able to use m:n. + + A function, not the module-level tuple it would rather be: `FineGrainedWeightGlobals` lives in + `integrations.finegrained`, which imports THIS module, so naming it at import time fails + whenever the integration is imported first (`cannot import name ... from partially initialized + module`). Called once per `WeightConverter`, at construction. + """ + from .integrations.finegrained import FineGrainedWeightGlobals + + return ( + ErnieFuseAndSplitTextVisionExperts, + ErnieSplitAndDecoupleTextVisionExperts, + FineGrainedWeightGlobals, + ) + + class WeightConverter(WeightTransform): __slots__ = ("operations", "force_cpu") @@ -1166,8 +1174,8 @@ def __init__( self.force_cpu = force_cpu if bool(len(self.source_patterns) - 1) + bool(len(self.target_patterns) - 1) >= 2: - # We allow many-to-many only if an operation declares that it handles the whole relation - if not any(op.supports_many_to_many for op in self.operations): + # We allow many-to-many only if we use an internal operation that can handle it + if not any(isinstance(op, _internal_many_to_many_conversions()) for op in self.operations): raise ValueError( f"source keys={self.source_patterns}, target_patterns={self.target_patterns} but you can only have one to many, one to one or many to one." ) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index ca3cb2c2f75a..3e5e13136bb4 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -145,7 +145,6 @@ def shard_tensor( has_strided_shard = any(not placement.is_shard() for _, placement in dim_placements) # finally fetch from the disk only the slices - # finally fetch from the disk only the slices if has_strided_shard: # Multi-interval dim: read each piece separately, then concatenate. return self._slice_and_cat(source, intervals_by_dim, device, dtype) @@ -159,16 +158,14 @@ def shard_tensor( # MoE path # tensor_idx identifies the axis-0 piece in param space (not in source.shape). + # if this rank owns expert `tensor_idx` along axis 0, we need to slice the inner dimensions, else we drop it + if not self._owns_expert(tensor_idx, dim_placements): + return None + normalized_dim_placements = [ (mesh_dim, placement, self._normalize_param_dim(placement.dim)) for mesh_dim, placement in dim_placements ] - # if this rank owns expert `tensor_idx` along axis 0, we need to slice the inner dimensions, else we drop it - has_axis0_shard = any(param_dim == 0 for _, _, param_dim in normalized_dim_placements) - owns_tensor_idx = self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size - if has_axis0_shard and not owns_tensor_idx: - return None - # `param_dim` indexes the full parameter layout [N, in, out] (expert axis first). # In per-expert loading, leading axis is absent ([in, out]). # Therefore we need to shift the dimensions by 1 to the left so that sharding operations can be applied to the inner dimensions. diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 82d93e6db906..5d153fc0ea15 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -37,14 +37,12 @@ import torch from ..utils import logging -from ..utils.deprecation import deprecate_kwarg from ..utils.import_utils import ( is_kernels_available, maybe_import_error, resolve_internal_import, ) from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel -from .tensor_parallel import to_local logger = logging.get_logger(__name__) @@ -571,7 +569,6 @@ def _combine_routed_output( # ── Public dispatches ────────────────────────────────────────────────────────── -@deprecate_kwarg("output_dtype", version="v5.16") def prefers_deepgemm_linear( weight: torch.Tensor, weight_scale_inv: torch.Tensor, @@ -598,7 +595,6 @@ def prefers_deepgemm_linear( # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. and weight_scale_inv.ndim <= 2 - and is_deepgemm_loadable() and activation_scale is None # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format and weight_global_scale is None @@ -606,6 +602,9 @@ def prefers_deepgemm_linear( and activation_format is None and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" + # last: an environment probe (arch, `CUDA_HOME/bin/nvcc`, nvcc version) rather than a + # tensor check, so the cheap gates above short-circuit before it + and is_deepgemm_loadable() ) @@ -653,11 +652,47 @@ def deepgemm_fp8_fp4_linear( return output +def _assert_dynamic_activations(self: torch.nn.Module) -> None: + """DeepGEMM quantizes activations inside its kernels, so a checkpoint's calibrated per-tensor + scale has nowhere to be applied.""" + if getattr(self, "activation_scheme", "dynamic") == "static": + raise NotImplementedError( + "DeepGEMM experts dispatch does not support activation_scheme='static'; use " + "`experts_implementation='grouped_mm'` (or 'batched_mm')." + ) + + +def _assert_affine_scales(module: torch.nn.Module) -> None: + """DeepGEMM reads row-major block scales. A module loaded for a triton backend holds them in the + SWIZZLE_32_4_4 layout (5-D); consuming that as affine returns garbage, so a backend switched to + DeepGEMM after such a load fails here instead.""" + for name in ("gate_up_proj_scale_inv", "up_proj_scale_inv", "down_proj_scale_inv"): + scale = getattr(module, name, None) + if scale is not None and scale.ndim == 5: + raise NotImplementedError( + f"DeepGEMM experts need row-major block scales, but `{name}` is held swizzled for the " + "triton experts backends. Switching to a DeepGEMM backend after load is not supported — " + "pass `experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + +def _assert_stacked_gate_up(module: torch.nn.Module) -> None: + """Mega MoE's ``transform_weights_for_mega_moe`` does its own gate/up interleave, so it takes the + stacked ``[gate; up]`` rows. A module loaded for a triton backend holds them interleaved; switched + to this backend after load, silently re-permuting would be a wrong answer rather than a failure.""" + if getattr(module, "has_gate", True) and getattr(module, "holds_interleaved_gate_up", False): + raise NotImplementedError( + "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " + "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — pass " + "`experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: """Mega MoE fuses the routing-weighted reduce into its kernel, so a model's per-expert output norm has nowhere to go — it would be dropped silently, so refuse instead. The other DeepGEMM arms reduce in `_combine_routed_output` and apply the norm on the rows just before it.""" - if getattr(self, "post_expert_norm", None) is not None: + if getattr(self, "has_post_expert_norm", False): raise NotImplementedError( "DeepGEMM Mega MoE cannot apply this model's per-expert output norm; use " "`experts_implementation='deepgemm'`, 'grouped_mm' or 'batched_mm'." @@ -666,46 +701,64 @@ def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: def _apply_post_expert_norm(self: torch.nn.Module, rows: torch.Tensor) -> torch.Tensor: """The model's per-expert output norm, on the routed rows before the routing weights — the - same place the reference forwards apply it. A no-op for a model that names none.""" - if getattr(self, "post_expert_norm_name", None) is None: + same place the reference forwards apply it. A no-op for a model that declares none.""" + if not getattr(self, "has_post_expert_norm", False): return rows return self._apply_post_norm(rows) -def deepgemm_experts_guards(*, affine_scales: bool = False, sm100: bool = False, post_expert_norm: bool = True): +def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: + """Every DeepGEMM experts arm builds its intermediates and its output in bfloat16, so anything + else silently changes the dtype the caller gets back.""" + if hidden_states.dtype != torch.bfloat16: + raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") + + +def deepgemm_experts_guards( + forward=None, + *, + affine_scales: bool = False, + sm100: bool = False, + post_expert_norm: bool = True, + stacked_gate_up: bool = False, +): """State an experts forward's requirements on the module, checked before any kernel work. - ``post_expert_norm=False`` refuses a model that names a per-expert output norm — Mega MoE - only, whose fused reduce has no seam for it. ``affine_scales`` additionally refuses the - SWIZZLE_32_4_4 scales a module loaded for a triton backend holds, and ``sm100`` fails before - the hub download + JIT when the device cannot serve these dtypes.""" + Usable bare (`@deepgemm_experts_guards`) for an arm that only needs the common checks. + + ``post_expert_norm=False`` refuses a model that declares a per-expert output norm — Mega MoE + only, whose fused reduce has no seam for it. ``affine_scales`` refuses the SWIZZLE_32_4_4 + scales a module loaded for a triton backend holds, ``stacked_gate_up`` refuses the interleaved + rows that same load produces, and ``sm100`` fails before the hub download + JIT when the device + cannot serve these dtypes.""" def decorate(forward): @functools.wraps(forward) - def guarded(self, *args, **kwargs): + def guarded(self, hidden_states, *args, **kwargs): + _assert_bf16_hidden_states(hidden_states) + _assert_dynamic_activations(self) if not post_expert_norm: _assert_no_post_expert_norm(self) if affine_scales: _assert_affine_scales(self) + if stacked_gate_up: + _assert_stacked_gate_up(self) if sm100: _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) - return forward(self, *args, **kwargs) + return forward(self, hidden_states, *args, **kwargs) return guarded - return decorate + return decorate(forward) if forward is not None else decorate -@deepgemm_experts_guards() +@deepgemm_experts_guards def deepgemm_bf16_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: - if hidden_states.dtype != torch.bfloat16: - raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") - deepgemm = load_deepgemm_kernel() # Non-transposed weights (E, N, K) → NT kernel; transposed (E, K, N) → NN kernel. grouped_bf16_matmul = deepgemm.grouped_bf16_matmul_nn if self.is_transposed else deepgemm.grouped_bf16_matmul_nt @@ -766,32 +819,6 @@ def deepgemm_bf16_experts_forward( ) -def _assert_affine_scales(module: torch.nn.Module) -> None: - """DeepGEMM reads row-major block scales. A module loaded for a triton backend holds them in the - SWIZZLE_32_4_4 layout (5-D); consuming that as affine returns garbage, so a backend switched to - DeepGEMM after such a load fails here instead.""" - for name in ("gate_up_proj_scale_inv", "up_proj_scale_inv", "down_proj_scale_inv"): - scale = getattr(module, name, None) - if scale is not None and scale.ndim == 5: - raise RuntimeError( - f"DeepGEMM experts need row-major block scales, but `{name}` is held swizzled for the " - "triton experts backends. Switching to a DeepGEMM backend after load is not supported — " - "pass `experts_implementation=...` to `from_pretrained` so the model loads for it." - ) - - -def _assert_stacked_gate_up(module: torch.nn.Module) -> None: - """Mega MoE's ``transform_weights_for_mega_moe`` does its own gate/up interleave, so it takes the - stacked ``[gate; up]`` rows. A module loaded for a triton backend holds them interleaved; switched - to this backend after load, silently re-permuting would be a wrong answer rather than a failure.""" - if getattr(module, "has_gate", True) and getattr(module, "holds_interleaved_gate_up", False): - raise RuntimeError( - "Mega MoE needs gate|up stacked, but this module was loaded interleaved for the triton " - "experts backend. Switching to 'deepgemm_megamoe' after load is not supported — pass " - "`experts_implementation=...` to `from_pretrained` so the model loads for it." - ) - - @deepgemm_experts_guards(affine_scales=True, sm100=True) def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, @@ -808,11 +835,6 @@ def deepgemm_fp8_fp4_experts_forward( "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) - if self.activation_scheme == "static": - raise NotImplementedError("DeepGEMM experts dispatch does not support static activation quantization.") - if hidden_states.dtype != torch.bfloat16: - raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") - deepgemm = load_deepgemm_kernel() grouped_fp8_fp4_matmul = ( @@ -902,15 +924,9 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: the ``[E_local, 2*I, *]`` leading dims so downstream ``.size(...)`` reads stay valid. - Unwraps any ``DTensor`` wrappers FSDP2/EP may have placed around the loader- - side Parameters — the kernel takes raw pointers. + Runs inside the first forward, where the parameters are already local: the module carries + ``_hf_quantized_needs_local_tp``, so the TP layer hands local shards rather than DTensors. """ - # A module loaded for a triton backend holds swizzled block scales, which DeepGEMM would read as - # affine and turn into garbage — refuse before the (hub-download + JIT) kernel load. - _assert_affine_scales(module) - # Same for its interleaved gate|up rows: `transform_weights_for_mega_moe` interleaves stacked rows - # itself, and re-permuting already-interleaved ones would be a silent wrong answer. - _assert_stacked_gate_up(module) # The UE8M0 scale grid is group-32 along both dims, so the kernel's SF layout needs both divisible. intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts @@ -925,8 +941,8 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: gate_up_sf_raw = module.gate_up_proj_scale_inv.data down_sf_raw = module.down_proj_scale_inv.data # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. - gate_up_w = to_local(module.gate_up_proj.data).view(torch.int8).contiguous() - down_w = to_local(module.down_proj.data).view(torch.int8).contiguous() + gate_up_w = module.gate_up_proj.data.view(torch.int8).contiguous() + down_w = module.down_proj.data.view(torch.int8).contiguous() gate_up_sf = deepgemm.transform_sf_into_required_layout( gate_up_sf_raw.float(), @@ -952,7 +968,7 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) -@deepgemm_experts_guards(affine_scales=True, sm100=True, post_expert_norm=False) +@deepgemm_experts_guards(affine_scales=True, sm100=True, post_expert_norm=False, stacked_gate_up=True) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index ada31c7da7ef..b2348098565d 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -20,7 +20,6 @@ from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass -from types import MethodType import torch import torch.nn as nn @@ -65,11 +64,15 @@ def _get_ue8m0_dtype() -> torch.dtype: return torch.float8_e8m0fnu -def _first_attr(obj, *names): +def _first_attr(obj, *names, raise_error: bool = True): + """The first of `names` the object defines; `None` when it defines none and the caller passes + `raise_error=False` because it has its own fallback.""" for name in names: if hasattr(obj, name): return getattr(obj, name) - raise AttributeError(f"{type(obj).__name__} has none of: {names}") + if raise_error: + raise AttributeError(f"{type(obj).__name__} has none of: {names}") + return None @dataclass(frozen=True) @@ -186,16 +189,18 @@ class WeightFormat: """Storage layout of one quantized weight format — the single source both module classes derive their parameter shapes from. ``scale_dtype`` ``None`` defers to the config's ``scale_fmt`` (block-FP8 ships fp32 or UE8M0 containers; the group formats pin theirs). - ``scale_group`` ``None`` means the block comes from the quant config's ``weight_block_size``.""" + ``scale_group`` ``None`` means the block comes from the quant config's ``weight_block_size``. + ``global_scale_dtype`` ``None`` means the format has no second level; NVFP4 names the dtype its + per-matrix global is held in.""" weight_dtype: torch.dtype values_per_byte: int = 1 scale_dtype: torch.dtype | None = None scale_group: tuple[int, int] | None = None - has_global_scale: bool = False + global_scale_dtype: torch.dtype | None = None -def _weight_formats() -> dict[str, WeightFormat]: +def weight_formats() -> dict[str, WeightFormat]: return { # block-scaled E4M3, block from the quant config, fp32/UE8M0 scale container "fp8": WeightFormat(weight_dtype=_FP8_DTYPE), @@ -211,7 +216,7 @@ def _weight_formats() -> dict[str, WeightFormat]: values_per_byte=2, scale_dtype=_FP8_DTYPE, scale_group=(1, 16), - has_global_scale=True, + global_scale_dtype=torch.float32, ), } @@ -223,14 +228,14 @@ def resolve_weight_format( ) -> tuple[WeightFormat, torch.dtype, tuple[int, int] | None]: """``(format, scale_dtype, (sf_gran_n, sf_gran_k))`` for one format name, with the config's ``scale_fmt``/``weight_block_size`` filling the slots the format leaves open.""" - formats = _weight_formats() + formats = weight_formats() if weight_format not in formats: raise ValueError(f"unknown weight_format {weight_format!r}; expected one of {sorted(formats)}") - fmt = formats[weight_format] - scale_dtype = fmt.scale_dtype + format_spec = formats[weight_format] + scale_dtype = format_spec.scale_dtype if scale_dtype is None: scale_dtype = _get_ue8m0_dtype() if scale_fmt == "ue8m0" else torch.float32 - return fmt, scale_dtype, fmt.scale_group if fmt.scale_group is not None else block_size + return format_spec, scale_dtype, format_spec.scale_group if format_spec.scale_group is not None else block_size def _set_optional_parameter(module: nn.Module, name: str, tensor: torch.Tensor | None) -> None: @@ -243,7 +248,7 @@ def _alloc_expert_proj( num_experts: int, proj_out: int, proj_in: int, - fmt: WeightFormat, + format_spec: WeightFormat, scale_dtype: torch.dtype, scale_group: tuple[int, int] | None, min_scale_out: int = 1, @@ -257,7 +262,7 @@ def _alloc_expert_proj( layout instead — 128-row blocks x 4-column groups of ``(2, 256)`` byte tiles, the expert axis leading so EP shards and gathers it on dim 0 — which needs whole blocks; a projection that does not tile stays affine, which the kernels read directly.""" - weight = torch.empty(num_experts, proj_out, proj_in // fmt.values_per_byte, dtype=fmt.weight_dtype) + weight = torch.empty(num_experts, proj_out, proj_in // format_spec.values_per_byte, dtype=format_spec.weight_dtype) if scale_group is None: scale_shape = (num_experts, max(1, min_scale_out), 1) else: @@ -272,16 +277,17 @@ def _alloc_expert_proj( ) -def _swizzles_scales(config, fmt: WeightFormat, activation_format: str | None) -> bool: +def _swizzles_scales(config, format_spec: WeightFormat, activation_format: str | None) -> bool: """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05 scaled-MMA reads directly (plain row-major forces a per-tile gather that caps the scaled dot below the fp8/fp4 peak): SM100, a triton dispatch (the fused forwards and the eager loop both read it; the DeepGEMM backends read affine scales), a group-scaled format (MX group-32, NVFP4 - group-16 — block-FP8's grid never reaches a scaled-MMA, even with UE8M0 scales) and a chain that - quantizes activations (weight-only reads scales per group affinely).""" + group-16; block-FP8 reaches the scaled-MMA too under UE8M0, but by broadcasting one 128-block + scalar in-register across the group-32 columns it consumes, so it has no per-group grid to lay + out) and a chain that quantizes activations (weight-only reads scales per group affinely).""" return ( getattr(config, "_experts_implementation", None) not in ("deepgemm", "deepgemm_megamoe") - and fmt.scale_group is not None + and format_spec.scale_group is not None and activation_format != "bf16" and torch.cuda.is_available() and is_sm100() @@ -349,8 +355,8 @@ def finegrained_linear( Args: input: (..., K) bf16/fp16 activations. weight: (N, K) `float8_e4m3fn` or (N, K // 2) `int8` (FP4-packed). - weight_scale_inv: per-block weight scales — `float32` (V3-style) or `float8_e8m0fnu` - (V4-style; reinterpreted as int32 at the DeepGEMM kernel boundary). + weight_scale_inv: per-block weight scales — `float32` (as DeepSeek-V3 ships them) or + `float8_e8m0fnu` (DeepSeek-V4; reinterpreted as int32 at the DeepGEMM kernel boundary). block_size: [block_n, block_k] for FP8 block-wise quant, or None/[N, K] for per-tensor. Ignored for FP4 weights (the kernel infers SF granularity from the dtype). bias: optional bias added to the matmul output. @@ -436,7 +442,6 @@ def __init__( has_bias: bool = False, weight_format: str = "fp8", activation_format: str | None = None, - has_global_scale: bool | None = None, ): super().__init__(in_features, out_features) @@ -447,29 +452,29 @@ def __init__( self.activation_format = activation_format # the format table decides storage: value dtype (packed E2M1 = 2 values per int8 byte # along K), scale dtype/granularity, and whether a two-level global exists - fmt, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) - if has_global_scale is None: - has_global_scale = fmt.has_global_scale - in_storage = in_features // fmt.values_per_byte + format_spec, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + # the format's own dtype both decides whether there is a second level and says what it + # is held in + global_scale_dtype = format_spec.global_scale_dtype + in_storage = in_features // format_spec.values_per_byte self.weight = torch.nn.Parameter( - torch.empty(out_features, in_storage, dtype=fmt.weight_dtype), - requires_grad=fmt.weight_dtype.is_floating_point, + torch.empty(out_features, in_storage, dtype=format_spec.weight_dtype), + requires_grad=format_spec.weight_dtype.is_floating_point, ) # NVFP4 two-level: the per-tensor fp32 global the kernel recovers on the accumulator _set_optional_parameter( - self, "weight_global_scale", torch.tensor(1.0, dtype=torch.float32) if has_global_scale else None + self, + "weight_global_scale", + torch.tensor(1.0, dtype=global_scale_dtype) if global_scale_dtype is not None else None, ) - if scale_group is None: - # no group and no block: one per-tensor scale - self.weight_scale_inv = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)) - else: - scale_out_features = _cdiv(out_features, scale_group[0]) - scale_in_features = _cdiv(in_features, scale_group[1]) - self.weight_scale_inv = nn.Parameter( - torch.empty(scale_out_features, scale_in_features, dtype=sf_dtype), - requires_grad=sf_dtype.is_floating_point, - ) + # no group and no block: one per-tensor scale + scale = ( + torch.tensor(1.0, dtype=torch.float32) + if scale_group is None + else torch.empty(_cdiv(out_features, scale_group[0]), _cdiv(in_features, scale_group[1]), dtype=sf_dtype) + ) + self.weight_scale_inv = nn.Parameter(scale, requires_grad=scale.dtype.is_floating_point) static = self.activation_scheme == "static" _set_optional_parameter(self, "activation_scale", torch.tensor(1.0, dtype=torch.float32) if static else None) @@ -510,16 +515,14 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: class FineGrainedGroupedLinear(FineGrainedLinear): - """FP8 drop-in for block-diagonal grouped linears. - - The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)` - weight; logically that's `n_groups` independent `(out_per_group, in_per_group)` - sub-matrices, each consuming a disjoint slice of the input's last-but-one dim. - Forward expects input of shape `(..., n_groups, in_per_group)` and returns - `(..., n_groups, out_per_group)` — same contract as the vanilla bf16 grouped - linear it replaces. + """Quantized drop-in for block-diagonal grouped linears, in any format `FineGrainedLinear` + serves — block-FP8, MXFP8, MXFP4, NVFP4 or weight-only. - """ + The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)` weight; + logically that is `n_groups` independent `(out_per_group, in_per_group)` sub-matrices, each + consuming a disjoint slice of the input's last-but-one dim. Forward takes `(..., n_groups, + in_per_group)` and returns `(..., n_groups, out_per_group)` — the same contract as the bf16 + grouped linear it replaces.""" def __init__( self, @@ -582,9 +585,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _norm_eps(norm: nn.Module) -> float: """A norm module's epsilon, under either name transformers modules give it. ``nn.RMSNorm`` leaves it ``None`` and resolves it to the dtype's own, so this does the same.""" - eps = getattr(norm, "eps", None) - if eps is None: - eps = getattr(norm, "variance_epsilon", None) + eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) return float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps @@ -603,12 +604,12 @@ def _moe_operands(kernel, module) -> dict: else: act_fn = module._apply_gate if module.has_gate else module.act_fn - norm, norm_eps, norm_weight = module.post_expert_norm, 1e-6, None - if norm is None: + norm_eps, norm_weight = 1e-6, None + if not module.has_post_expert_norm: post_expert_norm = None elif module.post_expert_norm_name in kernel.get_supported_norms(): post_expert_norm = module.post_expert_norm_name - norm_weight, norm_eps = norm.weight, _norm_eps(norm) + norm_weight, norm_eps = module.post_expert_norm.weight, _norm_eps(module.post_expert_norm) else: post_expert_norm = module._apply_post_norm @@ -665,18 +666,12 @@ def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top class FineGrainedExperts(_FineGrainedModule, nn.Module): - # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. - # The plan dicts store `{module-path-pattern: parallel-layer-kind}`; this maps an - # old kind to a new kind, and the quantizer rewrites every plan VALUE that matches. - # The default `MoeTensorParalellExperts` kind is impl-agnostic; some impls need a - # distinct TP layer (e.g. megamoe needs no gradient-sync hooks and an EP - # `process_group` injection). Declared here so the quantizer doesn't have to know - # about impl-specific TP needs — extend this dict when adding new impls. + # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. The default + # kind is impl-agnostic; some impls need a distinct one (megamoe needs no gradient-sync hooks + # and an EP `process_group` injection). Declared here so the quantizer does not have to know + # about impl-specific TP needs — extend when adding an impl. _impl_tp_layer_overrides: dict[str, dict[str, str]] = { - "deepgemm_megamoe": { - "moe_tp_experts": "megamoe_experts", - "ep_router": "megamoe_router", - }, + "deepgemm_megamoe": {"moe_tp_experts": "megamoe_experts", "ep_router": "megamoe_router"}, } def __init__( @@ -689,7 +684,6 @@ def __init__( has_gate: bool = True, weight_format: str | None = None, activation_format: str | None = None, - has_global_scale: bool | None = None, is_concatenated: bool = True, ): super().__init__() @@ -720,10 +714,10 @@ def __init__( if weight_format is None: weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" self.weight_format = weight_format - fmt, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) - # the forwards gate on this attribute, so it carries the format-resolved value (a None ctor - # arg would silently drop the nvfp4 global at every forward) - self.has_global_scale = fmt.has_global_scale if has_global_scale is None else has_global_scale + format_spec, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + # the format decides the second level: whether there is one, and the dtype it is held in + self.global_scale_dtype = format_spec.global_scale_dtype + self.has_global_scale = self.global_scale_dtype is not None # What the module HOLDS, as the loader's conversion ops deliver it (and their reverses # restore for the checkpoint): gate|up rows in the kernels' interleaved order unless the @@ -732,7 +726,7 @@ def __init__( # loader interleaves where the two disagree, and everything at runtime reads this one. impl = getattr(config, "_experts_implementation", None) self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" - swizzled = _swizzles_scales(config, fmt, activation_format) + swizzled = _swizzles_scales(config, format_spec, activation_format) up_name = "gate_up_proj" if self.has_gate else "up_proj" up_rows = (2 if self.has_gate else 1) * self.intermediate_dim @@ -741,10 +735,10 @@ def __init__( ("down_proj", self.hidden_dim, self.intermediate_dim, 1), ): weight, scale = _alloc_expert_proj( - self.num_experts, rows, in_dim, fmt, scale_dtype, scale_group, min_scale_out, swizzled + self.num_experts, rows, in_dim, format_spec, scale_dtype, scale_group, min_scale_out, swizzled ) - setattr(self, proj, weight) - setattr(self, f"{proj}_scale_inv", scale) + self.register_parameter(proj, weight) + self.register_parameter(f"{proj}_scale_inv", scale) # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype _set_optional_parameter( @@ -758,7 +752,7 @@ def __init__( _set_optional_parameter( self, f"{proj}_weight_global_scale", - torch.ones(self.num_experts, dtype=torch.float32) if self.has_global_scale else None, + torch.ones(self.num_experts, dtype=self.global_scale_dtype) if self.has_global_scale else None, ) calibrated = self.has_global_scale and self.activation_format != "bf16" _set_optional_parameter( @@ -771,18 +765,22 @@ def __init__( self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) self.down_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) - # The model's per-expert output norm and the name it gave the form (its - # `use_experts_implementation(post_expert_norm=...)`), both carried over by the swap. - # A submodule slot, not a class attribute: a class attribute of the same name shadows - # what `nn.Module.__setattr__` registers, so the norm would never be read back. + # The model's per-expert output norm, carried over by the swap. A submodule slot, not a + # class attribute: a class attribute of the same name shadows what + # `nn.Module.__setattr__` registers, so the norm would never be read back. self.post_expert_norm = None + self.has_post_expert_norm = False + # The form it implements, when the model names one it wants a backend to FUSE + # (`get_supported_norms`); unnamed norms still run through `_apply_post_norm`. self.post_expert_norm_name = None + def _apply_post_norm(self, expert_out: torch.Tensor) -> torch.Tensor: + """The standard hook every backend applies: the model's norm on one expert application's + rows, before the routing weights.""" + return self.post_expert_norm(expert_out) + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: - # the interleaved rows give alternating output columns, mirroring the fused epilogue's - # `split_gate_up` that the fused/unfused parity tests compare against. Only the chains - # that hold that layout come through here — DeepGEMM Mega MoE, the one backend keeping - # the stack, runs its own forward end to end. + # interleaved rows -> alternating output columns, the same split the fused epilogue does gate, up = gate_up[..., 0::2], gate_up[..., 1::2] if self.swiglu_alpha is not None: # Clamped SwiGLU-OAI gate (same math as the model's non-quantized experts). @@ -827,7 +825,7 @@ def forward( self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None ) proj_out = self.linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) - if self.post_expert_norm is not None: + if self.has_post_expert_norm: proj_out = self._apply_post_norm(proj_out) routing_weights = top_k_weights[token_idx, top_k_pos, None] weighted_out = proj_out * routing_weights.to(proj_out.dtype) @@ -936,12 +934,7 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No # The checkpoint's quant_method IS the weight format ("fp8", "mxfp8", "mxfp4", "nvfp4"). # Under the "fp8" key pass None: dsv4-style checkpoints declare fp4 EXPERTS through the # legacy `config.expert_dtype` model-config side-channel, which the ctor falls back to. - quant_method = getattr(quantization_config, "quant_method", None) - quant_method = getattr(quant_method, "value", quant_method) - # modelopt exports are NVFP4-only (validated by the config translation) - if quant_method == "modelopt": - quant_method = "nvfp4" - weight_format = quant_method if quant_method in ("mxfp8", "mxfp4", "nvfp4") else None + weight_format = next((f for f in ("mxfp8", "mxfp4", "nvfp4") if quantization_config.quant_method == f), None) has_been_replaced = False for module_name, module in model.named_modules(): @@ -975,27 +968,15 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No activation_format=getattr(quantization_config, "activation_format", None), is_concatenated=is_concatenated, ) - # GPT-OSS hardcodes its clamped-SwiGLU parameters as MODULE attributes - # (alpha=1.702, limit=7.0) rather than config fields — carry them over so - # `_apply_gate` (and the fused epilogue) keep the reference numerics. - if getattr(module, "alpha", None) is not None and new_module.swiglu_alpha is None: - new_module.swiglu_alpha = module.alpha - new_module.swiglu_limit = getattr(module, "limit", None) - # A per-expert output norm belongs to the model, not the quantization: carry the - # submodule, the name, and the model's own `_apply_post_norm` over, or the swap - # would drop the norm or quietly substitute different math for applying it. The - # method is rebound to THIS module, which holds the same submodule — binding the - # original's would pin the unquantized weights alive. + # A per-expert output norm belongs to the model, not the quantization, so the + # swap carries the submodule over; `_apply_post_norm` is the standard hook and + # needs no rebinding. A named form is what a backend can fuse. if getattr(module, "post_expert_norm", None) is not None: - apply_post_norm = getattr(type(module), "_apply_post_norm", None) - if apply_post_norm is None: - raise TypeError( - f"{type(module).__name__} holds a post-expert norm but does not define " - "`_apply_post_norm(self, expert_out)`; see `use_experts_implementation`." - ) new_module.post_expert_norm = module.post_expert_norm + # `use_experts_implementation` owns this flag on every experts class and sets + # it from the class declaration, so the swap updates it rather than deriving it + new_module.has_post_expert_norm = True new_module.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) - new_module._apply_post_norm = MethodType(apply_post_norm, new_module) elif type(module) is nn.Linear: # Vanilla `nn.Linear` → standard FineGrainedLinear swap. new_module = FineGrainedLinear( @@ -1025,13 +1006,18 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No has_bias=module.bias is not None, ) if new_module is not None: + # The kernels take raw pointers, so every operand must be a plain tensor. This is + # what tells the TP layer to hand local shards rather than DTensors + # (`should_use_local_tensors`); without it a dense linear sees DTensors whenever + # grad is enabled, since its inference fast path only covers `not is_grad_enabled`. + new_module._hf_quantized_needs_local_tp = True model.set_submodule(module_name, new_module) has_been_replaced = True if not has_been_replaced: logger.warning( - "You are loading your model using fp8 but no linear modules were found in your model." - " Please double check your model architecture." + "You are loading a fine-grained quantized model but no linear or experts modules were " + "found to convert. Please double check your model architecture." ) return model @@ -1072,14 +1058,6 @@ def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) -def _checkpoint_scale_container(model) -> torch.dtype | None: - """The scale container the checkpoint shipped, as recorded on its finegrained modules at load - (one checkpoint, one container); None when it was already the held dtype.""" - if model is None: - return None - return next((m.scale_container_dtype for m in model.modules() if getattr(m, "scale_container_dtype", None)), None) - - class FineGrainedInterleaveGateUp(ConversionOps): """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order (weight, scale grid and bias share the row axis) — the core ``Interleave`` along dim 1. Runs on load and @@ -1104,12 +1082,10 @@ def reverse_op(self) -> ConversionOps: class FineGrainedScaleContainer(ConversionOps): - """A block scale into the dtype its module holds. Two containers exist for UE8M0: the same - exponent byte under ``uint8`` (MiniMax-style mxfp8 checkpoints) — a bit reinterpretation into - ``float8_e8m0fnu`` — and the power-of-two value in a float32 tensor (dsv4-flash-base) — an exact - numeric cast. Other tensors, and scales already in the held dtype, pass through. The module - records the checkpoint's container (``scale_container_dtype``), so the reverse restores it - exactly on save.""" + """A block scale into the dtype its module holds. UE8M0 ships in two containers: the exponent + byte under ``uint8`` (MiniMax mxfp8), reinterpreted into ``float8_e8m0fnu``, and its + power-of-two value in float32 (dsv4-flash-base), cast exactly. Anything else passes through. + The module records which one the checkpoint used, so the reverse restores it on save.""" def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer @@ -1120,8 +1096,10 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= for key, value in _keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value if self.inverse: - container = _checkpoint_scale_container(model) - # e8m0 is only ever a scale, so the dtype alone picks the tensors to restore + # one checkpoint, one container, recorded on the modules at load; e8m0 is only + # ever a scale, so the dtype alone picks the tensors to restore + modules = model.modules() if model is not None else () + container = next((c for m in modules if (c := getattr(m, "scale_container_dtype", None))), None) if container is not None and value.dtype == _get_ue8m0_dtype(): value = _recontain(value, container) else: @@ -1138,11 +1116,10 @@ def reverse_op(self) -> ConversionOps: class FineGrainedSwizzleScales(ConversionOps): - """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D scale - Parameter, see ``_swizzles_scales``): one triton launch per rank-local shard, values unchanged. - Every other tensor the converter carries (weights, biases, the scales of a module kept affine) - passes through. The reverse restores the checkpoint's affine grid from a 5-D tensor's own shape - (whole 128-row blocks and 4-column groups, so no padding to trim), so saving is exact.""" + """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D Parameter, + see ``_swizzles_scales``): one triton launch per rank-local shard, values unchanged. Anything + else the converter carries passes through. The swizzle only ever covers whole 128-row blocks + and 4-column groups, so the reverse reads the affine grid back off the 5-D shape.""" def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer @@ -1178,9 +1155,8 @@ def reverse_op(self) -> ConversionOps: class FineGrainedPackedBlocks(ConversionOps): """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 — two low-nibble-first E2M1 values per byte, 16 bytes per group of 32 — as the packed ``(E, N, K/2)`` int8 the kernels read: the same - bytes, grouped per row. The reverse folds the rows back into 16-byte groups under uint8. The - ``{proj}_scales`` companion goes through ``FineGrainedScaleContainer`` (exponent bytes -> - ``float8_e8m0fnu``) in its own converter.""" + bytes, grouped per row, folded back into 16-byte groups on the way out. Its ``{proj}_scales`` + companion rides ``FineGrainedScaleContainer`` in its own converter.""" def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer @@ -1204,10 +1180,9 @@ def reverse_op(self) -> ConversionOps: class FineGrainedViewPackedInt8(ConversionOps): - """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules - store (same bytes — ``copy_`` into an int8 param would numerically CONVERT and corrupt - values >= 128). Passes non-uint8 tensors through untouched, so it can ride converters - whose pattern also matches unquantized modules.""" + """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store: + ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 + tensors pass through, so this can ride converters that also match unquantized modules.""" def __init__(self, hf_quantizer=None): self.hf_quantizer = hf_quantizer @@ -1221,27 +1196,20 @@ def convert(self, input_dict, target_patterns=None, **kwargs): class FineGrainedWeightGlobals(ConversionOps): - """modelopt's second-level globals, and the calibrated ``input_scale`` they have to travel - with, as the kernels index them: one fp32 global per expert per projection. - - A gate|up stack arrives with TWO per expert, because modelopt calibrates the two projections - separately — they differ by up to 1.6x per expert on Muse-Spark, and agree on GLM-5.2. The - stack keeps the gate's and the up half's folds onto the DOWN projection, which is why this one - op owns every global of a layer: ``inter = silu(gate) * up`` is linear in the up half, so - running the stack on the gate's global alone scales the expert's output by ``1 / ratio``, and - the down's weight global scales it back. The down's calibrated input global moves the other - way, because it is what the fused epilogue normalizes the intermediate by before requantizing - it: without that the requant drifts off the calibrated range and a large enough ratio pushes - the e4m3 block scales past their 448 ceiling. - - Rescaling the up half's e4m3 block scales instead — the other way to merge — would re-round - them against codes already chosen for the old global, moving every value in a block by up to - 6%. - - Sources are either the per-expert keys (``experts.*.gate_proj.weight_scale_2``, merged into a - list per expert) or one stacked ``(E, 2)`` tensor per layer (the fused vLLM layout).""" - - supports_many_to_many = True + """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: + one fp32 global per expert per projection, from either the per-expert checkpoint keys + (``experts.*.gate_proj.weight_scale_2``) or one stacked ``(E, 2)`` tensor per layer. + + A gate|up stack arrives with two per expert, calibrated separately (up to 1.6x apart on + Muse-Spark, equal on GLM-5.2). Merging them is not local, which is why one op owns a layer's + whole set: the stack keeps the gate's, and the up half's leaves as a ratio. That is sound + because ``inter = silu(gate) * up`` is LINEAR in the up half — running the stack on the gate's + global alone scales the expert output by ``1 / ratio``, which the down projection's weight + global takes back. The down's calibrated input global moves the other way, since it is what + the fused epilogue normalizes the intermediate by before requantizing: without it the requant + drifts off the calibrated range and a large enough ratio pushes the e4m3 block scales past + their 448 ceiling. Rescaling the up half's block scales instead would re-round them against + codes chosen for the old global, moving every value in a block by up to 6%.""" def __init__(self, hf_quantizer=None): self.hf_quantizer = hf_quantizer @@ -1365,20 +1333,36 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key module = holder[0] if holder is not None else None held_scale = getattr(module, holder[1]) if holder is not None else None - fmt = _weight_formats()[module.weight_format] if module is not None else None - if fmt is None or fmt.scale_group is None: - block = tuple(module.block_size) if module is not None and module.block_size else self._config_block(value) + format_spec = weight_formats()[module.weight_format] if module is not None else None + if format_spec is None or format_spec.scale_group is None: + if module is not None and module.block_size: + block = tuple(module.block_size) + else: + # no module (direct invocation): take the block from the config, and failing that + # put one scale over the whole matrix + config = self.hf_quantizer.quantization_config if self.hf_quantizer is not None else None + config_block = ( + config.get("weight_block_size") + if isinstance(config, dict) + else getattr(config, "weight_block_size", None) + ) + block = tuple(config_block) if config_block else (value.shape[-2], value.shape[-1]) ue8m0 = ( held_scale.dtype == _get_ue8m0_dtype() if held_scale is not None else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" ) - quantized = _quantize_block_fp8(value, block, ue8m0) + quantized = self._quantize_block_fp8(value, block, ue8m0) return {key: value} if quantized is None else {key: quantized[0], f"{prefix}_scale_inv": quantized[1]} - if value.device.type != "cuda": - raise ValueError(f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on CUDA") - weight, scale, global_scale = _quantize_group(value, fmt) + if value.device.type not in ("cuda", "xpu"): + # the kernels' quantizers are triton launches; the package builds for cuda, rocm + # (which torch reports as "cuda") and xpu + raise ValueError( + f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on an " + f"accelerator, but the weight is on {value.device.type}" + ) + weight, scale, global_scale = self._quantize_group(value, format_spec) scale = _recontain(scale, held_scale.dtype) if held_scale.ndim == 5: scale = load_finegrained_kernel().swizzle_mx_scales(scale) @@ -1392,83 +1376,67 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc out[f"{prefix}_input_global_scale"] = torch.ones_like(held_input) return out - def _config_block(self, value: torch.Tensor) -> tuple[int, int]: - config = self.hf_quantizer.quantization_config - block = ( - config.get("weight_block_size") if isinstance(config, dict) else getattr(config, "weight_block_size", None) + @staticmethod + def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): + """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or + ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two + before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" + block_m, block_n = block + rows, cols = value.shape[-2], value.shape[-1] + if rows % block_m or cols % block_n: + return None + tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) + max_abs = tiles.abs().amax(dim=(-3, -1)) + inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) + if ue8m0: + inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) + scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) + return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) + + @staticmethod + def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): + """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the + kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has + one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the + smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" + kernel = load_finegrained_kernel() + rows = value.reshape(-1, value.shape[-1]) + global_scale = None + if format_spec.global_scale_dtype is not None: + per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) + global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) + rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) + packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) + global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale + else: + quantize = kernel.mxfp4_act_quant if format_spec.values_per_byte == 2 else kernel.mxfp8_act_quant + packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) + return ( + packed.view(format_spec.weight_dtype).reshape(*value.shape[:-1], -1), + scale.reshape(*value.shape[:-1], -1), + global_scale, ) - return tuple(block) if block is not None else (value.shape[-2], value.shape[-1]) @property def reverse_op(self) -> ConversionOps: return FineGrainedDequantize(self.hf_quantizer) -def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): - """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or - ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two - before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" - block_m, block_n = block - rows, cols = value.shape[-2], value.shape[-1] - if rows % block_m or cols % block_n: - return None - tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) - max_abs = tiles.abs().amax(dim=(-3, -1)) - inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) - if ue8m0: - inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) - scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) - quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) - return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) - - -def _quantize_group(value: torch.Tensor, fmt: WeightFormat): - """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the - kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has - one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the - smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" - kernel = load_finegrained_kernel() - rows = value.reshape(-1, value.shape[-1]) - global_scale = None - if fmt.has_global_scale: - per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) - global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) - rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) - packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) - global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale - else: - quantize = kernel.mxfp4_act_quant if fmt.values_per_byte == 2 else kernel.mxfp8_act_quant - packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) - return ( - packed.view(fmt.weight_dtype).reshape(*value.shape[:-1], -1), - scale.reshape(*value.shape[:-1], -1), - global_scale, - ) - - class FineGrainedDequantize(ConversionOps): - """Dequantize FP8 weights using their per-block ``weight_scale_inv``. - - Designed to run as the *first* op in any :class:`WeightConverter` chain when - loading with ``dequantize=True`` — :meth:`update_weight_conversions` on the - FP8 quantizer attaches it to each existing model-specific converter so that - per-expert (weight, scale) pairs are folded into full-precision tensors before - the chain's merge / concat ops collapse the per-expert structure. - - Pattern semantics - Input ``input_dict`` carries one entry per source pattern; each value is a - list of tensors (one per ``*`` match). For every weight pattern that has a - sibling ``*.weight_scale_inv`` pattern in the dict, this op pairs them up by - index, dequantizes per-pair, and emits the dequantized list under the - original *weight* key. Scale entries are dropped from the output so the - remaining ops only see weights. + """A quantized weight folded back to full precision against its per-block scale grid. + + Runs FIRST in its converter under ``dequantize=True``, which is why + :meth:`update_weight_conversions` attaches it to each of the model's own converters: the + merge / concat ops after it collapse the per-expert structure the (weight, scale) pairing + needs. It pairs each weight pattern with its sibling scale pattern by index and emits the + result under the weight key, dropping the scales so the rest of the chain sees weights only. """ def __init__(self, hf_quantizer): self.hf_quantizer = hf_quantizer def _scale_pattern_for(self, weight_pattern: str) -> str: - # Strip the optional ``$`` regex anchor so we can match the underlying name. anchored = weight_pattern.endswith("$") base = weight_pattern[:-1] if anchored else weight_pattern if base.endswith("_blocks"): @@ -1482,9 +1450,7 @@ def _scale_pattern_for(self, weight_pattern: str) -> str: scale = base + "_scale_inv" return scale + "$" if anchored else scale - # E2M1 (FP4) value table — checkpoints sometimes ship MoE experts as packed FP4 - # (two e2m1 nibbles per int8 byte), so the "weight" dtype lands as ``int8`` / - # ``float4_e2m1fn_x2`` and we have to unpack before applying the scale grid. + # packed-FP4 experts arrive as int8 / float4_e2m1fn_x2, two e2m1 nibbles per byte _FP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: @@ -1499,17 +1465,15 @@ def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: def _dequantize_one( self, quantized: torch.Tensor, scales: torch.Tensor, output_dtype: torch.dtype | None = None ) -> torch.Tensor: - # FP4 path: int8 / float4_e2m1fn_x2 stores two nibbles per byte. Unpack to fp32 - # first so the rest of the routine sees a normal (rows, cols) float matrix. + # unpacked first, so the rest of the routine sees a normal (rows, cols) float matrix fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) if quantized.dtype == torch.int8 or (fp4_dtype is not None and quantized.dtype == fp4_dtype): quantized_fp32 = self._unpack_fp4(quantized) else: quantized_fp32 = quantized.to(torch.float32) rows, cols = quantized_fp32.shape[-2:] - # Derive block size from the scale grid rather than the global config: MoE experts - # ship MXFP4 with a ``[1, 32]`` block, dense linears ship FP8 with ``[128, 128]``, - # and the same dequant has to handle both within one checkpoint. + # the block comes from the scale grid, not the config: one checkpoint holds both a + # ``[1, 32]`` MXFP4 expert block and a ``[128, 128]`` FP8 dense one try: scale_rows, scale_cols = scales.shape[-2:] except Exception: @@ -1521,17 +1485,13 @@ def _dequantize_one( ) block_m = rows // scale_rows block_n = cols // scale_cols - # ``ue8m0`` (``float8_e8m0fnu``) scales have no CUDA ``mul`` kernel, and casting - # the FP8 weight to that dtype loses precision. Promote both sides to fp32 for - # the math; prefer the destination parameter's dtype when known so eager modules - # (e.g. plain ``nn.Linear``) keep the model's compute dtype after load. + # the math runs in fp32 either way (``float8_e8m0fnu`` has no ``mul`` kernel); the + # destination parameter's dtype wins when known, so an eager module keeps the model's if output_dtype is None: output_dtype = ( scales.dtype if scales.dtype.is_floating_point and scales.element_size() >= 2 else torch.bfloat16 ) - # MXFP8 checkpoints ship E8M0 exponents stored as ``torch.uint8`` (one byte per - # block) — the actual scale is `2 ** (byte - 127)`. Interpreting the raw bytes - # as scalar multipliers would be silently wrong, so unpack to fp32 here. + # an MXFP8 checkpoint's uint8 scale is a biased E8M0 exponent, not a multiplier if scales.dtype == torch.uint8: s_fp32 = (scales.to(torch.float32) - 127.0).exp2() else: @@ -1578,7 +1538,6 @@ def convert( continue # consumed by the dequant; drop from the chain scale_key = self._scale_pattern_for(key) if scale_key not in input_dict: - # No scale to apply (e.g. unrelated entry) — pass through untouched. result[key] = value continue weights = value if isinstance(value, list) else [value] @@ -1594,7 +1553,6 @@ def convert( @property def reverse_op(self) -> ConversionOps: - # Round-trip: dequantize on load -> re-quantize on save, so the saved - # checkpoint preserves the FP8 format (weight + per-block ``weight_scale_inv``) - # whether the in-memory state stayed quantized or was dequantized for compute. + # a save re-quantizes, so the checkpoint keeps its format whether the in-memory + # state stayed quantized or was dequantized for compute return FineGrainedQuantize(self.hf_quantizer) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index e8337d9dff19..ac1dafc83808 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -15,7 +15,6 @@ import functools import os -import warnings from collections.abc import Callable from dataclasses import dataclass @@ -39,16 +38,14 @@ from .moe import ExpertsInterface, use_experts_implementation -warnings.warn( - "finegrained_fp8 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - logger = logging.get_logger(__name__) +logger.warning( + "finegrained_fp8 is frozen for backward compatibility and no longer receives new recipes; the " + "fine-grained quantization machinery lives in transformers.integrations.finegrained " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." +) + _FP8_DTYPE = torch.float8_e4m3fn _FP8_MIN = torch.finfo(_FP8_DTYPE).min diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 5ada29dc4831..f29c33d4999f 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -730,7 +730,7 @@ def register_kernel_mapping_transformers(*args, **kwargs): "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); # exports the same three matmuls plus the fused-MoE forwards and per-recipe quant helpers - "finegrained-kernels": {"repo_id": "kernels-staging/finegrained-kernels", "revision": "pr-1018"}, + "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 0}, # dense NVFP4 GEMM behind `NVFP4Linear`; MoE experts route to finegrained-kernels instead "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, } diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index fcd3d8d39ae5..9f23817d91df 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -165,9 +165,9 @@ def batched_mm_experts_forward( ) # (S, hidden_dim) # Normalize each expert application, where such a norm is defined, before it is weighted. - # `getattr`: an experts module built without the decorator (or predating the norm) has no such - # attribute, and this forward is the shared one every backend adapts onto. - if getattr(self, "post_expert_norm_name", None) is not None: + # `getattr`: an experts module built without the decorator has no such attribute, and this + # forward is the shared one every backend adapts onto. + if getattr(self, "has_post_expert_norm", False): proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) # Apply routing weights @@ -482,9 +482,9 @@ def grouped_mm_experts_forward( proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Normalize each expert application, where such a norm is defined, before it is weighted. - # `getattr`: an experts module built without the decorator (or predating the norm) has no such - # attribute, and this forward is the shared one every backend adapts onto. - if getattr(self, "post_expert_norm_name", None) is not None: + # `getattr`: an experts module built without the decorator has no such attribute, and this + # forward is the shared one every backend adapts onto. + if getattr(self, "has_post_expert_norm", False): proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) # Apply routing weights @@ -554,7 +554,7 @@ def use_experts_implementation( is_transposed: bool = False, has_bias: bool = False, has_gate: bool = True, - post_expert_norm: str | None = None, + has_post_expert_norm: bool = False, ) -> type[torch.nn.Module]: """Decorator to modify experts class to support different experts implementations. @@ -573,12 +573,10 @@ def use_experts_implementation( has_gate (`bool`, *optional*, defaults to `True`): Whether the experts use a gating mechanism or not. Whether it has gate_up_proj weights or just up_proj weights. - post_expert_norm (`str`, *optional*): - Names the normalization the experts apply to the down output before the routing - weights, when they apply one — the way `hidden_act` names the activation - `_apply_gate` applies. Naming one requires the class to define - `_apply_post_norm(self, expert_out)`, which is where its own math lives; a backend - that implements the named form can fuse it instead of calling that. + has_post_expert_norm (`bool`, *optional*, defaults to `False`): + Whether the experts normalize the down output before the routing weights. The class + must then define `_apply_post_norm(self, expert_out)` — the standard name every + backend applies, so none of them needs to know the model's own math. Returns: `type[torch.nn.Module]`: The modified experts class. @@ -596,7 +594,7 @@ def __init__(self, config, *args, **kwargs): self.has_bias = has_bias self.is_transposed = is_transposed self.is_concatenated = is_concatenated - self.post_expert_norm_name = post_expert_norm + self.has_post_expert_norm = has_post_expert_norm self._is_expert_parallel = False @wraps(original_forward) @@ -607,9 +605,9 @@ def forward(self, *args, **kwargs): if not hasattr(experts_class, "_apply_gate"): experts_class._apply_gate = _default_apply_gate - if post_expert_norm is not None and not hasattr(experts_class, "_apply_post_norm"): + if has_post_expert_norm and not hasattr(experts_class, "_apply_post_norm"): raise TypeError( - f"{experts_class.__name__} names a post-expert norm ({post_expert_norm!r}) but does not " + f"{experts_class.__name__} declares a post-expert norm but does not " "define `_apply_post_norm(self, expert_out)`, which applies it to one expert " "application's rows before the routing weights." ) diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 896782d7e30c..67c74eda798c 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -64,14 +64,20 @@ def shard_and_distribute_module(*args, **kwargs): def to_local(t): - """Unwrap a `DTensor` to its local shard if needed; pass through otherwise. + """Deprecated: unwrap a `DTensor` to its local shard, or pass through. - Custom kernels (CUTLASS, CuteDSL, Triton) take raw tensor pointers and don't - understand `DTensor`, so weights wrapped by FSDP2 / EP need this unwrap before - they can be fed to the kernel. ``to_local()`` is autograd-aware on the train - path: backward rewraps the gradient as a DTensor matching each parameter's - placements. + Kept for callers that predate `_hf_quantized_needs_local_tp`. A quantized module that sets + it is handed local tensors by the TP loader, so the kernels — which take raw pointers and do + not understand `DTensor` — need no unwrap of their own; nothing in transformers calls this. + `DTensor.to_local()` is the direct replacement and is autograd-aware on the train path, + rewrapping the gradient to match each parameter's placements. """ + warnings.warn( + "`to_local` is deprecated. Set `_hf_quantized_needs_local_tp = True` on the module so the " + "tensor-parallel loader hands it local tensors, or call `DTensor.to_local()` directly.", + FutureWarning, + stacklevel=2, + ) from ..distributed.utils import is_dtensor if is_dtensor(t): diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index 47c029a5bca9..94e01398159d 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -22,6 +22,13 @@ @auto_docstring(checkpoint="openai/gpt-oss-20b") @strict class GptOssConfig(PreTrainedConfig): + r""" + swiglu_alpha (`float`, *optional*, defaults to 1.702): + Sigmoid gain of the clamped SwiGLU the experts apply. + swiglu_limit (`float`, *optional*, defaults to 7.0): + Clamp applied to the gate (above) and the up projection (both sides) before that SwiGLU. + """ + model_type = "gpt_oss" attribute_map = { "num_experts": "num_local_experts", @@ -58,6 +65,8 @@ class GptOssConfig(PreTrainedConfig): rope_parameters: dict | None = None attention_dropout: float | int = 0.0 num_experts_per_tok: int = 4 + swiglu_alpha: float = 1.702 + swiglu_limit: float = 7.0 router_aux_loss_coef: float = 0.001 output_router_logits: bool = False use_cache: bool = True diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 495b5c2182a1..944728a18251 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -76,8 +76,8 @@ def __init__(self, config): self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_size)) self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.intermediate_size, self.hidden_size))) self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) - self.alpha = 1.702 - self.limit = 7.0 + self.alpha = config.swiglu_alpha + self.limit = config.swiglu_limit def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: gate, up = gate_up[..., ::2], gate_up[..., 1::2] diff --git a/src/transformers/models/gpt_oss/modular_gpt_oss.py b/src/transformers/models/gpt_oss/modular_gpt_oss.py index 7c3357a393ad..ec39b9bb8a08 100644 --- a/src/transformers/models/gpt_oss/modular_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modular_gpt_oss.py @@ -73,8 +73,8 @@ def __init__(self, config): self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_size)) self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.intermediate_size, self.hidden_size))) self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) - self.alpha = 1.702 - self.limit = 7.0 + self.alpha = config.swiglu_alpha + self.limit = config.swiglu_limit def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: gate, up = gate_up[..., ::2], gate_up[..., 1::2] diff --git a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py index e7aaefde4bca..78498a2f7b4f 100644 --- a/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/configuration_openai_privacy_filter.py @@ -46,6 +46,13 @@ @auto_docstring(checkpoint="openai/privacy-filter") @strict class OpenAIPrivacyFilterConfig(PreTrainedConfig): + r""" + swiglu_alpha (`float`, *optional*, defaults to 1.702): + Sigmoid gain of the clamped SwiGLU the experts apply. + swiglu_limit (`float`, *optional*, defaults to 7.0): + Clamp applied to the gate (above) and the up projection (both sides) before that SwiGLU. + """ + model_type = "openai_privacy_filter" attribute_map = { "num_experts": "num_local_experts", @@ -80,6 +87,8 @@ class OpenAIPrivacyFilterConfig(PreTrainedConfig): rope_parameters: dict | None = None attention_dropout: float | int = 0.0 num_experts_per_tok: int = 4 + swiglu_alpha: float = 1.702 + swiglu_limit: float = 7.0 router_aux_loss_coef: float = 0.001 output_router_logits: bool = False use_cache: bool = True diff --git a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py index 426df7bd2a03..1b437886ec90 100644 --- a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py @@ -261,8 +261,8 @@ def __init__(self, config): self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_size)) self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.intermediate_size, self.hidden_size))) self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) - self.alpha = 1.702 - self.limit = 7.0 + self.alpha = config.swiglu_alpha + self.limit = config.swiglu_limit def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: # Concatenated layout instead of interleaving diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 72925356ee7d..545f8550abdc 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -14,7 +14,13 @@ import re from typing import TYPE_CHECKING -from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging +from ..utils import ( + is_accelerate_available, + is_torch_available, + is_torch_distributed_available, + is_torch_xpu_available, + logging, +) from .base import HfQuantizer from .quantizers_utils import get_module_from_name @@ -47,7 +53,7 @@ def _assert_dequantize_supported(self) -> None: """The dequantize chain folds a weight into full precision with its per-block scale alone. NVFP4 carries a second level (modelopt's `weight_scale_2`) that the chain has no slot for, so refuse rather than hand back a weight scaled by `1 / global`.""" - if self._quant_method() in ("nvfp4", "modelopt"): + if self._quant_method() == "nvfp4": raise NotImplementedError( f"`dequantize=True` is not supported for {self._quant_method()!r} checkpoints: their " "two-level scales cannot be folded into a full-precision weight by this path. Load " @@ -56,7 +62,7 @@ def _assert_dequantize_supported(self) -> None: def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): - raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") + raise ImportError("Loading a fine-grained quantized model requires accelerate (`pip install accelerate`)") if self.quantization_config.dequantize: self._assert_dequantize_supported() @@ -66,12 +72,13 @@ def validate_environment(self, *args, **kwargs): if self.pre_quantized: self._assert_dequantize_supported() logger.warning_once( - "Using FP8 quantized models requires a GPU or XPU, we will default to dequantizing the model to bf16 since no GPU or XPU is available" + "Fine-grained quantized models require a GPU or XPU; dequantizing to bf16 since neither is " + "available." ) self.quantization_config.dequantize = True return else: - raise RuntimeError("No GPU or XPU found. A GPU or XPU is needed for FP8 quantization.") + raise RuntimeError("No GPU or XPU found. A GPU or XPU is needed for fine-grained quantization.") if torch.cuda.is_available(): compute_capability = torch.cuda.get_device_capability() @@ -79,9 +86,9 @@ def validate_environment(self, *args, **kwargs): if (major < 8) or (major == 8 and minor < 9): self._assert_dequantize_supported() logger.warning_once( - "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)" - f", actual = `{major}.{minor}`. We will default to dequantizing the model to bf16. Feel free " - f"to use a different quantization method like bitsandbytes or torchao" + "Fine-grained quantized models are only supported on GPUs with compute capability >= 8.9 " + f"(e.g 4090/H100), actual = `{major}.{minor}`. We will default to dequantizing the model to " + "bf16. Feel free to use a different quantization method like bitsandbytes or torchao" ) self.quantization_config.dequantize = True return @@ -89,9 +96,9 @@ def validate_environment(self, *args, **kwargs): device_map = kwargs.get("device_map") if device_map is None: logger.warning_once( - "You have loaded an FP8 model on CPU and have a CUDA or XPU device available, make sure to set " - "your model on a GPU or XPU device in order to run your model. To remove this warning, " - "pass device_map = 'cuda' or 'xpu'. " + "You have loaded a fine-grained quantized model on CPU and have a CUDA or XPU device available, " + "make sure to set your model on a GPU or XPU device in order to run your model. To remove this " + "warning, pass device_map = 'cuda' or 'xpu'." ) elif isinstance(device_map, dict): if ( @@ -100,13 +107,14 @@ def validate_environment(self, *args, **kwargs): and ("cpu" in device_map.values() or "disk" in device_map.values()) ): raise ValueError( - "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device." + "You are attempting to load a fine-grained quantized model with a device_map that contains " + "a cpu/disk device." "This is not supported when the model is quantized on the fly. " "Please use a quantized checkpoint or remove the cpu/disk device from the device_map." ) def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: - # `FP8GroupedLinear` is a subclass of `FineGrainedLinear`, so the tuple covers it implicitly. + # `FineGrainedGroupedLinear` subclasses `FineGrainedLinear`, so the tuple covers it implicitly. from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear module, tensor_name = get_module_from_name(model, param_name) @@ -157,6 +165,13 @@ def _process_model_before_weight_loading( # GPT-OSS MXFP4 runs weight-only (W4A16): raw bf16 activations against packed # weights. The kernels' weight-native default would quantize activations to mxfp4. self.quantization_config.activation_format = "bf16" + if getattr(self.quantization_config, "activation_format", None) == "bf16": + # Weight-only holds no activation global, so a calibrated checkpoint's `input_scale` + # has no slot to load into. Not an unexpected key — one this run has no use for, which + # a later W4A4 run of the same checkpoint reads. + model._keys_to_ignore_on_load_unexpected = set(model._keys_to_ignore_on_load_unexpected or []) | { + r"input_scale$" + } self.modules_to_not_convert = self.get_modules_to_not_convert( model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules ) @@ -168,13 +183,6 @@ def _process_model_before_weight_loading( replace_with_finegrained_layer( model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config ) - if getattr(self.quantization_config, "activation_format", None) == "bf16": - # Weight-only: the experts hold no activation global, so a calibrated checkpoint's - # `input_scale` has no slot to load into. It is not an unexpected key, it is one this - # run has no use for — a later W4A4 run of the same checkpoint reads it. - model._keys_to_ignore_on_load_unexpected = set(model._keys_to_ignore_on_load_unexpected or []) | { - r"input_scale$" - } def _process_model_after_weight_loading(self, model, **kwargs): from ..integrations.finegrained import disable_deepgemm_on_multi_device @@ -182,6 +190,23 @@ def _process_model_after_weight_loading(self, model, **kwargs): disable_deepgemm_on_multi_device(model) return model + @staticmethod + def _expert_shard(dim: int, packed: bool = False) -> str: + """Register a param-only style that shards an expert tensor on `dim`, and return the name + a plan carries it under. A style name is the only way a plan can express a placement, and + the built-in styles locate their axis by counting from the END of the shape — which lands + inside a swizzled scale grid's inner tile rather than on its row blocks, and on the expert + axis of a bias.""" + from torch.distributed.tensor import Shard + from torch.distributed.tensor.placement_types import _StridedShard + + from ..distributed.tensor_parallel import ALL_PARALLEL_STYLES, MoEParamShard + + placement = _StridedShard(dim=dim, split_factor=2) if packed else Shard(dim) + name = f"moe_experts_{'packed_' if packed else ''}shard{dim}" + ALL_PARALLEL_STYLES.register(name, MoEParamShard(placement)) + return name + def update_tp_plan(self, config): if "Qwen3" in config.__class__.__name__: text_plan = { @@ -203,20 +228,65 @@ def update_tp_plan(self, config): config.base_model_tp_plan = text_plan - # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes - # on top of any plan written above (e.g. the Qwen3 dense plan). Models carry the - # experts mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite - # both. See `FineGrainedExperts._impl_tp_layer_overrides`. + # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes on + # top of any plan written above (e.g. the Qwen3 dense plan). Models carry the experts + # mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite both. from ..integrations.finegrained import FineGrainedExperts impl = getattr(config, "_experts_implementation", None) - layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl) - if layer_overrides: - for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): - base_plan = getattr(config, plan_attr, None) or {} - updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} - if updated_plan != base_plan: - setattr(config, plan_attr, updated_plan) + interleaved_gate_up = impl != "deepgemm_megamoe" + layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) + for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): + base_plan = getattr(config, plan_attr, None) or {} + updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} + + # Every companion beside a projection weight — block scales, bias, NVFP4 globals, a + # static activation scale — is expert-indexed too, so each shards with the weight. The + # matcher keys on the exact parameter name and falls back only to the owning MODULE, + # whose entry shards nothing, so without these the weight is this rank's expert slice + # while its scales are every rank's. gate_up's input global is per-tensor: replicated. + for key, style in list(updated_plan.items()): + # only the experts' own projections: a dense `self_attn.q_proj` is `colwise` too + projection = key.rpartition(".")[2] + if projection not in ("gate_up_proj", "up_proj", "down_proj"): + continue + if style == "grouped_gemm": + companions = ["_scale_inv", "_bias", "_weight_global_scale", "_activation_scale"] + if projection == "down_proj": + companions.append("_input_global_scale") + for suffix in companions: + updated_plan.setdefault(f"{key}{suffix}", style) + continue + + # ...and only a STACKED experts module: `mlp.up_proj`, and the dense + # `mlp.shared_experts.up_proj`, end in the same word but are plain 2-D linears + # that colwise/rowwise already shard correctly. The model's own entry for the + # owning module is what distinguishes them. + if updated_plan.get(key.rpartition(".")[0]) not in ("moe_tp_experts", "megamoe_experts"): + continue + if not is_torch_distributed_available(): + continue # the styles below need a distributed build; without one, no TP + + # Intra-expert TP: the experts stay whole and the projection's own axis splits, so + # the per-expert globals, activation scales and the down bias (added after the + # row-reduce) stay replicated — but the scale grid still has to follow its weight. + # dim 1 is the output axis and dim 2 the input axis in BOTH the affine and the + # 5-D swizzled scale container, which is why naming those two suffices. + if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): + updated_plan.setdefault(f"{key}_scale_inv", self._expert_shard(2)) + elif projection != "down_proj" and style in ("packed_colwise", "colwise"): + # `packed_colwise` splits the output axis as two packed halves — the + # `[gate; up]` STACK. A quantized module holds the kernels' INTERLEAVED rows, + # where that split hands a rank gate|up pairs for channels its own down + # projection does not serve, so the weight's own style changes with it. + rows = self._expert_shard(1, packed=not interleaved_gate_up) + if interleaved_gate_up: + updated_plan[key] = rows + updated_plan.setdefault(f"{key}_scale_inv", rows) + updated_plan.setdefault(f"{key}_bias", rows) + + if updated_plan != base_plan: + setattr(config, plan_attr, updated_plan) return config @@ -241,42 +311,27 @@ def _quant_method(self) -> str: return getattr(method, "value", method) def get_weight_conversions(self): - from ..core_model_loading import WeightConverter - from ..integrations.finegrained import ( - FineGrainedDequantize, - FineGrainedInterleaveGateUp, - FineGrainedPackedBlocks, - FineGrainedScaleContainer, - FineGrainedSwizzleScales, - ) - - # {proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed - # weight, the exponent bytes take the scale container op; then the usual layout ops. Bare - # patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`. - if self.pre_quantized and self._quant_method() == "mxfp4" and not self.quantization_config.dequantize: - converters = [] - for proj in ("gate_up_proj", "down_proj"): - interleave = [FineGrainedInterleaveGateUp(self)] if proj == "gate_up_proj" else [] - converters += [ - WeightConverter( - source_patterns=rf"{proj}_blocks$", - target_patterns=proj, - operations=[FineGrainedPackedBlocks(self), *interleave, FineGrainedSwizzleScales(self)], - ), - WeightConverter( - source_patterns=rf"{proj}_scales$", - target_patterns=f"{proj}_scale_inv", - operations=[*interleave, FineGrainedScaleContainer(self), FineGrainedSwizzleScales(self)], - ), - ] - return converters + """The converters this CHECKPOINT's own key layout needs, on top of the model's plan.""" + if not self.pre_quantized: + return [] + if self.quantization_config.dequantize: + return self._dequantize_conversions() + if self._quant_method() == "mxfp4": + return self._mxfp4_conversions() + if self._quant_method() == "nvfp4": + return self._nvfp4_conversions() + return [] - # The same GPT-OSS keys under `dequantize=True`: the blocks regroup into the packed rows - # and the exponent-byte scales come back out of them as bf16, in the (E, hidden, 2I) - # orientation the unquantized experts hold. - if self.pre_quantized and self._quant_method() == "mxfp4" and self.quantization_config.dequantize: - from ..core_model_loading import Transpose + def _dequantize_conversions(self): + """Every quantized key folded back to a full-precision weight. One converter per key + layout: the packed `{proj}_blocks` + `{proj}_scales` pair GPT-OSS ships, and the + `weight` + `weight_scale_inv` pair everything else does.""" + from ..core_model_loading import Transpose, WeightConverter + from ..integrations.finegrained import FineGrainedDequantize, FineGrainedPackedBlocks + if self._quant_method() == "mxfp4": + # the blocks regroup into the packed rows and the exponent-byte scales come back out + # of them as bf16, in the (E, hidden, 2I) orientation the unquantized experts hold return [ WeightConverter( source_patterns=[rf"{proj}_blocks$", rf"{proj}_scales$"], @@ -286,132 +341,159 @@ def get_weight_conversions(self): for proj in ("gate_up_proj", "down_proj") ] - # modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts - # a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already - # stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and - # the other never fires. Routed experts only — a generic `weight_scale*` rename would run - # before converter collection and mangle these keys. - if self.pre_quantized and self._quant_method() == "modelopt" and not self.quantization_config.dequantize: - # A weight-only run keeps activations bf16, so its experts hold no activation global - # and the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys - # to the unexpected-key filter below rather than converting them onto a module slot - # that does not exist. - calibrated = getattr(self.quantization_config, "activation_format", None) != "bf16" - from ..core_model_loading import Concatenate, MergeModulelist - from ..integrations.finegrained import ( - FineGrainedInputGlobals, - FineGrainedScaleContainer, - FineGrainedViewPackedInt8, - FineGrainedWeightGlobals, + # anchored `weight$` so the scale keys land in the scale slots; the activation + # scales are collected only to be dropped + return [ + WeightConverter( + source_patterns=["weight$", "weight_scale_inv", "activation_scale"], + target_patterns="weight", + operations=[FineGrainedDequantize(self)], ) + ] + + def _mxfp4_conversions(self): + """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed + weight, the exponent bytes take the scale container op; then the usual layout ops. Bare + patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained import ( + FineGrainedInterleaveGateUp, + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) - # MergeModulelist is what stamps each source with its expert index, which is what lets - # expert parallelism keep only this rank's experts (see `tensor_idx` in - # core_model_loading). Without it every rank collects all E values and the forward - # asserts on the per-expert count. - per_expert = [ + converters = [] + for proj in ("gate_up_proj", "down_proj"): + interleave = [FineGrainedInterleaveGateUp(self)] if proj == "gate_up_proj" else [] + converters += [ WeightConverter( - source_patterns=[ - "mlp.experts.*.gate_proj.weight_scale$", - "mlp.experts.*.up_proj.weight_scale$", - ], - target_patterns="mlp.experts.gate_up_proj_scale_inv", - operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + source_patterns=rf"{proj}_blocks$", + target_patterns=proj, + operations=[FineGrainedPackedBlocks(self), *interleave, FineGrainedSwizzleScales(self)], ), WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight_scale$", - target_patterns="mlp.experts.down_proj_scale_inv", - operations=[MergeModulelist(dim=0)], + source_patterns=rf"{proj}_scales$", + target_patterns=f"{proj}_scale_inv", + operations=[*interleave, FineGrainedScaleContainer(self), FineGrainedSwizzleScales(self)], ), - # every second-level global of a layer in ONE converter: the gate|up stack's two - # calibrated halves merge into one global per expert by folding the up half's onto - # the down projection, so the weight globals and the down's input scale have to be - # decided together (`FineGrainedWeightGlobals`) - WeightConverter( - source_patterns=[ - "mlp.experts.*.gate_proj.weight_scale_2", - "mlp.experts.*.up_proj.weight_scale_2", - "mlp.experts.*.down_proj.weight_scale_2", - *(["mlp.experts.*.down_proj.input_scale"] if calibrated else []), - ], - target_patterns=[ - "mlp.experts.gate_up_proj_weight_global_scale", - "mlp.experts.down_proj_weight_global_scale", - *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), - ], - operations=[MergeModulelist(dim=0), FineGrainedWeightGlobals(self)], - ), - ] - if calibrated: - per_expert.append( - WeightConverter( - source_patterns=[ - "mlp.experts.*.gate_proj.input_scale", - "mlp.experts.*.up_proj.input_scale", - ], - target_patterns="mlp.experts.gate_up_proj_input_global_scale", - operations=[MergeModulelist(dim=0), FineGrainedInputGlobals(self)], - ) - ) - fused = [ - # the container cast the module's scale dtype needs; `_with_expert_layout_ops` - # adds the interleave and the swizzle on top (it re-adds the cast, which is - # idempotent) - WeightConverter( - source_patterns=rf"experts\.{proj}_weight_scale$", - target_patterns=f"experts.{proj}_scale_inv", - operations=[FineGrainedScaleContainer(self)], - ) - for proj in ("gate_up_proj", "down_proj") ] - fused += [ - # the same one-converter rule as the per-expert layout above, over the keys the - # fused layout ships + return converters + + def _nvfp4_conversions(self): + """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts + a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already + stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and + the other never fires. Routed experts only — a generic `weight_scale*` rename would run + before converter collection and mangle these keys.""" + from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter + from ..integrations.finegrained import ( + FineGrainedInputGlobals, + FineGrainedScaleContainer, + FineGrainedViewPackedInt8, + FineGrainedWeightGlobals, + ) + + # A weight-only run keeps activations bf16, so its experts hold no activation global and + # the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys to the + # unexpected-key filter rather than converting them onto a module slot that does not exist. + calibrated = getattr(self.quantization_config, "activation_format", None) != "bf16" + + # MergeModulelist is what stamps each source with its expert index, which is what lets + # expert parallelism keep only this rank's experts (see `tensor_idx` in + # core_model_loading). Without it every rank collects all E values and the forward + # asserts on the per-expert count. + per_expert = [ + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.weight_scale$", + "mlp.experts.*.up_proj.weight_scale$", + ], + target_patterns="mlp.experts.gate_up_proj_scale_inv", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight_scale$", + target_patterns="mlp.experts.down_proj_scale_inv", + operations=[MergeModulelist(dim=0)], + ), + # every second-level global of a layer in ONE converter: the gate|up stack's two + # calibrated halves merge into one global per expert by folding the up half's onto + # the down projection, so the weight globals and the down's input scale have to be + # decided together (`FineGrainedWeightGlobals`) + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.weight_scale_2", + "mlp.experts.*.up_proj.weight_scale_2", + "mlp.experts.*.down_proj.weight_scale_2", + *(["mlp.experts.*.down_proj.input_scale"] if calibrated else []), + ], + target_patterns=[ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[MergeModulelist(dim=0), FineGrainedWeightGlobals(self)], + ), + ] + if calibrated: + per_expert.append( WeightConverter( source_patterns=[ - r"experts\.gate_up_proj_weight_scale_2$", - r"experts\.down_proj_weight_scale_2$", - *([r"experts\.down_proj_input_scale$"] if calibrated else []), - ], - target_patterns=[ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - *(["experts.down_proj_input_global_scale"] if calibrated else []), + "mlp.experts.*.gate_proj.input_scale", + "mlp.experts.*.up_proj.input_scale", ], - operations=[FineGrainedWeightGlobals(self)], - ), - ] - if calibrated: - fused.append( - WeightConverter( - source_patterns=r"experts\.gate_up_proj_input_scale$", - target_patterns="experts.gate_up_proj_input_global_scale", - operations=[FineGrainedInputGlobals(self)], - ) - ) - # the fused expert WEIGHTS take the packed uint8 -> int8 view the `.weight`-anchored - # rule below gives the per-expert layout; the layout ops ride on top of both - fused += [ - WeightConverter( - source_patterns=rf"experts\.{proj}$", - target_patterns=f"experts.{proj}", - operations=[FineGrainedViewPackedInt8(self)], + target_patterns="mlp.experts.gate_up_proj_input_global_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputGlobals(self)], ) - for proj in ("gate_up_proj", "down_proj") - ] - return per_expert + fused - - if self.pre_quantized and self.quantization_config.dequantize: - return [ - # anchored `weight$` so the scale keys land in the scale slots; the activation - # scales are collected only to be dropped + ) + fused = [ + # the container cast the module's scale dtype needs; `_with_expert_layout_ops` + # adds the interleave and the swizzle on top (it re-adds the cast, which is + # idempotent) + WeightConverter( + source_patterns=rf"experts\.{proj}_weight_scale$", + target_patterns=f"experts.{proj}_scale_inv", + operations=[FineGrainedScaleContainer(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + fused += [ + # the same one-converter rule as the per-expert layout above, over the keys the + # fused layout ships + WeightConverter( + source_patterns=[ + r"experts\.gate_up_proj_weight_scale_2$", + r"experts\.down_proj_weight_scale_2$", + *([r"experts\.down_proj_input_scale$"] if calibrated else []), + ], + target_patterns=[ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[FineGrainedWeightGlobals(self)], + ), + ] + if calibrated: + fused.append( WeightConverter( - source_patterns=["weight$", "weight_scale_inv", "activation_scale"], - target_patterns="weight", - operations=[FineGrainedDequantize(self)], + source_patterns=r"experts\.gate_up_proj_input_scale$", + target_patterns="experts.gate_up_proj_input_global_scale", + operations=[FineGrainedInputGlobals(self)], ) - ] - return [] + ) + # the fused expert WEIGHTS take the packed uint8 -> int8 view the `.weight`-anchored + # rule below gives the per-expert layout; the layout ops ride on top of both + fused += [ + WeightConverter( + source_patterns=rf"experts\.{proj}$", + target_patterns=f"experts.{proj}", + operations=[FineGrainedViewPackedInt8(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + return per_expert + fused def update_weight_conversions(self, weight_conversions): """Rewrite the model's conversion plan for this checkpoint: @@ -458,7 +540,7 @@ def update_weight_conversions(self, weight_conversions): updated.append(conv) return updated + self.get_weight_conversions() - if self.pre_quantized and self._quant_method() == "modelopt": + if self.pre_quantized and self._quant_method() == "nvfp4": updated = [] for conv in weight_conversions: if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): @@ -505,11 +587,10 @@ def layout_ops(target: str) -> list: t for t in targets if re.search(r"experts\.(gate_up_proj|down_proj)(_scale_inv|_bias)?\$?$", t) ] if expert_targets: - conv = WeightConverter( - source_patterns=conv.source_patterns, - target_patterns=conv.target_patterns, - operations=list(conv.operations) + layout_ops(expert_targets[0]), - ) + # Append to the EXISTING converter rather than rebuild it: `scope_prefix`, + # `base_model_prefix` and `force_cpu` are set on it before this hook runs and + # are not `__init__` arguments, so a fresh one silently loses them. + conv.operations = list(conv.operations) + layout_ops(expert_targets[0]) updated.append(conv) # dense linears have no converter: their scale keys take the container cast alone updated.append( diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 42036ff4416f..109af54715fa 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -23,6 +23,22 @@ def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: return module, tensor_name +def subtree_pattern_to_regex(pattern: str) -> str: + """A glob-style skip entry as the bounded regex `should_convert_module` matches. + + Checkpoints name whole subtrees with a trailing star ("model.layers.0*", "model.layers.1.*") + or bare ("self_attn", "model.layers.0."). Those are globs, not regexes, and handing them over + as-is is wrong twice: the dots match any character, and the star is greedy, so + "model.layers.1.*" also swallows layers 10-19. The bounded form matches the named module and + everything under it, nothing else. + """ + prefix = pattern[:-2] if pattern.endswith(".*") else pattern.rstrip("*").rstrip(".") + # anchored on both sides at a path boundary: a bare entry names that module wherever it sits + # ("self_attn" reaches "model.layers.2.self_attn.q_proj") without "layers.1" reaching + # "layers.16" or "attn" reaching "self_attn" + return r"(?:^|.*\.)" + re.escape(prefix) + r"(\..*)?$" + + def should_convert_module(full_name, patterns: list[str] | None = None): if patterns is None: return True diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 8868190034ba..588a269340b4 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -18,7 +18,6 @@ import importlib.metadata import json import os -import re from dataclasses import dataclass from enum import Enum from typing import Any, Optional, Union @@ -1687,6 +1686,12 @@ def post_init(self): raise TypeError("shapes must be a dict") +# What a modelopt export's `quant_algo` means in our vocabulary: the format it becomes, and the +# activation format that algo implies when the config does not name one. modelopt exports other +# algos (FP8, INT4 AWQ, W4A8, MXFP4); they are refused rather than silently mis-read. +_MODELOPT_ALGOS = {"NVFP4": (QuantizationMethod.NVFP4, "nvfp4")} + + @dataclass class FineGrainedConfig(QuantizationConfigMixin): """ @@ -1743,35 +1748,29 @@ def __init__( # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") - # NVIDIA modelopt exports: `quant_algo` names the format, the skip list is a glob-style - # `ignore` (or `exclude_modules`, modelopt's own spelling), and `kv_cache_scheme` is - # dropped — the KV cache stays in the compute dtype. The calibrated `input_scale` TENSORS - # are the loader's business, not this config's. + # NVIDIA modelopt exports: `quant_algo` names the format (modelopt covers FP8, INT4 AWQ, + # W4A8 and others; `_MODELOPT_ALGOS` is the subset this path serves) and the skip list is + # a glob-style `ignore`, or `exclude_modules` in modelopt's own spelling. The calibrated + # `input_scale` TENSORS are the loader's business, not this config's. if str(self.quant_method) == "modelopt" or kwargs.get("quant_algo") is not None: quant_algo = kwargs.pop("quant_algo", None) - kwargs.pop("config_groups", None) - kwargs.pop("kv_cache_scheme", None) - kwargs.pop("producer", None) - if quant_algo != "NVFP4": - raise ValueError(f"modelopt checkpoints are supported for quant_algo='NVFP4' only; got {quant_algo!r}") - self.activation_format = activation_format or "nvfp4" + if quant_algo not in _MODELOPT_ALGOS: + raise ValueError( + f"modelopt checkpoints are supported for quant_algo in " + f"{sorted(_MODELOPT_ALGOS)}; got {quant_algo!r}" + ) + # "modelopt" names the PRODUCER, not a format, so the algo it exported becomes the + # format here and nothing downstream has to ask where a config came from. + self.quant_method, algo_activation_format = _MODELOPT_ALGOS[quant_algo] + self.activation_format = activation_format or algo_activation_format ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) if modules_to_not_convert is None and ignore is not None: - # modelopt ships glob-style subtree entries ("model.layers.0*", - # "model.layers.1.*"); translate each to a BOUNDED regex — a naive - # "*" -> ".*" leaves the dots unescaped and "layers.1..*" swallows - # layers 10-19 (and ``should_convert_module``'s bare-prefix clause - # would even match "layers.16" from "layers.1") - - # (``exclude_modules`` entries are bare subtrees — "self_attn", "layers.0." — - # which land on the same bounded form) - def _subtree_regex(glob): - prefix = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") - return re.escape(prefix) + r"(\..*)?$" - - modules_to_not_convert = [_subtree_regex(g) for g in ignore] + # modelopt ships glob-style subtree entries, under `ignore` as "model.layers.0*" / + # "model.layers.1.*" and under `exclude_modules` bare ("self_attn", "layers.0.") + from ..quantizers.quantizers_utils import subtree_pattern_to_regex + + modules_to_not_convert = [subtree_pattern_to_regex(g) for g in ignore] self.modules_to_not_convert = modules_to_not_convert - # TODO: check overlap with not to convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme self.weight_block_size = weight_block_size @@ -1837,7 +1836,6 @@ def __init__( if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") self.modules_to_not_convert = modules_to_not_convert - # TODO: check overlap with not to convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme self.weight_block_size = weight_block_size diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 391133ef7a4c..264d7cbad1ee 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -46,7 +46,6 @@ from test_utils import make_experts, make_fp8_experts import transformers.integrations.deepgemm as dg -from transformers.distributed.utils import is_dtensor from transformers.integrations.deepgemm import ( deepgemm_bf16_experts_forward, deepgemm_fp8_fp4_experts_forward, @@ -230,21 +229,6 @@ def run(x): run(torch.zeros(3, device=torch_device)) # a graph break / traced probe would raise here - def test_to_local_is_compile_safe(self): - # Regression guard: every experts forward here unwraps its weights through `to_local`, so it runs - # inside the traced region. It calls `is_dtensor`, whose torch-distributed availability check must - # therefore fold to a constant instead of being traced — its body reaches `importlib.metadata`, - # which dynamo cannot follow (and follows differently across Python versions). `@lru_cache` is no - # protection: dynamo steps past cache wrappers, so the marker has to sit on the wrapped function. - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return x.to_local() + 1 if is_dtensor(x) else x + 1 - - out = run(torch.zeros(3, device=torch_device)) # a graph break / traced probe would raise here - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - # ── Capturing fake DeepGEMM bundle ───────────────────────────────────────────── # @@ -375,6 +359,26 @@ def _bundle(self, *, is_sm100): # ── deepgemm_fp8_fp4_linear ──────────────────────────────────────────────── + def test_every_arm_refuses_hidden_states_that_are_not_bfloat16(self): + """Each arm builds its intermediates and its output in bfloat16, so anything else changes + the dtype the caller gets back. The check belongs to the guard, not to one forward: it + used to be copied into two of the three and missing from Mega MoE, which allocates its + output bfloat16 regardless.""" + import transformers.integrations.deepgemm as dg + + class _Experts(torch.nn.Module): + activation_scheme = "dynamic" + + arms = ( + dg.deepgemm_bf16_experts_forward, + dg.deepgemm_fp8_fp4_experts_forward, + dg.deepgemm_fp8_fp4_megamoe_experts_forward, + ) + for forward in arms: + with self.subTest(forward=forward.__name__), self.assertRaises(ValueError) as caught: + forward(_Experts(), torch.zeros(2, 4, dtype=torch.float16), None, None) + self.assertIn("bfloat16", str(caught.exception)) + def test_post_expert_norm_runs_on_the_reducing_arms_and_is_refused_by_megamoe(self): """Mega MoE fuses the routing-weighted reduce, so a per-expert output norm has nowhere to go and is refused. The other arms reduce in `_combine_routed_output`, so the norm rides the @@ -384,8 +388,7 @@ def test_post_expert_norm_runs_on_the_reducing_arms_and_is_refused_by_megamoe(se applied = [] class _Experts(torch.nn.Module): - post_expert_norm_name = "input_scaled_rms_norm" - post_expert_norm = object() + has_post_expert_norm = True def _apply_post_norm(self, rows): applied.append(tuple(rows.shape)) @@ -397,7 +400,7 @@ def _apply_post_norm(self, rows): self.assertIs(dg._apply_post_expert_norm(module, rows), rows) self.assertEqual(applied, [(4, 8)]) - # a model naming no norm is untouched + # a model declaring no norm is untouched plain = torch.nn.Module() self.assertIs(dg._apply_post_expert_norm(plain, rows), rows) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 997a034650a2..b7245d078e17 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -15,7 +15,12 @@ mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` (the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" +import os +import socket +import subprocess +import tempfile import unittest +from contextlib import ExitStack, contextmanager from dataclasses import dataclass from unittest import mock @@ -29,7 +34,13 @@ finegrained_linear, load_finegrained_kernel, ) -from transformers.testing_utils import require_torch, require_torch_gpu +from transformers.testing_utils import ( + TestCasePlus, + require_torch, + require_torch_gpu, + require_torch_multi_accelerator, + slow, +) @dataclass @@ -341,8 +352,7 @@ def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): A form the kernels do not implement rides as the module's own ``_apply_post_norm``.""" kernel, rec = _fake_bundle() m = self._experts(has_gate=True) - norm = torch.nn.LayerNorm(m.hidden_dim) - m.post_expert_norm, m._apply_post_norm = norm, norm + m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) @@ -357,20 +367,20 @@ def forward(self, x): self.rows += x.shape[0] return x - m.post_expert_norm = m._apply_post_norm = _Recording() + m.post_expert_norm = _Recording() with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): m(*self._route()) self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") - def test_naming_a_post_expert_norm_requires_the_hook_that_applies_it(self): - """The shared experts interface: a class that names a post-expert norm has to define + def test_declaring_a_post_expert_norm_requires_the_hook_that_applies_it(self): + """The shared experts interface: a class that declares a post-expert norm has to define `_apply_post_norm`, where its own math lives — there is no default passthrough to fall back on, the way `act_fn` has no default activation.""" from transformers.integrations.moe import use_experts_implementation with self.assertRaises(TypeError): - @use_experts_implementation(post_expert_norm="rms_norm") + @use_experts_implementation(has_post_expert_norm=True) class _NoHook(torch.nn.Module): def forward(self, hidden_states, top_k_index, top_k_weights): return hidden_states @@ -381,7 +391,7 @@ def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): folds it into its reduce instead of calling the module.""" kernel, rec = _fake_bundle() m = self._experts(has_gate=True) - m.post_expert_norm = m._apply_post_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True p1, p2, p3, p4 = _loaded(kernel) for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): m.post_expert_norm_name = name @@ -427,18 +437,189 @@ def test_unfusable_act_fn_is_passed_as_the_module_glu(self): class FrozenFp8ShimTest(unittest.TestCase): def test_frozen_module_warns_and_is_self_contained(self): import importlib - import warnings import transformers.integrations.finegrained_fp8 as frozen - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") + # the notice rides the library logger, so it honours transformers' verbosity settings + with self.assertLogs("transformers.integrations.finegrained_fp8", level="WARNING") as captured: importlib.reload(frozen) - self.assertTrue(any(issubclass(x.category, DeprecationWarning) for x in w)) + self.assertTrue(any("frozen for backward compatibility" in line for line in captured.output)) # distinct machinery: the frozen classes are not the live ones self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) +@require_torch +class FineGrainedValidateEnvironmentTest(unittest.TestCase): + """`validate_environment` on the LIVE quantizer. The only tests that covered a + `validate_environment` sat on `FineGrainedFP8HfQuantizer`, which no `from_pretrained` can + reach any more (`AUTO_QUANTIZER_MAPPING["fp8"]` is `FineGrainedHfQuantizer`).""" + + @contextmanager + def _no_accelerator(self): + with ExitStack() as stack: + stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) + stack.enter_context( + mock.patch("transformers.quantizers.quantizer_finegrained.is_torch_xpu_available", return_value=False) + ) + yield + + def _quantizer(self, pre_quantized, **cfg_kwargs): + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**cfg_kwargs)) + quantizer.pre_quantized = pre_quantized + return quantizer + + def test_no_accelerator_refuses_to_quantize(self): + # quantizing a bf16 checkpoint needs the hardware; there is nothing to fall back to + quantizer = self._quantizer(pre_quantized=False) + with self._no_accelerator(), self.assertRaises(RuntimeError): + quantizer.validate_environment() + + def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): + quantizer = self._quantizer(pre_quantized=True) + with self._no_accelerator(): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + def test_nvfp4_refuses_to_dequantize(self): + # two-level scales have no slot in the dequantize chain, so the CPU fallback that rescues + # block-FP8 must refuse here rather than hand back a weight scaled by `1 / global` + for pre_quantized, kwargs in ((True, {}), (False, {"dequantize": True})): + with self.subTest(pre_quantized=pre_quantized): + quantizer = self._quantizer(pre_quantized=pre_quantized, quant_method="nvfp4", **kwargs) + with self._no_accelerator(), self.assertRaises(NotImplementedError): + quantizer.validate_environment() + + +class FineGrainedParallelPlanTest(unittest.TestCase): + """`update_tp_plan`: a model's plan as a quantized experts module needs it.""" + + BASE_EP = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise", + } + + def _planned(self, impl=None, raw=None): + from types import SimpleNamespace + + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = FineGrainedHfQuantizer(FineGrainedConfig()) + config = SimpleNamespace(base_model_tp_plan=dict(raw or self.BASE_EP), _experts_implementation=impl) + quantizer.update_tp_plan(config) + return config.base_model_tp_plan + + def test_every_companion_shards_with_its_expert_weight(self): + """The plan matcher keys on the exact parameter name and falls back only to the owning + module, whose entry shards nothing — so a scale with no entry of its own stays whole while + its weight is this rank's expert slice, and every rank but the first reads another rank's + scales.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + plan = self._planned() + style = lambda name: _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) # noqa: E731 + + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_weight_global_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + for name in ("down_proj_scale_inv", "down_proj_bias", "down_proj_input_global_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + # one per-tensor value for the pre-routing hidden states: replicated, so no entry of its own + self.assertEqual(style("gate_up_proj_input_global_scale"), "moe_tp_experts") + self.assertEqual(plan["layers.*.self_attn.q_proj"], "colwise") + + def test_intra_expert_entries_follow_the_projection_axis(self): + """Under intra-expert TP the experts stay whole and the projection's own axis splits, so + the scale grid must follow its weight while the per-expert globals stay replicated. And + `packed_colwise` splits the output axis as two packed halves — the `[gate; up]` STACK — + where a quantized module holds interleaved rows, so the weight's own style changes too.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + base = { + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + plan = self._planned(raw=base) + style = lambda n: _get_parameter_tp_plan(f"layers.3.mlp.experts.{n}", plan) # noqa: E731 + + # interleaved rows take a CONTIGUOUS split of the output axis, and the scale/bias with them + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias"): + self.assertEqual(style(name), "moe_experts_shard1", name) + # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike + self.assertEqual(style("down_proj"), "rowwise") + self.assertEqual(style("down_proj_scale_inv"), "moe_experts_shard2") + # TP does not split experts, so anything indexed per expert stays whole + for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): + self.assertEqual(style(name), "moe_tp_experts", name) + + # the one backend holding the stack keeps the packed split, weight and companions alike + stacked = self._planned(raw=base, impl="deepgemm_megamoe") + self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", stacked), "packed_colwise") + self.assertEqual( + _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", stacked), + "moe_experts_packed_shard1", + ) + + def test_dense_projections_that_share_a_name_are_left_alone(self): + """`mlp.up_proj` and `mlp.shared_experts.down_proj` end in the same words as the stacked + expert projections but are plain 2-D linears that colwise/rowwise already shard right. + Rewriting them to an expert-axis style sharded a dim they do not have — a CUDA illegal + memory access in the dense FP8 matmul, seen on DeepSeek-V3 under TP.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + plan = self._planned( + raw={ + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + } + ) + for dense, expected in ( + ("layers.*.mlp.up_proj", "colwise"), + ("layers.*.mlp.down_proj", "rowwise"), + ("layers.*.mlp.shared_experts.up_proj", "colwise"), + ("layers.*.mlp.shared_experts.down_proj", "rowwise"), + ): + self.assertEqual(plan[dense], expected, dense) + # and no scale entry was invented for a dense linear, whose scale `rowwise`/`colwise` + # already reaches through the module fallback + strays = [k for k in plan if k.endswith("_scale_inv") and ".experts." not in k] + self.assertEqual(strays, [], f"invented dense scale entries: {strays}") + # the stacked experts still get theirs + self.assertEqual( + _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), + "moe_experts_shard1", + ) + + def test_the_impl_rewrites_the_layer_kinds(self): + megamoe = self._planned("deepgemm_megamoe") + self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") + self.assertEqual(megamoe["layers.*.mlp.experts"], "megamoe_experts") + + +class SubtreePatternTest(unittest.TestCase): + def test_a_checkpoints_glob_skip_list_is_bounded(self): + """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any + character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" + from transformers.quantizers.quantizers_utils import should_convert_module, subtree_pattern_to_regex + + skip = [subtree_pattern_to_regex(g) for g in ("model.layers.0*", "model.layers.1.*", "self_attn")] + for name in ("model.layers.0.mlp", "model.layers.1.mlp.down_proj", "model.layers.2.self_attn.q_proj"): + self.assertFalse(should_convert_module(name, skip), name) + for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): + self.assertTrue(should_convert_module(name, skip), name) + + @require_torch class FineGrainedMxfp4ConverterTest(unittest.TestCase): """The gpt-oss checkpoint layout through the real ConversionOps: blocks (E, N, K/32, 16) @@ -606,7 +787,7 @@ def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): experts.gate_up_proj_scale_inv.shape, (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), ) - # block-FP8's (N/128, K/128) grid never reaches a scaled-MMA, whatever its scale dtype + # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block _, experts = self._experts("fp8") self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) with mock.patch.object(fg, "_swizzles_scales", return_value=False): @@ -1360,7 +1541,7 @@ def relative(out, reference): ): norm.weight.data.uniform_(0.5, 1.5) experts.post_expert_norm, experts.post_expert_norm_name = norm, norm_name - experts._apply_post_norm = norm # what the swap carries from the model's class + experts.has_post_expert_norm = True normed = self._moe_reference(dequantized, x, idx, wts, norm) self.assertGreater(relative(normed, reference), 0.05, "the norm has to change the output") for name, forward in forwards: @@ -1428,3 +1609,185 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): for proj, grid in affine.items(): restored = reverse.convert({proj: getattr(experts, f"{proj}_scale_inv").data})[proj] self.assertTrue(torch.equal(restored.view(torch.uint8), grid.view(torch.uint8)), proj) + + +# Run inside each rank; rank 0 writes its logits out for the parent to compare against the +# reference it computes in-process. +_WORKER = """ +import os, sys +import torch +from transformers import AutoModelForCausalLM +from transformers.distributed import DistributedConfig + +model_dir, out_dir = sys.argv[1], sys.argv[2] +world = int(os.environ["WORLD_SIZE"]) +ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + +# both modes in ONE launch: they share the tp_size=2 mesh, so the process group is created once +# and only the plan differs — a second torchrun would pay another interpreter + CUDA start +for mode, expert_parallel in (("ep", True), ("tp", False)): + model = AutoModelForCausalLM.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=expert_parallel), + ).eval() + with torch.no_grad(): + logits = model(ids.to(model.device)).logits.float().cpu() + if int(os.environ["RANK"]) == 0: + # the local shard of a stacked expert weight: the witness that this leg sharded at all + gate_up = model.get_parameter("model.layers.0.mlp.experts.gate_up_proj") + local = gate_up.to_local() if hasattr(gate_up, "to_local") else gate_up + torch.save({"logits": logits, "expert_local": tuple(local.shape)}, os.path.join(out_dir, mode + ".pt")) + del model + torch.cuda.empty_cache() +""" + + +@slow +@require_torch_multi_accelerator +class FineGrainedDistributedEquivalenceTest(TestCasePlus): + """Sharded output must match unsharded output. Catches wrong AXES, which shape checks cannot.""" + + @classmethod + def setUpClass(cls): + import torch + + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + + cls._tmp = tempfile.TemporaryDirectory() + cls.model_dir = os.path.join(cls._tmp.name, "tiny_moe") + # The kernels autotune per shape from a cold cache, which dwarfs everything else here + # (8+ minutes at the stock 100-trial budget). This test only needs a CORRECT config, not + # a fast one, and the reference is tuned under the same budget — so cut the search and + # share one cache across the reference and the ranks. + cls._env = { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), + } + os.environ.update(cls._env) + # V3, not V4: V4 ships an EMPTY `base_model_tp_plan`, so its TP leg sharded nothing. + # V3 carries `experts.gate_up_proj: packed_colwise` + `down_proj: rowwise` AND dense + # `shared_experts.*`, covering EP, intra-expert TP and the dense siblings the companion + # rules must leave alone. Dims are multiples of the 128 block so a 2-way split stays + # block-aligned and the on-the-fly scales are identical whole or sharded. + cfg = DeepseekV3Config( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=1, + num_attention_heads=4, + # V3 defaults this to 128, which makes `num_key_value_groups` 0 and zeroes out the + # heads `repeat_kv` produces under the eager attention this test pins + num_key_value_heads=4, + n_routed_experts=4, + num_experts_per_tok=2, + n_shared_experts=1, + n_group=1, + topk_group=1, + first_k_dense_replace=0, + max_position_embeddings=32, + q_lora_rank=None, + kv_lora_rank=32, + qk_nope_head_dim=32, + qk_rope_head_dim=16, + v_head_dim=32, + ) + torch.manual_seed(0) + bf16_dir = os.path.join(cls._tmp.name, "bf16") + DeepseekV3ForCausalLM(cfg).save_pretrained(bf16_dir, safe_serialization=True) + + # Quantize ONCE and save, so every load below is `pre_quantized`. Quantizing on the fly + # instead would defeat the whole test: each rank derives its scales from the weight shard + # it already holds, so they come out correctly sized whatever the plan says, and a plan + # that shards no scale at all still produces the right answer. Only a checkpoint scale — + # which arrives whole and must be SPLIT — can catch a missing companion rule. + from transformers import AutoModelForCausalLM + from transformers.utils.quantization_config import FineGrainedFP8Config + + quantized = AutoModelForCausalLM.from_pretrained( + bf16_dir, + dtype="auto", + attn_implementation="eager", + quantization_config=FineGrainedFP8Config(), + device_map="cuda:0", + ) + quantized.save_pretrained(cls.model_dir, safe_serialization=True) + del quantized + torch.cuda.empty_cache() + + @classmethod + def tearDownClass(cls): + cls._tmp.cleanup() + + @staticmethod + def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + def _sharded_logits(self): + """`{mode: logits}` from one 2-rank `torchrun`, through the real TP/EP load path.""" + import torch + + script = os.path.join(self._tmp.name, "worker.py") + with open(script, "w") as fh: + fh.write(_WORKER) + subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={self._free_port()}", + script, + self.model_dir, + self._tmp.name, + ], + check=True, + env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, + ) + return {mode: torch.load(os.path.join(self._tmp.name, f"{mode}.pt")) for mode in ("ep", "tp")} + + def test_ep_and_tp_match_the_unsharded_model(self): + import torch + + from transformers import AutoModelForCausalLM + + # the reference runs here rather than in a third subprocess — one less interpreter start. + # No `quantization_config`: the checkpoint carries it, so this loads pre-quantized. + model = AutoModelForCausalLM.from_pretrained( + self.model_dir, + dtype="auto", + attn_implementation="eager", + device_map="cuda:0", + ).eval() + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + reference = model(ids).logits.float().cpu() + del model + torch.cuda.empty_cache() + + # global stacked shape: (experts, 2 * moe_intermediate, hidden) + whole = (4, 512, 256) + # EP splits the expert axis, intra-expert TP splits the output rows; either way the + # local shard must be SMALLER than the whole, or the leg proved nothing + expected_local = {"ep": (2, 512, 256), "tp": (4, 256, 256)} + + for tag, payload in self._sharded_logits().items(): + with self.subTest(mode=tag): + self.assertEqual( + payload["expert_local"], + expected_local[tag], + f"{tag}: experts were not sharded ({payload['expert_local']} of {whole}) — this " + "leg cannot catch a bad axis, check the model's plan carries expert entries", + ) + sharded = payload["logits"] + self.assertEqual(sharded.shape, reference.shape) + # a mis-paired axis does not drift, it garbles: the tolerance only has to + # absorb reduction-order differences between one GEMM and two + torch.testing.assert_close(sharded, reference, rtol=2e-2, atol=2e-2) + self.assertTrue( + torch.equal(sharded.argmax(-1), reference.argmax(-1)), + f"{tag}: argmax diverged from the unsharded model", + ) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 2a04eeddcf50..7af6bb79591f 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -475,7 +475,7 @@ def test_is_kernel_regex(self): def test_lazy_load_kernel_success_and_cache(self): sentinel = types.ModuleType("sentinel_kernel_module") - def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False): + def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False, check_arch=True): self.assertIn(repo_id, {"kernels-community/causal-conv1d"}) self.assertFalse(allow_all_kernels) return sentinel @@ -520,7 +520,7 @@ def test_lazy_load_kernel_version(self): sentinel_mod = types.ModuleType("sentinel_kernel_module") call_count = {"n": 0} - def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False): + def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False, check_arch=True): call_count["n"] += 1 self.assertEqual(repo_id, "kernels-community/causal-conv1d") self.assertIsNone(revision) diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index 4e0d03d0340e..8b9efa7ca361 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -309,7 +309,8 @@ def test_quantized_model_offload(self): Simple test that checks if the quantized model returns an error when loading with cpu/disk offloaded """ with self.assertRaisesRegex( - ValueError, "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device." + ValueError, + "You are attempting to load a fine-grained quantized model with a device_map that contains a cpu/disk", ): AutoModelForCausalLM.from_pretrained( self.model_name, device_map=self.offload_device_map, quantization_config=self.quantization_config diff --git a/tests/utils/test_import_utils.py b/tests/utils/test_import_utils.py index 74a02676e66b..3f99611b513c 100644 --- a/tests/utils/test_import_utils.py +++ b/tests/utils/test_import_utils.py @@ -364,9 +364,11 @@ def _compile_constant_helpers(): @parameterized.expand(_compile_constant_helpers()) def test_availability_helpers_are_compile_safe(helper_name: str, args: tuple): """ - These helpers get called from inside `torch.compile`d regions — e.g. `is_dtensor`, which every MoE - kernel integration reaches through `to_local`. Each carries `@_make_compile_constant`, so dynamo evaluates - it once at trace time and never enters the body; this checks the marker actually takes effect. + These helpers get called from inside `torch.compile`d regions — e.g. `is_deepgemm_loadable`, which + `prefers_deepgemm_linear` reaches from the traced `fp8_linear` dispatch gate. Each carries + `@_make_compile_constant`, so dynamo evaluates it once at trace time and never enters the body + (whose probes reach `importlib.metadata`, which dynamo cannot follow); this checks the marker + actually takes effect. Folding rather than keeping the bodies traceable is deliberate. Most bottom out in `_is_package_available`, whose `importlib.metadata` lookup dynamo cannot follow — and follows From fe0b71825d831eada4fe44a54339003f82f10be7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 17 Sep 2026 09:19:57 -0700 Subject: [PATCH 29/68] fixes for the new model --- src/transformers/integrations/finegrained.py | 12 +- .../quantizers/quantizer_finegrained.py | 40 +++--- tests/kernels/test_finegrained.py | 123 +++++++++++++++--- 3 files changed, 136 insertions(+), 39 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index b2348098565d..98662f69cc0f 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -607,7 +607,13 @@ def _moe_operands(kernel, module) -> dict: norm_eps, norm_weight = 1e-6, None if not module.has_post_expert_norm: post_expert_norm = None - elif module.post_expert_norm_name in kernel.get_supported_norms(): + elif module.post_expert_norm_name in kernel.get_supported_norms() and not getattr( + module.post_expert_norm, "_hf_tp_input_reduce", False + ): + # fusing puts the norm INSIDE the launch that produced the rows, so a collective that has + # to land between the two has nowhere to go: under intra-expert TP those rows are a + # partial sum over the sharded intermediate. Fall back to the module, whose wrapped + # forward does the all-reduce first. post_expert_norm = module.post_expert_norm_name norm_weight, norm_eps = module.post_expert_norm.weight, _norm_eps(module.post_expert_norm) else: @@ -1217,7 +1223,9 @@ def __init__(self, hf_quantizer=None): def convert(self, input_dict, target_patterns=None, model=None, **kwargs): sources = defaultdict(dict) for key, value in input_dict.items(): - sources[_global_role(key)][key] = value + # the loader hands a converter's tensors in a list; unwrapped, a fused `(E, 2)` pair + # reads as `(1, 2E)` and the per-half fold below silently does not fire + sources[_global_role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value globals_ = {} for target in target_patterns: if role_sources := sources.get(_global_role(target)): diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 545f8550abdc..c9351b13eed5 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -234,7 +234,6 @@ def update_tp_plan(self, config): from ..integrations.finegrained import FineGrainedExperts impl = getattr(config, "_experts_implementation", None) - interleaved_gate_up = impl != "deepgemm_megamoe" layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): base_plan = getattr(config, plan_attr, None) or {} @@ -275,13 +274,13 @@ def update_tp_plan(self, config): if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): updated_plan.setdefault(f"{key}_scale_inv", self._expert_shard(2)) elif projection != "down_proj" and style in ("packed_colwise", "colwise"): - # `packed_colwise` splits the output axis as two packed halves — the - # `[gate; up]` STACK. A quantized module holds the kernels' INTERLEAVED rows, - # where that split hands a rank gate|up pairs for channels its own down - # projection does not serve, so the weight's own style changes with it. - rows = self._expert_shard(1, packed=not interleaved_gate_up) - if interleaved_gate_up: - updated_plan[key] = rows + # The companions take the WEIGHT's own split. The model declares which one it + # needs and is right: the interleave into the kernels' row order runs on each + # rank's shard, AFTER sharding, so the split has to match the layout the + # CHECKPOINT holds. `packed_colwise` is the strided split a concatenated + # `[gate; up]` keeps pairs together with; `colwise` the contiguous one a + # natively interleaved layout (GPT-OSS, `is_concatenated=False`) needs. + rows = self._expert_shard(1, packed=style == "packed_colwise") updated_plan.setdefault(f"{key}_scale_inv", rows) updated_plan.setdefault(f"{key}_bias", rows) @@ -406,14 +405,14 @@ def _nvfp4_conversions(self): per_expert = [ WeightConverter( source_patterns=[ - "mlp.experts.*.gate_proj.weight_scale$", - "mlp.experts.*.up_proj.weight_scale$", + r"mlp\.experts\..*\.gate_proj\.weight_scale$", + r"mlp\.experts\..*\.up_proj\.weight_scale$", ], target_patterns="mlp.experts.gate_up_proj_scale_inv", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight_scale$", + source_patterns=r"mlp\.experts\..*\.down_proj\.weight_scale$", target_patterns="mlp.experts.down_proj_scale_inv", operations=[MergeModulelist(dim=0)], ), @@ -423,10 +422,10 @@ def _nvfp4_conversions(self): # decided together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ - "mlp.experts.*.gate_proj.weight_scale_2", - "mlp.experts.*.up_proj.weight_scale_2", - "mlp.experts.*.down_proj.weight_scale_2", - *(["mlp.experts.*.down_proj.input_scale"] if calibrated else []), + r"mlp\.experts\..*\.gate_proj\.weight_scale_2", + r"mlp\.experts\..*\.up_proj\.weight_scale_2", + r"mlp\.experts\..*\.down_proj\.weight_scale_2", + *([r"mlp\.experts\..*\.down_proj\.input_scale"] if calibrated else []), ], target_patterns=[ "mlp.experts.gate_up_proj_weight_global_scale", @@ -440,8 +439,8 @@ def _nvfp4_conversions(self): per_expert.append( WeightConverter( source_patterns=[ - "mlp.experts.*.gate_proj.input_scale", - "mlp.experts.*.up_proj.input_scale", + r"mlp\.experts\..*\.gate_proj\.input_scale", + r"mlp\.experts\..*\.up_proj\.input_scale", ], target_patterns="mlp.experts.gate_up_proj_input_global_scale", operations=[MergeModulelist(dim=0), FineGrainedInputGlobals(self)], @@ -581,6 +580,7 @@ def layout_ops(target: str) -> list: return ops updated = [] + covered: set[str] = set() for conv in weight_conversions: targets = conv.target_patterns if isinstance(conv, WeightConverter) else [] expert_targets = [ @@ -591,6 +591,7 @@ def layout_ops(target: str) -> list: # `base_model_prefix` and `force_cpu` are set on it before this hook runs and # are not `__init__` arguments, so a fresh one silently loses them. conv.operations = list(conv.operations) + layout_ops(expert_targets[0]) + covered.update(re.sub(r".*experts\.|\$$", "", t) for t in expert_targets) updated.append(conv) # dense linears have no converter: their scale keys take the container cast alone updated.append( @@ -601,7 +602,10 @@ def layout_ops(target: str) -> list: ) ) for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias", "down_proj_scale_inv"): - if layout_ops(name): + # only where no converter already produces that target: a second converter for the + # same key is appended LATER and shadows the first, dropping whatever ops it carried + # that these do not (the packed uint8 -> int8 view, on a fused modelopt checkpoint) + if name not in covered and layout_ops(name): # the target is replacement text (op outputs are keyed by it), so no regex escapes updated.append( WeightConverter( diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index b7245d078e17..f5e0eb41eb06 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -16,6 +16,7 @@ (the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" import os +import re import socket import subprocess import tempfile @@ -433,6 +434,55 @@ def test_unfusable_act_fn_is_passed_as_the_module_glu(self): self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) +@require_torch +class FineGrainedFusedNormGateTest(unittest.TestCase): + """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. + + Fusing runs the norm inside the launch that produced the rows, so an all-reduce that has to + land between the two has nowhere to go — and under intra-expert TP those rows are a partial + sum over the sharded intermediate, so the fused norm would normalize a fraction of the value. + The wrapped module does the reduce in its own forward, so the fallback is the correct path. + """ + + class _Norm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4)) + self.eps = 1e-6 + + def _operands(self, *, input_reduce: bool): + norm = self._Norm() + if input_reduce: + norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks + + module = mock.Mock() + module.has_gate = True + module.act_fn_name = "silu" + module.swiglu_alpha = None + module.has_post_expert_norm = True + module.post_expert_norm = norm + module.post_expert_norm_name = "input_scaled_rms_norm" + module._apply_post_norm = lambda x: x + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) + return fg._moe_operands(kernel, module), module + + def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): + operands, _ = self._operands(input_reduce=False) + self.assertEqual(operands["post_expert_norm"], "input_scaled_rms_norm") + self.assertIsNotNone(operands["post_expert_norm_weight"], "the fused form needs the norm's weight") + + def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): + operands, module = self._operands(input_reduce=True) + self.assertIs( + operands["post_expert_norm"], + module._apply_post_norm, + "fused past a collective: the kernel would normalize this rank's partial sum", + ) + self.assertIsNone(operands["post_expert_norm_weight"]) + + @require_torch class FrozenFp8ShimTest(unittest.TestCase): def test_frozen_module_warns_and_is_self_contained(self): @@ -548,9 +598,12 @@ def test_intra_expert_entries_follow_the_projection_axis(self): plan = self._planned(raw=base) style = lambda n: _get_parameter_tp_plan(f"layers.3.mlp.experts.{n}", plan) # noqa: E731 - # interleaved rows take a CONTIGUOUS split of the output axis, and the scale/bias with them - for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_bias"): - self.assertEqual(style(name), "moe_experts_shard1", name) + # the weight keeps the style the MODEL declared — `packed_colwise` is the strided split a + # concatenated `[gate; up]` checkpoint needs, since the interleave into the kernels' row + # order runs on each rank's shard, after sharding. The companions follow it. + self.assertEqual(style("gate_up_proj"), "packed_colwise") + for name in ("gate_up_proj_scale_inv", "gate_up_proj_bias"): + self.assertEqual(style(name), "moe_experts_packed_shard1", name) # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike self.assertEqual(style("down_proj"), "rowwise") self.assertEqual(style("down_proj_scale_inv"), "moe_experts_shard2") @@ -558,12 +611,13 @@ def test_intra_expert_entries_follow_the_projection_axis(self): for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): self.assertEqual(style(name), "moe_tp_experts", name) - # the one backend holding the stack keeps the packed split, weight and companions alike - stacked = self._planned(raw=base, impl="deepgemm_megamoe") - self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", stacked), "packed_colwise") + # a natively INTERLEAVED layout declares `colwise`, and its companions split contiguously + inter = dict(base, **{"layers.*.mlp.experts.gate_up_proj": "colwise"}) + plan_i = self._planned(raw=inter) + self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", plan_i), "colwise") self.assertEqual( - _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", stacked), - "moe_experts_packed_shard1", + _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), + "moe_experts_shard1", ) def test_dense_projections_that_share_a_name_are_left_alone(self): @@ -598,7 +652,7 @@ def test_dense_projections_that_share_a_name_are_left_alone(self): # the stacked experts still get theirs self.assertEqual( _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), - "moe_experts_shard1", + "moe_experts_packed_shard1", ) def test_the_impl_rewrites_the_layer_kinds(self): @@ -1154,7 +1208,8 @@ def test_global_scale_converters_carry_an_expert_index(self): globals_converters = [ c for c in self._modelopt_conversions() - if any("global_scale" in t for t in _targets(c)) and any("experts.*." in p for p in c.source_patterns) + if any("global_scale" in t for t in _targets(c)) + and any(re.search(p, "model.layers.0.mlp.experts.7.gate_proj.weight_scale_2") for p in c.source_patterns) ] self.assertTrue(globals_converters, "no global-scale converter found") for conv in globals_converters: @@ -1169,15 +1224,45 @@ def test_both_modelopt_scale_layouts_have_converters(self): as one already-stacked tensor per layer (the fused vLLM layout, Muse-Spark). Both are converted; the block scale takes the module's layout ops either way.""" sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} - for suffix in ("weight_scale$", "weight_scale_2", "input_scale"): - self.assertTrue( - any(p.startswith("mlp.experts.*.gate_proj.") and p.endswith(suffix) for p in sources), - f"no per-expert converter for {suffix}", - ) - self.assertTrue( - any("gate_up_proj_" in p and p.endswith(suffix.rstrip("$") + "$") for p in sources), - f"no fused converter for {suffix}", - ) + for suffix in ("weight_scale", "weight_scale_2", "input_scale"): + per_expert = f"model.layers.0.mlp.experts.7.gate_proj.{suffix}" + fused = f"model.layers.0.mlp.experts.gate_up_proj_{suffix}" + self.assertTrue(any(re.search(p, per_expert) for p in sources), f"no per-expert converter for {suffix}") + self.assertTrue(any(re.search(p, fused) for p in sources), f"no fused converter for {suffix}") + # the per-expert patterns are REGEXES: unescaped dots matched `_` too, so + # `mlp.experts.*.up_proj.weight_scale_2` also claimed the fused key and the + # gate|up globals were flattened instead of folded to one per expert + claimed = [p for p in sources if re.search(p, fused) and re.search(p, per_expert)] + self.assertEqual(claimed, [], f"one pattern claims BOTH layouts for {suffix}: {claimed}") + + def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): + """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair + reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals + where the kernels assert one per expert.""" + import torch + + from transformers.integrations.finegrained import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + experts = 8 + for wrap in (False, True): + pair = torch.rand(experts, 2) + 1 + src = { + "experts.gate_up_proj_weight_scale_2": [pair] if wrap else pair, + "experts.down_proj_weight_scale_2": torch.rand(experts) + 1, + "experts.down_proj_input_scale": torch.rand(experts) + 1, + } + out = FineGrainedWeightGlobals(None).convert(src, target_patterns=targets, model=None) + with self.subTest(list_wrapped=wrap): + self.assertEqual( + out["experts.gate_up_proj_weight_global_scale"].shape, + torch.Size([experts]), + "the gate|up halves were not folded to one global per expert", + ) def test_one_converter_merges_a_layer_s_globals(self): """modelopt calibrates the gate|up halves separately. One converter owns every global of From 22b1001c3ecbbea07a4f002c61782bf1c0560db8 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 17 Sep 2026 12:00:39 -0700 Subject: [PATCH 30/68] fix multimodal and add some testing --- src/transformers/integrations/finegrained.py | 6 +- .../quantizers/quantizer_finegrained.py | 13 +- tests/kernels/test_finegrained.py | 591 ++++++++++++++---- 3 files changed, 481 insertions(+), 129 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 98662f69cc0f..d0962117768b 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -652,8 +652,12 @@ def _fused_experts_forward(module, kernel_forward: str, hidden_states, top_k_ind """The kernels' fused MoE chain (gate_up with the fused GLU epilogue and intermediate requant where supported -> down -> the routing-weighted top-k reduce) over the module's tensors.""" if module.activation_scheme == "static": + # the grouped static kernel takes ONE calibrated scalar for the whole matmul, while an + # experts module holds one per expert; a per-expert tensor would miss the static arm + # (it gates on `As.numel() == 1`) and be read as per-block activation scales instead raise NotImplementedError( - f"the {kernel_forward} experts dispatch does not support activation_scheme='static'; use " + f"the {kernel_forward} experts dispatch does not support activation_scheme='static': the " + "grouped static kernel takes a per-tensor scale and the experts hold one per expert; use " "experts_implementation='eager' or activation_scheme='dynamic'." ) kernel = load_finegrained_kernel() diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index c9351b13eed5..e9596a9554c4 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -235,8 +235,15 @@ def update_tp_plan(self, config): impl = getattr(config, "_experts_implementation", None) layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) - for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): - base_plan = getattr(config, plan_attr, None) or {} + # A multimodal model keeps its experts' plans on a SUB-config: the outer one carries a + # few projector entries and no `base_model_ep_plan` at all, so reading only what we were + # handed adds no companion anywhere while the weights still shard from the sub-config's + # own plan — the scale stays whole against a sharded weight. + sub_configs = [c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None))] + for plan_owner, plan_attr in ( + (owner, attr) for owner in (config, *sub_configs) for attr in ("base_model_tp_plan", "base_model_ep_plan") + ): + base_plan = getattr(plan_owner, plan_attr, None) or {} updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} # Every companion beside a projection weight — block scales, bias, NVFP4 globals, a @@ -285,7 +292,7 @@ def update_tp_plan(self, config): updated_plan.setdefault(f"{key}_bias", rows) if updated_plan != base_plan: - setattr(config, plan_attr, updated_plan) + setattr(plan_owner, plan_attr, updated_plan) return config diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index f5e0eb41eb06..294451bf94f1 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -420,6 +420,31 @@ def test_biases_ride_the_kernel_chain(self): self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) + def test_static_activation_scheme_is_refused_by_the_fused_chain(self): + """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, held + as a parameter and consumed by the eager loop. The fused chain refuses it, and the reason + is a shape mismatch rather than a missing kernel: `w8a8_block_static_fp8_matmul_grouped` + takes ONE calibrated scalar for the whole matmul, while an experts module holds one per + expert. Passing the per-expert tensor would miss the static arm entirely — the dispatch + gates on `As.numel() == 1` — and it would be read as per-block activation scales, which + computes silently wrong results. Lifting this needs the grouped kernel to accept `(E,)`. + """ + static, dynamic = ( + FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) + for scheme in ("static", "dynamic") + ) + held = getattr(static, "gate_up_proj_activation_scale", None) + self.assertIsNotNone(held, "the eager loop needs the scale it holds") + self.assertEqual(held.numel(), _Cfg.num_local_experts, "held PER EXPERT, which is the mismatch") + # a dynamic module holds no slot at all, rather than an unused one + self.assertIsNone(getattr(dynamic, "gate_up_proj_activation_scale", None)) + + hidden = torch.zeros(2, _Cfg.hidden_size) + index = torch.zeros(2, 1, dtype=torch.long) + weights = torch.ones(2, 1) + with self.assertRaises(NotImplementedError): + fg._fused_experts_forward(static, "moe_fused_grouped", hidden, index, weights) + def test_unfusable_act_fn_is_passed_as_the_module_glu(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True, has_bias=True) @@ -655,6 +680,65 @@ def test_dense_projections_that_share_a_name_are_left_alone(self): "moe_experts_packed_shard1", ) + def test_a_multimodal_models_expert_plans_are_reached(self): + """A multimodal model keeps its experts' plans on a SUB-config; the config the quantizer + is handed carries a few projector entries and often no `base_model_ep_plan` at all. + Rewriting only what we were handed adds no companion anywhere, while the weights still + shard from the sub-config's own plan — so the scale stays whole against a sharded weight, + which is a wrong answer rather than a crash.""" + from transformers import Glm4vMoeConfig + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + config = Glm4vMoeConfig() + self.assertFalse( + getattr(config, "base_model_ep_plan", None), + "this model no longer nests its plans, so it cannot guard the sub-config walk", + ) + before = dict(getattr(config.text_config, "base_model_ep_plan", None) or {}) + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + after = getattr(config.text_config, "base_model_ep_plan", None) or {} + + companions = { + k + for k in after + if k.endswith(("_scale_inv", "_bias", "_weight_global_scale", "_input_global_scale", "_activation_scale")) + } + self.assertTrue(companions, "no companion reached the sub-config's plan") + self.assertEqual(set(before), set(after) - companions, "the model's own entries were disturbed") + + def test_dequantize_folds_every_quantized_key_to_full_precision(self): + """`dequantize=True` is the escape hatch for hardware that cannot serve the format. It + must claim BOTH checkpoint key layouts — the `{proj}_blocks` + `{proj}_scales` pair + GPT-OSS ships and the `weight` + `weight_scale_inv` pair everything else does — and + land on the plain weight, since a module that was never swapped has no scale slot to + write into.""" + import re + + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for method, sources in ( + ("mxfp4", ("experts.gate_up_proj_blocks", "experts.gate_up_proj_scales")), + ("fp8", ("mlp.down_proj.weight", "mlp.down_proj.weight_scale_inv")), + ): + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(quant_method=method, dequantize=True)) + quantizer.pre_quantized = True + converters = quantizer.get_weight_conversions() + with self.subTest(quant_method=method): + self.assertTrue(converters, f"{method}: dequantize produced no converter") + for source in sources: + self.assertTrue( + any( + re.search(str(pattern), source) + for c in converters + for pattern in ( + c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns] + ) + ), + f"{method}: nothing dequantizes {source}", + ) + def test_the_impl_rewrites_the_layer_kinds(self): megamoe = self._planned("deepgemm_megamoe") self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") @@ -1235,6 +1319,45 @@ def test_both_modelopt_scale_layouts_have_converters(self): claimed = [p for p in sources if re.search(p, fused) and re.search(p, per_expert)] self.assertEqual(claimed, [], f"one pattern claims BOTH layouts for {suffix}: {claimed}") + def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): + """Two converters for one key is not additive — the later one WINS and silently drops + whatever ops the first carried. That is how a fused modelopt checkpoint lost its packed + uint8 -> int8 view (a catch-all duplicated the fused converter) and how the per-expert + patterns, which are REGEXES whose unescaped dots also match `_`, claimed the fused + globals and flattened the gate|up pair instead of folding it.""" + import re + + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for quant_kwargs in ( + {"quant_method": "fp8"}, + {"quant_method": "mxfp4"}, + {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + ): + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**quant_kwargs)) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + for key in ( + "model.layers.3.mlp.experts.gate_up_proj", + "model.layers.3.mlp.experts.down_proj", + "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", + "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", + ): + claimed = [ + c + for c in converters + for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) + if re.search(str(pattern), key) + ] + with self.subTest(quant=quant_kwargs["quant_method"], key=key.rsplit(".", 1)[1]): + self.assertLessEqual( + len(claimed), + 1, + f"{len(claimed)} converters claim this key; the last one wins and drops " + f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", + ) + def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals @@ -1696,34 +1819,49 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): self.assertTrue(torch.equal(restored.view(torch.uint8), grid.view(torch.uint8)), proj) -# Run inside each rank; rank 0 writes its logits out for the parent to compare against the -# reference it computes in-process. -_WORKER = """ -import os, sys -import torch -from transformers import AutoModelForCausalLM +def _checkpoint_shapes(path): + """`{key: shape}` of a saved checkpoint, for comparing one save against another.""" + import glob + + from safetensors import safe_open + + shapes = {} + for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shapes[key] = tuple(handle.get_slice(key).get_shape()) + return shapes + + +_SHARDING_WORKER = """ +import importlib, os, sys, torch from transformers.distributed import DistributedConfig -model_dir, out_dir = sys.argv[1], sys.argv[2] +model_dir, out_dir, modes, cls_path = sys.argv[1], sys.argv[2], sys.argv[3].split(","), sys.argv[4] +module_name, cls_name = cls_path.split(":") +model_cls = getattr(importlib.import_module(module_name), cls_name) world = int(os.environ["WORLD_SIZE"]) ids = torch.arange(16, dtype=torch.long).unsqueeze(0) -# both modes in ONE launch: they share the tp_size=2 mesh, so the process group is created once -# and only the plan differs — a second torchrun would pay another interpreter + CUDA start -for mode, expert_parallel in (("ep", True), ("tp", False)): - model = AutoModelForCausalLM.from_pretrained( +# both modes in ONE launch: they share the mesh, so the process group is created once and only +# the plan differs -- a second torchrun would pay another interpreter + CUDA start +for mode in modes: + model = model_cls.from_pretrained( model_dir, dtype="auto", attn_implementation="eager", - distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=expert_parallel), + distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=(mode == "ep")), ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, "gate_up_proj", None) + if weight is None: + weight = experts.up_proj + local = weight.to_local() if hasattr(weight, "to_local") else weight with torch.no_grad(): logits = model(ids.to(model.device)).logits.float().cpu() if int(os.environ["RANK"]) == 0: - # the local shard of a stacked expert weight: the witness that this leg sharded at all - gate_up = model.get_parameter("model.layers.0.mlp.experts.gate_up_proj") - local = gate_up.to_local() if hasattr(gate_up, "to_local") else gate_up - torch.save({"logits": logits, "expert_local": tuple(local.shape)}, os.path.join(out_dir, mode + ".pt")) + torch.save({"logits": logits, "expert_local": tuple(local.shape)}, + os.path.join(out_dir, mode + ".pt")) del model torch.cuda.empty_cache() """ @@ -1731,77 +1869,230 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): @slow @require_torch_multi_accelerator -class FineGrainedDistributedEquivalenceTest(TestCasePlus): - """Sharded output must match unsharded output. Catches wrong AXES, which shape checks cannot.""" +class FineGrainedLoadPathEquivalenceTest(TestCasePlus): + """Expert parallelism and intra-expert tensor parallelism must both reproduce what + `device_map` gives, which is the baseline because it places modules across devices without + splitting a single tensor — no process group, no DTensor, no collectives. So it computes the + unsharded answer, on the multi-GPU path people actually deploy. + + Each model is built in the format it actually ships, because the conversion path differs by + format and by checkpoint layout — a per-expert FP8 checkpoint and a fused NVFP4 one reach the + experts through different converters, and bugs have hidden in exactly that gap. + + Three properties make this able to fail, all of which earlier versions lacked: + * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the + shard it already holds, so they come out correctly sized whatever the plan says and a + plan that shards no scale at all still produces the right answer. + * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty + (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is skipped EXPLICITLY rather + than passing vacuously. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + """ - @classmethod - def setUpClass(cls): - import torch + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}` — each model + in the format it actually ships, with the quantization config that format arrives under, + not just its name: block-FP8 carries a `weight_block_size`, and an NVFP4 checkpoint comes + from modelopt under `quant_algo`, which is remapped on construction. Dims are multiples + of the 128 block so a 2-way split stays block-aligned, and the expert count divides the + mesh.""" + from transformers import ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + DeepseekV4Config, + DeepseekV4ForCausalLM, + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + Glm4vMoeTextConfig, + Glm4vMoeVisionConfig, + GptOssConfig, + GptOssForCausalLM, + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLConfig, + ) + from transformers.utils.quantization_config import FineGrainedConfig - from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config - from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM + return { + "deepseek_v3-fp8": ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "n_group": 1, + "topk_group": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "q_lora_rank": None, + "kv_lora_rank": 32, + "qk_nope_head_dim": 32, + "qk_rope_head_dim": 16, + "v_head_dim": 32, + }, + # DeepSeek-V3 ships block-FP8: 128x128 weight blocks, activations quantized + # per token at run time + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), + ), + # MULTIMODAL: the experts' plans live on `text_config`, not on the config the + # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer + # one adds no companion while the weights still shard, which is how a real + # multimodal MoE ends up with whole scales against sharded weights. + "glm4v_moe-nvfp4": ( + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + { + "text_config": Glm4vMoeTextConfig( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, + rope_theta=10000, + tie_word_embeddings=True, + bos_token_id=0, + eos_token_id=0, + pad_token_id=0, + n_routed_experts=4, + n_shared_experts=1, + n_group=1, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ), + "vision_config": Glm4vMoeVisionConfig( + depth=2, + hidden_size=48, + out_hidden_size=256, + intermediate_size=22, + patch_size=14, + spatial_merge_size=1, + temporal_patch_size=2, + ), + }, + # the GLM NVFP4 checkpoints are modelopt exports: `quant_algo` names the format + # and `FineGrainedConfig` remaps it to nvfp4 at construction + FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4"), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases — and + # an EP plan that already names those biases, so the companion rules meet entries + # the model wrote itself + "gpt_oss-mxfp4": ( + GptOssConfig, + GptOssForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + }, + # GPT-OSS ships weight-only: raw bf16 activations against packed fp4 weights + FineGrainedConfig(quant_method="mxfp4", activation_format="bf16"), + ), + # MIXED precision, and the only entry whose expert format is not the quantization + # config's: `expert_dtype` is a model-config side-channel that makes the EXPERTS + # mxfp4 while the dense and attention paths stay block-FP8 — with scales in UE8M0 + # containers rather than fp32, the other `scale_fmt` + "deepseek_v4-fp4_experts+fp8_dense": ( + DeepseekV4Config, + DeepseekV4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "expert_dtype": "fp4", + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second MULTIMODAL nesting, in the group-32 MX format. The sub-config shape + # follows this model's own tester; only the MoE dims are raised to a multiple of + # the 128 block so a 2-way split stays block-aligned. + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + "intermediate_size": 256, + "dense_intermediate_size": 256, + "shared_intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "rotary_dim": 32, + "vocab_size": 64, + "max_position_embeddings": 32, + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "moe_layer_freq": [0, 1], + "layer_types": ["full_attention", "minimax_m3_sparse"], + "tie_word_embeddings": False, + "index_n_heads": 2, + "index_head_dim": 16, + "index_block_size": 8, + "index_topk_blocks": 4, + "index_local_blocks": 1, + }, + "vision_config": { + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 14, + "patch_size": 14, + "temporal_patch_size": 2, + "spatial_merge_size": 1, + }, + "image_token_index": 4, + "video_token_index": 5, + "projector_hidden_size": 256, + "pad_token_id": 2, + }, + FineGrainedConfig(quant_method="mxfp8"), + ), + } + @classmethod + def setUpClass(cls): cls._tmp = tempfile.TemporaryDirectory() - cls.model_dir = os.path.join(cls._tmp.name, "tiny_moe") - # The kernels autotune per shape from a cold cache, which dwarfs everything else here - # (8+ minutes at the stock 100-trial budget). This test only needs a CORRECT config, not - # a fast one, and the reference is tuned under the same budget — so cut the search and - # share one cache across the reference and the ranks. + # the kernels autotune per shape from a cold cache, which dwarfs everything else here cls._env = { "FINEGRAINED_AUTOTUNE_TRIALS": "1", "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), } os.environ.update(cls._env) - # V3, not V4: V4 ships an EMPTY `base_model_tp_plan`, so its TP leg sharded nothing. - # V3 carries `experts.gate_up_proj: packed_colwise` + `down_proj: rowwise` AND dense - # `shared_experts.*`, covering EP, intra-expert TP and the dense siblings the companion - # rules must leave alone. Dims are multiples of the 128 block so a 2-way split stays - # block-aligned and the on-the-fly scales are identical whole or sharded. - cfg = DeepseekV3Config( - vocab_size=64, - hidden_size=256, - intermediate_size=256, - moe_intermediate_size=256, - num_hidden_layers=1, - num_attention_heads=4, - # V3 defaults this to 128, which makes `num_key_value_groups` 0 and zeroes out the - # heads `repeat_kv` produces under the eager attention this test pins - num_key_value_heads=4, - n_routed_experts=4, - num_experts_per_tok=2, - n_shared_experts=1, - n_group=1, - topk_group=1, - first_k_dense_replace=0, - max_position_embeddings=32, - q_lora_rank=None, - kv_lora_rank=32, - qk_nope_head_dim=32, - qk_rope_head_dim=16, - v_head_dim=32, - ) - torch.manual_seed(0) - bf16_dir = os.path.join(cls._tmp.name, "bf16") - DeepseekV3ForCausalLM(cfg).save_pretrained(bf16_dir, safe_serialization=True) - - # Quantize ONCE and save, so every load below is `pre_quantized`. Quantizing on the fly - # instead would defeat the whole test: each rank derives its scales from the weight shard - # it already holds, so they come out correctly sized whatever the plan says, and a plan - # that shards no scale at all still produces the right answer. Only a checkpoint scale — - # which arrives whole and must be SPLIT — can catch a missing companion rule. - from transformers import AutoModelForCausalLM - from transformers.utils.quantization_config import FineGrainedFP8Config - - quantized = AutoModelForCausalLM.from_pretrained( - bf16_dir, - dtype="auto", - attn_implementation="eager", - quantization_config=FineGrainedFP8Config(), - device_map="cuda:0", - ) - quantized.save_pretrained(cls.model_dir, safe_serialization=True) - del quantized - torch.cuda.empty_cache() @classmethod def tearDownClass(cls): @@ -1813,66 +2104,116 @@ def _free_port() -> int: sock.bind(("", 0)) return sock.getsockname()[1] - def _sharded_logits(self): - """`{mode: logits}` from one 2-rank `torchrun`, through the real TP/EP load path.""" + def _fixture(self, label): + """A PRE-QUANTIZED checkpoint of one tiny model, in the format that model ships.""" + import torch + + config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] + config = config_cls(**kwargs) + + bf16_dir = os.path.join(self._tmp.name, label, "bf16") + quant_dir = os.path.join(self._tmp.name, label, "quantized") + torch.manual_seed(0) + model_cls(config).save_pretrained(bf16_dir, safe_serialization=True) + quantized = model_cls.from_pretrained( + bf16_dir, + dtype="auto", + attn_implementation="eager", + quantization_config=quantization_config, + device_map="cuda:0", + ) + quantized.save_pretrained(quant_dir, safe_serialization=True) + del quantized + torch.cuda.empty_cache() + + # SAVE must restore the checkpoint's own layout: the layout ops (gate|up interleave, + # scale container, swizzle) each have a reverse, and a quantized model reloaded and + # written again has to land on the same keys and shapes. A broken reverse writes a + # corrupt checkpoint silently — the module still reads back whatever it wrote. + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map="cuda:0") + round_trip = os.path.join(self._tmp.name, label, "round_trip") + reloaded.save_pretrained(round_trip, safe_serialization=True) + del reloaded + torch.cuda.empty_cache() + self.assertEqual( + _checkpoint_shapes(quant_dir), _checkpoint_shapes(round_trip), f"{label}: save did not round-trip" + ) + return quant_dir, config, model_cls + + def _logits(self, model_dir, model_cls, **load_kwargs): import torch - script = os.path.join(self._tmp.name, "worker.py") + model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + logits = model(ids).logits.float().cpu() + del model + torch.cuda.empty_cache() + return logits + + def _sharded(self, model_dir, model_cls, modes): + """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" + import torch + + script = os.path.join(self._tmp.name, "sharding_worker.py") with open(script, "w") as fh: - fh.write(_WORKER) + fh.write(_SHARDING_WORKER) + out = os.path.join(self._tmp.name, "out") + os.makedirs(out, exist_ok=True) subprocess.run( [ "torchrun", "--nproc_per_node=2", f"--master_port={self._free_port()}", script, - self.model_dir, - self._tmp.name, + model_dir, + out, + ",".join(modes), + # the model's own class: `AutoModelForCausalLM` cannot resolve a multimodal + # `ForConditionalGeneration` from its config + f"{model_cls.__module__}:{model_cls.__name__}", ], check=True, env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, ) - return {mode: torch.load(os.path.join(self._tmp.name, f"{mode}.pt")) for mode in ("ep", "tp")} + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} - def test_ep_and_tp_match_the_unsharded_model(self): + def test_every_load_path_agrees(self): import torch - from transformers import AutoModelForCausalLM - - # the reference runs here rather than in a third subprocess — one less interpreter start. - # No `quantization_config`: the checkpoint carries it, so this loads pre-quantized. - model = AutoModelForCausalLM.from_pretrained( - self.model_dir, - dtype="auto", - attn_implementation="eager", - device_map="cuda:0", - ).eval() - ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) - with torch.no_grad(): - reference = model(ids).logits.float().cpu() - del model - torch.cuda.empty_cache() - - # global stacked shape: (experts, 2 * moe_intermediate, hidden) - whole = (4, 512, 256) - # EP splits the expert axis, intra-expert TP splits the output rows; either way the - # local shard must be SMALLER than the whole, or the leg proved nothing - expected_local = {"ep": (2, 512, 256), "tp": (4, 256, 256)} - - for tag, payload in self._sharded_logits().items(): - with self.subTest(mode=tag): - self.assertEqual( - payload["expert_local"], - expected_local[tag], - f"{tag}: experts were not sharded ({payload['expert_local']} of {whole}) — this " - "leg cannot catch a bad axis, check the model's plan carries expert entries", - ) - sharded = payload["logits"] - self.assertEqual(sharded.shape, reference.shape) - # a mis-paired axis does not drift, it garbles: the tolerance only has to - # absorb reduction-order differences between one GEMM and two - torch.testing.assert_close(sharded, reference, rtol=2e-2, atol=2e-2) - self.assertTrue( - torch.equal(sharded.argmax(-1), reference.argmax(-1)), - f"{tag}: argmax diverged from the unsharded model", - ) + for label in self._model_table(): + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # the baseline, spread across both devices: `max_memory` forces a real split, + # since `auto` alone fits this whole model on one GPU and would quietly compare + # the sharded legs against a single-device load + reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) + + # Only the modes this model's OWN plans shard experts under. A mode whose plan + # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places + # nothing, so running it would pass without being evidence of anything. A + # multimodal model keeps these on a sub-config. + owners = [config] + [ + c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) + ] + modes = [ + mode + for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) + if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) + ] + self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") + whole = None + for mode, payload in self._sharded(model_dir, model_cls, modes).items(): + with self.subTest(model=label, mode=mode): + local = payload["expert_local"] + whole = whole or local + self.assertLess( + local[0] * local[1], + whole[0] * whole[1] * 2, + f"{label}/{mode}: experts were not sharded ({local}) — this leg cannot catch a bad axis", + ) + torch.testing.assert_close(payload["logits"], reference, rtol=2e-2, atol=2e-2) + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) From f9302cc2156e96c2f137a967fe13664c332ca322 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 18 Sep 2026 00:23:36 -0700 Subject: [PATCH 31/68] static --- src/transformers/integrations/finegrained.py | 21 ++++------ tests/kernels/test_finegrained.py | 43 +++++++++++--------- 2 files changed, 31 insertions(+), 33 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index d0962117768b..85aa25ed9982 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -627,11 +627,13 @@ def _moe_operands(kernel, module) -> dict: "gate_up_proj_scale_inv": getattr(module, f"{up}_scale_inv"), "gate_up_proj_weight_global_scale": getattr(module, f"{up}_weight_global_scale"), "gate_up_proj_input_global_scale": getattr(module, f"{up}_input_global_scale"), + "gate_up_proj_activation_scale": getattr(module, f"{up}_activation_scale"), "gate_up_proj_bias": getattr(module, f"{up}_bias"), "down_proj": module.down_proj, "down_proj_scale_inv": module.down_proj_scale_inv, "down_proj_weight_global_scale": module.down_proj_weight_global_scale, "down_proj_input_global_scale": module.down_proj_input_global_scale, + "down_proj_activation_scale": module.down_proj_activation_scale, "down_proj_bias": module.down_proj_bias, # a supported activation NAME is fused into the gate_up epilogue; any other callable # leaves that GEMM plain and runs on the host between the two @@ -651,15 +653,6 @@ def _moe_operands(kernel, module) -> dict: def _fused_experts_forward(module, kernel_forward: str, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: """The kernels' fused MoE chain (gate_up with the fused GLU epilogue and intermediate requant where supported -> down -> the routing-weighted top-k reduce) over the module's tensors.""" - if module.activation_scheme == "static": - # the grouped static kernel takes ONE calibrated scalar for the whole matmul, while an - # experts module holds one per expert; a per-expert tensor would miss the static arm - # (it gates on `As.numel() == 1`) and be read as per-block activation scales instead - raise NotImplementedError( - f"the {kernel_forward} experts dispatch does not support activation_scheme='static': the " - "grouped static kernel takes a per-tensor scale and the experts hold one per expert; use " - "experts_implementation='eager' or activation_scheme='dynamic'." - ) kernel = load_finegrained_kernel() return getattr(kernel, kernel_forward)(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, module)) @@ -770,10 +763,12 @@ def __init__( f"{proj}_input_global_scale", torch.ones(1 if proj == up_name else self.num_experts, dtype=torch.float32) if calibrated else None, ) - - if self.activation_scheme == "static": - self.gate_up_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) - self.down_proj_activation_scale = nn.Parameter(torch.ones(self.num_experts, dtype=torch.float32)) + static = self.activation_scheme == "static" + _set_optional_parameter( + self, + f"{proj}_activation_scale", + torch.ones(self.num_experts, dtype=torch.float32) if static else None, + ) # The model's per-expert output norm, carried over by the swap. A submodule slot, not a # class attribute: a class attribute of the same name shadows what diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 294451bf94f1..cbde20aca5f4 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -420,30 +420,33 @@ def test_biases_ride_the_kernel_chain(self): self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) - def test_static_activation_scheme_is_refused_by_the_fused_chain(self): - """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, held - as a parameter and consumed by the eager loop. The fused chain refuses it, and the reason - is a shape mismatch rather than a missing kernel: `w8a8_block_static_fp8_matmul_grouped` - takes ONE calibrated scalar for the whole matmul, while an experts module holds one per - expert. Passing the per-expert tensor would miss the static arm entirely — the dispatch - gates on `As.numel() == 1` — and it would be read as per-block activation scales, which - computes silently wrong results. Lifting this needs the grouped kernel to accept `(E,)`. - """ + def test_static_activation_scheme_reaches_the_fused_chain(self): + """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a + MoE calibrates each expert separately, so the module holds one per expert. Both fused + chains take them per projection and the kernels apply each tile's own expert scale in + register — a host pre-quant could not, since top-k routes one row to several experts + whose scales differ. A dynamic module holds no such slot and passes `None`.""" static, dynamic = ( FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) for scheme in ("static", "dynamic") ) - held = getattr(static, "gate_up_proj_activation_scale", None) - self.assertIsNotNone(held, "the eager loop needs the scale it holds") - self.assertEqual(held.numel(), _Cfg.num_local_experts, "held PER EXPERT, which is the mismatch") - # a dynamic module holds no slot at all, rather than an unused one - self.assertIsNone(getattr(dynamic, "gate_up_proj_activation_scale", None)) - - hidden = torch.zeros(2, _Cfg.hidden_size) - index = torch.zeros(2, 1, dtype=torch.long) - weights = torch.ones(2, 1) - with self.assertRaises(NotImplementedError): - fg._fused_experts_forward(static, "moe_fused_grouped", hidden, index, weights) + held = static.gate_up_proj_activation_scale + self.assertIsNotNone(held, "the calibrated scale the chain needs") + self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") + # a dynamic module holds the slot empty, like every other optional parameter + self.assertIsNone(dynamic.gate_up_proj_activation_scale) + + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = () + for module, calibrated in ((static, True), (dynamic, False)): + operands = fg._moe_operands(kernel, module) + for projection in ("gate_up_proj", "down_proj"): + with self.subTest(scheme=module.activation_scheme, projection=projection): + scale = operands[f"{projection}_activation_scale"] + self.assertEqual(scale is not None, calibrated) + if calibrated: + self.assertEqual(scale.numel(), _Cfg.num_local_experts) def test_unfusable_act_fn_is_passed_as_the_module_glu(self): kernel, rec = _fake_bundle() From daa2420508e5c3e40672ddb42bf56a11b7487357 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 18 Sep 2026 03:51:45 -0700 Subject: [PATCH 32/68] Load a calibrated checkpoint's input_scale, and fix what per-tensor weights reach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every `input_scale` rule lived in the modelopt converters, where the key means NVFP4's second-level global. A plain-FP8 static checkpoint had no rule at all, so Ministral-3 and Mistral-4 would have loaded with the registered default of 1.0 and quantized every activation against the wrong scale, silently. `_static_activation_conversions` routes it to the slot each module holds: one per expert for both MoE layouts, a rename for dense linears. `FineGrainedInputGlobals` serves both meanings now and is named for what it carries; only the NVFP4 gate_up global collapses to one value, because that global is a split of the block scale and a static scale has no block level to absorb an inflated one. Adding the first per-tensor model to the load-path table then reached four defects that nothing else does. A 0-dim scale was sharded as `Shard(ndim - 2)` and had no axis-0 offset to take; the intra-expert rule split an `(E, 1, 1)` scale on a length-1 axis; on-the-fly quantize emitted the dense NVFP4 global as `weight_weight_global_scale` and left the expert input globals on meta, since `ones_like` inherits the device of an unmaterialized slot; and an expert target is a prefix of its own companions', so on save the weight converter claimed `..._input_global_scale` and tried to split one value into gate|up halves. The table and its round-trip assertion are newer than the last run of them: `@slow` kept the suite that covers all of this out of every fast gate. Dropped it — `@require_torch_multi_accelerator` already gates it on the hardware it needs. Co-Authored-By: Claude Opus 5 --- .../distributed/sharding_utils.py | 6 +- .../distributed/tensor_parallel.py | 20 +++- src/transformers/integrations/finegrained.py | 40 +++++-- .../quantizers/quantizer_finegrained.py | 73 +++++++++++-- tests/kernels/test_finegrained.py | 102 ++++++++++++++++-- 5 files changed, 209 insertions(+), 32 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 3e5e13136bb4..5b8e9297ea77 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -88,8 +88,10 @@ def __init__(self, param: DTensor): local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) # Axis-0 range owned by this rank (used to filter per-expert pieces) # [_axis0_offset, _axis0_offset + _axis0_local_size) - self._axis0_offset = offsets[0] - self._axis0_local_size = local_shape[0] + # a 0-dim (per-tensor) parameter is replicated whole — see `shard_tensor` — so it has no + # axis 0 to take an offset on + self._axis0_offset = offsets[0] if offsets else 0 + self._axis0_local_size = local_shape[0] if local_shape else 1 def shard_tensor( self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 3151f6264a12..7fa3bbbeea54 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -192,7 +192,13 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) + # a 0-dim scale is one value for the whole tensor: no axis to split, every rank keeps it + if meta.ndim == 0: + placement = Replicate() + elif isinstance(module, torch.nn.Embedding): + placement = Shard(1) + else: + placement = Shard(meta.ndim - 2) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -263,7 +269,10 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - if isinstance(module, torch.nn.Embedding): + # a 0-dim scale is one value for the whole tensor: no axis to split, every rank keeps it + if meta.ndim == 0: + placement = Replicate() + elif isinstance(module, torch.nn.Embedding): placement = Shard(0) else: # bias is replicated (added after the row-reduce); weight shards on input dim (-1) @@ -469,7 +478,10 @@ def shard_param(self, module, param, mesh): return shard_dim = self._packed_output_shard_dim(meta.ndim) # Wrap as a DTensor placeholder. Runs on meta — distribute_tensor builds metadata only. - if meta.ndim == 1: + # a 0-dim scale is one value for the whole tensor: no axis to split, every rank keeps it + if meta.ndim == 0: + placement = Replicate() + elif meta.ndim == 1: placement = Shard(shard_dim) else: placement = _StridedShard(dim=shard_dim, split_factor=self.split_factor) @@ -510,7 +522,7 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - placement = Replicate() if meta.ndim == 1 else _StridedShard(dim=-1, split_factor=self.split_factor) + placement = Replicate() if meta.ndim <= 1 else _StridedShard(dim=-1, split_factor=self.split_factor) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 85aa25ed9982..a5f5602a4e07 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1079,7 +1079,16 @@ def convert(self, input_dict, model=None, target_patterns=None, **kwargs): if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): return input_dict interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved - return {key: interleave.convert({key: value}, None, [key])[key] for key, value in input_dict.items()} + out = {} + for key, value in input_dict.items(): + probe = value[0] if isinstance(value, list) and len(value) == 1 else value + # a per-tensor scale is ONE value covering the whole gate|up stack — no row axis to + # reorder, and reshaping its length-1 axis into pairs is what would fail + if getattr(probe, "ndim", 0) < 2 or probe.shape[1] < 2: + out[key] = value + else: + out[key] = interleave.convert({key: value}, None, [key])[key] + return out @property def reverse_op(self) -> ConversionOps: @@ -1279,12 +1288,15 @@ def _stack_globals(sources: dict) -> torch.Tensor: return torch.cat(columns, dim=1) -class FineGrainedInputGlobals(ConversionOps): - """modelopt's calibrated ``input_scale`` in the layout the module holds: one value for the - gate_up (its rows are the hidden states, quantized once BEFORE routing, so the halves and - experts reduce to their max — the global is a split of the block scale, exact for any value - the accumulator multiplies back), and one per expert for the down, whose rows are per expert - by construction.""" +class FineGrainedInputScales(ConversionOps): + """A calibrated checkpoint's ``input_scale`` in the layout the module holds. One value per + quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max — + both halves read the same rows. + + The NVFP4 gate_up global collapses further, to ONE value: its rows are the hidden states, + quantized once BEFORE routing, and the global is a split of the block scale, exact for any + value the accumulator multiplies back. A static activation scale IS the quantization scale, + with no block level to absorb an inflated one, so it keeps its per-expert form.""" def __init__(self, hf_quantizer=None): self.hf_quantizer = hf_quantizer @@ -1296,7 +1308,7 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): values.append(value.float()) stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) per_expert = stacked.reshape(stacked.shape[0], -1).amax(dim=1) - one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") + one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} @@ -1375,12 +1387,20 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc scale = load_finegrained_kernel().swizzle_mx_scales(scale) out = {key: weight, f"{prefix}_scale_inv": scale} if global_scale is not None: - out[f"{prefix}_weight_global_scale"] = global_scale + # the experts name their global `_weight_global_scale`, a dense linear just + # `weight_global_scale` — both sit beside a `_scale_inv`, so ask for the slot + stem = holder[1].removesuffix("_scale_inv") + suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" + out[f"{prefix}{suffix}"] = global_scale # activations are quantized dynamically against the block scales alone here — there is no # calibration pass — so the module's activation globals are the identity held_input = getattr(module, holder[1].replace("_scale_inv", "_input_global_scale"), None) if held_input is not None: - out[f"{prefix}_input_global_scale"] = torch.ones_like(held_input) + # on the weight's device, not the slot's: the module is still on meta here, and + # `ones_like` would inherit that and leave the parameter unmaterialized + out[f"{prefix}_input_global_scale"] = torch.ones( + held_input.shape, dtype=held_input.dtype, device=value.device + ) return out @staticmethod diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index e9596a9554c4..8d55cbc5e88a 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -277,9 +277,13 @@ def update_tp_plan(self, config): # the per-expert globals, activation scales and the down bias (added after the # row-reduce) stay replicated — but the scale grid still has to follow its weight. # dim 1 is the output axis and dim 2 the input axis in BOTH the affine and the - # 5-D swizzled scale container, which is why naming those two suffices. + # 5-D swizzled scale container, which is why naming those two suffices. A + # per-tensor scale has no such grid — one value per expert, `(E, 1, 1)`, whose + # axes are both 1 — so it joins the replicated companions and only its bias splits. + blocked = self.quantization_config.weight_block_size is not None if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): - updated_plan.setdefault(f"{key}_scale_inv", self._expert_shard(2)) + if blocked: + updated_plan.setdefault(f"{key}_scale_inv", self._expert_shard(2)) elif projection != "down_proj" and style in ("packed_colwise", "colwise"): # The companions take the WEIGHT's own split. The model declares which one it # needs and is right: the interleave into the kernels' row order runs on each @@ -288,7 +292,8 @@ def update_tp_plan(self, config): # `[gate; up]` keeps pairs together with; `colwise` the contiguous one a # natively interleaved layout (GPT-OSS, `is_concatenated=False`) needs. rows = self._expert_shard(1, packed=style == "packed_colwise") - updated_plan.setdefault(f"{key}_scale_inv", rows) + if blocked: + updated_plan.setdefault(f"{key}_scale_inv", rows) updated_plan.setdefault(f"{key}_bias", rows) if updated_plan != base_plan: @@ -326,6 +331,8 @@ def get_weight_conversions(self): return self._mxfp4_conversions() if self._quant_method() == "nvfp4": return self._nvfp4_conversions() + if self.quantization_config.activation_scheme == "static": + return self._static_activation_conversions() return [] def _dequantize_conversions(self): @@ -394,7 +401,7 @@ def _nvfp4_conversions(self): before converter collection and mangle these keys.""" from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter from ..integrations.finegrained import ( - FineGrainedInputGlobals, + FineGrainedInputScales, FineGrainedScaleContainer, FineGrainedViewPackedInt8, FineGrainedWeightGlobals, @@ -450,7 +457,7 @@ def _nvfp4_conversions(self): r"mlp\.experts\..*\.up_proj\.input_scale", ], target_patterns="mlp.experts.gate_up_proj_input_global_scale", - operations=[MergeModulelist(dim=0), FineGrainedInputGlobals(self)], + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], ) ) fused = [ @@ -486,7 +493,7 @@ def _nvfp4_conversions(self): WeightConverter( source_patterns=r"experts\.gate_up_proj_input_scale$", target_patterns="experts.gate_up_proj_input_global_scale", - operations=[FineGrainedInputGlobals(self)], + operations=[FineGrainedInputScales(self)], ) ) # the fused expert WEIGHTS take the packed uint8 -> int8 view the `.weight`-anchored @@ -501,6 +508,50 @@ def _nvfp4_conversions(self): ] return per_expert + fused + def _static_activation_conversions(self): + """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in. One + value per quantized module: a dense linear brings one (Ministral-3), and a MoE brings one + per expert, each expert being its own quantized module (Mistral-4) — the gate|up pair + reduces to one per expert, both halves reading the same routed rows. Both expert layouts, + per expert per projection and already stacked per layer; a checkpoint matches one and the + other never fires. NVFP4 consumes ``input_scale`` as its activation GLOBAL instead + (:meth:`_nvfp4_conversions`) — a second level over a block scale, not the scale itself.""" + from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming + from ..integrations.finegrained import FineGrainedInputScales + + per_expert = [ + WeightConverter( + source_patterns=[ + r"mlp\.experts\..*\.gate_proj\.input_scale", + r"mlp\.experts\..*\.up_proj\.input_scale", + ], + target_patterns="mlp.experts.gate_up_proj_activation_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], + ), + WeightConverter( + source_patterns=r"mlp\.experts\..*\.down_proj\.input_scale", + target_patterns="mlp.experts.down_proj_activation_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], + ), + ] + fused = [ + WeightConverter( + source_patterns=rf"experts\.{proj}_input_scale$", + target_patterns=f"experts.{proj}_activation_scale", + operations=[FineGrainedInputScales(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + # every other calibrated module is a dense linear holding the one value itself — a plain + # rename. The lookahead keeps it off the expert keys above, which stack per layer instead + dense = [ + WeightRenaming( + source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", + target_patterns=r"\1.activation_scale", + ) + ] + return per_expert + fused + dense + def update_weight_conversions(self, weight_conversions): """Rewrite the model's conversion plan for this checkpoint: @@ -523,12 +574,16 @@ def update_weight_conversions(self, weight_conversions): scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") weight_conversions = [scale_rename, *weight_conversions] - # a converter's targets become source patterns on save: anchor sharded `.weight` targets - # (Qwen4-Exp's `ngram_embedding`) so they don't also match their scale keys then + # A converter's targets become source patterns on SAVE, and a parameter name is a prefix + # of its own companions' — `experts.gate_up_proj` matches `experts.gate_up_proj_scale_inv` + # and `..._input_global_scale` too. Unanchored, the weight converter claims them and tries + # to split a one-value global into gate|up halves. Anchor every target that names a + # parameter outright: a `.weight` (Qwen4-Exp's `ngram_embedding`) or an expert projection. + anchored = (".weight", "gate_up_proj", "up_proj", "down_proj") for conv in weight_conversions: if isinstance(conv, WeightConverter): conv._original_target_patterns = [ - f"{p}$" if p.endswith(".weight") else p for p in conv._original_target_patterns + f"{p}$" if p.endswith(anchored) else p for p in conv._original_target_patterns ] if self.pre_quantized and self.quantization_config.dequantize: diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index cbde20aca5f4..bef21c84ab9d 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -40,7 +40,6 @@ require_torch, require_torch_gpu, require_torch_multi_accelerator, - slow, ) @@ -1361,6 +1360,71 @@ def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", ) + def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): + """A static checkpoint calibrates one `input_scale` per quantized module. Nothing routed + those keys for plain FP8 — only NVFP4 read them, as its second-level global — so the two + shipped static models (Ministral-3 dense, Mistral-4 MoE) loaded with the registered + default of 1.0: every activation quantized against the wrong scale, silently. + + Each key must reach the slot its module holds, in that module's shape: one value on a + dense linear, one per expert on the stacked experts, and the gate|up pair reduced to one + per expert rather than flattened to 2E.""" + import re + + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + + def claim(key): + return [ + c + for c in converters + for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) + if re.search(str(pattern), key) + ] + + for key, expected in ( + ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), + ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.up_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.down_proj.input_scale", "mlp.experts.down_proj_activation_scale"), + ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), + ): + with self.subTest(key=key.rsplit(".", 2)[-2]): + claimed = claim(key) + self.assertEqual(len(claimed), 1, f"{len(claimed)} converters claim {key}") + # a renaming's target carries a backreference; resolve it against the key + targets = [ + re.sub(str(claimed[0].source_patterns[0]), str(t), key) if "\\1" in str(t) else str(t) + for t in claimed[0]._original_target_patterns + ] + self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") + # a dynamic checkpoint has no such key to route, and a weight-only one ignores it + dynamic = FineGrainedHfQuantizer(FineGrainedConfig(quant_method="fp8")) + dynamic.pre_quantized = True + self.assertEqual(dynamic.get_weight_conversions(), []) + + def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): + """`FineGrainedInputScales` collapses the gate_up to ONE value for the NVFP4 global — the + global is a split of the block scale, so an inflated one is exact. A static scale IS the + quantization scale, with no block level to absorb it, so collapsing it would quantize + every expert against the largest one's range. It stays per expert.""" + import torch + + from transformers.integrations.finegrained import FineGrainedInputScales + + sources = { + "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), + "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), + } + target = "mlp.experts.gate_up_proj_activation_scale" + scale = FineGrainedInputScales().convert(dict(sources), full_layer_name=target)[target] + torch.testing.assert_close(scale, torch.tensor([2.0, 3.0])) # per expert, over the pair + def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals @@ -1465,9 +1529,9 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated `input_scale` reduces to one value; the down's rows are per expert, so its stays per expert — the requant epilogue normalizes each row by its own expert's value.""" - from transformers.integrations.finegrained import FineGrainedInputGlobals + from transformers.integrations.finegrained import FineGrainedInputScales - up = FineGrainedInputGlobals().convert( + up = FineGrainedInputScales().convert( { "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), @@ -1475,7 +1539,7 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d full_layer_name="mlp.experts.gate_up_proj_input_global_scale", )["mlp.experts.gate_up_proj_input_global_scale"] torch.testing.assert_close(up, torch.tensor([3.0])) - down = FineGrainedInputGlobals().convert( + down = FineGrainedInputScales().convert( {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, full_layer_name="mlp.experts.down_proj_input_global_scale", )["mlp.experts.down_proj_input_global_scale"] @@ -1870,7 +1934,6 @@ def _checkpoint_shapes(path): """ -@slow @require_torch_multi_accelerator class FineGrainedLoadPathEquivalenceTest(TestCasePlus): """Expert parallelism and intra-expert tensor parallelism must both reproduce what @@ -1913,6 +1976,8 @@ def _model_table(): GptOssForCausalLM, MiniMaxM3SparseForConditionalGeneration, MiniMaxM3VLConfig, + Mistral4Config, + Mistral4ForCausalLM, ) from transformers.utils.quantization_config import FineGrainedConfig @@ -1945,6 +2010,28 @@ def _model_table(): # per token at run time FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), ), + # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a + # calibrated activation scale per quantized module, which for a MoE is one per expert. + # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense + # counterpart of the same export. + "mistral4-fp8_tensor_static": ( + Mistral4Config, + Mistral4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static"), + ), # MULTIMODAL: the experts' plans live on `text_config`, not on the config the # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer # one adds no companion while the weights still shard, which is how a real @@ -1977,9 +2064,10 @@ def _model_table(): ), "vision_config": Glm4vMoeVisionConfig( depth=2, - hidden_size=48, + num_heads=4, + hidden_size=64, out_hidden_size=256, - intermediate_size=22, + intermediate_size=64, patch_size=14, spatial_merge_size=1, temporal_patch_size=2, From b4dcb7b0b5d930f2ecb98c8dd79156694aba1b0f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 18 Sep 2026 05:01:05 -0700 Subject: [PATCH 33/68] fix --- src/transformers/integrations/finegrained.py | 40 ++++++++++++++++++- .../quantizers/quantizer_finegrained.py | 31 ++++++++------ 2 files changed, 56 insertions(+), 15 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index a5f5602a4e07..64bbc800981f 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -1198,16 +1198,24 @@ class FineGrainedViewPackedInt8(ConversionOps): ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 tensors pass through, so this can ride converters that also match unquantized modules.""" - def __init__(self, hf_quantizer=None): + def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer + self.inverse = inverse def convert(self, input_dict, target_patterns=None, **kwargs): + held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) out = {} for key, value in _keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value - out[key] = value.view(torch.int8) if torch.is_tensor(value) and value.dtype == torch.uint8 else value + out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value return out + @property + def reverse_op(self) -> ConversionOps: + # a bitcast is its own inverse; on save the int8 the module holds goes back out as the + # uint8 bytes the checkpoint packs fp4 into + return FineGrainedViewPackedInt8(self.hf_quantizer, inverse=not self.inverse) + class FineGrainedWeightGlobals(ConversionOps): """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: @@ -1311,6 +1319,34 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScalesSplit(self.hf_quantizer) + + +class FineGrainedInputScalesSplit(ConversionOps): + """Save reverse of :class:`FineGrainedInputScales`: the module's scale back onto every + projection key the checkpoint calibrated one for. The merge took a max — over the gate|up + pair, and for the NVFP4 gate_up global over the experts too — so the folded halves are gone + and each projection is written the value that covers it, which is what the module reads back. + A collapsed global re-expands per expert, since that is the shape the checkpoint holds.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + value = next(iter(input_dict.values())) + value = value[0] if isinstance(value, list) else value + if value.numel() == 1 and model is not None: + experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) + if experts is not None: + value = value.reshape(1).expand(experts.num_experts).contiguous() + return dict.fromkeys(target_patterns or list(input_dict), value) + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScales(self.hf_quantizer) + class FineGrainedQuantize(ConversionOps): """Quantize a full-precision weight on load into the format its module holds, emitting the diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 8d55cbc5e88a..5df106cca868 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -574,17 +574,22 @@ def update_weight_conversions(self, weight_conversions): scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") weight_conversions = [scale_rename, *weight_conversions] - # A converter's targets become source patterns on SAVE, and a parameter name is a prefix - # of its own companions' — `experts.gate_up_proj` matches `experts.gate_up_proj_scale_inv` - # and `..._input_global_scale` too. Unanchored, the weight converter claims them and tries - # to split a one-value global into gate|up halves. Anchor every target that names a - # parameter outright: a `.weight` (Qwen4-Exp's `ngram_embedding`) or an expert projection. - anchored = (".weight", "gate_up_proj", "up_proj", "down_proj") - for conv in weight_conversions: - if isinstance(conv, WeightConverter): - conv._original_target_patterns = [ - f"{p}$" if p.endswith(anchored) else p for p in conv._original_target_patterns - ] + + def anchor(conversions): + """A converter's targets become its source patterns on SAVE, and a parameter name is a + prefix of its own companions' — `experts.gate_up_proj` matches + `experts.gate_up_proj_scale_inv` and `..._input_global_scale` too. Unanchored, the + weight converter claims them on the way out and tries to split a one-value global into + gate|up halves. Anchor every target that names a parameter outright: a `.weight` + (Qwen4-Exp's `ngram_embedding`) or an expert projection. Runs LAST, since + `_with_expert_layout_ops` rebuilds converters and would drop it.""" + for conv in conversions: + if isinstance(conv, WeightConverter): + conv._original_target_patterns = [ + f"{p}$" if p.endswith((".weight", "gate_up_proj", "up_proj", "down_proj")) else p + for p in conv._original_target_patterns + ] + return conversions if self.pre_quantized and self.quantization_config.dequantize: updated = [] @@ -599,7 +604,7 @@ def update_weight_conversions(self, weight_conversions): operations=[FineGrainedDequantize(self), *conv.operations], ) updated.append(conv) - return updated + self.get_weight_conversions() + return anchor(updated + self.get_weight_conversions()) if self.pre_quantized and self._quant_method() == "nvfp4": updated = [] @@ -613,7 +618,7 @@ def update_weight_conversions(self, weight_conversions): updated.append(conv) weight_conversions = updated - return self._with_expert_layout_ops(weight_conversions + self.get_weight_conversions()) + return anchor(self._with_expert_layout_ops(weight_conversions + self.get_weight_conversions())) def _with_expert_layout_ops(self, weight_conversions): """Give every converter that produces expert tensors the layout ``FineGrainedExperts`` holds From f7081ecf4cd629f7f2a49e0c520d8b705e5304fe Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 21 Sep 2026 05:14:56 -0700 Subject: [PATCH 34/68] Finish the finegrained integration: reload, TP sharding, dsv4 norms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reload: an on-the-fly quantized model lost its swizzle reverse, so saving one wrote swizzled scales a reload could not read back. keep_swizzle_reverse_for_save retains it, and the post-load pass is split into the two delegations it always was, with `covered` computed from source-pattern matches against a fused probe name rather than assumed. TP: a tied lm_head IS the input embedding's parameter, so the embedding_rowwise entry that PretrainedConfig adds for tied weights shards the tensor lm_head reads, while lm_head itself sits outside base_model_tp_plan and gets no style — leaving F.linear with a plain input against a DTensor weight. FSDP resolves the pair; TP had nothing. init_parallel_plans now gives it colwise_rep. It reads the config flag rather than `head.weight is embed.weight` because tying happens later in post_init, where the identity test silently misses. glm4v_moe shipped no expert entries, so its experts were never sharded. Added packed_colwise / rowwise / moe_tp_experts plus shared_experts. DeepseekV4RMSNorm returns the input dtype. V4 pins these norms to FP32 via _keep_in_fp32_modules_strict, so `weight * hidden_states` promoted the result and handed FP32 to compressor projections that ship BF16. The cast-only form is bit-identical where the weight already matches (torch.equal). Tests: pin the fixture to bf16, which exposed three defects (two of them upstream, reproducing with no quantizer at all); add a rounding floor that perturbs every sharded block by one ULP, so the load-path comparison measures against accumulated reduction-order drift instead of an absolute tolerance; size shards against the checkpoint rather than against the other leg. _checkpoint_expert_shape also learns DeepSeek's w1/w3 spelling — it only knew gate_proj/up_proj and the fused layout, so the V4 fp4-experts fixture matched neither and failed outright. Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 12 ++ src/transformers/integrations/finegrained.py | 67 ++++++++- .../deepseek_v4/modeling_deepseek_v4.py | 5 +- .../models/deepseek_v4/modular_deepseek_v4.py | 10 +- .../glm4v_moe/configuration_glm4v_moe.py | 6 + .../models/glm4v_moe/modular_glm4v_moe.py | 6 + .../quantizers/quantizer_finegrained.py | 20 ++- tests/kernels/test_finegrained.py | 141 ++++++++++++++++-- 8 files changed, 250 insertions(+), 17 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 646b8ac102cc..8e0969529025 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,6 +80,18 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + # A tied lm_head IS the input embedding's parameter, so the `embedding_rowwise` entry + # `PretrainedConfig` adds for tied weights shards the tensor it reads — while lm_head + # itself sits outside any `base_model_tp_plan` and gets no style, leaving `F.linear` with + # a plain input against a DTensor weight. FSDP resolves the pair; TP had no equivalent. + # the CONFIG flag, not `head.weight is embed.weight`: tying happens later in `post_init`, + # so the parameters are still distinct here and the identity test silently misses + tied = getattr(self.config.get_text_config(), "tie_word_embeddings", False) + if tied and "embedding_rowwise" in self._tp_plan.values(): + head = self.get_output_embeddings() + if head is not None: + self._tp_plan.setdefault(next(n for n, m in self.named_modules() if m is head), "colwise_rep") + @property def tp_plan(self) -> dict[str, str]: """The full tp plan for the model's modules.""" diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 64bbc800981f..3da69395c87a 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -884,6 +884,48 @@ class FineGrainedExpertsInterface(ExpertsInterface): ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() +def keep_swizzle_reverse_for_save(model, hf_quantizer) -> None: + """Keep the swizzle's reverse reachable at save time for a model quantized ON THE FLY. + + `_weight_conversions` retains only converters that matched a checkpoint weight, and + quantizing a bf16 checkpoint CREATES the scale keys — so the converter carrying the + layout ops never matched, and saving would write the module's 5-D SWIZZLE_32_4_4 grid + where the checkpoint format is the affine one (silently: the module reads back whatever + it wrote, so only a re-save catches it). + + Adds a converter that renames nothing and carries the swizzle alone, for scales no + retained converter already covers — reversing it is the unswizzle, which reads the + affine grid off the self-describing 5-D shape. Scoped that way because a converter + appended here is matched BEFORE the retained ones (the save reverses the list), so a + broader pattern would take keys away from the real converter and skip the rest of its + chain. + """ + from ..core_model_loading import WeightConverter + + conversions = list(getattr(model, "_weight_conversions", None) or []) + covered = { + target + for conv in conversions + if isinstance(conv, WeightConverter) + for target in (getattr(conv, "target_patterns", None) or []) + } + held_swizzled = { + name.rsplit(".", 1)[-1] + for name, param in model.named_parameters() + if param.ndim == 5 and name.rsplit(".", 1)[-1].endswith("_scale_inv") + } + missing = sorted(name for name in held_swizzled if not any(name in target for target in covered)) + if not missing: + return + for name in missing: + conv = WeightConverter( + source_patterns=rf"{name}$", target_patterns=name, operations=[FineGrainedSwizzleScales(hf_quantizer)] + ) + conv._was_used = True # it describes a layout this quantizer applied, not one a checkpoint carried + conversions.append(conv) + model._weight_conversions = conversions + + def disable_deepgemm_on_multi_device(model: nn.Module) -> None: """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` @@ -1315,7 +1357,14 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): value = torch.stack(value, dim=0) if isinstance(value, list) else value values.append(value.float()) stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) - per_expert = stacked.reshape(stacked.shape[0], -1).amax(dim=1) + # A scale is a MAGNITUDE (`amax(|x|) / 448`), so reduce over absolute values and keep it + # off zero. A signed reduce returns whatever the largest signed entry is, and the two + # degenerate results it admits are both fatal downstream: a negative scale flips the sign + # of every dequantized row, and a zero one divides the activations by zero — NaN logits + # on a SINGLE device, before any sharding is involved. Real calibrated checkpoints ship + # positive scales, so this is a no-op for them; it is the models whose scales are built + # rather than measured that reach here with a sign or a zero. + per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1).clamp(min=1e-12) one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} @@ -1368,6 +1417,21 @@ def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> result.update(self._quantize_one(key, tensor, self._weight_holder(model, key))) return result + @staticmethod + def _as_expert_rows(module, name: str, value: torch.Tensor) -> torch.Tensor: + """An expert stack with its contraction axis LAST, swapped to the ``(E, rows, K)`` the + module's slot and the kernels take. A model may store its experts transposed (GPT-OSS: + ``(E, H, 2I)``), and quantizing before the swap packs the wrong axis — the row count and + the byte-halved K then disagree, which the grouped op rejects at its first launch.""" + if module is None or value.ndim != 3 or name not in ("gate_up_proj", "up_proj", "down_proj"): + return value + if name == "down_proj": + rows, in_dim = module.hidden_dim, module.intermediate_dim + else: + rows, in_dim = (2 if module.has_gate else 1) * module.intermediate_dim, module.hidden_dim + # only when the orientation is unambiguous: a square stack reads the same both ways + return value.transpose(1, 2).contiguous() if value.shape[1:] == (in_dim, rows) != (rows, in_dim) else value + @staticmethod def _weight_holder(model, key: str): """``(module, scale_param_name)`` when ``key`` is a finegrained module's weight, else ``None``.""" @@ -1387,6 +1451,7 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc return {key: value} prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key module = holder[0] if holder is not None else None + value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) held_scale = getattr(module, holder[1]) if holder is not None else None format_spec = weight_formats()[module.weight_format] if module is not None else None if format_spec is None or format_spec.scale_group is None: diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 99584b51fd64..9d3518ee7662 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -53,11 +53,14 @@ def __init__(self, hidden_size, eps: float = 1e-6) -> None: self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return the input dtype. V4 pins these norms to FP32 (`_keep_in_fp32_modules_strict`), + so `weight * hidden_states` promotes the result and hands FP32 to the compressor + projections, which ship BF16. A no-op where the weight already matches.""" input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) + return (self.weight * hidden_states.to(input_dtype)).to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" diff --git a/src/transformers/models/deepseek_v4/modular_deepseek_v4.py b/src/transformers/models/deepseek_v4/modular_deepseek_v4.py index dd034c26e58c..159491590fdb 100644 --- a/src/transformers/models/deepseek_v4/modular_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modular_deepseek_v4.py @@ -64,7 +64,15 @@ def apply_rotary_pos_emb( class DeepseekV4RMSNorm(DeepseekV3RMSNorm): - pass + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return the input dtype. V4 pins these norms to FP32 (`_keep_in_fp32_modules_strict`), + so `weight * hidden_states` promotes the result and hands FP32 to the compressor + projections, which ship BF16. A no-op where the weight already matches.""" + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states.to(input_dtype)).to(input_dtype) class DeepseekV4UnweightedRMSNorm(nn.Module): diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 1cfb69cf24b1..6954139a1ac6 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -57,6 +57,12 @@ class Glm4vMoeTextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 9e3f67c75826..3e8c672788eb 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -90,6 +90,12 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 5df106cca868..c60dd9c9567b 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -185,8 +185,9 @@ def _process_model_before_weight_loading( ) def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import disable_deepgemm_on_multi_device + from ..integrations.finegrained import disable_deepgemm_on_multi_device, keep_swizzle_reverse_for_save + keep_swizzle_reverse_for_save(model, self) disable_deepgemm_on_multi_device(model) return model @@ -658,7 +659,22 @@ def layout_ops(target: str) -> list: # `base_model_prefix` and `force_cpu` are set on it before this hook runs and # are not `__init__` arguments, so a fresh one silently loses them. conv.operations = list(conv.operations) + layout_ops(expert_targets[0]) - covered.update(re.sub(r".*experts\.|\$$", "", t) for t in expert_targets) + # A name counts as covered only when this converter can SOURCE the fused + # transformers-format key, not merely target it. A checkpoint that already ships + # fused (our own save, and any transformers-format one) names the scale + # `experts.gate_up_proj_scale_inv`, which the modelopt-style sources here + # (`...gate_proj.weight_scale`) never match — so keying on the target let those + # converters suppress the catch-all below and the scale arrived as a plain rename, + # carrying none of the layout ops. The module still ALLOCATES the swizzled grid, + # so the affine tensor lands in a 5-D slot: invisible on one device (the parameter + # is simply replaced and the kernels read the shape they are given) and a + # malformed DTensor under sharding, whose local is affine while its global says + # swizzled. + fused_probe = "model.layers.0.mlp.experts.{}" + for target in expert_targets: + fused = re.sub(r".*experts\.|\$$", "", target) + if any(re.search(pattern, fused_probe.format(fused)) for pattern in (conv.source_patterns or [])): + covered.add(fused) updated.append(conv) # dense linears have no converter: their scale keys take the container cast alone updated.append( diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index bef21c84ab9d..c73360eaa15e 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -1934,6 +1934,44 @@ def _checkpoint_shapes(path): """ +def _checkpoint_expert_shape(model_dir): + """The UNSHARDED `(experts, rows)` of the stacked gate|up projection the module holds. + + Read from the checkpoint, which ships the experts either already fused + (`experts.gate_up_proj`) or one per expert (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`, + whose rows are gate and up separately). The baseline a sharded leg is measured against has to come from outside + the sharded runs: taking it from whichever leg ran first compares one mode to another, and a + mode that placed nothing then still looks sharded. + """ + import glob + import json + import os + + from safetensors import safe_open + + shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + index = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.exists(index): + shards = sorted({os.path.join(model_dir, f) for f in json.load(open(index))["weight_map"].values()}) + + experts: set[int] = set() + rows = 0 + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shape = None + if re.search(r"\.experts\.gate_up_proj$", key): + shape = handle.get_slice(key).get_shape() + if len(shape) == 3: + return (shape[0], shape[1]) + if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): + experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) + rows = max(rows, handle.get_slice(key).get_shape()[0]) + if not experts or not rows: + raise AssertionError(f"no expert projection in {model_dir} to size the shard against") + return (len(experts), 2 * rows) # the module stacks gate and up into one row extent + + @require_torch_multi_accelerator class FineGrainedLoadPathEquivalenceTest(TestCasePlus): """Expert parallelism and intra-expert tensor parallelism must both reproduce what @@ -2130,7 +2168,9 @@ def _model_table(): { "text_config": { "hidden_size": 256, - "intermediate_size": 256, + # 512 so the 2-way split leaves 256: the experts' down projection + # contracts over this, and a sharded 128 has no 128-wide swizzled tile + "intermediate_size": 512, "dense_intermediate_size": 256, "shared_intermediate_size": 256, "num_hidden_layers": 2, @@ -2156,8 +2196,11 @@ def _model_table(): "index_local_blocks": 1, }, "vision_config": { - "hidden_size": 32, - "intermediate_size": 64, + # 256 so the 2-way SPLIT is still 128-aligned: an MXFP8 weight with + # pre-swizzled scales is read in 128-wide K tiles, and a sharded 128 dim + # leaves 64 — which has none to offer, and the tuner has no config at all + "hidden_size": 256, + "intermediate_size": 256, "num_hidden_layers": 2, "num_attention_heads": 4, "num_channels": 3, @@ -2205,7 +2248,12 @@ def _fixture(self, label): bf16_dir = os.path.join(self._tmp.name, label, "bf16") quant_dir = os.path.join(self._tmp.name, label, "quantized") torch.manual_seed(0) - model_cls(config).save_pretrained(bf16_dir, safe_serialization=True) + # BF16 is what every model here ships, and the dtype decides which kernel arms run: a + # float32 checkpoint (torch's default for a freshly built model, which `dtype="auto"` + # then faithfully reloads) sends the weight-only formats down arms `tl.dot_scaled` + # cannot serve at all, so the suite would exercise a dtype nobody deploys and leave the + # real one uncovered. Saving in bf16 makes `"auto"` mean bf16 for every load below. + model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) quantized = model_cls.from_pretrained( bf16_dir, dtype="auto", @@ -2242,6 +2290,50 @@ def _logits(self, model_dir, model_cls, **load_kwargs): torch.cuda.empty_cache() return logits + def _rounding_floor(self, model_dir, model_cls, reference): + """How far this model's logits move when every sharded block is perturbed by BF16 rounding. + + Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different + order, so the first sharded op differs by an ULP. How far that travels is a property of + the MODEL, not of the sharding — a chain of MoE layers can amplify it a hundredfold while + a dense stack barely moves. Injecting the same magnitude on ONE device measures that + amplification directly, giving each model a budget its own conditioning earns. + """ + import torch + + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map="cuda:0" + ).eval() + # Perturb EVERY block a sharded reduction passes through, not one: TP re-orders the sum + # in each of them, so the rounding accumulates down the stack instead of cancelling. + # Perturbing a single site with random noise measured LESS movement than sharding caused + # even at 14x the magnitude, which is what accumulation looks like from the wrong model. + blocks = [ + m + for n, m in model.named_modules() + if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 + ] + generator = torch.Generator(device=model.device).manual_seed(0) + + def perturb(module, args, output): + tensor = output[0] if isinstance(output, tuple) else output + if not torch.is_tensor(tensor): + return output + ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() + noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp + tensor = tensor + noise + return (tensor,) + output[1:] if isinstance(output, tuple) else tensor + + handles = [b.register_forward_hook(perturb) for b in blocks] + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + perturbed = model(ids).logits.float().cpu() + for handle in handles: + handle.remove() + del model + torch.cuda.empty_cache() + return (perturbed - reference).abs().max().item() + def _sharded(self, model_dir, model_cls, modes): """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" import torch @@ -2279,6 +2371,7 @@ def test_every_load_path_agrees(self): # since `auto` alone fits this whole model on one GPU and would quietly compare # the sharded legs against a single-device load reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) + noise_floor = self._rounding_floor(model_dir, model_cls, reference) # Only the modes this model's OWN plans shard experts under. A mode whose plan # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places @@ -2293,18 +2386,42 @@ def test_every_load_path_agrees(self): if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) ] self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") - whole = None + whole = _checkpoint_expert_shape(model_dir) for mode, payload in self._sharded(model_dir, model_cls, modes).items(): with self.subTest(model=label, mode=mode): local = payload["expert_local"] - whole = whole or local + # against the CHECKPOINT's own shape, never another mode's shard: seeding + # this from the first leg made the second leg check EP against TP, so a + # mode that placed nothing still looked sharded (glm4v's TP left its + # experts replicated and passed). self.assertLess( local[0] * local[1], - whole[0] * whole[1] * 2, - f"{label}/{mode}: experts were not sharded ({local}) — this leg cannot catch a bad axis", + whole[0] * whole[1], + f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " + f"this leg cannot catch a bad axis", ) - torch.testing.assert_close(payload["logits"], reference, rtol=2e-2, atol=2e-2) - self.assertTrue( - torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), - f"{label}/{mode}: argmax diverged from the unsharded model", + # Sharding reorders reductions (a rowwise all-reduce sums the same terms + # in another order), so the legs differ by BF16 rounding at the first + # sharded op whatever the axes. What that becomes at the logits is the + # MODEL's business: the NVFP4 fixture amplifies it ~100x across two MoE + # layers (measured: a 1e-3 perturbation moves its logits by 0.99 on ONE + # device), while the dense fixtures barely amplify at all. A fixed + # tolerance therefore asks the impossible of one model and nothing of + # another. Compare against the model's OWN floor instead: perturb the + # unsharded run by the same rounding and take the resulting logit shift as + # the budget. A wrong shard axis lands orders of magnitude above it. + budget = max(2e-2, noise_floor) + gap = (payload["logits"] - reference).abs().max().item() + self.assertLessEqual( + gap, + budget, + f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " + f"model's own rounding budget {budget:.4g} — that is a sharding bug, " + f"not reduction order", ) + if noise_floor <= 2e-2: + # only meaningful where rounding does NOT already move the argmax + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) From 9ddb2140adfda489d736ffada632370254bba295 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 21 Sep 2026 13:55:59 -0700 Subject: [PATCH 35/68] Fix a destructive Qwen3 TP-plan override; split out the conversions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `update_tp_plan` carried `if "Qwen3" in config.__class__.__name__:` and then ASSIGNED a hand-written dense plan over `config.base_model_tp_plan`. Three things were wrong with it: - redundant: the module-keyed entries Qwen3 already ships reach a parameter's scale through the plan matcher's module fallback — `q_proj.weight_scale_inv` resolves to `colwise`, `down_proj.weight_scale_inv` to `rowwise` — so the parameter-keyed duplicates bought nothing; - destructive: assigning dropped what the model ships, including `q_norm`/`k_norm` `replicated_with_grad_allreduce`, and on Qwen3-MoE the whole expert plan (`experts: moe_tp_experts`, `experts.gate_up_proj: packed_colwise`, ...). The companion pass immediately below keys on `moe_tp_experts`, so it then found nothing and the expert scales got no entry either; - over-broad: 25 config classes contain "Qwen3". On the multimodal ones, whose plans live on a sub-config and whose outer config ships none, it stamped a dense text plan onto the outer config where those paths match nothing. Removed. Qwen3 and Qwen3-MoE now keep every shipped entry, Qwen3-MoE additionally gains the three expert companions it was being denied, and a multimodal config's outer plan stays empty. Three assertions cover it, against REAL configs — the existing plan tests build a SimpleNamespace, which a branch keyed on the config CLASS is structurally invisible to. Structure: - the conversion ops move to `integrations/finegrained_conversions.py` (finegrained.py 1631 -> 985). The dependency is one-way; `keep_swizzle_reverse_for_save` goes with them, which is what removes the cycle rather than needing a local import to break it. - `_nvfp4_conversions` splits per checkpoint layout. The two never interact — a checkpoint matches one set and the other never fires — so they were adjacent, not related. - `replace_with_finegrained_layer` hoists the kwargs every swap shares, leaving each branch showing only what is its own. Review comments: - TP styles are registered properly (`moe_experts_colwise` / `moe_experts_packed_colwise` / `moe_experts_rowwise`) instead of `ALL_PARALLEL_STYLES.register` at plan-build time. - the frozen-fp8 deprecation notices fired at IMPORT, for anyone importing `transformers.integrations` at all, one of them through `warnings.warn` rather than the logger. Now one `logger.warning_once` on quantizer construction. - `has_global_scale` dropped: `global_scale_dtype: torch.dtype | None` already answered it. - defensive `getattr` on declared config fields replaced with direct access. - `quantizers_utils.py`, `integrations/tensor_parallel.py` and `integrations/finegrained_fp8.py` are back to their upstream contents and leave this diff entirely. Co-Authored-By: Claude Opus 5 --- src/transformers/core_model_loading.py | 2 +- .../distributed/tensor_parallel.py | 7 + src/transformers/integrations/deepgemm.py | 3 +- src/transformers/integrations/finegrained.py | 859 ++---------------- .../integrations/finegrained_conversions.py | 674 ++++++++++++++ .../integrations/finegrained_fp8.py | 6 - .../integrations/tensor_parallel.py | 23 - .../quantizers/quantizer_finegrained.py | 192 ++-- .../quantizers/quantizer_finegrained_fp8.py | 15 +- .../quantizers/quantizers_utils.py | 16 - src/transformers/utils/quantization_config.py | 19 +- tests/kernels/test_finegrained.py | 120 ++- 12 files changed, 950 insertions(+), 986 deletions(-) create mode 100644 src/transformers/integrations/finegrained_conversions.py diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1526b825ea64..1a0465c7f426 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1150,7 +1150,7 @@ def _internal_many_to_many_conversions() -> tuple[type[ConversionOps], ...]: whenever the integration is imported first (`cannot import name ... from partially initialized module`). Called once per `WeightConverter`, at construction. """ - from .integrations.finegrained import FineGrainedWeightGlobals + from .integrations.finegrained_conversions import FineGrainedWeightGlobals return ( ErnieFuseAndSplitTextVisionExperts, diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 7fa3bbbeea54..40f39ecedeff 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -787,6 +787,13 @@ class ParallelInterface(GeneralInterface): "packed_rowwise": PackedRowwiseParallel(), "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), + # a companion of an expert-stacked (E, N, K) projection — block scales, bias, NVFP4 + # globals — shards on the projection's own axis while the experts stay whole, named + # for the weight split it follows. The built-in styles count axes from the END of the + # shape, which lands inside a swizzled scale grid's inner tile rather than on its rows. + "moe_experts_colwise": MoEParamShard(Shard(1)), + "moe_experts_packed_colwise": MoEParamShard(_StridedShard(dim=1, split_factor=2)), + "moe_experts_rowwise": MoEParamShard(Shard(2)), "ep_router": EpRouterParallel(), "megamoe_router": RouterParallelMegaMoe(), "moe_tp_experts": MoeExpertsParallel(), diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 5d153fc0ea15..7f908b398846 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -924,8 +924,7 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: the ``[E_local, 2*I, *]`` leading dims so downstream ``.size(...)`` reads stay valid. - Runs inside the first forward, where the parameters are already local: the module carries - ``_hf_quantized_needs_local_tp``, so the TP layer hands local shards rather than DTensors. + Runs inside the first forward, where the parameters are already local. """ # The UE8M0 scale grid is group-32 along both dims, so the kernel's SF layout needs both divisible. intermediate_hidden = module.intermediate_dim diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 3da69395c87a..efcb0460ed65 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -17,7 +17,6 @@ import importlib import os import sys -from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass @@ -26,8 +25,7 @@ from torch.nn import functional as F from ..activations import ACT2FN -from ..core_model_loading import ConversionOps, Interleave -from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module +from ..quantizers.quantizers_utils import should_convert_module from ..utils import logging from ..utils.import_utils import is_kernels_available from .deepgemm import ( @@ -83,13 +81,12 @@ class FineGrained: the two MoE forwards, with the weight format resolved off the weight/scale dtypes inside the kernel and the module's ``activation_format`` passed through unchanged — the kernels speak the same vocabulary (``None`` = the weights' format, ``"bf16"`` = weight-only). The experts - forwards hand the kernels' MoE forwards the module's tensors and let them run the chain - (scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce); an - activation outside ``get_supported_act_fns()`` is passed as the module's own callable and runs - on the host between the two GEMMs; a per-expert output norm the kernels implement - (``get_supported_norms()``) rides the same way, by name, and is fused into the chain's reduce. The quantizers serve - on-the-fly quantization into the group formats. All symbols are required — a build missing any - raises at load with the full list. + forwards hand the kernels' MoE forwards the module's tensors and let them run the whole chain: + scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce. What they + cannot fuse runs in between — an activation outside ``get_supported_act_fns()`` as the module's + own callable, a per-expert output norm by name where ``get_supported_norms()`` has it. The + quantizers serve on-the-fly quantization into the group formats. All symbols are required; a + build missing any raises at load with the full list. """ matmul_2d: Callable @@ -256,12 +253,10 @@ def _alloc_expert_proj( ) -> tuple[nn.Parameter, nn.Parameter]: """``(weight, weight_scale_inv)`` Parameters for one expert projection: the weight in the format's storage (FP4 packs two values per byte along K), the scale grid at the format's - granularity (``None`` = one scale per expert). ``min_scale_out`` floors the grid's output dim — - the fused gate_up projection keeps room for both halves (``2``) even when ``proj_out`` is - smaller than the block. ``swizzled`` allocates the grid in the ``SWIZZLE_32_4_4`` artifact - layout instead — 128-row blocks x 4-column groups of ``(2, 256)`` byte tiles, the expert axis - leading so EP shards and gathers it on dim 0 — which needs whole blocks; a projection that - does not tile stays affine, which the kernels read directly.""" + granularity (``None`` = one scale per expert). ``min_scale_out`` floors the grid's output dim + so a fused gate_up keeps room for both halves. ``swizzled`` allocates the grid in the + ``SWIZZLE_32_4_4`` layout instead — expert axis leading, so EP shards it on dim 0 — which needs + whole blocks; a projection that does not tile stays affine.""" weight = torch.empty(num_experts, proj_out, proj_in // format_spec.values_per_byte, dtype=format_spec.weight_dtype) if scale_group is None: scale_shape = (num_experts, max(1, min_scale_out), 1) @@ -278,58 +273,19 @@ def _alloc_expert_proj( def _swizzles_scales(config, format_spec: WeightFormat, activation_format: str | None) -> bool: - """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the Blackwell - tcgen05 scaled-MMA reads directly (plain row-major forces a per-tile gather that caps the scaled - dot below the fp8/fp4 peak): SM100, a triton dispatch (the fused forwards and the eager loop both - read it; the DeepGEMM backends read affine scales), a group-scaled format (MX group-32, NVFP4 - group-16; block-FP8 reaches the scaled-MMA too under UE8M0, but by broadcasting one 128-block - scalar in-register across the group-32 columns it consumes, so it has no per-group grid to lay - out) and a chain that quantizes activations (weight-only reads scales per group affinely).""" + """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the + Blackwell tcgen05 scaled-MMA reads directly (row-major forces a per-tile gather). Needs SM100, + a triton dispatch (the DeepGEMM backends read affine scales), a group-scaled format and a chain + that quantizes activations. Block-FP8 reaches the scaled-MMA under UE8M0 by broadcasting one + 128-block scalar in-register, so it has no per-group grid to lay out.""" return ( getattr(config, "_experts_implementation", None) not in ("deepgemm", "deepgemm_megamoe") and format_spec.scale_group is not None and activation_format != "bf16" - and torch.cuda.is_available() and is_sm100() ) -def finegrained_triton_linear( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale_inv: torch.Tensor, - bias: torch.Tensor | None = None, - activation_scale: torch.Tensor | None = None, - weight_global_scale: torch.Tensor | None = None, - input_global_scale: torch.Tensor | None = None, - activation_format: str | None = None, -) -> torch.Tensor: - """Triton fine-grained linear: fused act-quant + matmul, then optional bias add. - - The weight format is read off the tensors (see ``WeightFormat``); what the caller decides is - the activation side. ``activation_scale`` ``None`` quantizes inline per token, a scalar - quantizes statically against it, and ``activation_format`` picks the format where the weights - leave it open (``"bf16"`` = weight-only). The two NVFP4 globals are the checkpoint's - ``weight_scale_2`` and ``input_scale``; both come back on the accumulator. - """ - kernel = load_finegrained_kernel() - original_shape = input.shape - output = kernel.matmul_2d( - input.reshape(-1, original_shape[-1]), - weight, - activation_scale, - weight_scale_inv, - activation_format=activation_format, - output_dtype=input.dtype, - a_global_scale=input_global_scale, - b_global_scale=weight_global_scale, - ) - output = output.reshape(*original_shape[:-1], output.shape[-1]) - if bias is not None: - output.add_(bias) - return output - - def finegrained_linear( input: torch.Tensor, weight: torch.Tensor, @@ -345,11 +301,10 @@ def finegrained_linear( """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales: - 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — on the pre-SM100 shapes it supports (FP4, - UE8M0 SFs, 128×128 block FP8). Never preferred on SM100: the Triton path is tuned - per shape there and is the only one that reads the pre-swizzled SWIZZLE_32_4_4 - scales, and DeepGEMM's context-bound kernels are the multi-device hazard below. - 2. Triton finegrained fallback — everywhere else: SM100, an ``activation_scale`` + 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) where it can actually run and is faster — the + pre-SM100 shapes it supports (FP4, UE8M0 SFs, 128×128 block FP8). Never on SM100, + where Triton is tuned per shape and is the only path reading pre-swizzled scales. + 2. Triton finegrained fallback everywhere else: SM100, an ``activation_scale`` (DeepGEMM is dynamic-only), or any shape DeepGEMM declined. Args: @@ -370,8 +325,8 @@ def finegrained_linear( (None for other formats). input_global_scale: the NVFP4 calibrated activation global, the checkpoint's ``input_scale`` (None = quantize against the block scales alone). - activation_format: the activation format where the weights leave it open (see - ``finegrained_triton_linear``). + activation_format: the activation format where the weights leave it open + (``"bf16"`` = weight-only). """ if prefers_deepgemm_linear( weight, @@ -405,16 +360,25 @@ def finegrained_linear( "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely." ) - return finegrained_triton_linear( - input, + # Triton arm: fused act-quant + matmul, then optional bias add. `activation_scale` None + # quantizes inline per token, a scalar quantizes statically against it; the two NVFP4 globals + # both come back on the accumulator. + kernel = load_finegrained_kernel() + original_shape = input.shape + output = kernel.matmul_2d( + input.reshape(-1, original_shape[-1]), weight, - weight_scale_inv, - bias, activation_scale, - weight_global_scale=weight_global_scale, - input_global_scale=input_global_scale, + weight_scale_inv, activation_format=activation_format, + output_dtype=input.dtype, + a_global_scale=input_global_scale, + b_global_scale=weight_global_scale, ) + output = output.reshape(*original_shape[:-1], output.shape[-1]) + if bias is not None: + output.add_(bias) + return output class _FineGrainedModule: @@ -515,8 +479,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: class FineGrainedGroupedLinear(FineGrainedLinear): - """Quantized drop-in for block-diagonal grouped linears, in any format `FineGrainedLinear` - serves — block-FP8, MXFP8, MXFP4, NVFP4 or weight-only. + """Quantized drop-in for block-diagonal grouped linears. NOT the MoE experts path + (`FineGrainedExperts`): every group runs on every token, with no routing. The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)` weight; logically that is `n_groups` independent `(out_per_group, in_per_group)` sub-matrices, each @@ -582,13 +546,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -def _norm_eps(norm: nn.Module) -> float: - """A norm module's epsilon, under either name transformers modules give it. ``nn.RMSNorm`` - leaves it ``None`` and resolves it to the dtype's own, so this does the same.""" - eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) - return float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps - - def _moe_operands(kernel, module) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their scales (held swizzled for dot_scaled chains), the per-expert globals and biases, @@ -615,7 +572,12 @@ def _moe_operands(kernel, module) -> dict: # partial sum over the sharded intermediate. Fall back to the module, whose wrapped # forward does the all-reduce first. post_expert_norm = module.post_expert_norm_name - norm_weight, norm_eps = module.post_expert_norm.weight, _norm_eps(module.post_expert_norm) + norm = module.post_expert_norm + norm_weight = norm.weight + # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves it + # None and resolves to the dtype's own + eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) + norm_eps = float(eps) if eps is not None else torch.finfo(norm_weight.dtype).eps else: post_expert_norm = module._apply_post_norm @@ -711,22 +673,22 @@ def __init__( # Expert weight storage is declared by the QUANT config's format key, not the model # config: `weight_format` arrives from `replace_with_finegrained_layer` as the - # checkpoint's quant_method ("fp8", "mxfp8", "mxfp4", "nvfp4"). The DeepSeek V4-style - # `config.expert_dtype = "fp4"` model-config side-channel predates it and is kept as - # the legacy fallback when no format is passed (dsv4 fp4 experts under the "fp8" key). + # checkpoint's quant_method. `config.expert_dtype` is the legacy DeepSeek-V4 side-channel, + # kept as the fallback when no format is passed. if weight_format is None: weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" self.weight_format = weight_format format_spec, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) - # the format decides the second level: whether there is one, and the dtype it is held in + # the format decides the second level: None means there is none, else the dtype it is in self.global_scale_dtype = format_spec.global_scale_dtype - self.has_global_scale = self.global_scale_dtype is not None # What the module HOLDS, as the loader's conversion ops deliver it (and their reverses # restore for the checkpoint): gate|up rows in the kernels' interleaved order unless the # backend packs gate|up itself (DeepGEMM Mega MoE), block scales swizzled where the # scaled-MMA reads them. `is_concatenated` above is the checkpoint's own order: the # loader interleaves where the two disagree, and everything at runtime reads this one. + # FROZEN at load, not a property over `_experts_implementation`: it describes the bytes + # the tensors now hold, so re-deriving it would claim a layout they do not have. impl = getattr(config, "_experts_implementation", None) self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" swizzled = _swizzles_scales(config, format_spec, activation_format) @@ -755,9 +717,11 @@ def __init__( _set_optional_parameter( self, f"{proj}_weight_global_scale", - torch.ones(self.num_experts, dtype=self.global_scale_dtype) if self.has_global_scale else None, + torch.ones(self.num_experts, dtype=self.global_scale_dtype) + if self.global_scale_dtype is not None + else None, ) - calibrated = self.has_global_scale and self.activation_format != "bf16" + calibrated = self.global_scale_dtype is not None and self.activation_format != "bf16" _set_optional_parameter( self, f"{proj}_input_global_scale", @@ -884,48 +848,6 @@ class FineGrainedExpertsInterface(ExpertsInterface): ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() -def keep_swizzle_reverse_for_save(model, hf_quantizer) -> None: - """Keep the swizzle's reverse reachable at save time for a model quantized ON THE FLY. - - `_weight_conversions` retains only converters that matched a checkpoint weight, and - quantizing a bf16 checkpoint CREATES the scale keys — so the converter carrying the - layout ops never matched, and saving would write the module's 5-D SWIZZLE_32_4_4 grid - where the checkpoint format is the affine one (silently: the module reads back whatever - it wrote, so only a re-save catches it). - - Adds a converter that renames nothing and carries the swizzle alone, for scales no - retained converter already covers — reversing it is the unswizzle, which reads the - affine grid off the self-describing 5-D shape. Scoped that way because a converter - appended here is matched BEFORE the retained ones (the save reverses the list), so a - broader pattern would take keys away from the real converter and skip the rest of its - chain. - """ - from ..core_model_loading import WeightConverter - - conversions = list(getattr(model, "_weight_conversions", None) or []) - covered = { - target - for conv in conversions - if isinstance(conv, WeightConverter) - for target in (getattr(conv, "target_patterns", None) or []) - } - held_swizzled = { - name.rsplit(".", 1)[-1] - for name, param in model.named_parameters() - if param.ndim == 5 and name.rsplit(".", 1)[-1].endswith("_scale_inv") - } - missing = sorted(name for name in held_swizzled if not any(name in target for target in covered)) - if not missing: - return - for name in missing: - conv = WeightConverter( - source_patterns=rf"{name}$", target_patterns=name, operations=[FineGrainedSwizzleScales(hf_quantizer)] - ) - conv._was_used = True # it describes a layout this quantizer applied, not one a checkpoint carried - conversions.append(conv) - model._weight_conversions = conversions - - def disable_deepgemm_on_multi_device(model: nn.Module) -> None: """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` @@ -983,6 +905,13 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No # legacy `config.expert_dtype` model-config side-channel, which the ctor falls back to. weight_format = next((f for f in ("mxfp8", "mxfp4", "nvfp4") if quantization_config.quant_method == f), None) + # every swap reads the same three off the quant config; the branches differ only in shape + storage = { + "block_size": quantization_config.weight_block_size, + "activation_scheme": quantization_config.activation_scheme, + "scale_fmt": quantization_config.scale_fmt, + } + has_been_replaced = False for module_name, module in model.named_modules(): if not should_convert_module(module_name, modules_to_not_convert): @@ -991,29 +920,24 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No new_module = None with torch.device("meta"): if module_name.endswith(".experts"): - has_gate = getattr(module, "has_gate", True) - has_bias = getattr(module, "has_bias", False) - is_concatenated = getattr(module, "is_concatenated", True) - config = getattr(module, "config", model.config.get_text_config()) - # the decorator stamps these flags on the instance after __init__, so the model's - # own gate|up convention has to travel through it too + # the decorator stamps these on the instance after __init__, so the model's own + # gate|up convention has to travel through it as well as into the ctor + flags = { + "has_gate": getattr(module, "has_gate", True), + "has_bias": getattr(module, "has_bias", False), + "is_concatenated": getattr(module, "is_concatenated", True), + } new_class = use_experts_implementation( experts_class=FineGrainedExperts, experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, - has_bias=has_bias, - has_gate=has_gate, - is_concatenated=is_concatenated, + **flags, ) new_module = new_class( - config=config, - block_size=quantization_config.weight_block_size, - activation_scheme=quantization_config.activation_scheme, - scale_fmt=quantization_config.scale_fmt, - has_bias=has_bias, - has_gate=has_gate, + config=getattr(module, "config", model.config.get_text_config()), weight_format=weight_format, - activation_format=getattr(quantization_config, "activation_format", None), - is_concatenated=is_concatenated, + activation_format=quantization_config.activation_format, + **storage, + **flags, ) # A per-expert output norm belongs to the model, not the quantization, so the # swap carries the submodule over; `_apply_post_norm` is the standard hook and @@ -1025,32 +949,24 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No new_module.has_post_expert_norm = True new_module.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) elif type(module) is nn.Linear: - # Vanilla `nn.Linear` → standard FineGrainedLinear swap. new_module = FineGrainedLinear( in_features=module.in_features, out_features=module.out_features, - block_size=quantization_config.weight_block_size, - activation_scheme=quantization_config.activation_scheme, - scale_fmt=quantization_config.scale_fmt, has_bias=module.bias is not None, weight_format=weight_format or "fp8", - activation_format=getattr(quantization_config, "activation_format", None), + activation_format=quantization_config.activation_format, + **storage, ) elif isinstance(module, nn.Linear) and "GroupedLinear" in type(module).__name__: - # Block-diagonal grouped linear (e.g. DSv4's `DeepseekV4GroupedLinear`): - # one underlying weight conceptually split into `n_groups` independent - # sub-matmuls fed by disjoint input slices. Vanilla `FineGrainedLinear` would - # collapse those groups into one giant linear and yield the wrong - # output dim, so swap to `FineGrainedGroupedLinear` which keeps the per-group - # bmm contract and runs each block as its own FP8 matmul. + # Block-diagonal grouped linear (e.g. DSv4's `DeepseekV4GroupedLinear`): a + # plain `FineGrainedLinear` would collapse the groups into one giant linear and + # yield the wrong output dim. new_module = FineGrainedGroupedLinear( in_features_per_group=module.in_features, out_features=module.out_features, n_groups=module.n_groups, - block_size=quantization_config.weight_block_size, - activation_scheme=quantization_config.activation_scheme, - scale_fmt=quantization_config.scale_fmt, has_bias=module.bias is not None, + **storage, ) if new_module is not None: # The kernels take raw pointers, so every operand must be a plain tensor. This is @@ -1067,620 +983,3 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No "found to convert. Please double check your model architecture." ) return model - - -# --------------------------------------------------------------------------------------------------- -# Conversion ops. The loader runs a converter's ops on the rank-local tensors it collected for one -# target, and `save_pretrained` runs their reverses on the gathered tensors, so everything below -# maps a checkpoint layout onto the layout the modules above hold — and back. -# --------------------------------------------------------------------------------------------------- - - -def _keyed_by_target(input_dict: dict, target_patterns) -> dict: - """A converter's ops see their single tensor under the SOURCE pattern until an op re-keys it to the - target (the core ops do; the loader then expands the target into the full name) — so a - layout op that is first in its converter re-keys the same way. Multi-tensor dicts (a - deserializer's fully named outputs) pass through.""" - if target_patterns and len(input_dict) == 1: - value = next(iter(input_dict.values())) - return {target_patterns[0]: value[0] if isinstance(value, list) else value} - return input_dict - - -def _held_scale(model, full_layer_name, key): - """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else - ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's - name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" - param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[2] - if not param_name.endswith("_scale_inv") or model is None: - return None, None - module = model.get_submodule(full_layer_name.rpartition(".")[0]) - if not isinstance(module, _FineGrainedModule): - return None, None - return module, getattr(module, param_name, None) - - -def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" - return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) - - -class FineGrainedInterleaveGateUp(ConversionOps): - """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order (weight, - scale grid and bias share the row axis) — the core ``Interleave`` along dim 1. Runs on load and - on save alike wherever the checkpoint's row order (``is_concatenated``) differs from the one - the experts hold (``holds_interleaved_gate_up``).""" - - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - - def convert(self, input_dict, model=None, target_patterns=None, **kwargs): - input_dict = _keyed_by_target(input_dict, target_patterns) - experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None - if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): - return input_dict - interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved - out = {} - for key, value in input_dict.items(): - probe = value[0] if isinstance(value, list) and len(value) == 1 else value - # a per-tensor scale is ONE value covering the whole gate|up stack — no row axis to - # reorder, and reshaping its length-1 axis into pairs is what would fail - if getattr(probe, "ndim", 0) < 2 or probe.shape[1] < 2: - out[key] = value - else: - out[key] = interleave.convert({key: value}, None, [key])[key] - return out - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedInterleaveGateUp(self.hf_quantizer, inverse=not self.inverse) - - -class FineGrainedScaleContainer(ConversionOps): - """A block scale into the dtype its module holds. UE8M0 ships in two containers: the exponent - byte under ``uint8`` (MiniMax mxfp8), reinterpreted into ``float8_e8m0fnu``, and its - power-of-two value in float32 (dsv4-flash-base), cast exactly. Anything else passes through. - The module records which one the checkpoint used, so the reverse restores it on save.""" - - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - - def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): - out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): - value = value[0] if isinstance(value, list) else value - if self.inverse: - # one checkpoint, one container, recorded on the modules at load; e8m0 is only - # ever a scale, so the dtype alone picks the tensors to restore - modules = model.modules() if model is not None else () - container = next((c for m in modules if (c := getattr(m, "scale_container_dtype", None))), None) - if container is not None and value.dtype == _get_ue8m0_dtype(): - value = _recontain(value, container) - else: - module, held = _held_scale(model, full_layer_name, key) - if held is not None and value.dtype != held.dtype: - module.scale_container_dtype = value.dtype - value = _recontain(value, held.dtype) - out[key] = value - return out - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedScaleContainer(self.hf_quantizer, inverse=not self.inverse) - - -class FineGrainedSwizzleScales(ConversionOps): - """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D Parameter, - see ``_swizzles_scales``): one triton launch per rank-local shard, values unchanged. Anything - else the converter carries passes through. The swizzle only ever covers whole 128-row blocks - and 4-column groups, so the reverse reads the affine grid back off the 5-D shape.""" - - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - - def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): - out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): - value = value[0] if isinstance(value, list) else value - out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) - return out - - @staticmethod - def _unswizzle(value): - if value.ndim != 5: - return value - experts, row_blocks, col_groups = value.shape[:3] - kernel = load_finegrained_kernel() - return kernel.unswizzle_mx_scales(value, row_blocks * 128, col_groups * 4, num_experts=experts) - - @staticmethod - def _swizzle(key, value, model, full_layer_name): - _, held = _held_scale(model, full_layer_name, key) - if held is None or held.ndim != 5 or value.ndim == 5: - return value - return load_finegrained_kernel().swizzle_mx_scales(value) - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedSwizzleScales(self.hf_quantizer, inverse=not self.inverse) - - -class FineGrainedPackedBlocks(ConversionOps): - """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 — two low-nibble-first E2M1 values per - byte, 16 bytes per group of 32 — as the packed ``(E, N, K/2)`` int8 the kernels read: the same - bytes, grouped per row, folded back into 16-byte groups on the way out. Its ``{proj}_scales`` - companion rides ``FineGrainedScaleContainer`` in its own converter.""" - - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - - def convert(self, input_dict, target_patterns=None, **kwargs): - out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): - value = value[0] if isinstance(value, list) else value - if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) - out[key] = value - elif self.inverse: - out[key] = value.view(torch.uint8).reshape(*value.shape[:-1], value.shape[-1] // 16, 16) - else: - out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) - return out - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedPackedBlocks(self.hf_quantizer, inverse=not self.inverse) - - -class FineGrainedViewPackedInt8(ConversionOps): - """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store: - ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 - tensors pass through, so this can ride converters that also match unquantized modules.""" - - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - - def convert(self, input_dict, target_patterns=None, **kwargs): - held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) - out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): - value = value[0] if isinstance(value, list) else value - out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value - return out - - @property - def reverse_op(self) -> ConversionOps: - # a bitcast is its own inverse; on save the int8 the module holds goes back out as the - # uint8 bytes the checkpoint packs fp4 into - return FineGrainedViewPackedInt8(self.hf_quantizer, inverse=not self.inverse) - - -class FineGrainedWeightGlobals(ConversionOps): - """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: - one fp32 global per expert per projection, from either the per-expert checkpoint keys - (``experts.*.gate_proj.weight_scale_2``) or one stacked ``(E, 2)`` tensor per layer. - - A gate|up stack arrives with two per expert, calibrated separately (up to 1.6x apart on - Muse-Spark, equal on GLM-5.2). Merging them is not local, which is why one op owns a layer's - whole set: the stack keeps the gate's, and the up half's leaves as a ratio. That is sound - because ``inter = silu(gate) * up`` is LINEAR in the up half — running the stack on the gate's - global alone scales the expert output by ``1 / ratio``, which the down projection's weight - global takes back. The down's calibrated input global moves the other way, since it is what - the fused epilogue normalizes the intermediate by before requantizing: without it the requant - drifts off the calibrated range and a large enough ratio pushes the e4m3 block scales past - their 448 ceiling. Rescaling the up half's block scales instead would re-round them against - codes chosen for the old global, moving every value in a block by up to 6%.""" - - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - - def convert(self, input_dict, target_patterns=None, model=None, **kwargs): - sources = defaultdict(dict) - for key, value in input_dict.items(): - # the loader hands a converter's tensors in a list; unwrapped, a fused `(E, 2)` pair - # reads as `(1, 2E)` and the per-half fold below silently does not fire - sources[_global_role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value - globals_ = {} - for target in target_patterns: - if role_sources := sources.get(_global_role(target)): - globals_[target] = _stack_globals(role_sources) - targets = {_global_role(target): target for target in globals_} - stack = targets.get(("gate_up", "weight")) - - # a stack calibrated per half keeps the gate's global; the up half's leaves as a ratio, - # which the down's weight global takes on (scaling the expert output back up) and its - # input global gives back (keeping the requantized intermediate on the calibrated range) - if stack is not None and globals_[stack].shape[1] == 2: - self._assert_no_gate_up_bias(model) - gate, up = globals_[stack][:, 0], globals_[stack][:, 1] - globals_[stack] = gate - ratio = (up / gate)[:, None] - for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): - if target := targets.get(role): - globals_[target] = globals_[target] * factor - return {target: value.reshape(-1).contiguous() for target, value in globals_.items()} - - @staticmethod - def _assert_no_gate_up_bias(model) -> None: - """A gate|up bias is added AFTER the global, so merging the halves would have to divide - its up rows by the same ratio — in the bias converter, which the layout ops own. No - checkpoint pairs the two, so refuse rather than carry a silent half-correction.""" - if model is not None and any( - isinstance(module, FineGrainedExperts) and module.has_bias for module in model.modules() - ): - raise NotImplementedError( - "this checkpoint calibrates the gate|up halves separately AND carries an expert " - "bias; the up half's bias would have to be folded with the globals." - ) - - -def _global_role(key: str) -> tuple[str, str]: - """What a global-scale key holds, for a checkpoint key and a module target alike: which - projection it belongs to, and which of modelopt's two levels it is — the weight's second-level - global (``weight_scale_2``) or the calibrated activation one (``input_scale``).""" - projection = "down" if "down_proj" in key or "w2" in key else "gate_up" - level = "input" if "input_scale" in key or "input_global" in key else "weight" - return projection, level - - -def _stack_globals(sources: dict) -> torch.Tensor: - """One ``(E, calibrated projections)`` fp32 tensor from a layer's globals, gate column first: - two sources are the stack's halves, one source is the fused ``(E, 2)`` pair or a single - projection's ``(E,)``.""" - columns = [] - for key, value in sorted(sources.items(), key=lambda kv: 0 if "gate_proj" in kv[0] else 1): - value = torch.stack(value, dim=0) if isinstance(value, list) else value - columns.append(value.float().reshape(value.shape[0], -1)) - return torch.cat(columns, dim=1) - - -class FineGrainedInputScales(ConversionOps): - """A calibrated checkpoint's ``input_scale`` in the layout the module holds. One value per - quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max — - both halves read the same rows. - - The NVFP4 gate_up global collapses further, to ONE value: its rows are the hidden states, - quantized once BEFORE routing, and the global is a split of the block scale, exact for any - value the accumulator multiplies back. A static activation scale IS the quantization scale, - with no block level to absorb an inflated one, so it keeps its per-expert form.""" - - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - - def convert(self, input_dict, full_layer_name=None, **kwargs): - values = [] - for value in input_dict.values(): - value = torch.stack(value, dim=0) if isinstance(value, list) else value - values.append(value.float()) - stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) - # A scale is a MAGNITUDE (`amax(|x|) / 448`), so reduce over absolute values and keep it - # off zero. A signed reduce returns whatever the largest signed entry is, and the two - # degenerate results it admits are both fatal downstream: a negative scale flips the sign - # of every dequantized row, and a zero one divides the activations by zero — NaN logits - # on a SINGLE device, before any sharding is involved. Real calibrated checkpoints ship - # positive scales, so this is a no-op for them; it is the models whose scales are built - # rather than measured that reach here with a sign or a zero. - per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1).clamp(min=1e-12) - one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only - return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedInputScalesSplit(self.hf_quantizer) - - -class FineGrainedInputScalesSplit(ConversionOps): - """Save reverse of :class:`FineGrainedInputScales`: the module's scale back onto every - projection key the checkpoint calibrated one for. The merge took a max — over the gate|up - pair, and for the NVFP4 gate_up global over the experts too — so the folded halves are gone - and each projection is written the value that covers it, which is what the module reads back. - A collapsed global re-expands per expert, since that is the shape the checkpoint holds.""" - - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - - def convert(self, input_dict, model=None, target_patterns=None, **kwargs): - value = next(iter(input_dict.values())) - value = value[0] if isinstance(value, list) else value - if value.numel() == 1 and model is not None: - experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) - if experts is not None: - value = value.reshape(1).expand(experts.num_experts).contiguous() - return dict.fromkeys(target_patterns or list(input_dict), value) - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedInputScales(self.hf_quantizer) - - -class FineGrainedQuantize(ConversionOps): - """Quantize a full-precision weight on load into the format its module holds, emitting the - scale (and the NVFP4 global) in the module's layout — the held container dtype, the swizzled - artifact where the module holds one. Block-FP8 is computed in torch (per-block E4M3, fp32 or - power-of-two UE8M0 inverse scales); the group formats run the kernels' row-wise quantizers - (``mxfp8_act_quant`` / ``mxfp4_act_quant`` / ``nvfp4_act_quant``), NVFP4 after normalizing by the - canonical per-tensor / per-expert global ``amax / (6 * 448)``. - Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) - pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" - - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer - - def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: - result: dict[str, torch.Tensor] = {} - for key, value in input_dict.items(): - tensor = value[0] if isinstance(value, list) else value - result.update(self._quantize_one(key, tensor, self._weight_holder(model, key))) - return result - - @staticmethod - def _as_expert_rows(module, name: str, value: torch.Tensor) -> torch.Tensor: - """An expert stack with its contraction axis LAST, swapped to the ``(E, rows, K)`` the - module's slot and the kernels take. A model may store its experts transposed (GPT-OSS: - ``(E, H, 2I)``), and quantizing before the swap packs the wrong axis — the row count and - the byte-halved K then disagree, which the grouped op rejects at its first launch.""" - if module is None or value.ndim != 3 or name not in ("gate_up_proj", "up_proj", "down_proj"): - return value - if name == "down_proj": - rows, in_dim = module.hidden_dim, module.intermediate_dim - else: - rows, in_dim = (2 if module.has_gate else 1) * module.intermediate_dim, module.hidden_dim - # only when the orientation is unambiguous: a square stack reads the same both ways - return value.transpose(1, 2).contiguous() if value.shape[1:] == (in_dim, rows) != (rows, in_dim) else value - - @staticmethod - def _weight_holder(model, key: str): - """``(module, scale_param_name)`` when ``key`` is a finegrained module's weight, else ``None``.""" - if model is None: - return None - module, name = get_module_from_name(model, key) - if not isinstance(module, _FineGrainedModule) or getattr(module, name, None) is None: - return None - if name == "weight": - return module, "weight_scale_inv" - if name in ("gate_up_proj", "up_proj", "down_proj"): - return module, f"{name}_scale_inv" - return None - - def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torch.Tensor]: - if value.ndim < 2: - return {key: value} - prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key - module = holder[0] if holder is not None else None - value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) - held_scale = getattr(module, holder[1]) if holder is not None else None - format_spec = weight_formats()[module.weight_format] if module is not None else None - if format_spec is None or format_spec.scale_group is None: - if module is not None and module.block_size: - block = tuple(module.block_size) - else: - # no module (direct invocation): take the block from the config, and failing that - # put one scale over the whole matrix - config = self.hf_quantizer.quantization_config if self.hf_quantizer is not None else None - config_block = ( - config.get("weight_block_size") - if isinstance(config, dict) - else getattr(config, "weight_block_size", None) - ) - block = tuple(config_block) if config_block else (value.shape[-2], value.shape[-1]) - ue8m0 = ( - held_scale.dtype == _get_ue8m0_dtype() - if held_scale is not None - else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" - ) - quantized = self._quantize_block_fp8(value, block, ue8m0) - return {key: value} if quantized is None else {key: quantized[0], f"{prefix}_scale_inv": quantized[1]} - - if value.device.type not in ("cuda", "xpu"): - # the kernels' quantizers are triton launches; the package builds for cuda, rocm - # (which torch reports as "cuda") and xpu - raise ValueError( - f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on an " - f"accelerator, but the weight is on {value.device.type}" - ) - weight, scale, global_scale = self._quantize_group(value, format_spec) - scale = _recontain(scale, held_scale.dtype) - if held_scale.ndim == 5: - scale = load_finegrained_kernel().swizzle_mx_scales(scale) - out = {key: weight, f"{prefix}_scale_inv": scale} - if global_scale is not None: - # the experts name their global `_weight_global_scale`, a dense linear just - # `weight_global_scale` — both sit beside a `_scale_inv`, so ask for the slot - stem = holder[1].removesuffix("_scale_inv") - suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" - out[f"{prefix}{suffix}"] = global_scale - # activations are quantized dynamically against the block scales alone here — there is no - # calibration pass — so the module's activation globals are the identity - held_input = getattr(module, holder[1].replace("_scale_inv", "_input_global_scale"), None) - if held_input is not None: - # on the weight's device, not the slot's: the module is still on meta here, and - # `ones_like` would inherit that and leave the parameter unmaterialized - out[f"{prefix}_input_global_scale"] = torch.ones( - held_input.shape, dtype=held_input.dtype, device=value.device - ) - return out - - @staticmethod - def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): - """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or - ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two - before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" - block_m, block_n = block - rows, cols = value.shape[-2], value.shape[-1] - if rows % block_m or cols % block_n: - return None - tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) - max_abs = tiles.abs().amax(dim=(-3, -1)) - inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) - if ue8m0: - inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) - scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) - quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) - return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) - - @staticmethod - def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): - """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the - kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has - one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the - smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" - kernel = load_finegrained_kernel() - rows = value.reshape(-1, value.shape[-1]) - global_scale = None - if format_spec.global_scale_dtype is not None: - per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) - global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) - rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) - packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) - global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale - else: - quantize = kernel.mxfp4_act_quant if format_spec.values_per_byte == 2 else kernel.mxfp8_act_quant - packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) - return ( - packed.view(format_spec.weight_dtype).reshape(*value.shape[:-1], -1), - scale.reshape(*value.shape[:-1], -1), - global_scale, - ) - - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedDequantize(self.hf_quantizer) - - -class FineGrainedDequantize(ConversionOps): - """A quantized weight folded back to full precision against its per-block scale grid. - - Runs FIRST in its converter under ``dequantize=True``, which is why - :meth:`update_weight_conversions` attaches it to each of the model's own converters: the - merge / concat ops after it collapse the per-expert structure the (weight, scale) pairing - needs. It pairs each weight pattern with its sibling scale pattern by index and emits the - result under the weight key, dropping the scales so the rest of the chain sees weights only. - """ - - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer - - def _scale_pattern_for(self, weight_pattern: str) -> str: - anchored = weight_pattern.endswith("$") - base = weight_pattern[:-1] if anchored else weight_pattern - if base.endswith("_blocks"): - # GPT-OSS packs its experts as `{proj}_blocks` + `{proj}_scales` - scale = base[: -len("_blocks")] + "_scales" - elif base.endswith(".weight"): - scale = base[: -len(".weight")] + ".weight_scale_inv" - elif base == "weight": - scale = "weight_scale_inv" - else: - scale = base + "_scale_inv" - return scale + "$" if anchored else scale - - # packed-FP4 experts arrive as int8 / float4_e2m1fn_x2, two e2m1 nibbles per byte - _FP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) - - def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: - """Two ``e2m1`` FP4 values per byte → float32 tensor twice as wide on the last dim.""" - lut = torch.tensor(self._FP4_E2M1_LUT, dtype=torch.float32, device=packed.device) - u8 = packed.contiguous().view(torch.uint8) - low = (u8 & 0xF).long() - high = ((u8 >> 4) & 0xF).long() - unpacked = torch.stack([lut[low], lut[high]], dim=-1) - return unpacked.reshape(*packed.shape[:-1], 2 * packed.shape[-1]) - - def _dequantize_one( - self, quantized: torch.Tensor, scales: torch.Tensor, output_dtype: torch.dtype | None = None - ) -> torch.Tensor: - # unpacked first, so the rest of the routine sees a normal (rows, cols) float matrix - fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) - if quantized.dtype == torch.int8 or (fp4_dtype is not None and quantized.dtype == fp4_dtype): - quantized_fp32 = self._unpack_fp4(quantized) - else: - quantized_fp32 = quantized.to(torch.float32) - rows, cols = quantized_fp32.shape[-2:] - # the block comes from the scale grid, not the config: one checkpoint holds both a - # ``[1, 32]`` MXFP4 expert block and a ``[128, 128]`` FP8 dense one - try: - scale_rows, scale_cols = scales.shape[-2:] - except Exception: - # scale can be a single tensor in extreme cases where it was not wrapped properly but is [1,0]. - scale_rows, scale_cols = 1, 1 - if rows % scale_rows or cols % scale_cols: - raise ValueError( - f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." - ) - block_m = rows // scale_rows - block_n = cols // scale_cols - # the math runs in fp32 either way (``float8_e8m0fnu`` has no ``mul`` kernel); the - # destination parameter's dtype wins when known, so an eager module keeps the model's - if output_dtype is None: - output_dtype = ( - scales.dtype if scales.dtype.is_floating_point and scales.element_size() >= 2 else torch.bfloat16 - ) - # an MXFP8 checkpoint's uint8 scale is a biased E8M0 exponent, not a multiplier - if scales.dtype == torch.uint8: - s_fp32 = (scales.to(torch.float32) - 127.0).exp2() - else: - s_fp32 = scales.to(torch.float32) - original_shape = quantized_fp32.shape - q = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) - s = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) - return (q * s).to(output_dtype).reshape(original_shape) - - def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: - if model is None or full_layer_name is None: - return None - module, tensor_name = get_module_from_name(model, full_layer_name) - param = getattr(module, tensor_name, None) - return getattr(param, "dtype", None) - - def convert( - self, - input_dict: dict[str, list[torch.Tensor] | torch.Tensor], - full_layer_name: str | None = None, - model: torch.nn.Module | None = None, - **kwargs, - ) -> dict[str, list[torch.Tensor] | torch.Tensor]: - output_dtype = self._get_target_dtype(model, full_layer_name) - # The dense converter (``["weight$", "weight_scale_inv", "activation_scale"] -> weight``) - # hands one weight, with or without a scale (RMSNorm weights match ``weight$`` too). - if "weight$" in input_dict: - # the loader derives prefix/suffix from the output KEY; without a `full_layer_name` - # (direct invocation) key it by the converter's target - target_key = full_layer_name if full_layer_name is not None else "weight" - quantized = input_dict["weight$"] - quantized = quantized[0] if isinstance(quantized, list) else quantized - if "weight_scale_inv" in input_dict: - scales = input_dict["weight_scale_inv"] - scales = scales[0] if isinstance(scales, list) else scales - return {target_key: self._dequantize_one(quantized, scales, output_dtype=output_dtype)} - return {target_key: quantized} - - # Generic chain path: dequantize every weight pattern that has a sibling scale. - consumed = {self._scale_pattern_for(key) for key in input_dict} - result: dict[str, list[torch.Tensor] | torch.Tensor] = {} - for key, value in input_dict.items(): - if "activation_scale" in key or key in consumed: - continue # consumed by the dequant; drop from the chain - scale_key = self._scale_pattern_for(key) - if scale_key not in input_dict: - result[key] = value - continue - weights = value if isinstance(value, list) else [value] - scales = input_dict[scale_key] - scales = scales if isinstance(scales, list) else [scales] - if len(weights) != len(scales): - raise ValueError( - f"FineGrainedDequantize: weight/scale count mismatch for {key} " - f"({len(weights)} weights vs {len(scales)} scales)." - ) - result[key] = [self._dequantize_one(w, s, output_dtype=output_dtype) for w, s in zip(weights, scales)] - return result - - @property - def reverse_op(self) -> ConversionOps: - # a save re-quantizes, so the checkpoint keeps its format whether the in-memory - # state stayed quantized or was dequantized for compute - return FineGrainedQuantize(self.hf_quantizer) diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained_conversions.py new file mode 100644 index 000000000000..626ca478877f --- /dev/null +++ b/src/transformers/integrations/finegrained_conversions.py @@ -0,0 +1,674 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Loader/saver conversion ops for the fine-grained family. + +Each op moves ONE checkpoint-side layout to the one the `finegrained` modules hold, and carries +its own `reverse_op` so `save_pretrained` writes the checkpoint's layout back. They run inside the +`WeightConverter` chains the quantizer builds; nothing here touches a forward. +""" + +from __future__ import annotations + +from collections import defaultdict + +import torch + +from ..core_model_loading import ConversionOps, Interleave +from ..quantizers.quantizers_utils import get_module_from_name +from .finegrained import ( + _FP8_DTYPE, + _FP8_MAX, + _FP8_MIN, + FineGrainedExperts, + WeightFormat, + _FineGrainedModule, + _get_ue8m0_dtype, + load_finegrained_kernel, + weight_formats, +) + + +def _keyed_by_target(input_dict: dict, target_patterns) -> dict: + """A converter's ops see their single tensor under the SOURCE pattern until an op re-keys it to the + target (the core ops do; the loader then expands the target into the full name) — so a + layout op that is first in its converter re-keys the same way. Multi-tensor dicts (a + deserializer's fully named outputs) pass through.""" + if target_patterns and len(input_dict) == 1: + value = next(iter(input_dict.values())) + return {target_patterns[0]: value[0] if isinstance(value, list) else value} + return input_dict + + +def _held_scale(model, full_layer_name, key): + """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else + ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's + name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" + param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[2] + if not param_name.endswith("_scale_inv") or model is None: + return None, None + module = model.get_submodule(full_layer_name.rpartition(".")[0]) + if not isinstance(module, _FineGrainedModule): + return None, None + return module, getattr(module, param_name, None) + + +def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" + return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) + + +class FineGrainedInterleaveGateUp(ConversionOps): + """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order — the + core ``Interleave`` along dim 1, over the weight, scale grid and bias alike. Runs on load and + save wherever the checkpoint's row order differs from the one the experts hold.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + input_dict = _keyed_by_target(input_dict, target_patterns) + experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None + if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): + return input_dict + interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved + out = {} + for key, value in input_dict.items(): + probe = value[0] if isinstance(value, list) and len(value) == 1 else value + # a per-tensor scale is ONE value covering the whole gate|up stack — no row axis to + # reorder, and reshaping its length-1 axis into pairs is what would fail + if getattr(probe, "ndim", 0) < 2 or probe.shape[1] < 2: + out[key] = value + else: + out[key] = interleave.convert({key: value}, None, [key])[key] + return out + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInterleaveGateUp(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedScaleContainer(ConversionOps): + """A block scale into the dtype its module holds. UE8M0 ships either as the exponent byte + under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast exactly); the + module records which the checkpoint used so the reverse restores it on save.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + if self.inverse: + # one checkpoint, one container, recorded on the modules at load; e8m0 is only + # ever a scale, so the dtype alone picks the tensors to restore + modules = model.modules() if model is not None else () + container = next((c for m in modules if (c := getattr(m, "scale_container_dtype", None))), None) + if container is not None and value.dtype == _get_ue8m0_dtype(): + value = _recontain(value, container) + else: + module, held = _held_scale(model, full_layer_name, key) + if held is not None and value.dtype != held.dtype: + module.scale_container_dtype = value.dtype + value = _recontain(value, held.dtype) + out[key] = value + return out + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedScaleContainer(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedSwizzleScales(ConversionOps): + """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D Parameter): + one triton launch per rank-local shard, values unchanged. The swizzle covers whole 128-row + blocks and 4-column groups, so the reverse reads the affine grid back off the 5-D shape.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) + return out + + @staticmethod + def _unswizzle(value): + if value.ndim != 5: + return value + experts, row_blocks, col_groups = value.shape[:3] + kernel = load_finegrained_kernel() + return kernel.unswizzle_mx_scales(value, row_blocks * 128, col_groups * 4, num_experts=experts) + + @staticmethod + def _swizzle(key, value, model, full_layer_name): + _, held = _held_scale(model, full_layer_name, key) + if held is None or held.ndim != 5 or value.ndim == 5: + return value + return load_finegrained_kernel().swizzle_mx_scales(value) + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedSwizzleScales(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedPackedBlocks(ConversionOps): + """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 as the packed ``(E, N, K/2)`` int8 + the kernels read: the same bytes regrouped per row, folded back into 16-byte groups on the way + out. Two low-nibble-first E2M1 values per byte.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, target_patterns=None, **kwargs): + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) + out[key] = value + elif self.inverse: + out[key] = value.view(torch.uint8).reshape(*value.shape[:-1], value.shape[-1] // 16, 16) + else: + out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) + return out + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedPackedBlocks(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedViewPackedInt8(ConversionOps): + """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store: + ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 + tensors pass through, so this can ride converters that also match unquantized modules.""" + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + def convert(self, input_dict, target_patterns=None, **kwargs): + held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) + out = {} + for key, value in _keyed_by_target(input_dict, target_patterns).items(): + value = value[0] if isinstance(value, list) else value + out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value + return out + + @property + def reverse_op(self) -> ConversionOps: + # a bitcast is its own inverse; on save the int8 the module holds goes back out as the + # uint8 bytes the checkpoint packs fp4 into + return FineGrainedViewPackedInt8(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedWeightGlobals(ConversionOps): + """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: + one fp32 global per expert per projection, from either the per-expert checkpoint keys + (``experts.*.gate_proj.weight_scale_2``) or one stacked ``(E, 2)`` tensor per layer. + + A gate|up stack has two per expert and merging them is not local, so one op owns a layer's + whole set: the stack keeps the gate's and the up half's leaves as a ratio. ``inter = + silu(gate) * up`` is linear in the up half, so that only scales the expert output by + ``1 / ratio`` — which the down projection's weight global takes back, and its calibrated input + global keeps the requant on the range the e4m3 block scales were chosen for. Rescaling the up + half's block scales instead would re-round them against codes chosen for the old global.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, target_patterns=None, model=None, **kwargs): + sources = defaultdict(dict) + for key, value in input_dict.items(): + # the loader hands a converter's tensors in a list; unwrapped, a fused `(E, 2)` pair + # reads as `(1, 2E)` and the per-half fold below silently does not fire + sources[_global_role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value + globals_ = {} + for target in target_patterns: + if role_sources := sources.get(_global_role(target)): + # one (E, calibrated projections) fp32 tensor, gate column first: two sources are + # the stack's halves, one is the fused (E, 2) pair or a single projection's (E,) + columns = [] + for key, value in sorted(role_sources.items(), key=lambda kv: 0 if "gate_proj" in kv[0] else 1): + value = torch.stack(value, dim=0) if isinstance(value, list) else value + columns.append(value.float().reshape(value.shape[0], -1)) + globals_[target] = torch.cat(columns, dim=1) + targets = {_global_role(target): target for target in globals_} + stack = targets.get(("gate_up", "weight")) + + # a stack calibrated per half keeps the gate's global; the up half's leaves as a ratio, + # which the down's weight global takes on (scaling the expert output back up) and its + # input global gives back (keeping the requantized intermediate on the calibrated range) + if stack is not None and globals_[stack].shape[1] == 2: + self._assert_no_gate_up_bias(model) + gate, up = globals_[stack][:, 0], globals_[stack][:, 1] + globals_[stack] = gate + ratio = (up / gate)[:, None] + for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): + if target := targets.get(role): + globals_[target] = globals_[target] * factor + return {target: value.reshape(-1).contiguous() for target, value in globals_.items()} + + @staticmethod + def _assert_no_gate_up_bias(model) -> None: + """A gate|up bias is added AFTER the global, so merging the halves would have to divide + its up rows by the same ratio — in the bias converter, which the layout ops own. No + checkpoint pairs the two, so refuse rather than carry a silent half-correction.""" + if model is not None and any( + isinstance(module, FineGrainedExperts) and module.has_bias for module in model.modules() + ): + raise NotImplementedError( + "this checkpoint calibrates the gate|up halves separately AND carries an expert " + "bias; the up half's bias would have to be folded with the globals." + ) + + +def _global_role(key: str) -> tuple[str, str]: + """What a global-scale key holds, for a checkpoint key and a module target alike: which + projection it belongs to, and which of modelopt's two levels it is — the weight's second-level + global (``weight_scale_2``) or the calibrated activation one (``input_scale``).""" + projection = "down" if "down_proj" in key or "w2" in key else "gate_up" + level = "input" if "input_scale" in key or "input_global" in key else "weight" + return projection, level + + +class FineGrainedInputScales(ConversionOps): + """A calibrated checkpoint's ``input_scale`` in the layout the module holds: one value per + quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max since + both halves read the same rows. The NVFP4 gate_up global collapses to ONE value — its rows are + the hidden states, quantized once before routing. A static activation scale IS the + quantization scale, with no block level to absorb an inflated one, so it stays per-expert.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, full_layer_name=None, **kwargs): + values = [] + for value in input_dict.values(): + value = torch.stack(value, dim=0) if isinstance(value, list) else value + values.append(value.float()) + stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) + # A scale is a MAGNITUDE (`amax(|x|) / 448`), so reduce over absolute values and keep it + # off zero. A signed reduce returns whatever the largest signed entry is, and the two + # degenerate results it admits are both fatal downstream: a negative scale flips the sign + # of every dequantized row, and a zero one divides the activations by zero — NaN logits + # on a SINGLE device, before any sharding is involved. Real calibrated checkpoints ship + # positive scales, so this is a no-op for them; it is the models whose scales are built + # rather than measured that reach here with a sign or a zero. + per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1).clamp(min=1e-12) + one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only + return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScalesSplit(self.hf_quantizer) + + +class FineGrainedInputScalesSplit(ConversionOps): + """Save reverse of :class:`FineGrainedInputScales`: the module's scale back onto every + projection key the checkpoint calibrated one for. The merge took a max, so the folded halves + are gone and each projection is written the value that covers it. A collapsed global + re-expands per expert, the shape the checkpoint holds.""" + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + value = next(iter(input_dict.values())) + value = value[0] if isinstance(value, list) else value + if value.numel() == 1 and model is not None: + experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) + if experts is not None: + value = value.reshape(1).expand(experts.num_experts).contiguous() + return dict.fromkeys(target_patterns or list(input_dict), value) + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScales(self.hf_quantizer) + + +class FineGrainedQuantize(ConversionOps): + """Quantize a full-precision weight on load into the format its module holds, emitting the + scale (and the NVFP4 global) in the module's layout — the held container dtype, the swizzled + artifact where the module holds one. Block-FP8 is computed in torch (per-block E4M3, fp32 or + power-of-two UE8M0 inverse scales); the group formats run the kernels' row-wise quantizers + (``mxfp8_act_quant`` / ``mxfp4_act_quant`` / ``nvfp4_act_quant``), NVFP4 after normalizing by the + canonical per-tensor / per-expert global ``amax / (6 * 448)``. + Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) + pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: + result: dict[str, torch.Tensor] = {} + for key, value in input_dict.items(): + tensor = value[0] if isinstance(value, list) else value + result.update(self._quantize_one(key, tensor, self._weight_holder(model, key))) + return result + + @staticmethod + def _as_expert_rows(module, name: str, value: torch.Tensor) -> torch.Tensor: + """An expert stack with its contraction axis LAST, swapped to the ``(E, rows, K)`` the + module's slot and the kernels take. A model may store its experts transposed (GPT-OSS: + ``(E, H, 2I)``), and quantizing before the swap packs the wrong axis — the row count and + the byte-halved K then disagree, which the grouped op rejects at its first launch.""" + if module is None or value.ndim != 3 or name not in ("gate_up_proj", "up_proj", "down_proj"): + return value + if name == "down_proj": + rows, in_dim = module.hidden_dim, module.intermediate_dim + else: + rows, in_dim = (2 if module.has_gate else 1) * module.intermediate_dim, module.hidden_dim + # only when the orientation is unambiguous: a square stack reads the same both ways + return value.transpose(1, 2).contiguous() if value.shape[1:] == (in_dim, rows) != (rows, in_dim) else value + + @staticmethod + def _weight_holder(model, key: str): + """``(module, scale_param_name)`` when ``key`` is a finegrained module's weight, else ``None``.""" + if model is None: + return None + module, name = get_module_from_name(model, key) + if not isinstance(module, _FineGrainedModule) or getattr(module, name, None) is None: + return None + if name == "weight": + return module, "weight_scale_inv" + if name in ("gate_up_proj", "up_proj", "down_proj"): + return module, f"{name}_scale_inv" + return None + + def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torch.Tensor]: + if value.ndim < 2: + return {key: value} + prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key + module = holder[0] if holder is not None else None + value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) + held_scale = getattr(module, holder[1]) if holder is not None else None + format_spec = weight_formats()[module.weight_format] if module is not None else None + if format_spec is None or format_spec.scale_group is None: + if module is not None and module.block_size: + block = tuple(module.block_size) + else: + # no module (direct invocation): take the block from the config, and failing that + # put one scale over the whole matrix + config = self.hf_quantizer.quantization_config if self.hf_quantizer is not None else None + config_block = ( + config.get("weight_block_size") + if isinstance(config, dict) + else getattr(config, "weight_block_size", None) + ) + block = tuple(config_block) if config_block else (value.shape[-2], value.shape[-1]) + ue8m0 = ( + held_scale.dtype == _get_ue8m0_dtype() + if held_scale is not None + else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" + ) + quantized = self._quantize_block_fp8(value, block, ue8m0) + return {key: value} if quantized is None else {key: quantized[0], f"{prefix}_scale_inv": quantized[1]} + + if value.device.type not in ("cuda", "xpu"): + # the kernels' quantizers are triton launches; the package builds for cuda, rocm + # (which torch reports as "cuda") and xpu + raise ValueError( + f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on an " + f"accelerator, but the weight is on {value.device.type}" + ) + weight, scale, global_scale = self._quantize_group(value, format_spec) + scale = _recontain(scale, held_scale.dtype) + if held_scale.ndim == 5: + scale = load_finegrained_kernel().swizzle_mx_scales(scale) + out = {key: weight, f"{prefix}_scale_inv": scale} + if global_scale is not None: + # the experts name their global `_weight_global_scale`, a dense linear just + # `weight_global_scale` — both sit beside a `_scale_inv`, so ask for the slot + stem = holder[1].removesuffix("_scale_inv") + suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" + out[f"{prefix}{suffix}"] = global_scale + # activations are quantized dynamically against the block scales alone here — there is no + # calibration pass — so the module's activation globals are the identity + held_input = getattr(module, holder[1].replace("_scale_inv", "_input_global_scale"), None) + if held_input is not None: + # on the weight's device, not the slot's: the module is still on meta here, and + # `ones_like` would inherit that and leave the parameter unmaterialized + out[f"{prefix}_input_global_scale"] = torch.ones( + held_input.shape, dtype=held_input.dtype, device=value.device + ) + return out + + @staticmethod + def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): + """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or + ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two + before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" + block_m, block_n = block + rows, cols = value.shape[-2], value.shape[-1] + if rows % block_m or cols % block_n: + return None + tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) + max_abs = tiles.abs().amax(dim=(-3, -1)) + inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) + if ue8m0: + inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) + scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) + return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) + + @staticmethod + def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): + """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the + kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has + one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the + smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" + kernel = load_finegrained_kernel() + rows = value.reshape(-1, value.shape[-1]) + global_scale = None + if format_spec.global_scale_dtype is not None: + per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) + global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) + rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) + packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) + global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale + else: + quantize = kernel.mxfp4_act_quant if format_spec.values_per_byte == 2 else kernel.mxfp8_act_quant + packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) + return ( + packed.view(format_spec.weight_dtype).reshape(*value.shape[:-1], -1), + scale.reshape(*value.shape[:-1], -1), + global_scale, + ) + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedDequantize(self.hf_quantizer) + + +class FineGrainedDequantize(ConversionOps): + """A quantized weight folded back to full precision against its per-block scale grid. + + Runs FIRST in its converter under ``dequantize=True``, which is why + :meth:`update_weight_conversions` attaches it to each of the model's own converters: the + merge / concat ops after it collapse the per-expert structure the (weight, scale) pairing + needs. It pairs each weight pattern with its sibling scale pattern by index and emits the + result under the weight key, dropping the scales so the rest of the chain sees weights only. + """ + + def __init__(self, hf_quantizer): + self.hf_quantizer = hf_quantizer + + def _scale_pattern_for(self, weight_pattern: str) -> str: + anchored = weight_pattern.endswith("$") + base = weight_pattern[:-1] if anchored else weight_pattern + if base.endswith("_blocks"): + # GPT-OSS packs its experts as `{proj}_blocks` + `{proj}_scales` + scale = base[: -len("_blocks")] + "_scales" + elif base.endswith(".weight"): + scale = base[: -len(".weight")] + ".weight_scale_inv" + elif base == "weight": + scale = "weight_scale_inv" + else: + scale = base + "_scale_inv" + return scale + "$" if anchored else scale + + # packed-FP4 experts arrive as int8 / float4_e2m1fn_x2, two e2m1 nibbles per byte + _FP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: + """Two ``e2m1`` FP4 values per byte → float32 tensor twice as wide on the last dim.""" + lut = torch.tensor(self._FP4_E2M1_LUT, dtype=torch.float32, device=packed.device) + u8 = packed.contiguous().view(torch.uint8) + low = (u8 & 0xF).long() + high = ((u8 >> 4) & 0xF).long() + unpacked = torch.stack([lut[low], lut[high]], dim=-1) + return unpacked.reshape(*packed.shape[:-1], 2 * packed.shape[-1]) + + def _dequantize_one( + self, quantized: torch.Tensor, scales: torch.Tensor, output_dtype: torch.dtype | None = None + ) -> torch.Tensor: + # unpacked first, so the rest of the routine sees a normal (rows, cols) float matrix + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if quantized.dtype == torch.int8 or (fp4_dtype is not None and quantized.dtype == fp4_dtype): + quantized_fp32 = self._unpack_fp4(quantized) + else: + quantized_fp32 = quantized.to(torch.float32) + rows, cols = quantized_fp32.shape[-2:] + # the block comes from the scale grid, not the config: one checkpoint holds both a + # ``[1, 32]`` MXFP4 expert block and a ``[128, 128]`` FP8 dense one + try: + scale_rows, scale_cols = scales.shape[-2:] + except Exception: + # scale can be a single tensor in extreme cases where it was not wrapped properly but is [1,0]. + scale_rows, scale_cols = 1, 1 + if rows % scale_rows or cols % scale_cols: + raise ValueError( + f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." + ) + block_m = rows // scale_rows + block_n = cols // scale_cols + # the math runs in fp32 either way (``float8_e8m0fnu`` has no ``mul`` kernel); the + # destination parameter's dtype wins when known, so an eager module keeps the model's + if output_dtype is None: + output_dtype = ( + scales.dtype if scales.dtype.is_floating_point and scales.element_size() >= 2 else torch.bfloat16 + ) + # an MXFP8 checkpoint's uint8 scale is a biased E8M0 exponent, not a multiplier + if scales.dtype == torch.uint8: + s_fp32 = (scales.to(torch.float32) - 127.0).exp2() + else: + s_fp32 = scales.to(torch.float32) + original_shape = quantized_fp32.shape + q = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) + s = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) + return (q * s).to(output_dtype).reshape(original_shape) + + def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: + if model is None or full_layer_name is None: + return None + module, tensor_name = get_module_from_name(model, full_layer_name) + param = getattr(module, tensor_name, None) + return getattr(param, "dtype", None) + + def convert( + self, + input_dict: dict[str, list[torch.Tensor] | torch.Tensor], + full_layer_name: str | None = None, + model: torch.nn.Module | None = None, + **kwargs, + ) -> dict[str, list[torch.Tensor] | torch.Tensor]: + output_dtype = self._get_target_dtype(model, full_layer_name) + # The dense converter (``["weight$", "weight_scale_inv", "activation_scale"] -> weight``) + # hands one weight, with or without a scale (RMSNorm weights match ``weight$`` too). + if "weight$" in input_dict: + # the loader derives prefix/suffix from the output KEY; without a `full_layer_name` + # (direct invocation) key it by the converter's target + target_key = full_layer_name if full_layer_name is not None else "weight" + quantized = input_dict["weight$"] + quantized = quantized[0] if isinstance(quantized, list) else quantized + if "weight_scale_inv" in input_dict: + scales = input_dict["weight_scale_inv"] + scales = scales[0] if isinstance(scales, list) else scales + return {target_key: self._dequantize_one(quantized, scales, output_dtype=output_dtype)} + return {target_key: quantized} + + # Generic chain path: dequantize every weight pattern that has a sibling scale. + consumed = {self._scale_pattern_for(key) for key in input_dict} + result: dict[str, list[torch.Tensor] | torch.Tensor] = {} + for key, value in input_dict.items(): + if "activation_scale" in key or key in consumed: + continue # consumed by the dequant; drop from the chain + scale_key = self._scale_pattern_for(key) + if scale_key not in input_dict: + result[key] = value + continue + weights = value if isinstance(value, list) else [value] + scales = input_dict[scale_key] + scales = scales if isinstance(scales, list) else [scales] + if len(weights) != len(scales): + raise ValueError( + f"FineGrainedDequantize: weight/scale count mismatch for {key} " + f"({len(weights)} weights vs {len(scales)} scales)." + ) + result[key] = [self._dequantize_one(w, s, output_dtype=output_dtype) for w, s in zip(weights, scales)] + return result + + @property + def reverse_op(self) -> ConversionOps: + # a save re-quantizes, so the checkpoint keeps its format whether the in-memory + # state stayed quantized or was dequantized for compute + return FineGrainedQuantize(self.hf_quantizer) + + +def keep_swizzle_reverse_for_save(model, hf_quantizer) -> None: + """Keep the swizzle's reverse reachable at save time for a model quantized ON THE FLY. + + `_weight_conversions` retains only converters that matched a checkpoint weight, and + quantizing a bf16 checkpoint CREATES the scale keys — so the converter carrying the + layout ops never matched, and saving would write the module's 5-D SWIZZLE_32_4_4 grid + where the checkpoint format is the affine one (silently: the module reads back whatever + it wrote, so only a re-save catches it). + + Adds a converter that renames nothing and carries the swizzle alone, for scales no + retained converter already covers — reversing it is the unswizzle, which reads the + affine grid off the self-describing 5-D shape. Scoped that way because a converter + appended here is matched BEFORE the retained ones (the save reverses the list), so a + broader pattern would take keys away from the real converter and skip the rest of its + chain. + """ + from ..core_model_loading import WeightConverter + + conversions = list(getattr(model, "_weight_conversions", None) or []) + covered = { + target + for conv in conversions + if isinstance(conv, WeightConverter) + for target in (getattr(conv, "target_patterns", None) or []) + } + held_swizzled = { + name.rsplit(".", 1)[-1] + for name, param in model.named_parameters() + if param.ndim == 5 and name.rsplit(".", 1)[-1].endswith("_scale_inv") + } + missing = sorted(name for name in held_swizzled if not any(name in target for target in covered)) + if not missing: + return + for name in missing: + conv = WeightConverter( + source_patterns=rf"{name}$", target_patterns=name, operations=[FineGrainedSwizzleScales(hf_quantizer)] + ) + conv._was_used = True # it describes a layout this quantizer applied, not one a checkpoint carried + conversions.append(conv) + model._weight_conversions = conversions diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index ac1dafc83808..084b17733b16 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -40,12 +40,6 @@ logger = logging.get_logger(__name__) -logger.warning( - "finegrained_fp8 is frozen for backward compatibility and no longer receives new recipes; the " - "fine-grained quantization machinery lives in transformers.integrations.finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." -) - _FP8_DTYPE = torch.float8_e4m3fn _FP8_MIN = torch.finfo(_FP8_DTYPE).min diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py index 67c74eda798c..567684b72b74 100644 --- a/src/transformers/integrations/tensor_parallel.py +++ b/src/transformers/integrations/tensor_parallel.py @@ -63,30 +63,7 @@ def shard_and_distribute_module(*args, **kwargs): raise RuntimeError("`shard_and_distribute_module` is unavailable with the DTensor tensor-parallel loading path.") -def to_local(t): - """Deprecated: unwrap a `DTensor` to its local shard, or pass through. - - Kept for callers that predate `_hf_quantized_needs_local_tp`. A quantized module that sets - it is handed local tensors by the TP loader, so the kernels — which take raw pointers and do - not understand `DTensor` — need no unwrap of their own; nothing in transformers calls this. - `DTensor.to_local()` is the direct replacement and is autograd-aware on the train path, - rewrapping the gradient to match each parameter's placements. - """ - warnings.warn( - "`to_local` is deprecated. Set `_hf_quantized_needs_local_tp = True` on the module so the " - "tensor-parallel loader hands it local tensors, or call `DTensor.to_local()` directly.", - FutureWarning, - stacklevel=2, - ) - from ..distributed.utils import is_dtensor - - if is_dtensor(t): - return t.to_local() - return t - - __all__ = [ - "to_local", "ALL_PARALLEL_STYLES", "MoeIdentityExpertParallel", "MoeTensorParalellExperts", diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index c60dd9c9567b..49f8f127fb01 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -161,11 +161,11 @@ def _process_model_before_weight_loading( from ..integrations.finegrained import replace_with_finegrained_embedding, replace_with_finegrained_layer self._normalize_modules_to_not_convert(model) - if self._quant_method() == "mxfp4" and getattr(self.quantization_config, "activation_format", None) is None: + if self._quant_method() == "mxfp4" and self.quantization_config.activation_format is None: # GPT-OSS MXFP4 runs weight-only (W4A16): raw bf16 activations against packed # weights. The kernels' weight-native default would quantize activations to mxfp4. self.quantization_config.activation_format = "bf16" - if getattr(self.quantization_config, "activation_format", None) == "bf16": + if self.quantization_config.activation_format == "bf16": # Weight-only holds no activation global, so a calibrated checkpoint's `input_scale` # has no slot to load into. Not an unexpected key — one this run has no use for, which # a later W4A4 run of the same checkpoint reads. @@ -185,50 +185,14 @@ def _process_model_before_weight_loading( ) def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import disable_deepgemm_on_multi_device, keep_swizzle_reverse_for_save + from ..integrations.finegrained import disable_deepgemm_on_multi_device + from ..integrations.finegrained_conversions import keep_swizzle_reverse_for_save keep_swizzle_reverse_for_save(model, self) disable_deepgemm_on_multi_device(model) return model - @staticmethod - def _expert_shard(dim: int, packed: bool = False) -> str: - """Register a param-only style that shards an expert tensor on `dim`, and return the name - a plan carries it under. A style name is the only way a plan can express a placement, and - the built-in styles locate their axis by counting from the END of the shape — which lands - inside a swizzled scale grid's inner tile rather than on its row blocks, and on the expert - axis of a bias.""" - from torch.distributed.tensor import Shard - from torch.distributed.tensor.placement_types import _StridedShard - - from ..distributed.tensor_parallel import ALL_PARALLEL_STYLES, MoEParamShard - - placement = _StridedShard(dim=dim, split_factor=2) if packed else Shard(dim) - name = f"moe_experts_{'packed_' if packed else ''}shard{dim}" - ALL_PARALLEL_STYLES.register(name, MoEParamShard(placement)) - return name - def update_tp_plan(self, config): - if "Qwen3" in config.__class__.__name__: - text_plan = { - "layers.*.self_attn.q_proj.weight": "colwise", - "layers.*.self_attn.q_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.k_proj.weight": "colwise", - "layers.*.self_attn.k_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.v_proj.weight": "colwise", - "layers.*.self_attn.v_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.o_proj.weight": "rowwise", - "layers.*.self_attn.o_proj.weight_scale_inv": "rowwise", - "layers.*.mlp.gate_proj.weight": "colwise", - "layers.*.mlp.gate_proj.weight_scale_inv": "colwise", - "layers.*.mlp.up_proj.weight": "colwise", - "layers.*.mlp.up_proj.weight_scale_inv": "colwise", - "layers.*.mlp.down_proj.weight": "rowwise", - "layers.*.mlp.down_proj.weight_scale_inv": "rowwise", - } - - config.base_model_tp_plan = text_plan - # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes on # top of any plan written above (e.g. the Qwen3 dense plan). Models carry the experts # mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite both. @@ -276,23 +240,18 @@ def update_tp_plan(self, config): # Intra-expert TP: the experts stay whole and the projection's own axis splits, so # the per-expert globals, activation scales and the down bias (added after the - # row-reduce) stay replicated — but the scale grid still has to follow its weight. - # dim 1 is the output axis and dim 2 the input axis in BOTH the affine and the - # 5-D swizzled scale container, which is why naming those two suffices. A - # per-tensor scale has no such grid — one value per expert, `(E, 1, 1)`, whose - # axes are both 1 — so it joins the replicated companions and only its bias splits. + # row-reduce) stay replicated — only the scale grid follows its weight. A + # per-tensor scale is `(E, 1, 1)`: no grid to split, so it stays replicated too. blocked = self.quantization_config.weight_block_size is not None if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): if blocked: - updated_plan.setdefault(f"{key}_scale_inv", self._expert_shard(2)) + updated_plan.setdefault(f"{key}_scale_inv", "moe_experts_rowwise") elif projection != "down_proj" and style in ("packed_colwise", "colwise"): - # The companions take the WEIGHT's own split. The model declares which one it - # needs and is right: the interleave into the kernels' row order runs on each - # rank's shard, AFTER sharding, so the split has to match the layout the - # CHECKPOINT holds. `packed_colwise` is the strided split a concatenated - # `[gate; up]` keeps pairs together with; `colwise` the contiguous one a - # natively interleaved layout (GPT-OSS, `is_concatenated=False`) needs. - rows = self._expert_shard(1, packed=style == "packed_colwise") + # Companions take the WEIGHT's split, which is the CHECKPOINT's layout: the + # interleave into the kernels' row order runs on each rank's shard afterwards. + # `packed_colwise` keeps a concatenated `[gate; up]`'s pairs together; + # `colwise` suits an already-interleaved layout (GPT-OSS). + rows = "moe_experts_packed_colwise" if style == "packed_colwise" else "moe_experts_colwise" if blocked: updated_plan.setdefault(f"{key}_scale_inv", rows) updated_plan.setdefault(f"{key}_bias", rows) @@ -314,13 +273,13 @@ def is_compileable(self) -> bool: return True def get_quantize_ops(self): - from ..integrations.finegrained import FineGrainedQuantize + from ..integrations.finegrained_conversions import FineGrainedQuantize return FineGrainedQuantize(self) def _quant_method(self) -> str: - method = getattr(self.quantization_config, "quant_method", None) - return getattr(method, "value", method) + method = self.quantization_config.quant_method + return getattr(method, "value", method) # a QuantizationMethod enum, or already its str def get_weight_conversions(self): """The converters this CHECKPOINT's own key layout needs, on top of the model's plan.""" @@ -341,7 +300,7 @@ def _dequantize_conversions(self): layout: the packed `{proj}_blocks` + `{proj}_scales` pair GPT-OSS ships, and the `weight` + `weight_scale_inv` pair everything else does.""" from ..core_model_loading import Transpose, WeightConverter - from ..integrations.finegrained import FineGrainedDequantize, FineGrainedPackedBlocks + from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks if self._quant_method() == "mxfp4": # the blocks regroup into the packed rows and the exponent-byte scales come back out @@ -370,7 +329,7 @@ def _mxfp4_conversions(self): weight, the exponent bytes take the scale container op; then the usual layout ops. Bare patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`.""" from ..core_model_loading import WeightConverter - from ..integrations.finegrained import ( + from ..integrations.finegrained_conversions import ( FineGrainedInterleaveGateUp, FineGrainedPackedBlocks, FineGrainedScaleContainer, @@ -400,36 +359,34 @@ def _nvfp4_conversions(self): stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and the other never fires. Routed experts only — a generic `weight_scale*` rename would run before converter collection and mangle these keys.""" - from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter - from ..integrations.finegrained import ( - FineGrainedInputScales, - FineGrainedScaleContainer, - FineGrainedViewPackedInt8, - FineGrainedWeightGlobals, - ) - # A weight-only run keeps activations bf16, so its experts hold no activation global and # the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys to the # unexpected-key filter rather than converting them onto a module slot that does not exist. - calibrated = getattr(self.quantization_config, "activation_format", None) != "bf16" + calibrated = self.quantization_config.activation_format != "bf16" + return self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) - # MergeModulelist is what stamps each source with its expert index, which is what lets - # expert parallelism keep only this rank's experts (see `tensor_idx` in - # core_model_loading). Without it every rank collects all E values and the forward - # asserts on the per-expert count. - per_expert = [ + def _nvfp4_per_expert_conversions(self, calibrated: bool): + """One key per expert per projection, as GLM-5.2-NVFP4 ships them. + + `MergeModulelist` stamps each source with its expert index, which is what lets expert + parallelism keep only this rank's experts (`tensor_idx` in `core_model_loading`); without + it every rank collects all E values and the forward asserts on the per-expert count. + """ + from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter + from ..integrations.finegrained_conversions import FineGrainedInputScales, FineGrainedWeightGlobals + + merge = [MergeModulelist(dim=0)] + expert = r"mlp\.experts\..*\." + converters = [ WeightConverter( - source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.weight_scale$", - r"mlp\.experts\..*\.up_proj\.weight_scale$", - ], + source_patterns=[rf"{expert}gate_proj\.weight_scale$", rf"{expert}up_proj\.weight_scale$"], target_patterns="mlp.experts.gate_up_proj_scale_inv", - operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + operations=[*merge, Concatenate(dim=1)], ), WeightConverter( - source_patterns=r"mlp\.experts\..*\.down_proj\.weight_scale$", + source_patterns=rf"{expert}down_proj\.weight_scale$", target_patterns="mlp.experts.down_proj_scale_inv", - operations=[MergeModulelist(dim=0)], + operations=merge, ), # every second-level global of a layer in ONE converter: the gate|up stack's two # calibrated halves merge into one global per expert by folding the up half's onto @@ -437,31 +394,42 @@ def _nvfp4_conversions(self): # decided together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.weight_scale_2", - r"mlp\.experts\..*\.up_proj\.weight_scale_2", - r"mlp\.experts\..*\.down_proj\.weight_scale_2", - *([r"mlp\.experts\..*\.down_proj\.input_scale"] if calibrated else []), + rf"{expert}gate_proj\.weight_scale_2", + rf"{expert}up_proj\.weight_scale_2", + rf"{expert}down_proj\.weight_scale_2", + *([rf"{expert}down_proj\.input_scale"] if calibrated else []), ], target_patterns=[ "mlp.experts.gate_up_proj_weight_global_scale", "mlp.experts.down_proj_weight_global_scale", *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), ], - operations=[MergeModulelist(dim=0), FineGrainedWeightGlobals(self)], + operations=[*merge, FineGrainedWeightGlobals(self)], ), ] if calibrated: - per_expert.append( + converters.append( WeightConverter( - source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.input_scale", - r"mlp\.experts\..*\.up_proj\.input_scale", - ], + source_patterns=[rf"{expert}gate_proj\.input_scale", rf"{expert}up_proj\.input_scale"], target_patterns="mlp.experts.gate_up_proj_input_global_scale", - operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], + operations=[*merge, FineGrainedInputScales(self)], ) ) - fused = [ + return converters + + def _nvfp4_fused_conversions(self, calibrated: bool): + """Already stacked per layer, as the vLLM fused layout and Muse-Spark ship them — so no + `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the + `.weight`-anchored rule gives the per-expert layout.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained_conversions import ( + FineGrainedInputScales, + FineGrainedScaleContainer, + FineGrainedViewPackedInt8, + FineGrainedWeightGlobals, + ) + + converters = [ # the container cast the module's scale dtype needs; `_with_expert_layout_ops` # adds the interleave and the swizzle on top (it re-adds the cast, which is # idempotent) @@ -472,9 +440,8 @@ def _nvfp4_conversions(self): ) for proj in ("gate_up_proj", "down_proj") ] - fused += [ - # the same one-converter rule as the per-expert layout above, over the keys the - # fused layout ships + converters.append( + # the same one-converter rule as the per-expert layout, over the fused keys WeightConverter( source_patterns=[ r"experts\.gate_up_proj_weight_scale_2$", @@ -487,19 +454,17 @@ def _nvfp4_conversions(self): *(["experts.down_proj_input_global_scale"] if calibrated else []), ], operations=[FineGrainedWeightGlobals(self)], - ), - ] + ) + ) if calibrated: - fused.append( + converters.append( WeightConverter( source_patterns=r"experts\.gate_up_proj_input_scale$", target_patterns="experts.gate_up_proj_input_global_scale", operations=[FineGrainedInputScales(self)], ) ) - # the fused expert WEIGHTS take the packed uint8 -> int8 view the `.weight`-anchored - # rule below gives the per-expert layout; the layout ops ride on top of both - fused += [ + converters += [ WeightConverter( source_patterns=rf"experts\.{proj}$", target_patterns=f"experts.{proj}", @@ -507,7 +472,7 @@ def _nvfp4_conversions(self): ) for proj in ("gate_up_proj", "down_proj") ] - return per_expert + fused + return converters def _static_activation_conversions(self): """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in. One @@ -518,7 +483,7 @@ def _static_activation_conversions(self): other never fires. NVFP4 consumes ``input_scale`` as its activation GLOBAL instead (:meth:`_nvfp4_conversions`) — a second level over a block scale, not the scale itself.""" from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming - from ..integrations.finegrained import FineGrainedInputScales + from ..integrations.finegrained_conversions import FineGrainedInputScales per_expert = [ WeightConverter( @@ -571,7 +536,7 @@ def update_weight_conversions(self, weight_conversions): :meth:`get_weight_conversions` is appended in every mode.""" from ..core_model_loading import WeightConverter, WeightRenaming - from ..integrations.finegrained import FineGrainedDequantize, FineGrainedViewPackedInt8 + from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedViewPackedInt8 scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") weight_conversions = [scale_rename, *weight_conversions] @@ -630,7 +595,7 @@ def _with_expert_layout_ops(self, weight_conversions): the fused names (a plain rename, a transformers-format checkpoint) and dense linears' scale keys, so they get the same treatment.""" from ..core_model_loading import WeightConverter - from ..integrations.finegrained import ( + from ..integrations.finegrained_conversions import ( FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales, @@ -659,20 +624,15 @@ def layout_ops(target: str) -> list: # `base_model_prefix` and `force_cpu` are set on it before this hook runs and # are not `__init__` arguments, so a fresh one silently loses them. conv.operations = list(conv.operations) + layout_ops(expert_targets[0]) - # A name counts as covered only when this converter can SOURCE the fused - # transformers-format key, not merely target it. A checkpoint that already ships - # fused (our own save, and any transformers-format one) names the scale - # `experts.gate_up_proj_scale_inv`, which the modelopt-style sources here - # (`...gate_proj.weight_scale`) never match — so keying on the target let those - # converters suppress the catch-all below and the scale arrived as a plain rename, - # carrying none of the layout ops. The module still ALLOCATES the swizzled grid, - # so the affine tensor lands in a 5-D slot: invisible on one device (the parameter - # is simply replaced and the kernels read the shape they are given) and a - # malformed DTensor under sharding, whose local is affine while its global says - # swizzled. + # Covered only when this converter can SOURCE the fused transformers-format key, + # not merely target it: modelopt-style sources (`...gate_proj.weight_scale`) never + # match an already-fused `experts.gate_up_proj_scale_inv`, so keying on the target + # let them suppress the catch-all below and the scale arrived as a plain rename + # with none of the layout ops — an affine tensor in the swizzled 5-D slot the + # module allocated, which is a malformed DTensor under sharding. fused_probe = "model.layers.0.mlp.experts.{}" for target in expert_targets: - fused = re.sub(r".*experts\.|\$$", "", target) + fused = target.rpartition("experts.")[2].removesuffix("$") if any(re.search(pattern, fused_probe.format(fused)) for pattern in (conv.source_patterns or [])): covered.add(fused) updated.append(conv) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 40d98c730f2e..70653a5b759f 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -1,4 +1,3 @@ -import warnings from typing import TYPE_CHECKING from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging @@ -9,15 +8,6 @@ if is_torch_available(): import torch - -warnings.warn( - "quantizer_finegrained_fp8 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - if TYPE_CHECKING: from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import FineGrainedFP8Config @@ -36,6 +26,11 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) + logger.warning_once( + "`FineGrainedFP8Config` is frozen for backward compatibility and receives no new " + "recipes. `FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, " + "weight-only)." + ) def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 109af54715fa..42036ff4416f 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -23,22 +23,6 @@ def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: return module, tensor_name -def subtree_pattern_to_regex(pattern: str) -> str: - """A glob-style skip entry as the bounded regex `should_convert_module` matches. - - Checkpoints name whole subtrees with a trailing star ("model.layers.0*", "model.layers.1.*") - or bare ("self_attn", "model.layers.0."). Those are globs, not regexes, and handing them over - as-is is wrong twice: the dots match any character, and the star is greedy, so - "model.layers.1.*" also swallows layers 10-19. The bounded form matches the named module and - everything under it, nothing else. - """ - prefix = pattern[:-2] if pattern.endswith(".*") else pattern.rstrip("*").rstrip(".") - # anchored on both sides at a path boundary: a bare entry names that module wherever it sits - # ("self_attn" reaches "model.layers.2.self_attn.q_proj") without "layers.1" reaching - # "layers.16" or "attn" reaching "self_attn" - return r"(?:^|.*\.)" + re.escape(prefix) + r"(\..*)?$" - - def should_convert_module(full_name, patterns: list[str] | None = None): if patterns is None: return True diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 588a269340b4..169bec0af60c 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -18,6 +18,7 @@ import importlib.metadata import json import os +import re from dataclasses import dataclass from enum import Enum from typing import Any, Optional, Union @@ -1762,14 +1763,21 @@ def __init__( # "modelopt" names the PRODUCER, not a format, so the algo it exported becomes the # format here and nothing downstream has to ask where a config came from. self.quant_method, algo_activation_format = _MODELOPT_ALGOS[quant_algo] + logger.info( + f"modelopt checkpoint exported with quant_algo={quant_algo!r}; loading it as " + f"{str(self.quant_method)!r}." + ) self.activation_format = activation_format or algo_activation_format ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) if modules_to_not_convert is None and ignore is not None: - # modelopt ships glob-style subtree entries, under `ignore` as "model.layers.0*" / - # "model.layers.1.*" and under `exclude_modules` bare ("self_attn", "layers.0.") - from ..quantizers.quantizers_utils import subtree_pattern_to_regex - - modules_to_not_convert = [subtree_pattern_to_regex(g) for g in ignore] + # modelopt names skipped subtrees with GLOBS — "model.layers.0*", "model.layers.1.*", + # or bare "self_attn". Handed on as regexes they are wrong twice: the dots match any + # character, and the star is greedy, so "model.layers.1.*" also takes layers 10-19. + # Anchor each at a path boundary instead: the named module and everything under it. + modules_to_not_convert = [] + for glob in ignore: + subtree = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") + modules_to_not_convert.append(r"(?:^|.*\.)" + re.escape(subtree) + r"(\..*)?$") self.modules_to_not_convert = modules_to_not_convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme @@ -1836,6 +1844,7 @@ def __init__( if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") self.modules_to_not_convert = modules_to_not_convert + # TODO: check overlap with not to convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme self.weight_block_size = weight_block_size diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index c73360eaa15e..7e8fc08b6197 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -21,6 +21,7 @@ import subprocess import tempfile import unittest +import warnings from contextlib import ExitStack, contextmanager from dataclasses import dataclass from unittest import mock @@ -309,7 +310,7 @@ def test_nvfp4_experts_thread_per_expert_globals(self): m = self._experts(has_gate=True, weight_format="nvfp4") # the format table must resolve the ATTRIBUTE the forwards gate on, not just the # param allocation — a None here silently drops the global at every forward - self.assertIs(m.has_global_scale, True) + self.assertIsNotNone(m.global_scale_dtype) p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_batched_mm_experts_forward(m, *self._route()) @@ -516,10 +517,19 @@ def test_frozen_module_warns_and_is_self_contained(self): import importlib import transformers.integrations.finegrained_fp8 as frozen + from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + from transformers.utils.quantization_config import FineGrainedFP8Config - # the notice rides the library logger, so it honours transformers' verbosity settings - with self.assertLogs("transformers.integrations.finegrained_fp8", level="WARNING") as captured: + # the notice fires when someone actually asks for the frozen scheme, not on import: + # importing is what every `transformers.integrations` consumer does + with warnings.catch_warnings(record=True) as raised: + warnings.simplefilter("always") importlib.reload(frozen) + self.assertEqual([w for w in raised if "frozen" in str(w.message)], []) + + # ...and it rides the library logger, so it honours transformers' verbosity settings + with self.assertLogs("transformers.quantizers.quantizer_finegrained_fp8", level="WARNING") as captured: + FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) self.assertTrue(any("frozen for backward compatibility" in line for line in captured.output)) # distinct machinery: the frozen classes are not the live ones self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) @@ -630,10 +640,10 @@ def test_intra_expert_entries_follow_the_projection_axis(self): # order runs on each rank's shard, after sharding. The companions follow it. self.assertEqual(style("gate_up_proj"), "packed_colwise") for name in ("gate_up_proj_scale_inv", "gate_up_proj_bias"): - self.assertEqual(style(name), "moe_experts_packed_shard1", name) + self.assertEqual(style(name), "moe_experts_packed_colwise", name) # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike self.assertEqual(style("down_proj"), "rowwise") - self.assertEqual(style("down_proj_scale_inv"), "moe_experts_shard2") + self.assertEqual(style("down_proj_scale_inv"), "moe_experts_rowwise") # TP does not split experts, so anything indexed per expert stays whole for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): self.assertEqual(style(name), "moe_tp_experts", name) @@ -644,7 +654,7 @@ def test_intra_expert_entries_follow_the_projection_axis(self): self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", plan_i), "colwise") self.assertEqual( _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), - "moe_experts_shard1", + "moe_experts_colwise", ) def test_dense_projections_that_share_a_name_are_left_alone(self): @@ -679,7 +689,7 @@ def test_dense_projections_that_share_a_name_are_left_alone(self): # the stacked experts still get theirs self.assertEqual( _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), - "moe_experts_packed_shard1", + "moe_experts_packed_colwise", ) def test_a_multimodal_models_expert_plans_are_reached(self): @@ -746,14 +756,64 @@ def test_the_impl_rewrites_the_layer_kinds(self): self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") self.assertEqual(megamoe["layers.*.mlp.experts"], "megamoe_experts") + def test_a_models_own_plan_survives_the_update(self): + """`update_tp_plan` ADDS companions; it must never drop what the model ships. + + The config here is a REAL one, which the `SimpleNamespace` the tests above use cannot + substitute for: a branch keyed on the config CLASS is invisible to a fake. A hardcoded + `"Qwen3" in config.__class__.__name__` override used to assign a dense plan over + `base_model_tp_plan`, dropping `q_norm`/`k_norm` on Qwen3 and — because the check is a + substring — the whole expert plan on Qwen3-MoE. + """ + from transformers import Qwen3Config, Qwen3MoeConfig, Qwen3VLMoeConfig + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for config_cls in (Qwen3Config, Qwen3MoeConfig): + with self.subTest(config=config_cls.__name__): + config = config_cls() + shipped = dict(config.base_model_tp_plan) + self.assertTrue(shipped, f"{config_cls.__name__} ships no plan to preserve") + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + for key, style in shipped.items(): + self.assertEqual(config.base_model_tp_plan.get(key), style, key) + + # and a multimodal one keeps its plan on the SUB-config: the outer carries none, so + # nothing may be invented there. 25 config classes contain "Qwen3" — a name-substring + # branch reaches every one of them, including these. + multimodal = Qwen3VLMoeConfig() + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(multimodal) + self.assertFalse(getattr(multimodal, "base_model_tp_plan", None) or {}) + self.assertTrue(multimodal.text_config.base_model_tp_plan) + + def test_a_real_moe_models_experts_gain_their_companions(self): + """The other half of the same invariant: preserving the model's entries is what lets the + companion pass find `moe_tp_experts` and give the expert scales an entry of their own.""" + from transformers import Qwen3MoeConfig + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + config = Qwen3MoeConfig() + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + plan = config.base_model_tp_plan + self.assertEqual(plan["layers.*.mlp.experts"], "moe_tp_experts") + for companion in ("gate_up_proj_scale_inv", "down_proj_scale_inv", "gate_up_proj_bias"): + self.assertIn(f"layers.*.mlp.experts.{companion}", plan, companion) + class SubtreePatternTest(unittest.TestCase): def test_a_checkpoints_glob_skip_list_is_bounded(self): """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" - from transformers.quantizers.quantizers_utils import should_convert_module, subtree_pattern_to_regex + from transformers.quantizers.quantizers_utils import should_convert_module + from transformers.utils.quantization_config import FineGrainedConfig - skip = [subtree_pattern_to_regex(g) for g in ("model.layers.0*", "model.layers.1.*", "self_attn")] + # as a modelopt checkpoint ships them, through the config that normalizes them + skip = FineGrainedConfig( + quant_method="modelopt", + quant_algo="NVFP4", + ignore=["model.layers.0*", "model.layers.1.*", "self_attn"], + ).modules_to_not_convert for name in ("model.layers.0.mlp", "model.layers.1.mlp.down_proj", "model.layers.2.self_attn.q_proj"): self.assertFalse(should_convert_module(name, skip), name) for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): @@ -778,7 +838,10 @@ def _reference_dequant(self, blocks, scales): return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): - from transformers.integrations.finegrained import FineGrainedPackedBlocks, FineGrainedScaleContainer + from transformers.integrations.finegrained_conversions import ( + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + ) torch.manual_seed(0) E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows @@ -949,7 +1012,7 @@ def _op_kernel(self): return kernel def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): - from transformers.integrations.finegrained import FineGrainedSwizzleScales + from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales model, experts = self._experts("mxfp8") kernel = self._op_kernel() @@ -992,7 +1055,7 @@ def test_scale_container_op_round_trips_every_container(self): the op brings both into the held e8m0 (exact cast / same bytes), records the container on the module, and its reverse restores that container on save; weights and native scales pass both ways.""" - from transformers.integrations.finegrained import FineGrainedScaleContainer + from transformers.integrations.finegrained_conversions import FineGrainedScaleContainer op = FineGrainedScaleContainer(hf_quantizer=None) native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) @@ -1024,7 +1087,7 @@ def test_scale_container_op_round_trips_every_container(self): self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): - from transformers.integrations.finegrained import FineGrainedSwizzleScales + from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales kernel = self._op_kernel() reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op @@ -1067,7 +1130,7 @@ def _arch_converters(self): def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained import ( + from transformers.integrations.finegrained_conversions import ( FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales, @@ -1097,7 +1160,10 @@ def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): def test_blocks_scales_converters_carry_the_layout_ops(self): from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained import FineGrainedInterleaveGateUp, FineGrainedSwizzleScales + from transformers.integrations.finegrained_conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedSwizzleScales, + ) converters = self._quantizer("mxfp4")([]) by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} @@ -1178,7 +1244,7 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself and keeps the stack; every other case gets the interleaved order, and the reverse restores the stack.""" - from transformers.integrations.finegrained import FineGrainedInterleaveGateUp + from transformers.integrations.finegrained_conversions import FineGrainedInterleaveGateUp stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] op = FineGrainedInterleaveGateUp(hf_quantizer=None) @@ -1204,7 +1270,7 @@ def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation the unquantized experts hold — so it is compared against that integration directly.""" from transformers.core_model_loading import Transpose - from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedPackedBlocks + from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks from transformers.integrations.mxfp4 import convert_moe_packed_tensors torch.manual_seed(0) @@ -1221,7 +1287,7 @@ def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) def test_block_fp8_round_trips_within_its_floor(self): - from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedQuantize + from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1415,7 +1481,7 @@ def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): every expert against the largest one's range. It stays per expert.""" import torch - from transformers.integrations.finegrained import FineGrainedInputScales + from transformers.integrations.finegrained_conversions import FineGrainedInputScales sources = { "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), @@ -1431,7 +1497,7 @@ def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): where the kernels assert one per expert.""" import torch - from transformers.integrations.finegrained import FineGrainedWeightGlobals + from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1461,7 +1527,7 @@ def test_one_converter_merges_a_layer_s_globals(self): global and the down's weight global scales the expert output back. The down's calibrated input scale moves the other way, keeping the requantized intermediate on the range the checkpoint calibrated.""" - from transformers.integrations.finegrained import FineGrainedWeightGlobals + from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) @@ -1487,7 +1553,7 @@ def test_one_converter_merges_a_layer_s_globals(self): def test_globals_with_one_calibrated_projection_pass_through(self): """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a single global) has nothing to merge, so every global reaches its module unchanged.""" - from transformers.integrations.finegrained import FineGrainedWeightGlobals + from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1509,7 +1575,7 @@ def test_globals_with_one_calibrated_projection_pass_through(self): def test_the_fused_layout_ships_both_halves_in_one_tensor(self): """Muse-Spark's `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two keys, and merges the same way.""" - from transformers.integrations.finegrained import FineGrainedWeightGlobals + from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1529,7 +1595,7 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated `input_scale` reduces to one value; the down's rows are per expert, so its stays per expert — the requant epilogue normalizes each row by its own expert's value.""" - from transformers.integrations.finegrained import FineGrainedInputScales + from transformers.integrations.finegrained_conversions import FineGrainedInputScales up = FineGrainedInputScales().convert( { @@ -1648,7 +1714,7 @@ def test_on_the_fly_group_formats_round_trip_within_their_floor(self): """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in the module's layout (swizzled where held so); dequantizing them back lands within each format's floor.""" - from transformers.integrations.finegrained import FineGrainedDequantize, FineGrainedQuantize + from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1689,7 +1755,7 @@ def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes globals. Returns the module, the exact dequantized weights (so the reference is the quantization floor rather than the pre-quant weight), and the dequantization a reader that applied ONE of the two globals to both halves would get.""" - from transformers.integrations.finegrained import FineGrainedWeightGlobals + from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals kernel = load_finegrained_kernel() experts = FineGrainedExperts(cfg, weight_format="nvfp4", activation_format=activation_format).cuda() @@ -1828,7 +1894,7 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op fills them from the affine grid, the fused forward matches the affine module's, and the reverse op hands the affine grid back bitwise.""" - from transformers.integrations.finegrained import FineGrainedSwizzleScales + from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales cfg = _Cfg() cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = ( From e83eadf5cb7c64631f42ba44ec315b85934c6605 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 21 Sep 2026 16:06:07 -0700 Subject: [PATCH 36/68] Collapse the conversion-op boilerplate; shorten the longest units MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ten conversion ops were five self-inverting (reverse is themselves with the direction flipped) and five paired (reverse names a different class), and every one of them spelled out its own `__init__` and `reverse_op`. A single `_FineGrainedOp` base carries the quantizer and defaults `reverse_op` to `type(self)(..., inverse=not self.inverse)`; the paired ops override it and ignore the flag. Twelve lines replace fifty-two, and `type(self)` cannot go stale the way five hardcoded class names could. - `_moe_operands`: the twelve projection entries were `{kernel_name: getattr(module, held_name)}` twelve times over. One comprehension states the rule instead — the kernels always say `gate_up_proj`, the module says `up_proj` where the model has no gate. - `FineGrainedExperts.__init__` (104 lines) was reading config and allocating parameters in one breath. Allocation moves to `_register_projection`, called once per projection, with the four `_set_optional_parameter` calls folded into a loop and two loop-invariants hoisted out. - `_with_expert_layout_ops`: a six-line comment explaining a condition becomes `sources_the_fused_key`, a named predicate carrying that reasoning; the loop body goes from fourteen lines to three. The expert-target regex compiles once. - the `expert_dtype` comment now says why the side-channel exists: DeepSeek-V4 is MIXED (mxfp4 experts under an "fp8" quant config) and a quant config carries a single `quant_method` with no way to express it. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 140 +++++++++--------- .../integrations/finegrained_conversions.py | 87 ++++------- .../quantizers/quantizer_finegrained.py | 42 +++--- 3 files changed, 121 insertions(+), 148 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index efcb0460ed65..04c1fa42298a 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -561,42 +561,33 @@ def _moe_operands(kernel, module) -> dict: else: act_fn = module._apply_gate if module.has_gate else module.act_fn - norm_eps, norm_weight = 1e-6, None - if not module.has_post_expert_norm: - post_expert_norm = None - elif module.post_expert_norm_name in kernel.get_supported_norms() and not getattr( - module.post_expert_norm, "_hf_tp_input_reduce", False - ): - # fusing puts the norm INSIDE the launch that produced the rows, so a collective that has - # to land between the two has nowhere to go: under intra-expert TP those rows are a - # partial sum over the sharded intermediate. Fall back to the module, whose wrapped - # forward does the all-reduce first. - post_expert_norm = module.post_expert_norm_name - norm = module.post_expert_norm - norm_weight = norm.weight - # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves it - # None and resolves to the dtype's own - eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) - norm_eps = float(eps) if eps is not None else torch.finfo(norm_weight.dtype).eps - else: - post_expert_norm = module._apply_post_norm + # a `get_supported_norms()` NAME is folded into the routing-weighted reduce; any other + # callable runs on the routed rows before it. The name is only usable where nothing has to + # happen in between: fusing puts the norm INSIDE the launch that produced the rows, so under + # intra-expert TP — where they are a partial sum over the sharded intermediate — the + # all-reduce would have nowhere to land, and the module's wrapped forward must do it first. + post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 + if module.has_post_expert_norm: + fusable = module.post_expert_norm_name in kernel.get_supported_norms() + if not fusable or getattr(module.post_expert_norm, "_hf_tp_input_reduce", False): + post_expert_norm = module._apply_post_norm + else: + norm = module.post_expert_norm + post_expert_norm, norm_weight = module.post_expert_norm_name, norm.weight + # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves + # it None and resolves to the dtype's own + eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) + norm_eps = float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps # The kernels always say `gate_up_proj`; the module names that weight `up_proj` when the model # has no gate, which is the only reason these are not plain attribute reads. A slot the module # does not hold reads back as None — the argument the kernels take for an absent one. return { - "gate_up_proj": getattr(module, up), - "gate_up_proj_scale_inv": getattr(module, f"{up}_scale_inv"), - "gate_up_proj_weight_global_scale": getattr(module, f"{up}_weight_global_scale"), - "gate_up_proj_input_global_scale": getattr(module, f"{up}_input_global_scale"), - "gate_up_proj_activation_scale": getattr(module, f"{up}_activation_scale"), - "gate_up_proj_bias": getattr(module, f"{up}_bias"), - "down_proj": module.down_proj, - "down_proj_scale_inv": module.down_proj_scale_inv, - "down_proj_weight_global_scale": module.down_proj_weight_global_scale, - "down_proj_input_global_scale": module.down_proj_input_global_scale, - "down_proj_activation_scale": module.down_proj_activation_scale, - "down_proj_bias": module.down_proj_bias, + **{ + f"{kernel_proj}{slot}": getattr(module, f"{held_proj}{slot}") + for kernel_proj, held_proj in (("gate_up_proj", up), ("down_proj", "down_proj")) + for slot in ("", "_scale_inv", "_weight_global_scale", "_input_global_scale", "_activation_scale", "_bias") + }, # a supported activation NAME is fused into the gate_up epilogue; any other callable # leaves that GEMM plain and runs on the host between the two "act_fn": act_fn, @@ -671,10 +662,11 @@ def __init__( self.swiglu_limit = getattr(config, "swiglu_limit", None) self.act_fn = ACT2FN[self.act_fn_name] - # Expert weight storage is declared by the QUANT config's format key, not the model - # config: `weight_format` arrives from `replace_with_finegrained_layer` as the - # checkpoint's quant_method. `config.expert_dtype` is the legacy DeepSeek-V4 side-channel, - # kept as the fallback when no format is passed. + # `weight_format` is the checkpoint's quant_method, as `replace_with_finegrained_layer` + # passes it. DeepSeek-V4 is MIXED though — mxfp4 experts under an "fp8" quant config — + # and declares that on the MODEL config (`expert_dtype`), because a quant config carries + # one `quant_method` and has no way to say "the experts are a different format". That + # side-channel is the fallback here when no format is passed. if weight_format is None: weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" self.weight_format = weight_format @@ -695,43 +687,19 @@ def __init__( up_name = "gate_up_proj" if self.has_gate else "up_proj" up_rows = (2 if self.has_gate else 1) * self.intermediate_dim - for proj, rows, in_dim, min_scale_out in ( - (up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1), - ("down_proj", self.hidden_dim, self.intermediate_dim, 1), + # gate_up's activation global is ONE value — its rows are the pre-routing hidden states, + # quantized once before routing; down's rows belong to an expert each. + for proj, rows, in_dim, min_scale_out, input_globals in ( + (up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, 1), + ("down_proj", self.hidden_dim, self.intermediate_dim, 1, self.num_experts), ): - weight, scale = _alloc_expert_proj( - self.num_experts, rows, in_dim, format_spec, scale_dtype, scale_group, min_scale_out, swizzled - ) - self.register_parameter(proj, weight) - self.register_parameter(f"{proj}_scale_inv", scale) - # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it - # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype - _set_optional_parameter( - self, f"{proj}_bias", torch.empty(self.num_experts, rows) if self.has_bias else None - ) - # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator, one per - # expert for the weight (`FineGrainedWeightGlobals` merges a separately calibrated - # gate|up stack down to that). The activation's is the checkpoint's `input_scale` — - # one value for gate_up, whose rows are the pre-routing hidden states, and one per - # expert for down, whose rows are per expert. - _set_optional_parameter( - self, - f"{proj}_weight_global_scale", - torch.ones(self.num_experts, dtype=self.global_scale_dtype) - if self.global_scale_dtype is not None - else None, - ) - calibrated = self.global_scale_dtype is not None and self.activation_format != "bf16" - _set_optional_parameter( - self, - f"{proj}_input_global_scale", - torch.ones(1 if proj == up_name else self.num_experts, dtype=torch.float32) if calibrated else None, - ) - static = self.activation_scheme == "static" - _set_optional_parameter( - self, - f"{proj}_activation_scale", - torch.ones(self.num_experts, dtype=torch.float32) if static else None, + self._register_projection( + proj, + rows, + in_dim, + min_scale_out, + input_globals, + (format_spec, scale_dtype, scale_group, swizzled), ) # The model's per-expert output norm, carried over by the swap. A submodule slot, not a @@ -743,6 +711,36 @@ def __init__( # (`get_supported_norms`); unnamed norms still run through `_apply_post_norm`. self.post_expert_norm_name = None + def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, storage): + """One expert projection: the packed weight and its scale grid, then the slots a + checkpoint may or may not fill — a bias, the NVFP4 second-level globals, a calibrated + activation scale. `storage` is the layout `resolve_weight_format` settled on.""" + format_spec, scale_dtype, scale_group, swizzled = storage + weight, scale = _alloc_expert_proj( + self.num_experts, rows, in_dim, format_spec, scale_dtype, scale_group, min_scale_out, swizzled + ) + self.register_parameter(proj, weight) + self.register_parameter(f"{proj}_scale_inv", scale) + + # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it + # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype + bias = torch.empty(self.num_experts, rows) if self.has_bias else None + # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator, one per expert + # for the weight (`FineGrainedWeightGlobals` merges a separately calibrated gate|up stack + # down to that). The activation's is the checkpoint's `input_scale`. + two_level = self.global_scale_dtype is not None + calibrated = two_level and self.activation_format != "bf16" + for name, tensor in ( + ("bias", bias), + ( + "weight_global_scale", + torch.ones(self.num_experts, dtype=self.global_scale_dtype) if two_level else None, + ), + ("input_global_scale", torch.ones(input_globals, dtype=torch.float32) if calibrated else None), + ("activation_scale", torch.ones(self.num_experts) if self.activation_scheme == "static" else None), + ): + _set_optional_parameter(self, f"{proj}_{name}", tensor) + def _apply_post_norm(self, expert_out: torch.Tensor) -> torch.Tensor: """The standard hook every backend applies: the model's norm on one expert application's rows, before the routing weights.""" diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained_conversions.py index 626ca478877f..0debee5e2cb6 100644 --- a/src/transformers/integrations/finegrained_conversions.py +++ b/src/transformers/integrations/finegrained_conversions.py @@ -68,15 +68,29 @@ def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) -class FineGrainedInterleaveGateUp(ConversionOps): - """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order — the - core ``Interleave`` along dim 1, over the weight, scale grid and bias alike. Runs on load and - save wherever the checkpoint's row order differs from the one the experts hold.""" +class _FineGrainedOp(ConversionOps): + """Base for the ops below. Each carries the quantizer that built it, which is how an op reaches + the config (the activation format, the block size) while it runs. + + The default reverse is the op ITSELF with `inverse` flipped: it applies the layout the modules + hold on load and puts the checkpoint's back on save. The pairs that reverse into a *different* + class — merge/split, quantize/dequantize — override `reverse_op` and ignore the flag. + """ def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer self.inverse = inverse + @property + def reverse_op(self) -> ConversionOps: + return type(self)(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedInterleaveGateUp(_FineGrainedOp): + """Stacked ``[gate; up]`` expert rows into the kernels' ``[g0, u0, g1, u1, ...]`` order — the + core ``Interleave`` along dim 1, over the weight, scale grid and bias alike. Runs on load and + save wherever the checkpoint's row order differs from the one the experts hold.""" + def convert(self, input_dict, model=None, target_patterns=None, **kwargs): input_dict = _keyed_by_target(input_dict, target_patterns) experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None @@ -94,20 +108,12 @@ def convert(self, input_dict, model=None, target_patterns=None, **kwargs): out[key] = interleave.convert({key: value}, None, [key])[key] return out - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedInterleaveGateUp(self.hf_quantizer, inverse=not self.inverse) - -class FineGrainedScaleContainer(ConversionOps): +class FineGrainedScaleContainer(_FineGrainedOp): """A block scale into the dtype its module holds. UE8M0 ships either as the exponent byte under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast exactly); the module records which the checkpoint used so the reverse restores it on save.""" - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): out = {} for key, value in _keyed_by_target(input_dict, target_patterns).items(): @@ -127,20 +133,12 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= out[key] = value return out - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedScaleContainer(self.hf_quantizer, inverse=not self.inverse) - -class FineGrainedSwizzleScales(ConversionOps): +class FineGrainedSwizzleScales(_FineGrainedOp): """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D Parameter): one triton launch per rank-local shard, values unchanged. The swizzle covers whole 128-row blocks and 4-column groups, so the reverse reads the affine grid back off the 5-D shape.""" - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): out = {} for key, value in _keyed_by_target(input_dict, target_patterns).items(): @@ -163,20 +161,12 @@ def _swizzle(key, value, model, full_layer_name): return value return load_finegrained_kernel().swizzle_mx_scales(value) - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedSwizzleScales(self.hf_quantizer, inverse=not self.inverse) - -class FineGrainedPackedBlocks(ConversionOps): +class FineGrainedPackedBlocks(_FineGrainedOp): """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 as the packed ``(E, N, K/2)`` int8 the kernels read: the same bytes regrouped per row, folded back into 16-byte groups on the way out. Two low-nibble-first E2M1 values per byte.""" - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - def convert(self, input_dict, target_patterns=None, **kwargs): out = {} for key, value in _keyed_by_target(input_dict, target_patterns).items(): @@ -189,20 +179,12 @@ def convert(self, input_dict, target_patterns=None, **kwargs): out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) return out - @property - def reverse_op(self) -> ConversionOps: - return FineGrainedPackedBlocks(self.hf_quantizer, inverse=not self.inverse) - -class FineGrainedViewPackedInt8(ConversionOps): +class FineGrainedViewPackedInt8(_FineGrainedOp): """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store: ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 tensors pass through, so this can ride converters that also match unquantized modules.""" - def __init__(self, hf_quantizer=None, inverse: bool = False): - self.hf_quantizer = hf_quantizer - self.inverse = inverse - def convert(self, input_dict, target_patterns=None, **kwargs): held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) out = {} @@ -211,14 +193,8 @@ def convert(self, input_dict, target_patterns=None, **kwargs): out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value return out - @property - def reverse_op(self) -> ConversionOps: - # a bitcast is its own inverse; on save the int8 the module holds goes back out as the - # uint8 bytes the checkpoint packs fp4 into - return FineGrainedViewPackedInt8(self.hf_quantizer, inverse=not self.inverse) - -class FineGrainedWeightGlobals(ConversionOps): +class FineGrainedWeightGlobals(_FineGrainedOp): """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: one fp32 global per expert per projection, from either the per-expert checkpoint keys (``experts.*.gate_proj.weight_scale_2``) or one stacked ``(E, 2)`` tensor per layer. @@ -230,9 +206,6 @@ class FineGrainedWeightGlobals(ConversionOps): global keeps the requant on the range the e4m3 block scales were chosen for. Rescaling the up half's block scales instead would re-round them against codes chosen for the old global.""" - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - def convert(self, input_dict, target_patterns=None, model=None, **kwargs): sources = defaultdict(dict) for key, value in input_dict.items(): @@ -288,16 +261,13 @@ def _global_role(key: str) -> tuple[str, str]: return projection, level -class FineGrainedInputScales(ConversionOps): +class FineGrainedInputScales(_FineGrainedOp): """A calibrated checkpoint's ``input_scale`` in the layout the module holds: one value per quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max since both halves read the same rows. The NVFP4 gate_up global collapses to ONE value — its rows are the hidden states, quantized once before routing. A static activation scale IS the quantization scale, with no block level to absorb an inflated one, so it stays per-expert.""" - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - def convert(self, input_dict, full_layer_name=None, **kwargs): values = [] for value in input_dict.values(): @@ -320,15 +290,12 @@ def reverse_op(self) -> ConversionOps: return FineGrainedInputScalesSplit(self.hf_quantizer) -class FineGrainedInputScalesSplit(ConversionOps): +class FineGrainedInputScalesSplit(_FineGrainedOp): """Save reverse of :class:`FineGrainedInputScales`: the module's scale back onto every projection key the checkpoint calibrated one for. The merge took a max, so the folded halves are gone and each projection is written the value that covers it. A collapsed global re-expands per expert, the shape the checkpoint holds.""" - def __init__(self, hf_quantizer=None): - self.hf_quantizer = hf_quantizer - def convert(self, input_dict, model=None, target_patterns=None, **kwargs): value = next(iter(input_dict.values())) value = value[0] if isinstance(value, list) else value @@ -343,7 +310,7 @@ def reverse_op(self) -> ConversionOps: return FineGrainedInputScales(self.hf_quantizer) -class FineGrainedQuantize(ConversionOps): +class FineGrainedQuantize(_FineGrainedOp): """Quantize a full-precision weight on load into the format its module holds, emitting the scale (and the NVFP4 global) in the module's layout — the held container dtype, the swizzled artifact where the module holds one. Block-FP8 is computed in torch (per-block E4M3, fp32 or @@ -497,7 +464,7 @@ def reverse_op(self) -> ConversionOps: return FineGrainedDequantize(self.hf_quantizer) -class FineGrainedDequantize(ConversionOps): +class FineGrainedDequantize(_FineGrainedOp): """A quantized weight folded back to full precision against its per-block scale grid. Runs FIRST in its converter under ``dequantize=True``, which is why diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 49f8f127fb01..dc28302f68b8 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -560,11 +560,13 @@ def anchor(conversions): if self.pre_quantized and self.quantization_config.dequantize: updated = [] for conv in weight_conversions: - weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] - if isinstance(conv, WeightConverter) and weight_sources: + # collect the sibling `.weight_scale_inv` alongside each `.weight` so the pair folds + # to full precision BEFORE merge/concat ops collapse the per-expert structure + weights = [p for p in conv.source_patterns if p.endswith(".weight")] + if isinstance(conv, WeightConverter) and weights: conv = WeightConverter( - source_patterns=[p + "$" for p in weight_sources] - + [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] + source_patterns=[p + "$" for p in weights] + + [p.removesuffix(".weight") + ".weight_scale_inv$" for p in weights] + [p for p in conv.source_patterns if not p.endswith(".weight")], target_patterns=conv._original_target_patterns, operations=[FineGrainedDequantize(self), *conv.operations], @@ -575,6 +577,7 @@ def anchor(conversions): if self.pre_quantized and self._quant_method() == "nvfp4": updated = [] for conv in weight_conversions: + # anchored so the `.weight` source cannot swallow `weight_scale` / `weight_scale_2` if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): conv = WeightConverter( source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], @@ -612,28 +615,32 @@ def layout_ops(target: str) -> list: ] # the swizzle reads 1-byte scales return ops + def sources_the_fused_key(conv, fused: str) -> bool: + """Whether `conv` can SOURCE the fused transformers-format key, not merely target it. + + modelopt-style sources (`...gate_proj.weight_scale`) never match an already-fused + `experts.gate_up_proj_scale_inv`, so keying on the target would let them suppress the + catch-all below: the scale then arrives as a plain rename with none of the layout ops + — an affine tensor in the swizzled 5-D slot the module allocated, which is a malformed + DTensor under sharding. + """ + probe = f"model.layers.0.mlp.experts.{fused}" + return any(re.search(pattern, probe) for pattern in (conv.source_patterns or [])) + + expert_target = re.compile(r"experts\.(gate_up_proj|down_proj)(_scale_inv|_bias)?\$?$") updated = [] covered: set[str] = set() for conv in weight_conversions: targets = conv.target_patterns if isinstance(conv, WeightConverter) else [] - expert_targets = [ - t for t in targets if re.search(r"experts\.(gate_up_proj|down_proj)(_scale_inv|_bias)?\$?$", t) - ] + expert_targets = [t for t in targets if expert_target.search(t)] if expert_targets: # Append to the EXISTING converter rather than rebuild it: `scope_prefix`, # `base_model_prefix` and `force_cpu` are set on it before this hook runs and # are not `__init__` arguments, so a fresh one silently loses them. conv.operations = list(conv.operations) + layout_ops(expert_targets[0]) - # Covered only when this converter can SOURCE the fused transformers-format key, - # not merely target it: modelopt-style sources (`...gate_proj.weight_scale`) never - # match an already-fused `experts.gate_up_proj_scale_inv`, so keying on the target - # let them suppress the catch-all below and the scale arrived as a plain rename - # with none of the layout ops — an affine tensor in the swizzled 5-D slot the - # module allocated, which is a malformed DTensor under sharding. - fused_probe = "model.layers.0.mlp.experts.{}" for target in expert_targets: fused = target.rpartition("experts.")[2].removesuffix("$") - if any(re.search(pattern, fused_probe.format(fused)) for pattern in (conv.source_patterns or [])): + if sources_the_fused_key(conv, fused): covered.add(fused) updated.append(conv) # dense linears have no converter: their scale keys take the container cast alone @@ -648,13 +655,14 @@ def layout_ops(target: str) -> list: # only where no converter already produces that target: a second converter for the # same key is appended LATER and shadows the first, dropping whatever ops it carried # that these do not (the packed uint8 -> int8 view, on a fused modelopt checkpoint) - if name not in covered and layout_ops(name): + ops = layout_ops(name) + if name not in covered and ops: # the target is replacement text (op outputs are keyed by it), so no regex escapes updated.append( WeightConverter( source_patterns=rf"experts\.{name}$", target_patterns=f"experts.{name}", - operations=layout_ops(name), + operations=ops, ) ) return updated From 60f904e6182235d5660e44a934d7b5186b5f3ea6 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 22 Sep 2026 02:31:33 -0700 Subject: [PATCH 37/68] Key on capability not name; drop the last op boilerplate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two branches decided behaviour by matching a NAME, which is how the Qwen3 plan override stayed invisible — the default arm silently absorbs a miss: - the grouped-linear swap tested `"GroupedLinear" in type(module).__name__`. It now tests `hasattr(module, "n_groups")`, the attribute the branch consumes, so a grouped linear named anything else is still swapped and one named right but shaped wrong no longer reaches an AttributeError. ConvBERT's `GroupedLinearLayer` was never caught (it is an `nn.Module`), so this is hardening rather than a fix. - `_global_role` classified with `"w2" in key`, an unanchored substring for what is a path SEGMENT. Since the default is gate_up, a stray match merged a down global into the wrong bucket silently. Matched on segments now. Readability: - `FineGrainedQuantize` and `FineGrainedDequantize` still carried their own `__init__`; the base-class collapse only matched the form with a default. - `_weight_holder`'s pair was indexed as `holder[0]` / `holder[1]` in four places, and is now unpacked into `module, scale_name`. - `_recontain` -> `_as_container` and `_swizzles_scales` -> `_holds_swizzled_scales`: both were coined, and the codebase already says "container" (`scale_container_dtype`) and "holds" (`holds_interleaved_gate_up`). - `q` / `s` -> `blocked` / `per_block` in the dequant reshape. - each file opens with a short statement of what it is and where its neighbour lives, rather than none at all. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 18 ++++++--- .../integrations/finegrained_conversions.py | 40 +++++++++---------- .../quantizers/quantizer_finegrained.py | 9 +++-- tests/kernels/test_finegrained.py | 36 ++++++++++++++++- 4 files changed, 71 insertions(+), 32 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 04c1fa42298a..0f409829ada9 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -11,6 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Quantized modules and forwards for the fine-grained family — block-FP8, MXFP8, MXFP4, NVFP4 +and weight-only — served by `kernels-community/finegrained-kernels`. The checkpoint's layout is +moved onto these modules by `finegrained_conversions`. +""" + from __future__ import annotations import functools @@ -272,7 +277,7 @@ def _alloc_expert_proj( ) -def _swizzles_scales(config, format_spec: WeightFormat, activation_format: str | None) -> bool: +def _holds_swizzled_scales(config, format_spec: WeightFormat, activation_format: str | None) -> bool: """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05 scaled-MMA reads directly (row-major forces a per-tile gather). Needs SM100, a triton dispatch (the DeepGEMM backends read affine scales), a group-scaled format and a chain @@ -683,7 +688,7 @@ def __init__( # the tensors now hold, so re-deriving it would claim a layout they do not have. impl = getattr(config, "_experts_implementation", None) self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" - swizzled = _swizzles_scales(config, format_spec, activation_format) + swizzled = _holds_swizzled_scales(config, format_spec, activation_format) up_name = "gate_up_proj" if self.has_gate else "up_proj" up_rows = (2 if self.has_gate else 1) * self.intermediate_dim @@ -955,10 +960,11 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No activation_format=quantization_config.activation_format, **storage, ) - elif isinstance(module, nn.Linear) and "GroupedLinear" in type(module).__name__: - # Block-diagonal grouped linear (e.g. DSv4's `DeepseekV4GroupedLinear`): a - # plain `FineGrainedLinear` would collapse the groups into one giant linear and - # yield the wrong output dim. + elif isinstance(module, nn.Linear) and hasattr(module, "n_groups"): + # Block-diagonal grouped linear (DSv4's `DeepseekV4GroupedLinear`), recognised by + # the attribute the swap needs rather than by its class NAME: a plain + # `FineGrainedLinear` would collapse the groups into one giant linear and yield + # the wrong output dim. new_module = FineGrainedGroupedLinear( in_features_per_group=module.in_features, out_features=module.out_features, diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained_conversions.py index 0debee5e2cb6..8ce5ba63e23b 100644 --- a/src/transformers/integrations/finegrained_conversions.py +++ b/src/transformers/integrations/finegrained_conversions.py @@ -13,9 +13,9 @@ # limitations under the License. """Loader/saver conversion ops for the fine-grained family. -Each op moves ONE checkpoint-side layout to the one the `finegrained` modules hold, and carries -its own `reverse_op` so `save_pretrained` writes the checkpoint's layout back. They run inside the -`WeightConverter` chains the quantizer builds; nothing here touches a forward. +Each op moves ONE checkpoint-side layout to the one the `finegrained` modules hold, and carries a +`reverse_op` so `save_pretrained` writes the checkpoint's layout back. Nothing here runs in a +forward. """ from __future__ import annotations @@ -63,7 +63,7 @@ def _held_scale(model, full_layer_name, key): return module, getattr(module, param_name, None) -def _recontain(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: +def _as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) @@ -124,12 +124,12 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= modules = model.modules() if model is not None else () container = next((c for m in modules if (c := getattr(m, "scale_container_dtype", None))), None) if container is not None and value.dtype == _get_ue8m0_dtype(): - value = _recontain(value, container) + value = _as_container(value, container) else: module, held = _held_scale(model, full_layer_name, key) if held is not None and value.dtype != held.dtype: module.scale_container_dtype = value.dtype - value = _recontain(value, held.dtype) + value = _as_container(value, held.dtype) out[key] = value return out @@ -255,10 +255,12 @@ def _assert_no_gate_up_bias(model) -> None: def _global_role(key: str) -> tuple[str, str]: """What a global-scale key holds, for a checkpoint key and a module target alike: which projection it belongs to, and which of modelopt's two levels it is — the weight's second-level - global (``weight_scale_2``) or the calibrated activation one (``input_scale``).""" - projection = "down" if "down_proj" in key or "w2" in key else "gate_up" + global (``weight_scale_2``) or the calibrated activation one (``input_scale``). Matched on path + SEGMENTS, since the default is gate_up and a loose `"w2" in key` would silently bucket a down + global as one.""" + down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) level = "input" if "input_scale" in key or "input_global" in key else "weight" - return projection, level + return ("down" if down else "gate_up"), level class FineGrainedInputScales(_FineGrainedOp): @@ -320,8 +322,6 @@ class FineGrainedQuantize(_FineGrainedOp): Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: result: dict[str, torch.Tensor] = {} @@ -363,9 +363,9 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc if value.ndim < 2: return {key: value} prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key - module = holder[0] if holder is not None else None + module, scale_name = holder if holder is not None else (None, None) value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) - held_scale = getattr(module, holder[1]) if holder is not None else None + held_scale = getattr(module, scale_name) if module is not None else None format_spec = weight_formats()[module.weight_format] if module is not None else None if format_spec is None or format_spec.scale_group is None: if module is not None and module.block_size: @@ -396,19 +396,19 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc f"accelerator, but the weight is on {value.device.type}" ) weight, scale, global_scale = self._quantize_group(value, format_spec) - scale = _recontain(scale, held_scale.dtype) + scale = _as_container(scale, held_scale.dtype) if held_scale.ndim == 5: scale = load_finegrained_kernel().swizzle_mx_scales(scale) out = {key: weight, f"{prefix}_scale_inv": scale} if global_scale is not None: # the experts name their global `_weight_global_scale`, a dense linear just # `weight_global_scale` — both sit beside a `_scale_inv`, so ask for the slot - stem = holder[1].removesuffix("_scale_inv") + stem = scale_name.removesuffix("_scale_inv") suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" out[f"{prefix}{suffix}"] = global_scale # activations are quantized dynamically against the block scales alone here — there is no # calibration pass — so the module's activation globals are the identity - held_input = getattr(module, holder[1].replace("_scale_inv", "_input_global_scale"), None) + held_input = getattr(module, scale_name.replace("_scale_inv", "_input_global_scale"), None) if held_input is not None: # on the weight's device, not the slot's: the module is still on meta here, and # `ones_like` would inherit that and leave the parameter unmaterialized @@ -474,8 +474,6 @@ class FineGrainedDequantize(_FineGrainedOp): result under the weight key, dropping the scales so the rest of the chain sees weights only. """ - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer def _scale_pattern_for(self, weight_pattern: str) -> str: anchored = weight_pattern.endswith("$") @@ -538,9 +536,9 @@ def _dequantize_one( else: s_fp32 = scales.to(torch.float32) original_shape = quantized_fp32.shape - q = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) - s = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) - return (q * s).to(output_dtype).reshape(original_shape) + blocked = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) + per_block = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) + return (blocked * per_block).to(output_dtype).reshape(original_shape) def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: if model is None or full_layer_name is None: diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index dc28302f68b8..0a591dadc5cc 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -11,6 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""The `HfQuantizer` for the fine-grained family: converter chains, the parallel plan, and the +environment checks. The modules are in `integrations/finegrained`. +""" + import re from typing import TYPE_CHECKING @@ -359,9 +363,8 @@ def _nvfp4_conversions(self): stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and the other never fires. Routed experts only — a generic `weight_scale*` rename would run before converter collection and mangle these keys.""" - # A weight-only run keeps activations bf16, so its experts hold no activation global and - # the checkpoint's calibrated `input_scale` has nowhere to land: leave those keys to the - # unexpected-key filter rather than converting them onto a module slot that does not exist. + # weight-only keeps activations bf16, so there is no activation global for a calibrated + # `input_scale` to land on — leave those keys to the unexpected-key filter calibrated = self.quantization_config.activation_format != "bf16" return self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 7e8fc08b6197..4c0d3864f7b3 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -186,6 +186,38 @@ def test_module_forward_threads_everything(self): @require_torch +@require_torch +class FineGrainedGroupedLinearSwapTest(unittest.TestCase): + def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): + """`replace_with_finegrained_layer` picks `FineGrainedGroupedLinear` for a block-diagonal + linear — a plain one would collapse the groups into a single giant matmul and return the + wrong output dim. The branch keys on `n_groups`, the attribute the swap consumes: a class + NAME substring would miss a grouped linear named anything else, and would match (then + `AttributeError` on) one that carries its group count under another name. + """ + from types import SimpleNamespace + + import torch.nn as nn + + from transformers.integrations.finegrained import FineGrainedGroupedLinear, replace_with_finegrained_layer + from transformers.utils.quantization_config import FineGrainedConfig + + class OddlyNamedBlockDiagonal(nn.Linear): + def __init__(self): + super().__init__(8, 16, bias=False) + self.n_groups = 2 + + model = nn.Module() + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + model.grouped = OddlyNamedBlockDiagonal() + model.plain = nn.Linear(8, 16, bias=False) + + with torch.device("meta"): + replace_with_finegrained_layer(model, None, FineGrainedConfig(quant_method="fp8")) + self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) + self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) + + class FineGrainedEmbeddingTest(unittest.TestCase): """A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`.""" @@ -993,7 +1025,7 @@ def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block _, experts = self._experts("fp8") self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) - with mock.patch.object(fg, "_swizzles_scales", return_value=False): + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): _, experts = self._experts("mxfp8") self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) @@ -1904,7 +1936,7 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): ) # K >= 256: the swizzled arm has tiles cfg._experts_implementation = "grouped_mm" experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() - with mock.patch.object(fg, "_swizzles_scales", return_value=False): + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): affine_experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() model = torch.nn.Module() model.experts = experts From 4cb5d9cbe93030c1e1aeff4831182a68c8fddf56 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 22 Sep 2026 05:05:51 -0700 Subject: [PATCH 38/68] Quantize per module group; split the arms by producer, not by format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two axes that were tangled into one `quant_method`. GROUPS. A checkpoint can be more than one format — DeepSeek-V4 is W4A4 mxfp4 experts over block-FP8 linears — and a flat config cannot say so, which is why its expert format arrives on the MODEL config as `expert_dtype` and is read back out deep inside `FineGrainedExperts.__init__`. `FineGrainedGroup` names a format per module subset; `group_for(name)` resolves one; `replace_with_finegrained_layer` asks per module, so the side-channel read is gone. Nothing about this is DeepSeek-shaped. Verified across three models and three combinations, swap and in-flight quantization both: Qwen3-MoE with nvfp4 experts over mxfp8 dense, and GPT-OSS with nvfp4 attention, static block-FP8 experts and mxfp8 for the rest — three groups, the exotic format on the dense side, and differing activation SCHEMES per group. Backwards compatibility is normalisation, the same move `quant_algo` already gets: the flat fields become a single group, and a legacy `expert_dtype` becomes two (`split_out_experts`). Existing checkpoints load unchanged and nothing downstream learns which spelling it came from. `config_groups` is normalised too — GLM-5.2-NVFP4 ships it and we were dropping it, along with the `input_activations` and `dynamic: false` it carries. A group naming a format we do not serve falls back to the flat fields rather than guessing at a partial translation. Resolution is order-INDEPENDENT: `to_json_string` sorts keys, so a config that round-tripped through `config.json` has lost its declaration order. Targeted groups win over the catch-all, and two targeted groups claiming one module raises. ARMS. Split per PRODUCER key layout rather than per format, which is what actually differs: GPT-OSS's packed `_blocks`, modelopt's two-level scales, the calibrated `input_scale`. MXFP8 ships the plain `weight`/`weight_scale_inv` pair and gets no arm — the base is its handler. The base keeps no `_quant_method() == ...` test at all. Also: a finegrained module holding a full-precision weight now fails the load instead of silently computing in bf16 and handing back a model the caller believes is quantized. Checked once in the post-load pass, not per forward — the dtype cannot change between them. No shipping checkpoint relied on the old fallback: gpt-oss-120b lists attention in `modules_to_not_convert`, so those modules are never swapped. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 86 +++--- .../integrations/finegrained_conversions.py | 2 - src/transformers/quantizers/auto.py | 20 +- .../quantizers/quantizer_finegrained.py | 273 ++---------------- .../quantizer_finegrained_blockfp8.py | 82 ++++++ .../quantizers/quantizer_finegrained_mxfp4.py | 83 ++++++ .../quantizers/quantizer_finegrained_nvfp4.py | 183 ++++++++++++ src/transformers/utils/quantization_config.py | 138 ++++++++- tests/kernels/test_finegrained.py | 32 +- 9 files changed, 590 insertions(+), 309 deletions(-) create mode 100644 src/transformers/quantizers/quantizer_finegrained_blockfp8.py create mode 100644 src/transformers/quantizers/quantizer_finegrained_mxfp4.py create mode 100644 src/transformers/quantizers/quantizer_finegrained_nvfp4.py diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 0f409829ada9..28dcb4c19ac1 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -27,7 +27,6 @@ import torch import torch.nn as nn -from torch.nn import functional as F from ..activations import ACT2FN from ..quantizers.quantizers_utils import should_convert_module @@ -450,9 +449,6 @@ def __init__( _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) def forward(self, input: torch.Tensor) -> torch.Tensor: - if self.weight.element_size() > 1: - return F.linear(input, self.weight, self.bias) - return finegrained_linear( input, self.weight, @@ -517,15 +513,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: input_shape = x.shape[:-2] hidden_dim = x.shape[-1] - if self.weight.element_size() > 1: - w = self.weight.view(self.n_groups, -1, hidden_dim).transpose(1, 2) - x = x.reshape(-1, self.n_groups, hidden_dim).transpose(0, 1) - y = torch.bmm(x, w).transpose(0, 1) - y = y.reshape(*input_shape, self.n_groups, -1) - if self.has_bias: - y.add_(self.bias.view(self.n_groups, -1)) - return y - w = self.weight.view(self.n_groups, -1, hidden_dim) scale_inv = self.weight_scale_inv x = x.movedim(-2, 0).reshape(-1, hidden_dim) @@ -643,7 +630,7 @@ def __init__( scale_fmt: str = "float", has_bias: bool = False, has_gate: bool = True, - weight_format: str | None = None, + weight_format: str = "fp8", activation_format: str | None = None, is_concatenated: bool = True, ): @@ -667,13 +654,9 @@ def __init__( self.swiglu_limit = getattr(config, "swiglu_limit", None) self.act_fn = ACT2FN[self.act_fn_name] - # `weight_format` is the checkpoint's quant_method, as `replace_with_finegrained_layer` - # passes it. DeepSeek-V4 is MIXED though — mxfp4 experts under an "fp8" quant config — - # and declares that on the MODEL config (`expert_dtype`), because a quant config carries - # one `quant_method` and has no way to say "the experts are a different format". That - # side-channel is the fallback here when no format is passed. - if weight_format is None: - weight_format = "mxfp4" if getattr(config, "expert_dtype", "fp8") == "fp4" else "fp8" + # the format of THIS module's group, as `replace_with_finegrained_layer` resolved it. A + # mixed checkpoint (DeepSeek-V4: mxfp4 experts, block-FP8 linears) is two groups, so the + # experts no longer read a `config.expert_dtype` side-channel to find out. self.weight_format = weight_format format_spec, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) # the format decides the second level: None means there is none, else the dtype it is in @@ -813,8 +796,6 @@ def linear( weight = getattr(self, proj)[expert_idx] bias = getattr(self, f"{proj}_bias") bias = bias[expert_idx] if bias is not None else None - if weight.element_size() > 1: - return F.linear(input, weight, bias) scale = getattr(self, f"{proj}_scale_inv") weight_globals = getattr(self, f"{proj}_weight_global_scale") input_globals = getattr(self, f"{proj}_input_global_scale") @@ -851,6 +832,28 @@ class FineGrainedExpertsInterface(ExpertsInterface): ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() +def assert_modules_are_quantized(model: nn.Module) -> None: + """Every finegrained module must hold a quantized weight once the checkpoint is in. + + A full-precision weight in a quantized slot means the swap and the checkpoint disagree: the + module was allocated quantized storage and the loader replaced it with what the checkpoint + actually ships. Serving that in full precision would hand back a model the caller believes is + quantized and is not, so it is a load failure. Checked once here rather than per forward — + the dtype cannot change between them. + """ + for name, module in model.named_modules(): + if not isinstance(module, _FineGrainedModule): + continue + for attr in ("weight", "gate_up_proj", "up_proj", "down_proj"): + weight = getattr(module, attr, None) + if weight is not None and weight.element_size() > 1: + raise ValueError( + f"{name}.{attr} was converted for quantized compute but holds a {weight.dtype} " + "weight — this checkpoint does not quantize it. Name the module in the " + "quantization config's `modules_to_not_convert` to leave it as it is." + ) + + def disable_deepgemm_on_multi_device(model: nn.Module) -> None: """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` @@ -903,17 +906,18 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No if quantization_config.dequantize: return model - # The checkpoint's quant_method IS the weight format ("fp8", "mxfp8", "mxfp4", "nvfp4"). - # Under the "fp8" key pass None: dsv4-style checkpoints declare fp4 EXPERTS through the - # legacy `config.expert_dtype` model-config side-channel, which the ctor falls back to. - weight_format = next((f for f in ("mxfp8", "mxfp4", "nvfp4") if quantization_config.quant_method == f), None) - - # every swap reads the same three off the quant config; the branches differ only in shape - storage = { - "block_size": quantization_config.weight_block_size, - "activation_scheme": quantization_config.activation_scheme, - "scale_fmt": quantization_config.scale_fmt, - } + def storage_for(module_name: str) -> dict: + """The format THIS module is quantized in. A single-format checkpoint has one group + covering everything; DeepSeek-V4 is mxfp4 experts over block-FP8 linears, so its experts + and its attention projections resolve differently.""" + group = quantization_config.group_for(module_name) + return { + "weight_format": group.quant_method, + "activation_format": group.activation_format, + "block_size": group.weight_block_size, + "activation_scheme": group.activation_scheme, + "scale_fmt": group.scale_fmt, + } has_been_replaced = False for module_name, module in model.named_modules(): @@ -937,9 +941,7 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No ) new_module = new_class( config=getattr(module, "config", model.config.get_text_config()), - weight_format=weight_format, - activation_format=quantization_config.activation_format, - **storage, + **storage_for(module_name), **flags, ) # A per-expert output norm belongs to the model, not the quantization, so the @@ -956,9 +958,7 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No in_features=module.in_features, out_features=module.out_features, has_bias=module.bias is not None, - weight_format=weight_format or "fp8", - activation_format=quantization_config.activation_format, - **storage, + **storage_for(module_name), ) elif isinstance(module, nn.Linear) and hasattr(module, "n_groups"): # Block-diagonal grouped linear (DSv4's `DeepseekV4GroupedLinear`), recognised by @@ -970,7 +970,11 @@ def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | No out_features=module.out_features, n_groups=module.n_groups, has_bias=module.bias is not None, - **storage, + **{ + k: v + for k, v in storage_for(module_name).items() + if k in ("block_size", "activation_scheme", "scale_fmt") + }, ) if new_module is not None: # The kernels take raw pointers, so every operand must be a plain tensor. This is diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained_conversions.py index 8ce5ba63e23b..c411078230e5 100644 --- a/src/transformers/integrations/finegrained_conversions.py +++ b/src/transformers/integrations/finegrained_conversions.py @@ -322,7 +322,6 @@ class FineGrainedQuantize(_FineGrainedOp): Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" - def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: result: dict[str, torch.Tensor] = {} for key, value in input_dict.items(): @@ -474,7 +473,6 @@ class FineGrainedDequantize(_FineGrainedOp): result under the weight key, dropping the scales so the rest of the chain sees weights only. """ - def _scale_pattern_for(self, weight_pattern: str) -> str: anchored = weight_pattern.endswith("$") base = weight_pattern[:-1] if anchored else weight_pattern diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index f00f4865bb6d..508b63481cad 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -56,6 +56,9 @@ from .quantizer_eetq import EetqHfQuantizer from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer from .quantizer_finegrained import FineGrainedHfQuantizer +from .quantizer_finegrained_blockfp8 import FineGrainedBlockFp8HfQuantizer +from .quantizer_finegrained_mxfp4 import FineGrainedMxfp4HfQuantizer +from .quantizer_finegrained_nvfp4 import FineGrainedNvfp4HfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer from .quantizer_gemma import GemmaQuantizer @@ -94,14 +97,15 @@ "gguf": GgufHfQuantizer, "metal": MetalHfQuantizer, "auto-round": AutoRoundQuantizer, - # the finegrained quantizer serves every block/group-scaled format — block-FP8, MXFP8, MXFP4, - # NVFP4 and modelopt's NVFP4 export — for dense linears and MoE experts alike; the format is - # resolved off the checkpoint tensors - "fp8": FineGrainedHfQuantizer, - "mxfp8": FineGrainedHfQuantizer, - "mxfp4": FineGrainedHfQuantizer, - "nvfp4": FineGrainedHfQuantizer, - "modelopt": FineGrainedHfQuantizer, + # An arm exists where a PRODUCER's key layout needs one — GPT-OSS's packed `_blocks`, + # modelopt's two-level scales, the calibrated `input_scale`. MXFP8 ships the plain + # `weight` / `weight_scale_inv` pair, so the base serves it. Which FORMAT each module is + # quantized in is the config's groups, not the arm. + "fp8": FineGrainedBlockFp8HfQuantizer, + "mxfp8": FineGrainedHfQuantizer, # no distinct key layout — the base is its handler + "mxfp4": FineGrainedMxfp4HfQuantizer, + "nvfp4": FineGrainedNvfp4HfQuantizer, + "modelopt": FineGrainedNvfp4HfQuantizer, "sinq": SinqHfQuantizer, "gemma": GemmaQuantizer, } diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 0a591dadc5cc..211a5df5a0b9 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -13,6 +13,12 @@ # limitations under the License. """The `HfQuantizer` for the fine-grained family: converter chains, the parallel plan, and the environment checks. The modules are in `integrations/finegrained`. + +Two axes, deliberately separate. A per-scheme arm subclasses this one for what a PRODUCER's key +layout needs — modelopt's two-level scales, GPT-OSS's packed `_blocks` — and the registry picks +one by `quant_method`. What FORMAT each module is quantized in comes from the config's groups +instead (`group_for`), because a checkpoint can be several: DeepSeek-V4 is mxfp4 experts over +block-FP8 linears, and one arm per checkpoint could never express that. """ import re @@ -53,16 +59,14 @@ class FineGrainedHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedConfig" + def _default_activation_format(self) -> str | None: + """What to quantize activations to when the config leaves it open. `None` keeps the + kernels' weight-native choice; the MXFP4 arm overrides this to run weight-only.""" + return None + def _assert_dequantize_supported(self) -> None: - """The dequantize chain folds a weight into full precision with its per-block scale alone. - NVFP4 carries a second level (modelopt's `weight_scale_2`) that the chain has no slot for, - so refuse rather than hand back a weight scaled by `1 / global`.""" - if self._quant_method() == "nvfp4": - raise NotImplementedError( - f"`dequantize=True` is not supported for {self._quant_method()!r} checkpoints: their " - "two-level scales cannot be folded into a full-precision weight by this path. Load " - "them quantized on a GPU that serves NVFP4, or start from a bf16 checkpoint." - ) + """Whether `dequantize=True` can fold this format's scales into a full-precision weight. + Every one-level format can; the NVFP4 arm overrides this to refuse.""" def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): @@ -165,10 +169,11 @@ def _process_model_before_weight_loading( from ..integrations.finegrained import replace_with_finegrained_embedding, replace_with_finegrained_layer self._normalize_modules_to_not_convert(model) - if self._quant_method() == "mxfp4" and self.quantization_config.activation_format is None: - # GPT-OSS MXFP4 runs weight-only (W4A16): raw bf16 activations against packed - # weights. The kernels' weight-native default would quantize activations to mxfp4. - self.quantization_config.activation_format = "bf16" + # the one place both configs are in hand: a legacy checkpoint declares a second format + # for its experts on the MODEL config, and it becomes a group here + self.quantization_config.split_out_experts(getattr(model.config.get_text_config(), "expert_dtype", None)) + if self.quantization_config.activation_format is None: + self.quantization_config.activation_format = self._default_activation_format() if self.quantization_config.activation_format == "bf16": # Weight-only holds no activation global, so a calibrated checkpoint's `input_scale` # has no slot to load into. Not an unexpected key — one this run has no use for, which @@ -189,11 +194,13 @@ def _process_model_before_weight_loading( ) def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import disable_deepgemm_on_multi_device + from ..integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device from ..integrations.finegrained_conversions import keep_swizzle_reverse_for_save - keep_swizzle_reverse_for_save(model, self) + assert_modules_are_quantized(model) disable_deepgemm_on_multi_device(model) + keep_swizzle_reverse_for_save(model, self) + return model def update_tp_plan(self, config): @@ -291,32 +298,14 @@ def get_weight_conversions(self): return [] if self.quantization_config.dequantize: return self._dequantize_conversions() - if self._quant_method() == "mxfp4": - return self._mxfp4_conversions() - if self._quant_method() == "nvfp4": - return self._nvfp4_conversions() - if self.quantization_config.activation_scheme == "static": - return self._static_activation_conversions() return [] def _dequantize_conversions(self): - """Every quantized key folded back to a full-precision weight. One converter per key - layout: the packed `{proj}_blocks` + `{proj}_scales` pair GPT-OSS ships, and the - `weight` + `weight_scale_inv` pair everything else does.""" - from ..core_model_loading import Transpose, WeightConverter - from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks - - if self._quant_method() == "mxfp4": - # the blocks regroup into the packed rows and the exponent-byte scales come back out - # of them as bf16, in the (E, hidden, 2I) orientation the unquantized experts hold - return [ - WeightConverter( - source_patterns=[rf"{proj}_blocks$", rf"{proj}_scales$"], - target_patterns=proj, - operations=[FineGrainedPackedBlocks(self), FineGrainedDequantize(self), Transpose(1, 2)], - ) - for proj in ("gate_up_proj", "down_proj") - ] + """Every quantized key folded back to a full-precision weight, from the + `weight` + `weight_scale_inv` pair. The MXFP4 arm overrides this for GPT-OSS's + `{proj}_blocks` + `{proj}_scales`.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained_conversions import FineGrainedDequantize # anchored `weight$` so the scale keys land in the scale slots; the activation # scales are collected only to be dropped @@ -328,199 +317,6 @@ def _dequantize_conversions(self): ) ] - def _mxfp4_conversions(self): - """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed - weight, the exponent bytes take the scale container op; then the usual layout ops. Bare - patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedPackedBlocks, - FineGrainedScaleContainer, - FineGrainedSwizzleScales, - ) - - converters = [] - for proj in ("gate_up_proj", "down_proj"): - interleave = [FineGrainedInterleaveGateUp(self)] if proj == "gate_up_proj" else [] - converters += [ - WeightConverter( - source_patterns=rf"{proj}_blocks$", - target_patterns=proj, - operations=[FineGrainedPackedBlocks(self), *interleave, FineGrainedSwizzleScales(self)], - ), - WeightConverter( - source_patterns=rf"{proj}_scales$", - target_patterns=f"{proj}_scale_inv", - operations=[*interleave, FineGrainedScaleContainer(self), FineGrainedSwizzleScales(self)], - ), - ] - return converters - - def _nvfp4_conversions(self): - """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts - a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already - stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and - the other never fires. Routed experts only — a generic `weight_scale*` rename would run - before converter collection and mangle these keys.""" - # weight-only keeps activations bf16, so there is no activation global for a calibrated - # `input_scale` to land on — leave those keys to the unexpected-key filter - calibrated = self.quantization_config.activation_format != "bf16" - return self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) - - def _nvfp4_per_expert_conversions(self, calibrated: bool): - """One key per expert per projection, as GLM-5.2-NVFP4 ships them. - - `MergeModulelist` stamps each source with its expert index, which is what lets expert - parallelism keep only this rank's experts (`tensor_idx` in `core_model_loading`); without - it every rank collects all E values and the forward asserts on the per-expert count. - """ - from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter - from ..integrations.finegrained_conversions import FineGrainedInputScales, FineGrainedWeightGlobals - - merge = [MergeModulelist(dim=0)] - expert = r"mlp\.experts\..*\." - converters = [ - WeightConverter( - source_patterns=[rf"{expert}gate_proj\.weight_scale$", rf"{expert}up_proj\.weight_scale$"], - target_patterns="mlp.experts.gate_up_proj_scale_inv", - operations=[*merge, Concatenate(dim=1)], - ), - WeightConverter( - source_patterns=rf"{expert}down_proj\.weight_scale$", - target_patterns="mlp.experts.down_proj_scale_inv", - operations=merge, - ), - # every second-level global of a layer in ONE converter: the gate|up stack's two - # calibrated halves merge into one global per expert by folding the up half's onto - # the down projection, so the weight globals and the down's input scale have to be - # decided together (`FineGrainedWeightGlobals`) - WeightConverter( - source_patterns=[ - rf"{expert}gate_proj\.weight_scale_2", - rf"{expert}up_proj\.weight_scale_2", - rf"{expert}down_proj\.weight_scale_2", - *([rf"{expert}down_proj\.input_scale"] if calibrated else []), - ], - target_patterns=[ - "mlp.experts.gate_up_proj_weight_global_scale", - "mlp.experts.down_proj_weight_global_scale", - *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), - ], - operations=[*merge, FineGrainedWeightGlobals(self)], - ), - ] - if calibrated: - converters.append( - WeightConverter( - source_patterns=[rf"{expert}gate_proj\.input_scale", rf"{expert}up_proj\.input_scale"], - target_patterns="mlp.experts.gate_up_proj_input_global_scale", - operations=[*merge, FineGrainedInputScales(self)], - ) - ) - return converters - - def _nvfp4_fused_conversions(self, calibrated: bool): - """Already stacked per layer, as the vLLM fused layout and Muse-Spark ship them — so no - `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the - `.weight`-anchored rule gives the per-expert layout.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import ( - FineGrainedInputScales, - FineGrainedScaleContainer, - FineGrainedViewPackedInt8, - FineGrainedWeightGlobals, - ) - - converters = [ - # the container cast the module's scale dtype needs; `_with_expert_layout_ops` - # adds the interleave and the swizzle on top (it re-adds the cast, which is - # idempotent) - WeightConverter( - source_patterns=rf"experts\.{proj}_weight_scale$", - target_patterns=f"experts.{proj}_scale_inv", - operations=[FineGrainedScaleContainer(self)], - ) - for proj in ("gate_up_proj", "down_proj") - ] - converters.append( - # the same one-converter rule as the per-expert layout, over the fused keys - WeightConverter( - source_patterns=[ - r"experts\.gate_up_proj_weight_scale_2$", - r"experts\.down_proj_weight_scale_2$", - *([r"experts\.down_proj_input_scale$"] if calibrated else []), - ], - target_patterns=[ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - *(["experts.down_proj_input_global_scale"] if calibrated else []), - ], - operations=[FineGrainedWeightGlobals(self)], - ) - ) - if calibrated: - converters.append( - WeightConverter( - source_patterns=r"experts\.gate_up_proj_input_scale$", - target_patterns="experts.gate_up_proj_input_global_scale", - operations=[FineGrainedInputScales(self)], - ) - ) - converters += [ - WeightConverter( - source_patterns=rf"experts\.{proj}$", - target_patterns=f"experts.{proj}", - operations=[FineGrainedViewPackedInt8(self)], - ) - for proj in ("gate_up_proj", "down_proj") - ] - return converters - - def _static_activation_conversions(self): - """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in. One - value per quantized module: a dense linear brings one (Ministral-3), and a MoE brings one - per expert, each expert being its own quantized module (Mistral-4) — the gate|up pair - reduces to one per expert, both halves reading the same routed rows. Both expert layouts, - per expert per projection and already stacked per layer; a checkpoint matches one and the - other never fires. NVFP4 consumes ``input_scale`` as its activation GLOBAL instead - (:meth:`_nvfp4_conversions`) — a second level over a block scale, not the scale itself.""" - from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming - from ..integrations.finegrained_conversions import FineGrainedInputScales - - per_expert = [ - WeightConverter( - source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.input_scale", - r"mlp\.experts\..*\.up_proj\.input_scale", - ], - target_patterns="mlp.experts.gate_up_proj_activation_scale", - operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], - ), - WeightConverter( - source_patterns=r"mlp\.experts\..*\.down_proj\.input_scale", - target_patterns="mlp.experts.down_proj_activation_scale", - operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], - ), - ] - fused = [ - WeightConverter( - source_patterns=rf"experts\.{proj}_input_scale$", - target_patterns=f"experts.{proj}_activation_scale", - operations=[FineGrainedInputScales(self)], - ) - for proj in ("gate_up_proj", "down_proj") - ] - # every other calibrated module is a dense linear holding the one value itself — a plain - # rename. The lookahead keeps it off the expert keys above, which stack per layer instead - dense = [ - WeightRenaming( - source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", - target_patterns=r"\1.activation_scale", - ) - ] - return per_expert + fused + dense - def update_weight_conversions(self, weight_conversions): """Rewrite the model's conversion plan for this checkpoint: @@ -539,7 +335,7 @@ def update_weight_conversions(self, weight_conversions): :meth:`get_weight_conversions` is appended in every mode.""" from ..core_model_loading import WeightConverter, WeightRenaming - from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedViewPackedInt8 + from ..integrations.finegrained_conversions import FineGrainedDequantize scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") weight_conversions = [scale_rename, *weight_conversions] @@ -577,19 +373,6 @@ def anchor(conversions): updated.append(conv) return anchor(updated + self.get_weight_conversions()) - if self.pre_quantized and self._quant_method() == "nvfp4": - updated = [] - for conv in weight_conversions: - # anchored so the `.weight` source cannot swallow `weight_scale` / `weight_scale_2` - if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): - conv = WeightConverter( - source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], - target_patterns=conv.target_patterns, - operations=[*conv.operations, FineGrainedViewPackedInt8(self)], - ) - updated.append(conv) - weight_conversions = updated - return anchor(self._with_expert_layout_ops(weight_conversions + self.get_weight_conversions())) def _with_expert_layout_ops(self, weight_conversions): diff --git a/src/transformers/quantizers/quantizer_finegrained_blockfp8.py b/src/transformers/quantizers/quantizer_finegrained_blockfp8.py new file mode 100644 index 000000000000..5795ff07868e --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained_blockfp8.py @@ -0,0 +1,82 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The BLOCKFP8 arm of the fine-grained quantizer.""" + +from .quantizer_finegrained import FineGrainedHfQuantizer + + +class FineGrainedBlockFp8HfQuantizer(FineGrainedHfQuantizer): + """Block-FP8 and per-tensor FP8 — the only formats that take a CALIBRATED activation scale. + + A static scale is one value standing in for the quantization the other formats do inline, so + it only means anything where activations are quantized per tensor or per block. The group + formats hand the kernels a scale GRID instead (`As` is per group-32 / group-16), and NVFP4 + spends the checkpoint's `input_scale` as its second-level activation global, so neither can + consume one. Keeping the converters here is what makes that structural. + + Named `blockfp8` because `quantizer_finegrained_fp8` is the frozen legacy module. + """ + + def get_weight_conversions(self): + if not self.pre_quantized: + return [] + if self.quantization_config.dequantize: + return self._dequantize_conversions() + if self.quantization_config.activation_scheme == "static": + return self._static_activation_conversions() + return [] + + def _static_activation_conversions(self): + """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in. One + value per quantized module: a dense linear brings one (Ministral-3), and a MoE brings one + per expert, each expert being its own quantized module (Mistral-4) — the gate|up pair + reduces to one per expert, both halves reading the same routed rows. Both expert layouts, + per expert per projection and already stacked per layer; a checkpoint matches one and the + other never fires. NVFP4 consumes ``input_scale`` as its activation GLOBAL instead + (the NVFP4 arm's converters) — a second level over a block scale, not the scale itself.""" + from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming + from ..integrations.finegrained_conversions import FineGrainedInputScales + + per_expert = [ + WeightConverter( + source_patterns=[ + r"mlp\.experts\..*\.gate_proj\.input_scale", + r"mlp\.experts\..*\.up_proj\.input_scale", + ], + target_patterns="mlp.experts.gate_up_proj_activation_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], + ), + WeightConverter( + source_patterns=r"mlp\.experts\..*\.down_proj\.input_scale", + target_patterns="mlp.experts.down_proj_activation_scale", + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], + ), + ] + fused = [ + WeightConverter( + source_patterns=rf"experts\.{proj}_input_scale$", + target_patterns=f"experts.{proj}_activation_scale", + operations=[FineGrainedInputScales(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + # every other calibrated module is a dense linear holding the one value itself — a plain + # rename. The lookahead keeps it off the expert keys above, which stack per layer instead + dense = [ + WeightRenaming( + source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", + target_patterns=r"\1.activation_scale", + ) + ] + return per_expert + fused + dense diff --git a/src/transformers/quantizers/quantizer_finegrained_mxfp4.py b/src/transformers/quantizers/quantizer_finegrained_mxfp4.py new file mode 100644 index 000000000000..66b8d553f973 --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained_mxfp4.py @@ -0,0 +1,83 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The MXFP4 arm of the fine-grained quantizer: GPT-OSS's packed `_blocks` / `_scales` layout, +and the weight-only default that layout implies. +""" + +from .quantizer_finegrained import FineGrainedHfQuantizer + + +class FineGrainedMxfp4HfQuantizer(FineGrainedHfQuantizer): + """MXFP4 weights, W4A16 or W4A4. + + One thing is MXFP4's alone: GPT-OSS names each projection `{proj}_blocks` + `{proj}_scales` + rather than a weight/scale pair. + + The weight-only default is NOT a property of the format — real W4A4 (mxfp4 activations) is + supported and is what DeepSeek-V4's experts run. It only fills the gap when a config leaves + `activation_format` open, because the checkpoint that ships this key layout (GPT-OSS) is + W4A16; an explicit `activation_format="mxfp4"` still gets W4A4. + """ + + def _default_activation_format(self) -> str: + return "bf16" + + def get_weight_conversions(self): + if not self.pre_quantized: + return [] + return self._dequantize_conversions() if self.quantization_config.dequantize else self._mxfp4_conversions() + + def _dequantize_conversions(self): + """The blocks regroup into the packed rows and the exponent-byte scales come back out of + them as bf16, in the `(E, hidden, 2I)` orientation the unquantized experts hold.""" + from ..core_model_loading import Transpose, WeightConverter + from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks + + return [ + WeightConverter( + source_patterns=[rf"{proj}_blocks$", rf"{proj}_scales$"], + target_patterns=proj, + operations=[FineGrainedPackedBlocks(self), FineGrainedDequantize(self), Transpose(1, 2)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + + def _mxfp4_conversions(self): + """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed + weight, the exponent bytes take the scale container op; then the usual layout ops. Bare + patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained_conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) + + converters = [] + for proj in ("gate_up_proj", "down_proj"): + interleave = [FineGrainedInterleaveGateUp(self)] if proj == "gate_up_proj" else [] + converters += [ + WeightConverter( + source_patterns=rf"{proj}_blocks$", + target_patterns=proj, + operations=[FineGrainedPackedBlocks(self), *interleave, FineGrainedSwizzleScales(self)], + ), + WeightConverter( + source_patterns=rf"{proj}_scales$", + target_patterns=f"{proj}_scale_inv", + operations=[*interleave, FineGrainedScaleContainer(self), FineGrainedSwizzleScales(self)], + ), + ] + return converters diff --git a/src/transformers/quantizers/quantizer_finegrained_nvfp4.py b/src/transformers/quantizers/quantizer_finegrained_nvfp4.py new file mode 100644 index 000000000000..ca8a3c1e0b68 --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained_nvfp4.py @@ -0,0 +1,183 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The NVFP4 arm of the fine-grained quantizer: everything modelopt's two-level scales need +that the other formats do not. Everything else is inherited. +""" + +from ..utils import logging +from .quantizer_finegrained import FineGrainedHfQuantizer + + +logger = logging.get_logger(__name__) + + +class FineGrainedNvfp4HfQuantizer(FineGrainedHfQuantizer): + """NVFP4, in both key layouts modelopt ships: per expert per projection, or stacked per layer. + + What is NVFP4's alone, and why the base does not carry it: + + * the two-level globals. `weight_scale_2` has no slot in a one-level fold, so `dequantize=True` + is refused here rather than tested for in the shared path. + * the packed E2M1 weights, which need a uint8 -> int8 BITCAST (`copy_` would convert and + corrupt anything >= 128) and an anchored `.weight` source so it cannot swallow + `weight_scale` / `weight_scale_2`. + """ + + def _assert_dequantize_supported(self) -> None: + raise NotImplementedError( + "`dequantize=True` is not supported for nvfp4 checkpoints: their two-level scales " + "cannot be folded into a full-precision weight by this path. Load them quantized on a " + "GPU that serves NVFP4, or start from a bf16 checkpoint." + ) + + def get_weight_conversions(self): + return self._nvfp4_conversions() if self.pre_quantized else [] + + def update_weight_conversions(self, weight_conversions): + """The base chain, with every `.weight` source anchored and bitcast to the int8 view.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained_conversions import FineGrainedViewPackedInt8 + + if self.pre_quantized: + updated = [] + for conv in weight_conversions: + if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): + conv = WeightConverter( + source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], + target_patterns=conv.target_patterns, + operations=[*conv.operations, FineGrainedViewPackedInt8(self)], + ) + updated.append(conv) + weight_conversions = updated + return super().update_weight_conversions(weight_conversions) + + def _nvfp4_conversions(self): + """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts + a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already + stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and + the other never fires. Routed experts only — a generic `weight_scale*` rename would run + before converter collection and mangle these keys.""" + # weight-only keeps activations bf16, so there is no activation global for a calibrated + # `input_scale` to land on — leave those keys to the unexpected-key filter + calibrated = self.quantization_config.activation_format != "bf16" + return self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) + + def _nvfp4_per_expert_conversions(self, calibrated: bool): + """One key per expert per projection, as GLM-5.2-NVFP4 ships them. + + `MergeModulelist` stamps each source with its expert index, which is what lets expert + parallelism keep only this rank's experts (`tensor_idx` in `core_model_loading`); without + it every rank collects all E values and the forward asserts on the per-expert count. + """ + from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter + from ..integrations.finegrained_conversions import FineGrainedInputScales, FineGrainedWeightGlobals + + merge = [MergeModulelist(dim=0)] + expert = r"mlp\.experts\..*\." + converters = [ + WeightConverter( + source_patterns=[rf"{expert}gate_proj\.weight_scale$", rf"{expert}up_proj\.weight_scale$"], + target_patterns="mlp.experts.gate_up_proj_scale_inv", + operations=[*merge, Concatenate(dim=1)], + ), + WeightConverter( + source_patterns=rf"{expert}down_proj\.weight_scale$", + target_patterns="mlp.experts.down_proj_scale_inv", + operations=merge, + ), + # every second-level global of a layer in ONE converter: the gate|up stack's two + # calibrated halves merge into one global per expert by folding the up half's onto + # the down projection, so the weight globals and the down's input scale have to be + # decided together (`FineGrainedWeightGlobals`) + WeightConverter( + source_patterns=[ + rf"{expert}gate_proj\.weight_scale_2", + rf"{expert}up_proj\.weight_scale_2", + rf"{expert}down_proj\.weight_scale_2", + *([rf"{expert}down_proj\.input_scale"] if calibrated else []), + ], + target_patterns=[ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[*merge, FineGrainedWeightGlobals(self)], + ), + ] + if calibrated: + converters.append( + WeightConverter( + source_patterns=[rf"{expert}gate_proj\.input_scale", rf"{expert}up_proj\.input_scale"], + target_patterns="mlp.experts.gate_up_proj_input_global_scale", + operations=[*merge, FineGrainedInputScales(self)], + ) + ) + return converters + + def _nvfp4_fused_conversions(self, calibrated: bool): + """Already stacked per layer, as the vLLM fused layout and Muse-Spark ship them — so no + `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the + `.weight`-anchored rule gives the per-expert layout.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained_conversions import ( + FineGrainedInputScales, + FineGrainedScaleContainer, + FineGrainedViewPackedInt8, + FineGrainedWeightGlobals, + ) + + converters = [ + # the container cast the module's scale dtype needs; `_with_expert_layout_ops` + # adds the interleave and the swizzle on top (it re-adds the cast, which is + # idempotent) + WeightConverter( + source_patterns=rf"experts\.{proj}_weight_scale$", + target_patterns=f"experts.{proj}_scale_inv", + operations=[FineGrainedScaleContainer(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + converters.append( + # the same one-converter rule as the per-expert layout, over the fused keys + WeightConverter( + source_patterns=[ + r"experts\.gate_up_proj_weight_scale_2$", + r"experts\.down_proj_weight_scale_2$", + *([r"experts\.down_proj_input_scale$"] if calibrated else []), + ], + target_patterns=[ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_input_global_scale"] if calibrated else []), + ], + operations=[FineGrainedWeightGlobals(self)], + ) + ) + if calibrated: + converters.append( + WeightConverter( + source_patterns=r"experts\.gate_up_proj_input_scale$", + target_patterns="experts.gate_up_proj_input_global_scale", + operations=[FineGrainedInputScales(self)], + ) + ) + converters += [ + WeightConverter( + source_patterns=rf"experts\.{proj}$", + target_patterns=f"experts.{proj}", + operations=[FineGrainedViewPackedInt8(self)], + ) + for proj in ("gate_up_proj", "down_proj") + ] + return converters diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 169bec0af60c..be8b726ab993 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -19,7 +19,7 @@ import json import os import re -from dataclasses import dataclass +from dataclasses import asdict, dataclass, field, replace from enum import Enum from typing import Any, Optional, Union @@ -1693,6 +1693,75 @@ def post_init(self): _MODELOPT_ALGOS = {"NVFP4": (QuantizationMethod.NVFP4, "nvfp4")} +_CATCH_ALL = [".*"] # the targets of a group that takes whatever no other group claimed + + +@dataclass +class FineGrainedGroup: + """One set of modules and the format they are quantized in. + + A single-format checkpoint is one group covering everything, which is what the flat + `FineGrainedConfig` fields spell. A checkpoint that is more than one format needs more than + one: DeepSeek-V4 is W4A4 mxfp4 EXPERTS over block-FP8 linears. + """ + + quant_method: str + targets: list[str] = field(default_factory=lambda: list(_CATCH_ALL)) + activation_format: str | None = None + activation_scheme: str = "dynamic" + weight_block_size: tuple[int, int] | None = None + scale_fmt: str = "float" + + def __post_init__(self): + if self.weight_block_size is not None: + self.weight_block_size = tuple(self.weight_block_size) + + def matches(self, module_name: str) -> bool: + return any(re.search(target, module_name) for target in self.targets) + + +# (num_bits, type, group_size) as compressed-tensors / modelopt spell a format, to the name our +# kernels dispatch on. `group_size` is what separates the two 4-bit formats: nvfp4's block scale +# is E4M3 over 16 values, mxfp4's is E8M0 over 32. +_CT_FORMATS = { + (4, "float", 16): "nvfp4", + (4, "float", 32): "mxfp4", + (8, "float", 32): "mxfp8", + (8, "float", None): "fp8", +} + + +def _group_from_config_groups(spec: dict) -> FineGrainedGroup | None: + """One `config_groups` entry as a `FineGrainedGroup`, or `None` if it names a format we do not + serve — in which case the caller falls back to the flat fields rather than guessing. + + compressed-tensors describes a format by its parameters; we name it. `targets` is theirs too: + a bare entry is a CLASS name (GLM-5.2 ships `["Linear"]`, i.e. every linear), and an `re:` + prefix introduces a pattern. + """ + weights = spec.get("weights") or {} + key = (weights.get("num_bits"), weights.get("type"), weights.get("group_size")) + quant_method = _CT_FORMATS.get(key) + if quant_method is None: + return None + targets = [] + for target in spec.get("targets") or []: + if target.startswith("re:"): + targets.append(target[3:]) + elif target == "Linear": # every linear — the catch-all, in their spelling + targets.extend(_CATCH_ALL) + else: # a class name we cannot resolve to module paths here + return None + activations = spec.get("input_activations") or {} + act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) + return FineGrainedGroup( + quant_method=quant_method, + targets=targets or list(_CATCH_ALL), + activation_format=_CT_FORMATS.get(act_key) if activations else None, + activation_scheme="dynamic" if activations.get("dynamic", True) else "static", + ) + + @dataclass class FineGrainedConfig(QuantizationConfigMixin): """ @@ -1742,6 +1811,7 @@ def __init__( modules_to_convert: list | None = None, scale_fmt: str = "float", activation_format: str | None = None, + groups: dict[str, FineGrainedGroup] | None = None, **kwargs, ): self.quant_method = kwargs.pop("quant_method", QuantizationMethod.FP8) @@ -1768,6 +1838,13 @@ def __init__( f"{str(self.quant_method)!r}." ) self.activation_format = activation_format or algo_activation_format + # modelopt also describes the format parametrically under `config_groups`; take it + # when every entry maps, so a multi-group export is not silently flattened to one + ct_groups = kwargs.pop("config_groups", None) + if groups is None and ct_groups: + mapped = {name: _group_from_config_groups(spec) for name, spec in ct_groups.items()} + if all(group is not None for group in mapped.values()): + groups = mapped ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) if modules_to_not_convert is None and ignore is not None: # modelopt names skipped subtrees with GLOBS — "model.layers.0*", "model.layers.1.*", @@ -1784,8 +1861,67 @@ def __init__( self.weight_block_size = weight_block_size self.dequantize = dequantize self.scale_fmt = scale_fmt + # Every config is GROUPED internally. The flat fields above are the one-group spelling a + # single-format checkpoint ships, normalised here the way `quant_algo` is, so that + # downstream asks `group_for(name)` and never learns which spelling it came from. + if groups: + # from `config.json` a group is a plain dict + self.groups = { + name: group if isinstance(group, FineGrainedGroup) else FineGrainedGroup(**group) + for name, group in groups.items() + } + else: + self.groups = { + "default": FineGrainedGroup( + quant_method=str(getattr(self.quant_method, "value", self.quant_method)), + activation_format=self.activation_format, + activation_scheme=self.activation_scheme, + weight_block_size=self.weight_block_size, + scale_fmt=self.scale_fmt, + ) + } self.post_init() + def to_dict(self): + """`config.json`-ready: the groups go out as plain dicts so the config round-trips.""" + out = super().to_dict() + out["groups"] = {name: asdict(group) for name, group in self.groups.items()} + return out + + def group_for(self, module_name: str) -> FineGrainedGroup: + """The group a module is quantized by. + + A targeted group wins over the catch-all whatever order they were declared in, because + `to_json_string` sorts keys and a config that round-tripped through `config.json` has lost + that order. Two targeted groups claiming the same module is ambiguous, and says so. + """ + claimed = [ + name for name, group in self.groups.items() if group.targets != _CATCH_ALL and group.matches(module_name) + ] + if len(claimed) > 1: + raise ValueError(f"{module_name!r} is claimed by more than one group: {claimed}") + if claimed: + return self.groups[claimed[0]] + for group in self.groups.values(): + if group.targets == _CATCH_ALL: + return group + raise KeyError(f"no group covers {module_name!r}; groups: {list(self.groups)}") + + def split_out_experts(self, expert_dtype: str | None) -> None: + """Fold a legacy `config.expert_dtype` into the group it means. + + DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL config, + because a flat quant config carries one format and cannot say "the experts differ". Split + that into the two groups it describes, so nothing downstream reads the side-channel. + """ + if expert_dtype != "fp4" or len(self.groups) > 1: + return + flat = self.groups["default"] + self.groups = { + "experts": replace(flat, quant_method="mxfp4", activation_format="mxfp4", targets=[r"\.experts($|\.)"]), + "dense": flat, + } + def post_init(self): r""" Safety checker that arguments are correct diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 4c0d3864f7b3..17782b860b84 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -113,6 +113,15 @@ def _loaded(kernel): ) +def _quantizer_for(config): + """The quantizer a `from_pretrained` would pick for this config. Since the arms were split + per scheme, constructing the base directly only exercises what every scheme shares.""" + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING + + method = getattr(config.quant_method, "value", config.quant_method) + return AUTO_QUANTIZER_MAPPING[method](config) + + @require_torch class FineGrainedLoaderTest(unittest.TestCase): def test_loader_requires_every_symbol(self): @@ -583,10 +592,13 @@ def _no_accelerator(self): yield def _quantizer(self, pre_quantized, **cfg_kwargs): - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + # through the registry, so the test sees the class a `from_pretrained` would pick + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING from transformers.utils.quantization_config import FineGrainedConfig - quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**cfg_kwargs)) + config = FineGrainedConfig(**cfg_kwargs) + method = getattr(config.quant_method, "value", config.quant_method) + quantizer = AUTO_QUANTIZER_MAPPING[method](config) quantizer.pre_quantized = pre_quantized return quantizer @@ -759,14 +771,13 @@ def test_dequantize_folds_every_quantized_key_to_full_precision(self): write into.""" import re - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig for method, sources in ( ("mxfp4", ("experts.gate_up_proj_blocks", "experts.gate_up_proj_scales")), ("fp8", ("mlp.down_proj.weight", "mlp.down_proj.weight_scale_inv")), ): - quantizer = FineGrainedHfQuantizer(FineGrainedConfig(quant_method=method, dequantize=True)) + quantizer = _quantizer_for(FineGrainedConfig(quant_method=method, dequantize=True)) quantizer.pre_quantized = True converters = quantizer.get_weight_conversions() with self.subTest(quant_method=method): @@ -1134,10 +1145,9 @@ def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) def _quantizer(self, quant_method): - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig - return FineGrainedHfQuantizer(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions + return _quantizer_for(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions def _arch_converters(self): from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter @@ -1352,12 +1362,11 @@ class FineGrainedModeloptConverterTest(unittest.TestCase): rank collects all E globals and the forward asserts on the per-expert count.""" def _modelopt_conversions(self, **cfg_kwargs): - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.quantizer_finegrained_nvfp4 import FineGrainedNvfp4HfQuantizer from transformers.utils.quantization_config import FineGrainedConfig cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) - quantizer = FineGrainedHfQuantizer(cfg) - return quantizer.get_weight_conversions() + return FineGrainedNvfp4HfQuantizer(cfg).get_weight_conversions() def test_every_declared_target_is_a_slot_the_module_holds(self): """A converter target that no module holds is a load failure, and the activation global is @@ -1469,10 +1478,9 @@ def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): per expert rather than flattened to 2E.""" import re - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig - quantizer = FineGrainedHfQuantizer(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) + quantizer = _quantizer_for(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) quantizer.pre_quantized = True converters = quantizer.update_weight_conversions([]) @@ -1502,7 +1510,7 @@ def claim(key): ] self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") # a dynamic checkpoint has no such key to route, and a weight-only one ignores it - dynamic = FineGrainedHfQuantizer(FineGrainedConfig(quant_method="fp8")) + dynamic = _quantizer_for(FineGrainedConfig(quant_method="fp8")) dynamic.pre_quantized = True self.assertEqual(dynamic.get_weight_conversions(), []) From a67e7acfcd816d4444dee70cbecb02d6be8749f3 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Tue, 22 Sep 2026 07:51:40 -0700 Subject: [PATCH 39/68] Carry the quantized-load fixes onto the per-scheme arms, and the same simplification The three defects found on `finegrained-unified` were all in code this branch shares, so they were all live here -- and the middle one bites harder, because this branch has no `element_size() > 1` fallback and `assert_modules_are_quantized` turns it into a failed load rather than a silent bf16 forward: - the experts' static `activation_scale` was allocated without an explicit dtype and came out bf16 under `from_pretrained`; - it was never WRITTEN. The loader materialises a key no checkpoint supplies with `torch.empty_like` and `_init_weights` has no branch for a scale, so it reached the kernels as uninitialised memory -- a zero divides the activations by zero; - `_quantize_block_fp8` refused a shape the block does not divide, though the format's own producers do not: DeepSeek-V3 ships `kv_a_proj_with_mqa` as (576, 7168) E4M3 against a 128x128 block with a (5, 56) grid. It pads now. That refusal was also the filter keeping 2-D expert biases out, so they are turned away against `_weight_holder` instead. With shapes no longer refused, `assert_modules_are_quantized` stops being a trap: what is left for it to catch is a checkpoint that targets a module and ships no scales, which is the thing it was written for. The swap simplification lands here too, so the two branches stay comparable: `FineGrainedGroupedLinear` takes the same storage mapping as the dense linear (it was hardcoding fp8), and `_quantized_experts` lifts 27 lines out of the three-way dispatch. Co-Authored-By: Claude Opus 5 --- src/transformers/integrations/finegrained.py | 70 ++++++----- .../integrations/finegrained_conversions.py | 75 +++++++---- tests/kernels/test_finegrained.py | 116 ++++++++++++++++++ 3 files changed, 207 insertions(+), 54 deletions(-) diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained.py index 28dcb4c19ac1..5795e9c757aa 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained.py @@ -498,6 +498,8 @@ def __init__( activation_scheme: str = "dynamic", scale_fmt: str = "float", has_bias: bool = False, + weight_format: str = "fp8", + activation_format: str | None = None, ): super().__init__( in_features=in_features_per_group, @@ -506,6 +508,8 @@ def __init__( activation_scheme=activation_scheme, scale_fmt=scale_fmt, has_bias=has_bias, + weight_format=weight_format, + activation_format=activation_format, ) self.n_groups = n_groups @@ -725,7 +729,10 @@ def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, torch.ones(self.num_experts, dtype=self.global_scale_dtype) if two_level else None, ), ("input_global_scale", torch.ones(input_globals, dtype=torch.float32) if calibrated else None), - ("activation_scale", torch.ones(self.num_experts) if self.activation_scheme == "static" else None), + ( + "activation_scale", + torch.ones(self.num_experts, dtype=torch.float32) if self.activation_scheme == "static" else None, + ), ): _set_optional_parameter(self, f"{proj}_{name}", tensor) @@ -897,6 +904,34 @@ def replace_with_finegrained_embedding(model, patterns: list[str], modules_to_no return model +def _quantized_experts(module: nn.Module, model: nn.Module, storage: dict) -> nn.Module: + """A `FineGrainedExperts` holding `module`'s weights, in the model's own gate|up convention. + + The flags travel twice over: `use_experts_implementation` stamps them on the instance after + `__init__`, so the class it builds needs them as well as the constructor. A per-expert output + norm belongs to the model rather than to the quantization, so it comes across as it is. + """ + flags = { + "has_gate": getattr(module, "has_gate", True), + "has_bias": getattr(module, "has_bias", False), + "is_concatenated": getattr(module, "is_concatenated", True), + } + experts_class = use_experts_implementation( + experts_class=FineGrainedExperts, + experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + **flags, + ) + experts = experts_class(config=getattr(module, "config", model.config.get_text_config()), **storage, **flags) + if getattr(module, "post_expert_norm", None) is not None: + # `_apply_post_norm` is the standard hook and needs no rebinding; a NAMED form is what a + # backend can fuse. `use_experts_implementation` owns the flag and sets it from the class + # declaration, so this updates it rather than deriving it. + experts.post_expert_norm = module.post_expert_norm + experts.has_post_expert_norm = True + experts.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) + return experts + + def replace_with_finegrained_layer(model, modules_to_not_convert: list[str] | None = None, quantization_config=None): """Swap the model's ``nn.Linear`` (and grouped-linear) modules for ``FineGrainedLinear`` and its ``.experts`` modules for ``FineGrainedExperts``, both allocated in the quantized storage the @@ -927,32 +962,7 @@ def storage_for(module_name: str) -> dict: new_module = None with torch.device("meta"): if module_name.endswith(".experts"): - # the decorator stamps these on the instance after __init__, so the model's own - # gate|up convention has to travel through it as well as into the ctor - flags = { - "has_gate": getattr(module, "has_gate", True), - "has_bias": getattr(module, "has_bias", False), - "is_concatenated": getattr(module, "is_concatenated", True), - } - new_class = use_experts_implementation( - experts_class=FineGrainedExperts, - experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, - **flags, - ) - new_module = new_class( - config=getattr(module, "config", model.config.get_text_config()), - **storage_for(module_name), - **flags, - ) - # A per-expert output norm belongs to the model, not the quantization, so the - # swap carries the submodule over; `_apply_post_norm` is the standard hook and - # needs no rebinding. A named form is what a backend can fuse. - if getattr(module, "post_expert_norm", None) is not None: - new_module.post_expert_norm = module.post_expert_norm - # `use_experts_implementation` owns this flag on every experts class and sets - # it from the class declaration, so the swap updates it rather than deriving it - new_module.has_post_expert_norm = True - new_module.post_expert_norm_name = getattr(module, "post_expert_norm_name", None) + new_module = _quantized_experts(module, model, storage_for(module_name)) elif type(module) is nn.Linear: new_module = FineGrainedLinear( in_features=module.in_features, @@ -970,11 +980,7 @@ def storage_for(module_name: str) -> dict: out_features=module.out_features, n_groups=module.n_groups, has_bias=module.bias is not None, - **{ - k: v - for k, v in storage_for(module_name).items() - if k in ("block_size", "activation_scheme", "scale_fmt") - }, + **storage_for(module_name), ) if new_module is not None: # The kernels take raw pointers, so every operand must be a plain tensor. This is diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained_conversions.py index c411078230e5..e21067ea9336 100644 --- a/src/transformers/integrations/finegrained_conversions.py +++ b/src/transformers/integrations/finegrained_conversions.py @@ -319,14 +319,22 @@ class FineGrainedQuantize(_FineGrainedOp): power-of-two UE8M0 inverse scales); the group formats run the kernels' row-wise quantizers (``mxfp8_act_quant`` / ``mxfp4_act_quant`` / ``nvfp4_act_quant``), NVFP4 after normalizing by the canonical per-tensor / per-expert global ``amax / (6 * 448)``. - Tensors that are not a finegrained module's weight (1-D norms, biases, shapes that don't tile) - pass through. Without a model (direct invocation) it quantizes block-FP8 from the config.""" + A tensor that is not a finegrained module's weight passes through — `_weight_holder` decides, + since rank alone does not (an expert bias is 2-D). Without a model, a direct invocation has + nothing to ask and quantizes block-FP8 from the config.""" def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: result: dict[str, torch.Tensor] = {} for key, value in input_dict.items(): tensor = value[0] if isinstance(value, list) else value - result.update(self._quantize_one(key, tensor, self._weight_holder(model, key))) + holder = self._weight_holder(model, key) + if holder is None and model is not None: + # not a weight, whatever its rank: an expert bias is 2-D and would otherwise be + # quantized into a `_bias_scale_inv` no module holds. Only a direct + # invocation, with no model to ask, quantizes without a holder. + result[key] = tensor + continue + result.update(self._quantize_one(key, tensor, holder)) return result @staticmethod @@ -384,8 +392,12 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc if held_scale is not None else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" ) - quantized = self._quantize_block_fp8(value, block, ue8m0) - return {key: value} if quantized is None else {key: quantized[0], f"{prefix}_scale_inv": quantized[1]} + weight, scale = self._quantize_block_fp8(value, block, ue8m0) + return { + key: weight, + f"{prefix}_scale_inv": scale, + **self._identity_activation_scales(module, prefix, scale_name, value.device), + } if value.device.type not in ("cuda", "xpu"): # the kernels' quantizers are triton launches; the package builds for cuda, rocm @@ -405,34 +417,53 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc stem = scale_name.removesuffix("_scale_inv") suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" out[f"{prefix}{suffix}"] = global_scale - # activations are quantized dynamically against the block scales alone here — there is no - # calibration pass — so the module's activation globals are the identity - held_input = getattr(module, scale_name.replace("_scale_inv", "_input_global_scale"), None) - if held_input is not None: - # on the weight's device, not the slot's: the module is still on meta here, and - # `ones_like` would inherit that and leave the parameter unmaterialized - out[f"{prefix}_input_global_scale"] = torch.ones( - held_input.shape, dtype=held_input.dtype, device=value.device - ) + out.update(self._identity_activation_scales(module, prefix, scale_name, value.device)) + return out + + @staticmethod + def _identity_activation_scales(module, prefix: str, scale_name: str | None, device) -> dict: + """This module's activation-side scales, as the identity, because quantizing on the fly + is not a calibration pass. They have to be WRITTEN: the loader materializes a key no + checkpoint supplies with `torch.empty_like` and `_init_weights` has no branch for a + scale, so a slot left out here reaches the kernels as uninitialized memory.""" + if module is None: + return {} + stem = scale_name.removesuffix("_scale_inv") + # the experts prefix each slot with the projection; a dense linear, whose stem IS + # `weight`, names them bare + proj = "" if stem == "weight" else f"{stem}_" + out = {} + for slot in (f"{proj}activation_scale", f"{proj}input_global_scale"): + held = getattr(module, slot, None) + if held is not None: + # on the weight's device, not the slot's: the module is still on meta here, and + # `ones_like` would inherit that and leave the parameter unmaterialized + out[prefix.removesuffix(stem) + slot] = torch.ones(held.shape, dtype=held.dtype, device=device) return out @staticmethod def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): - """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``, or - ``None`` when the shape does not tile. UE8M0 rounds the inverse scale up to a power of two - before quantizing, so dequant multiplies by exactly the scale the weight was divided by.""" + """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``. + + A trailing PARTIAL block is padded, not refused — DeepSeek-V3 ships `kv_a_proj_with_mqa` + as `(576, 7168)` against a 128x128 block with a `(5, 56)` grid, so the format's own + producers round up and the module allocates to match. The zeros cannot move a block's + amax, so the short block is scaled by its real values. UE8M0 rounds the inverse scale up + to a power of two before quantizing, so dequant multiplies by exactly what it divided by. + """ block_m, block_n = block rows, cols = value.shape[-2], value.shape[-1] - if rows % block_m or cols % block_n: - return None - tiles = value.float().reshape(*value.shape[:-2], rows // block_m, block_m, cols // block_n, block_n) + padded = torch.nn.functional.pad(value.float(), (0, -cols % block_n, 0, -rows % block_m)) + grid_m, grid_n = padded.shape[-2] // block_m, padded.shape[-1] // block_n + tiles = padded.reshape(*padded.shape[:-2], grid_m, block_m, grid_n, block_n) max_abs = tiles.abs().amax(dim=(-3, -1)) inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) if ue8m0: inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) - quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(value.shape) - return quantized, inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(padded.shape) + # a no-op slice on a shape that tiled, and `contiguous` then returns the tensor itself + return quantized[..., :rows, :cols].contiguous(), inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) @staticmethod def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 17782b860b84..1242ed506572 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -504,6 +504,54 @@ def test_unfusable_act_fn_is_passed_as_the_module_glu(self): @require_torch +class FineGrainedScaleDtypeTest(unittest.TestCase): + """A scale's dtype is the format's, never the ambient default. + + `from_pretrained` sets the default dtype to the checkpoint's for the duration of model + construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it + as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches + that, whatever the format decides each scale should be. + """ + + def _dtypes(self, build, default): + previous = torch.get_default_dtype() + torch.set_default_dtype(default) + try: + with torch.device("meta"): + module = build() + return {name: p.dtype for name, p in module.named_parameters() if p is not None} + finally: + torch.set_default_dtype(previous) + + def _assert_default_dtype_invariant(self, build): + under_fp32 = self._dtypes(build, torch.float32) + under_bf16 = self._dtypes(build, torch.bfloat16) + self.assertTrue(under_fp32) + # the weight and bias DO follow the model dtype; every scale beside them must not + scales = {name for name in under_fp32 if "scale" in name} + self.assertTrue(scales) + for name in sorted(scales): + self.assertEqual(under_bf16[name], under_fp32[name], f"{name} followed the default dtype") + + def test_linear_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedLinear( + in_features=64, out_features=32, block_size=(4, 4), weight_format=f, activation_scheme=s + ) + ) + + def test_experts_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedExperts( + _Cfg(), block_size=(4, 4), weight_format=f, activation_scheme=s + ) + ) + + class FineGrainedFusedNormGateTest(unittest.TestCase): """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. @@ -1328,6 +1376,74 @@ def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) + def test_a_static_scheme_gets_its_activation_scales_written(self): + """A calibration-fed slot the checkpoint does not supply is still WRITTEN, as the identity. + + The loader materializes a missing key with `torch.empty_like` and `_init_weights` has no + branch for a scale, so a slot this op leaves out reaches the kernels as uninitialized + memory: a zero divides the activations by zero and a negative flips their sign, which + showed up as NaN logits from one fixture and not another, run to run. + """ + from transformers.integrations.finegrained_conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts( + cfg, block_size=(128, 128), weight_format="fp8", activation_scheme="static" + ) + model.proj = fg.FineGrainedLinear( + in_features=256, + out_features=256, + block_size=(128, 128), + weight_format="fp8", + activation_scheme="static", + ) + op = FineGrainedQuantize(hf_quantizer=None) + for key, tensor, slot in ( + ("experts.gate_up_proj", torch.randn(4, 256, 256), "experts.gate_up_proj_activation_scale"), + ("proj.weight", torch.randn(256, 256), "proj.activation_scale"), + ): + with self.subTest(key=key): + out = op.convert({key: tensor}, model=model) + self.assertIn(slot, out, f"{key}: the static activation scale was not written") + held = model.get_parameter(slot) + self.assertEqual((out[slot].shape, out[slot].dtype), (held.shape, held.dtype)) + torch.testing.assert_close(out[slot], torch.ones_like(out[slot]), rtol=0, atol=0) + + def test_an_expert_bias_is_passed_through_not_quantized(self): + """A GPT-OSS expert bias is `(E, rows)` — 2-D, like a dense weight — so the rank guard + alone lets it through. It has no scale slot, and emitting one gives the loader a + `_bias_scale_inv` no module holds.""" + from transformers.integrations.finegrained_conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", has_bias=True) + bias = torch.randn(4, 256) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.down_proj_bias": bias}, model=model) + self.assertEqual(list(out), ["experts.down_proj_bias"]) + self.assertIs(out["experts.down_proj_bias"], bias) + + def test_a_partial_trailing_block_is_padded_not_refused(self): + """DeepSeek-V3 ships `kv_a_proj_with_mqa` as `(576, 7168)` against a 128x128 block with a + `(5, 56)` scale grid, so the format rounds the grid UP and quantizes the short block on + its own values. Refusing the shape instead left the weight full precision.""" + from transformers.integrations.finegrained_conversions import FineGrainedQuantize + + torch.manual_seed(0) + for rows, cols in ((576, 256), (192, 256), (256, 256)): + with self.subTest(shape=(rows, cols)): + weight, scale = FineGrainedQuantize._quantize_block_fp8( + torch.randn(rows, cols), (128, 128), ue8m0=False + ) + self.assertEqual(weight.shape, (rows, cols)) + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(scale.shape, (-(-rows // 128), -(-cols // 128))) + def test_block_fp8_round_trips_within_its_floor(self): from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedQuantize From 9fc683de91b624f3f3d4bb7d0adf464bc4bcce8f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 09:24:28 -0700 Subject: [PATCH 40/68] Address the per-scheme review Structure, as asked twice in the thread: integrations/finegrained/{__init__,core,conversions}.py quantizers/finegrained/{__init__,base,blockfp8,mxfp4,mxfp8,nvfp4}.py `quantizer_finegrained_fp8.py` stays where it is: that is the deprecated `FineGrainedFP8HfQuantizer`, which this PR does not own. Each package re-exports its public names, so `_import_structure` and every `from ...finegrained import X` are untouched; the tests name `finegrained.core` directly, because they patch module internals and a re-export cannot stand in for that. `_apply_post_norm` is gone. `post_expert_norm` is an `nn.Module` and so was already callable, no backend overrode the hook, and `_apply_gate` one branch above was already DEFAULTED where this one raised `TypeError` -- the asymmetry was the tell. Every call site now calls the norm, and the callable handed to the kernels for an unfusable norm is the norm itself. Comments: the blocks marked "suggested change" are deleted rather than shortened, as are the ones called a nothingburger and self-explanatory, plus three more that only restated their own line and two that were teaching `nn.Module.__setattr__` and `Tensor.contiguous`. The rest are trimmed where the thread asked for trimming. What survives says something the code cannot: why a lookahead excludes the expert keys, why each companion needs its own plan entry, why a `shared_experts.up_proj` is skipped despite matching the word. Smaller ones: `_fused_experts_forward` removed and both arms spelled out; two `for name, tensor` loops unrolled; the gate_up/up name resolved once as `gate_up_name`; `is_concatenated` moved in with the layout metadata; the `input_global_scale` lookup collapsed from three copies of one condition to one (weight-only never registers the slot, so the `bf16` check was testing a fact the parameter already encoded); helpers hoisted to the top of `conversions.py`; `as_container` and `group_from_config_groups` made public; `rpartition(".")[-1]`; `chunk` for the gate|up split; `_CT_FORMATS` and `_MODELOPT_ALGOS` folded into their one reader; `_read_modelopt` and a general `subtree_patterns` split out; `split_out_experts` replaced by a pure `groups_with_expert_dtype`; `supports_post_expert_norm` and a `supports_dequantize` bool; one named quantizer per key in `auto.py`; four import-time `warnings.warn` shims moved to `logger.warning_once` at construction. Tests: the two `core_model_loading` fp8 tests now exercise this integration instead of the frozen module, which surfaced three bugs in them -- a `PermuteForRope()` built with no `permute_layer_names` (it iterates them unconditionally, so it raised on any input, and every shipped call site passes `["q_proj", "k_proj"]`), an expectation that `v_proj` is rope-permuted, and state-dict keys prefixed `model.` against a model that never nested that way. The float-path test keeps its bit-equality guard but against OUR formula: we store `amax / MAX` and divide by it, where the frozen path stores `1.0 / (MAX / amax)` and lands off `amax / MAX` on 51 of 200 random blocks. Real DeepSeek-V3 agrees with ours -- recomputing `amax(q * scale) / 448` on a block reproduces its shipped `weight_scale_inv` exactly. `TestZeroDimParameterSharding` becomes two methods on `TestDtensorShardOperation` beside the other `shard_tensor` tests, using their builder. It needs no ranks and no devices, so behind `@is_tensor_parallel_test` it was skipped in CI outright. The triton skip on the qkv test goes too: that path is pure torch. Not done, and answered on the thread instead: no new `QuantizationMethod` member. `quant_method` is at once the quantizer key, the config key and the weight FORMAT fed to `resolve_weight_format`, so a `FINEGRAINED` member has no format to resolve to; splitting those two meanings is its own change. Separately, pointing `fp8` and `mxfp8` at the new arms leaves `FineGrainedFP8HfQuantizer` orphaned -- defined, never imported, never mapped -- so nothing mis-routes and what remains is dead code someone else should decide about. Co-Authored-By: Claude Opus 5 --- src/transformers/core_model_loading.py | 2 +- .../distributed/sharding_utils.py | 15 +- .../distributed/tensor_parallel.py | 4 - src/transformers/integrations/deepgemm.py | 12 +- .../integrations/finegrained/__init__.py | 45 +++++ .../conversions.py} | 72 ++++---- .../{finegrained.py => finegrained/core.py} | 164 +++++++----------- src/transformers/integrations/moe.py | 16 +- src/transformers/integrations/mxfp4.py | 9 - src/transformers/integrations/nvfp4.py | 9 - .../deepseek_v4/modeling_deepseek_v4.py | 5 +- .../models/deepseek_v4/modular_deepseek_v4.py | 5 +- src/transformers/quantizers/auto.py | 18 +- .../quantizers/finegrained/__init__.py | 31 ++++ .../base.py} | 84 ++++----- .../blockfp8.py} | 30 ++-- .../mxfp4.py} | 10 +- .../quantizers/finegrained/mxfp8.py | 22 +++ .../nvfp4.py} | 26 ++- .../quantizers/quantizer_mxfp4.py | 13 +- .../quantizers/quantizer_nvfp4.py | 21 +-- src/transformers/utils/quantization_config.py | 141 ++++++++------- tests/kernels/test_deepgemm.py | 11 +- tests/kernels/test_finegrained.py | 84 +++++---- tests/tensor_parallel/test_tensor_parallel.py | 33 ---- tests/utils/test_core_model_loading.py | 84 +++++---- .../utils/test_distributed_sharding_utils.py | 19 ++ 27 files changed, 499 insertions(+), 486 deletions(-) create mode 100644 src/transformers/integrations/finegrained/__init__.py rename src/transformers/integrations/{finegrained_conversions.py => finegrained/conversions.py} (94%) rename src/transformers/integrations/{finegrained.py => finegrained/core.py} (87%) create mode 100644 src/transformers/quantizers/finegrained/__init__.py rename src/transformers/quantizers/{quantizer_finegrained.py => finegrained/base.py} (86%) rename src/transformers/quantizers/{quantizer_finegrained_blockfp8.py => finegrained/blockfp8.py} (61%) rename src/transformers/quantizers/{quantizer_finegrained_mxfp4.py => finegrained/mxfp4.py} (90%) create mode 100644 src/transformers/quantizers/finegrained/mxfp8.py rename src/transformers/quantizers/{quantizer_finegrained_nvfp4.py => finegrained/nvfp4.py} (90%) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1a0465c7f426..7d4ee840ee69 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1150,7 +1150,7 @@ def _internal_many_to_many_conversions() -> tuple[type[ConversionOps], ...]: whenever the integration is imported first (`cannot import name ... from partially initialized module`). Called once per `WeightConverter`, at construction. """ - from .integrations.finegrained_conversions import FineGrainedWeightGlobals + from .integrations.finegrained.conversions import FineGrainedWeightGlobals return ( ErnieFuseAndSplitTextVisionExperts, diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 5b8e9297ea77..4af642ca42ae 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -112,16 +112,11 @@ def shard_tensor( (mesh_dim, placement) for mesh_dim, placement in enumerate(self.placements) if hasattr(placement, "dim") ] - # A 0-dim tensor has no axis to shard, so every rank that owns it takes the whole value - # (`source[...]`, since a lazy safetensors slice rejects `source[()]`). - if not source_shape: - if tensor_idx is not None and not self._owns_expert(tensor_idx, dim_placements): - return None - return source[...].to(device=device, dtype=dtype) - # Dense path if tensor_idx is None: - if not dim_placements: + # nothing to slice: a 0-dim tensor has no axis, and no placement names one + # (`source[...]`, since a lazy safetensors slice rejects `source[()]`) + if not source_shape or not dim_placements: return source[...].to(device=device, dtype=dtype) # Determine for each tensor dimension, which type of sharding operations to apply (_StridedShard or Shard) and which rank to apply it to. @@ -164,6 +159,10 @@ def shard_tensor( if not self._owns_expert(tensor_idx, dim_placements): return None + # a 0-dim per-expert value has no axis to slice, so this rank takes it whole + if not source_shape: + return source[...].to(device=device, dtype=dtype) + normalized_dim_placements = [ (mesh_dim, placement, self._normalize_param_dim(placement.dim)) for mesh_dim, placement in dim_placements ] diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 40f39ecedeff..a2339932ea31 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -787,10 +787,6 @@ class ParallelInterface(GeneralInterface): "packed_rowwise": PackedRowwiseParallel(), "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), - # a companion of an expert-stacked (E, N, K) projection — block scales, bias, NVFP4 - # globals — shards on the projection's own axis while the experts stay whole, named - # for the weight split it follows. The built-in styles count axes from the END of the - # shape, which lands inside a swizzled scale grid's inner tile rather than on its rows. "moe_experts_colwise": MoEParamShard(Shard(1)), "moe_experts_packed_colwise": MoEParamShard(_StridedShard(dim=1, split_factor=2)), "moe_experts_rowwise": MoEParamShard(Shard(2)), diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 7f908b398846..ac2a2a90bae6 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -704,7 +704,7 @@ def _apply_post_expert_norm(self: torch.nn.Module, rows: torch.Tensor) -> torch. same place the reference forwards apply it. A no-op for a model that declares none.""" if not getattr(self, "has_post_expert_norm", False): return rows - return self._apply_post_norm(rows) + return self.post_expert_norm(rows) def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: @@ -719,15 +719,15 @@ def deepgemm_experts_guards( *, affine_scales: bool = False, sm100: bool = False, - post_expert_norm: bool = True, + supports_post_expert_norm: bool = True, stacked_gate_up: bool = False, ): """State an experts forward's requirements on the module, checked before any kernel work. Usable bare (`@deepgemm_experts_guards`) for an arm that only needs the common checks. - ``post_expert_norm=False`` refuses a model that declares a per-expert output norm — Mega MoE - only, whose fused reduce has no seam for it. ``affine_scales`` refuses the SWIZZLE_32_4_4 + ``supports_post_expert_norm=False`` refuses a model that declares a per-expert output norm — + Mega MoE only, whose fused reduce has no seam for it. ``affine_scales`` refuses the SWIZZLE_32_4_4 scales a module loaded for a triton backend holds, ``stacked_gate_up`` refuses the interleaved rows that same load produces, and ``sm100`` fails before the hub download + JIT when the device cannot serve these dtypes.""" @@ -737,7 +737,7 @@ def decorate(forward): def guarded(self, hidden_states, *args, **kwargs): _assert_bf16_hidden_states(hidden_states) _assert_dynamic_activations(self) - if not post_expert_norm: + if not supports_post_expert_norm: _assert_no_post_expert_norm(self) if affine_scales: _assert_affine_scales(self) @@ -967,7 +967,7 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) -@deepgemm_experts_guards(affine_scales=True, sm100=True, post_expert_norm=False, stacked_gate_up=True) +@deepgemm_experts_guards(affine_scales=True, sm100=True, supports_post_expert_norm=False, stacked_gate_up=True) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, diff --git a/src/transformers/integrations/finegrained/__init__.py b/src/transformers/integrations/finegrained/__init__.py new file mode 100644 index 000000000000..eb08ca99af47 --- /dev/null +++ b/src/transformers/integrations/finegrained/__init__.py @@ -0,0 +1,45 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The fine-grained quantization integration: `core` holds the modules and their forwards, +`conversions` the loader ops that fill them. `conversions` is imported by path, so that reading +a checkpoint's layout never drags the kernels in.""" + +from .core import ( + ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedGroupedLinear, + FineGrainedLinear, + assert_modules_are_quantized, + disable_deepgemm_on_multi_device, + finegrained_linear, + load_finegrained_kernel, + replace_with_finegrained_embedding, + replace_with_finegrained_layer, +) + + +__all__ = [ + "ALL_FINEGRAINED_EXPERTS_FUNCTIONS", + "FineGrainedEmbedding", + "FineGrainedExperts", + "FineGrainedGroupedLinear", + "FineGrainedLinear", + "assert_modules_are_quantized", + "disable_deepgemm_on_multi_device", + "finegrained_linear", + "load_finegrained_kernel", + "replace_with_finegrained_embedding", + "replace_with_finegrained_layer", +] diff --git a/src/transformers/integrations/finegrained_conversions.py b/src/transformers/integrations/finegrained/conversions.py similarity index 94% rename from src/transformers/integrations/finegrained_conversions.py rename to src/transformers/integrations/finegrained/conversions.py index e21067ea9336..487fb6789a8b 100644 --- a/src/transformers/integrations/finegrained_conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -24,9 +24,9 @@ import torch -from ..core_model_loading import ConversionOps, Interleave -from ..quantizers.quantizers_utils import get_module_from_name -from .finegrained import ( +from ...core_model_loading import ConversionOps, Interleave +from ...quantizers.quantizers_utils import get_module_from_name +from .core import ( _FP8_DTYPE, _FP8_MAX, _FP8_MIN, @@ -54,7 +54,7 @@ def _held_scale(model, full_layer_name, key): """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" - param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[2] + param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[-1] if not param_name.endswith("_scale_inv") or model is None: return None, None module = model.get_submodule(full_layer_name.rpartition(".")[0]) @@ -63,11 +63,25 @@ def _held_scale(model, full_layer_name, key): return module, getattr(module, param_name, None) -def _as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: +def as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) +def _global_role(key: str) -> tuple[str, str]: + """What a global-scale key holds, for a checkpoint key and a module target alike: which + projection it belongs to, and which of modelopt's two levels it is — the weight's second-level + global (``weight_scale_2``) or the calibrated activation one (``input_scale``). + + This runs while the converters are being BUILT, against raw checkpoint keys, so the vLLM-style + `w1`/`w2`/`w3` names are still in play and have not been renamed to `gate/down/up_proj` yet. + Matched on path SEGMENTS: the default is gate_up, so a loose `"w2" in key` would silently + bucket a down global as one.""" + down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) + level = "input" if "input_scale" in key or "input_global" in key else "weight" + return ("down" if down else "gate_up"), level + + class _FineGrainedOp(ConversionOps): """Base for the ops below. Each carries the quantizer that built it, which is how an op reaches the config (the activation format, the block size) while it runs. @@ -124,12 +138,12 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= modules = model.modules() if model is not None else () container = next((c for m in modules if (c := getattr(m, "scale_container_dtype", None))), None) if container is not None and value.dtype == _get_ue8m0_dtype(): - value = _as_container(value, container) + value = as_container(value, container) else: module, held = _held_scale(model, full_layer_name, key) if held is not None and value.dtype != held.dtype: module.scale_container_dtype = value.dtype - value = _as_container(value, held.dtype) + value = as_container(value, held.dtype) out[key] = value return out @@ -230,7 +244,8 @@ def convert(self, input_dict, target_patterns=None, model=None, **kwargs): # input global gives back (keeping the requantized intermediate on the calibrated range) if stack is not None and globals_[stack].shape[1] == 2: self._assert_no_gate_up_bias(model) - gate, up = globals_[stack][:, 0], globals_[stack][:, 1] + gate, up = globals_[stack].chunk(2, dim=1) + gate, up = gate.squeeze(1), up.squeeze(1) globals_[stack] = gate ratio = (up / gate)[:, None] for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): @@ -252,17 +267,6 @@ def _assert_no_gate_up_bias(model) -> None: ) -def _global_role(key: str) -> tuple[str, str]: - """What a global-scale key holds, for a checkpoint key and a module target alike: which - projection it belongs to, and which of modelopt's two levels it is — the weight's second-level - global (``weight_scale_2``) or the calibrated activation one (``input_scale``). Matched on path - SEGMENTS, since the default is gate_up and a loose `"w2" in key` would silently bucket a down - global as one.""" - down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) - level = "input" if "input_scale" in key or "input_global" in key else "weight" - return ("down" if down else "gate_up"), level - - class FineGrainedInputScales(_FineGrainedOp): """A calibrated checkpoint's ``input_scale`` in the layout the module holds: one value per quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max since @@ -276,13 +280,9 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): value = torch.stack(value, dim=0) if isinstance(value, list) else value values.append(value.float()) stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) - # A scale is a MAGNITUDE (`amax(|x|) / 448`), so reduce over absolute values and keep it - # off zero. A signed reduce returns whatever the largest signed entry is, and the two - # degenerate results it admits are both fatal downstream: a negative scale flips the sign - # of every dequantized row, and a zero one divides the activations by zero — NaN logits - # on a SINGLE device, before any sharding is involved. Real calibrated checkpoints ship - # positive scales, so this is a no-op for them; it is the models whose scales are built - # rather than measured that reach here with a sign or a zero. + # A scale is a MAGNITUDE, so reduce over absolute values and keep it off zero: a negative + # flips every dequantized row's sign and a zero divides the activations by zero. A no-op + # for a real calibrated checkpoint; it catches scales that were built rather than measured. per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1).clamp(min=1e-12) one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} @@ -314,14 +314,12 @@ def reverse_op(self) -> ConversionOps: class FineGrainedQuantize(_FineGrainedOp): """Quantize a full-precision weight on load into the format its module holds, emitting the - scale (and the NVFP4 global) in the module's layout — the held container dtype, the swizzled - artifact where the module holds one. Block-FP8 is computed in torch (per-block E4M3, fp32 or - power-of-two UE8M0 inverse scales); the group formats run the kernels' row-wise quantizers - (``mxfp8_act_quant`` / ``mxfp4_act_quant`` / ``nvfp4_act_quant``), NVFP4 after normalizing by the - canonical per-tensor / per-expert global ``amax / (6 * 448)``. - A tensor that is not a finegrained module's weight passes through — `_weight_holder` decides, - since rank alone does not (an expert bias is 2-D). Without a model, a direct invocation has - nothing to ask and quantizes block-FP8 from the config.""" + scale (and the NVFP4 global) in that module's layout. + + Block-FP8 is computed in torch; the group formats run the kernels' row-wise quantizers, NVFP4 + after normalizing by the canonical global ``amax / (6 * 448)``. A tensor that is not a + finegrained module's weight passes through — `_weight_holder` decides, since rank alone does + not (an expert bias is 2-D).""" def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: result: dict[str, torch.Tensor] = {} @@ -407,7 +405,7 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc f"accelerator, but the weight is on {value.device.type}" ) weight, scale, global_scale = self._quantize_group(value, format_spec) - scale = _as_container(scale, held_scale.dtype) + scale = as_container(scale, held_scale.dtype) if held_scale.ndim == 5: scale = load_finegrained_kernel().swizzle_mx_scales(scale) out = {key: weight, f"{prefix}_scale_inv": scale} @@ -462,7 +460,7 @@ def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(padded.shape) - # a no-op slice on a shape that tiled, and `contiguous` then returns the tensor itself + # a no-op when the shape tiled return quantized[..., :rows, :cols].contiguous(), inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) @staticmethod @@ -642,7 +640,7 @@ def keep_swizzle_reverse_for_save(model, hf_quantizer) -> None: broader pattern would take keys away from the real converter and skip the rest of its chain. """ - from ..core_model_loading import WeightConverter + from ...core_model_loading import WeightConverter conversions = list(getattr(model, "_weight_conversions", None) or []) covered = { diff --git a/src/transformers/integrations/finegrained.py b/src/transformers/integrations/finegrained/core.py similarity index 87% rename from src/transformers/integrations/finegrained.py rename to src/transformers/integrations/finegrained/core.py index 5795e9c757aa..4386cca0e86f 100644 --- a/src/transformers/integrations/finegrained.py +++ b/src/transformers/integrations/finegrained/core.py @@ -28,19 +28,19 @@ import torch import torch.nn as nn -from ..activations import ACT2FN -from ..quantizers.quantizers_utils import should_convert_module -from ..utils import logging -from ..utils.import_utils import is_kernels_available -from .deepgemm import ( +from ...activations import ACT2FN +from ...quantizers.quantizers_utils import should_convert_module +from ...utils import logging +from ...utils.import_utils import is_kernels_available +from ..deepgemm import ( deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, deepgemm_fp8_fp4_megamoe_experts_forward, is_sm100, prefers_deepgemm_linear, ) -from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel -from .moe import ExpertsInterface, use_experts_implementation +from ..hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel +from ..moe import ExpertsInterface, use_experts_implementation logger = logging.get_logger(__name__) @@ -364,9 +364,6 @@ def finegrained_linear( "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely." ) - # Triton arm: fused act-quant + matmul, then optional bias add. `activation_scale` None - # quantizes inline per token, a scalar quantizes statically against it; the two NVFP4 globals - # both come back on the accumulator. kernel = load_finegrained_kernel() original_shape = input.shape output = kernel.matmul_2d( @@ -545,10 +542,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _moe_operands(kernel, module) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their scales (held swizzled for dot_scaled chains), the per-expert globals and biases, - and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the module's - own GLU as a callable the kernels run on the host, so any activation works without a kernel - change.""" - up = "gate_up_proj" if module.has_gate else "up_proj" + and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the + module's own GLU as a callable the kernels run on the host, so any activation works without a + kernel change.""" + up = module.gate_up_name act_name = module.act_fn_name fused = act_name in kernel.get_supported_act_fns() and (module.swiglu_alpha is None or act_name == "silu") @@ -557,16 +554,14 @@ def _moe_operands(kernel, module) -> dict: else: act_fn = module._apply_gate if module.has_gate else module.act_fn - # a `get_supported_norms()` NAME is folded into the routing-weighted reduce; any other - # callable runs on the routed rows before it. The name is only usable where nothing has to - # happen in between: fusing puts the norm INSIDE the launch that produced the rows, so under - # intra-expert TP — where they are a partial sum over the sharded intermediate — the - # all-reduce would have nowhere to land, and the module's wrapped forward must do it first. + # A named norm is fused into the reduce; anything else runs as a callable on the routed rows. + # Fusing is off under intra-expert TP: those rows are a partial sum, so the all-reduce the + # module's wrapped forward does has to happen first (`_hf_tp_input_reduce` marks it). post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 if module.has_post_expert_norm: fusable = module.post_expert_norm_name in kernel.get_supported_norms() if not fusable or getattr(module.post_expert_norm, "_hf_tp_input_reduce", False): - post_expert_norm = module._apply_post_norm + post_expert_norm = module.post_expert_norm else: norm = module.post_expert_norm post_expert_norm, norm_weight = module.post_expert_norm_name, norm.weight @@ -575,9 +570,6 @@ def _moe_operands(kernel, module) -> dict: eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) norm_eps = float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps - # The kernels always say `gate_up_proj`; the module names that weight `up_proj` when the model - # has no gate, which is the only reason these are not plain attribute reads. A slot the module - # does not hold reads back as None — the argument the kernels take for an absent one. return { **{ f"{kernel_proj}{slot}": getattr(module, f"{held_proj}{slot}") @@ -599,22 +591,17 @@ def _moe_operands(kernel, module) -> dict: } -def _fused_experts_forward(module, kernel_forward: str, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: - """The kernels' fused MoE chain (gate_up with the fused GLU epilogue and intermediate requant - where supported -> down -> the routing-weighted top-k reduce) over the module's tensors.""" - kernel = load_finegrained_kernel() - return getattr(kernel, kernel_forward)(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, module)) - - def finegrained_batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: """Batched (decode) experts forward: one program per routed row.""" - return _fused_experts_forward(self, "moe_fused_batched", hidden_states, top_k_index, top_k_weights) + kernel = load_finegrained_kernel() + return kernel.moe_fused_batched(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, self)) def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: """Grouped (prefill) experts forward: one on-device routing pass, then the chain over the expert-sorted schedule, scattering back to routed rows (EP-sentinel rows skipped).""" - return _fused_experts_forward(self, "moe_fused_grouped", hidden_states, top_k_index, top_k_weights) + kernel = load_finegrained_kernel() + return kernel.moe_fused_grouped(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, self)) class FineGrainedExperts(_FineGrainedModule, nn.Module): @@ -647,7 +634,6 @@ def __init__( # The model's own gate|up row order (transformers' experts flag): stacked ``[gate; up]`` # or, like GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``. It describes the # CHECKPOINT only — what this module holds is `holds_interleaved_gate_up` below. - self.is_concatenated = is_concatenated self.activation_format = activation_format self.activation_scheme = activation_scheme self.hidden_dim = _first_attr(config, "moe_hidden_size", "hidden_size") @@ -658,49 +644,32 @@ def __init__( self.swiglu_limit = getattr(config, "swiglu_limit", None) self.act_fn = ACT2FN[self.act_fn_name] - # the format of THIS module's group, as `replace_with_finegrained_layer` resolved it. A - # mixed checkpoint (DeepSeek-V4: mxfp4 experts, block-FP8 linears) is two groups, so the - # experts no longer read a `config.expert_dtype` side-channel to find out. self.weight_format = weight_format format_spec, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) - # the format decides the second level: None means there is none, else the dtype it is in self.global_scale_dtype = format_spec.global_scale_dtype - # What the module HOLDS, as the loader's conversion ops deliver it (and their reverses - # restore for the checkpoint): gate|up rows in the kernels' interleaved order unless the - # backend packs gate|up itself (DeepGEMM Mega MoE), block scales swizzled where the - # scaled-MMA reads them. `is_concatenated` above is the checkpoint's own order: the - # loader interleaves where the two disagree, and everything at runtime reads this one. - # FROZEN at load, not a property over `_experts_implementation`: it describes the bytes - # the tensors now hold, so re-deriving it would claim a layout they do not have. + # The layouts: what the CHECKPOINT ships (`is_concatenated`) against what the module now + # HOLDS, which the loader's conversion ops interleave and swizzle into where the two + # disagree. Frozen at load rather than derived on read, since it describes the bytes. + self.is_concatenated = is_concatenated impl = getattr(config, "_experts_implementation", None) self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" swizzled = _holds_swizzled_scales(config, format_spec, activation_format) - up_name = "gate_up_proj" if self.has_gate else "up_proj" + self.gate_up_name = "gate_up_proj" if self.has_gate else "up_proj" up_rows = (2 if self.has_gate else 1) * self.intermediate_dim - # gate_up's activation global is ONE value — its rows are the pre-routing hidden states, - # quantized once before routing; down's rows belong to an expert each. - for proj, rows, in_dim, min_scale_out, input_globals in ( - (up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, 1), - ("down_proj", self.hidden_dim, self.intermediate_dim, 1, self.num_experts), - ): - self._register_projection( - proj, - rows, - in_dim, - min_scale_out, - input_globals, - (format_spec, scale_dtype, scale_group, swizzled), - ) + storage = (format_spec, scale_dtype, scale_group, swizzled) + # gate_up takes ONE activation global — its rows are the pre-routing hidden states, + # quantized once before routing; down's rows belong to an expert each + gate_up_globals, down_globals = 1, self.num_experts + self._register_projection( + self.gate_up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, gate_up_globals, storage + ) + self._register_projection("down_proj", self.hidden_dim, self.intermediate_dim, 1, down_globals, storage) - # The model's per-expert output norm, carried over by the swap. A submodule slot, not a - # class attribute: a class attribute of the same name shadows what - # `nn.Module.__setattr__` registers, so the norm would never be read back. + # the model's per-expert output norm, filled by the swap self.post_expert_norm = None self.has_post_expert_norm = False - # The form it implements, when the model names one it wants a backend to FUSE - # (`get_supported_norms`); unnamed norms still run through `_apply_post_norm`. self.post_expert_norm_name = None def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, storage): @@ -722,24 +691,14 @@ def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, # down to that). The activation's is the checkpoint's `input_scale`. two_level = self.global_scale_dtype is not None calibrated = two_level and self.activation_format != "bf16" - for name, tensor in ( - ("bias", bias), - ( - "weight_global_scale", - torch.ones(self.num_experts, dtype=self.global_scale_dtype) if two_level else None, - ), - ("input_global_scale", torch.ones(input_globals, dtype=torch.float32) if calibrated else None), - ( - "activation_scale", - torch.ones(self.num_experts, dtype=torch.float32) if self.activation_scheme == "static" else None, - ), - ): - _set_optional_parameter(self, f"{proj}_{name}", tensor) - - def _apply_post_norm(self, expert_out: torch.Tensor) -> torch.Tensor: - """The standard hook every backend applies: the model's norm on one expert application's - rows, before the routing weights.""" - return self.post_expert_norm(expert_out) + static = self.activation_scheme == "static" + weight_global = torch.ones(self.num_experts, dtype=self.global_scale_dtype) if two_level else None + input_global = torch.ones(input_globals, dtype=torch.float32) if calibrated else None + activation = torch.ones(self.num_experts, dtype=torch.float32) if static else None + _set_optional_parameter(self, f"{proj}_bias", bias) + _set_optional_parameter(self, f"{proj}_weight_global_scale", weight_global) + _set_optional_parameter(self, f"{proj}_input_global_scale", input_global) + _set_optional_parameter(self, f"{proj}_activation_scale", activation) def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: # interleaved rows -> alternating output columns, the same split the fused epilogue does @@ -767,7 +726,8 @@ def forward( expert_mask = expert_mask.permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1) - for expert_idx in expert_hit: + # one sync here instead of one per iteration: indexing with a 0-dim cuda tensor syncs + for expert_idx in expert_hit.tolist(): if expert_idx == self.num_experts: continue @@ -776,26 +736,23 @@ def forward( gate_up_act_scale = ( self.gate_up_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None ) - proj_out = self.linear( - current_state, - "gate_up_proj" if self.has_gate else "up_proj", - expert_idx, - activation_scale=gate_up_act_scale, + proj_out = self.expert_linear( + current_state, self.gate_up_name, expert_idx, activation_scale=gate_up_act_scale ) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) down_act_scale = ( self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None ) - proj_out = self.linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) + proj_out = self.expert_linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) if self.has_post_expert_norm: - proj_out = self._apply_post_norm(proj_out) + proj_out = self.post_expert_norm(proj_out) routing_weights = top_k_weights[token_idx, top_k_pos, None] weighted_out = proj_out * routing_weights.to(proj_out.dtype) final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) return final_hidden_states.to(hidden_states.dtype) - def linear( - self, input: torch.Tensor, proj: str, expert_idx, activation_scale: torch.Tensor | None = None + def expert_linear( + self, input: torch.Tensor, proj: str, expert_idx: int, activation_scale: torch.Tensor | None = None ) -> torch.Tensor: """One expert's ``proj`` as a dense linear: the same weight, scale (one expert's slice of a swizzled stack is the ``(1, ...)`` artifact ``matmul_2d`` reads), bias, NVFP4 global and @@ -804,9 +761,16 @@ def linear( bias = getattr(self, f"{proj}_bias") bias = bias[expert_idx] if bias is not None else None scale = getattr(self, f"{proj}_scale_inv") - weight_globals = getattr(self, f"{proj}_weight_global_scale") - input_globals = getattr(self, f"{proj}_input_global_scale") - weight_global = None if weight_globals is None else weight_globals[expert_idx] + weight_global_scale = getattr(self, f"{proj}_weight_global_scale") + weight_global_scale = weight_global_scale[expert_idx] if weight_global_scale is not None else None + # None under weight-only, which registers no input global at all. `% numel` picks this + # expert's, or the single one gate_up holds because its rows are quantized pre-routing. + input_global_scale = getattr(self, f"{proj}_input_global_scale") + input_global_scale = ( + input_global_scale.reshape(-1)[expert_idx % input_global_scale.numel()].reshape(1) + if input_global_scale is not None + else None + ) return finegrained_linear( input, weight, @@ -815,12 +779,8 @@ def linear( bias=bias, activation_scale=activation_scale, allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=weight_global, - input_global_scale=( - None - if input_globals is None or self.activation_format == "bf16" - else input_globals.reshape(-1)[expert_idx if input_globals.numel() > 1 else 0].reshape(1) - ), + weight_global_scale=weight_global_scale, + input_global_scale=input_global_scale, activation_format=self.activation_format, ) @@ -923,7 +883,7 @@ def _quantized_experts(module: nn.Module, model: nn.Module, storage: dict) -> nn ) experts = experts_class(config=getattr(module, "config", model.config.get_text_config()), **storage, **flags) if getattr(module, "post_expert_norm", None) is not None: - # `_apply_post_norm` is the standard hook and needs no rebinding; a NAMED form is what a + # the norm is called directly wherever it is not fused; a NAMED form is what a # backend can fuse. `use_experts_implementation` owns the flag and sets it from the class # declaration, so this updates it rather than deriving it. experts.post_expert_norm = module.post_expert_norm diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 9f23817d91df..af95f980c9bd 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -168,7 +168,7 @@ def batched_mm_experts_forward( # `getattr`: an experts module built without the decorator has no such attribute, and this # forward is the shared one every backend adapts onto. if getattr(self, "has_post_expert_norm", False): - proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) + proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) # Apply routing weights if sentinel_mask is not None: @@ -485,7 +485,7 @@ def grouped_mm_experts_forward( # `getattr`: an experts module built without the decorator has no such attribute, and this # forward is the shared one every backend adapts onto. if getattr(self, "has_post_expert_norm", False): - proj_out = self._apply_post_norm(proj_out) # (S, hidden_dim) + proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) @@ -574,9 +574,8 @@ def use_experts_implementation( Whether the experts use a gating mechanism or not. Whether it has gate_up_proj weights or just up_proj weights. has_post_expert_norm (`bool`, *optional*, defaults to `False`): - Whether the experts normalize the down output before the routing weights. The class - must then define `_apply_post_norm(self, expert_out)` — the standard name every - backend applies, so none of them needs to know the model's own math. + Whether the experts normalize the down output before the routing weights. The norm + itself is the module under `post_expert_norm`, which every backend calls directly. Returns: `type[torch.nn.Module]`: The modified experts class. @@ -605,13 +604,6 @@ def forward(self, *args, **kwargs): if not hasattr(experts_class, "_apply_gate"): experts_class._apply_gate = _default_apply_gate - if has_post_expert_norm and not hasattr(experts_class, "_apply_post_norm"): - raise TypeError( - f"{experts_class.__name__} declares a post-expert norm but does not " - "define `_apply_post_norm(self, expert_out)`, which applies it to one expert " - "application's rows before the routing weights." - ) - experts_class.__init__ = __init__ experts_class.forward = forward return experts_class diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index b5efea30f582..acfe79de4f93 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import warnings from ..utils import is_torch_available, logging @@ -29,14 +28,6 @@ logger = logging.get_logger(__name__) -warnings.warn( - "mxfp4 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - FP4_VALUES = [ +0.0, +0.5, diff --git a/src/transformers/integrations/nvfp4.py b/src/transformers/integrations/nvfp4.py index 51570d197768..476e459df405 100644 --- a/src/transformers/integrations/nvfp4.py +++ b/src/transformers/integrations/nvfp4.py @@ -14,7 +14,6 @@ from __future__ import annotations import functools -import warnings from collections.abc import Callable from dataclasses import dataclass @@ -34,14 +33,6 @@ import torch.nn as nn -warnings.warn( - "nvfp4 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.integrations.finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - # NVFP4 uses one FP8 scale factor for every 16 consecutive weight values. SF_VEC_SIZE = 16 diff --git a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py index 9d3518ee7662..c9e3bff56449 100644 --- a/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modeling_deepseek_v4.py @@ -53,13 +53,12 @@ def __init__(self, hidden_size, eps: float = 1e-6) -> None: self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Return the input dtype. V4 pins these norms to FP32 (`_keep_in_fp32_modules_strict`), - so `weight * hidden_states` promotes the result and hands FP32 to the compressor - projections, which ship BF16. A no-op where the weight already matches.""" input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # back to the input dtype: V4 pins these norms to FP32, and `weight * hidden_states` + # would otherwise hand FP32 to the compressor projections, which ship BF16 return (self.weight * hidden_states.to(input_dtype)).to(input_dtype) def extra_repr(self): diff --git a/src/transformers/models/deepseek_v4/modular_deepseek_v4.py b/src/transformers/models/deepseek_v4/modular_deepseek_v4.py index 159491590fdb..f71d0782b1fc 100644 --- a/src/transformers/models/deepseek_v4/modular_deepseek_v4.py +++ b/src/transformers/models/deepseek_v4/modular_deepseek_v4.py @@ -65,13 +65,12 @@ def apply_rotary_pos_emb( class DeepseekV4RMSNorm(DeepseekV3RMSNorm): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Return the input dtype. V4 pins these norms to FP32 (`_keep_in_fp32_modules_strict`), - so `weight * hidden_states` promotes the result and hands FP32 to the compressor - projections, which ship BF16. A no-op where the weight already matches.""" input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + # back to the input dtype: V4 pins these norms to FP32, and `weight * hidden_states` + # would otherwise hand FP32 to the compressor projections, which ship BF16 return (self.weight * hidden_states.to(input_dtype)).to(input_dtype) diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 508b63481cad..48047595d1aa 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -46,6 +46,12 @@ VptqConfig, ) from .base import HfQuantizer +from .finegrained import ( + FineGrainedBlockFp8HfQuantizer, + FineGrainedMxfp4HfQuantizer, + FineGrainedMxfp8HfQuantizer, + FineGrainedNvfp4HfQuantizer, +) from .quantizer_aqlm import AqlmHfQuantizer from .quantizer_auto_round import AutoRoundQuantizer from .quantizer_awq import AwqQuantizer @@ -55,10 +61,6 @@ from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer from .quantizer_eetq import EetqHfQuantizer from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer -from .quantizer_finegrained import FineGrainedHfQuantizer -from .quantizer_finegrained_blockfp8 import FineGrainedBlockFp8HfQuantizer -from .quantizer_finegrained_mxfp4 import FineGrainedMxfp4HfQuantizer -from .quantizer_finegrained_nvfp4 import FineGrainedNvfp4HfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer from .quantizer_gemma import GemmaQuantizer @@ -97,12 +99,10 @@ "gguf": GgufHfQuantizer, "metal": MetalHfQuantizer, "auto-round": AutoRoundQuantizer, - # An arm exists where a PRODUCER's key layout needs one — GPT-OSS's packed `_blocks`, - # modelopt's two-level scales, the calibrated `input_scale`. MXFP8 ships the plain - # `weight` / `weight_scale_inv` pair, so the base serves it. Which FORMAT each module is - # quantized in is the config's groups, not the arm. + # one arm per checkpoint PRODUCER, since that is what decides the key layout; which FORMAT + # each module is in comes from the config's groups, not from the arm "fp8": FineGrainedBlockFp8HfQuantizer, - "mxfp8": FineGrainedHfQuantizer, # no distinct key layout — the base is its handler + "mxfp8": FineGrainedMxfp8HfQuantizer, "mxfp4": FineGrainedMxfp4HfQuantizer, "nvfp4": FineGrainedNvfp4HfQuantizer, "modelopt": FineGrainedNvfp4HfQuantizer, diff --git a/src/transformers/quantizers/finegrained/__init__.py b/src/transformers/quantizers/finegrained/__init__.py new file mode 100644 index 000000000000..b7574e4b876a --- /dev/null +++ b/src/transformers/quantizers/finegrained/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""One quantizer per checkpoint PRODUCER. `base` serves the plain `weight` / `weight_scale_inv` +pair; an arm exists where a producer's key layout needs one, and says which format each module +is in through the config's groups, not through the arm.""" + +from .base import FineGrainedHfQuantizer +from .blockfp8 import FineGrainedBlockFp8HfQuantizer +from .mxfp4 import FineGrainedMxfp4HfQuantizer +from .mxfp8 import FineGrainedMxfp8HfQuantizer +from .nvfp4 import FineGrainedNvfp4HfQuantizer + + +__all__ = [ + "FineGrainedBlockFp8HfQuantizer", + "FineGrainedHfQuantizer", + "FineGrainedMxfp4HfQuantizer", + "FineGrainedMxfp8HfQuantizer", + "FineGrainedNvfp4HfQuantizer", +] diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/finegrained/base.py similarity index 86% rename from src/transformers/quantizers/quantizer_finegrained.py rename to src/transformers/quantizers/finegrained/base.py index 211a5df5a0b9..8b2664540fc0 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -24,23 +24,24 @@ import re from typing import TYPE_CHECKING -from ..utils import ( +from ...utils import ( is_accelerate_available, is_torch_available, is_torch_distributed_available, is_torch_xpu_available, logging, ) -from .base import HfQuantizer -from .quantizers_utils import get_module_from_name +from ...utils.quantization_config import groups_with_expert_dtype +from ..base import HfQuantizer +from ..quantizers_utils import get_module_from_name if is_torch_available(): import torch if TYPE_CHECKING: - from ..modeling_utils import PreTrainedModel - from ..utils.quantization_config import FineGrainedConfig + from ...modeling_utils import PreTrainedModel + from ...utils.quantization_config import FineGrainedConfig logger = logging.get_logger(__name__) @@ -64,9 +65,18 @@ def _default_activation_format(self) -> str | None: kernels' weight-native choice; the MXFP4 arm overrides this to run weight-only.""" return None - def _assert_dequantize_supported(self) -> None: + @property + def supports_dequantize(self) -> bool: """Whether `dequantize=True` can fold this format's scales into a full-precision weight. - Every one-level format can; the NVFP4 arm overrides this to refuse.""" + Every one-level format can; the NVFP4 arm says no.""" + return True + + def _assert_dequantize_supported(self) -> None: + if not self.supports_dequantize: + raise NotImplementedError( + f"`dequantize=True` is not supported for {self._quant_method()} checkpoints. Load them " + "quantized on a GPU that serves the format, or start from a bf16 checkpoint." + ) def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): @@ -123,7 +133,7 @@ def validate_environment(self, *args, **kwargs): def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: # `FineGrainedGroupedLinear` subclasses `FineGrainedLinear`, so the tuple covers it implicitly. - from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear + from ...integrations.finegrained import FineGrainedExperts, FineGrainedLinear module, tensor_name = get_module_from_name(model, param_name) if isinstance(module, (FineGrainedLinear, FineGrainedExperts)): @@ -150,7 +160,7 @@ def _normalize_modules_to_not_convert(self, model: "PreTrainedModel"): if not skip: return - from ..conversion_mapping import get_model_conversion_mapping + from ...conversion_mapping import get_model_conversion_mapping renamings = get_model_conversion_mapping(model) remapped = [] @@ -166,12 +176,13 @@ def _process_model_before_weight_loading( model: "PreTrainedModel", **kwargs, ): - from ..integrations.finegrained import replace_with_finegrained_embedding, replace_with_finegrained_layer + from ...integrations.finegrained import replace_with_finegrained_embedding, replace_with_finegrained_layer self._normalize_modules_to_not_convert(model) # the one place both configs are in hand: a legacy checkpoint declares a second format # for its experts on the MODEL config, and it becomes a group here - self.quantization_config.split_out_experts(getattr(model.config.get_text_config(), "expert_dtype", None)) + expert_dtype = getattr(model.config.get_text_config(), "expert_dtype", None) + self.quantization_config.groups = groups_with_expert_dtype(self.quantization_config.groups, expert_dtype) if self.quantization_config.activation_format is None: self.quantization_config.activation_format = self._default_activation_format() if self.quantization_config.activation_format == "bf16": @@ -194,8 +205,8 @@ def _process_model_before_weight_loading( ) def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device - from ..integrations.finegrained_conversions import keep_swizzle_reverse_for_save + from ...integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device + from ...integrations.finegrained.conversions import keep_swizzle_reverse_for_save assert_modules_are_quantized(model) disable_deepgemm_on_multi_device(model) @@ -207,14 +218,10 @@ def update_tp_plan(self, config): # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes on # top of any plan written above (e.g. the Qwen3 dense plan). Models carry the experts # mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite both. - from ..integrations.finegrained import FineGrainedExperts + from ...integrations.finegrained import FineGrainedExperts impl = getattr(config, "_experts_implementation", None) layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) - # A multimodal model keeps its experts' plans on a SUB-config: the outer one carries a - # few projector entries and no `base_model_ep_plan` at all, so reading only what we were - # handed adds no companion anywhere while the weights still shard from the sub-config's - # own plan — the scale stays whole against a sharded weight. sub_configs = [c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None))] for plan_owner, plan_attr in ( (owner, attr) for owner in (config, *sub_configs) for attr in ("base_model_tp_plan", "base_model_ep_plan") @@ -222,11 +229,9 @@ def update_tp_plan(self, config): base_plan = getattr(plan_owner, plan_attr, None) or {} updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} - # Every companion beside a projection weight — block scales, bias, NVFP4 globals, a - # static activation scale — is expert-indexed too, so each shards with the weight. The - # matcher keys on the exact parameter name and falls back only to the owning MODULE, - # whose entry shards nothing, so without these the weight is this rank's expert slice - # while its scales are every rank's. gate_up's input global is per-tensor: replicated. + # every companion beside a projection weight is expert-indexed too, so each needs its + # own entry: the matcher keys on the parameter name and the module's entry shards + # nothing, which would leave this rank's weight slice beside every rank's scales for key, style in list(updated_plan.items()): # only the experts' own projections: a dense `self_attn.q_proj` is `colwise` too projection = key.rpartition(".")[2] @@ -240,28 +245,23 @@ def update_tp_plan(self, config): updated_plan.setdefault(f"{key}{suffix}", style) continue - # ...and only a STACKED experts module: `mlp.up_proj`, and the dense - # `mlp.shared_experts.up_proj`, end in the same word but are plain 2-D linears - # that colwise/rowwise already shard correctly. The model's own entry for the - # owning module is what distinguishes them. + # ...and only a STACKED experts module: a dense `mlp.shared_experts.up_proj` ends + # in the same word but colwise/rowwise already shards it correctly if updated_plan.get(key.rpartition(".")[0]) not in ("moe_tp_experts", "megamoe_experts"): continue if not is_torch_distributed_available(): continue # the styles below need a distributed build; without one, no TP - # Intra-expert TP: the experts stay whole and the projection's own axis splits, so - # the per-expert globals, activation scales and the down bias (added after the - # row-reduce) stay replicated — only the scale grid follows its weight. A - # per-tensor scale is `(E, 1, 1)`: no grid to split, so it stays replicated too. + # Intra-expert TP: the experts stay whole and the projection's own axis splits, + # so only the scale GRID follows its weight — the per-expert globals, activation + # scales and the down bias stay replicated, as does a per-tensor `(E, 1, 1)`. blocked = self.quantization_config.weight_block_size is not None if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): if blocked: updated_plan.setdefault(f"{key}_scale_inv", "moe_experts_rowwise") elif projection != "down_proj" and style in ("packed_colwise", "colwise"): - # Companions take the WEIGHT's split, which is the CHECKPOINT's layout: the - # interleave into the kernels' row order runs on each rank's shard afterwards. - # `packed_colwise` keeps a concatenated `[gate; up]`'s pairs together; - # `colwise` suits an already-interleaved layout (GPT-OSS). + # the CHECKPOINT's split, since the interleave into the kernels' row order + # runs on each rank's shard afterwards rows = "moe_experts_packed_colwise" if style == "packed_colwise" else "moe_experts_colwise" if blocked: updated_plan.setdefault(f"{key}_scale_inv", rows) @@ -284,7 +284,7 @@ def is_compileable(self) -> bool: return True def get_quantize_ops(self): - from ..integrations.finegrained_conversions import FineGrainedQuantize + from ...integrations.finegrained.conversions import FineGrainedQuantize return FineGrainedQuantize(self) @@ -304,8 +304,8 @@ def _dequantize_conversions(self): """Every quantized key folded back to a full-precision weight, from the `weight` + `weight_scale_inv` pair. The MXFP4 arm overrides this for GPT-OSS's `{proj}_blocks` + `{proj}_scales`.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import FineGrainedDequantize + from ...core_model_loading import WeightConverter + from ...integrations.finegrained.conversions import FineGrainedDequantize # anchored `weight$` so the scale keys land in the scale slots; the activation # scales are collected only to be dropped @@ -334,8 +334,8 @@ def update_weight_conversions(self, weight_conversions): - otherwise the expert converters get the module layout ops (``_with_expert_layout_ops``). :meth:`get_weight_conversions` is appended in every mode.""" - from ..core_model_loading import WeightConverter, WeightRenaming - from ..integrations.finegrained_conversions import FineGrainedDequantize + from ...core_model_loading import WeightConverter, WeightRenaming + from ...integrations.finegrained.conversions import FineGrainedDequantize scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") weight_conversions = [scale_rename, *weight_conversions] @@ -383,8 +383,8 @@ def _with_expert_layout_ops(self, weight_conversions): the checkpoint layout. Catch-all converters at the end take keys that arrive already under the fused names (a plain rename, a transformers-format checkpoint) and dense linears' scale keys, so they get the same treatment.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import ( + from ...core_model_loading import WeightConverter + from ...integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales, diff --git a/src/transformers/quantizers/quantizer_finegrained_blockfp8.py b/src/transformers/quantizers/finegrained/blockfp8.py similarity index 61% rename from src/transformers/quantizers/quantizer_finegrained_blockfp8.py rename to src/transformers/quantizers/finegrained/blockfp8.py index 5795ff07868e..15cc3c285614 100644 --- a/src/transformers/quantizers/quantizer_finegrained_blockfp8.py +++ b/src/transformers/quantizers/finegrained/blockfp8.py @@ -13,19 +13,15 @@ # limitations under the License. """The BLOCKFP8 arm of the fine-grained quantizer.""" -from .quantizer_finegrained import FineGrainedHfQuantizer +from .base import FineGrainedHfQuantizer class FineGrainedBlockFp8HfQuantizer(FineGrainedHfQuantizer): """Block-FP8 and per-tensor FP8 — the only formats that take a CALIBRATED activation scale. - A static scale is one value standing in for the quantization the other formats do inline, so - it only means anything where activations are quantized per tensor or per block. The group - formats hand the kernels a scale GRID instead (`As` is per group-32 / group-16), and NVFP4 - spends the checkpoint's `input_scale` as its second-level activation global, so neither can - consume one. Keeping the converters here is what makes that structural. - - Named `blockfp8` because `quantizer_finegrained_fp8` is the frozen legacy module. + The group formats hand the kernels a scale GRID instead, and NVFP4 spends `input_scale` as + its activation global, so neither can consume one. Named `blockfp8` because + `quantizer_finegrained_fp8` is the frozen legacy module. """ def get_weight_conversions(self): @@ -38,15 +34,12 @@ def get_weight_conversions(self): return [] def _static_activation_conversions(self): - """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in. One - value per quantized module: a dense linear brings one (Ministral-3), and a MoE brings one - per expert, each expert being its own quantized module (Mistral-4) — the gate|up pair - reduces to one per expert, both halves reading the same routed rows. Both expert layouts, - per expert per projection and already stacked per layer; a checkpoint matches one and the - other never fires. NVFP4 consumes ``input_scale`` as its activation GLOBAL instead - (the NVFP4 arm's converters) — a second level over a block scale, not the scale itself.""" - from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming - from ..integrations.finegrained_conversions import FineGrainedInputScales + """A calibrated checkpoint's ``input_scale`` onto the slot its module holds it in: one + value per quantized module, so a MoE brings one per expert. Both expert layouts are + covered — per expert per projection, and already stacked per layer — and a checkpoint + matches one while the other never fires.""" + from ...core_model_loading import MergeModulelist, WeightConverter, WeightRenaming + from ...integrations.finegrained.conversions import FineGrainedInputScales per_expert = [ WeightConverter( @@ -71,8 +64,7 @@ def _static_activation_conversions(self): ) for proj in ("gate_up_proj", "down_proj") ] - # every other calibrated module is a dense linear holding the one value itself — a plain - # rename. The lookahead keeps it off the expert keys above, which stack per layer instead + # a dense linear holds the one value itself; the lookahead keeps this off the expert keys dense = [ WeightRenaming( source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", diff --git a/src/transformers/quantizers/quantizer_finegrained_mxfp4.py b/src/transformers/quantizers/finegrained/mxfp4.py similarity index 90% rename from src/transformers/quantizers/quantizer_finegrained_mxfp4.py rename to src/transformers/quantizers/finegrained/mxfp4.py index 66b8d553f973..403b15548286 100644 --- a/src/transformers/quantizers/quantizer_finegrained_mxfp4.py +++ b/src/transformers/quantizers/finegrained/mxfp4.py @@ -15,7 +15,7 @@ and the weight-only default that layout implies. """ -from .quantizer_finegrained import FineGrainedHfQuantizer +from .base import FineGrainedHfQuantizer class FineGrainedMxfp4HfQuantizer(FineGrainedHfQuantizer): @@ -41,8 +41,8 @@ def get_weight_conversions(self): def _dequantize_conversions(self): """The blocks regroup into the packed rows and the exponent-byte scales come back out of them as bf16, in the `(E, hidden, 2I)` orientation the unquantized experts hold.""" - from ..core_model_loading import Transpose, WeightConverter - from ..integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks + from ...core_model_loading import Transpose, WeightConverter + from ...integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks return [ WeightConverter( @@ -57,8 +57,8 @@ def _mxfp4_conversions(self): """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed weight, the exponent bytes take the scale container op; then the usual layout ops. Bare patterns (no `experts.` scope), so they ride here rather than via `_with_expert_layout_ops`.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import ( + from ...core_model_loading import WeightConverter + from ...integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, FineGrainedPackedBlocks, FineGrainedScaleContainer, diff --git a/src/transformers/quantizers/finegrained/mxfp8.py b/src/transformers/quantizers/finegrained/mxfp8.py new file mode 100644 index 000000000000..1456fc72c14d --- /dev/null +++ b/src/transformers/quantizers/finegrained/mxfp8.py @@ -0,0 +1,22 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The MXFP8 arm of the fine-grained quantizer.""" + +from .base import FineGrainedHfQuantizer + + +class FineGrainedMxfp8HfQuantizer(FineGrainedHfQuantizer): + """MXFP8 ships the plain `weight` / `weight_scale_inv` pair, so it needs none of the key + surgery the other producers do and adds nothing to the base. It is named all the same, so + `AUTO_QUANTIZER_MAPPING` reads one class per key.""" diff --git a/src/transformers/quantizers/quantizer_finegrained_nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py similarity index 90% rename from src/transformers/quantizers/quantizer_finegrained_nvfp4.py rename to src/transformers/quantizers/finegrained/nvfp4.py index ca8a3c1e0b68..10fc37738726 100644 --- a/src/transformers/quantizers/quantizer_finegrained_nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -15,8 +15,8 @@ that the other formats do not. Everything else is inherited. """ -from ..utils import logging -from .quantizer_finegrained import FineGrainedHfQuantizer +from ...utils import logging +from .base import FineGrainedHfQuantizer logger = logging.get_logger(__name__) @@ -34,20 +34,18 @@ class FineGrainedNvfp4HfQuantizer(FineGrainedHfQuantizer): `weight_scale` / `weight_scale_2`. """ - def _assert_dequantize_supported(self) -> None: - raise NotImplementedError( - "`dequantize=True` is not supported for nvfp4 checkpoints: their two-level scales " - "cannot be folded into a full-precision weight by this path. Load them quantized on a " - "GPU that serves NVFP4, or start from a bf16 checkpoint." - ) + @property + def supports_dequantize(self) -> bool: + """No: NVFP4's two-level scales do not fold into a single full-precision weight here.""" + return False def get_weight_conversions(self): return self._nvfp4_conversions() if self.pre_quantized else [] def update_weight_conversions(self, weight_conversions): """The base chain, with every `.weight` source anchored and bitcast to the int8 view.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import FineGrainedViewPackedInt8 + from ...core_model_loading import WeightConverter + from ...integrations.finegrained.conversions import FineGrainedViewPackedInt8 if self.pre_quantized: updated = [] @@ -80,8 +78,8 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): parallelism keep only this rank's experts (`tensor_idx` in `core_model_loading`); without it every rank collects all E values and the forward asserts on the per-expert count. """ - from ..core_model_loading import Concatenate, MergeModulelist, WeightConverter - from ..integrations.finegrained_conversions import FineGrainedInputScales, FineGrainedWeightGlobals + from ...core_model_loading import Concatenate, MergeModulelist, WeightConverter + from ...integrations.finegrained.conversions import FineGrainedInputScales, FineGrainedWeightGlobals merge = [MergeModulelist(dim=0)] expert = r"mlp\.experts\..*\." @@ -129,8 +127,8 @@ def _nvfp4_fused_conversions(self, calibrated: bool): """Already stacked per layer, as the vLLM fused layout and Muse-Spark ship them — so no `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the `.weight`-anchored rule gives the per-expert layout.""" - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_conversions import ( + from ...core_model_loading import WeightConverter + from ...integrations.finegrained.conversions import ( FineGrainedInputScales, FineGrainedScaleContainer, FineGrainedViewPackedInt8, diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index fee0b8b6af68..f18ffac9a539 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import warnings from typing import TYPE_CHECKING from .base import HfQuantizer @@ -37,14 +36,6 @@ from ..core_model_loading import WeightConverter -warnings.warn( - "quantizer_mxfp4 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - logger = logging.get_logger(__name__) triton_kernels_hub = None @@ -59,6 +50,10 @@ class Mxfp4HfQuantizer(HfQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) + logger.warning_once( + "`Mxfp4Config` is frozen for backward compatibility and receives no new recipes. " + "`FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." + ) self.triton_kernels_hub = None def _lazy_import_kernels(self): diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index 1906afe944b2..78f2d6a494ec 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -13,15 +13,17 @@ # limitations under the License. from __future__ import annotations -import warnings from typing import TYPE_CHECKING -from ..utils import is_kernels_available, is_torch_available +from ..utils import is_kernels_available, is_torch_available, logging from ..utils.import_utils import KERNELS_MAX_VERSION, KERNELS_MIN_VERSION from .base import HfQuantizer from .quantizers_utils import get_module_from_name +logger = logging.get_logger(__name__) + + if is_torch_available(): import torch @@ -30,14 +32,6 @@ from ..utils.quantization_config import NVFP4Config -warnings.warn( - "quantizer_nvfp4 is frozen for backward compatibility and no longer " - "receives new recipes; the fine-grained quantization machinery lives in transformers.quantizers.quantizer_finegrained " - "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only).", - DeprecationWarning, -) - - def _as_cuda_device(device) -> torch.device | None: if isinstance(device, int): return torch.device("cuda", device) @@ -54,6 +48,13 @@ class NVFP4HfQuantizer(HfQuantizer): requires_calibration = False quantization_config: NVFP4Config + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + logger.warning_once( + "`NVFP4Config` is frozen for backward compatibility and receives no new recipes. " + "`FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." + ) + def validate_environment(self, device_map, **kwargs): if self.pre_quantized: raise ValueError("Loading pre-quantized NVFP4 checkpoints is not supported yet.") diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index be8b726ab993..cdf57095b900 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1687,10 +1687,18 @@ def post_init(self): raise TypeError("shapes must be a dict") -# What a modelopt export's `quant_algo` means in our vocabulary: the format it becomes, and the -# activation format that algo implies when the config does not name one. modelopt exports other -# algos (FP8, INT4 AWQ, W4A8, MXFP4); they are refused rather than silently mis-read. -_MODELOPT_ALGOS = {"NVFP4": (QuantizationMethod.NVFP4, "nvfp4")} +def subtree_patterns(globs: list[str]) -> list[str]: + """Glob-style module names ("model.layers.0*", "model.layers.1.*", bare "self_attn") as + regexes that match the named module and everything under it. + + Handed on as regexes the globs are wrong twice: the dots match any character, and the star is + greedy, so "model.layers.1.*" also takes layers 10-19. Anchor each at a path boundary instead. + """ + patterns = [] + for glob in globs: + subtree = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") + patterns.append(r"(?:^|.*\.)" + re.escape(subtree) + r"(\..*)?$") + return patterns _CATCH_ALL = [".*"] # the targets of a group that takes whatever no other group claimed @@ -1720,28 +1728,43 @@ def matches(self, module_name: str) -> bool: return any(re.search(target, module_name) for target in self.targets) -# (num_bits, type, group_size) as compressed-tensors / modelopt spell a format, to the name our -# kernels dispatch on. `group_size` is what separates the two 4-bit formats: nvfp4's block scale -# is E4M3 over 16 values, mxfp4's is E8M0 over 32. -_CT_FORMATS = { - (4, "float", 16): "nvfp4", - (4, "float", 32): "mxfp4", - (8, "float", 32): "mxfp8", - (8, "float", None): "fp8", -} +def groups_with_expert_dtype( + groups: dict[str, FineGrainedGroup], expert_dtype: str | None +) -> dict[str, FineGrainedGroup]: + """`groups` with a legacy `config.expert_dtype` folded into the group it means, unchanged if + there is nothing to fold. + DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL config, + because a flat quant config carries one format and cannot say "the experts differ". The + caller holds both configs; this only says what the two of them describe. + """ + if expert_dtype != "fp4" or len(groups) > 1: + return groups + flat = groups["default"] + return { + "experts": replace(flat, quant_method="mxfp4", activation_format="mxfp4", targets=[r"\.experts($|\.)"]), + "dense": flat, + } -def _group_from_config_groups(spec: dict) -> FineGrainedGroup | None: + +def group_from_config_groups(spec: dict) -> FineGrainedGroup | None: """One `config_groups` entry as a `FineGrainedGroup`, or `None` if it names a format we do not serve — in which case the caller falls back to the flat fields rather than guessing. - compressed-tensors describes a format by its parameters; we name it. `targets` is theirs too: - a bare entry is a CLASS name (GLM-5.2 ships `["Linear"]`, i.e. every linear), and an `re:` - prefix introduces a pattern. + compressed-tensors describes a format by its parameters; we name it. `group_size` is what + separates the two 4-bit formats: nvfp4's block scale is E4M3 over 16 values, mxfp4's is E8M0 + over 32. `targets` is theirs too: a bare entry is a CLASS name (GLM-5.2 ships `["Linear"]`, + i.e. every linear), and an `re:` prefix introduces a pattern. """ + formats = { + (4, "float", 16): "nvfp4", + (4, "float", 32): "mxfp4", + (8, "float", 32): "mxfp8", + (8, "float", None): "fp8", + } weights = spec.get("weights") or {} key = (weights.get("num_bits"), weights.get("type"), weights.get("group_size")) - quant_method = _CT_FORMATS.get(key) + quant_method = formats.get(key) if quant_method is None: return None targets = [] @@ -1757,7 +1780,7 @@ def _group_from_config_groups(spec: dict) -> FineGrainedGroup | None: return FineGrainedGroup( quant_method=quant_method, targets=targets or list(_CATCH_ALL), - activation_format=_CT_FORMATS.get(act_key) if activations else None, + activation_format=formats.get(act_key) if activations else None, activation_scheme="dynamic" if activations.get("dynamic", True) else "static", ) @@ -1819,42 +1842,13 @@ def __init__( # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") - # NVIDIA modelopt exports: `quant_algo` names the format (modelopt covers FP8, INT4 AWQ, - # W4A8 and others; `_MODELOPT_ALGOS` is the subset this path serves) and the skip list is - # a glob-style `ignore`, or `exclude_modules` in modelopt's own spelling. The calibrated - # `input_scale` TENSORS are the loader's business, not this config's. + # "modelopt" names the PRODUCER, not a format; read what it exported into our fields if str(self.quant_method) == "modelopt" or kwargs.get("quant_algo") is not None: - quant_algo = kwargs.pop("quant_algo", None) - if quant_algo not in _MODELOPT_ALGOS: - raise ValueError( - f"modelopt checkpoints are supported for quant_algo in " - f"{sorted(_MODELOPT_ALGOS)}; got {quant_algo!r}" - ) - # "modelopt" names the PRODUCER, not a format, so the algo it exported becomes the - # format here and nothing downstream has to ask where a config came from. - self.quant_method, algo_activation_format = _MODELOPT_ALGOS[quant_algo] - logger.info( - f"modelopt checkpoint exported with quant_algo={quant_algo!r}; loading it as " - f"{str(self.quant_method)!r}." - ) + self.quant_method, algo_activation_format, algo_groups, ignore = self._read_modelopt(kwargs) self.activation_format = activation_format or algo_activation_format - # modelopt also describes the format parametrically under `config_groups`; take it - # when every entry maps, so a multi-group export is not silently flattened to one - ct_groups = kwargs.pop("config_groups", None) - if groups is None and ct_groups: - mapped = {name: _group_from_config_groups(spec) for name, spec in ct_groups.items()} - if all(group is not None for group in mapped.values()): - groups = mapped - ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) + groups = groups or algo_groups if modules_to_not_convert is None and ignore is not None: - # modelopt names skipped subtrees with GLOBS — "model.layers.0*", "model.layers.1.*", - # or bare "self_attn". Handed on as regexes they are wrong twice: the dots match any - # character, and the star is greedy, so "model.layers.1.*" also takes layers 10-19. - # Anchor each at a path boundary instead: the named module and everything under it. - modules_to_not_convert = [] - for glob in ignore: - subtree = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") - modules_to_not_convert.append(r"(?:^|.*\.)" + re.escape(subtree) + r"(\..*)?$") + modules_to_not_convert = subtree_patterns(ignore) self.modules_to_not_convert = modules_to_not_convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme @@ -1882,6 +1876,34 @@ def __init__( } self.post_init() + @staticmethod + def _read_modelopt(kwargs: dict) -> tuple[str, str | None, dict | None, list | None]: + """`(format, activation format, groups, ignore list)` out of an NVIDIA modelopt export. + + `quant_algo` names what it exported; modelopt also covers FP8, INT4 AWQ, W4A8 and MXFP4, + which are refused rather than silently mis-read. It may describe the same thing + parametrically under `config_groups` — taken only when every entry maps, so a multi-group + export is not flattened to one. The calibrated `input_scale` TENSORS are the loader's + business, not this config's. + """ + algos = {"NVFP4": (QuantizationMethod.NVFP4, "nvfp4")} + quant_algo = kwargs.pop("quant_algo", None) + if quant_algo not in algos: + raise ValueError( + f"modelopt checkpoints are supported for quant_algo in {sorted(algos)}; got {quant_algo!r}" + ) + quant_method, activation_format = algos[quant_algo] + logger.info( + f"modelopt checkpoint exported with quant_algo={quant_algo!r}; loading it as {str(quant_method)!r}." + ) + groups = None + if ct_groups := kwargs.pop("config_groups", None): + mapped = {name: group_from_config_groups(spec) for name, spec in ct_groups.items()} + if all(group is not None for group in mapped.values()): + groups = mapped + ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) + return quant_method, activation_format, groups, ignore + def to_dict(self): """`config.json`-ready: the groups go out as plain dicts so the config round-trips.""" out = super().to_dict() @@ -1907,21 +1929,6 @@ def group_for(self, module_name: str) -> FineGrainedGroup: return group raise KeyError(f"no group covers {module_name!r}; groups: {list(self.groups)}") - def split_out_experts(self, expert_dtype: str | None) -> None: - """Fold a legacy `config.expert_dtype` into the group it means. - - DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL config, - because a flat quant config carries one format and cannot say "the experts differ". Split - that into the two groups it describes, so nothing downstream reads the side-channel. - """ - if expert_dtype != "fp4" or len(self.groups) > 1: - return - flat = self.groups["default"] - self.groups = { - "experts": replace(flat, quant_method="mxfp4", activation_format="mxfp4", targets=[r"\.experts($|\.)"]), - "dense": flat, - } - def post_init(self): r""" Safety checker that arguments are correct diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 264d7cbad1ee..aa2b3ea1403a 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -387,12 +387,17 @@ def test_post_expert_norm_runs_on_the_reducing_arms_and_is_refused_by_megamoe(se applied = [] + class _Norm(torch.nn.Module): + def forward(self, rows): + applied.append(tuple(rows.shape)) + return rows + class _Experts(torch.nn.Module): has_post_expert_norm = True - def _apply_post_norm(self, rows): - applied.append(tuple(rows.shape)) - return rows + def __init__(self): + super().__init__() + self.post_expert_norm = _Norm() module = _Experts() rows = torch.randn(4, 8) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 1242ed506572..60e938e78b9d 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -29,7 +29,7 @@ import torch import transformers.integrations.deepgemm as deepgemm -import transformers.integrations.finegrained as fg +import transformers.integrations.finegrained.core as fg from transformers.integrations.finegrained import ( FineGrainedExperts, FineGrainedLinear, @@ -391,14 +391,14 @@ def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): """A model whose experts norm the down output before the routing weights (Muse-Spark): the swap carries the norm over, the kernel chain runs it on the routed rows, and the eager loop applies it per expert application — the same place the reference forwards do. - A form the kernels do not implement rides as the module's own ``_apply_post_norm``.""" + A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" kernel, rec = _fake_bundle() m = self._experts(has_gate=True) m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True p1, p2, p3, p4 = _loaded(kernel) with p1, p2, p3, p4: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m._apply_post_norm) + self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) class _Recording(torch.nn.Module): def __init__(self): @@ -414,18 +414,17 @@ def forward(self, x): m(*self._route()) self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") - def test_declaring_a_post_expert_norm_requires_the_hook_that_applies_it(self): - """The shared experts interface: a class that declares a post-expert norm has to define - `_apply_post_norm`, where its own math lives — there is no default passthrough to fall - back on, the way `act_fn` has no default activation.""" + def test_a_post_expert_norm_needs_nothing_but_the_norm(self): + """A class that declares a post-expert norm only has to HOLD it: every backend calls + `post_expert_norm` directly, so there is no per-class hook to define alongside it.""" from transformers.integrations.moe import use_experts_implementation - with self.assertRaises(TypeError): + @use_experts_implementation(has_post_expert_norm=True) + class _JustTheNorm(torch.nn.Module): + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states - @use_experts_implementation(has_post_expert_norm=True) - class _NoHook(torch.nn.Module): - def forward(self, hidden_states, top_k_index, top_k_weights): - return hidden_states + self.assertFalse(hasattr(_JustTheNorm, "_apply_post_norm")) def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels @@ -445,7 +444,7 @@ def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) else: - self.assertEqual(call.kwargs["post_expert_norm"], m._apply_post_norm) + self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) self.assertIsNone(call.kwargs["post_expert_norm_weight"]) def test_biases_ride_the_kernel_chain(self): @@ -579,7 +578,6 @@ def _operands(self, *, input_reduce: bool): module.has_post_expert_norm = True module.post_expert_norm = norm module.post_expert_norm_name = "input_scaled_rms_norm" - module._apply_post_norm = lambda x: x kernel = mock.Mock() kernel.get_supported_act_fns.return_value = ("silu",) kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) @@ -594,7 +592,7 @@ def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): operands, module = self._operands(input_reduce=True) self.assertIs( operands["post_expert_norm"], - module._apply_post_norm, + module.post_expert_norm, "fused past a collective: the kernel would normalize this rank's partial sum", ) self.assertIsNone(operands["post_expert_norm_weight"]) @@ -635,7 +633,7 @@ def _no_accelerator(self): with ExitStack() as stack: stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) stack.enter_context( - mock.patch("transformers.quantizers.quantizer_finegrained.is_torch_xpu_available", return_value=False) + mock.patch("transformers.quantizers.finegrained.base.is_torch_xpu_available", return_value=False) ) yield @@ -686,7 +684,7 @@ class FineGrainedParallelPlanTest(unittest.TestCase): def _planned(self, impl=None, raw=None): from types import SimpleNamespace - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig quantizer = FineGrainedHfQuantizer(FineGrainedConfig()) @@ -791,7 +789,7 @@ def test_a_multimodal_models_expert_plans_are_reached(self): shard from the sub-config's own plan — so the scale stays whole against a sharded weight, which is a wrong answer rather than a crash.""" from transformers import Glm4vMoeConfig - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig config = Glm4vMoeConfig() @@ -857,7 +855,7 @@ def test_a_models_own_plan_survives_the_update(self): substring — the whole expert plan on Qwen3-MoE. """ from transformers import Qwen3Config, Qwen3MoeConfig, Qwen3VLMoeConfig - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig for config_cls in (Qwen3Config, Qwen3MoeConfig): @@ -881,7 +879,7 @@ def test_a_real_moe_models_experts_gain_their_companions(self): """The other half of the same invariant: preserving the model's entries is what lets the companion pass find `moe_tp_experts` and give the expert scales an entry of their own.""" from transformers import Qwen3MoeConfig - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig config = Qwen3MoeConfig() @@ -929,7 +927,7 @@ def _reference_dequant(self, blocks, scales): return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): - from transformers.integrations.finegrained_conversions import ( + from transformers.integrations.finegrained.conversions import ( FineGrainedPackedBlocks, FineGrainedScaleContainer, ) @@ -1103,7 +1101,7 @@ def _op_kernel(self): return kernel def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): - from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales model, experts = self._experts("mxfp8") kernel = self._op_kernel() @@ -1146,7 +1144,7 @@ def test_scale_container_op_round_trips_every_container(self): the op brings both into the held e8m0 (exact cast / same bytes), records the container on the module, and its reverse restores that container on save; weights and native scales pass both ways.""" - from transformers.integrations.finegrained_conversions import FineGrainedScaleContainer + from transformers.integrations.finegrained.conversions import FineGrainedScaleContainer op = FineGrainedScaleContainer(hf_quantizer=None) native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) @@ -1178,7 +1176,7 @@ def test_scale_container_op_round_trips_every_container(self): self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): - from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales kernel = self._op_kernel() reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op @@ -1220,7 +1218,7 @@ def _arch_converters(self): def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained_conversions import ( + from transformers.integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales, @@ -1250,7 +1248,7 @@ def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): def test_blocks_scales_converters_carry_the_layout_ops(self): from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained_conversions import ( + from transformers.integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, FineGrainedSwizzleScales, ) @@ -1334,7 +1332,7 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself and keeps the stack; every other case gets the interleaved order, and the reverse restores the stack.""" - from transformers.integrations.finegrained_conversions import FineGrainedInterleaveGateUp + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] op = FineGrainedInterleaveGateUp(hf_quantizer=None) @@ -1360,7 +1358,7 @@ def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation the unquantized experts hold — so it is compared against that integration directly.""" from transformers.core_model_loading import Transpose - from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedPackedBlocks + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks from transformers.integrations.mxfp4 import convert_moe_packed_tensors torch.manual_seed(0) @@ -1384,7 +1382,7 @@ def test_a_static_scheme_gets_its_activation_scales_written(self): memory: a zero divides the activations by zero and a negative flips their sign, which showed up as NaN logits from one fixture and not another, run to run. """ - from transformers.integrations.finegrained_conversions import FineGrainedQuantize + from transformers.integrations.finegrained.conversions import FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1416,7 +1414,7 @@ def test_an_expert_bias_is_passed_through_not_quantized(self): """A GPT-OSS expert bias is `(E, rows)` — 2-D, like a dense weight — so the rank guard alone lets it through. It has no scale slot, and emitting one gives the loader a `_bias_scale_inv` no module holds.""" - from transformers.integrations.finegrained_conversions import FineGrainedQuantize + from transformers.integrations.finegrained.conversions import FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1432,7 +1430,7 @@ def test_a_partial_trailing_block_is_padded_not_refused(self): """DeepSeek-V3 ships `kv_a_proj_with_mqa` as `(576, 7168)` against a 128x128 block with a `(5, 56)` scale grid, so the format rounds the grid UP and quantizes the short block on its own values. Refusing the shape instead left the weight full precision.""" - from transformers.integrations.finegrained_conversions import FineGrainedQuantize + from transformers.integrations.finegrained.conversions import FineGrainedQuantize torch.manual_seed(0) for rows, cols in ((576, 256), (192, 256), (256, 256)): @@ -1445,7 +1443,7 @@ def test_a_partial_trailing_block_is_padded_not_refused(self): self.assertEqual(scale.shape, (-(-rows // 128), -(-cols // 128))) def test_block_fp8_round_trips_within_its_floor(self): - from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedQuantize + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1478,7 +1476,7 @@ class FineGrainedModeloptConverterTest(unittest.TestCase): rank collects all E globals and the forward asserts on the per-expert count.""" def _modelopt_conversions(self, **cfg_kwargs): - from transformers.quantizers.quantizer_finegrained_nvfp4 import FineGrainedNvfp4HfQuantizer + from transformers.quantizers.finegrained.nvfp4 import FineGrainedNvfp4HfQuantizer from transformers.utils.quantization_config import FineGrainedConfig cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) @@ -1552,7 +1550,7 @@ def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): globals and flattened the gate|up pair instead of folding it.""" import re - from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer from transformers.utils.quantization_config import FineGrainedConfig for quant_kwargs in ( @@ -1637,7 +1635,7 @@ def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): every expert against the largest one's range. It stays per expert.""" import torch - from transformers.integrations.finegrained_conversions import FineGrainedInputScales + from transformers.integrations.finegrained.conversions import FineGrainedInputScales sources = { "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), @@ -1653,7 +1651,7 @@ def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): where the kernels assert one per expert.""" import torch - from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1683,7 +1681,7 @@ def test_one_converter_merges_a_layer_s_globals(self): global and the down's weight global scales the expert output back. The down's calibrated input scale moves the other way, keeping the requantized intermediate on the range the checkpoint calibrated.""" - from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) @@ -1709,7 +1707,7 @@ def test_one_converter_merges_a_layer_s_globals(self): def test_globals_with_one_calibrated_projection_pass_through(self): """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a single global) has nothing to merge, so every global reaches its module unchanged.""" - from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1731,7 +1729,7 @@ def test_globals_with_one_calibrated_projection_pass_through(self): def test_the_fused_layout_ships_both_halves_in_one_tensor(self): """Muse-Spark's `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two keys, and merges the same way.""" - from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals targets = [ "experts.gate_up_proj_weight_global_scale", @@ -1751,7 +1749,7 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated `input_scale` reduces to one value; the down's rows are per expert, so its stays per expert — the requant epilogue normalizes each row by its own expert's value.""" - from transformers.integrations.finegrained_conversions import FineGrainedInputScales + from transformers.integrations.finegrained.conversions import FineGrainedInputScales up = FineGrainedInputScales().convert( { @@ -1870,7 +1868,7 @@ def test_on_the_fly_group_formats_round_trip_within_their_floor(self): """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in the module's layout (swizzled where held so); dequantizing them back lands within each format's floor.""" - from transformers.integrations.finegrained_conversions import FineGrainedDequantize, FineGrainedQuantize + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize torch.manual_seed(0) cfg = _Cfg() @@ -1911,7 +1909,7 @@ def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes globals. Returns the module, the exact dequantized weights (so the reference is the quantization floor rather than the pre-quant weight), and the dequantization a reader that applied ONE of the two globals to both halves would get.""" - from transformers.integrations.finegrained_conversions import FineGrainedWeightGlobals + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals kernel = load_finegrained_kernel() experts = FineGrainedExperts(cfg, weight_format="nvfp4", activation_format=activation_format).cuda() @@ -2050,7 +2048,7 @@ def test_experts_scale_layout_op_keeps_the_forward_correct(self): """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op fills them from the affine grid, the fused forward matches the affine module's, and the reverse op hands the affine grid back bitwise.""" - from transformers.integrations.finegrained_conversions import FineGrainedSwizzleScales + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales cfg = _Cfg() cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = ( diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 954258411fd3..c83741361fd5 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -365,36 +365,3 @@ def test_sharding_does_not_create_unrelated_module_attributes(self): self.assertEqual(module.random_attr, 123) self.assertFalse(hasattr(module, "num_experts")) - - -@is_tensor_parallel_test -class TestZeroDimParameterSharding(TestCasePlus): - """A 0-dim parameter has no axis to shard, so every owning rank takes the whole value. - - Reachable in practice: quantized checkpoints carry per-tensor scalar scales, and a ModelOpt - NVFP4 model ships one `weight_scale_2` / `input_scale` per expert projection — thousands of - 0-dim tensors. Without the guard, `shard_tensor` builds an empty slice tuple and `source[()]` - is not a form a lazy safetensors slice accepts. - """ - - @staticmethod - def _op(placements, axis0_offset=0, axis0_local=2): - op = DtensorShardOperation.__new__(DtensorShardOperation) - op.placements = placements - op.param_ndim = 3 - op._axis0_offset = axis0_offset - op._axis0_local_size = axis0_local - return op - - def test_dense_scalar_replicates(self): - out = self._op(placements=()).shard_tensor(torch.tensor(3.5)) - self.assertEqual(out.shape, torch.Size([])) - self.assertEqual(out.item(), 3.5) - - def test_moe_scalar_follows_expert_ownership(self): - from torch.distributed.tensor.placement_types import Shard - - op = self._op(placements=(Shard(0),), axis0_offset=2, axis0_local=2) # owns experts [2, 4) - owned = op.shard_tensor(torch.tensor(3.5), tensor_idx=3) - self.assertEqual(owned.item(), 3.5) - self.assertIsNone(op.shard_tensor(torch.tensor(3.5), tensor_idx=1)) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 309879ab5091..174d1dce8cb6 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -46,7 +46,6 @@ spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig -from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts from .test_distributed_sharding_utils import FakeMesh, _make_dtensor_shard_op @@ -515,10 +514,10 @@ def test_moe_and_qkv_conversion_reversed(self): self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict)) def test_qkv_chunk_rope_permute_with_fp8_quantization(self): - if is_triton_available(): - from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize - else: - self.skipTest("Fine-grained FP8 integration tests require Triton to be installed.") + # block-FP8 quantize/dequantize is pure torch here; no kernel is launched + from transformers.integrations.finegrained import FineGrainedLinear + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + n_heads = 2 head_dim = 4 in_dim = 4 @@ -526,17 +525,15 @@ def test_qkv_chunk_rope_permute_with_fp8_quantization(self): block_size = (4, 4) class RopeProjector(nn.Module): - def __init__(self, *, with_scale: bool = False): + def __init__(self): super().__init__() self.weight = nn.Parameter(torch.zeros(out_dim, in_dim)) - if with_scale: - scale_shape = (out_dim // block_size[0], in_dim // block_size[1]) - self.weight_scale_inv = nn.Parameter(torch.ones(scale_shape)) class RopeSelfAttn(nn.Module): def __init__(self): super().__init__() - self.q_proj = RopeProjector(with_scale=True) + # only q_proj is quantized, so only it is a finegrained module holding a scale + self.q_proj = FineGrainedLinear(in_dim, out_dim, block_size=block_size, weight_format="fp8") self.k_proj = RopeProjector() self.v_proj = RopeProjector() @@ -553,6 +550,9 @@ def __init__(self, config): self.layers = nn.ModuleList([RopeLayer()]) self.post_init() + def _init_weights(self, module): + pass # every weight here comes from the state dict, and `normal_` has no fp8 kernel + config = PreTrainedConfig() config.num_attention_heads = n_heads model = RopeModel(config) @@ -573,17 +573,17 @@ def __init__(self, config): raw_k = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) raw_v = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + 100.0 raw_qkv = torch.cat([raw_q, raw_k, raw_v], dim=0) - state_dict = {"model.layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} + state_dict = {"layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} quantizer_cls = type( - "FineGrainedFP8HfQuantizer", + "FineGrainedHfQuantizer", (), { "__init__": lambda self, bs=block_size: setattr( - self, "quantization_config", SimpleNamespace(weight_block_size=bs) + self, "quantization_config", SimpleNamespace(weight_block_size=bs, scale_fmt="float") ), "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), - "get_quantize_ops": lambda self: Fp8Quantize(self), + "get_quantize_ops": lambda self: FineGrainedQuantize(self), "pre_quantized": False, }, ) @@ -597,7 +597,7 @@ def __init__(self, config): "self_attn.k_proj.weight", "self_attn.v_proj.weight", ], - operations=[Chunk(dim=0), PermuteForRope()], + operations=[Chunk(dim=0), PermuteForRope(permute_layer_names=["q_proj", "k_proj"])], ) ] load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) @@ -608,19 +608,19 @@ def __init__(self, config): self.assertEqual(loading_info.mismatched_keys, set()) self.assertEqual(loading_info.conversion_errors, {}) - permute_op = PermuteForRope() + permute_op = PermuteForRope(permute_layer_names=["q_proj", "k_proj"]) permute_op.config = model.config expected_q = permute_op._apply(raw_q) expected_k = permute_op._apply(raw_k) - expected_v = permute_op._apply(raw_v) + expected_v = raw_v # V carries no RoPE, so the permute names only q and k model_state = model.state_dict() self.assertFalse(torch.allclose(raw_k, expected_k)) - torch.testing.assert_close(model_state["model.layers.0.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(model_state["model.layers.0.self_attn.v_proj.weight"], expected_v) + torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) - q_weight_key = "model.layers.0.self_attn.q_proj.weight" - scale_key = "model.layers.0.self_attn.q_proj.weight_scale_inv" + q_weight_key = "layers.0.self_attn.q_proj.weight" + scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" self.assertIn(scale_key, model_state) expected_dtype = torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else torch.int8 self.assertEqual(model_state[q_weight_key].dtype, expected_dtype) @@ -631,10 +631,8 @@ def __init__(self, config): torch.Size((out_dim // block_size[0], in_dim // block_size[1])), ) - dequant = Fp8Dequantize(block_size=block_size) - dequantized_q = dequant.convert( - [model_state[q_weight_key], model_state[scale_key]], - context={"quantization_config": quantizer.quantization_config}, + dequantized_q = FineGrainedDequantize(None)._dequantize_one( + model_state[q_weight_key], model_state[scale_key], output_dtype=torch.float32 ) torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) @@ -642,7 +640,7 @@ def test_fp8_ue8m0_quantize_dequantize_round_trip(self): # ``scale_fmt="ue8m0"`` rounds weight_scale_inv to a power of two; the weight has to be # quantized with that same rounded scale, or ``weight * weight_scale_inv`` at dequant # disagrees with the scale the weight was divided by (a silent per-block error up to an octave). - from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize if not hasattr(torch, "float8_e8m0fnu"): self.skipTest("ue8m0 storage requires torch.float8_e8m0fnu") @@ -653,29 +651,39 @@ def test_fp8_ue8m0_quantize_dequantize_round_trip(self): torch.manual_seed(0) weight = torch.randn(128, 128, dtype=torch.float32) - quantized = Fp8Quantize(quantizer)._quantize_one("layer.weight", weight) - recovered = Fp8Dequantize(quantizer)._dequantize_one( - quantized["layer.weight"], quantized["layer.weight_scale_inv"], output_dtype=torch.float32 + # no model to ask, so the op takes the block and the scale format from the config + quantized = FineGrainedQuantize(quantizer)._quantize_one("layer.weight", weight, None) + recovered = FineGrainedDequantize(quantizer)._dequantize_one( + quantized["layer.weight"], quantized["layer.weight_scale_inv"].float(), output_dtype=torch.float32 ) rel_err = ((recovered - weight).abs().sum() / weight.abs().sum()).item() self.assertLess(rel_err, 5e-2) # fp8 round-trip is ~2e-2; a scale mismatch inflates it past 0.2 def test_fp8_float_scale_fmt_quantization_unchanged(self): - # The ue8m0 reordering must leave the default scale_fmt="float" path bit-identical to the - # original single-division formula (scales = _FP8_MAX / max_abs). Any divergence only shows - # on rare fp32-ULP block maxima, so scan many random 128x128 blocks rather than one. - from transformers.integrations.finegrained_fp8 import _FP8_DTYPE, _FP8_MAX, _FP8_MIN, Fp8Quantize + """The default `scale_fmt="float"` path stays bit-identical to its formula: the inverse + scale is `amax / MAX` and the weight is DIVIDED by it. + + Reordering that into the algebraically equal `weight * (MAX / amax)` is not bit-equal — + measured, it moves 29 of 100 random blocks' scales and 1 of their weights — so a refactor + that "simplifies" the arithmetic changes every checkpoint this path writes. Scan many + blocks, since the divergence only shows on particular fp32 maxima. + """ + from transformers.integrations.finegrained.conversions import ( + _FP8_DTYPE, + _FP8_MAX, + _FP8_MIN, + FineGrainedQuantize, + ) - quantizer = Fp8Quantize( + quantizer = FineGrainedQuantize( SimpleNamespace(quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float")) ) for seed in range(100): torch.manual_seed(seed) weight = torch.randn(128, 128, dtype=torch.float32) - out = quantizer._quantize_one("layer.weight", weight) - scales = _FP8_MAX / weight.abs().amax() - ref_q = torch.clamp(weight * scales, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) - ref_inv = (1.0 / scales).to(torch.float32).reshape(1, 1) + out = quantizer._quantize_one("layer.weight", weight, None) + ref_inv = (weight.abs().amax() / _FP8_MAX).to(torch.float32).reshape(1, 1) + ref_q = torch.clamp(weight / ref_inv, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) self.assertTrue(torch.equal(out["layer.weight"], ref_q), f"float weight diverged at seed {seed}") self.assertTrue( torch.equal(out["layer.weight_scale_inv"], ref_inv), f"float scale diverged at seed {seed}" diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index 9be7b2ebd753..b6a6d4d26093 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -294,6 +294,25 @@ def test_moe_2D_shard_with_negative_dim_indices(self): else: torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + def test_0dim_dense_is_replicated(self): + """A 0-dim parameter has no axis to shard, so every rank takes the whole value. + + Reachable in practice: a per-tensor scalar scale, and a ModelOpt NVFP4 model ships one + `weight_scale_2` / `input_scale` per expert projection — thousands of 0-dim tensors. + """ + mesh = FakeMesh(2, rank=0) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 4), local_shape=(2, 4)) + out = op.shard_tensor(torch.tensor(3.5)) + self.assertEqual(out.shape, torch.Size([])) + self.assertEqual(out.item(), 3.5) + + def test_0dim_per_expert_follows_expert_ownership(self): + """Per-expert 0-dim scales are dropped on the ranks that do not own the expert.""" + mesh = FakeMesh(2, rank=1) # owns experts [2, 4) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 8), local_shape=(2, 8)) + self.assertEqual(op.shard_tensor(torch.tensor(3.5), tensor_idx=3).item(), 3.5) + self.assertIsNone(op.shard_tensor(torch.tensor(3.5), tensor_idx=1)) + def test_compute_strided_slice(self): # Direct tests for _compute_strided_slice(intervals, rank, world_size, split_factor). # Keys: (input_interval, rank, world_size, split_factor) -> expected output list. From 95b3e44f55ecce8a76b04f5eef7a5c54e2bbc62d Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 10:56:33 -0700 Subject: [PATCH 41/68] trim --- src/transformers/integrations/finegrained/conversions.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 487fb6789a8b..55ea2cf9a43a 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -280,10 +280,8 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): value = torch.stack(value, dim=0) if isinstance(value, list) else value values.append(value.float()) stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) - # A scale is a MAGNITUDE, so reduce over absolute values and keep it off zero: a negative - # flips every dequantized row's sign and a zero divides the activations by zero. A no-op - # for a real calibrated checkpoint; it catches scales that were built rather than measured. - per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1).clamp(min=1e-12) + # a scale is a magnitude, so the reduce is over absolute values + per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1) one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} From 2392e6a61e39f934f62ee4867f3e9c97853b9a1b Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 11:34:39 -0700 Subject: [PATCH 42/68] fixes --- src/transformers/quantizers/finegrained/blockfp8.py | 4 +++- src/transformers/quantizers/finegrained/nvfp4.py | 4 ++-- tests/kernels/test_finegrained.py | 12 ++++++------ 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/transformers/quantizers/finegrained/blockfp8.py b/src/transformers/quantizers/finegrained/blockfp8.py index 15cc3c285614..dbfeee6e4379 100644 --- a/src/transformers/quantizers/finegrained/blockfp8.py +++ b/src/transformers/quantizers/finegrained/blockfp8.py @@ -62,7 +62,9 @@ def _static_activation_conversions(self): target_patterns=f"experts.{proj}_activation_scale", operations=[FineGrainedInputScales(self)], ) - for proj in ("gate_up_proj", "down_proj") + # `up_proj` because an UNGATED experts module holds its scale under that name; the + # converter whose source no checkpoint has simply never fires + for proj in ("gate_up_proj", "up_proj", "down_proj") ] # a dense linear holds the one value itself; the lookahead keeps this off the expert keys dense = [ diff --git a/src/transformers/quantizers/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py index 10fc37738726..d2d85ec44e02 100644 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -63,7 +63,7 @@ def update_weight_conversions(self, weight_conversions): def _nvfp4_conversions(self): """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already - stacked per layer (the fused vLLM layout, Muse-Spark). A checkpoint matches one set and + stacked per layer (the fused vLLM layout). A checkpoint matches one set and the other never fires. Routed experts only — a generic `weight_scale*` rename would run before converter collection and mangle these keys.""" # weight-only keeps activations bf16, so there is no activation global for a calibrated @@ -124,7 +124,7 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): return converters def _nvfp4_fused_conversions(self, calibrated: bool): - """Already stacked per layer, as the vLLM fused layout and Muse-Spark ship them — so no + """Already stacked per layer, as the vLLM fused layout ships them — so no `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the `.weight`-anchored rule gives the per-expert layout.""" from ...core_model_loading import WeightConverter diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 60e938e78b9d..926d95005f52 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -388,7 +388,7 @@ def test_weight_only_experts_hold_no_activation_global(self): self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): - """A model whose experts norm the down output before the routing weights (Muse-Spark): + """A model whose experts norm the down output before the routing weights: the swap carries the norm over, the kernel chain runs it on the routed rows, and the eager loop applies it per expert application — the same place the reference forwards do. A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" @@ -1528,7 +1528,7 @@ def test_global_scale_converters_carry_an_expert_index(self): def test_both_modelopt_scale_layouts_have_converters(self): """modelopt ships the scales either as one tensor per expert per projection (GLM-5.2) or - as one already-stacked tensor per layer (the fused vLLM layout, Muse-Spark). Both are + as one already-stacked tensor per layer (the fused vLLM layout). Both are converted; the block scale takes the module's layout ops either way.""" sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} for suffix in ("weight_scale", "weight_scale_2", "input_scale"): @@ -1727,7 +1727,7 @@ def test_globals_with_one_calibrated_projection_pass_through(self): torch.testing.assert_close(out[targets[2]], torch.tensor([0.5, 0.25])) def test_the_fused_layout_ships_both_halves_in_one_tensor(self): - """Muse-Spark's `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two + """A stacked `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two keys, and merges the same way.""" from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals @@ -1767,7 +1767,7 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d class _InputScaledRMSNorm(torch.nn.Module): - """Muse-Spark's post-expert norm: scale by ``1 + weight`` BEFORE normalizing, so the row's + """A post-expert norm that scales by ``1 + weight`` BEFORE normalizing, so the row's mean square is taken on the scaled values (not the same function as normalizing first).""" def __init__(self, dim, eps=1e-6): @@ -1905,7 +1905,7 @@ def test_on_the_fly_group_formats_round_trip_within_their_floor(self): def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes=(8.0, 1.0)): """NVFP4 experts in the layout a modelopt checkpoint delivers: every (expert, half) of the gate|up stack quantized against its OWN global (``weight_scale_2`` is per projection, and - the two differ per expert on Muse-Spark), the down per expert, plus calibrated activation + the two can differ per expert), the down per expert, plus calibrated activation globals. Returns the module, the exact dequantized weights (so the reference is the quantization floor rather than the pre-quant weight), and the dequantization a reader that applied ONE of the two globals to both halves would get.""" @@ -1986,7 +1986,7 @@ def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): return out def test_modelopt_nvfp4_experts_merge_their_globals_and_fuse_the_post_norm(self): - """The Muse-Spark checkpoint shape end to end through the integration: a gate|up stack + """The stacked checkpoint shape end to end through the integration: a gate|up stack calibrated per half (merged at load), calibrated activation globals, and a per-expert output norm the binder fuses, on all three forwards. Weight-only pins the numbers against the dequantized weights (nothing else rounds); the W4A4 chain rides the 4-bit activation floor, so it is pinned by being far From f7a8e543c52f8a019f4c73e80ed42336b2f3b5bd Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 11:50:04 -0700 Subject: [PATCH 43/68] last --- src/transformers/integrations/finegrained/conversions.py | 3 +-- src/transformers/integrations/finegrained/core.py | 9 ++------- src/transformers/quantizers/finegrained/base.py | 5 +++++ src/transformers/utils/quantization_config.py | 7 +++++++ 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 55ea2cf9a43a..71ae8ac5bd03 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -396,8 +396,7 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc } if value.device.type not in ("cuda", "xpu"): - # the kernels' quantizers are triton launches; the package builds for cuda, rocm - # (which torch reports as "cuda") and xpu + # rocm counts: torch reports it as "cuda" raise ValueError( f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on an " f"accelerator, but the weight is on {value.device.type}" diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 4386cca0e86f..22d1b374e852 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -352,13 +352,8 @@ def finegrained_linear( bias=bias, ) except (ImportError, NotImplementedError, ValueError) as e: - # DeepGEMM is loadable but declined this specific input; fall back to Triton, which is more - # permissive (handles FP8/FP4 with float32 or UE8M0 scales on any arch, plus input dtypes - # DeepGEMM rejects). If Triton can't serve it either, it raises its own error. - # - NotImplementedError: an arch/input combo DeepGEMM has no kernel for (FP4 on Hopper — - # `is_deepgemm_loadable` is dtype-agnostic and passes there — or float32 scales on Blackwell); - # - ValueError: an input DeepGEMM rejects but Triton supports (e.g. activations it won't quantize); - # - ImportError: a symbol/build gap. + # Triton is the more permissive of the two, so a decline here is not fatal: it serves + # arch/input combos DeepGEMM has no kernel for, and raises its own error if it cannot. logger.warning_once( f"DeepGEMM declined this call, falling back to Triton. Reason: {e} " "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely." diff --git a/src/transformers/quantizers/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py index 8b2664540fc0..08a4ad1019f0 100644 --- a/src/transformers/quantizers/finegrained/base.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -197,6 +197,11 @@ def _process_model_before_weight_loading( ) if self.pre_quantized and self.quantization_config.modules_to_convert: + if self._quant_method() != "fp8": + logger.warning_once( + f"Embedding tables are quantized to FP8; {self._quant_method()} has no embedding path. " + f"{self.quantization_config.modules_to_convert} will hold FP8 rows with one per-tensor scale." + ) replace_with_finegrained_embedding( model, self.quantization_config.modules_to_convert, self.modules_to_not_convert ) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index cdf57095b900..e02d51b1a8a4 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1816,6 +1816,8 @@ class FineGrainedConfig(QuantizationConfigMixin): A list of module names that should not be converted during quantization. modules_to_convert (`list`, *optional*): A list of additional module names, such as embedding tables, that should be converted during quantization. + An embedding table is quantized to FP8 whatever the rest of the config says: it holds one + per-tensor scale, and the group formats have no embedding path. scale_fmt (`str`, *optional*, defaults to `"float"`): Storage dtype of the per-block weight scales: `"float"` (fp32, V3-style) or `"ue8m0"` (1-byte `torch.float8_e8m0fnu`, V4-style). @@ -1823,6 +1825,11 @@ class FineGrainedConfig(QuantizationConfigMixin): Activation quantization format, needed only where the weights leave it ambiguous (MXFP4 weights run as W4A16 with `"bf16"`, W4A8 with `"mxfp8"`, W4A4 with `"mxfp4"`). `None` (default) matches the kernels' weight-native choice. + groups (`dict[str, transformers.utils.quantization_config.FineGrainedGroup] | None`, *optional*): + Which format covers which modules, for a checkpoint that is more than one (DeepSeek-V4 + is MXFP4 experts over block-FP8 linears). Each group names its modules with regexes and + carries its own copy of the fields above. Leave unset for a single-format checkpoint: + the flat fields are the one-group spelling. """ def __init__( From 4cbbe7ff193a52c081a84159b10ed20f5b8e351f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 12:19:28 -0700 Subject: [PATCH 44/68] Cover the groups feature, and export what it asks callers to build `FineGrainedGroupsTest` was written against the unified branch and never came across, so `groups` -- the thing this PR is for -- had no unit test here at all: nothing touched `FineGrainedGroup`, `group_for` or `group_from_config_groups`, and only the seven-model load-path run would have caught a regression. Six tests now cover the flat-config normalisation, the legacy `expert_dtype` split, the `config.json` round trip (where `to_json_string` sorts keys and a catch-all that won over a targeted group would swallow everything), an ambiguous claim, a producer's `config_groups`, and a format we do not serve. `FineGrainedGroup` was not exported, so a caller could not build what `groups=` asks for without reaching into `transformers.utils.quantization_config` -- and the autodoc entry beside `FineGrainedConfig` would have failed the docs build. The two loader compile-safety tests come over from the frozen suite: the opaque loader node must return None under `torch.compile`, cold and warm, and nothing here was checking it. Its sentinel and in-place-bias tests do NOT come over -- that glue moved into the kernels package when the experts forwards became a single call, so there is no longer anything on this side to test. Also two `PLW1514` fixes CI catches and a per-file `ruff check` does not. Co-Authored-By: Claude Opus 5 --- docs/source/en/main_classes/quantization.md | 4 + src/transformers/__init__.py | 2 + tests/kernels/test_finegrained.py | 130 +++++++++++++++++++- 3 files changed, 134 insertions(+), 2 deletions(-) diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index 2a2a8c27b7aa..ea00c4c8eb29 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -102,6 +102,10 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] FineGrainedConfig +## FineGrainedGroup + +[[autodoc]] FineGrainedGroup + ## FineGrainedFP8Config [[autodoc]] FineGrainedFP8Config diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index c2d05c546d9f..8189220c43a9 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -261,6 +261,7 @@ "FbgemmFp8Config", "FineGrainedConfig", "FineGrainedFP8Config", + "FineGrainedGroup", "FourOverSixConfig", "FPQuantConfig", "GemmaQuantizationConfig", @@ -793,6 +794,7 @@ from .utils.quantization_config import FbgemmFp8Config as FbgemmFp8Config from .utils.quantization_config import FineGrainedConfig as FineGrainedConfig from .utils.quantization_config import FineGrainedFP8Config as FineGrainedFP8Config + from .utils.quantization_config import FineGrainedGroup as FineGrainedGroup from .utils.quantization_config import FourOverSixConfig as FourOverSixConfig from .utils.quantization_config import FPQuantConfig as FPQuantConfig from .utils.quantization_config import GemmaQuantizationConfig as GemmaQuantizationConfig diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 926d95005f52..b0b631df3f8c 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -134,6 +134,40 @@ def test_loader_requires_every_symbol(self): load_finegrained_kernel() self.assertIn("matmul_grouped", str(ctx.exception)) + def test_loader_is_compile_safe_cold(self): + """Cold path: the compiled call is first to load, so the opaque loader node runs its full + body under compile and must return None, never the bundle (`torch.* op returned + non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" + kernel, _ = _fake_bundle() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_is_compile_safe_when_warm(self): + """Warm path, which is the production order: eager warm-up, then compile. The loader hits + its short-circuit at trace time — the branch that must also return None.""" + kernel, _ = _fake_bundle() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + load_finegrained_kernel() + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + def test_loader_binds_all_symbols(self): kernel, _ = _fake_bundle() p1, p2, p3, p4 = _loaded(kernel) @@ -890,6 +924,96 @@ def test_a_real_moe_models_experts_gain_their_companions(self): self.assertIn(f"layers.*.mlp.experts.{companion}", plan, companion) +class FineGrainedGroupsTest(unittest.TestCase): + """A quant config names a format per module SUBSET, because a checkpoint can be more than one.""" + + def _config(self, **kwargs): + from transformers.utils.quantization_config import FineGrainedConfig + + return FineGrainedConfig(**kwargs) + + def test_a_flat_config_is_one_group_over_everything(self): + groups = self._config(quant_method="mxfp8").groups + self.assertEqual({name: group.quant_method for name, group in groups.items()}, {"default": "mxfp8"}) + + def test_a_legacy_expert_dtype_becomes_the_two_groups_it_meant(self): + """DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL + config, because a flat config cannot say "the experts differ". `shared_experts` is DENSE + despite the name, which is why the pattern matches a path segment and not a substring.""" + from transformers.utils.quantization_config import groups_with_expert_dtype + + config = self._config(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0") + config.groups = groups_with_expert_dtype(config.groups, "fp4") + for module, expected in ( + ("model.layers.3.mlp.experts", "mxfp4"), + ("model.layers.3.mlp.experts.gate_up_proj", "mxfp4"), + ("model.layers.3.self_attn.q_proj", "fp8"), + ("model.layers.3.mlp.shared_experts.up_proj", "fp8"), + ): + self.assertEqual(config.group_for(module).quant_method, expected, module) + + def test_resolution_survives_a_round_trip_through_config_json(self): + """`to_json_string` SORTS keys, so a config read back has lost the order it was written + in. Resolution must not depend on it: the catch-all sorts first here and would swallow + every module if targeted groups were not preferred outright.""" + import json + + from transformers.utils.quantization_config import FineGrainedConfig, FineGrainedGroup + + config = FineGrainedConfig( + groups={ + "experts": FineGrainedGroup(quant_method="nvfp4", targets=[r"\.experts($|\.)"]), + "dense": FineGrainedGroup(quant_method="fp8", weight_block_size=(128, 128)), + } + ) + back = FineGrainedConfig.from_dict(json.loads(config.to_json_string())) + self.assertEqual(back.group_for("model.layers.0.mlp.experts").quant_method, "nvfp4") + self.assertEqual(back.group_for("model.layers.0.self_attn.q_proj").quant_method, "fp8") + # a list survives JSON where a tuple does not, and the block size is compared as a tuple + self.assertEqual(back.groups["dense"].weight_block_size, (128, 128)) + + def test_two_targeted_groups_claiming_one_module_is_an_error(self): + from transformers.utils.quantization_config import FineGrainedGroup + + config = self._config( + groups={ + "a": FineGrainedGroup(quant_method="fp8", targets=[r"\.experts"]), + "b": FineGrainedGroup(quant_method="nvfp4", targets=[r"mlp\."]), + } + ) + with self.assertRaises(ValueError): + config.group_for("model.layers.0.mlp.experts") + + def test_a_producers_config_groups_are_normalized(self): + """modelopt describes a format by its parameters; we name it. GLM-5.2-NVFP4 ships exactly + this, and `targets: ["Linear"]` is its spelling of the catch-all.""" + config = self._config( + quant_method="modelopt", + quant_algo="NVFP4", + config_groups={ + "group_0": { + "weights": {"num_bits": 4, "type": "float", "group_size": 16}, + "input_activations": {"num_bits": 4, "type": "float", "group_size": 16, "dynamic": False}, + "targets": ["Linear"], + } + }, + ) + group = config.group_for("model.layers.0.self_attn.q_proj") + self.assertEqual((group.quant_method, group.activation_format), ("nvfp4", "nvfp4")) + self.assertEqual(group.activation_scheme, "static") # `dynamic: False` + + def test_a_format_we_do_not_serve_falls_back_rather_than_guessing(self): + """Half a translation would quantize modules by a rule the producer did not write.""" + config = self._config( + quant_method="modelopt", + quant_algo="NVFP4", + config_groups={ + "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} + }, + ) + self.assertEqual(list(config.groups), ["default"]) + + class SubtreePatternTest(unittest.TestCase): def test_a_checkpoints_glob_skip_list_is_bounded(self): """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any @@ -2172,7 +2296,9 @@ def _checkpoint_expert_shape(model_dir): shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) index = os.path.join(model_dir, "model.safetensors.index.json") if os.path.exists(index): - shards = sorted({os.path.join(model_dir, f) for f in json.load(open(index))["weight_map"].values()}) + with open(index, encoding="utf-8") as fh: + weight_map = json.load(fh)["weight_map"] + shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) experts: set[int] = set() rows = 0 @@ -2559,7 +2685,7 @@ def _sharded(self, model_dir, model_cls, modes): import torch script = os.path.join(self._tmp.name, "sharding_worker.py") - with open(script, "w") as fh: + with open(script, "w", encoding="utf-8") as fh: fh.write(_SHARDING_WORKER) out = os.path.join(self._tmp.name, "out") os.makedirs(out, exist_ok=True) From f68786045eec00d77f840db57719454e293e5ee7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 12:35:07 -0700 Subject: [PATCH 45/68] Point everything at the fine-grained integration; deprecate the frozen files Stacked on the review PR so that one stays reviewable. The frozen modules stay where they are and keep working -- this only moves the callers over and says loudly that the files are on the way out. Deprecation moves from the CLASS to the FILE. A module-level `DeprecationWarning` fires when anything is imported out of `integrations/{finegrained_fp8,mxfp4,nvfp4}` or `quantizers/quantizer_{finegrained_fp8,mxfp4,nvfp4}`, which is the scope that matches the intent -- the whole file is frozen, not one config class. That is only safe because `_LazyModule` never executes these unless someone asks for them: `import transformers` does not reach them, nor does `from transformers import Mxfp4Config` (the config lives in `utils/quantization_config`), nor does `integrations.compressed_tensors`, which has its own `CompressedTensorsFP8Linear`. All four verified. The per-class notices go, and `FrozenFp8ShimTest` flips to the new contract. Callers, everywhere the frozen path was still in use for a model this integration serves: - `tests/models/{deepseek_v4,glm_moe_dsa,gpt_oss,qwen4_exp}` build `FineGrainedConfig`. The qwen4-exp embedding test is REPLACED rather than duplicated -- the frozen path keeps its own coverage in `tests/quantization/finegrained_fp8`, and a second copy of it in a model file was only ever testing the deprecated route. It also gains the case that the table is FP8 under any format, which used to be silent. - the ministral3 and mistral4 conversion scripts. Both already built a `FineGrainedConfig` -- `"fp8"` maps to it now -- and then swapped modules with the frozen `replace_with_fp8_linear`, so they wrote checkpoints whose config said fine-grained while the modules were `FP8Linear`. Now both halves agree. - the docs: the model pages, `experts_interface`, and a warning at the top of each frozen method page pointing at `FineGrainedConfig`. Not done here: `quantization/finegrained.md` does not exist. The three frozen methods each have a user guide and the integration replacing them has none, so the warnings point somewhere there is nothing to read. Co-Authored-By: Claude Opus 5 --- docs/source/en/experts_interface.md | 8 ++-- docs/source/en/model_doc/deepseek_v32.md | 4 +- docs/source/en/model_doc/minimax_m3_vl.md | 4 +- .../source/en/quantization/finegrained_fp8.md | 7 +++ docs/source/en/quantization/mxfp4.md | 7 +++ docs/source/en/quantization/nvfp4.md | 7 +++ .../integrations/finegrained_fp8.py | 10 +++++ src/transformers/integrations/mxfp4.py | 11 +++++ src/transformers/integrations/nvfp4.py | 10 +++++ .../convert_ministral3_weights_to_hf.py | 6 +-- .../mistral4/convert_mistral4_weight_to_hf.py | 4 +- .../quantizers/quantizer_finegrained_fp8.py | 15 ++++--- .../quantizers/quantizer_mxfp4.py | 14 ++++-- .../quantizers/quantizer_nvfp4.py | 14 ++++-- tests/kernels/test_finegrained.py | 17 ++++---- .../deepseek_v4/test_modeling_deepseek_v4.py | 10 ++--- .../glm_moe_dsa/test_modeling_glm_moe_dsa.py | 4 +- tests/models/gpt_oss/test_modeling_gpt_oss.py | 4 +- .../qwen4_exp/test_modeling_qwen4_exp.py | 43 +++++++++++++++---- 19 files changed, 148 insertions(+), 51 deletions(-) diff --git a/docs/source/en/experts_interface.md b/docs/source/en/experts_interface.md index baada16c77d7..ab673c3d802a 100644 --- a/docs/source/en/experts_interface.md +++ b/docs/source/en/experts_interface.md @@ -117,7 +117,7 @@ The kernel is loaded lazily on the first forward. ### FP8 and FP4 quantized experts -DeepSeek-style checkpoints are usually pre-quantized and carry their own quantization config, so you don't need to pass a [`FineGrainedFP8Config`]. The `"deepgemm"` backend automatically picks the FP8 (or FP4 on Blackwell) grouped-GEMM kernel. DeepGEMM requires dynamic per-row activation scales (`activation_scheme="dynamic"`) and rejects static (per-tensor) activation quantization. +DeepSeek-style checkpoints are usually pre-quantized and carry their own quantization config, so you don't need to pass a [`FineGrainedConfig`]. The `"deepgemm"` backend automatically picks the FP8 (or FP4 on Blackwell) grouped-GEMM kernel. DeepGEMM requires dynamic per-row activation scales (`activation_scheme="dynamic"`) and rejects static (per-tensor) activation quantization. ```py from transformers import AutoModelForCausalLM @@ -133,14 +133,14 @@ For FP4-packed expert weights (DeepSeek V4-style), the GPU must be SM100+ (Black > [!NOTE] > On Blackwell (SM100+), the `"deepgemm"` and `"deepgemm_megamoe"` experts kernels require power-of-two UE8M0 expert scales. A checkpoint quantized with plain `float32` scales (`scale_fmt="float"`) raises a `ValueError` on the first forward instead of silently corrupting the output. Load a checkpoint quantized with `scale_fmt="ue8m0"`, or switch to `grouped_mm` or `batched_mm`, which consume `float32` block scales directly. Hopper (SM90+) consumes `float32` scales on the DeepGEMM path without conversion. -The main reason to pass a [`FineGrainedFP8Config`] for a pre-quantized checkpoint is to dequantize it back to `bfloat16`, in which case the experts run in `bfloat16` rather than on the FP8/FP4 DeepGEMM path. +The main reason to pass a [`FineGrainedConfig`] for a pre-quantized checkpoint is to dequantize it back to `bfloat16`, in which case the experts run in `bfloat16` rather than on the FP8/FP4 DeepGEMM path. ```py -from transformers import AutoModelForCausalLM, FineGrainedFP8Config +from transformers import AutoModelForCausalLM, FineGrainedConfig model = AutoModelForCausalLM.from_pretrained( "deepseek-ai/DeepSeek-V3", - quantization_config=FineGrainedFP8Config(dequantize=True), + quantization_config=FineGrainedConfig(dequantize=True), experts_implementation="deepgemm", ) ``` diff --git a/docs/source/en/model_doc/deepseek_v32.md b/docs/source/en/model_doc/deepseek_v32.md index ce932526d48b..c7cfdaf1c84a 100644 --- a/docs/source/en/model_doc/deepseek_v32.md +++ b/docs/source/en/model_doc/deepseek_v32.md @@ -58,11 +58,11 @@ In DeepSeek-V3.2 **every layer runs its own indexer** — there is no cross-laye DeepSeek-V3.2-Exp is distributed as an FP8 checkpoint. The indexer projections are kept out of FP8 quantization, since the checkpoint stores them in bf16/fp32: ```python -from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer +from transformers import FineGrainedConfig, AutoModelForCausalLM, AutoTokenizer import torch model_name = "deepseek-ai/DeepSeek-V3.2-Exp" -quantization_config = FineGrainedFP8Config( +quantization_config = FineGrainedConfig( modules_to_not_convert=["model.layers.*.mlp.gate.*", "*.self_attn.indexer.weights_proj.*"], weight_block_size=(128, 128), ) diff --git a/docs/source/en/model_doc/minimax_m3_vl.md b/docs/source/en/model_doc/minimax_m3_vl.md index 671d59922a44..a90c61ab4f26 100644 --- a/docs/source/en/model_doc/minimax_m3_vl.md +++ b/docs/source/en/model_doc/minimax_m3_vl.md @@ -153,7 +153,7 @@ from transformers import ( AutoModelForImageTextToText, AutoTokenizer, CompileConfig, - FineGrainedFP8Config, + FineGrainedConfig, ) from transformers.distributed import DistributedConfig @@ -164,7 +164,7 @@ model = AutoModelForImageTextToText.from_pretrained( "MiniMaxAI/MiniMax-M3-preview", dtype=torch.bfloat16, # Dequantize the native MXFP8 weights to bf16 at load (the speed win); needs even TP/EP sharding. - quantization_config=FineGrainedFP8Config(dequantize=True), + quantization_config=FineGrainedConfig(dequantize=True), distributed_config=DistributedConfig( tp_size=int(os.environ["WORLD_SIZE"]), enable_expert_parallel=True, diff --git a/docs/source/en/quantization/finegrained_fp8.md b/docs/source/en/quantization/finegrained_fp8.md index efc601c1cdae..afb55a2a7414 100644 --- a/docs/source/en/quantization/finegrained_fp8.md +++ b/docs/source/en/quantization/finegrained_fp8.md @@ -16,6 +16,13 @@ rendered properly in your Markdown viewer. # Fine-grained FP8 +> [!WARNING] +> `FineGrainedFP8Config` is frozen and receives no new recipes. [`FineGrainedConfig`] supersedes it, serving +> block-FP8, MXFP8, MXFP4 and NVFP4 through one config for dense linears, embeddings and MoE +> experts alike. Existing checkpoints keep loading — the `fp8`, `mxfp8`, `mxfp4` and `nvfp4` +> `quant_method` keys now route to the fine-grained integration — so this page is here for +> reference while the frozen classes remain. They will be removed in a future release. + Fine-grained FP8 quantization quantizes the weights and activations to fp8. - The weights are quantized to 8-bits for each 2D block (`weight_block_size=(128, 128)`). diff --git a/docs/source/en/quantization/mxfp4.md b/docs/source/en/quantization/mxfp4.md index f96833c0cc09..09ac3c73a56d 100644 --- a/docs/source/en/quantization/mxfp4.md +++ b/docs/source/en/quantization/mxfp4.md @@ -16,6 +16,13 @@ rendered properly in your Markdown viewer. # MXFP4 +> [!WARNING] +> `Mxfp4Config` is frozen and receives no new recipes. [`FineGrainedConfig`] supersedes it, serving +> block-FP8, MXFP8, MXFP4 and NVFP4 through one config for dense linears, embeddings and MoE +> experts alike. Existing checkpoints keep loading — the `fp8`, `mxfp8`, `mxfp4` and `nvfp4` +> `quant_method` keys now route to the fine-grained integration — so this page is here for +> reference while the frozen classes remain. They will be removed in a future release. + Note: MXFP4 quantization currently only works for OpenAI GPT-OSS 120b and 20b. MXFP4 is a 4-bit floating point format that dramatically reduces the memory requirements of large models. Large models (GPT-OSS-120B) can fit on a single 80GB GPU and smaller models (GPT-OSS-20B) only require 16GB of memory. It uses blockwise scaling to preserve its range and accuracy, which typically becomes degraded at lower precisions. diff --git a/docs/source/en/quantization/nvfp4.md b/docs/source/en/quantization/nvfp4.md index 42a2d31cfa5a..1ea263317c15 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -16,6 +16,13 @@ rendered properly in your Markdown viewer. # NVFP4 +> [!WARNING] +> `NVFP4Config` is frozen and receives no new recipes. [`FineGrainedConfig`] supersedes it, serving +> block-FP8, MXFP8, MXFP4 and NVFP4 through one config for dense linears, embeddings and MoE +> experts alike. Existing checkpoints keep loading — the `fp8`, `mxfp8`, `mxfp4` and `nvfp4` +> `quant_method` keys now route to the fine-grained integration — so this page is here for +> reference while the frozen classes remain. They will be removed in a future release. + NVFP4 quantization packs full-precision linear weights into NVIDIA's 4-bit floating-point format while a model is loaded. [`NVFP4Config`] replaces eligible bias-free `torch.nn.Linear` modules, whose `in_features` and `out_features` are both divisible by 16, with an NVFP4 linear implementation from the [NVFP4 Hub kernel](https://huggingface.co/kernels-community/nvfp4-gemm). The model's attention and MLP interfaces are diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 084b17733b16..5cc3bcdae74d 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -15,6 +15,7 @@ import functools import os +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -38,6 +39,15 @@ from .moe import ExpertsInterface, use_experts_implementation +warnings.warn( + "`transformers.integrations.finegrained_fp8` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index acfe79de4f93..b8a01bcae243 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -13,6 +13,8 @@ # limitations under the License. +import warnings + from ..utils import is_torch_available, logging @@ -25,6 +27,15 @@ from ..quantizers.quantizers_utils import get_module_from_name, on_device, should_convert_module +warnings.warn( + "`transformers.integrations.mxfp4` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) diff --git a/src/transformers/integrations/nvfp4.py b/src/transformers/integrations/nvfp4.py index 476e459df405..cb74485ea95b 100644 --- a/src/transformers/integrations/nvfp4.py +++ b/src/transformers/integrations/nvfp4.py @@ -14,6 +14,7 @@ from __future__ import annotations import functools +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -28,6 +29,15 @@ from .hub_kernels import lazy_load_kernel +warnings.warn( + "`transformers.integrations.nvfp4` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + if is_torch_available(): import torch import torch.nn as nn diff --git a/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py b/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py index 5546d83c3291..f270a0c31c9f 100644 --- a/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py +++ b/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py @@ -29,7 +29,7 @@ PixtralProcessor, PixtralVisionConfig, ) -from transformers.integrations.finegrained_fp8 import replace_with_fp8_linear +from transformers.integrations.finegrained import replace_with_finegrained_layer from transformers.integrations.mistral import convert_tekken_tokenizer from transformers.quantizers.auto import AutoQuantizationConfig @@ -254,9 +254,9 @@ def convert_and_write_model(input_dir: str, output_dir: str, max_position_embedd else: raise ValueError(f"Unknown config type {type(config)}.") - # let's swap nn.Linear to FP8 Linear before loading + # swap nn.Linear for the quantized module before loading if hasattr(model.config, "quantization_config"): - model = replace_with_fp8_linear( + model = replace_with_finegrained_layer( model, model.config.quantization_config.modules_to_not_convert, model.config.quantization_config ) diff --git a/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py b/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py index b3c3b7714fca..4997f4f633b4 100644 --- a/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py +++ b/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py @@ -37,7 +37,7 @@ MergeModulelist, WeightRenaming, ) -from transformers.integrations.finegrained_fp8 import replace_with_fp8_linear +from transformers.integrations.finegrained import replace_with_finegrained_layer from transformers.integrations.mistral import convert_tekken_tokenizer from transformers.models.mistral4.modeling_mistral4 import Mistral4ForCausalLM from transformers.quantizers.auto import AutoQuantizationConfig @@ -536,7 +536,7 @@ def convert_and_write_model( if output_fp8 and hasattr(model.config, "quantization_config"): qconfig = model.config.quantization_config - model = replace_with_fp8_linear(model, qconfig.modules_to_not_convert, qconfig) + model = replace_with_finegrained_layer(model, qconfig.modules_to_not_convert, qconfig) model.load_state_dict(full_state_dict, strict=True, assign=True) model.save_pretrained(str(output_dir)) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 70653a5b759f..dde5b3691bcc 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -1,3 +1,4 @@ +import warnings from typing import TYPE_CHECKING from ..utils import is_accelerate_available, is_torch_available, is_torch_xpu_available, logging @@ -5,6 +6,15 @@ from .quantizers_utils import get_module_from_name +warnings.warn( + "`transformers.quantizers.quantizer_finegrained_fp8` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + if is_torch_available(): import torch @@ -26,11 +36,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) - logger.warning_once( - "`FineGrainedFP8Config` is frozen for backward compatibility and receives no new " - "recipes. `FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, " - "weight-only)." - ) def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index f18ffac9a539..14222c668fcd 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import warnings from typing import TYPE_CHECKING from .base import HfQuantizer @@ -31,6 +32,15 @@ from .quantizers_utils import get_module_from_name +warnings.warn( + "`transformers.quantizers.quantizer_mxfp4` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + if is_torch_available(): import torch @@ -50,10 +60,6 @@ class Mxfp4HfQuantizer(HfQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) - logger.warning_once( - "`Mxfp4Config` is frozen for backward compatibility and receives no new recipes. " - "`FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." - ) self.triton_kernels_hub = None def _lazy_import_kernels(self): diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index 78f2d6a494ec..3215d28777a8 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import warnings from typing import TYPE_CHECKING from ..utils import is_kernels_available, is_torch_available, logging @@ -21,6 +22,15 @@ from .quantizers_utils import get_module_from_name +warnings.warn( + "`transformers.quantizers.quantizer_nvfp4` is frozen and receives no new recipes. `FineGrainedConfig` supersedes it " + "(block-FP8, MXFP8, MXFP4, NVFP4, weight-only); this module will be removed in a future " + "release.", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) @@ -50,10 +60,6 @@ class NVFP4HfQuantizer(HfQuantizer): def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) - logger.warning_once( - "`NVFP4Config` is frozen for backward compatibility and receives no new recipes. " - "`FineGrainedConfig` supersedes it (block-FP8, MXFP8, MXFP4, NVFP4, weight-only)." - ) def validate_environment(self, device_map, **kwargs): if self.pre_quantized: diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index b0b631df3f8c..1dc00e88e3ae 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -641,17 +641,18 @@ def test_frozen_module_warns_and_is_self_contained(self): from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer from transformers.utils.quantization_config import FineGrainedFP8Config - # the notice fires when someone actually asks for the frozen scheme, not on import: - # importing is what every `transformers.integrations` consumer does + # the whole FILE is deprecated, so importing anything out of it is what warns. Safe to + # do at module scope only because `_LazyModule` never executes these unless asked: a + # plain `import transformers` does not reach them. with warnings.catch_warnings(record=True) as raised: warnings.simplefilter("always") importlib.reload(frozen) - self.assertEqual([w for w in raised if "frozen" in str(w.message)], []) - - # ...and it rides the library logger, so it honours transformers' verbosity settings - with self.assertLogs("transformers.quantizers.quantizer_finegrained_fp8", level="WARNING") as captured: - FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) - self.assertTrue(any("frozen for backward compatibility" in line for line in captured.output)) + self.assertTrue( + [w for w in raised if issubclass(w.category, DeprecationWarning) and "frozen" in str(w.message)], + [str(w.message) for w in raised], + ) + # the quantizer is still constructible — frozen means no new recipes, not removed + FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) # distinct machinery: the frozen classes are not the live ones self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) diff --git a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py index 413b92841d7c..ca13e05a7813 100644 --- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py +++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py @@ -38,7 +38,7 @@ AutoModelForCausalLM, AutoTokenizer, DeepseekV4Model, - FineGrainedFP8Config, + FineGrainedConfig, ) from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester @@ -239,7 +239,7 @@ class DeepseekV4IntegrationTest(unittest.TestCase): """End-to-end check on the published DeepSeek-V4-Flash checkpoint. Loads the real 43-layer FP8 weights, dequantizes on the fly via - :class:`FineGrainedFP8Config`, and greedy-generates a continuation of a fixed + :class:`FineGrainedConfig`, and greedy-generates a continuation of a fixed prompt. The forward path that this test covers is everything past the typical tiny-config tests can reach: the per-layer FP8 dequant in ``update_weight_conversions``, the ``compress_ratios → layer_types`` config @@ -270,7 +270,7 @@ class DeepseekV4IntegrationTest(unittest.TestCase): ) ) def test_v4_flash_dequantized_generation(self): - quantization_config = FineGrainedFP8Config(dequantize=True) + quantization_config = FineGrainedConfig(dequantize=True) config = AutoConfig.from_pretrained(self.model_id) tokenizer = AutoTokenizer.from_pretrained(self.model_id) model = AutoModelForCausalLM.from_pretrained( @@ -388,7 +388,7 @@ def test_v4_flash_dequantized_chat_seven_prompts(self): ), ] - quantization_config = FineGrainedFP8Config(dequantize=True) + quantization_config = FineGrainedConfig(dequantize=True) config = AutoConfig.from_pretrained(self.model_id) tokenizer = AutoTokenizer.from_pretrained(self.model_id) model = AutoModelForCausalLM.from_pretrained( @@ -432,7 +432,7 @@ def test_v4_flash_dequantized_chat_seven_prompts(self): from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed import DistributedConfig -from transformers.utils.quantization_config import FineGrainedFP8Config +from transformers.utils.quantization_config import FineGrainedConfig LOADTIME_DISPATCH = {loadtime_dispatch!r} diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py index a21a81b5f677..c5335f780b41 100644 --- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py +++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py @@ -22,7 +22,7 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, - FineGrainedFP8Config, + FineGrainedConfig, GlmMoeDsaConfig, is_torch_available, set_seed, @@ -164,7 +164,7 @@ def test_glm_moe_dsa_fp8_inference(self): set_seed(0) # different ranks need the same seed model_id = "zai-org/GLM-5-FP8" - quantization_config = FineGrainedFP8Config( + quantization_config = FineGrainedConfig( modules_to_not_convert=[ "model.layers.*.mlp.gate$", "model.layers.*.self_attn.indexer.weights_proj$", diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 936131b71f1d..fd6c1b0c2518 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -28,7 +28,7 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, - Mxfp4Config, + FineGrainedConfig, is_torch_available, ) from transformers.testing_utils import ( @@ -547,7 +547,7 @@ def test_training_step(self, quantized, model, kernels, attn_impl, mode): device_map="auto", attn_implementation=attn_impl, use_kernels=kernels, - quantization_config=Mxfp4Config(dequantize=True), + quantization_config=FineGrainedConfig(quant_method="mxfp4", dequantize=True), ) model_obj.train() diff --git a/tests/models/qwen4_exp/test_modeling_qwen4_exp.py b/tests/models/qwen4_exp/test_modeling_qwen4_exp.py index 640d89a67416..e4d899ec3d74 100644 --- a/tests/models/qwen4_exp/test_modeling_qwen4_exp.py +++ b/tests/models/qwen4_exp/test_modeling_qwen4_exp.py @@ -35,7 +35,6 @@ from transformers import ( AutoModelForCausalLM, DynamicCache, - FineGrainedFP8Config, Qwen4ExpConfig, Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration, @@ -46,8 +45,9 @@ StaticCache, ) from transformers.distributed.fsdp import verify_fsdp_plan - from transformers.integrations import FP8Embedding - from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + from transformers.integrations.finegrained import FineGrainedEmbedding + from transformers.quantizers.finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig class Qwen4ExpTextModelTester(CausalLMModelTester): @@ -196,17 +196,44 @@ def test_ple_layers_must_use_linear_attention(self): layer_types=["linear_attention", "qwen_sparse_attention"], ) - def test_finegrained_fp8_embedding_conversion(self): + def test_finegrained_embedding_conversion(self): + """The n-gram table through the CURRENT integration, not the frozen one: a table too big + to dequantize at load is held in FP8 with one per-tensor scale, and the rescale runs on + the rows a lookup gathers.""" config = self.model_tester.get_config() with torch.device("meta"): model = Qwen4ExpForCausalLM(config) - quantization_config = FineGrainedFP8Config( - modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[] + quantizer = FineGrainedHfQuantizer( + FineGrainedConfig(modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[]) ) - FineGrainedFP8HfQuantizer(quantization_config).preprocess_model(model) + quantizer.pre_quantized = True + quantizer.preprocess_model(model) - self.assertIsInstance(model.model.layers[0].ple.ple_embedding.ngram_embedding, FP8Embedding) + table = model.model.layers[0].ple.ple_embedding.ngram_embedding + self.assertIsInstance(table, FineGrainedEmbedding) + self.assertEqual(table.weight.dtype, torch.float8_e4m3fn) + self.assertEqual(table.weight_scale.shape, torch.Size([1])) + + def test_a_non_fp8_format_says_the_table_is_still_fp8(self): + """`FineGrainedEmbedding` takes no weight format — a table is FP8 whatever was asked for. + It used to be silent, so an NVFP4 run returned an FP8 table with no way to tell.""" + config = self.model_tester.get_config() + with torch.device("meta"): + model = Qwen4ExpForCausalLM(config) + + quantizer = FineGrainedHfQuantizer( + FineGrainedConfig( + quant_method="nvfp4", + modules_to_convert=["ple.ple_embedding.ngram_embedding"], + modules_to_not_convert=[], + ) + ) + quantizer.pre_quantized = True + with self.assertLogs("transformers.quantizers.finegrained.base", level="WARNING") as logs: + quantizer.preprocess_model(model) + self.assertTrue(any("no embedding path" in line for line in logs.output), logs.output) + self.assertEqual(model.model.layers[0].ple.ple_embedding.ngram_embedding.weight.dtype, torch.float8_e4m3fn) def test_ple_padding_and_static_cache_match_unpadded_sequence(self): torch.manual_seed(0) From fab7f74160cea3a77349831c2c25380a06fcf622 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 13:02:05 -0700 Subject: [PATCH 46/68] Bring two comment blocks under the noisy-comment limits `utils/check_noisy_comments.py` caps a block at 5 lines and 500 characters, and these were 10/679 and 6/504. Both keep the part that is not in the code: why the sharded comparison uses each model's measured floor rather than a fixed tolerance, and why the tied-head check reads the config flag instead of testing `head.weight is embed.weight` (tying happens later, in `post_init`). Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 10 ++++------ tests/kernels/test_finegrained.py | 15 +++++---------- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 8e0969529025..7ea273e3b0f6 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,12 +80,10 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) - # A tied lm_head IS the input embedding's parameter, so the `embedding_rowwise` entry - # `PretrainedConfig` adds for tied weights shards the tensor it reads — while lm_head - # itself sits outside any `base_model_tp_plan` and gets no style, leaving `F.linear` with - # a plain input against a DTensor weight. FSDP resolves the pair; TP had no equivalent. - # the CONFIG flag, not `head.weight is embed.weight`: tying happens later in `post_init`, - # so the parameters are still distinct here and the identity test silently misses + # A tied lm_head IS the embedding's parameter, so the `embedding_rowwise` entry shards + # the tensor it reads while lm_head itself gets no style — leaving `F.linear` a plain + # input against a DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: + # tying happens later in `post_init`, so an identity test here silently misses. tied = getattr(self.config.get_text_config(), "tie_word_embeddings", False) if tied and "embedding_rowwise" in self._tp_plan.values(): head = self.get_output_embeddings() diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index b0b631df3f8c..243676fc5506 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -2746,16 +2746,11 @@ def test_every_load_path_agrees(self): f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " f"this leg cannot catch a bad axis", ) - # Sharding reorders reductions (a rowwise all-reduce sums the same terms - # in another order), so the legs differ by BF16 rounding at the first - # sharded op whatever the axes. What that becomes at the logits is the - # MODEL's business: the NVFP4 fixture amplifies it ~100x across two MoE - # layers (measured: a 1e-3 perturbation moves its logits by 0.99 on ONE - # device), while the dense fixtures barely amplify at all. A fixed - # tolerance therefore asks the impossible of one model and nothing of - # another. Compare against the model's OWN floor instead: perturb the - # unsharded run by the same rounding and take the resulting logit shift as - # the budget. A wrong shard axis lands orders of magnitude above it. + # Sharding reorders reductions, so the legs differ by BF16 rounding + # whatever the axes — and how much that shows at the logits is the + # model's own business (the NVFP4 fixture amplifies ~100x, the dense + # ones barely). Hence each model's measured floor, not a fixed + # tolerance. A wrong shard axis lands orders of magnitude above it. budget = max(2e-2, noise_floor) gap = (payload["logits"] - reference).abs().max().item() self.assertLessEqual( From 21c0a070999f5ca04b470192179402b19841059d Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 13:08:04 -0700 Subject: [PATCH 47/68] Bring two comment blocks under the noisy-comment limits `utils/check_noisy_comments.py` caps a block at 5 lines and 500 characters, and these were 10/679 and 6/504. Both keep the part that is not in the code: why the sharded comparison uses each model's measured floor rather than a fixed tolerance, and why the tied-head check reads the config flag instead of testing `head.weight is embed.weight` (tying happens later, in `post_init`). Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 10 ++++------ tests/kernels/test_finegrained.py | 15 +++++---------- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 8e0969529025..7ea273e3b0f6 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,12 +80,10 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) - # A tied lm_head IS the input embedding's parameter, so the `embedding_rowwise` entry - # `PretrainedConfig` adds for tied weights shards the tensor it reads — while lm_head - # itself sits outside any `base_model_tp_plan` and gets no style, leaving `F.linear` with - # a plain input against a DTensor weight. FSDP resolves the pair; TP had no equivalent. - # the CONFIG flag, not `head.weight is embed.weight`: tying happens later in `post_init`, - # so the parameters are still distinct here and the identity test silently misses + # A tied lm_head IS the embedding's parameter, so the `embedding_rowwise` entry shards + # the tensor it reads while lm_head itself gets no style — leaving `F.linear` a plain + # input against a DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: + # tying happens later in `post_init`, so an identity test here silently misses. tied = getattr(self.config.get_text_config(), "tie_word_embeddings", False) if tied and "embedding_rowwise" in self._tp_plan.values(): head = self.get_output_embeddings() diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 1dc00e88e3ae..c239bdef79a9 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -2747,16 +2747,11 @@ def test_every_load_path_agrees(self): f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " f"this leg cannot catch a bad axis", ) - # Sharding reorders reductions (a rowwise all-reduce sums the same terms - # in another order), so the legs differ by BF16 rounding at the first - # sharded op whatever the axes. What that becomes at the logits is the - # MODEL's business: the NVFP4 fixture amplifies it ~100x across two MoE - # layers (measured: a 1e-3 perturbation moves its logits by 0.99 on ONE - # device), while the dense fixtures barely amplify at all. A fixed - # tolerance therefore asks the impossible of one model and nothing of - # another. Compare against the model's OWN floor instead: perturb the - # unsharded run by the same rounding and take the resulting logit shift as - # the budget. A wrong shard axis lands orders of magnitude above it. + # Sharding reorders reductions, so the legs differ by BF16 rounding + # whatever the axes — and how much that shows at the logits is the + # model's own business (the NVFP4 fixture amplifies ~100x, the dense + # ones barely). Hence each model's measured floor, not a fixed + # tolerance. A wrong shard axis lands orders of magnitude above it. budget = max(2e-2, noise_floor) gap = (payload["logits"] - reference).abs().max().item() self.assertLessEqual( From 12f5aa8e5c12294abc826eb2ae9608eae2788ad1 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 13:41:44 -0700 Subject: [PATCH 48/68] Do not let the tied-head plan raise in post_init `init_parallel_plans` runs from `post_init`, so it is on the constructor path of every model in the library -- and it called `config.get_text_config()` bare. That RAISES on a composite model carrying more than one text sub-config: musicgen has `text_encoder` and `decoder`, so every musicgen and musicgen-melody test died before doing anything, along with `TestGetEncoder`. Falls back to the top-level config when the call cannot disambiguate, which is the same answer for `tie_word_embeddings` in every case that used to work. Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 7ea273e3b0f6..f228f356e7f1 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,11 +80,16 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) - # A tied lm_head IS the embedding's parameter, so the `embedding_rowwise` entry shards - # the tensor it reads while lm_head itself gets no style — leaving `F.linear` a plain - # input against a DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: - # tying happens later in `post_init`, so an identity test here silently misses. - tied = getattr(self.config.get_text_config(), "tie_word_embeddings", False) + # A tied lm_head IS the embedding's parameter, so `embedding_rowwise` shards the tensor + # it reads while lm_head gets no style — leaving `F.linear` a plain input against a + # DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: tying happens in + # `post_init`, after this. And `get_text_config` RAISES on a composite model with two + # text sub-configs (musicgen), which this would hit in every model's `post_init`. + try: + text_config = self.config.get_text_config() + except ValueError: + text_config = self.config + tied = getattr(text_config, "tie_word_embeddings", False) if tied and "embedding_rowwise" in self._tp_plan.values(): head = self.get_output_embeddings() if head is not None: From 7abcedf9ba0907b6c1410cb18bd73fe5440a7e9d Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 13:41:44 -0700 Subject: [PATCH 49/68] Do not let the tied-head plan raise in post_init `init_parallel_plans` runs from `post_init`, so it is on the constructor path of every model in the library -- and it called `config.get_text_config()` bare. That RAISES on a composite model carrying more than one text sub-config: musicgen has `text_encoder` and `decoder`, so every musicgen and musicgen-melody test died before doing anything, along with `TestGetEncoder`. Falls back to the top-level config when the call cannot disambiguate, which is the same answer for `tie_word_embeddings` in every case that used to work. Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 7ea273e3b0f6..f228f356e7f1 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,11 +80,16 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) - # A tied lm_head IS the embedding's parameter, so the `embedding_rowwise` entry shards - # the tensor it reads while lm_head itself gets no style — leaving `F.linear` a plain - # input against a DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: - # tying happens later in `post_init`, so an identity test here silently misses. - tied = getattr(self.config.get_text_config(), "tie_word_embeddings", False) + # A tied lm_head IS the embedding's parameter, so `embedding_rowwise` shards the tensor + # it reads while lm_head gets no style — leaving `F.linear` a plain input against a + # DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: tying happens in + # `post_init`, after this. And `get_text_config` RAISES on a composite model with two + # text sub-configs (musicgen), which this would hit in every model's `post_init`. + try: + text_config = self.config.get_text_config() + except ValueError: + text_config = self.config + tied = getattr(text_config, "tie_word_embeddings", False) if tied and "embedding_rowwise" in self._tp_plan.values(): head = self.get_output_embeddings() if head is not None: From db85a446ea7f4baeb19a49dfedafb81fd219b1d3 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 23 Sep 2026 23:57:31 -0700 Subject: [PATCH 50/68] Address the second review round Conversions: * `_dequantize_one` inverted the scale grid to recover the block, which is only valid when the shape tiles. DeepSeek-V3's `(576, 7168)` against a `(5, 56)` grid raised, and `(192, 256)` silently used block 96 in place of 128. It now takes the block from the config when that config's ceil-grid reproduces the scale it was handed, and pads-then-crops the way the quantizer already did. * helpers used across ops are public (`keyed_by_target`, `held_scale`, `as_container`); `_global_role` is a static method of the one op that reads it. Tests move to `tests/integrations/finegrained/`, next to the integration they cover rather than under `tests/kernels/`, split by what each file answers: `test_core` what the integration passes, `test_conversions` what it loads, `test_forwards` what the kernels return, `test_models` what it shards. The frozen fp8 integration's tests follow their module to `tests/integrations/`. `test_forwards` ran nowhere it should have: `@require_torch_gpu` plus a capability check kept it off Hopper, where these kernels work, and off XPU, which the integration supports; and `setUpClass` turned any `ImportError` into a skip, so a broken loader or a missing build passed silently. It is `@require_kernels` + `@require_torch_accelerator` now, device-agnostic, with the load unguarded. Also: `_CATCH_ALL` inlined, `_owns_expert` folded into the placement check it duplicated, the wording the review suggested, and a note in the quantization docs. Co-Authored-By: Claude Opus 5 --- docs/source/en/main_classes/quantization.md | 5 + .../distributed/sharding_utils.py | 19 +- .../integrations/finegrained/conversions.py | 83 +- .../integrations/finegrained/core.py | 14 +- src/transformers/utils/quantization_config.py | 13 +- tests/integrations/finegrained/__init__.py | 0 .../finegrained/test_conversions.py | 694 +++++ tests/integrations/finegrained/test_core.py | 1402 +++++++++ .../integrations/finegrained/test_forwards.py | 372 +++ tests/integrations/finegrained/test_models.py | 561 ++++ .../test_finegrained_fp8.py | 0 tests/kernels/test_finegrained.py | 2769 ----------------- tests/utils/test_core_model_loading.py | 178 -- 13 files changed, 3099 insertions(+), 3011 deletions(-) create mode 100644 tests/integrations/finegrained/__init__.py create mode 100644 tests/integrations/finegrained/test_conversions.py create mode 100644 tests/integrations/finegrained/test_core.py create mode 100644 tests/integrations/finegrained/test_forwards.py create mode 100644 tests/integrations/finegrained/test_models.py rename tests/{kernels => integrations}/test_finegrained_fp8.py (100%) delete mode 100644 tests/kernels/test_finegrained.py diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index ea00c4c8eb29..d9da976a992c 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -100,6 +100,11 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui ## FineGrainedConfig +> [!WARNING] +> `FineGrainedConfig` supersedes `FineGrainedFP8Config`, `Mxfp4Config` and `NVFP4Config`, which are +> frozen and will be removed in a future release. The `fp8`, `mxfp8`, `mxfp4` and `nvfp4` +> `quant_method` keys already route here, so existing checkpoints need no change. + [[autodoc]] FineGrainedConfig ## FineGrainedGroup diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 4af642ca42ae..21839e307af4 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -155,18 +155,20 @@ def shard_tensor( # MoE path # tensor_idx identifies the axis-0 piece in param space (not in source.shape). + normalized_dim_placements = [ + (mesh_dim, placement, self._normalize_param_dim(placement.dim)) for mesh_dim, placement in dim_placements + ] + # if this rank owns expert `tensor_idx` along axis 0, we need to slice the inner dimensions, else we drop it - if not self._owns_expert(tensor_idx, dim_placements): + has_axis0_shard = any(param_dim == 0 for _, _, param_dim in normalized_dim_placements) + owns_tensor_idx = self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size + if has_axis0_shard and not owns_tensor_idx: return None # a 0-dim per-expert value has no axis to slice, so this rank takes it whole if not source_shape: return source[...].to(device=device, dtype=dtype) - normalized_dim_placements = [ - (mesh_dim, placement, self._normalize_param_dim(placement.dim)) for mesh_dim, placement in dim_placements - ] - # `param_dim` indexes the full parameter layout [N, in, out] (expert axis first). # In per-expert loading, leading axis is absent ([in, out]). # Therefore we need to shift the dimensions by 1 to the left so that sharding operations can be applied to the inner dimensions. @@ -322,13 +324,6 @@ def _normalize_param_dim(self, dim: int) -> int: # if dim is negative, it should be normalized to the last axis return dim if dim >= 0 else self.param_ndim + dim - def _owns_expert(self, tensor_idx: int, dim_placements: list) -> bool: - """Whether this rank holds expert `tensor_idx`. True when nothing shards axis 0 — the - experts are replicated, so every rank owns every one of them.""" - if not any(self._normalize_param_dim(placement.dim) == 0 for _, placement in dim_placements): - return True - return self._axis0_offset <= tensor_idx < self._axis0_offset + self._axis0_local_size - def _dtensor_from_local_like(local_tensor: torch.Tensor, ref: DTensor) -> DTensor: """Wrap `local_tensor` as a DTensor that mirrors `ref`'s mesh, placements, diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 71ae8ac5bd03..f46fe46f30b6 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -32,6 +32,7 @@ _FP8_MIN, FineGrainedExperts, WeightFormat, + _cdiv, _FineGrainedModule, _get_ue8m0_dtype, load_finegrained_kernel, @@ -39,7 +40,7 @@ ) -def _keyed_by_target(input_dict: dict, target_patterns) -> dict: +def keyed_by_target(input_dict: dict, target_patterns) -> dict: """A converter's ops see their single tensor under the SOURCE pattern until an op re-keys it to the target (the core ops do; the loader then expands the target into the full name) — so a layout op that is first in its converter re-keys the same way. Multi-tensor dicts (a @@ -50,7 +51,7 @@ def _keyed_by_target(input_dict: dict, target_patterns) -> dict: return input_dict -def _held_scale(model, full_layer_name, key): +def held_scale(model, full_layer_name, key): """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" @@ -68,20 +69,6 @@ def as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) -def _global_role(key: str) -> tuple[str, str]: - """What a global-scale key holds, for a checkpoint key and a module target alike: which - projection it belongs to, and which of modelopt's two levels it is — the weight's second-level - global (``weight_scale_2``) or the calibrated activation one (``input_scale``). - - This runs while the converters are being BUILT, against raw checkpoint keys, so the vLLM-style - `w1`/`w2`/`w3` names are still in play and have not been renamed to `gate/down/up_proj` yet. - Matched on path SEGMENTS: the default is gate_up, so a loose `"w2" in key` would silently - bucket a down global as one.""" - down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) - level = "input" if "input_scale" in key or "input_global" in key else "weight" - return ("down" if down else "gate_up"), level - - class _FineGrainedOp(ConversionOps): """Base for the ops below. Each carries the quantizer that built it, which is how an op reaches the config (the activation format, the block size) while it runs. @@ -106,7 +93,7 @@ class FineGrainedInterleaveGateUp(_FineGrainedOp): save wherever the checkpoint's row order differs from the one the experts hold.""" def convert(self, input_dict, model=None, target_patterns=None, **kwargs): - input_dict = _keyed_by_target(input_dict, target_patterns) + input_dict = keyed_by_target(input_dict, target_patterns) experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): return input_dict @@ -130,7 +117,7 @@ class FineGrainedScaleContainer(_FineGrainedOp): def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): + for key, value in keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value if self.inverse: # one checkpoint, one container, recorded on the modules at load; e8m0 is only @@ -140,7 +127,7 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= if container is not None and value.dtype == _get_ue8m0_dtype(): value = as_container(value, container) else: - module, held = _held_scale(model, full_layer_name, key) + module, held = held_scale(model, full_layer_name, key) if held is not None and value.dtype != held.dtype: module.scale_container_dtype = value.dtype value = as_container(value, held.dtype) @@ -155,7 +142,7 @@ class FineGrainedSwizzleScales(_FineGrainedOp): def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): + for key, value in keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) return out @@ -170,7 +157,7 @@ def _unswizzle(value): @staticmethod def _swizzle(key, value, model, full_layer_name): - _, held = _held_scale(model, full_layer_name, key) + _, held = held_scale(model, full_layer_name, key) if held is None or held.ndim != 5 or value.ndim == 5: return value return load_finegrained_kernel().swizzle_mx_scales(value) @@ -183,7 +170,7 @@ class FineGrainedPackedBlocks(_FineGrainedOp): def convert(self, input_dict, target_patterns=None, **kwargs): out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): + for key, value in keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) out[key] = value @@ -202,7 +189,7 @@ class FineGrainedViewPackedInt8(_FineGrainedOp): def convert(self, input_dict, target_patterns=None, **kwargs): held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) out = {} - for key, value in _keyed_by_target(input_dict, target_patterns).items(): + for key, value in keyed_by_target(input_dict, target_patterns).items(): value = value[0] if isinstance(value, list) else value out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value return out @@ -220,15 +207,29 @@ class FineGrainedWeightGlobals(_FineGrainedOp): global keeps the requant on the range the e4m3 block scales were chosen for. Rescaling the up half's block scales instead would re-round them against codes chosen for the old global.""" + @staticmethod + def role(key: str) -> tuple[str, str]: + """What a global-scale key holds, for a checkpoint key and a module target alike: which + projection it belongs to, and which of modelopt's two levels it is — the weight's + second-level global (``weight_scale_2``) or the calibrated activation one (``input_scale``). + + This runs while the converters are being BUILT, against raw checkpoint keys, so the + vLLM-style `w1`/`w2`/`w3` names are still in play and have not been renamed to + `gate/down/up_proj` yet. Matched on path SEGMENTS: the default is gate_up, so a loose + `"w2" in key` would silently bucket a down global as one.""" + down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) + level = "input" if "input_scale" in key or "input_global" in key else "weight" + return ("down" if down else "gate_up"), level + def convert(self, input_dict, target_patterns=None, model=None, **kwargs): sources = defaultdict(dict) for key, value in input_dict.items(): # the loader hands a converter's tensors in a list; unwrapped, a fused `(E, 2)` pair # reads as `(1, 2E)` and the per-half fold below silently does not fire - sources[_global_role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value + sources[self.role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value globals_ = {} for target in target_patterns: - if role_sources := sources.get(_global_role(target)): + if role_sources := sources.get(self.role(target)): # one (E, calibrated projections) fp32 tensor, gate column first: two sources are # the stack's halves, one is the fused (E, 2) pair or a single projection's (E,) columns = [] @@ -236,7 +237,7 @@ def convert(self, input_dict, target_patterns=None, model=None, **kwargs): value = torch.stack(value, dim=0) if isinstance(value, list) else value columns.append(value.float().reshape(value.shape[0], -1)) globals_[target] = torch.cat(columns, dim=1) - targets = {_global_role(target): target for target in globals_} + targets = {self.role(target): target for target in globals_} stack = targets.get(("gate_up", "weight")) # a stack calibrated per half keeps the gate's global; the up half's leaves as a ratio, @@ -368,7 +369,7 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key module, scale_name = holder if holder is not None else (None, None) value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) - held_scale = getattr(module, scale_name) if module is not None else None + held = getattr(module, scale_name) if module is not None else None format_spec = weight_formats()[module.weight_format] if module is not None else None if format_spec is None or format_spec.scale_group is None: if module is not None and module.block_size: @@ -384,8 +385,8 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc ) block = tuple(config_block) if config_block else (value.shape[-2], value.shape[-1]) ue8m0 = ( - held_scale.dtype == _get_ue8m0_dtype() - if held_scale is not None + held.dtype == _get_ue8m0_dtype() + if held is not None else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" ) weight, scale = self._quantize_block_fp8(value, block, ue8m0) @@ -402,8 +403,8 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc f"accelerator, but the weight is on {value.device.type}" ) weight, scale, global_scale = self._quantize_group(value, format_spec) - scale = as_container(scale, held_scale.dtype) - if held_scale.ndim == 5: + scale = as_container(scale, held.dtype) + if held.ndim == 5: scale = load_finegrained_kernel().swizzle_mx_scales(scale) out = {key: weight, f"{prefix}_scale_inv": scale} if global_scale is not None: @@ -542,12 +543,18 @@ def _dequantize_one( except Exception: # scale can be a single tensor in extreme cases where it was not wrapped properly but is [1,0]. scale_rows, scale_cols = 1, 1 - if rows % scale_rows or cols % scale_cols: + # the grid ROUNDS UP, so it does not invert to the block (DSv3's `(576, 7168)` ships a + # `(5, 56)` grid against a 128x128 block); only guess when no config reproduces it + configured = getattr(getattr(self.hf_quantizer, "quantization_config", None), "weight_block_size", None) + if configured and (_cdiv(rows, configured[0]), _cdiv(cols, configured[1])) == (scale_rows, scale_cols): + block_m, block_n = configured + elif rows % scale_rows or cols % scale_cols: raise ValueError( f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." ) - block_m = rows // scale_rows - block_n = cols // scale_cols + else: + block_m = rows // scale_rows + block_n = cols // scale_cols # the math runs in fp32 either way (``float8_e8m0fnu`` has no ``mul`` kernel); the # destination parameter's dtype wins when known, so an eager module keeps the model's if output_dtype is None: @@ -560,9 +567,13 @@ def _dequantize_one( else: s_fp32 = scales.to(torch.float32) original_shape = quantized_fp32.shape - blocked = quantized_fp32.reshape(-1, scale_rows, block_m, scale_cols, block_n) + # pad out to whole blocks so a ceil-rounded grid lines up, then crop back + pad_m, pad_n = -rows % block_m, -cols % block_n + padded = torch.nn.functional.pad(quantized_fp32, (0, pad_n, 0, pad_m)) + blocked = padded.reshape(-1, scale_rows, block_m, scale_cols, block_n) per_block = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) - return (blocked * per_block).to(output_dtype).reshape(original_shape) + out = (blocked * per_block).reshape(*original_shape[:-2], rows + pad_m, cols + pad_n) + return out[..., :rows, :cols].to(output_dtype) def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: if model is None or full_layer_name is None: diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 22d1b374e852..6a488ba4cb3b 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -356,7 +356,7 @@ def finegrained_linear( # arch/input combos DeepGEMM has no kernel for, and raises its own error if it cannot. logger.warning_once( f"DeepGEMM declined this call, falling back to Triton. Reason: {e} " - "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for FP8 linear entirely." + "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for finegrained linear entirely." ) kernel = load_finegrained_kernel() @@ -823,7 +823,7 @@ def disable_deepgemm_on_multi_device(model: nn.Module) -> None: launches it against the wrong context and produces garbage. (Build-time fix: compile DeepGEMM with `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until our wheel picks that up we avoid single-process multi-device.) Setting `_deepgemm_disabled` routes both the - linear and experts paths through Triton/grouped_mm. A model that fits on one device keeps + linear and experts paths through our triton kernels. A model that fits on one device keeps DeepGEMM even with other GPUs visible; TP/EP put one device per process, so this is a no-op there.""" quantized_modules = [m for m in model.modules() if isinstance(m, _FineGrainedModule)] @@ -837,8 +837,8 @@ def disable_deepgemm_on_multi_device(model: nn.Module) -> None: for m in quantized_modules: m._deepgemm_disabled = True logger.warning_once( - "This FP8 model spans multiple CUDA devices in one process; routing its FP8 linear and experts " - "layers through Triton/grouped_mm instead of DeepGEMM (DeepGEMM's cached kernels are bound to a " + "This finegrained quantized model spans multiple CUDA devices in one process; routing its linear and experts " + "layers through our triton kernels instead of DeepGEMM (DeepGEMM's cached kernels are bound to a " "single CUDA context and corrupt across devices). Run tensor/expert parallel (one device per " "process) to use the faster DeepGEMM path." ) @@ -938,10 +938,8 @@ def storage_for(module_name: str) -> dict: **storage_for(module_name), ) if new_module is not None: - # The kernels take raw pointers, so every operand must be a plain tensor. This is - # what tells the TP layer to hand local shards rather than DTensors - # (`should_use_local_tensors`); without it a dense linear sees DTensors whenever - # grad is enabled, since its inference fast path only covers `not is_grad_enabled`. + # the kernels take raw pointers, so the TP layer must hand local shards and not + # DTensors, whatever `should_use_local_tensors` would decide from grad mode alone new_module._hf_quantized_needs_local_tp = True model.set_submodule(module_name, new_module) has_been_replaced = True diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index e02d51b1a8a4..eba1470b60ef 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1701,9 +1701,6 @@ def subtree_patterns(globs: list[str]) -> list[str]: return patterns -_CATCH_ALL = [".*"] # the targets of a group that takes whatever no other group claimed - - @dataclass class FineGrainedGroup: """One set of modules and the format they are quantized in. @@ -1714,7 +1711,7 @@ class FineGrainedGroup: """ quant_method: str - targets: list[str] = field(default_factory=lambda: list(_CATCH_ALL)) + targets: list[str] = field(default_factory=lambda: [".*"]) # default: whatever no other group claimed activation_format: str | None = None activation_scheme: str = "dynamic" weight_block_size: tuple[int, int] | None = None @@ -1772,14 +1769,14 @@ def group_from_config_groups(spec: dict) -> FineGrainedGroup | None: if target.startswith("re:"): targets.append(target[3:]) elif target == "Linear": # every linear — the catch-all, in their spelling - targets.extend(_CATCH_ALL) + targets.append(".*") else: # a class name we cannot resolve to module paths here return None activations = spec.get("input_activations") or {} act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) return FineGrainedGroup( quant_method=quant_method, - targets=targets or list(_CATCH_ALL), + targets=targets or [".*"], activation_format=formats.get(act_key) if activations else None, activation_scheme="dynamic" if activations.get("dynamic", True) else "static", ) @@ -1925,14 +1922,14 @@ def group_for(self, module_name: str) -> FineGrainedGroup: that order. Two targeted groups claiming the same module is ambiguous, and says so. """ claimed = [ - name for name, group in self.groups.items() if group.targets != _CATCH_ALL and group.matches(module_name) + name for name, group in self.groups.items() if group.targets != [".*"] and group.matches(module_name) ] if len(claimed) > 1: raise ValueError(f"{module_name!r} is claimed by more than one group: {claimed}") if claimed: return self.groups[claimed[0]] for group in self.groups.values(): - if group.targets == _CATCH_ALL: + if group.targets == [".*"]: return group raise KeyError(f"no group covers {module_name!r}; groups: {list(self.groups)}") diff --git a/tests/integrations/finegrained/__init__.py b/tests/integrations/finegrained/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py new file mode 100644 index 000000000000..a84d14fb9b23 --- /dev/null +++ b/tests/integrations/finegrained/test_conversions.py @@ -0,0 +1,694 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The `finegrained` conversion ops: what each one emits from a checkpoint layout, and how the +loader chains them. Pure torch, so no kernel is loaded and no GPU is needed.""" + +import re +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +import transformers.integrations.finegrained.core as fg +from transformers import PreTrainedConfig, PreTrainedModel +from transformers.core_model_loading import ( + Chunk, + PermuteForRope, + WeightConverter, + convert_and_load_state_dict_in_model, +) +from transformers.integrations.finegrained import FineGrainedExperts +from transformers.modeling_utils import LoadStateDictConfig +from transformers.testing_utils import require_torch + +from .test_core import _Cfg, _quantizer_for, _targets + + +@require_torch +class FineGrainedMxfp4ConverterTest(unittest.TestCase): + """The gpt-oss checkpoint layout through the real ConversionOps: blocks (E, N, K/32, 16) + uint8 low-nibble-first E2M1 + biased-127 exponent scales, gate|up rows INTERLEAVED — + deserialized output must dequantize identically to the reference LUT unpack.""" + + FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + + def _reference_dequant(self, blocks, scales): + # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) + lut = torch.tensor(self.FP4_VALUES) + lo = lut[(blocks & 0xF).long()] + hi = lut[(blocks >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) + exp = (scales.long() - 127).unsqueeze(-1) + return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) + + def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): + from transformers.integrations.finegrained.conversions import ( + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + ) + + torch.manual_seed(0) + E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows + blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) + scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) + + weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ + "gate_up_proj" + ] + self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) + back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] + self.assertTrue(torch.equal(back, blocks)) + + # the scales are the container op's job: a module holding e8m0 receives the exponent bytes as is + module = fg.FineGrainedExperts.__new__(fg.FineGrainedExperts) + torch.nn.Module.__init__(module) + module.gate_up_proj_scale_inv = torch.nn.Parameter( + torch.zeros(E, N, K // 32, dtype=torch.float8_e8m0fnu), requires_grad=False + ) + model = torch.nn.Module() + model.experts = module + scale_inv = FineGrainedScaleContainer().convert( + {"gate_up_proj_scales$": scales}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + target_patterns=["gate_up_proj_scale_inv"], + )["gate_up_proj_scale_inv"] + self.assertEqual(scale_inv.dtype, torch.float8_e8m0fnu) + + ref = self._reference_dequant(blocks, scales) # rows pass through interleaved + # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) + lut = torch.tensor(self.FP4_VALUES) + w_u8 = weight.view(torch.uint8) + lo = lut[(w_u8 & 0xF).long()] + hi = lut[(w_u8 >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) + exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) + got = vals * torch.pow(torch.tensor(2.0), exp) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + +@require_torch +class FineGrainedOnTheFlyQuantizeTest(unittest.TestCase): + """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run + the kernels' own quantizers, which need a GPU), emitted in the module's layout.""" + + def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): + """`dequantize=True` on a GPT-OSS MXFP4 checkpoint used to be the mxfp4 integration's job. + The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation + the unquantized experts hold — so it is compared against that integration directly.""" + from transformers.core_model_loading import Transpose + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks + from transformers.integrations.mxfp4 import convert_moe_packed_tensors + + torch.manual_seed(0) + num_experts, rows, hidden = 2, 8, 64 + blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) + scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) + reference = convert_moe_packed_tensors(blocks.clone(), scales.clone(), dtype=torch.bfloat16) + + sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] + tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} + for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): + tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) + self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain + torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) + + def test_a_static_scheme_gets_its_activation_scales_written(self): + """A calibration-fed slot the checkpoint does not supply is still WRITTEN, as the identity. + + The loader materializes a missing key with `torch.empty_like` and `_init_weights` has no + branch for a scale, so a slot this op leaves out reaches the kernels as uninitialized + memory — a zero divides the activations by zero, a negative flips their sign. + """ + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts( + cfg, block_size=(128, 128), weight_format="fp8", activation_scheme="static" + ) + model.proj = fg.FineGrainedLinear( + in_features=256, + out_features=256, + block_size=(128, 128), + weight_format="fp8", + activation_scheme="static", + ) + op = FineGrainedQuantize(hf_quantizer=None) + for key, tensor, slot in ( + ("experts.gate_up_proj", torch.randn(4, 256, 256), "experts.gate_up_proj_activation_scale"), + ("proj.weight", torch.randn(256, 256), "proj.activation_scale"), + ): + with self.subTest(key=key): + out = op.convert({key: tensor}, model=model) + self.assertIn(slot, out, f"{key}: the static activation scale was not written") + held = model.get_parameter(slot) + self.assertEqual((out[slot].shape, out[slot].dtype), (held.shape, held.dtype)) + torch.testing.assert_close(out[slot], torch.ones_like(out[slot]), rtol=0, atol=0) + + def test_an_expert_bias_is_passed_through_not_quantized(self): + """A GPT-OSS expert bias is `(E, rows)` — 2-D, like a dense weight — so the rank guard + alone lets it through. It has no scale slot, and emitting one gives the loader a + `_bias_scale_inv` no module holds.""" + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", has_bias=True) + bias = torch.randn(4, 256) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.down_proj_bias": bias}, model=model) + self.assertEqual(list(out), ["experts.down_proj_bias"]) + self.assertIs(out["experts.down_proj_bias"], bias) + + def test_a_partial_trailing_block_round_trips_against_the_configured_block(self): + """DeepSeek-V3 ships `kv_a_proj_with_mqa` as `(576, 7168)` against a 128x128 block, so the + grid rounds UP to `(5, 56)` and the short block quantizes on its own values. + + The grid cannot be inverted to recover the block: 576 // 5 is 116, and `(192, 256)` divides + evenly at 96 while the real block is still 128. A dequantize that guesses then reads the + wrong rows per scale, which only lifts the error to 4.4e-2 against a 2.6e-2 floor — hence + the tight bound, since a loose one passes the broken arm too.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + quantizer = SimpleNamespace( + quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float") + ) + for rows, cols in ((576, 256), (192, 256), (256, 256)): + with self.subTest(shape=(rows, cols)): + original = torch.randn(rows, cols) + weight, scale = FineGrainedQuantize._quantize_block_fp8(original, (128, 128), ue8m0=False) + self.assertEqual(weight.shape, (rows, cols)) + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(scale.shape, (-(-rows // 128), -(-cols // 128))) + + recovered = FineGrainedDequantize(quantizer)._dequantize_one( + weight, scale.float(), output_dtype=torch.float32 + ) + rel = ((recovered - original).norm() / original.norm()).item() + self.assertLess(rel, 3e-2, f"{(rows, cols)}: {rel:.4f}") + + def test_block_fp8_round_trips_within_its_floor(self): + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + for scale_fmt in ("float", "ue8m0"): + with self.subTest(scale_fmt=scale_fmt), mock.patch.object(fg, "is_sm100", return_value=False): + experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", scale_fmt=scale_fmt) + model = torch.nn.Module() + model.experts = experts + original = torch.randn(4, 256, 256) * 0.02 + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] + self.assertEqual( + (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) + ) + self.assertEqual( + (scale.dtype, scale.shape), + (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), + ) + deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) + rel = ((deq - original).norm() / original.norm()).item() + self.assertLess(rel, 4e-2, f"{scale_fmt}: {rel:.3f}") + + def test_ue8m0_quantizes_against_the_rounded_scale(self): + # ``scale_fmt="ue8m0"`` rounds weight_scale_inv to a power of two, and the weight has to be + # divided by that rounded value or the dequantized block is off by up to an octave + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + quantizer = SimpleNamespace( + quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="ue8m0") + ) + torch.manual_seed(0) + weight = torch.randn(128, 128, dtype=torch.float32) + + # no model to ask, so the op takes the block and the scale format from the config + quantized = FineGrainedQuantize(quantizer)._quantize_one("layer.weight", weight, None) + recovered = FineGrainedDequantize(quantizer)._dequantize_one( + quantized["layer.weight"], quantized["layer.weight_scale_inv"].float(), output_dtype=torch.float32 + ) + rel_err = ((recovered - weight).abs().sum() / weight.abs().sum()).item() + self.assertLess(rel_err, 5e-2) # ~2e-2 for a round trip; a scale mismatch inflates it past 0.2 + + def test_float_scale_fmt_stays_bit_identical_to_its_formula(self): + """The inverse scale is `amax / MAX` and the weight is DIVIDED by it. The algebraically + equal `weight * (MAX / amax)` is not bit-equal — it moves 29 of these 100 blocks' scales — + so the divergence only shows across many blocks, on particular fp32 maxima.""" + from transformers.integrations.finegrained.conversions import ( + _FP8_DTYPE, + _FP8_MAX, + _FP8_MIN, + FineGrainedQuantize, + ) + + quantizer = FineGrainedQuantize( + SimpleNamespace(quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float")) + ) + for seed in range(100): + torch.manual_seed(seed) + weight = torch.randn(128, 128, dtype=torch.float32) + out = quantizer._quantize_one("layer.weight", weight, None) + ref_inv = (weight.abs().amax() / _FP8_MAX).to(torch.float32).reshape(1, 1) + ref_q = torch.clamp(weight / ref_inv, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) + self.assertTrue(torch.equal(out["layer.weight"], ref_q), f"float weight diverged at seed {seed}") + self.assertTrue( + torch.equal(out["layer.weight_scale_inv"], ref_inv), f"float scale diverged at seed {seed}" + ) + + +@require_torch +class FineGrainedModeloptConverterTest(unittest.TestCase): + """The ModelOpt (NVFP4) weight conversions. `MergeModulelist` is load-bearing beyond the + merge itself: `core_model_loading` stamps a source with its expert index ONLY when one is + present in the chain, and expert parallelism selects experts by that index. Without it every + rank collects all E globals and the forward asserts on the per-expert count.""" + + def _modelopt_conversions(self, **cfg_kwargs): + from transformers.quantizers.finegrained.nvfp4 import FineGrainedNvfp4HfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) + return FineGrainedNvfp4HfQuantizer(cfg).get_weight_conversions() + + def test_every_declared_target_is_a_slot_the_module_holds(self): + """A converter target that no module holds is a load failure, and the activation global is + exactly such a target under a weight-only run: the experts allocate one only when they + quantize activations. Both runs are checked against a real module's parameter names, since + the declarations are the only thing standing between the two.""" + cfg = _Cfg() + for activation_format, holds_global in ((None, True), ("bf16", False)): + experts = FineGrainedExperts( + cfg, block_size=(4, 4), weight_format="nvfp4", activation_format=activation_format + ) + slots = {name for name, _ in experts.named_parameters()} | set(experts._buffers) + self.assertEqual( + any("input_global_scale" in slot for slot in slots), + holds_global, + f"module with activation_format={activation_format!r} disagrees about the slot", + ) + targets = { + t.split("experts.")[-1].rstrip("$") + for c in self._modelopt_conversions(activation_format=activation_format) + for t in _targets(c) + } + for target in targets: + self.assertIn(target, slots, f"activation_format={activation_format!r} converts onto a missing slot") + + def test_global_scale_converters_carry_an_expert_index(self): + from transformers.core_model_loading import MergeModulelist + + # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked + # tensor per layer has no per-expert sources to stamp — it is sharded like the fused + # expert weight next to it + globals_converters = [ + c + for c in self._modelopt_conversions() + if any("global_scale" in t for t in _targets(c)) + and any(re.search(p, "model.layers.0.mlp.experts.7.gate_proj.weight_scale_2") for p in c.source_patterns) + ] + self.assertTrue(globals_converters, "no global-scale converter found") + for conv in globals_converters: + self.assertTrue( + any(isinstance(op, MergeModulelist) for op in conv.operations), + f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " + "experts and every rank would collect all of them", + ) + + def test_both_modelopt_scale_layouts_have_converters(self): + """modelopt ships the scales either as one tensor per expert per projection (GLM-5.2) or + as one already-stacked tensor per layer (the fused vLLM layout). Both are + converted; the block scale takes the module's layout ops either way.""" + sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} + for suffix in ("weight_scale", "weight_scale_2", "input_scale"): + per_expert = f"model.layers.0.mlp.experts.7.gate_proj.{suffix}" + fused = f"model.layers.0.mlp.experts.gate_up_proj_{suffix}" + self.assertTrue(any(re.search(p, per_expert) for p in sources), f"no per-expert converter for {suffix}") + self.assertTrue(any(re.search(p, fused) for p in sources), f"no fused converter for {suffix}") + # the per-expert patterns are REGEXES: an unescaped dot also matches `_`, so + # `mlp.experts.*.up_proj.weight_scale_2` would claim the fused key as well + claimed = [p for p in sources if re.search(p, fused) and re.search(p, per_expert)] + self.assertEqual(claimed, [], f"one pattern claims BOTH layouts for {suffix}: {claimed}") + + def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): + """Two converters for one key is not additive — the later one WINS and silently drops + whatever ops the first carried. That is how a fused modelopt checkpoint lost its packed + uint8 -> int8 view (a catch-all duplicated the fused converter) and how the per-expert + patterns, which are REGEXES whose unescaped dots also match `_`, claimed the fused + globals and flattened the gate|up pair instead of folding it.""" + import re + + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for quant_kwargs in ( + {"quant_method": "fp8"}, + {"quant_method": "mxfp4"}, + {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + ): + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**quant_kwargs)) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + for key in ( + "model.layers.3.mlp.experts.gate_up_proj", + "model.layers.3.mlp.experts.down_proj", + "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", + "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", + ): + claimed = [ + c + for c in converters + for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) + if re.search(str(pattern), key) + ] + with self.subTest(quant=quant_kwargs["quant_method"], key=key.rsplit(".", 1)[1]): + self.assertLessEqual( + len(claimed), + 1, + f"{len(claimed)} converters claim this key; the last one wins and drops " + f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", + ) + + def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): + """A static checkpoint calibrates one `input_scale` per quantized module, and each key + must reach the slot its module holds, in that module's shape: one value on a dense linear, + one per expert on the stacked experts, and the gate|up pair reduced to one per expert + rather than flattened to 2E. Unrouted, the module keeps its registered default of 1.0 and + every activation quantizes against the wrong scale, silently.""" + import re + + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = _quantizer_for(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + + def claim(key): + return [ + c + for c in converters + for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) + if re.search(str(pattern), key) + ] + + for key, expected in ( + ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), + ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.up_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.down_proj.input_scale", "mlp.experts.down_proj_activation_scale"), + ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), + ): + with self.subTest(key=key.rsplit(".", 2)[-2]): + claimed = claim(key) + self.assertEqual(len(claimed), 1, f"{len(claimed)} converters claim {key}") + # a renaming's target carries a backreference; resolve it against the key + targets = [ + re.sub(str(claimed[0].source_patterns[0]), str(t), key) if "\\1" in str(t) else str(t) + for t in claimed[0]._original_target_patterns + ] + self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") + # a dynamic checkpoint has no such key to route, and a weight-only one ignores it + dynamic = _quantizer_for(FineGrainedConfig(quant_method="fp8")) + dynamic.pre_quantized = True + self.assertEqual(dynamic.get_weight_conversions(), []) + + def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): + """`FineGrainedInputScales` collapses the gate_up to ONE value for the NVFP4 global — the + global is a split of the block scale, so an inflated one is exact. A static scale IS the + quantization scale, with no block level to absorb it, so collapsing it would quantize + every expert against the largest one's range. It stays per expert.""" + import torch + + from transformers.integrations.finegrained.conversions import FineGrainedInputScales + + sources = { + "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), + "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), + } + target = "mlp.experts.gate_up_proj_activation_scale" + scale = FineGrainedInputScales().convert(dict(sources), full_layer_name=target)[target] + torch.testing.assert_close(scale, torch.tensor([2.0, 3.0])) # per expert, over the pair + + def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): + """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair + reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals + where the kernels assert one per expert.""" + import torch + + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + experts = 8 + for wrap in (False, True): + pair = torch.rand(experts, 2) + 1 + src = { + "experts.gate_up_proj_weight_scale_2": [pair] if wrap else pair, + "experts.down_proj_weight_scale_2": torch.rand(experts) + 1, + "experts.down_proj_input_scale": torch.rand(experts) + 1, + } + out = FineGrainedWeightGlobals(None).convert(src, target_patterns=targets, model=None) + with self.subTest(list_wrapped=wrap): + self.assertEqual( + out["experts.gate_up_proj_weight_global_scale"].shape, + torch.Size([experts]), + "the gate|up halves were not folded to one global per expert", + ) + + def test_one_converter_merges_a_layer_s_globals(self): + """modelopt calibrates the gate|up halves separately. One converter owns every global of + a layer, because merging them to the one-per-expert the kernels take moves the up half's + onto the down projection: SwiGLU is linear in the up half, so the stack keeps the gate's + global and the down's weight global scales the expert output back. The down's calibrated + input scale moves the other way, keeping the requantized intermediate on the range the + checkpoint calibrated.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) + down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) + targets = [ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + "mlp.experts.down_proj_input_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "mlp.experts.*.up_proj.weight_scale_2": up, + "mlp.experts.*.gate_proj.weight_scale_2": gate, + "mlp.experts.*.down_proj.weight_scale_2": down, + "mlp.experts.*.down_proj.input_scale": down_input, + }, + target_patterns=targets, + ) + ratio = up / gate + torch.testing.assert_close(out[targets[0]], gate) + torch.testing.assert_close(out[targets[1]], down * ratio) + torch.testing.assert_close(out[targets[2]], down_input / ratio) + + def test_globals_with_one_calibrated_projection_pass_through(self): + """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a + single global) has nothing to merge, so every global reaches its module unchanged.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0], [2.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([3.0, 4.0]), + "experts.down_proj_input_scale": torch.tensor([0.5, 0.25]), + }, + target_patterns=targets, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([3.0, 4.0])) + torch.testing.assert_close(out[targets[2]], torch.tensor([0.5, 0.25])) + + def test_the_fused_layout_ships_both_halves_in_one_tensor(self): + """A stacked `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two + keys, and merges the same way.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 3.0], [2.0, 4.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([5.0 * 3.0, 6.0 * 2.0])) + + def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_down(self): + """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated + `input_scale` reduces to one value; the down's rows are per expert, so its stays per + expert — the requant epilogue normalizes each row by its own expert's value.""" + from transformers.integrations.finegrained.conversions import FineGrainedInputScales + + up = FineGrainedInputScales().convert( + { + "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), + "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), + }, + full_layer_name="mlp.experts.gate_up_proj_input_global_scale", + )["mlp.experts.gate_up_proj_input_global_scale"] + torch.testing.assert_close(up, torch.tensor([3.0])) + down = FineGrainedInputScales().convert( + {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, + full_layer_name="mlp.experts.down_proj_input_global_scale", + )["mlp.experts.down_proj_input_global_scale"] + torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) + + +@require_torch +class FineGrainedLoaderChainTest(unittest.TestCase): + """A core conversion chain and an on-the-fly quantization in the same load: the permute has to + run on the full-precision tensor, and only the finegrained module of the three gets a scale.""" + + def test_a_chunked_rope_permute_quantizes_only_the_finegrained_target(self): + from transformers.integrations.finegrained import FineGrainedLinear + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + n_heads = 2 + head_dim = 4 + in_dim = 4 + out_dim = n_heads * head_dim + block_size = (4, 4) + + class RopeProjector(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(out_dim, in_dim)) + + class RopeSelfAttn(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = FineGrainedLinear(in_dim, out_dim, block_size=block_size, weight_format="fp8") + self.k_proj = RopeProjector() + self.v_proj = RopeProjector() + + class RopeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = RopeSelfAttn() + + class RopeModel(PreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + self.layers = torch.nn.ModuleList([RopeLayer()]) + self.post_init() + + def _init_weights(self, module): + pass # every weight here comes from the state dict, and `normal_` has no fp8 kernel + + config = PreTrainedConfig() + config.num_attention_heads = n_heads + model = RopeModel(config) + + raw_q = torch.tensor( + [ + [1.0, -1.0, 1.0, -1.0], + [0.5, -0.5, 0.5, -0.5], + [-1.0, 1.0, -1.0, 1.0], + [-0.5, 0.5, -0.5, 0.5], + [1.0, 1.0, -1.0, -1.0], + [0.5, 0.5, -0.5, -0.5], + [-1.0, -1.0, 1.0, 1.0], + [-0.5, -0.5, 0.5, 0.5], + ], + dtype=torch.float32, + ) + raw_k = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + raw_v = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + 100.0 + raw_qkv = torch.cat([raw_q, raw_k, raw_v], dim=0) + state_dict = {"layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} + + quantizer_cls = type( + "FineGrainedHfQuantizer", + (), + { + "__init__": lambda self, bs=block_size: setattr( + self, "quantization_config", SimpleNamespace(weight_block_size=bs, scale_fmt="float") + ), + "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), + "get_quantize_ops": lambda self: FineGrainedQuantize(self), + "pre_quantized": False, + }, + ) + quantizer = quantizer_cls() + + weight_mapping = [ + WeightConverter( + "self_attn.qkv_proj.weight", + [ + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + ], + operations=[Chunk(dim=0), PermuteForRope(permute_layer_names=["q_proj", "k_proj"])], + ) + ] + load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) + loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config) + + self.assertEqual(loading_info.missing_keys, set()) + self.assertEqual(loading_info.unexpected_keys, set()) + self.assertEqual(loading_info.mismatched_keys, set()) + self.assertEqual(loading_info.conversion_errors, {}) + + permute_op = PermuteForRope(permute_layer_names=["q_proj", "k_proj"]) + permute_op.config = model.config + expected_q = permute_op._apply(raw_q) + expected_k = permute_op._apply(raw_k) + expected_v = raw_v # V carries no RoPE, so the permute names only q and k + + model_state = model.state_dict() + self.assertFalse(torch.allclose(raw_k, expected_k)) + torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) + + q_weight_key = "layers.0.self_attn.q_proj.weight" + scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" + self.assertIn(scale_key, model_state) + self.assertEqual(model_state[q_weight_key].dtype, torch.float8_e4m3fn) + self.assertEqual(model_state[q_weight_key].shape, torch.Size((out_dim, in_dim))) + self.assertEqual(model_state[scale_key].dtype, torch.float32) + self.assertEqual( + model_state[scale_key].shape, + torch.Size((out_dim // block_size[0], in_dim // block_size[1])), + ) + + dequantized_q = FineGrainedDequantize(None)._dequantize_one( + model_state[q_weight_key], model_state[scale_key], output_dtype=torch.float32 + ) + torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py new file mode 100644 index 000000000000..94cb85bfb724 --- /dev/null +++ b/tests/integrations/finegrained/test_core.py @@ -0,0 +1,1402 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Marshalling tests for the multi-format ``finegrained`` integration: the kernel bundle is +mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` +(the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" + +import re +import unittest +import warnings +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from unittest import mock + +import torch + +import transformers.integrations.deepgemm as deepgemm +import transformers.integrations.finegrained.core as fg +from transformers.integrations.finegrained import ( + FineGrainedExperts, + FineGrainedLinear, + finegrained_linear, + load_finegrained_kernel, +) +from transformers.testing_utils import require_torch + + +@dataclass +class _Call: + args: tuple + kwargs: dict + + +class _Recorder: + """Records every call and returns a correctly-shaped zero tensor for matmuls.""" + + def __init__(self): + self.calls = {} + + def _op(self, name): + def op(*args, **kwargs): + self.calls.setdefault(name, []).append(_Call(args, kwargs)) + if name == "swizzle_mx_scales": + return args[0] + if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) + hidden, down = args[0], kwargs["down_proj"] + return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) + a, b = args[0], args[1] + n = b.shape[-2] + gather = kwargs.get("gather_idx") + rows = gather.shape[0] if gather is not None else a.shape[0] + return a.new_zeros(rows, n, dtype=torch.bfloat16) + + return op + + +def _fake_bundle(): + rec = _Recorder() + kernel = mock.Mock() + for name in ( + "matmul_2d", + "matmul_batched", + "matmul_grouped", + "moe_fused_batched", + "moe_fused_grouped", + "swizzle_mx_scales", + "unswizzle_mx_scales", + "mxfp8_act_quant", + "mxfp4_act_quant", + "nvfp4_act_quant", + ): + setattr(kernel, name, rec._op(name)) + kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") + kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") + return kernel, rec + + +class _Cfg: + hidden_size = 64 + num_local_experts = 4 + intermediate_size = 32 + hidden_act = "silu" + _experts_implementation = None + + +def _loaded(kernel): + return ( + mock.patch.object(fg, "_FINEGRAINED", None), + mock.patch.object(fg, "is_kernels_available", return_value=True), + mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), + # a locally importable checkout (FINEGRAINED_KERNELS_PATH / installed package) + # takes precedence over the hub loader; tests must stay hermetic to the fake bundle + mock.patch.object(fg, "_import_local_finegrained", return_value=None), + ) + + +def _quantizer_for(config): + """The quantizer a `from_pretrained` would pick for this config. Since the arms were split + per scheme, constructing the base directly only exercises what every scheme shares.""" + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING + + method = getattr(config.quant_method, "value", config.quant_method) + return AUTO_QUANTIZER_MAPPING[method](config) + + +def _targets(conv): + t = conv.target_patterns + return t if isinstance(t, (list, tuple)) else [t] + + +@require_torch +class FineGrainedLoaderTest(unittest.TestCase): + def test_loader_requires_every_symbol(self): + kernel, _ = _fake_bundle() + del kernel.matmul_grouped + # Mock auto-creates attributes; force the miss + kernel.matmul_grouped = None + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, self.assertRaises(ImportError) as ctx: + load_finegrained_kernel() + self.assertIn("matmul_grouped", str(ctx.exception)) + + def test_loader_is_compile_safe_cold(self): + """Cold path: the compiled call is first to load, so the opaque loader node runs its full + body under compile and must return None, never the bundle (`torch.* op returned + non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" + kernel, _ = _fake_bundle() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_is_compile_safe_when_warm(self): + """Warm path, which is the production order: eager warm-up, then compile. The loader hits + its short-circuit at trace time — the branch that must also return None.""" + kernel, _ = _fake_bundle() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + load_finegrained_kernel() + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_binds_all_symbols(self): + kernel, _ = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + bundle = load_finegrained_kernel() + self.assertIs(bundle.matmul_2d, kernel.matmul_2d) + self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) + + +@require_torch +class FineGrainedLinearMarshallingTest(unittest.TestCase): + def _run(self, **linear_kwargs): + kernel, rec = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + x = torch.randn(3, 5, 64, dtype=torch.bfloat16) + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + ws = torch.randn(1, 1, dtype=torch.float32) + out = finegrained_linear(x, w, ws, **linear_kwargs) + return out, rec.calls["matmul_2d"][-1] + + def test_scale_rides_positionally_and_block_size_dies(self): + out, call = self._run(block_size=[128, 128]) + # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here + self.assertEqual(call.args[0].shape, (15, 64)) + self.assertIsNone(call.args[2]) + self.assertEqual(call.args[3].dtype, torch.float32) + self.assertNotIn("block_size", call.kwargs) + self.assertEqual(out.shape, (3, 5, 32)) + + def test_static_activation_scale_is_As(self): + scale = torch.tensor(0.5) + _, call = self._run(activation_scale=scale) + self.assertIs(call.args[2], scale) + + def test_nvfp4_global_and_weight_only_format(self): + g = torch.tensor(2.0) + _, call = self._run(weight_global_scale=g, activation_format="bf16") + self.assertIs(call.kwargs["b_global_scale"], g) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_module_forward_threads_everything(self): + kernel, rec = _fake_bundle() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m = FineGrainedLinear( + 64, + 32, + block_size=None, + weight_format="nvfp4", + activation_format="bf16", + ) + m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) + out = m(torch.randn(2, 64, dtype=torch.bfloat16)) + call = rec.calls["matmul_2d"][-1] + self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) + self.assertEqual(call.kwargs["activation_format"], "bf16") + self.assertEqual(out.shape, (2, 32)) + + +@require_torch +class FineGrainedGroupedLinearSwapTest(unittest.TestCase): + def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): + """`replace_with_finegrained_layer` picks `FineGrainedGroupedLinear` for a block-diagonal + linear — a plain one would collapse the groups into a single giant matmul and return the + wrong output dim. The branch keys on `n_groups`, the attribute the swap consumes: a class + NAME substring would miss a grouped linear named anything else, and would match (then + `AttributeError` on) one that carries its group count under another name. + """ + from types import SimpleNamespace + + import torch.nn as nn + + from transformers.integrations.finegrained import FineGrainedGroupedLinear, replace_with_finegrained_layer + from transformers.utils.quantization_config import FineGrainedConfig + + class OddlyNamedBlockDiagonal(nn.Linear): + def __init__(self): + super().__init__(8, 16, bias=False) + self.n_groups = 2 + + model = nn.Module() + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + model.grouped = OddlyNamedBlockDiagonal() + model.plain = nn.Linear(8, 16, bias=False) + + with torch.device("meta"): + replace_with_finegrained_layer(model, None, FineGrainedConfig(quant_method="fp8")) + self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) + self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) + + +class FineGrainedEmbeddingTest(unittest.TestCase): + """A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, + rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`.""" + + def test_lookup_rescales_the_gathered_rows(self): + torch.manual_seed(0) + table = torch.randn(64, 16) + scale = table.abs().max() / 448.0 + embedding = fg.FineGrainedEmbedding(64, 16) + embedding.weight = torch.nn.Parameter((table / scale).to(torch.float8_e4m3fn), requires_grad=False) + embedding.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16).reshape(1), requires_grad=False) + ids = torch.tensor([[3, 7], [60, 0]]) + out = embedding(ids) + self.assertEqual(out.dtype, torch.bfloat16) + expected = (table / scale).to(torch.float8_e4m3fn)[ids].to(torch.bfloat16) * scale.to(torch.bfloat16) + torch.testing.assert_close(out, expected) + + def test_replacement_follows_the_patterns_and_the_skip_list(self): + model = torch.nn.Module() + model.ngram_embedding = torch.nn.Embedding(8, 4) + model.embed_tokens = torch.nn.Embedding(8, 4) + model.other = torch.nn.Embedding(8, 4) + fg.replace_with_finegrained_embedding(model, ["ngram_embedding", "other"], modules_to_not_convert=["other"]) + self.assertIsInstance(model.ngram_embedding, fg.FineGrainedEmbedding) + self.assertTrue(model.ngram_embedding._hf_quantized_needs_local_tp) + self.assertIs(type(model.embed_tokens), torch.nn.Embedding) + self.assertIs(type(model.other), torch.nn.Embedding) + + +@require_torch +class FineGrainedExpertsMarshallingTest(unittest.TestCase): + def _experts(self, **kw): + cfg = _Cfg() + weight_format = kw.pop("weight_format", "fp8") + m = FineGrainedExperts(cfg, block_size=(4, 4), weight_format=weight_format, **kw) + for name, p in m.named_parameters(): + if p is None: + continue + if p.dtype == torch.float8_e4m3fn: + p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) + elif p.dtype.is_floating_point: + p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) + return m + + def _route(self, tokens=6): + hs = torch.randn(tokens, 64, dtype=torch.bfloat16) + idx = torch.randint(0, 4, (tokens, 2)) + wts = torch.rand(tokens, 2) + return hs, idx, wts + + def test_eager_linear_threads_bias_global_and_activation_format(self): + """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: + the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" + kernel, rec = _fake_bundle() + m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") + p1, p2, p3, p4 = _loaded(kernel) + linear = mock.Mock(wraps=fg.finegrained_linear) + with ( + p1, + p2, + p3, + p4, + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), + mock.patch.object(fg, "finegrained_linear", linear), + ): + m(*self._route()) + self.assertTrue(linear.call_args_list) + for call in linear.call_args_list: + self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) + # this expert's NVFP4 weight global: one value for the down, and (1, 2) for the + # gate|up stack, whose halves modelopt calibrates separately + self.assertIn(tuple(call.kwargs["weight_global_scale"].shape), {(), (1, 2)}) + self.assertEqual(call.kwargs["activation_format"], "bf16") + # weight-only: the activations stay bf16, so no activation global is quantized against + self.assertIsNone(call.kwargs["input_global_scale"]) + for call in rec.calls["matmul_2d"]: + self.assertIsNotNone(call.kwargs["b_global_scale"]) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_batched_marshalling(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + hidden, top_k_index, top_k_weights = call.args + self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers + self.assertEqual(top_k_index.shape, (6, 2)) + self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) + # one scale tensor per projection: the affine Parameter, or the swizzled cache when the + # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) + self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) + self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) + self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name + self.assertIs(call.kwargs["gate"], True) + self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format + self.assertEqual(out.shape, (6, 64)) + + def test_grouped_marshalling(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIs(call.kwargs["down_proj"], m.down_proj) + self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) + self.assertNotIn("moe_fused_batched", rec.calls) + + def test_activation_format_passes_through(self): + kernel, rec = _fake_bundle() + for activation_format in (None, "bf16", "mxfp8"): + m = self._experts(has_gate=True, activation_format=activation_format) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) + + def test_nvfp4_experts_thread_per_expert_globals(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, weight_format="nvfp4") + # the format table must resolve the ATTRIBUTE the forwards gate on, not just the + # param allocation — a None here silently drops the global at every forward + self.assertIsNotNone(m.global_scale_dtype) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) + self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) + # one weight global per expert on both projections: a gate|up stack calibrated per half + # is merged to that at load, so nothing downstream carries the pair + self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) + # the calibrated activation globals (the checkpoint's `input_scale`): one value for the + # gate_up, whose rows are quantized once before routing, and one per expert for the down, + # whose rows the gate_up epilogue requantizes per expert + self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) + self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) + self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) + self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) + + def test_weight_only_experts_hold_no_activation_global(self): + """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation + quant for a calibrated global to normalize: the module never allocates one, and the ops + would refuse a two-level path without it.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIsNone(m.gate_up_proj_input_global_scale) + self.assertIsNone(m.down_proj_input_global_scale) + self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) + self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) + # the WEIGHT globals are the format's own second level and stay + self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) + + def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): + """A model whose experts norm the down output before the routing weights: + the swap carries the norm over, the kernel chain runs it on the routed rows, and the + eager loop applies it per expert application — the same place the reference forwards do. + A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) + + class _Recording(torch.nn.Module): + def __init__(self): + super().__init__() + self.rows = 0 + + def forward(self, x): + self.rows += x.shape[0] + return x + + m.post_expert_norm = _Recording() + with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m(*self._route()) + self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") + + def test_a_post_expert_norm_needs_nothing_but_the_norm(self): + """A class that declares a post-expert norm only has to HOLD it: every backend calls + `post_expert_norm` directly, so there is no per-class hook to define alongside it.""" + from transformers.integrations.moe import use_experts_implementation + + @use_experts_implementation(has_post_expert_norm=True) + class _JustTheNorm(torch.nn.Module): + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states + + self.assertFalse(hasattr(_JustTheNorm, "_apply_post_norm")) + + def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): + """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels + when they implement it, with the weight and epsilon the fused form needs, so the chain + folds it into its reduce instead of calling the module.""" + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True + p1, p2, p3, p4 = _loaded(kernel) + for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): + m.post_expert_norm_name = name + with p1, p2, p3, p4: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + call = rec.calls["moe_fused_grouped"][-1] + if fused: + self.assertEqual(call.kwargs["post_expert_norm"], name) + self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) + self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) + else: + self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) + self.assertIsNone(call.kwargs["post_expert_norm_weight"]) + + def test_biases_ride_the_kernel_chain(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, has_bias=True) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertTrue(torch.isfinite(out).all()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels add gate_up's bias before the gated split and down's before the + # routing-weight multiply; nothing is added host-side + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) + + def test_static_activation_scheme_reaches_the_fused_chain(self): + """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a + MoE calibrates each expert separately, so the module holds one per expert. Both fused + chains take them per projection and the kernels apply each tile's own expert scale in + register — a host pre-quant could not, since top-k routes one row to several experts + whose scales differ. A dynamic module holds no such slot and passes `None`.""" + static, dynamic = ( + FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) + for scheme in ("static", "dynamic") + ) + held = static.gate_up_proj_activation_scale + self.assertIsNotNone(held, "the calibrated scale the chain needs") + self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") + # a dynamic module holds the slot empty, like every other optional parameter + self.assertIsNone(dynamic.gate_up_proj_activation_scale) + + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = () + for module, calibrated in ((static, True), (dynamic, False)): + operands = fg._moe_operands(kernel, module) + for projection in ("gate_up_proj", "down_proj"): + with self.subTest(scheme=module.activation_scheme, projection=projection): + scale = operands[f"{projection}_activation_scale"] + self.assertEqual(scale is not None, calibrated) + if calibrated: + self.assertEqual(scale.numel(), _Cfg.num_local_experts) + + def test_unfusable_act_fn_is_passed_as_the_module_glu(self): + kernel, rec = _fake_bundle() + m = self._experts(has_gate=True, has_bias=True) + m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels run the module's own GLU on the host between the two GEMMs — a new + # activation never waits for a kernel release; the bias still rides the GEMM + self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + + +@require_torch +class FineGrainedScaleDtypeTest(unittest.TestCase): + """A scale's dtype is the format's, never the ambient default. + + `from_pretrained` sets the default dtype to the checkpoint's for the duration of model + construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it + as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches + that, whatever the format decides each scale should be. + """ + + def _dtypes(self, build, default): + previous = torch.get_default_dtype() + torch.set_default_dtype(default) + try: + with torch.device("meta"): + module = build() + return {name: p.dtype for name, p in module.named_parameters() if p is not None} + finally: + torch.set_default_dtype(previous) + + def _assert_default_dtype_invariant(self, build): + under_fp32 = self._dtypes(build, torch.float32) + under_bf16 = self._dtypes(build, torch.bfloat16) + self.assertTrue(under_fp32) + # the weight and bias DO follow the model dtype; every scale beside them must not + scales = {name for name in under_fp32 if "scale" in name} + self.assertTrue(scales) + for name in sorted(scales): + self.assertEqual(under_bf16[name], under_fp32[name], f"{name} followed the default dtype") + + def test_linear_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedLinear( + in_features=64, out_features=32, block_size=(4, 4), weight_format=f, activation_scheme=s + ) + ) + + def test_experts_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedExperts( + _Cfg(), block_size=(4, 4), weight_format=f, activation_scheme=s + ) + ) + + +class FineGrainedFusedNormGateTest(unittest.TestCase): + """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. + + Fusing runs the norm inside the launch that produced the rows, so an all-reduce that has to + land between the two has nowhere to go — and under intra-expert TP those rows are a partial + sum over the sharded intermediate, so the fused norm would normalize a fraction of the value. + The wrapped module does the reduce in its own forward, so the fallback is the correct path. + """ + + class _Norm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4)) + self.eps = 1e-6 + + def _operands(self, *, input_reduce: bool): + norm = self._Norm() + if input_reduce: + norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks + + module = mock.Mock() + module.has_gate = True + module.act_fn_name = "silu" + module.swiglu_alpha = None + module.has_post_expert_norm = True + module.post_expert_norm = norm + module.post_expert_norm_name = "input_scaled_rms_norm" + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) + return fg._moe_operands(kernel, module), module + + def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): + operands, _ = self._operands(input_reduce=False) + self.assertEqual(operands["post_expert_norm"], "input_scaled_rms_norm") + self.assertIsNotNone(operands["post_expert_norm_weight"], "the fused form needs the norm's weight") + + def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): + operands, module = self._operands(input_reduce=True) + self.assertIs( + operands["post_expert_norm"], + module.post_expert_norm, + "fused past a collective: the kernel would normalize this rank's partial sum", + ) + self.assertIsNone(operands["post_expert_norm_weight"]) + + +@require_torch +class FrozenFp8ShimTest(unittest.TestCase): + def test_frozen_module_warns_and_is_self_contained(self): + import importlib + + import transformers.integrations.finegrained_fp8 as frozen + from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + from transformers.utils.quantization_config import FineGrainedFP8Config + + # the whole FILE is deprecated, so importing anything out of it is what warns. Safe to + # do at module scope only because `_LazyModule` never executes these unless asked: a + # plain `import transformers` does not reach them. + with warnings.catch_warnings(record=True) as raised: + warnings.simplefilter("always") + importlib.reload(frozen) + self.assertTrue( + [w for w in raised if issubclass(w.category, DeprecationWarning) and "frozen" in str(w.message)], + [str(w.message) for w in raised], + ) + # the quantizer is still constructible — frozen means no new recipes, not removed + FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) + # distinct machinery: the frozen classes are not the live ones + self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) + + +@require_torch +class FineGrainedValidateEnvironmentTest(unittest.TestCase): + """`validate_environment` on the LIVE quantizer. The only tests that covered a + `validate_environment` sat on `FineGrainedFP8HfQuantizer`, which no `from_pretrained` can + reach any more (`AUTO_QUANTIZER_MAPPING["fp8"]` is `FineGrainedHfQuantizer`).""" + + @contextmanager + def _no_accelerator(self): + with ExitStack() as stack: + stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) + stack.enter_context( + mock.patch("transformers.quantizers.finegrained.base.is_torch_xpu_available", return_value=False) + ) + yield + + def _quantizer(self, pre_quantized, **cfg_kwargs): + # through the registry, so the test sees the class a `from_pretrained` would pick + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING + from transformers.utils.quantization_config import FineGrainedConfig + + config = FineGrainedConfig(**cfg_kwargs) + method = getattr(config.quant_method, "value", config.quant_method) + quantizer = AUTO_QUANTIZER_MAPPING[method](config) + quantizer.pre_quantized = pre_quantized + return quantizer + + def test_no_accelerator_refuses_to_quantize(self): + # quantizing a bf16 checkpoint needs the hardware; there is nothing to fall back to + quantizer = self._quantizer(pre_quantized=False) + with self._no_accelerator(), self.assertRaises(RuntimeError): + quantizer.validate_environment() + + def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): + quantizer = self._quantizer(pre_quantized=True) + with self._no_accelerator(): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + def test_nvfp4_refuses_to_dequantize(self): + # two-level scales have no slot in the dequantize chain, so the CPU fallback that rescues + # block-FP8 must refuse here rather than hand back a weight scaled by `1 / global` + for pre_quantized, kwargs in ((True, {}), (False, {"dequantize": True})): + with self.subTest(pre_quantized=pre_quantized): + quantizer = self._quantizer(pre_quantized=pre_quantized, quant_method="nvfp4", **kwargs) + with self._no_accelerator(), self.assertRaises(NotImplementedError): + quantizer.validate_environment() + + +class FineGrainedParallelPlanTest(unittest.TestCase): + """`update_tp_plan`: a model's plan as a quantized experts module needs it.""" + + BASE_EP = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise", + } + + def _planned(self, impl=None, raw=None): + from types import SimpleNamespace + + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = FineGrainedHfQuantizer(FineGrainedConfig()) + config = SimpleNamespace(base_model_tp_plan=dict(raw or self.BASE_EP), _experts_implementation=impl) + quantizer.update_tp_plan(config) + return config.base_model_tp_plan + + def test_every_companion_shards_with_its_expert_weight(self): + """The plan matcher keys on the exact parameter name and falls back only to the owning + module, whose entry shards nothing — so a scale with no entry of its own stays whole while + its weight is this rank's expert slice, and every rank but the first reads another rank's + scales.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + plan = self._planned() + style = lambda name: _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) # noqa: E731 + + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_weight_global_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + for name in ("down_proj_scale_inv", "down_proj_bias", "down_proj_input_global_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + # one per-tensor value for the pre-routing hidden states: replicated, so no entry of its own + self.assertEqual(style("gate_up_proj_input_global_scale"), "moe_tp_experts") + self.assertEqual(plan["layers.*.self_attn.q_proj"], "colwise") + + def test_intra_expert_entries_follow_the_projection_axis(self): + """Under intra-expert TP the experts stay whole and the projection's own axis splits, so + the scale grid must follow its weight while the per-expert globals stay replicated. And + `packed_colwise` splits the output axis as two packed halves — the `[gate; up]` STACK — + where a quantized module holds interleaved rows, so the weight's own style changes too.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + base = { + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + plan = self._planned(raw=base) + style = lambda n: _get_parameter_tp_plan(f"layers.3.mlp.experts.{n}", plan) # noqa: E731 + + # the weight keeps the style the MODEL declared — `packed_colwise` is the strided split a + # concatenated `[gate; up]` checkpoint needs, since the interleave into the kernels' row + # order runs on each rank's shard, after sharding. The companions follow it. + self.assertEqual(style("gate_up_proj"), "packed_colwise") + for name in ("gate_up_proj_scale_inv", "gate_up_proj_bias"): + self.assertEqual(style(name), "moe_experts_packed_colwise", name) + # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike + self.assertEqual(style("down_proj"), "rowwise") + self.assertEqual(style("down_proj_scale_inv"), "moe_experts_rowwise") + # TP does not split experts, so anything indexed per expert stays whole + for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): + self.assertEqual(style(name), "moe_tp_experts", name) + + # a natively INTERLEAVED layout declares `colwise`, and its companions split contiguously + inter = dict(base, **{"layers.*.mlp.experts.gate_up_proj": "colwise"}) + plan_i = self._planned(raw=inter) + self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", plan_i), "colwise") + self.assertEqual( + _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), + "moe_experts_colwise", + ) + + def test_dense_projections_that_share_a_name_are_left_alone(self): + """`mlp.up_proj` and `mlp.shared_experts.down_proj` end in the same words as the stacked + expert projections but are plain 2-D linears that colwise/rowwise already shard right. + Rewriting them to an expert-axis style sharded a dim they do not have — a CUDA illegal + memory access in the dense FP8 matmul, seen on DeepSeek-V3 under TP.""" + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + + plan = self._planned( + raw={ + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + } + ) + for dense, expected in ( + ("layers.*.mlp.up_proj", "colwise"), + ("layers.*.mlp.down_proj", "rowwise"), + ("layers.*.mlp.shared_experts.up_proj", "colwise"), + ("layers.*.mlp.shared_experts.down_proj", "rowwise"), + ): + self.assertEqual(plan[dense], expected, dense) + # and no scale entry was invented for a dense linear, whose scale `rowwise`/`colwise` + # already reaches through the module fallback + strays = [k for k in plan if k.endswith("_scale_inv") and ".experts." not in k] + self.assertEqual(strays, [], f"invented dense scale entries: {strays}") + # the stacked experts still get theirs + self.assertEqual( + _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), + "moe_experts_packed_colwise", + ) + + def test_a_multimodal_models_expert_plans_are_reached(self): + """A multimodal model keeps its experts' plans on a SUB-config; the config the quantizer + is handed carries a few projector entries and often no `base_model_ep_plan` at all. + Rewriting only what we were handed adds no companion anywhere, while the weights still + shard from the sub-config's own plan — so the scale stays whole against a sharded weight, + which is a wrong answer rather than a crash.""" + from transformers import Glm4vMoeConfig + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + config = Glm4vMoeConfig() + self.assertFalse( + getattr(config, "base_model_ep_plan", None), + "this model no longer nests its plans, so it cannot guard the sub-config walk", + ) + before = dict(getattr(config.text_config, "base_model_ep_plan", None) or {}) + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + after = getattr(config.text_config, "base_model_ep_plan", None) or {} + + companions = { + k + for k in after + if k.endswith(("_scale_inv", "_bias", "_weight_global_scale", "_input_global_scale", "_activation_scale")) + } + self.assertTrue(companions, "no companion reached the sub-config's plan") + self.assertEqual(set(before), set(after) - companions, "the model's own entries were disturbed") + + def test_dequantize_folds_every_quantized_key_to_full_precision(self): + """`dequantize=True` is the escape hatch for hardware that cannot serve the format. It + must claim BOTH checkpoint key layouts — the `{proj}_blocks` + `{proj}_scales` pair + GPT-OSS ships and the `weight` + `weight_scale_inv` pair everything else does — and + land on the plain weight, since a module that was never swapped has no scale slot to + write into.""" + + from transformers.utils.quantization_config import FineGrainedConfig + + for method, sources in ( + ("mxfp4", ("experts.gate_up_proj_blocks", "experts.gate_up_proj_scales")), + ("fp8", ("mlp.down_proj.weight", "mlp.down_proj.weight_scale_inv")), + ): + quantizer = _quantizer_for(FineGrainedConfig(quant_method=method, dequantize=True)) + quantizer.pre_quantized = True + converters = quantizer.get_weight_conversions() + with self.subTest(quant_method=method): + self.assertTrue(converters, f"{method}: dequantize produced no converter") + for source in sources: + self.assertTrue( + any( + re.search(str(pattern), source) + for c in converters + for pattern in ( + c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns] + ) + ), + f"{method}: nothing dequantizes {source}", + ) + + def test_the_impl_rewrites_the_layer_kinds(self): + megamoe = self._planned("deepgemm_megamoe") + self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") + self.assertEqual(megamoe["layers.*.mlp.experts"], "megamoe_experts") + + def test_a_models_own_plan_survives_the_update(self): + """`update_tp_plan` ADDS companions; it must never drop what the model ships. + + The config here is a REAL one, which the `SimpleNamespace` the tests above use cannot + substitute for: a branch keyed on the config CLASS is invisible to a fake, and a name + substring (`"Qwen3" in config.__class__.__name__`) also claims Qwen3-MoE. + """ + from transformers import Qwen3Config, Qwen3MoeConfig, Qwen3VLMoeConfig + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for config_cls in (Qwen3Config, Qwen3MoeConfig): + with self.subTest(config=config_cls.__name__): + config = config_cls() + shipped = dict(config.base_model_tp_plan) + self.assertTrue(shipped, f"{config_cls.__name__} ships no plan to preserve") + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + for key, style in shipped.items(): + self.assertEqual(config.base_model_tp_plan.get(key), style, key) + + # and a multimodal one keeps its plan on the SUB-config: the outer carries none, so + # nothing may be invented there. 25 config classes contain "Qwen3" — a name-substring + # branch reaches every one of them, including these. + multimodal = Qwen3VLMoeConfig() + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(multimodal) + self.assertFalse(getattr(multimodal, "base_model_tp_plan", None) or {}) + self.assertTrue(multimodal.text_config.base_model_tp_plan) + + def test_a_real_moe_models_experts_gain_their_companions(self): + """The other half of the same invariant: preserving the model's entries is what lets the + companion pass find `moe_tp_experts` and give the expert scales an entry of their own.""" + from transformers import Qwen3MoeConfig + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + config = Qwen3MoeConfig() + FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) + plan = config.base_model_tp_plan + self.assertEqual(plan["layers.*.mlp.experts"], "moe_tp_experts") + for companion in ("gate_up_proj_scale_inv", "down_proj_scale_inv", "gate_up_proj_bias"): + self.assertIn(f"layers.*.mlp.experts.{companion}", plan, companion) + + +class FineGrainedGroupsTest(unittest.TestCase): + """A quant config names a format per module SUBSET, because a checkpoint can be more than one.""" + + def _config(self, **kwargs): + from transformers.utils.quantization_config import FineGrainedConfig + + return FineGrainedConfig(**kwargs) + + def test_a_flat_config_is_one_group_over_everything(self): + groups = self._config(quant_method="mxfp8").groups + self.assertEqual({name: group.quant_method for name, group in groups.items()}, {"default": "mxfp8"}) + + def test_a_legacy_expert_dtype_becomes_the_two_groups_it_meant(self): + """DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL + config, because a flat config cannot say "the experts differ". `shared_experts` is DENSE + despite the name, which is why the pattern matches a path segment and not a substring.""" + from transformers.utils.quantization_config import groups_with_expert_dtype + + config = self._config(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0") + config.groups = groups_with_expert_dtype(config.groups, "fp4") + for module, expected in ( + ("model.layers.3.mlp.experts", "mxfp4"), + ("model.layers.3.mlp.experts.gate_up_proj", "mxfp4"), + ("model.layers.3.self_attn.q_proj", "fp8"), + ("model.layers.3.mlp.shared_experts.up_proj", "fp8"), + ): + self.assertEqual(config.group_for(module).quant_method, expected, module) + + def test_resolution_survives_a_round_trip_through_config_json(self): + """`to_json_string` SORTS keys, so a config read back has lost the order it was written + in. Resolution must not depend on it: the catch-all sorts first here and would swallow + every module if targeted groups were not preferred outright.""" + import json + + from transformers.utils.quantization_config import FineGrainedConfig, FineGrainedGroup + + config = FineGrainedConfig( + groups={ + "experts": FineGrainedGroup(quant_method="nvfp4", targets=[r"\.experts($|\.)"]), + "dense": FineGrainedGroup(quant_method="fp8", weight_block_size=(128, 128)), + } + ) + back = FineGrainedConfig.from_dict(json.loads(config.to_json_string())) + self.assertEqual(back.group_for("model.layers.0.mlp.experts").quant_method, "nvfp4") + self.assertEqual(back.group_for("model.layers.0.self_attn.q_proj").quant_method, "fp8") + # a list survives JSON where a tuple does not, and the block size is compared as a tuple + self.assertEqual(back.groups["dense"].weight_block_size, (128, 128)) + + def test_two_targeted_groups_claiming_one_module_is_an_error(self): + from transformers.utils.quantization_config import FineGrainedGroup + + config = self._config( + groups={ + "a": FineGrainedGroup(quant_method="fp8", targets=[r"\.experts"]), + "b": FineGrainedGroup(quant_method="nvfp4", targets=[r"mlp\."]), + } + ) + with self.assertRaises(ValueError): + config.group_for("model.layers.0.mlp.experts") + + def test_a_producers_config_groups_are_normalized(self): + """modelopt describes a format by its parameters; we name it. GLM-5.2-NVFP4 ships exactly + this, and `targets: ["Linear"]` is its spelling of the catch-all.""" + config = self._config( + quant_method="modelopt", + quant_algo="NVFP4", + config_groups={ + "group_0": { + "weights": {"num_bits": 4, "type": "float", "group_size": 16}, + "input_activations": {"num_bits": 4, "type": "float", "group_size": 16, "dynamic": False}, + "targets": ["Linear"], + } + }, + ) + group = config.group_for("model.layers.0.self_attn.q_proj") + self.assertEqual((group.quant_method, group.activation_format), ("nvfp4", "nvfp4")) + self.assertEqual(group.activation_scheme, "static") # `dynamic: False` + + def test_a_format_we_do_not_serve_falls_back_rather_than_guessing(self): + """Half a translation would quantize modules by a rule the producer did not write.""" + config = self._config( + quant_method="modelopt", + quant_algo="NVFP4", + config_groups={ + "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} + }, + ) + self.assertEqual(list(config.groups), ["default"]) + + +class SubtreePatternTest(unittest.TestCase): + def test_a_checkpoints_glob_skip_list_is_bounded(self): + """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any + character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" + from transformers.quantizers.quantizers_utils import should_convert_module + from transformers.utils.quantization_config import FineGrainedConfig + + # as a modelopt checkpoint ships them, through the config that normalizes them + skip = FineGrainedConfig( + quant_method="modelopt", + quant_algo="NVFP4", + ignore=["model.layers.0*", "model.layers.1.*", "self_attn"], + ).modules_to_not_convert + for name in ("model.layers.0.mlp", "model.layers.1.mlp.down_proj", "model.layers.2.self_attn.q_proj"): + self.assertFalse(should_convert_module(name, skip), name) + for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): + self.assertTrue(should_convert_module(name, skip), name) + + +@require_torch +class FineGrainedDeepGemmDispatchTest(unittest.TestCase): + """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled + scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine + and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 + scales have no swizzled layout, so the first never fires for the shape the second catches.""" + + def _routed_to(self, *, sm100, scale_ndim): + """Which backend a block-FP8 linear reaches, given arch and scale layout.""" + kernel, _ = _fake_bundle() + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) + p1, p2, p3, p4 = _loaded(kernel) + with ( + p1, + p2, + p3, + p4, + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), + mock.patch.object(deepgemm, "is_sm100", return_value=sm100), + mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, + ): + fg.finegrained_linear(torch.randn(4, 64, dtype=torch.bfloat16), w, s, block_size=[128, 128]) + return "deepgemm" if dg.called else "triton" + + def test_sm100_never_prefers_deepgemm(self): + self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") + + def test_pre_sm100_still_uses_deepgemm(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") + + def test_swizzled_scales_never_reach_deepgemm(self): + """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + + def test_deepgemm_experts_refuse_swizzled_scales(self): + """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM + experts backend afterwards must fail loudly rather than read the permuted buffer as affine.""" + from transformers.integrations.deepgemm import ( + deepgemm_fp8_fp4_experts_forward, + deepgemm_fp8_fp4_megamoe_experts_forward, + ) + + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" + with ( + mock.patch.object(fg, "is_sm100", return_value=True), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts(cfg, weight_format="mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) + hs = torch.randn(2, 256, dtype=torch.bfloat16) + idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.assertRaisesRegex(RuntimeError, "held swizzled"): + forward(experts, hs, idx, wts) + + +@require_torch +class FineGrainedScaleLayoutTest(unittest.TestCase): + """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 + scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the + scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's + `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" + + E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks + + def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E + cfg._experts_implementation = impl + with ( + mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts( + cfg, block_size=(128, 128), weight_format=weight_format, activation_format=activation_format + ) + model = torch.nn.Module() + model.config = cfg + model.experts = experts + return model, experts + + def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) + self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertTrue(experts.holds_interleaved_gate_up) + # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a + # float32 grid there; GPT-OSS's config has no scale_fmt at all) + _, experts = self._experts("mxfp4", activation_format="bf16") + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + + def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): + for kw in ( + {"weight_format": "mxfp8", "sm100": False}, + {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely + {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales + {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, + ): + with self.subTest(**kw): + _, experts = self._experts(**kw) + self.assertEqual( + experts.gate_up_proj_scale_inv.shape, + (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), + ) + # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block + _, experts = self._experts("fp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) + + def test_megamoe_holds_gate_up_stacked(self): + _, experts = self._experts("fp8", impl="deepgemm_megamoe") + self.assertFalse(experts.holds_interleaved_gate_up) + + def _op_kernel(self): + kernel, _ = _fake_bundle() + kernel.swizzle_mx_scales = mock.Mock( + side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) + ) + kernel.unswizzle_mx_scales = mock.Mock( + side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) + ) + return kernel + + def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + model, experts = self._experts("mxfp8") + kernel = self._op_kernel() + op = FineGrainedSwizzleScales(hf_quantizer=None) + grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + # the loader hands a single target under its pattern, the full name (with suffix) alongside + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" + ) + self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) + # the same converter's weight passes through + out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") + self.assertIs(out["mlp.experts.gate_up_proj"], weight) + # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully + out = op.convert( + { + "experts.down_proj": weight, + "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), + }, + model=model, + full_layer_name="experts.down_proj", + ) + self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) + self.assertIs(out["experts.down_proj"], weight) + # a module that holds affine scales keeps them + affine_model, _ = self._experts("mxfp4", activation_format="bf16") + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, + model=affine_model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + self.assertIs(out["mlp.experts.gate_up_proj"], grid) + self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) + + def test_scale_container_op_round_trips_every_container(self): + """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: + the op brings both into the held e8m0 (exact cast / same bytes), records the container on + the module, and its reverse restores that container on save; weights and native scales + pass both ways.""" + from transformers.integrations.finegrained.conversions import FineGrainedScaleContainer + + op = FineGrainedScaleContainer(hf_quantizer=None) + native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) + for container, shipped in ((torch.float32, native.float()), (torch.uint8, native.view(torch.uint8))): + with self.subTest(container=container): + model, experts = self._experts("mxfp8", sm100=False) # affine e8m0 scale Parameters + out = op.convert( + {"mlp.experts.gate_up_proj": shipped}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + held = out["mlp.experts.gate_up_proj"] + self.assertEqual(held.dtype, torch.float8_e8m0fnu) + self.assertTrue(torch.equal(held.view(torch.uint8), native.view(torch.uint8))) + self.assertIs(experts.scale_container_dtype, container) + back = op.reverse_op.convert({"x": held}, model=model)["x"] + self.assertEqual(back.dtype, container) + self.assertTrue(torch.equal(back, shipped)) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + self.assertIs(op.reverse_op.convert({"w": weight}, model=model)["w"], weight) + model, experts = self._experts("mxfp8", sm100=False) + self.assertIs( + op.convert({"mlp.experts.down_proj": native}, model=model, full_layer_name="experts.down_proj_scale_inv")[ + "mlp.experts.down_proj" + ], + native, + ) + self.assertIsNone(experts.scale_container_dtype) + self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays + + def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + kernel = self._op_kernel() + reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op + artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) + affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") + self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) + kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) + out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") + self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) + + def _quantizer(self, quant_method): + from transformers.utils.quantization_config import FineGrainedConfig + + return _quantizer_for(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions + + def _arch_converters(self): + from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter + + return [ + WeightConverter( + source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"], + target_patterns="mlp.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight", + target_patterns="mlp.experts.down_proj", + operations=[MergeModulelist(dim=0)], + ), + WeightConverter( + source_patterns=["mlp.fc1.weight", "mlp.fc2.weight"], + target_patterns="mlp.fc.weight", + operations=[Concatenate(dim=0)], + ), + ] + + def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): + from transformers.core_model_loading import WeightConverter + from transformers.integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) + + converters = self._quantizer("mxfp8")(self._arch_converters()) + by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} + gate_up_ops = by_target["mlp.experts.gate_up_proj"].operations + # stacked [gate; up] -> [g0, u0, ...], the scale into the held dtype, then the swizzle packs + # the final row order (it reads 1-byte scales) + self.assertIsInstance(gate_up_ops[-3], FineGrainedInterleaveGateUp) + self.assertFalse(gate_up_ops[-3].inverse) + self.assertIsInstance(gate_up_ops[-2], FineGrainedScaleContainer) + self.assertIsInstance(gate_up_ops[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["mlp.experts.down_proj"].operations[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["weight_scale_inv"].operations[0], FineGrainedScaleContainer) # dense linears + self.assertEqual(len(by_target["mlp.fc.weight"].operations), 1) # dense converters untouched + # keys arriving under the fused names get the same layout + self.assertIn("experts.gate_up_proj_bias", by_target) # patterns are stored anchor-stripped + self.assertIsInstance(by_target["experts.down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) + # ...and saving reverses in the opposite order: unswizzle first, then de-interleave + reverse_ops = by_target["mlp.experts.gate_up_proj"].reverse_transform().operations + self.assertIsInstance(reverse_ops[0], FineGrainedSwizzleScales) + self.assertTrue(reverse_ops[0].inverse) + self.assertIsInstance(reverse_ops[2], FineGrainedInterleaveGateUp) + self.assertTrue(reverse_ops[2].inverse) + + def test_blocks_scales_converters_carry_the_layout_ops(self): + from transformers.core_model_loading import WeightConverter + from transformers.integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedSwizzleScales, + ) + + converters = self._quantizer("mxfp4")([]) + by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} + ops = by_target["gate_up_proj"].operations + self.assertIsInstance(ops[-2], FineGrainedInterleaveGateUp) + self.assertIsInstance(ops[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["gate_up_proj_scale_inv"].operations[0], FineGrainedInterleaveGateUp) + self.assertIsInstance(by_target["down_proj"].operations[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) + + def test_catch_all_converters_deliver_the_full_parameter_name(self): + """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax + down_proj_scale_inv) goes through the real converter plumbing: the ops receive the tensor + under the source PATTERN and must hand back the target, which the loader expands to the + full name.""" + from transformers.core_model_loading import WeightConverter, rename_source_key + + model, _ = self._experts("mxfp8") + converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] + for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): + renamed, source_pattern = rename_source_key(key, [], converters) + self.assertEqual(renamed, key) + converter = next(c for c in converters if source_pattern in c.source_patterns) + tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) + converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) + kernel = self._op_kernel() + p1, p2, p3, p4 = _loaded(kernel) + with p1, p2, p3, p4: + out = converter.convert(renamed, model=model) + self.assertEqual(list(out), [key]) + + def test_replacement_keeps_the_models_gate_up_convention(self): + """`replace_with_finegrained_layer` goes through `use_experts_implementation`, which stamps + the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must + survive it (it decides whether the loader interleaves).""" + from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM + from transformers.utils.quantization_config import FineGrainedConfig + + gpt_oss = GptOssConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=32, + num_local_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + sliding_window=16, + ) + qwen = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + ) + for cls, cfg, quant_method, expected in ( + (GptOssForCausalLM, gpt_oss, "mxfp4", False), + (Qwen3MoeForCausalLM, qwen, "mxfp8", True), + ): + with self.subTest(model=cls.__name__), torch.device("meta"): + model = cls(cfg) + fg.replace_with_finegrained_layer( + model, modules_to_not_convert=[], quantization_config=FineGrainedConfig(quant_method=quant_method) + ) + experts = model.model.layers[0].mlp.experts + self.assertIsInstance(experts, fg.FineGrainedExperts) + self.assertIs(experts.is_concatenated, expected) + + def test_interleave_op_follows_the_model_and_the_experts_backend(self): + """The op reads the model's experts modules: a model whose own rows already come + interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself + and keeps the stack; every other case gets the interleaved order, and the reverse restores + the stack.""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + model, experts = self._experts("mxfp8") + out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) + back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(back, stacked) + experts.is_concatenated = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) + experts.is_concatenated = True + experts.holds_interleaved_gate_up = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) diff --git a/tests/integrations/finegrained/test_forwards.py b/tests/integrations/finegrained/test_forwards.py new file mode 100644 index 000000000000..b29574aee63c --- /dev/null +++ b/tests/integrations/finegrained/test_forwards.py @@ -0,0 +1,372 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The same chains as the marshalling tests, against the build the integration actually loads: +what those assert is passed to the kernels, these assert comes back out of them. Needs an +accelerator, since nothing here is mocked.""" + +import unittest +from unittest import mock + +import torch + +import transformers.integrations.finegrained.core as fg +from transformers.integrations.finegrained import ( + FineGrainedExperts, + finegrained_linear, + load_finegrained_kernel, +) +from transformers.testing_utils import ( + backend_synchronize, + require_kernels, + require_torch_accelerator, + torch_device, +) + +from .test_core import _Cfg + + +class _InputScaledRMSNorm(torch.nn.Module): + """A post-expert norm that scales by ``1 + weight`` BEFORE normalizing, so the row's + mean square is taken on the scaled values (not the same function as normalizing first).""" + + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16)) + + def forward(self, x): + out = x.float() * (1.0 + self.weight.float()) + return (out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) + + +class _UnfusableNorm(torch.nn.Module): + """A norm whose math is none of the fused forms (a bias the kernels have no slot for): the + binder has to leave it a module call rather than pick the nearest form.""" + + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim, dtype=torch.bfloat16)) + self.bias = torch.nn.Parameter(torch.full((dim,), 0.5, dtype=torch.bfloat16)) + + def forward(self, x): + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) + return (normed * self.weight.float() + self.bias.float()).type_as(x) + + +def _nvfp4_two_level(kernel, weight): + """Canonical two-level NVFP4 quant of one matrix: the global is ``amax / (6 * 448)`` — the + smallest one that keeps every e4m3 block scale in range — and the block quant runs on the + normalized values. Returns ``(packed_e2m1, e4m3 block scales, fp32 global)``.""" + global_scale = (weight.abs().amax() / (6.0 * 448.0)).clamp(min=torch.finfo(torch.float32).tiny) + packed, block = kernel.nvfp4_act_quant((weight.float() / global_scale).contiguous()) + return packed.view(torch.int8), block, global_scale + + +def _unpack_e2m1(packed): + """Packed-E2M1 bytes (two codes per byte along the last axis) back to fp32 values.""" + lut = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + device=packed.device, + ) + b = packed.view(torch.uint8) + pairs = torch.stack([lut[(b & 0x0F).long()], lut[(b >> 4).long()]], dim=-1) + return pairs.reshape(*packed.shape[:-1], packed.shape[-1] * 2) + + +@require_kernels +@require_torch_accelerator +class FineGrainedForwardTest(unittest.TestCase): + """`test_core.py` pins the calling contract; this pins that the contract produces correct + numbers, which no amount of argument-matching can show. + + The load is NOT wrapped in a skip: `kernels` missing is the one environment reason to sit these + out and the decorator states it, so anything else that stops the bundle loading — a broken + loader, no build for this torch, a bad hub entry — has to fail, not quietly pass.""" + + @classmethod + def setUpClass(cls): + load_finegrained_kernel() + + def _block_fp8_weight(self, N, K, *, ue8m0): + """(N, K) E4M3 + its (N/128, K/128) inv-scale grid, and the exact dequantized values the + kernel will see — so the reference is the quantization floor, not the pre-quant weight.""" + w = torch.randn(N, K, device=torch_device, dtype=torch.float32) + blocks = w.reshape(N // 128, 128, K // 128, 128) + amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12) + inv = amax / 448.0 + if ue8m0: # power-of-two scales: the tcgen05 dot_scaled format + inv = torch.pow(2.0, torch.ceil(torch.log2(inv))) + q = (blocks / inv).to(torch.float8_e4m3fn) + deq = (q.float() * inv).reshape(N, K) + return q.reshape(N, K), inv.reshape(N // 128, K // 128).contiguous(), deq + + def test_linear_forward_matches_the_quantization_floor(self): + """A real quantized linear against its own dequantized weight. Loose vs bf16 would pass + even if the kernel silently dropped a K-block; comparing to the floor does not.""" + M, N, K = 256, 512, 384 + for ue8m0 in (False, True): + with self.subTest(ue8m0=ue8m0): + w, s, deq = self._block_fp8_weight(N, K, ue8m0=ue8m0) + x = torch.randn(M, K, device=torch_device, dtype=torch.bfloat16) + out = finegrained_linear(x, w, s, block_size=[128, 128]) + ref = x.float() @ deq.t() + rel = ((out.float() - ref).norm() / ref.norm()).item() + self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") + + def test_on_the_fly_group_formats_round_trip_within_their_floor(self): + """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in + the module's layout (swizzled where held so); dequantizing them back lands within each + format's floor.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + for fmt, floor in (("mxfp8", 4e-2), ("mxfp4", 2e-1), ("nvfp4", 2e-1)): + with self.subTest(fmt=fmt): + experts = FineGrainedExperts(cfg, weight_format=fmt).to(torch_device) + model = torch.nn.Module() + model.experts = experts + original = (torch.randn(4, 512, 512, device=torch_device) * 0.02).to(torch.bfloat16) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] + self.assertEqual( + (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) + ) + self.assertEqual( + (scale.dtype, scale.shape), + (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), + ) + if scale.ndim == 5: + kernel = load_finegrained_kernel() + scale = kernel.unswizzle_mx_scales(scale, 512, scale.shape[2] * 4, num_experts=4) + deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) + if fmt == "nvfp4": + global_scale = out["experts.gate_up_proj_weight_global_scale"] + deq = deq * global_scale.reshape(-1, 1, 1) + else: + self.assertNotIn("experts.gate_up_proj_weight_global_scale", out) + rel = ((deq - original.float()).norm() / original.float().norm()).item() + self.assertLess(rel, floor, f"{fmt}: {rel:.3f}") + self.assertGreater(rel, 0.0) + + def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes=(8.0, 1.0)): + """NVFP4 experts in the layout a modelopt checkpoint delivers: every (expert, half) of the + gate|up stack quantized against its OWN global (``weight_scale_2`` is per projection, and + the two can differ per expert), the down per expert, plus calibrated activation + globals. Returns the module, the exact dequantized weights (so the reference is the + quantization floor rather than the pre-quant weight), and the dequantization a reader that + applied ONE of the two globals to both halves would get.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + kernel = load_finegrained_kernel() + experts = FineGrainedExperts(cfg, weight_format="nvfp4", activation_format=activation_format).to(torch_device) + dequantized, single_global, checkpoint_globals = {}, {}, {} + for proj, rows, cols in ( + ("gate_up_proj", 2 * cfg.intermediate_size, cfg.hidden_size), + ("down_proj", cfg.hidden_size, cfg.intermediate_size), + ): + reference = torch.randn(cfg.num_local_experts, rows, cols, device=torch_device) + if proj == "gate_up_proj": # the halves interleave by row; draw them far apart + reference[:, 0::2] *= magnitudes[0] + reference[:, 1::2] *= magnitudes[1] + packed = torch.empty(cfg.num_local_experts, rows, cols // 2, dtype=torch.int8, device=torch_device) + block = torch.empty( + cfg.num_local_experts, rows, cols // 16, dtype=torch.float8_e4m3fn, device=torch_device + ) + halves = 2 if proj == "gate_up_proj" else 1 + globals_ = torch.empty(cfg.num_local_experts, halves, device=torch_device) + deq, deq_one = torch.empty_like(reference), torch.empty_like(reference) + for e in range(cfg.num_local_experts): + for h in range(halves): + rows_h = slice(h, None, halves) + q, b, g = _nvfp4_two_level(kernel, reference[e, rows_h]) + packed[e, rows_h], block[e, rows_h], globals_[e, h] = q, b, g + values = _unpack_e2m1(q) * b.float().repeat_interleave(16, dim=-1) + deq[e, rows_h] = values * g + if halves == 2: + # the control: BOTH halves read the up half's global, which is what a kernel + # that ignored the layout does. It has to be the up one — putting the gate's + # global on both scales the expert's whole output uniformly, and the + # post-expert RMSNorm divides exactly that out; a wrong GATE scale goes + # through the SiLU and survives. + for h in (0, 1): + rows_h = slice(h, None, 2) + deq_one[e, rows_h] = deq[e, rows_h] * (globals_[e, 1] / globals_[e, h]) + scale = block if getattr(experts, f"{proj}_scale_inv").dim() != 5 else kernel.swizzle_mx_scales(block) + setattr(experts, proj, torch.nn.Parameter(packed, requires_grad=False)) + setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(scale, requires_grad=False)) + checkpoint_globals[f"experts.{proj}_weight_scale_2"] = globals_ + dequantized[proj], single_global[proj] = deq, deq_one + # calibrated input_scale: one value for the gate_up (quantized before routing), one per + # expert for the down (its rows are per expert). A weight-only module holds neither — + # nothing there quantizes activations — so the checkpoint's keys stay unused. + if activation_format != "bf16": + experts.gate_up_proj_input_global_scale = torch.nn.Parameter( + torch.full((1,), 0.05, device=torch_device), requires_grad=False + ) + checkpoint_globals["experts.down_proj_input_scale"] = torch.linspace( + 0.02, 0.08, cfg.num_local_experts, device=torch_device + ) + single_global["down_proj"] = dequantized["down_proj"] # the down has one global either way + # The globals go in through the loader's own converter, so the module holds exactly what a + # real checkpoint puts there: one per expert, with the gate|up pair merged onto the down. + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_input_global_scale", + ] + for target, value in FineGrainedWeightGlobals().convert(checkpoint_globals, target_patterns=targets).items(): + setattr(experts, target.split(".")[-1], torch.nn.Parameter(value, requires_grad=False)) + return experts, dequantized, single_global + + def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): + """The experts chain in torch over the dequantized weights: gate|up (interleaved columns), + SwiGLU, down, the optional per-expert output norm, then the routing-weighted sum.""" + out = torch.zeros_like(x, dtype=torch.float32) + for token in range(x.shape[0]): + for slot in range(idx.shape[1]): + e = int(idx[token, slot]) + pre = x[token].float() @ dequantized["gate_up_proj"][e].t() + inter = torch.nn.functional.silu(pre[0::2]) * pre[1::2] + row = inter @ dequantized["down_proj"][e].t() + if post_norm is not None: + row = post_norm(row.to(x.dtype)).float() + out[token] += row * float(wts[token, slot]) + return out + + def test_modelopt_nvfp4_experts_merge_their_globals_and_fuse_the_post_norm(self): + """The stacked checkpoint shape end to end through the integration: a gate|up stack + calibrated per half (merged at load), calibrated activation globals, and a per-expert + output norm the binder fuses, on all three forwards. Weight-only pins the numbers against the dequantized weights (nothing else + rounds); the W4A4 chain rides the 4-bit activation floor, so it is pinned by being far + closer to the right globals than to one global folded over both halves — what a kernel + that ignored the layout computes.""" + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + x = torch.randn(8, cfg.hidden_size, device=torch_device, dtype=torch.bfloat16) * 0.1 + idx = torch.randint(0, cfg.num_local_experts, (8, 2), device=torch_device, dtype=torch.long) + wts = torch.rand(8, 2, device=torch_device, dtype=torch.bfloat16) + forwards = ( + ("grouped", fg.finegrained_grouped_mm_experts_forward), + ("batched", fg.finegrained_batched_mm_experts_forward), + ("eager", lambda module, *args: module(*args)), + ) + + def relative(out, reference): + return ((out.float() - reference).norm() / reference.norm()).item() + + for activation_format, floor in (("bf16", 0.02), (None, 0.35)): + experts, dequantized, one_global = self._nvfp4_experts_modelopt_layout(cfg, activation_format) + reference = self._moe_reference(dequantized, x, idx, wts) + wrong = self._moe_reference(one_global, x, idx, wts) + for name, forward in forwards: + with self.subTest(forward=name, activation_format=activation_format): + out = forward(experts, x, idx, wts) + rel, rel_wrong = relative(out, reference), relative(out, wrong) + self.assertLess(rel, floor, f"{name} diverged from the dequantized reference: {rel:.3f}") + self.assertLess( + 3 * rel, + rel_wrong, + f"{name} is as close to ONE global over both gate|up halves ({rel_wrong:.3f}) " + f"as to the per-half ones ({rel:.3f}) — the fold is not reading the halves", + ) + + # the same chain with the model's per-expert output norm, which runs on the routed + # rows: every forward has to apply it where the reference forwards do. A name the + # kernels implement is folded into the chain's reduce and anything else is the module + # itself, so both have to land on the same numbers as the reference. + for norm, norm_name in ( + (torch.nn.RMSNorm(cfg.hidden_size, device=torch_device, dtype=torch.bfloat16), "rms_norm"), + (_InputScaledRMSNorm(cfg.hidden_size).to(torch_device), "input_scaled_rms_norm"), + # a form the kernels do not implement: the module is called on the rows + (_UnfusableNorm(cfg.hidden_size).to(torch_device), "a_models_own_norm"), + ): + norm.weight.data.uniform_(0.5, 1.5) + experts.post_expert_norm, experts.post_expert_norm_name = norm, norm_name + experts.has_post_expert_norm = True + normed = self._moe_reference(dequantized, x, idx, wts, norm) + self.assertGreater(relative(normed, reference), 0.05, "the norm has to change the output") + for name, forward in forwards: + with self.subTest(forward=name, activation_format=activation_format, norm=norm_name): + rel = relative(forward(experts, x, idx, wts), normed) + self.assertLess(rel, floor, f"{name} diverged with the post-expert norm: {rel:.3f}") + + def test_experts_scale_layout_op_keeps_the_forward_correct(self): + """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op + fills them from the affine grid, the fused forward matches the affine module's, and the + reverse op hands the affine grid back bitwise.""" + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = ( + 512, + 256, + 4, + ) # K >= 256: the swizzled arm has tiles + cfg._experts_implementation = "grouped_mm" + experts = FineGrainedExperts(cfg, weight_format="mxfp8").to(torch_device) + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): + affine_experts = FineGrainedExperts(cfg, weight_format="mxfp8").to(torch_device) + model = torch.nn.Module() + model.experts = experts + op = FineGrainedSwizzleScales(hf_quantizer=None) + + affine = {} + for proj, rows in (("gate_up_proj", 2 * cfg.intermediate_size), ("down_proj", cfg.hidden_size)): + cols = cfg.hidden_size if proj == "gate_up_proj" else cfg.intermediate_size + weight = torch.randn(cfg.num_local_experts, rows, cols, device=torch_device).to(torch.float8_e4m3fn) + grid = torch.randint( + 120, 134, (cfg.num_local_experts, rows, cols // 32), dtype=torch.uint8, device=torch_device + ) + affine[proj] = grid.view(torch.float8_e8m0fnu) + for module in (experts, affine_experts): + setattr(module, proj, torch.nn.Parameter(weight.clone(), requires_grad=False)) + setattr(affine_experts, f"{proj}_scale_inv", torch.nn.Parameter(affine[proj].clone(), requires_grad=False)) + self.assertEqual( + getattr(experts, f"{proj}_scale_inv").dim(), 5, "the module should hold this projection swizzled" + ) + out = op.convert( + {f"mlp.experts.{proj}": affine[proj]}, model=model, full_layer_name=f"experts.{proj}_scale_inv" + ) + swizzled = out[f"mlp.experts.{proj}"] + self.assertEqual(swizzled.shape, getattr(experts, f"{proj}_scale_inv").shape) + setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(swizzled, requires_grad=False)) + + x = torch.randn(8, cfg.hidden_size, device=torch_device, dtype=torch.bfloat16) + idx = torch.randint(0, cfg.num_local_experts, (8, 2), device=torch_device, dtype=torch.long) + wts = torch.rand(8, 2, device=torch_device, dtype=torch.bfloat16) + reference = fg.finegrained_grouped_mm_experts_forward(affine_experts, x, idx, wts) + out = fg.finegrained_grouped_mm_experts_forward(experts, x, idx, wts) + backend_synchronize(torch_device) + rel = ((out.float() - reference.float()).norm() / reference.float().norm().clamp(min=1e-9)).item() + # different tuned arms reorder the fp32 reduction (~1e-5); a layout bug is O(1) + self.assertLess(rel, 1e-3, f"swizzled forward diverged from the affine one: {rel:.2e}") + # the eager per-expert loop reads one expert's slice of the swizzled stack directly + eager_reference = affine_experts(x, idx, wts) + eager_out = experts(x, idx, wts) + rel = ( + (eager_out.float() - eager_reference.float()).norm() / eager_reference.float().norm().clamp(min=1e-9) + ).item() + self.assertLess(rel, 1e-3, f"eager forward on swizzled scales diverged from affine: {rel:.2e}") + + reverse = op.reverse_op + for proj, grid in affine.items(): + restored = reverse.convert({proj: getattr(experts, f"{proj}_scale_inv").data})[proj] + self.assertTrue(torch.equal(restored.view(torch.uint8), grid.view(torch.uint8)), proj) diff --git a/tests/integrations/finegrained/test_models.py b/tests/integrations/finegrained/test_models.py new file mode 100644 index 000000000000..b8d3d315ae96 --- /dev/null +++ b/tests/integrations/finegrained/test_models.py @@ -0,0 +1,561 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The sharded load paths on real checkpoints: expert parallelism and intra-expert tensor +parallelism must both reproduce what `device_map` computes, for every format a shipped model +quantizes with.""" + +import os +import re +import socket +import subprocess +import tempfile + +import torch + +from transformers.testing_utils import ( + TestCasePlus, + backend_empty_cache, + require_torch_multi_accelerator, + torch_device, +) + + +def _checkpoint_shapes(path): + """`{key: shape}` of a saved checkpoint, for comparing one save against another.""" + import glob + + from safetensors import safe_open + + shapes = {} + for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shapes[key] = tuple(handle.get_slice(key).get_shape()) + return shapes + + +_SHARDING_WORKER = """ +import importlib, os, sys, torch +from transformers.distributed import DistributedConfig + +model_dir, out_dir, modes, cls_path = sys.argv[1], sys.argv[2], sys.argv[3].split(","), sys.argv[4] +module_name, cls_name = cls_path.split(":") +model_cls = getattr(importlib.import_module(module_name), cls_name) +world = int(os.environ["WORLD_SIZE"]) +ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + +# both modes in ONE launch: they share the mesh, so the process group is created once and only +# the plan differs -- a second torchrun would pay another interpreter + CUDA start +for mode in modes: + model = model_cls.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=(mode == "ep")), + ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, "gate_up_proj", None) + if weight is None: + weight = experts.up_proj + local = weight.to_local() if hasattr(weight, "to_local") else weight + with torch.no_grad(): + logits = model(ids.to(model.device)).logits.float().cpu() + if int(os.environ["RANK"]) == 0: + torch.save({"logits": logits, "expert_local": tuple(local.shape)}, + os.path.join(out_dir, mode + ".pt")) + del model + torch.accelerator.empty_cache() +""" + + +def _checkpoint_expert_shape(model_dir): + """The UNSHARDED `(experts, rows)` of the stacked gate|up projection the module holds. + + Read from the checkpoint, which ships the experts either already fused (`experts.gate_up_proj`) + or one per expert (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`, whose rows are gate and + up separately). Taking the baseline from whichever leg ran first would compare one mode to + another, and a mode that placed nothing would still look sharded. + """ + import glob + import json + import os + + from safetensors import safe_open + + shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + index = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as fh: + weight_map = json.load(fh)["weight_map"] + shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) + + experts: set[int] = set() + rows = 0 + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shape = None + if re.search(r"\.experts\.gate_up_proj$", key): + shape = handle.get_slice(key).get_shape() + if len(shape) == 3: + return (shape[0], shape[1]) + if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): + experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) + rows = max(rows, handle.get_slice(key).get_shape()[0]) + if not experts or not rows: + raise AssertionError(f"no expert projection in {model_dir} to size the shard against") + return (len(experts), 2 * rows) # the module stacks gate and up into one row extent + + +@require_torch_multi_accelerator +class FineGrainedLoadPathEquivalenceTest(TestCasePlus): + """Expert parallelism and intra-expert tensor parallelism must both reproduce what + `device_map` gives, which is the baseline because it places modules across devices without + splitting a single tensor — no process group, no DTensor, no collectives. So it computes the + unsharded answer, on the multi-GPU path people actually deploy. + + Each model is built in the format it actually ships, because the conversion path differs by + format and by checkpoint layout — a per-expert FP8 checkpoint and a fused NVFP4 one reach the + experts through different converters, and bugs have hidden in exactly that gap. + + Three properties make this able to fail: + * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the + shard it already holds, so they come out correctly sized whatever the plan says and a + plan that shards no scale at all still produces the right answer. + * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty + (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is skipped EXPLICITLY rather + than passing vacuously. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + """ + + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}` — each model + in the format it actually ships, with the quantization config that format arrives under, + not just its name: block-FP8 carries a `weight_block_size`, and an NVFP4 checkpoint comes + from modelopt under `quant_algo`, which is remapped on construction. Dims are multiples + of the 128 block so a 2-way split stays block-aligned, and the expert count divides the + mesh.""" + from transformers import ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + DeepseekV4Config, + DeepseekV4ForCausalLM, + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + Glm4vMoeTextConfig, + Glm4vMoeVisionConfig, + GptOssConfig, + GptOssForCausalLM, + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLConfig, + Mistral4Config, + Mistral4ForCausalLM, + ) + from transformers.utils.quantization_config import FineGrainedConfig + + return { + "deepseek_v3-fp8": ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "n_group": 1, + "topk_group": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "q_lora_rank": None, + "kv_lora_rank": 32, + "qk_nope_head_dim": 32, + "qk_rope_head_dim": 16, + "v_head_dim": 32, + }, + # DeepSeek-V3 ships block-FP8: 128x128 weight blocks, activations quantized + # per token at run time + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), + ), + # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a + # calibrated activation scale per quantized module, which for a MoE is one per expert. + # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense + # counterpart of the same export. + "mistral4-fp8_tensor_static": ( + Mistral4Config, + Mistral4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static"), + ), + # MULTIMODAL: the experts' plans live on `text_config`, not on the config the + # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer + # one adds no companion while the weights still shard, which is how a real + # multimodal MoE ends up with whole scales against sharded weights. + "glm4v_moe-nvfp4": ( + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + { + "text_config": Glm4vMoeTextConfig( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, + rope_theta=10000, + tie_word_embeddings=True, + bos_token_id=0, + eos_token_id=0, + pad_token_id=0, + n_routed_experts=4, + n_shared_experts=1, + n_group=1, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ), + "vision_config": Glm4vMoeVisionConfig( + depth=2, + num_heads=4, + hidden_size=64, + out_hidden_size=256, + intermediate_size=64, + patch_size=14, + spatial_merge_size=1, + temporal_patch_size=2, + ), + }, + # the GLM NVFP4 checkpoints are modelopt exports: `quant_algo` names the format + # and `FineGrainedConfig` remaps it to nvfp4 at construction + FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4"), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases — and + # an EP plan that already names those biases, so the companion rules meet entries + # the model wrote itself + "gpt_oss-mxfp4": ( + GptOssConfig, + GptOssForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + }, + # GPT-OSS ships weight-only: raw bf16 activations against packed fp4 weights + FineGrainedConfig(quant_method="mxfp4", activation_format="bf16"), + ), + # MIXED precision, and the only entry whose expert format is not the quantization + # config's: `expert_dtype` is a model-config side-channel that makes the EXPERTS + # mxfp4 while the dense and attention paths stay block-FP8 — with scales in UE8M0 + # containers rather than fp32, the other `scale_fmt` + "deepseek_v4-fp4_experts+fp8_dense": ( + DeepseekV4Config, + DeepseekV4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "expert_dtype": "fp4", + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second MULTIMODAL nesting, in the group-32 MX format. The sub-config shape + # follows this model's own tester; only the MoE dims are raised to a multiple of + # the 128 block so a 2-way split stays block-aligned. + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + # 512 so the 2-way split leaves 256: the experts' down projection + # contracts over this, and a sharded 128 has no 128-wide swizzled tile + "intermediate_size": 512, + "dense_intermediate_size": 256, + "shared_intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "rotary_dim": 32, + "vocab_size": 64, + "max_position_embeddings": 32, + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "moe_layer_freq": [0, 1], + "layer_types": ["full_attention", "minimax_m3_sparse"], + "tie_word_embeddings": False, + "index_n_heads": 2, + "index_head_dim": 16, + "index_block_size": 8, + "index_topk_blocks": 4, + "index_local_blocks": 1, + }, + "vision_config": { + # 256 so the 2-way SPLIT is still 128-aligned: an MXFP8 weight with + # pre-swizzled scales is read in 128-wide K tiles, and a sharded 128 dim + # leaves 64 — which has none to offer, and the tuner has no config at all + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 14, + "patch_size": 14, + "temporal_patch_size": 2, + "spatial_merge_size": 1, + }, + "image_token_index": 4, + "video_token_index": 5, + "projector_hidden_size": 256, + "pad_token_id": 2, + }, + FineGrainedConfig(quant_method="mxfp8"), + ), + } + + @classmethod + def setUpClass(cls): + cls._tmp = tempfile.TemporaryDirectory() + # the kernels autotune per shape from a cold cache, which dwarfs everything else here + cls._env = { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), + } + os.environ.update(cls._env) + + @classmethod + def tearDownClass(cls): + cls._tmp.cleanup() + + @staticmethod + def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + def _fixture(self, label): + """A PRE-QUANTIZED checkpoint of one tiny model, in the format that model ships.""" + + config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] + config = config_cls(**kwargs) + + bf16_dir = os.path.join(self._tmp.name, label, "bf16") + quant_dir = os.path.join(self._tmp.name, label, "quantized") + torch.manual_seed(0) + # BF16 is what every model here ships, and the dtype decides which kernel arms run: a + # float32 checkpoint (torch's default for a freshly built model, which `dtype="auto"` + # then faithfully reloads) sends the weight-only formats down arms `tl.dot_scaled` + # cannot serve at all, so the suite would exercise a dtype nobody deploys and leave the + # real one uncovered. Saving in bf16 makes `"auto"` mean bf16 for every load below. + model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) + quantized = model_cls.from_pretrained( + bf16_dir, + dtype="auto", + attn_implementation="eager", + quantization_config=quantization_config, + device_map=f"{torch_device}:0", + ) + quantized.save_pretrained(quant_dir, safe_serialization=True) + del quantized + backend_empty_cache(torch_device) + + # SAVE must restore the checkpoint's own layout: the layout ops (gate|up interleave, + # scale container, swizzle) each have a reverse, and a quantized model reloaded and + # written again has to land on the same keys and shapes. A broken reverse writes a + # corrupt checkpoint silently — the module still reads back whatever it wrote. + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") + round_trip = os.path.join(self._tmp.name, label, "round_trip") + reloaded.save_pretrained(round_trip, safe_serialization=True) + del reloaded + backend_empty_cache(torch_device) + self.assertEqual( + _checkpoint_shapes(quant_dir), _checkpoint_shapes(round_trip), f"{label}: save did not round-trip" + ) + return quant_dir, config, model_cls + + def _logits(self, model_dir, model_cls, **load_kwargs): + model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + logits = model(ids).logits.float().cpu() + del model + backend_empty_cache(torch_device) + return logits + + def _rounding_floor(self, model_dir, model_cls, reference): + """How far this model's logits move when every sharded block is perturbed by BF16 rounding. + + Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different + order, so the first sharded op differs by an ULP. How far that travels is a property of + the MODEL, not of the sharding — a chain of MoE layers can amplify it a hundredfold while + a dense stack barely moves. Injecting the same magnitude on ONE device measures that + amplification directly, giving each model a budget its own conditioning earns. + """ + + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" + ).eval() + # Perturb EVERY block a sharded reduction passes through, not one: TP re-orders the sum + # in each of them, so the rounding accumulates down the stack instead of cancelling. + # Perturbing a single site with random noise measured LESS movement than sharding caused + # even at 14x the magnitude, which is what accumulation looks like from the wrong model. + blocks = [ + m + for n, m in model.named_modules() + if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 + ] + generator = torch.Generator(device=model.device).manual_seed(0) + + def perturb(module, args, output): + tensor = output[0] if isinstance(output, tuple) else output + if not torch.is_tensor(tensor): + return output + ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() + noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp + tensor = tensor + noise + return (tensor,) + output[1:] if isinstance(output, tuple) else tensor + + handles = [b.register_forward_hook(perturb) for b in blocks] + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + perturbed = model(ids).logits.float().cpu() + for handle in handles: + handle.remove() + del model + backend_empty_cache(torch_device) + return (perturbed - reference).abs().max().item() + + def _sharded(self, model_dir, model_cls, modes): + """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" + + script = os.path.join(self._tmp.name, "sharding_worker.py") + with open(script, "w", encoding="utf-8") as fh: + fh.write(_SHARDING_WORKER) + out = os.path.join(self._tmp.name, "out") + os.makedirs(out, exist_ok=True) + subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={self._free_port()}", + script, + model_dir, + out, + ",".join(modes), + # the model's own class: `AutoModelForCausalLM` cannot resolve a multimodal + # `ForConditionalGeneration` from its config + f"{model_cls.__module__}:{model_cls.__name__}", + ], + check=True, + env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, + ) + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} + + def test_every_load_path_agrees(self): + for label in self._model_table(): + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # the baseline, spread across both devices: `max_memory` forces a real split, + # since `auto` alone fits this whole model on one GPU and would quietly compare + # the sharded legs against a single-device load + reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) + noise_floor = self._rounding_floor(model_dir, model_cls, reference) + + # Only the modes this model's OWN plans shard experts under. A mode whose plan + # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places + # nothing, so running it would pass without being evidence of anything. A + # multimodal model keeps these on a sub-config. + owners = [config] + [ + c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) + ] + modes = [ + mode + for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) + if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) + ] + self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") + whole = _checkpoint_expert_shape(model_dir) + for mode, payload in self._sharded(model_dir, model_cls, modes).items(): + with self.subTest(model=label, mode=mode): + local = payload["expert_local"] + # against the CHECKPOINT's own shape, never another mode's shard: seeding + # this from the first leg made the second leg check EP against TP, so a + # mode that placed nothing still looked sharded (glm4v's TP left its + # experts replicated and passed). + self.assertLess( + local[0] * local[1], + whole[0] * whole[1], + f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " + f"this leg cannot catch a bad axis", + ) + # Sharding reorders reductions, so the legs differ by BF16 rounding + # whatever the axes — and how much that shows at the logits is the + # model's own business (the NVFP4 fixture amplifies ~100x, the dense + # ones barely). Hence each model's measured floor, not a fixed + # tolerance. A wrong shard axis lands orders of magnitude above it. + budget = max(2e-2, noise_floor) + gap = (payload["logits"] - reference).abs().max().item() + self.assertLessEqual( + gap, + budget, + f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " + f"model's own rounding budget {budget:.4g} — that is a sharding bug, " + f"not reduction order", + ) + if noise_floor <= 2e-2: + # only meaningful where rounding does NOT already move the argmax + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) diff --git a/tests/kernels/test_finegrained_fp8.py b/tests/integrations/test_finegrained_fp8.py similarity index 100% rename from tests/kernels/test_finegrained_fp8.py rename to tests/integrations/test_finegrained_fp8.py diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py deleted file mode 100644 index c239bdef79a9..000000000000 --- a/tests/kernels/test_finegrained.py +++ /dev/null @@ -1,2769 +0,0 @@ -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Marshalling tests for the multi-format ``finegrained`` integration: the kernel bundle is -mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` -(the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" - -import os -import re -import socket -import subprocess -import tempfile -import unittest -import warnings -from contextlib import ExitStack, contextmanager -from dataclasses import dataclass -from unittest import mock - -import torch - -import transformers.integrations.deepgemm as deepgemm -import transformers.integrations.finegrained.core as fg -from transformers.integrations.finegrained import ( - FineGrainedExperts, - FineGrainedLinear, - finegrained_linear, - load_finegrained_kernel, -) -from transformers.testing_utils import ( - TestCasePlus, - require_torch, - require_torch_gpu, - require_torch_multi_accelerator, -) - - -@dataclass -class _Call: - args: tuple - kwargs: dict - - -class _Recorder: - """Records every call and returns a correctly-shaped zero tensor for matmuls.""" - - def __init__(self): - self.calls = {} - - def _op(self, name): - def op(*args, **kwargs): - self.calls.setdefault(name, []).append(_Call(args, kwargs)) - if name == "swizzle_mx_scales": - return args[0] - if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) - hidden, down = args[0], kwargs["down_proj"] - return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) - a, b = args[0], args[1] - n = b.shape[-2] - gather = kwargs.get("gather_idx") - rows = gather.shape[0] if gather is not None else a.shape[0] - return a.new_zeros(rows, n, dtype=torch.bfloat16) - - return op - - -def _fake_bundle(): - rec = _Recorder() - kernel = mock.Mock() - for name in ( - "matmul_2d", - "matmul_batched", - "matmul_grouped", - "moe_fused_batched", - "moe_fused_grouped", - "swizzle_mx_scales", - "unswizzle_mx_scales", - "mxfp8_act_quant", - "mxfp4_act_quant", - "nvfp4_act_quant", - ): - setattr(kernel, name, rec._op(name)) - kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") - kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") - return kernel, rec - - -class _Cfg: - hidden_size = 64 - num_local_experts = 4 - intermediate_size = 32 - hidden_act = "silu" - _experts_implementation = None - - -def _loaded(kernel): - return ( - mock.patch.object(fg, "_FINEGRAINED", None), - mock.patch.object(fg, "is_kernels_available", return_value=True), - mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), - # a locally importable checkout (FINEGRAINED_KERNELS_PATH / installed package) - # takes precedence over the hub loader; tests must stay hermetic to the fake bundle - mock.patch.object(fg, "_import_local_finegrained", return_value=None), - ) - - -def _quantizer_for(config): - """The quantizer a `from_pretrained` would pick for this config. Since the arms were split - per scheme, constructing the base directly only exercises what every scheme shares.""" - from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING - - method = getattr(config.quant_method, "value", config.quant_method) - return AUTO_QUANTIZER_MAPPING[method](config) - - -@require_torch -class FineGrainedLoaderTest(unittest.TestCase): - def test_loader_requires_every_symbol(self): - kernel, _ = _fake_bundle() - del kernel.matmul_grouped - # Mock auto-creates attributes; force the miss - kernel.matmul_grouped = None - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, self.assertRaises(ImportError) as ctx: - load_finegrained_kernel() - self.assertIn("matmul_grouped", str(ctx.exception)) - - def test_loader_is_compile_safe_cold(self): - """Cold path: the compiled call is first to load, so the opaque loader node runs its full - body under compile and must return None, never the bundle (`torch.* op returned - non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" - kernel, _ = _fake_bundle() - kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return load_finegrained_kernel().matmul_2d(x) - - out = run(torch.zeros(3)) - self.assertTrue(torch.equal(out, torch.ones(3))) - - def test_loader_is_compile_safe_when_warm(self): - """Warm path, which is the production order: eager warm-up, then compile. The loader hits - its short-circuit at trace time — the branch that must also return None.""" - kernel, _ = _fake_bundle() - kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - load_finegrained_kernel() - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return load_finegrained_kernel().matmul_2d(x) - - out = run(torch.zeros(3)) - self.assertTrue(torch.equal(out, torch.ones(3))) - - def test_loader_binds_all_symbols(self): - kernel, _ = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - bundle = load_finegrained_kernel() - self.assertIs(bundle.matmul_2d, kernel.matmul_2d) - self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) - - -@require_torch -class FineGrainedLinearMarshallingTest(unittest.TestCase): - def _run(self, **linear_kwargs): - kernel, rec = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - x = torch.randn(3, 5, 64, dtype=torch.bfloat16) - w = torch.randn(32, 64).to(torch.float8_e4m3fn) - ws = torch.randn(1, 1, dtype=torch.float32) - out = finegrained_linear(x, w, ws, **linear_kwargs) - return out, rec.calls["matmul_2d"][-1] - - def test_scale_rides_positionally_and_block_size_dies(self): - out, call = self._run(block_size=[128, 128]) - # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here - self.assertEqual(call.args[0].shape, (15, 64)) - self.assertIsNone(call.args[2]) - self.assertEqual(call.args[3].dtype, torch.float32) - self.assertNotIn("block_size", call.kwargs) - self.assertEqual(out.shape, (3, 5, 32)) - - def test_static_activation_scale_is_As(self): - scale = torch.tensor(0.5) - _, call = self._run(activation_scale=scale) - self.assertIs(call.args[2], scale) - - def test_nvfp4_global_and_weight_only_format(self): - g = torch.tensor(2.0) - _, call = self._run(weight_global_scale=g, activation_format="bf16") - self.assertIs(call.kwargs["b_global_scale"], g) - self.assertEqual(call.kwargs["activation_format"], "bf16") - - def test_module_forward_threads_everything(self): - kernel, rec = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - m = FineGrainedLinear( - 64, - 32, - block_size=None, - weight_format="nvfp4", - activation_format="bf16", - ) - m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) - out = m(torch.randn(2, 64, dtype=torch.bfloat16)) - call = rec.calls["matmul_2d"][-1] - self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) - self.assertEqual(call.kwargs["activation_format"], "bf16") - self.assertEqual(out.shape, (2, 32)) - - -@require_torch -@require_torch -class FineGrainedGroupedLinearSwapTest(unittest.TestCase): - def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): - """`replace_with_finegrained_layer` picks `FineGrainedGroupedLinear` for a block-diagonal - linear — a plain one would collapse the groups into a single giant matmul and return the - wrong output dim. The branch keys on `n_groups`, the attribute the swap consumes: a class - NAME substring would miss a grouped linear named anything else, and would match (then - `AttributeError` on) one that carries its group count under another name. - """ - from types import SimpleNamespace - - import torch.nn as nn - - from transformers.integrations.finegrained import FineGrainedGroupedLinear, replace_with_finegrained_layer - from transformers.utils.quantization_config import FineGrainedConfig - - class OddlyNamedBlockDiagonal(nn.Linear): - def __init__(self): - super().__init__(8, 16, bias=False) - self.n_groups = 2 - - model = nn.Module() - model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) - model.grouped = OddlyNamedBlockDiagonal() - model.plain = nn.Linear(8, 16, bias=False) - - with torch.device("meta"): - replace_with_finegrained_layer(model, None, FineGrainedConfig(quant_method="fp8")) - self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) - self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) - - -class FineGrainedEmbeddingTest(unittest.TestCase): - """A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, - rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`.""" - - def test_lookup_rescales_the_gathered_rows(self): - torch.manual_seed(0) - table = torch.randn(64, 16) - scale = table.abs().max() / 448.0 - embedding = fg.FineGrainedEmbedding(64, 16) - embedding.weight = torch.nn.Parameter((table / scale).to(torch.float8_e4m3fn), requires_grad=False) - embedding.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16).reshape(1), requires_grad=False) - ids = torch.tensor([[3, 7], [60, 0]]) - out = embedding(ids) - self.assertEqual(out.dtype, torch.bfloat16) - expected = (table / scale).to(torch.float8_e4m3fn)[ids].to(torch.bfloat16) * scale.to(torch.bfloat16) - torch.testing.assert_close(out, expected) - - def test_replacement_follows_the_patterns_and_the_skip_list(self): - model = torch.nn.Module() - model.ngram_embedding = torch.nn.Embedding(8, 4) - model.embed_tokens = torch.nn.Embedding(8, 4) - model.other = torch.nn.Embedding(8, 4) - fg.replace_with_finegrained_embedding(model, ["ngram_embedding", "other"], modules_to_not_convert=["other"]) - self.assertIsInstance(model.ngram_embedding, fg.FineGrainedEmbedding) - self.assertTrue(model.ngram_embedding._hf_quantized_needs_local_tp) - self.assertIs(type(model.embed_tokens), torch.nn.Embedding) - self.assertIs(type(model.other), torch.nn.Embedding) - - -@require_torch -class FineGrainedExpertsMarshallingTest(unittest.TestCase): - def _experts(self, **kw): - cfg = _Cfg() - weight_format = kw.pop("weight_format", "fp8") - m = FineGrainedExperts(cfg, block_size=(4, 4), weight_format=weight_format, **kw) - for name, p in m.named_parameters(): - if p is None: - continue - if p.dtype == torch.float8_e4m3fn: - p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) - elif p.dtype.is_floating_point: - p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) - return m - - def _route(self, tokens=6): - hs = torch.randn(tokens, 64, dtype=torch.bfloat16) - idx = torch.randint(0, 4, (tokens, 2)) - wts = torch.rand(tokens, 2) - return hs, idx, wts - - def test_eager_linear_threads_bias_global_and_activation_format(self): - """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: - the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" - kernel, rec = _fake_bundle() - m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") - p1, p2, p3, p4 = _loaded(kernel) - linear = mock.Mock(wraps=fg.finegrained_linear) - with ( - p1, - p2, - p3, - p4, - mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), - mock.patch.object(fg, "finegrained_linear", linear), - ): - m(*self._route()) - self.assertTrue(linear.call_args_list) - for call in linear.call_args_list: - self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) - # this expert's NVFP4 weight global: one value for the down, and (1, 2) for the - # gate|up stack, whose halves modelopt calibrates separately - self.assertIn(tuple(call.kwargs["weight_global_scale"].shape), {(), (1, 2)}) - self.assertEqual(call.kwargs["activation_format"], "bf16") - # weight-only: the activations stay bf16, so no activation global is quantized against - self.assertIsNone(call.kwargs["input_global_scale"]) - for call in rec.calls["matmul_2d"]: - self.assertIsNotNone(call.kwargs["b_global_scale"]) - self.assertEqual(call.kwargs["activation_format"], "bf16") - - def test_batched_marshalling(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - hidden, top_k_index, top_k_weights = call.args - self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers - self.assertEqual(top_k_index.shape, (6, 2)) - self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) - # one scale tensor per projection: the affine Parameter, or the swizzled cache when the - # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) - self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) - self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) - self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name - self.assertIs(call.kwargs["gate"], True) - self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format - self.assertEqual(out.shape, (6, 64)) - - def test_grouped_marshalling(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_grouped"] - self.assertIs(call.kwargs["down_proj"], m.down_proj) - self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) - self.assertNotIn("moe_fused_batched", rec.calls) - - def test_activation_format_passes_through(self): - kernel, rec = _fake_bundle() - for activation_format in (None, "bf16", "mxfp8"): - m = self._experts(has_gate=True, activation_format=activation_format) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) - - def test_nvfp4_experts_thread_per_expert_globals(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, weight_format="nvfp4") - # the format table must resolve the ATTRIBUTE the forwards gate on, not just the - # param allocation — a None here silently drops the global at every forward - self.assertIsNotNone(m.global_scale_dtype) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) - self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) - # one weight global per expert on both projections: a gate|up stack calibrated per half - # is merged to that at load, so nothing downstream carries the pair - self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) - self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) - # the calibrated activation globals (the checkpoint's `input_scale`): one value for the - # gate_up, whose rows are quantized once before routing, and one per expert for the down, - # whose rows the gate_up epilogue requantizes per expert - self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) - self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) - self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) - self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) - - def test_weight_only_experts_hold_no_activation_global(self): - """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation - quant for a calibrated global to normalize: the module never allocates one, and the ops - would refuse a two-level path without it.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_grouped"] - self.assertIsNone(m.gate_up_proj_input_global_scale) - self.assertIsNone(m.down_proj_input_global_scale) - self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) - self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) - # the WEIGHT globals are the format's own second level and stay - self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) - - def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): - """A model whose experts norm the down output before the routing weights: - the swap carries the norm over, the kernel chain runs it on the routed rows, and the - eager loop applies it per expert application — the same place the reference forwards do. - A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) - - class _Recording(torch.nn.Module): - def __init__(self): - super().__init__() - self.rows = 0 - - def forward(self, x): - self.rows += x.shape[0] - return x - - m.post_expert_norm = _Recording() - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - m(*self._route()) - self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") - - def test_a_post_expert_norm_needs_nothing_but_the_norm(self): - """A class that declares a post-expert norm only has to HOLD it: every backend calls - `post_expert_norm` directly, so there is no per-class hook to define alongside it.""" - from transformers.integrations.moe import use_experts_implementation - - @use_experts_implementation(has_post_expert_norm=True) - class _JustTheNorm(torch.nn.Module): - def forward(self, hidden_states, top_k_index, top_k_weights): - return hidden_states - - self.assertFalse(hasattr(_JustTheNorm, "_apply_post_norm")) - - def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): - """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels - when they implement it, with the weight and epsilon the fused form needs, so the chain - folds it into its reduce instead of calling the module.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True - p1, p2, p3, p4 = _loaded(kernel) - for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): - m.post_expert_norm_name = name - with p1, p2, p3, p4: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - call = rec.calls["moe_fused_grouped"][-1] - if fused: - self.assertEqual(call.kwargs["post_expert_norm"], name) - self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) - self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) - else: - self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) - self.assertIsNone(call.kwargs["post_expert_norm_weight"]) - - def test_biases_ride_the_kernel_chain(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, has_bias=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) - self.assertTrue(torch.isfinite(out).all()) - (call,) = rec.calls["moe_fused_batched"] - # the kernels add gate_up's bias before the gated split and down's before the - # routing-weight multiply; nothing is added host-side - self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) - self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) - - def test_static_activation_scheme_reaches_the_fused_chain(self): - """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a - MoE calibrates each expert separately, so the module holds one per expert. Both fused - chains take them per projection and the kernels apply each tile's own expert scale in - register — a host pre-quant could not, since top-k routes one row to several experts - whose scales differ. A dynamic module holds no such slot and passes `None`.""" - static, dynamic = ( - FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) - for scheme in ("static", "dynamic") - ) - held = static.gate_up_proj_activation_scale - self.assertIsNotNone(held, "the calibrated scale the chain needs") - self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") - # a dynamic module holds the slot empty, like every other optional parameter - self.assertIsNone(dynamic.gate_up_proj_activation_scale) - - kernel = mock.Mock() - kernel.get_supported_act_fns.return_value = ("silu",) - kernel.get_supported_norms.return_value = () - for module, calibrated in ((static, True), (dynamic, False)): - operands = fg._moe_operands(kernel, module) - for projection in ("gate_up_proj", "down_proj"): - with self.subTest(scheme=module.activation_scheme, projection=projection): - scale = operands[f"{projection}_activation_scale"] - self.assertEqual(scale is not None, calibrated) - if calibrated: - self.assertEqual(scale.numel(), _Cfg.num_local_experts) - - def test_unfusable_act_fn_is_passed_as_the_module_glu(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, has_bias=True) - m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - # the kernels run the module's own GLU on the host between the two GEMMs — a new - # activation never waits for a kernel release; the bias still rides the GEMM - self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality - self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) - - -@require_torch -class FineGrainedScaleDtypeTest(unittest.TestCase): - """A scale's dtype is the format's, never the ambient default. - - `from_pretrained` sets the default dtype to the checkpoint's for the duration of model - construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it - as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches - that, whatever the format decides each scale should be. - """ - - def _dtypes(self, build, default): - previous = torch.get_default_dtype() - torch.set_default_dtype(default) - try: - with torch.device("meta"): - module = build() - return {name: p.dtype for name, p in module.named_parameters() if p is not None} - finally: - torch.set_default_dtype(previous) - - def _assert_default_dtype_invariant(self, build): - under_fp32 = self._dtypes(build, torch.float32) - under_bf16 = self._dtypes(build, torch.bfloat16) - self.assertTrue(under_fp32) - # the weight and bias DO follow the model dtype; every scale beside them must not - scales = {name for name in under_fp32 if "scale" in name} - self.assertTrue(scales) - for name in sorted(scales): - self.assertEqual(under_bf16[name], under_fp32[name], f"{name} followed the default dtype") - - def test_linear_scales_ignore_the_default_dtype(self): - for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): - with self.subTest(weight_format=weight_format, activation_scheme=scheme): - self._assert_default_dtype_invariant( - lambda f=weight_format, s=scheme: FineGrainedLinear( - in_features=64, out_features=32, block_size=(4, 4), weight_format=f, activation_scheme=s - ) - ) - - def test_experts_scales_ignore_the_default_dtype(self): - for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): - with self.subTest(weight_format=weight_format, activation_scheme=scheme): - self._assert_default_dtype_invariant( - lambda f=weight_format, s=scheme: FineGrainedExperts( - _Cfg(), block_size=(4, 4), weight_format=f, activation_scheme=s - ) - ) - - -class FineGrainedFusedNormGateTest(unittest.TestCase): - """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. - - Fusing runs the norm inside the launch that produced the rows, so an all-reduce that has to - land between the two has nowhere to go — and under intra-expert TP those rows are a partial - sum over the sharded intermediate, so the fused norm would normalize a fraction of the value. - The wrapped module does the reduce in its own forward, so the fallback is the correct path. - """ - - class _Norm(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(4)) - self.eps = 1e-6 - - def _operands(self, *, input_reduce: bool): - norm = self._Norm() - if input_reduce: - norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks - - module = mock.Mock() - module.has_gate = True - module.act_fn_name = "silu" - module.swiglu_alpha = None - module.has_post_expert_norm = True - module.post_expert_norm = norm - module.post_expert_norm_name = "input_scaled_rms_norm" - kernel = mock.Mock() - kernel.get_supported_act_fns.return_value = ("silu",) - kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) - return fg._moe_operands(kernel, module), module - - def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): - operands, _ = self._operands(input_reduce=False) - self.assertEqual(operands["post_expert_norm"], "input_scaled_rms_norm") - self.assertIsNotNone(operands["post_expert_norm_weight"], "the fused form needs the norm's weight") - - def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): - operands, module = self._operands(input_reduce=True) - self.assertIs( - operands["post_expert_norm"], - module.post_expert_norm, - "fused past a collective: the kernel would normalize this rank's partial sum", - ) - self.assertIsNone(operands["post_expert_norm_weight"]) - - -@require_torch -class FrozenFp8ShimTest(unittest.TestCase): - def test_frozen_module_warns_and_is_self_contained(self): - import importlib - - import transformers.integrations.finegrained_fp8 as frozen - from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer - from transformers.utils.quantization_config import FineGrainedFP8Config - - # the whole FILE is deprecated, so importing anything out of it is what warns. Safe to - # do at module scope only because `_LazyModule` never executes these unless asked: a - # plain `import transformers` does not reach them. - with warnings.catch_warnings(record=True) as raised: - warnings.simplefilter("always") - importlib.reload(frozen) - self.assertTrue( - [w for w in raised if issubclass(w.category, DeprecationWarning) and "frozen" in str(w.message)], - [str(w.message) for w in raised], - ) - # the quantizer is still constructible — frozen means no new recipes, not removed - FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) - # distinct machinery: the frozen classes are not the live ones - self.assertIsNot(frozen.FP8Linear, FineGrainedLinear) - - -@require_torch -class FineGrainedValidateEnvironmentTest(unittest.TestCase): - """`validate_environment` on the LIVE quantizer. The only tests that covered a - `validate_environment` sat on `FineGrainedFP8HfQuantizer`, which no `from_pretrained` can - reach any more (`AUTO_QUANTIZER_MAPPING["fp8"]` is `FineGrainedHfQuantizer`).""" - - @contextmanager - def _no_accelerator(self): - with ExitStack() as stack: - stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) - stack.enter_context( - mock.patch("transformers.quantizers.finegrained.base.is_torch_xpu_available", return_value=False) - ) - yield - - def _quantizer(self, pre_quantized, **cfg_kwargs): - # through the registry, so the test sees the class a `from_pretrained` would pick - from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING - from transformers.utils.quantization_config import FineGrainedConfig - - config = FineGrainedConfig(**cfg_kwargs) - method = getattr(config.quant_method, "value", config.quant_method) - quantizer = AUTO_QUANTIZER_MAPPING[method](config) - quantizer.pre_quantized = pre_quantized - return quantizer - - def test_no_accelerator_refuses_to_quantize(self): - # quantizing a bf16 checkpoint needs the hardware; there is nothing to fall back to - quantizer = self._quantizer(pre_quantized=False) - with self._no_accelerator(), self.assertRaises(RuntimeError): - quantizer.validate_environment() - - def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): - quantizer = self._quantizer(pre_quantized=True) - with self._no_accelerator(): - quantizer.validate_environment() - self.assertTrue(quantizer.quantization_config.dequantize) - - def test_nvfp4_refuses_to_dequantize(self): - # two-level scales have no slot in the dequantize chain, so the CPU fallback that rescues - # block-FP8 must refuse here rather than hand back a weight scaled by `1 / global` - for pre_quantized, kwargs in ((True, {}), (False, {"dequantize": True})): - with self.subTest(pre_quantized=pre_quantized): - quantizer = self._quantizer(pre_quantized=pre_quantized, quant_method="nvfp4", **kwargs) - with self._no_accelerator(), self.assertRaises(NotImplementedError): - quantizer.validate_environment() - - -class FineGrainedParallelPlanTest(unittest.TestCase): - """`update_tp_plan`: a model's plan as a quantized experts module needs it.""" - - BASE_EP = { - "layers.*.mlp.gate": "ep_router", - "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.self_attn.q_proj": "colwise", - } - - def _planned(self, impl=None, raw=None): - from types import SimpleNamespace - - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - quantizer = FineGrainedHfQuantizer(FineGrainedConfig()) - config = SimpleNamespace(base_model_tp_plan=dict(raw or self.BASE_EP), _experts_implementation=impl) - quantizer.update_tp_plan(config) - return config.base_model_tp_plan - - def test_every_companion_shards_with_its_expert_weight(self): - """The plan matcher keys on the exact parameter name and falls back only to the owning - module, whose entry shards nothing — so a scale with no entry of its own stays whole while - its weight is this rank's expert slice, and every rank but the first reads another rank's - scales.""" - from transformers.distributed.tensor_parallel import _get_parameter_tp_plan - - plan = self._planned() - style = lambda name: _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) # noqa: E731 - - for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_weight_global_scale"): - self.assertEqual(style(name), "grouped_gemm", name) - for name in ("down_proj_scale_inv", "down_proj_bias", "down_proj_input_global_scale"): - self.assertEqual(style(name), "grouped_gemm", name) - # one per-tensor value for the pre-routing hidden states: replicated, so no entry of its own - self.assertEqual(style("gate_up_proj_input_global_scale"), "moe_tp_experts") - self.assertEqual(plan["layers.*.self_attn.q_proj"], "colwise") - - def test_intra_expert_entries_follow_the_projection_axis(self): - """Under intra-expert TP the experts stay whole and the projection's own axis splits, so - the scale grid must follow its weight while the per-expert globals stay replicated. And - `packed_colwise` splits the output axis as two packed halves — the `[gate; up]` STACK — - where a quantized module holds interleaved rows, so the weight's own style changes too.""" - from transformers.distributed.tensor_parallel import _get_parameter_tp_plan - - base = { - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", - } - plan = self._planned(raw=base) - style = lambda n: _get_parameter_tp_plan(f"layers.3.mlp.experts.{n}", plan) # noqa: E731 - - # the weight keeps the style the MODEL declared — `packed_colwise` is the strided split a - # concatenated `[gate; up]` checkpoint needs, since the interleave into the kernels' row - # order runs on each rank's shard, after sharding. The companions follow it. - self.assertEqual(style("gate_up_proj"), "packed_colwise") - for name in ("gate_up_proj_scale_inv", "gate_up_proj_bias"): - self.assertEqual(style(name), "moe_experts_packed_colwise", name) - # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike - self.assertEqual(style("down_proj"), "rowwise") - self.assertEqual(style("down_proj_scale_inv"), "moe_experts_rowwise") - # TP does not split experts, so anything indexed per expert stays whole - for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): - self.assertEqual(style(name), "moe_tp_experts", name) - - # a natively INTERLEAVED layout declares `colwise`, and its companions split contiguously - inter = dict(base, **{"layers.*.mlp.experts.gate_up_proj": "colwise"}) - plan_i = self._planned(raw=inter) - self.assertEqual(_get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj", plan_i), "colwise") - self.assertEqual( - _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), - "moe_experts_colwise", - ) - - def test_dense_projections_that_share_a_name_are_left_alone(self): - """`mlp.up_proj` and `mlp.shared_experts.down_proj` end in the same words as the stacked - expert projections but are plain 2-D linears that colwise/rowwise already shard right. - Rewriting them to an expert-axis style sharded a dim they do not have — a CUDA illegal - memory access in the dense FP8 matmul, seen on DeepSeek-V3 under TP.""" - from transformers.distributed.tensor_parallel import _get_parameter_tp_plan - - plan = self._planned( - raw={ - "layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - "layers.*.mlp.shared_experts.up_proj": "colwise", - "layers.*.mlp.shared_experts.down_proj": "rowwise", - } - ) - for dense, expected in ( - ("layers.*.mlp.up_proj", "colwise"), - ("layers.*.mlp.down_proj", "rowwise"), - ("layers.*.mlp.shared_experts.up_proj", "colwise"), - ("layers.*.mlp.shared_experts.down_proj", "rowwise"), - ): - self.assertEqual(plan[dense], expected, dense) - # and no scale entry was invented for a dense linear, whose scale `rowwise`/`colwise` - # already reaches through the module fallback - strays = [k for k in plan if k.endswith("_scale_inv") and ".experts." not in k] - self.assertEqual(strays, [], f"invented dense scale entries: {strays}") - # the stacked experts still get theirs - self.assertEqual( - _get_parameter_tp_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), - "moe_experts_packed_colwise", - ) - - def test_a_multimodal_models_expert_plans_are_reached(self): - """A multimodal model keeps its experts' plans on a SUB-config; the config the quantizer - is handed carries a few projector entries and often no `base_model_ep_plan` at all. - Rewriting only what we were handed adds no companion anywhere, while the weights still - shard from the sub-config's own plan — so the scale stays whole against a sharded weight, - which is a wrong answer rather than a crash.""" - from transformers import Glm4vMoeConfig - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - config = Glm4vMoeConfig() - self.assertFalse( - getattr(config, "base_model_ep_plan", None), - "this model no longer nests its plans, so it cannot guard the sub-config walk", - ) - before = dict(getattr(config.text_config, "base_model_ep_plan", None) or {}) - FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) - after = getattr(config.text_config, "base_model_ep_plan", None) or {} - - companions = { - k - for k in after - if k.endswith(("_scale_inv", "_bias", "_weight_global_scale", "_input_global_scale", "_activation_scale")) - } - self.assertTrue(companions, "no companion reached the sub-config's plan") - self.assertEqual(set(before), set(after) - companions, "the model's own entries were disturbed") - - def test_dequantize_folds_every_quantized_key_to_full_precision(self): - """`dequantize=True` is the escape hatch for hardware that cannot serve the format. It - must claim BOTH checkpoint key layouts — the `{proj}_blocks` + `{proj}_scales` pair - GPT-OSS ships and the `weight` + `weight_scale_inv` pair everything else does — and - land on the plain weight, since a module that was never swapped has no scale slot to - write into.""" - import re - - from transformers.utils.quantization_config import FineGrainedConfig - - for method, sources in ( - ("mxfp4", ("experts.gate_up_proj_blocks", "experts.gate_up_proj_scales")), - ("fp8", ("mlp.down_proj.weight", "mlp.down_proj.weight_scale_inv")), - ): - quantizer = _quantizer_for(FineGrainedConfig(quant_method=method, dequantize=True)) - quantizer.pre_quantized = True - converters = quantizer.get_weight_conversions() - with self.subTest(quant_method=method): - self.assertTrue(converters, f"{method}: dequantize produced no converter") - for source in sources: - self.assertTrue( - any( - re.search(str(pattern), source) - for c in converters - for pattern in ( - c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns] - ) - ), - f"{method}: nothing dequantizes {source}", - ) - - def test_the_impl_rewrites_the_layer_kinds(self): - megamoe = self._planned("deepgemm_megamoe") - self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") - self.assertEqual(megamoe["layers.*.mlp.experts"], "megamoe_experts") - - def test_a_models_own_plan_survives_the_update(self): - """`update_tp_plan` ADDS companions; it must never drop what the model ships. - - The config here is a REAL one, which the `SimpleNamespace` the tests above use cannot - substitute for: a branch keyed on the config CLASS is invisible to a fake. A hardcoded - `"Qwen3" in config.__class__.__name__` override used to assign a dense plan over - `base_model_tp_plan`, dropping `q_norm`/`k_norm` on Qwen3 and — because the check is a - substring — the whole expert plan on Qwen3-MoE. - """ - from transformers import Qwen3Config, Qwen3MoeConfig, Qwen3VLMoeConfig - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - for config_cls in (Qwen3Config, Qwen3MoeConfig): - with self.subTest(config=config_cls.__name__): - config = config_cls() - shipped = dict(config.base_model_tp_plan) - self.assertTrue(shipped, f"{config_cls.__name__} ships no plan to preserve") - FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) - for key, style in shipped.items(): - self.assertEqual(config.base_model_tp_plan.get(key), style, key) - - # and a multimodal one keeps its plan on the SUB-config: the outer carries none, so - # nothing may be invented there. 25 config classes contain "Qwen3" — a name-substring - # branch reaches every one of them, including these. - multimodal = Qwen3VLMoeConfig() - FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(multimodal) - self.assertFalse(getattr(multimodal, "base_model_tp_plan", None) or {}) - self.assertTrue(multimodal.text_config.base_model_tp_plan) - - def test_a_real_moe_models_experts_gain_their_companions(self): - """The other half of the same invariant: preserving the model's entries is what lets the - companion pass find `moe_tp_experts` and give the expert scales an entry of their own.""" - from transformers import Qwen3MoeConfig - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - config = Qwen3MoeConfig() - FineGrainedHfQuantizer(FineGrainedConfig()).update_tp_plan(config) - plan = config.base_model_tp_plan - self.assertEqual(plan["layers.*.mlp.experts"], "moe_tp_experts") - for companion in ("gate_up_proj_scale_inv", "down_proj_scale_inv", "gate_up_proj_bias"): - self.assertIn(f"layers.*.mlp.experts.{companion}", plan, companion) - - -class FineGrainedGroupsTest(unittest.TestCase): - """A quant config names a format per module SUBSET, because a checkpoint can be more than one.""" - - def _config(self, **kwargs): - from transformers.utils.quantization_config import FineGrainedConfig - - return FineGrainedConfig(**kwargs) - - def test_a_flat_config_is_one_group_over_everything(self): - groups = self._config(quant_method="mxfp8").groups - self.assertEqual({name: group.quant_method for name, group in groups.items()}, {"default": "mxfp8"}) - - def test_a_legacy_expert_dtype_becomes_the_two_groups_it_meant(self): - """DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL - config, because a flat config cannot say "the experts differ". `shared_experts` is DENSE - despite the name, which is why the pattern matches a path segment and not a substring.""" - from transformers.utils.quantization_config import groups_with_expert_dtype - - config = self._config(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0") - config.groups = groups_with_expert_dtype(config.groups, "fp4") - for module, expected in ( - ("model.layers.3.mlp.experts", "mxfp4"), - ("model.layers.3.mlp.experts.gate_up_proj", "mxfp4"), - ("model.layers.3.self_attn.q_proj", "fp8"), - ("model.layers.3.mlp.shared_experts.up_proj", "fp8"), - ): - self.assertEqual(config.group_for(module).quant_method, expected, module) - - def test_resolution_survives_a_round_trip_through_config_json(self): - """`to_json_string` SORTS keys, so a config read back has lost the order it was written - in. Resolution must not depend on it: the catch-all sorts first here and would swallow - every module if targeted groups were not preferred outright.""" - import json - - from transformers.utils.quantization_config import FineGrainedConfig, FineGrainedGroup - - config = FineGrainedConfig( - groups={ - "experts": FineGrainedGroup(quant_method="nvfp4", targets=[r"\.experts($|\.)"]), - "dense": FineGrainedGroup(quant_method="fp8", weight_block_size=(128, 128)), - } - ) - back = FineGrainedConfig.from_dict(json.loads(config.to_json_string())) - self.assertEqual(back.group_for("model.layers.0.mlp.experts").quant_method, "nvfp4") - self.assertEqual(back.group_for("model.layers.0.self_attn.q_proj").quant_method, "fp8") - # a list survives JSON where a tuple does not, and the block size is compared as a tuple - self.assertEqual(back.groups["dense"].weight_block_size, (128, 128)) - - def test_two_targeted_groups_claiming_one_module_is_an_error(self): - from transformers.utils.quantization_config import FineGrainedGroup - - config = self._config( - groups={ - "a": FineGrainedGroup(quant_method="fp8", targets=[r"\.experts"]), - "b": FineGrainedGroup(quant_method="nvfp4", targets=[r"mlp\."]), - } - ) - with self.assertRaises(ValueError): - config.group_for("model.layers.0.mlp.experts") - - def test_a_producers_config_groups_are_normalized(self): - """modelopt describes a format by its parameters; we name it. GLM-5.2-NVFP4 ships exactly - this, and `targets: ["Linear"]` is its spelling of the catch-all.""" - config = self._config( - quant_method="modelopt", - quant_algo="NVFP4", - config_groups={ - "group_0": { - "weights": {"num_bits": 4, "type": "float", "group_size": 16}, - "input_activations": {"num_bits": 4, "type": "float", "group_size": 16, "dynamic": False}, - "targets": ["Linear"], - } - }, - ) - group = config.group_for("model.layers.0.self_attn.q_proj") - self.assertEqual((group.quant_method, group.activation_format), ("nvfp4", "nvfp4")) - self.assertEqual(group.activation_scheme, "static") # `dynamic: False` - - def test_a_format_we_do_not_serve_falls_back_rather_than_guessing(self): - """Half a translation would quantize modules by a rule the producer did not write.""" - config = self._config( - quant_method="modelopt", - quant_algo="NVFP4", - config_groups={ - "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} - }, - ) - self.assertEqual(list(config.groups), ["default"]) - - -class SubtreePatternTest(unittest.TestCase): - def test_a_checkpoints_glob_skip_list_is_bounded(self): - """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any - character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" - from transformers.quantizers.quantizers_utils import should_convert_module - from transformers.utils.quantization_config import FineGrainedConfig - - # as a modelopt checkpoint ships them, through the config that normalizes them - skip = FineGrainedConfig( - quant_method="modelopt", - quant_algo="NVFP4", - ignore=["model.layers.0*", "model.layers.1.*", "self_attn"], - ).modules_to_not_convert - for name in ("model.layers.0.mlp", "model.layers.1.mlp.down_proj", "model.layers.2.self_attn.q_proj"): - self.assertFalse(should_convert_module(name, skip), name) - for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): - self.assertTrue(should_convert_module(name, skip), name) - - -@require_torch -class FineGrainedMxfp4ConverterTest(unittest.TestCase): - """The gpt-oss checkpoint layout through the real ConversionOps: blocks (E, N, K/32, 16) - uint8 low-nibble-first E2M1 + biased-127 exponent scales, gate|up rows INTERLEAVED — - deserialized output must dequantize identically to the reference LUT unpack.""" - - FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] - - def _reference_dequant(self, blocks, scales): - # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) - lut = torch.tensor(self.FP4_VALUES) - lo = lut[(blocks & 0xF).long()] - hi = lut[(blocks >> 4).long()] - vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) - exp = (scales.long() - 127).unsqueeze(-1) - return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) - - def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): - from transformers.integrations.finegrained.conversions import ( - FineGrainedPackedBlocks, - FineGrainedScaleContainer, - ) - - torch.manual_seed(0) - E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows - blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) - scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) - - weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ - "gate_up_proj" - ] - self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) - back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] - self.assertTrue(torch.equal(back, blocks)) - - # the scales are the container op's job: a module holding e8m0 receives the exponent bytes as is - module = fg.FineGrainedExperts.__new__(fg.FineGrainedExperts) - torch.nn.Module.__init__(module) - module.gate_up_proj_scale_inv = torch.nn.Parameter( - torch.zeros(E, N, K // 32, dtype=torch.float8_e8m0fnu), requires_grad=False - ) - model = torch.nn.Module() - model.experts = module - scale_inv = FineGrainedScaleContainer().convert( - {"gate_up_proj_scales$": scales}, - model=model, - full_layer_name="experts.gate_up_proj_scale_inv", - target_patterns=["gate_up_proj_scale_inv"], - )["gate_up_proj_scale_inv"] - self.assertEqual(scale_inv.dtype, torch.float8_e8m0fnu) - - ref = self._reference_dequant(blocks, scales) # rows pass through interleaved - # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) - lut = torch.tensor(self.FP4_VALUES) - w_u8 = weight.view(torch.uint8) - lo = lut[(w_u8 & 0xF).long()] - hi = lut[(w_u8 >> 4).long()] - vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) - exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) - got = vals * torch.pow(torch.tensor(2.0), exp) - torch.testing.assert_close(got, ref, rtol=0, atol=0) - - -@require_torch -class FineGrainedDeepGemmDispatchTest(unittest.TestCase): - """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled - scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine - and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 - scales have no swizzled layout, so the first never fires for the shape the second catches.""" - - def _routed_to(self, *, sm100, scale_ndim): - """Which backend a block-FP8 linear reaches, given arch and scale layout.""" - kernel, _ = _fake_bundle() - w = torch.randn(32, 64).to(torch.float8_e4m3fn) - s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) - p1, p2, p3, p4 = _loaded(kernel) - with ( - p1, - p2, - p3, - p4, - mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), - mock.patch.object(deepgemm, "is_sm100", return_value=sm100), - mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, - ): - fg.finegrained_linear(torch.randn(4, 64, dtype=torch.bfloat16), w, s, block_size=[128, 128]) - return "deepgemm" if dg.called else "triton" - - def test_sm100_never_prefers_deepgemm(self): - self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") - - def test_pre_sm100_still_uses_deepgemm(self): - self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") - - def test_swizzled_scales_never_reach_deepgemm(self): - """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" - self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") - - def test_deepgemm_experts_refuse_swizzled_scales(self): - """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM - experts backend afterwards must fail loudly rather than read the permuted buffer as affine.""" - from transformers.integrations.deepgemm import ( - deepgemm_fp8_fp4_experts_forward, - deepgemm_fp8_fp4_megamoe_experts_forward, - ) - - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" - with ( - mock.patch.object(fg, "is_sm100", return_value=True), - mock.patch("torch.cuda.is_available", return_value=True), - ): - experts = fg.FineGrainedExperts(cfg, weight_format="mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) - hs = torch.randn(2, 256, dtype=torch.bfloat16) - idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) - for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): - with self.assertRaisesRegex(RuntimeError, "held swizzled"): - forward(experts, hs, idx, wts) - - -@require_torch -class FineGrainedScaleLayoutTest(unittest.TestCase): - """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 - scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the - scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's - `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" - - E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks - - def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E - cfg._experts_implementation = impl - with ( - mock.patch.object(fg, "is_sm100", return_value=sm100), - mock.patch("torch.cuda.is_available", return_value=True), - ): - experts = fg.FineGrainedExperts( - cfg, block_size=(128, 128), weight_format=weight_format, activation_format=activation_format - ) - model = torch.nn.Module() - model.config = cfg - model.experts = experts - return model, experts - - def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): - _, experts = self._experts("mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) - self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) - self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) - self.assertTrue(experts.holds_interleaved_gate_up) - # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a - # float32 grid there; GPT-OSS's config has no scale_fmt at all) - _, experts = self._experts("mxfp4", activation_format="bf16") - self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) - - def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): - for kw in ( - {"weight_format": "mxfp8", "sm100": False}, - {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely - {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales - {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, - ): - with self.subTest(**kw): - _, experts = self._experts(**kw) - self.assertEqual( - experts.gate_up_proj_scale_inv.shape, - (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), - ) - # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block - _, experts = self._experts("fp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) - with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): - _, experts = self._experts("mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) - - def test_megamoe_holds_gate_up_stacked(self): - _, experts = self._experts("fp8", impl="deepgemm_megamoe") - self.assertFalse(experts.holds_interleaved_gate_up) - - def _op_kernel(self): - kernel, _ = _fake_bundle() - kernel.swizzle_mx_scales = mock.Mock( - side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) - ) - kernel.unswizzle_mx_scales = mock.Mock( - side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) - ) - return kernel - - def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): - from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales - - model, experts = self._experts("mxfp8") - kernel = self._op_kernel() - op = FineGrainedSwizzleScales(hf_quantizer=None) - grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) - weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - # the loader hands a single target under its pattern, the full name (with suffix) alongside - out = op.convert( - {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" - ) - self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) - # the same converter's weight passes through - out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") - self.assertIs(out["mlp.experts.gate_up_proj"], weight) - # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully - out = op.convert( - { - "experts.down_proj": weight, - "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), - }, - model=model, - full_layer_name="experts.down_proj", - ) - self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) - self.assertIs(out["experts.down_proj"], weight) - # a module that holds affine scales keeps them - affine_model, _ = self._experts("mxfp4", activation_format="bf16") - out = op.convert( - {"mlp.experts.gate_up_proj": grid}, - model=affine_model, - full_layer_name="experts.gate_up_proj_scale_inv", - ) - self.assertIs(out["mlp.experts.gate_up_proj"], grid) - self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) - - def test_scale_container_op_round_trips_every_container(self): - """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: - the op brings both into the held e8m0 (exact cast / same bytes), records the container on - the module, and its reverse restores that container on save; weights and native scales - pass both ways.""" - from transformers.integrations.finegrained.conversions import FineGrainedScaleContainer - - op = FineGrainedScaleContainer(hf_quantizer=None) - native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) - for container, shipped in ((torch.float32, native.float()), (torch.uint8, native.view(torch.uint8))): - with self.subTest(container=container): - model, experts = self._experts("mxfp8", sm100=False) # affine e8m0 scale Parameters - out = op.convert( - {"mlp.experts.gate_up_proj": shipped}, - model=model, - full_layer_name="experts.gate_up_proj_scale_inv", - ) - held = out["mlp.experts.gate_up_proj"] - self.assertEqual(held.dtype, torch.float8_e8m0fnu) - self.assertTrue(torch.equal(held.view(torch.uint8), native.view(torch.uint8))) - self.assertIs(experts.scale_container_dtype, container) - back = op.reverse_op.convert({"x": held}, model=model)["x"] - self.assertEqual(back.dtype, container) - self.assertTrue(torch.equal(back, shipped)) - weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) - self.assertIs(op.reverse_op.convert({"w": weight}, model=model)["w"], weight) - model, experts = self._experts("mxfp8", sm100=False) - self.assertIs( - op.convert({"mlp.experts.down_proj": native}, model=model, full_layer_name="experts.down_proj_scale_inv")[ - "mlp.experts.down_proj" - ], - native, - ) - self.assertIsNone(experts.scale_container_dtype) - self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays - - def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): - from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales - - kernel = self._op_kernel() - reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op - artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) - affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") - self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) - kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) - out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") - self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) - - def _quantizer(self, quant_method): - from transformers.utils.quantization_config import FineGrainedConfig - - return _quantizer_for(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions - - def _arch_converters(self): - from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter - - return [ - WeightConverter( - source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"], - target_patterns="mlp.experts.gate_up_proj", - operations=[MergeModulelist(dim=0), Concatenate(dim=1)], - ), - WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight", - target_patterns="mlp.experts.down_proj", - operations=[MergeModulelist(dim=0)], - ), - WeightConverter( - source_patterns=["mlp.fc1.weight", "mlp.fc2.weight"], - target_patterns="mlp.fc.weight", - operations=[Concatenate(dim=0)], - ), - ] - - def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): - from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained.conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, - FineGrainedSwizzleScales, - ) - - converters = self._quantizer("mxfp8")(self._arch_converters()) - by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} - gate_up_ops = by_target["mlp.experts.gate_up_proj"].operations - # stacked [gate; up] -> [g0, u0, ...], the scale into the held dtype, then the swizzle packs - # the final row order (it reads 1-byte scales) - self.assertIsInstance(gate_up_ops[-3], FineGrainedInterleaveGateUp) - self.assertFalse(gate_up_ops[-3].inverse) - self.assertIsInstance(gate_up_ops[-2], FineGrainedScaleContainer) - self.assertIsInstance(gate_up_ops[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["mlp.experts.down_proj"].operations[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["weight_scale_inv"].operations[0], FineGrainedScaleContainer) # dense linears - self.assertEqual(len(by_target["mlp.fc.weight"].operations), 1) # dense converters untouched - # keys arriving under the fused names get the same layout - self.assertIn("experts.gate_up_proj_bias", by_target) # patterns are stored anchor-stripped - self.assertIsInstance(by_target["experts.down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) - # ...and saving reverses in the opposite order: unswizzle first, then de-interleave - reverse_ops = by_target["mlp.experts.gate_up_proj"].reverse_transform().operations - self.assertIsInstance(reverse_ops[0], FineGrainedSwizzleScales) - self.assertTrue(reverse_ops[0].inverse) - self.assertIsInstance(reverse_ops[2], FineGrainedInterleaveGateUp) - self.assertTrue(reverse_ops[2].inverse) - - def test_blocks_scales_converters_carry_the_layout_ops(self): - from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained.conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedSwizzleScales, - ) - - converters = self._quantizer("mxfp4")([]) - by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} - ops = by_target["gate_up_proj"].operations - self.assertIsInstance(ops[-2], FineGrainedInterleaveGateUp) - self.assertIsInstance(ops[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["gate_up_proj_scale_inv"].operations[0], FineGrainedInterleaveGateUp) - self.assertIsInstance(by_target["down_proj"].operations[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) - - def test_catch_all_converters_deliver_the_full_parameter_name(self): - """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax - down_proj_scale_inv) goes through the real converter plumbing: the ops receive the tensor - under the source PATTERN and must hand back the target, which the loader expands to the - full name.""" - from transformers.core_model_loading import WeightConverter, rename_source_key - - model, _ = self._experts("mxfp8") - converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] - for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): - renamed, source_pattern = rename_source_key(key, [], converters) - self.assertEqual(renamed, key) - converter = next(c for c in converters if source_pattern in c.source_patterns) - tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) - converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) - kernel = self._op_kernel() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: - out = converter.convert(renamed, model=model) - self.assertEqual(list(out), [key]) - - def test_replacement_keeps_the_models_gate_up_convention(self): - """`replace_with_finegrained_layer` goes through `use_experts_implementation`, which stamps - the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must - survive it (it decides whether the loader interleaves).""" - from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM - from transformers.utils.quantization_config import FineGrainedConfig - - gpt_oss = GptOssConfig( - vocab_size=64, - hidden_size=64, - intermediate_size=32, - num_local_experts=2, - num_experts_per_tok=1, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - sliding_window=16, - ) - qwen = Qwen3MoeConfig( - vocab_size=64, - hidden_size=64, - intermediate_size=64, - moe_intermediate_size=32, - num_experts=2, - num_experts_per_tok=1, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - ) - for cls, cfg, quant_method, expected in ( - (GptOssForCausalLM, gpt_oss, "mxfp4", False), - (Qwen3MoeForCausalLM, qwen, "mxfp8", True), - ): - with self.subTest(model=cls.__name__), torch.device("meta"): - model = cls(cfg) - fg.replace_with_finegrained_layer( - model, modules_to_not_convert=[], quantization_config=FineGrainedConfig(quant_method=quant_method) - ) - experts = model.model.layers[0].mlp.experts - self.assertIsInstance(experts, fg.FineGrainedExperts) - self.assertIs(experts.is_concatenated, expected) - - def test_interleave_op_follows_the_model_and_the_experts_backend(self): - """The op reads the model's experts modules: a model whose own rows already come - interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself - and keeps the stack; every other case gets the interleaved order, and the reverse restores - the stack.""" - from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp - - stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] - op = FineGrainedInterleaveGateUp(hf_quantizer=None) - model, experts = self._experts("mxfp8") - out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] - torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) - back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] - torch.testing.assert_close(back, stacked) - experts.is_concatenated = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) - experts.is_concatenated = True - experts.holds_interleaved_gate_up = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) - - -@require_torch -class FineGrainedOnTheFlyQuantizeTest(unittest.TestCase): - """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run the - kernels' quantizers — covered with real kernels below), emitted in the module's layout.""" - - def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): - """`dequantize=True` on a GPT-OSS MXFP4 checkpoint used to be the mxfp4 integration's job. - The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation - the unquantized experts hold — so it is compared against that integration directly.""" - from transformers.core_model_loading import Transpose - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks - from transformers.integrations.mxfp4 import convert_moe_packed_tensors - - torch.manual_seed(0) - num_experts, rows, hidden = 2, 8, 64 - blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) - scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) - reference = convert_moe_packed_tensors(blocks.clone(), scales.clone(), dtype=torch.bfloat16) - - sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] - tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} - for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): - tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) - self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain - torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) - - def test_a_static_scheme_gets_its_activation_scales_written(self): - """A calibration-fed slot the checkpoint does not supply is still WRITTEN, as the identity. - - The loader materializes a missing key with `torch.empty_like` and `_init_weights` has no - branch for a scale, so a slot this op leaves out reaches the kernels as uninitialized - memory: a zero divides the activations by zero and a negative flips their sign, which - showed up as NaN logits from one fixture and not another, run to run. - """ - from transformers.integrations.finegrained.conversions import FineGrainedQuantize - - torch.manual_seed(0) - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 - model = torch.nn.Module() - model.experts = fg.FineGrainedExperts( - cfg, block_size=(128, 128), weight_format="fp8", activation_scheme="static" - ) - model.proj = fg.FineGrainedLinear( - in_features=256, - out_features=256, - block_size=(128, 128), - weight_format="fp8", - activation_scheme="static", - ) - op = FineGrainedQuantize(hf_quantizer=None) - for key, tensor, slot in ( - ("experts.gate_up_proj", torch.randn(4, 256, 256), "experts.gate_up_proj_activation_scale"), - ("proj.weight", torch.randn(256, 256), "proj.activation_scale"), - ): - with self.subTest(key=key): - out = op.convert({key: tensor}, model=model) - self.assertIn(slot, out, f"{key}: the static activation scale was not written") - held = model.get_parameter(slot) - self.assertEqual((out[slot].shape, out[slot].dtype), (held.shape, held.dtype)) - torch.testing.assert_close(out[slot], torch.ones_like(out[slot]), rtol=0, atol=0) - - def test_an_expert_bias_is_passed_through_not_quantized(self): - """A GPT-OSS expert bias is `(E, rows)` — 2-D, like a dense weight — so the rank guard - alone lets it through. It has no scale slot, and emitting one gives the loader a - `_bias_scale_inv` no module holds.""" - from transformers.integrations.finegrained.conversions import FineGrainedQuantize - - torch.manual_seed(0) - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 - model = torch.nn.Module() - model.experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", has_bias=True) - bias = torch.randn(4, 256) - out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.down_proj_bias": bias}, model=model) - self.assertEqual(list(out), ["experts.down_proj_bias"]) - self.assertIs(out["experts.down_proj_bias"], bias) - - def test_a_partial_trailing_block_is_padded_not_refused(self): - """DeepSeek-V3 ships `kv_a_proj_with_mqa` as `(576, 7168)` against a 128x128 block with a - `(5, 56)` scale grid, so the format rounds the grid UP and quantizes the short block on - its own values. Refusing the shape instead left the weight full precision.""" - from transformers.integrations.finegrained.conversions import FineGrainedQuantize - - torch.manual_seed(0) - for rows, cols in ((576, 256), (192, 256), (256, 256)): - with self.subTest(shape=(rows, cols)): - weight, scale = FineGrainedQuantize._quantize_block_fp8( - torch.randn(rows, cols), (128, 128), ue8m0=False - ) - self.assertEqual(weight.shape, (rows, cols)) - self.assertEqual(weight.dtype, torch.float8_e4m3fn) - self.assertEqual(scale.shape, (-(-rows // 128), -(-cols // 128))) - - def test_block_fp8_round_trips_within_its_floor(self): - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize - - torch.manual_seed(0) - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 - for scale_fmt in ("float", "ue8m0"): - with self.subTest(scale_fmt=scale_fmt), mock.patch.object(fg, "is_sm100", return_value=False): - experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8", scale_fmt=scale_fmt) - model = torch.nn.Module() - model.experts = experts - original = torch.randn(4, 256, 256) * 0.02 - out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) - weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] - self.assertEqual( - (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) - ) - self.assertEqual( - (scale.dtype, scale.shape), - (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), - ) - deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) - rel = ((deq - original).norm() / original.norm()).item() - self.assertLess(rel, 4e-2, f"{scale_fmt}: {rel:.3f}") - - -@require_torch -class FineGrainedModeloptConverterTest(unittest.TestCase): - """The ModelOpt (NVFP4) weight conversions. `MergeModulelist` is load-bearing beyond the - merge itself: `core_model_loading` stamps a source with its expert index ONLY when one is - present in the chain, and expert parallelism selects experts by that index. Without it every - rank collects all E globals and the forward asserts on the per-expert count.""" - - def _modelopt_conversions(self, **cfg_kwargs): - from transformers.quantizers.finegrained.nvfp4 import FineGrainedNvfp4HfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) - return FineGrainedNvfp4HfQuantizer(cfg).get_weight_conversions() - - def test_every_declared_target_is_a_slot_the_module_holds(self): - """A converter target that no module holds is a load failure, and the activation global is - exactly such a target under a weight-only run: the experts allocate one only when they - quantize activations. Both runs are checked against a real module's parameter names, since - the declarations are the only thing standing between the two.""" - cfg = _Cfg() - for activation_format, holds_global in ((None, True), ("bf16", False)): - experts = FineGrainedExperts( - cfg, block_size=(4, 4), weight_format="nvfp4", activation_format=activation_format - ) - slots = {name for name, _ in experts.named_parameters()} | set(experts._buffers) - self.assertEqual( - any("input_global_scale" in slot for slot in slots), - holds_global, - f"module with activation_format={activation_format!r} disagrees about the slot", - ) - targets = { - t.split("experts.")[-1].rstrip("$") - for c in self._modelopt_conversions(activation_format=activation_format) - for t in _targets(c) - } - for target in targets: - self.assertIn(target, slots, f"activation_format={activation_format!r} converts onto a missing slot") - - def test_global_scale_converters_carry_an_expert_index(self): - from transformers.core_model_loading import MergeModulelist - - # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked - # tensor per layer has no per-expert sources to stamp — it is sharded like the fused - # expert weight next to it - globals_converters = [ - c - for c in self._modelopt_conversions() - if any("global_scale" in t for t in _targets(c)) - and any(re.search(p, "model.layers.0.mlp.experts.7.gate_proj.weight_scale_2") for p in c.source_patterns) - ] - self.assertTrue(globals_converters, "no global-scale converter found") - for conv in globals_converters: - self.assertTrue( - any(isinstance(op, MergeModulelist) for op in conv.operations), - f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " - "experts and every rank would collect all of them", - ) - - def test_both_modelopt_scale_layouts_have_converters(self): - """modelopt ships the scales either as one tensor per expert per projection (GLM-5.2) or - as one already-stacked tensor per layer (the fused vLLM layout). Both are - converted; the block scale takes the module's layout ops either way.""" - sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} - for suffix in ("weight_scale", "weight_scale_2", "input_scale"): - per_expert = f"model.layers.0.mlp.experts.7.gate_proj.{suffix}" - fused = f"model.layers.0.mlp.experts.gate_up_proj_{suffix}" - self.assertTrue(any(re.search(p, per_expert) for p in sources), f"no per-expert converter for {suffix}") - self.assertTrue(any(re.search(p, fused) for p in sources), f"no fused converter for {suffix}") - # the per-expert patterns are REGEXES: unescaped dots matched `_` too, so - # `mlp.experts.*.up_proj.weight_scale_2` also claimed the fused key and the - # gate|up globals were flattened instead of folded to one per expert - claimed = [p for p in sources if re.search(p, fused) and re.search(p, per_expert)] - self.assertEqual(claimed, [], f"one pattern claims BOTH layouts for {suffix}: {claimed}") - - def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): - """Two converters for one key is not additive — the later one WINS and silently drops - whatever ops the first carried. That is how a fused modelopt checkpoint lost its packed - uint8 -> int8 view (a catch-all duplicated the fused converter) and how the per-expert - patterns, which are REGEXES whose unescaped dots also match `_`, claimed the fused - globals and flattened the gate|up pair instead of folding it.""" - import re - - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig - - for quant_kwargs in ( - {"quant_method": "fp8"}, - {"quant_method": "mxfp4"}, - {"quant_method": "modelopt", "quant_algo": "NVFP4"}, - ): - quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**quant_kwargs)) - quantizer.pre_quantized = True - converters = quantizer.update_weight_conversions([]) - for key in ( - "model.layers.3.mlp.experts.gate_up_proj", - "model.layers.3.mlp.experts.down_proj", - "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", - "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", - ): - claimed = [ - c - for c in converters - for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) - if re.search(str(pattern), key) - ] - with self.subTest(quant=quant_kwargs["quant_method"], key=key.rsplit(".", 1)[1]): - self.assertLessEqual( - len(claimed), - 1, - f"{len(claimed)} converters claim this key; the last one wins and drops " - f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", - ) - - def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): - """A static checkpoint calibrates one `input_scale` per quantized module. Nothing routed - those keys for plain FP8 — only NVFP4 read them, as its second-level global — so the two - shipped static models (Ministral-3 dense, Mistral-4 MoE) loaded with the registered - default of 1.0: every activation quantized against the wrong scale, silently. - - Each key must reach the slot its module holds, in that module's shape: one value on a - dense linear, one per expert on the stacked experts, and the gate|up pair reduced to one - per expert rather than flattened to 2E.""" - import re - - from transformers.utils.quantization_config import FineGrainedConfig - - quantizer = _quantizer_for(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) - quantizer.pre_quantized = True - converters = quantizer.update_weight_conversions([]) - - def claim(key): - return [ - c - for c in converters - for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) - if re.search(str(pattern), key) - ] - - for key, expected in ( - ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), - ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.7.up_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.7.down_proj.input_scale", "mlp.experts.down_proj_activation_scale"), - ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), - ): - with self.subTest(key=key.rsplit(".", 2)[-2]): - claimed = claim(key) - self.assertEqual(len(claimed), 1, f"{len(claimed)} converters claim {key}") - # a renaming's target carries a backreference; resolve it against the key - targets = [ - re.sub(str(claimed[0].source_patterns[0]), str(t), key) if "\\1" in str(t) else str(t) - for t in claimed[0]._original_target_patterns - ] - self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") - # a dynamic checkpoint has no such key to route, and a weight-only one ignores it - dynamic = _quantizer_for(FineGrainedConfig(quant_method="fp8")) - dynamic.pre_quantized = True - self.assertEqual(dynamic.get_weight_conversions(), []) - - def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): - """`FineGrainedInputScales` collapses the gate_up to ONE value for the NVFP4 global — the - global is a split of the block scale, so an inflated one is exact. A static scale IS the - quantization scale, with no block level to absorb it, so collapsing it would quantize - every expert against the largest one's range. It stays per expert.""" - import torch - - from transformers.integrations.finegrained.conversions import FineGrainedInputScales - - sources = { - "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), - "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), - } - target = "mlp.experts.gate_up_proj_activation_scale" - scale = FineGrainedInputScales().convert(dict(sources), full_layer_name=target)[target] - torch.testing.assert_close(scale, torch.tensor([2.0, 3.0])) # per expert, over the pair - - def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): - """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair - reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals - where the kernels assert one per expert.""" - import torch - - from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals - - targets = [ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - "experts.down_proj_input_global_scale", - ] - experts = 8 - for wrap in (False, True): - pair = torch.rand(experts, 2) + 1 - src = { - "experts.gate_up_proj_weight_scale_2": [pair] if wrap else pair, - "experts.down_proj_weight_scale_2": torch.rand(experts) + 1, - "experts.down_proj_input_scale": torch.rand(experts) + 1, - } - out = FineGrainedWeightGlobals(None).convert(src, target_patterns=targets, model=None) - with self.subTest(list_wrapped=wrap): - self.assertEqual( - out["experts.gate_up_proj_weight_global_scale"].shape, - torch.Size([experts]), - "the gate|up halves were not folded to one global per expert", - ) - - def test_one_converter_merges_a_layer_s_globals(self): - """modelopt calibrates the gate|up halves separately. One converter owns every global of - a layer, because merging them to the one-per-expert the kernels take moves the up half's - onto the down projection: SwiGLU is linear in the up half, so the stack keeps the gate's - global and the down's weight global scales the expert output back. The down's calibrated - input scale moves the other way, keeping the requantized intermediate on the range the - checkpoint calibrated.""" - from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals - - gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) - down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) - targets = [ - "mlp.experts.gate_up_proj_weight_global_scale", - "mlp.experts.down_proj_weight_global_scale", - "mlp.experts.down_proj_input_global_scale", - ] - out = FineGrainedWeightGlobals().convert( - { - "mlp.experts.*.up_proj.weight_scale_2": up, - "mlp.experts.*.gate_proj.weight_scale_2": gate, - "mlp.experts.*.down_proj.weight_scale_2": down, - "mlp.experts.*.down_proj.input_scale": down_input, - }, - target_patterns=targets, - ) - ratio = up / gate - torch.testing.assert_close(out[targets[0]], gate) - torch.testing.assert_close(out[targets[1]], down * ratio) - torch.testing.assert_close(out[targets[2]], down_input / ratio) - - def test_globals_with_one_calibrated_projection_pass_through(self): - """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a - single global) has nothing to merge, so every global reaches its module unchanged.""" - from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals - - targets = [ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - "experts.down_proj_input_global_scale", - ] - out = FineGrainedWeightGlobals().convert( - { - "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0], [2.0]]), - "experts.down_proj_weight_scale_2": torch.tensor([3.0, 4.0]), - "experts.down_proj_input_scale": torch.tensor([0.5, 0.25]), - }, - target_patterns=targets, - ) - torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) - torch.testing.assert_close(out[targets[1]], torch.tensor([3.0, 4.0])) - torch.testing.assert_close(out[targets[2]], torch.tensor([0.5, 0.25])) - - def test_the_fused_layout_ships_both_halves_in_one_tensor(self): - """A stacked `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two - keys, and merges the same way.""" - from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals - - targets = [ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - ] - out = FineGrainedWeightGlobals().convert( - { - "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 3.0], [2.0, 4.0]]), - "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), - }, - target_patterns=targets, - ) - torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) - torch.testing.assert_close(out[targets[1]], torch.tensor([5.0 * 3.0, 6.0 * 2.0])) - - def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_down(self): - """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated - `input_scale` reduces to one value; the down's rows are per expert, so its stays per - expert — the requant epilogue normalizes each row by its own expert's value.""" - from transformers.integrations.finegrained.conversions import FineGrainedInputScales - - up = FineGrainedInputScales().convert( - { - "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), - "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), - }, - full_layer_name="mlp.experts.gate_up_proj_input_global_scale", - )["mlp.experts.gate_up_proj_input_global_scale"] - torch.testing.assert_close(up, torch.tensor([3.0])) - down = FineGrainedInputScales().convert( - {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, - full_layer_name="mlp.experts.down_proj_input_global_scale", - )["mlp.experts.down_proj_input_global_scale"] - torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) - - -class _InputScaledRMSNorm(torch.nn.Module): - """A post-expert norm that scales by ``1 + weight`` BEFORE normalizing, so the row's - mean square is taken on the scaled values (not the same function as normalizing first).""" - - def __init__(self, dim, eps=1e-6): - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.zeros(dim, dtype=torch.bfloat16)) - - def forward(self, x): - out = x.float() * (1.0 + self.weight.float()) - return (out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) - - -class _UnfusableNorm(torch.nn.Module): - """A norm whose math is none of the fused forms (a bias the kernels have no slot for): the - binder has to leave it a module call rather than pick the nearest form.""" - - def __init__(self, dim, eps=1e-6): - super().__init__() - self.eps = eps - self.weight = torch.nn.Parameter(torch.ones(dim, dtype=torch.bfloat16)) - self.bias = torch.nn.Parameter(torch.full((dim,), 0.5, dtype=torch.bfloat16)) - - def forward(self, x): - xf = x.float() - normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) - return (normed * self.weight.float() + self.bias.float()).type_as(x) - - -def _nvfp4_two_level(kernel, weight): - """Canonical two-level NVFP4 quant of one matrix: the global is ``amax / (6 * 448)`` — the - smallest one that keeps every e4m3 block scale in range — and the block quant runs on the - normalized values. Returns ``(packed_e2m1, e4m3 block scales, fp32 global)``.""" - global_scale = (weight.abs().amax() / (6.0 * 448.0)).clamp(min=torch.finfo(torch.float32).tiny) - packed, block = kernel.nvfp4_act_quant((weight.float() / global_scale).contiguous()) - return packed.view(torch.int8), block, global_scale - - -def _unpack_e2m1(packed): - """Packed-E2M1 bytes (two codes per byte along the last axis) back to fp32 values.""" - lut = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], - device=packed.device, - ) - b = packed.view(torch.uint8) - pairs = torch.stack([lut[(b & 0x0F).long()], lut[(b >> 4).long()]], dim=-1) - return pairs.reshape(*packed.shape[:-1], packed.shape[-1] * 2) - - -def _targets(conv): - t = conv.target_patterns - return t if isinstance(t, (list, tuple)) else [t] - - -def _sm100(): - return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 - - -@require_torch_gpu -@unittest.skipUnless(_sm100(), "the fine-grained kernels target SM100+") -class FineGrainedRealKernelTest(unittest.TestCase): - """Runs the REAL kernels. The mocked tests above pin the calling contract; these pin that the - contract produces correct numbers, which no amount of argument-matching can show.""" - - @classmethod - def setUpClass(cls): - try: - load_finegrained_kernel() - except ImportError as e: - raise unittest.SkipTest(f"finegrained kernels unavailable: {e}") - - def _block_fp8_weight(self, N, K, *, ue8m0): - """(N, K) E4M3 + its (N/128, K/128) inv-scale grid, and the exact dequantized values the - kernel will see — so the reference is the quantization floor, not the pre-quant weight.""" - w = torch.randn(N, K, device="cuda", dtype=torch.float32) - blocks = w.reshape(N // 128, 128, K // 128, 128) - amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12) - inv = amax / 448.0 - if ue8m0: # power-of-two scales: the tcgen05 dot_scaled format - inv = torch.pow(2.0, torch.ceil(torch.log2(inv))) - q = (blocks / inv).to(torch.float8_e4m3fn) - deq = (q.float() * inv).reshape(N, K) - return q.reshape(N, K), inv.reshape(N // 128, K // 128).contiguous(), deq - - def test_linear_forward_matches_the_quantization_floor(self): - """A real quantized linear against its own dequantized weight. Loose vs bf16 would pass - even if the kernel silently dropped a K-block; comparing to the floor does not.""" - M, N, K = 256, 512, 384 - for ue8m0 in (False, True): - with self.subTest(ue8m0=ue8m0): - w, s, deq = self._block_fp8_weight(N, K, ue8m0=ue8m0) - x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) - out = finegrained_linear(x, w, s, block_size=[128, 128]) - ref = x.float() @ deq.t() - rel = ((out.float() - ref).norm() / ref.norm()).item() - self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") - - def test_on_the_fly_group_formats_round_trip_within_their_floor(self): - """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in - the module's layout (swizzled where held so); dequantizing them back lands within each - format's floor.""" - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize - - torch.manual_seed(0) - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 - cfg._experts_implementation = "grouped_mm" - for fmt, floor in (("mxfp8", 4e-2), ("mxfp4", 2e-1), ("nvfp4", 2e-1)): - with self.subTest(fmt=fmt): - experts = FineGrainedExperts(cfg, weight_format=fmt).cuda() - model = torch.nn.Module() - model.experts = experts - original = (torch.randn(4, 512, 512, device="cuda") * 0.02).to(torch.bfloat16) - out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) - weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] - self.assertEqual( - (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) - ) - self.assertEqual( - (scale.dtype, scale.shape), - (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), - ) - if scale.ndim == 5: - kernel = load_finegrained_kernel() - scale = kernel.unswizzle_mx_scales(scale, 512, scale.shape[2] * 4, num_experts=4) - deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) - if fmt == "nvfp4": - global_scale = out["experts.gate_up_proj_weight_global_scale"] - deq = deq * global_scale.reshape(-1, 1, 1) - else: - self.assertNotIn("experts.gate_up_proj_weight_global_scale", out) - rel = ((deq - original.float()).norm() / original.float().norm()).item() - self.assertLess(rel, floor, f"{fmt}: {rel:.3f}") - self.assertGreater(rel, 0.0) - - def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes=(8.0, 1.0)): - """NVFP4 experts in the layout a modelopt checkpoint delivers: every (expert, half) of the - gate|up stack quantized against its OWN global (``weight_scale_2`` is per projection, and - the two can differ per expert), the down per expert, plus calibrated activation - globals. Returns the module, the exact dequantized weights (so the reference is the - quantization floor rather than the pre-quant weight), and the dequantization a reader that - applied ONE of the two globals to both halves would get.""" - from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals - - kernel = load_finegrained_kernel() - experts = FineGrainedExperts(cfg, weight_format="nvfp4", activation_format=activation_format).cuda() - dequantized, single_global, checkpoint_globals = {}, {}, {} - for proj, rows, cols in ( - ("gate_up_proj", 2 * cfg.intermediate_size, cfg.hidden_size), - ("down_proj", cfg.hidden_size, cfg.intermediate_size), - ): - reference = torch.randn(cfg.num_local_experts, rows, cols, device="cuda") - if proj == "gate_up_proj": # the halves interleave by row; draw them far apart - reference[:, 0::2] *= magnitudes[0] - reference[:, 1::2] *= magnitudes[1] - packed = torch.empty(cfg.num_local_experts, rows, cols // 2, dtype=torch.int8, device="cuda") - block = torch.empty(cfg.num_local_experts, rows, cols // 16, dtype=torch.float8_e4m3fn, device="cuda") - halves = 2 if proj == "gate_up_proj" else 1 - globals_ = torch.empty(cfg.num_local_experts, halves, device="cuda") - deq, deq_one = torch.empty_like(reference), torch.empty_like(reference) - for e in range(cfg.num_local_experts): - for h in range(halves): - rows_h = slice(h, None, halves) - q, b, g = _nvfp4_two_level(kernel, reference[e, rows_h]) - packed[e, rows_h], block[e, rows_h], globals_[e, h] = q, b, g - values = _unpack_e2m1(q) * b.float().repeat_interleave(16, dim=-1) - deq[e, rows_h] = values * g - if halves == 2: - # the control: BOTH halves read the up half's global, which is what a kernel - # that ignored the layout does. It has to be the up one — putting the gate's - # global on both scales the expert's whole output uniformly, and the - # post-expert RMSNorm divides exactly that out; a wrong GATE scale goes - # through the SiLU and survives. - for h in (0, 1): - rows_h = slice(h, None, 2) - deq_one[e, rows_h] = deq[e, rows_h] * (globals_[e, 1] / globals_[e, h]) - scale = block if getattr(experts, f"{proj}_scale_inv").dim() != 5 else kernel.swizzle_mx_scales(block) - setattr(experts, proj, torch.nn.Parameter(packed, requires_grad=False)) - setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(scale, requires_grad=False)) - checkpoint_globals[f"experts.{proj}_weight_scale_2"] = globals_ - dequantized[proj], single_global[proj] = deq, deq_one - # calibrated input_scale: one value for the gate_up (quantized before routing), one per - # expert for the down (its rows are per expert). A weight-only module holds neither — - # nothing there quantizes activations — so the checkpoint's keys stay unused. - if activation_format != "bf16": - experts.gate_up_proj_input_global_scale = torch.nn.Parameter( - torch.full((1,), 0.05, device="cuda"), requires_grad=False - ) - checkpoint_globals["experts.down_proj_input_scale"] = torch.linspace( - 0.02, 0.08, cfg.num_local_experts, device="cuda" - ) - single_global["down_proj"] = dequantized["down_proj"] # the down has one global either way - # The globals go in through the loader's own converter, so the module holds exactly what a - # real checkpoint puts there: one per expert, with the gate|up pair merged onto the down. - targets = [ - "experts.gate_up_proj_weight_global_scale", - "experts.down_proj_weight_global_scale", - "experts.down_proj_input_global_scale", - ] - for target, value in FineGrainedWeightGlobals().convert(checkpoint_globals, target_patterns=targets).items(): - setattr(experts, target.split(".")[-1], torch.nn.Parameter(value, requires_grad=False)) - return experts, dequantized, single_global - - def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): - """The experts chain in torch over the dequantized weights: gate|up (interleaved columns), - SwiGLU, down, the optional per-expert output norm, then the routing-weighted sum.""" - out = torch.zeros_like(x, dtype=torch.float32) - for token in range(x.shape[0]): - for slot in range(idx.shape[1]): - e = int(idx[token, slot]) - pre = x[token].float() @ dequantized["gate_up_proj"][e].t() - inter = torch.nn.functional.silu(pre[0::2]) * pre[1::2] - row = inter @ dequantized["down_proj"][e].t() - if post_norm is not None: - row = post_norm(row.to(x.dtype)).float() - out[token] += row * float(wts[token, slot]) - return out - - def test_modelopt_nvfp4_experts_merge_their_globals_and_fuse_the_post_norm(self): - """The stacked checkpoint shape end to end through the integration: a gate|up stack - calibrated per half (merged at load), calibrated activation globals, and a per-expert - output norm the binder fuses, on all three forwards. Weight-only pins the numbers against the dequantized weights (nothing else - rounds); the W4A4 chain rides the 4-bit activation floor, so it is pinned by being far - closer to the right globals than to one global folded over both halves — what a kernel - that ignored the layout computes.""" - torch.manual_seed(0) - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 - cfg._experts_implementation = "grouped_mm" - x = torch.randn(8, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) * 0.1 - idx = torch.randint(0, cfg.num_local_experts, (8, 2), device="cuda", dtype=torch.long) - wts = torch.rand(8, 2, device="cuda", dtype=torch.bfloat16) - forwards = ( - ("grouped", fg.finegrained_grouped_mm_experts_forward), - ("batched", fg.finegrained_batched_mm_experts_forward), - ("eager", lambda module, *args: module(*args)), - ) - - def relative(out, reference): - return ((out.float() - reference).norm() / reference.norm()).item() - - for activation_format, floor in (("bf16", 0.02), (None, 0.35)): - experts, dequantized, one_global = self._nvfp4_experts_modelopt_layout(cfg, activation_format) - reference = self._moe_reference(dequantized, x, idx, wts) - wrong = self._moe_reference(one_global, x, idx, wts) - for name, forward in forwards: - with self.subTest(forward=name, activation_format=activation_format): - out = forward(experts, x, idx, wts) - rel, rel_wrong = relative(out, reference), relative(out, wrong) - self.assertLess(rel, floor, f"{name} diverged from the dequantized reference: {rel:.3f}") - self.assertLess( - 3 * rel, - rel_wrong, - f"{name} is as close to ONE global over both gate|up halves ({rel_wrong:.3f}) " - f"as to the per-half ones ({rel:.3f}) — the fold is not reading the halves", - ) - - # the same chain with the model's per-expert output norm, which runs on the routed - # rows: every forward has to apply it where the reference forwards do. A name the - # kernels implement is folded into the chain's reduce and anything else is the module - # itself, so both have to land on the same numbers as the reference. - for norm, norm_name in ( - (torch.nn.RMSNorm(cfg.hidden_size, device="cuda", dtype=torch.bfloat16), "rms_norm"), - (_InputScaledRMSNorm(cfg.hidden_size).cuda(), "input_scaled_rms_norm"), - # a form the kernels do not implement: the module is called on the rows - (_UnfusableNorm(cfg.hidden_size).cuda(), "a_models_own_norm"), - ): - norm.weight.data.uniform_(0.5, 1.5) - experts.post_expert_norm, experts.post_expert_norm_name = norm, norm_name - experts.has_post_expert_norm = True - normed = self._moe_reference(dequantized, x, idx, wts, norm) - self.assertGreater(relative(normed, reference), 0.05, "the norm has to change the output") - for name, forward in forwards: - with self.subTest(forward=name, activation_format=activation_format, norm=norm_name): - rel = relative(forward(experts, x, idx, wts), normed) - self.assertLess(rel, floor, f"{name} diverged with the post-expert norm: {rel:.3f}") - - def test_experts_scale_layout_op_keeps_the_forward_correct(self): - """End-to-end over the integration: real MXFP8 experts hold swizzled scales, the loader's op - fills them from the affine grid, the fused forward matches the affine module's, and the - reverse op hands the affine grid back bitwise.""" - from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales - - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = ( - 512, - 256, - 4, - ) # K >= 256: the swizzled arm has tiles - cfg._experts_implementation = "grouped_mm" - experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() - with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): - affine_experts = FineGrainedExperts(cfg, weight_format="mxfp8").cuda() - model = torch.nn.Module() - model.experts = experts - op = FineGrainedSwizzleScales(hf_quantizer=None) - - affine = {} - for proj, rows in (("gate_up_proj", 2 * cfg.intermediate_size), ("down_proj", cfg.hidden_size)): - cols = cfg.hidden_size if proj == "gate_up_proj" else cfg.intermediate_size - weight = torch.randn(cfg.num_local_experts, rows, cols, device="cuda").to(torch.float8_e4m3fn) - grid = torch.randint(120, 134, (cfg.num_local_experts, rows, cols // 32), dtype=torch.uint8, device="cuda") - affine[proj] = grid.view(torch.float8_e8m0fnu) - for module in (experts, affine_experts): - setattr(module, proj, torch.nn.Parameter(weight.clone(), requires_grad=False)) - setattr(affine_experts, f"{proj}_scale_inv", torch.nn.Parameter(affine[proj].clone(), requires_grad=False)) - self.assertEqual( - getattr(experts, f"{proj}_scale_inv").dim(), 5, "the module should hold this projection swizzled" - ) - out = op.convert( - {f"mlp.experts.{proj}": affine[proj]}, model=model, full_layer_name=f"experts.{proj}_scale_inv" - ) - swizzled = out[f"mlp.experts.{proj}"] - self.assertEqual(swizzled.shape, getattr(experts, f"{proj}_scale_inv").shape) - setattr(experts, f"{proj}_scale_inv", torch.nn.Parameter(swizzled, requires_grad=False)) - - x = torch.randn(8, cfg.hidden_size, device="cuda", dtype=torch.bfloat16) - idx = torch.randint(0, cfg.num_local_experts, (8, 2), device="cuda", dtype=torch.long) - wts = torch.rand(8, 2, device="cuda", dtype=torch.bfloat16) - reference = fg.finegrained_grouped_mm_experts_forward(affine_experts, x, idx, wts) - out = fg.finegrained_grouped_mm_experts_forward(experts, x, idx, wts) - torch.cuda.synchronize() - rel = ((out.float() - reference.float()).norm() / reference.float().norm().clamp(min=1e-9)).item() - # different tuned arms reorder the fp32 reduction (~1e-5); a layout bug is O(1) - self.assertLess(rel, 1e-3, f"swizzled forward diverged from the affine one: {rel:.2e}") - # the eager per-expert loop reads one expert's slice of the swizzled stack directly - eager_reference = affine_experts(x, idx, wts) - eager_out = experts(x, idx, wts) - rel = ( - (eager_out.float() - eager_reference.float()).norm() / eager_reference.float().norm().clamp(min=1e-9) - ).item() - self.assertLess(rel, 1e-3, f"eager forward on swizzled scales diverged from affine: {rel:.2e}") - - reverse = op.reverse_op - for proj, grid in affine.items(): - restored = reverse.convert({proj: getattr(experts, f"{proj}_scale_inv").data})[proj] - self.assertTrue(torch.equal(restored.view(torch.uint8), grid.view(torch.uint8)), proj) - - -def _checkpoint_shapes(path): - """`{key: shape}` of a saved checkpoint, for comparing one save against another.""" - import glob - - from safetensors import safe_open - - shapes = {} - for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))): - with safe_open(shard, framework="pt") as handle: - for key in handle.keys(): - shapes[key] = tuple(handle.get_slice(key).get_shape()) - return shapes - - -_SHARDING_WORKER = """ -import importlib, os, sys, torch -from transformers.distributed import DistributedConfig - -model_dir, out_dir, modes, cls_path = sys.argv[1], sys.argv[2], sys.argv[3].split(","), sys.argv[4] -module_name, cls_name = cls_path.split(":") -model_cls = getattr(importlib.import_module(module_name), cls_name) -world = int(os.environ["WORLD_SIZE"]) -ids = torch.arange(16, dtype=torch.long).unsqueeze(0) - -# both modes in ONE launch: they share the mesh, so the process group is created once and only -# the plan differs -- a second torchrun would pay another interpreter + CUDA start -for mode in modes: - model = model_cls.from_pretrained( - model_dir, - dtype="auto", - attn_implementation="eager", - distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=(mode == "ep")), - ).eval() - experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) - weight = getattr(experts, "gate_up_proj", None) - if weight is None: - weight = experts.up_proj - local = weight.to_local() if hasattr(weight, "to_local") else weight - with torch.no_grad(): - logits = model(ids.to(model.device)).logits.float().cpu() - if int(os.environ["RANK"]) == 0: - torch.save({"logits": logits, "expert_local": tuple(local.shape)}, - os.path.join(out_dir, mode + ".pt")) - del model - torch.cuda.empty_cache() -""" - - -def _checkpoint_expert_shape(model_dir): - """The UNSHARDED `(experts, rows)` of the stacked gate|up projection the module holds. - - Read from the checkpoint, which ships the experts either already fused - (`experts.gate_up_proj`) or one per expert (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`, - whose rows are gate and up separately). The baseline a sharded leg is measured against has to come from outside - the sharded runs: taking it from whichever leg ran first compares one mode to another, and a - mode that placed nothing then still looks sharded. - """ - import glob - import json - import os - - from safetensors import safe_open - - shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) - index = os.path.join(model_dir, "model.safetensors.index.json") - if os.path.exists(index): - with open(index, encoding="utf-8") as fh: - weight_map = json.load(fh)["weight_map"] - shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) - - experts: set[int] = set() - rows = 0 - for shard in shards: - with safe_open(shard, framework="pt") as handle: - for key in handle.keys(): - shape = None - if re.search(r"\.experts\.gate_up_proj$", key): - shape = handle.get_slice(key).get_shape() - if len(shape) == 3: - return (shape[0], shape[1]) - if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): - experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) - rows = max(rows, handle.get_slice(key).get_shape()[0]) - if not experts or not rows: - raise AssertionError(f"no expert projection in {model_dir} to size the shard against") - return (len(experts), 2 * rows) # the module stacks gate and up into one row extent - - -@require_torch_multi_accelerator -class FineGrainedLoadPathEquivalenceTest(TestCasePlus): - """Expert parallelism and intra-expert tensor parallelism must both reproduce what - `device_map` gives, which is the baseline because it places modules across devices without - splitting a single tensor — no process group, no DTensor, no collectives. So it computes the - unsharded answer, on the multi-GPU path people actually deploy. - - Each model is built in the format it actually ships, because the conversion path differs by - format and by checkpoint layout — a per-expert FP8 checkpoint and a fused NVFP4 one reach the - experts through different converters, and bugs have hidden in exactly that gap. - - Three properties make this able to fail, all of which earlier versions lacked: - * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the - shard it already holds, so they come out correctly sized whatever the plan says and a - plan that shards no scale at all still produces the right answer. - * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty - (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is skipped EXPLICITLY rather - than passing vacuously. - * every leg reports the local expert shard, so a leg that placed nothing fails loudly. - """ - - @staticmethod - def _model_table(): - """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}` — each model - in the format it actually ships, with the quantization config that format arrives under, - not just its name: block-FP8 carries a `weight_block_size`, and an NVFP4 checkpoint comes - from modelopt under `quant_algo`, which is remapped on construction. Dims are multiples - of the 128 block so a 2-way split stays block-aligned, and the expert count divides the - mesh.""" - from transformers import ( - DeepseekV3Config, - DeepseekV3ForCausalLM, - DeepseekV4Config, - DeepseekV4ForCausalLM, - Glm4vMoeConfig, - Glm4vMoeForConditionalGeneration, - Glm4vMoeTextConfig, - Glm4vMoeVisionConfig, - GptOssConfig, - GptOssForCausalLM, - MiniMaxM3SparseForConditionalGeneration, - MiniMaxM3VLConfig, - Mistral4Config, - Mistral4ForCausalLM, - ) - from transformers.utils.quantization_config import FineGrainedConfig - - return { - "deepseek_v3-fp8": ( - DeepseekV3Config, - DeepseekV3ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "n_group": 1, - "topk_group": 1, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - "q_lora_rank": None, - "kv_lora_rank": 32, - "qk_nope_head_dim": 32, - "qk_rope_head_dim": 16, - "v_head_dim": 32, - }, - # DeepSeek-V3 ships block-FP8: 128x128 weight blocks, activations quantized - # per token at run time - FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), - ), - # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a - # calibrated activation scale per quantized module, which for a MoE is one per expert. - # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense - # counterpart of the same export. - "mistral4-fp8_tensor_static": ( - Mistral4Config, - Mistral4ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - }, - FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static"), - ), - # MULTIMODAL: the experts' plans live on `text_config`, not on the config the - # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer - # one adds no companion while the weights still shard, which is how a real - # multimodal MoE ends up with whole scales against sharded weights. - "glm4v_moe-nvfp4": ( - Glm4vMoeConfig, - Glm4vMoeForConditionalGeneration, - { - "text_config": Glm4vMoeTextConfig( - vocab_size=64, - hidden_size=256, - intermediate_size=256, - moe_intermediate_size=256, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=32, - rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, - rope_theta=10000, - tie_word_embeddings=True, - bos_token_id=0, - eos_token_id=0, - pad_token_id=0, - n_routed_experts=4, - n_shared_experts=1, - n_group=1, - topk_group=1, - num_experts_per_tok=2, - first_k_dense_replace=0, - ), - "vision_config": Glm4vMoeVisionConfig( - depth=2, - num_heads=4, - hidden_size=64, - out_hidden_size=256, - intermediate_size=64, - patch_size=14, - spatial_merge_size=1, - temporal_patch_size=2, - ), - }, - # the GLM NVFP4 checkpoints are modelopt exports: `quant_algo` names the format - # and `FineGrainedConfig` remaps it to nvfp4 at construction - FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4"), - ), - # interleaved rows (`is_concatenated=False`), transposed, with expert biases — and - # an EP plan that already names those biases, so the companion rules meet entries - # the model wrote itself - "gpt_oss-mxfp4": ( - GptOssConfig, - GptOssForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "num_local_experts": 4, - "num_experts_per_tok": 2, - "max_position_embeddings": 32, - }, - # GPT-OSS ships weight-only: raw bf16 activations against packed fp4 weights - FineGrainedConfig(quant_method="mxfp4", activation_format="bf16"), - ), - # MIXED precision, and the only entry whose expert format is not the quantization - # config's: `expert_dtype` is a model-config side-channel that makes the EXPERTS - # mxfp4 while the dense and attention paths stay block-FP8 — with scales in UE8M0 - # containers rather than fp32, the other `scale_fmt` - "deepseek_v4-fp4_experts+fp8_dense": ( - DeepseekV4Config, - DeepseekV4ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - "expert_dtype": "fp4", - }, - FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0"), - ), - # a second MULTIMODAL nesting, in the group-32 MX format. The sub-config shape - # follows this model's own tester; only the MoE dims are raised to a multiple of - # the 128 block so a 2-way split stays block-aligned. - "minimax_m3_vl-mxfp8": ( - MiniMaxM3VLConfig, - MiniMaxM3SparseForConditionalGeneration, - { - "text_config": { - "hidden_size": 256, - # 512 so the 2-way split leaves 256: the experts' down projection - # contracts over this, and a sharded 128 has no 128-wide swizzled tile - "intermediate_size": 512, - "dense_intermediate_size": 256, - "shared_intermediate_size": 256, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "head_dim": 64, - "rotary_dim": 32, - "vocab_size": 64, - "max_position_embeddings": 32, - "bos_token_id": 0, - "eos_token_id": 1, - "pad_token_id": 2, - "num_local_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "moe_layer_freq": [0, 1], - "layer_types": ["full_attention", "minimax_m3_sparse"], - "tie_word_embeddings": False, - "index_n_heads": 2, - "index_head_dim": 16, - "index_block_size": 8, - "index_topk_blocks": 4, - "index_local_blocks": 1, - }, - "vision_config": { - # 256 so the 2-way SPLIT is still 128-aligned: an MXFP8 weight with - # pre-swizzled scales is read in 128-wide K tiles, and a sharded 128 dim - # leaves 64 — which has none to offer, and the tuner has no config at all - "hidden_size": 256, - "intermediate_size": 256, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_channels": 3, - "image_size": 14, - "patch_size": 14, - "temporal_patch_size": 2, - "spatial_merge_size": 1, - }, - "image_token_index": 4, - "video_token_index": 5, - "projector_hidden_size": 256, - "pad_token_id": 2, - }, - FineGrainedConfig(quant_method="mxfp8"), - ), - } - - @classmethod - def setUpClass(cls): - cls._tmp = tempfile.TemporaryDirectory() - # the kernels autotune per shape from a cold cache, which dwarfs everything else here - cls._env = { - "FINEGRAINED_AUTOTUNE_TRIALS": "1", - "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), - } - os.environ.update(cls._env) - - @classmethod - def tearDownClass(cls): - cls._tmp.cleanup() - - @staticmethod - def _free_port() -> int: - with socket.socket() as sock: - sock.bind(("", 0)) - return sock.getsockname()[1] - - def _fixture(self, label): - """A PRE-QUANTIZED checkpoint of one tiny model, in the format that model ships.""" - import torch - - config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] - config = config_cls(**kwargs) - - bf16_dir = os.path.join(self._tmp.name, label, "bf16") - quant_dir = os.path.join(self._tmp.name, label, "quantized") - torch.manual_seed(0) - # BF16 is what every model here ships, and the dtype decides which kernel arms run: a - # float32 checkpoint (torch's default for a freshly built model, which `dtype="auto"` - # then faithfully reloads) sends the weight-only formats down arms `tl.dot_scaled` - # cannot serve at all, so the suite would exercise a dtype nobody deploys and leave the - # real one uncovered. Saving in bf16 makes `"auto"` mean bf16 for every load below. - model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) - quantized = model_cls.from_pretrained( - bf16_dir, - dtype="auto", - attn_implementation="eager", - quantization_config=quantization_config, - device_map="cuda:0", - ) - quantized.save_pretrained(quant_dir, safe_serialization=True) - del quantized - torch.cuda.empty_cache() - - # SAVE must restore the checkpoint's own layout: the layout ops (gate|up interleave, - # scale container, swizzle) each have a reverse, and a quantized model reloaded and - # written again has to land on the same keys and shapes. A broken reverse writes a - # corrupt checkpoint silently — the module still reads back whatever it wrote. - reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map="cuda:0") - round_trip = os.path.join(self._tmp.name, label, "round_trip") - reloaded.save_pretrained(round_trip, safe_serialization=True) - del reloaded - torch.cuda.empty_cache() - self.assertEqual( - _checkpoint_shapes(quant_dir), _checkpoint_shapes(round_trip), f"{label}: save did not round-trip" - ) - return quant_dir, config, model_cls - - def _logits(self, model_dir, model_cls, **load_kwargs): - import torch - - model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() - ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) - with torch.no_grad(): - logits = model(ids).logits.float().cpu() - del model - torch.cuda.empty_cache() - return logits - - def _rounding_floor(self, model_dir, model_cls, reference): - """How far this model's logits move when every sharded block is perturbed by BF16 rounding. - - Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different - order, so the first sharded op differs by an ULP. How far that travels is a property of - the MODEL, not of the sharding — a chain of MoE layers can amplify it a hundredfold while - a dense stack barely moves. Injecting the same magnitude on ONE device measures that - amplification directly, giving each model a budget its own conditioning earns. - """ - import torch - - model = model_cls.from_pretrained( - model_dir, dtype="auto", attn_implementation="eager", device_map="cuda:0" - ).eval() - # Perturb EVERY block a sharded reduction passes through, not one: TP re-orders the sum - # in each of them, so the rounding accumulates down the stack instead of cancelling. - # Perturbing a single site with random noise measured LESS movement than sharding caused - # even at 14x the magnitude, which is what accumulation looks like from the wrong model. - blocks = [ - m - for n, m in model.named_modules() - if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 - ] - generator = torch.Generator(device=model.device).manual_seed(0) - - def perturb(module, args, output): - tensor = output[0] if isinstance(output, tuple) else output - if not torch.is_tensor(tensor): - return output - ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() - noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp - tensor = tensor + noise - return (tensor,) + output[1:] if isinstance(output, tuple) else tensor - - handles = [b.register_forward_hook(perturb) for b in blocks] - ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) - with torch.no_grad(): - perturbed = model(ids).logits.float().cpu() - for handle in handles: - handle.remove() - del model - torch.cuda.empty_cache() - return (perturbed - reference).abs().max().item() - - def _sharded(self, model_dir, model_cls, modes): - """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" - import torch - - script = os.path.join(self._tmp.name, "sharding_worker.py") - with open(script, "w", encoding="utf-8") as fh: - fh.write(_SHARDING_WORKER) - out = os.path.join(self._tmp.name, "out") - os.makedirs(out, exist_ok=True) - subprocess.run( - [ - "torchrun", - "--nproc_per_node=2", - f"--master_port={self._free_port()}", - script, - model_dir, - out, - ",".join(modes), - # the model's own class: `AutoModelForCausalLM` cannot resolve a multimodal - # `ForConditionalGeneration` from its config - f"{model_cls.__module__}:{model_cls.__name__}", - ], - check=True, - env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, - ) - return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} - - def test_every_load_path_agrees(self): - import torch - - for label in self._model_table(): - with self.subTest(model=label): - model_dir, config, model_cls = self._fixture(label) - # the baseline, spread across both devices: `max_memory` forces a real split, - # since `auto` alone fits this whole model on one GPU and would quietly compare - # the sharded legs against a single-device load - reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) - noise_floor = self._rounding_floor(model_dir, model_cls, reference) - - # Only the modes this model's OWN plans shard experts under. A mode whose plan - # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places - # nothing, so running it would pass without being evidence of anything. A - # multimodal model keeps these on a sub-config. - owners = [config] + [ - c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) - ] - modes = [ - mode - for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) - if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) - ] - self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") - whole = _checkpoint_expert_shape(model_dir) - for mode, payload in self._sharded(model_dir, model_cls, modes).items(): - with self.subTest(model=label, mode=mode): - local = payload["expert_local"] - # against the CHECKPOINT's own shape, never another mode's shard: seeding - # this from the first leg made the second leg check EP against TP, so a - # mode that placed nothing still looked sharded (glm4v's TP left its - # experts replicated and passed). - self.assertLess( - local[0] * local[1], - whole[0] * whole[1], - f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " - f"this leg cannot catch a bad axis", - ) - # Sharding reorders reductions, so the legs differ by BF16 rounding - # whatever the axes — and how much that shows at the logits is the - # model's own business (the NVFP4 fixture amplifies ~100x, the dense - # ones barely). Hence each model's measured floor, not a fixed - # tolerance. A wrong shard axis lands orders of magnitude above it. - budget = max(2e-2, noise_floor) - gap = (payload["logits"] - reference).abs().max().item() - self.assertLessEqual( - gap, - budget, - f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " - f"model's own rounding budget {budget:.4g} — that is a sharding bug, " - f"not reduction order", - ) - if noise_floor <= 2e-2: - # only meaningful where rounding does NOT already move the argmax - self.assertTrue( - torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), - f"{label}/{mode}: argmax diverged from the unsharded model", - ) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 174d1dce8cb6..24dd238a7963 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -13,7 +13,6 @@ # limitations under the License. import copy import unittest -from types import SimpleNamespace import torch import torch.nn as nn @@ -33,7 +32,6 @@ GroupWeightRename, LinearToConv3d, MergeModulelist, - PermuteForRope, PrefixChange, VisionFuseAndPermuteForRope, VisionUnfuseAndPermuteForRope, @@ -513,182 +511,6 @@ def test_moe_and_qkv_conversion_reversed(self): # Make sure both saved state_dict are identical self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict)) - def test_qkv_chunk_rope_permute_with_fp8_quantization(self): - # block-FP8 quantize/dequantize is pure torch here; no kernel is launched - from transformers.integrations.finegrained import FineGrainedLinear - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize - - n_heads = 2 - head_dim = 4 - in_dim = 4 - out_dim = n_heads * head_dim - block_size = (4, 4) - - class RopeProjector(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.zeros(out_dim, in_dim)) - - class RopeSelfAttn(nn.Module): - def __init__(self): - super().__init__() - # only q_proj is quantized, so only it is a finegrained module holding a scale - self.q_proj = FineGrainedLinear(in_dim, out_dim, block_size=block_size, weight_format="fp8") - self.k_proj = RopeProjector() - self.v_proj = RopeProjector() - - class RopeLayer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = RopeSelfAttn() - - class RopeModel(PreTrainedModel): - base_model_prefix = "model" - - def __init__(self, config): - super().__init__(config) - self.layers = nn.ModuleList([RopeLayer()]) - self.post_init() - - def _init_weights(self, module): - pass # every weight here comes from the state dict, and `normal_` has no fp8 kernel - - config = PreTrainedConfig() - config.num_attention_heads = n_heads - model = RopeModel(config) - - raw_q = torch.tensor( - [ - [1.0, -1.0, 1.0, -1.0], - [0.5, -0.5, 0.5, -0.5], - [-1.0, 1.0, -1.0, 1.0], - [-0.5, 0.5, -0.5, 0.5], - [1.0, 1.0, -1.0, -1.0], - [0.5, 0.5, -0.5, -0.5], - [-1.0, -1.0, 1.0, 1.0], - [-0.5, -0.5, 0.5, 0.5], - ], - dtype=torch.float32, - ) - raw_k = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) - raw_v = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + 100.0 - raw_qkv = torch.cat([raw_q, raw_k, raw_v], dim=0) - state_dict = {"layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} - - quantizer_cls = type( - "FineGrainedHfQuantizer", - (), - { - "__init__": lambda self, bs=block_size: setattr( - self, "quantization_config", SimpleNamespace(weight_block_size=bs, scale_fmt="float") - ), - "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), - "get_quantize_ops": lambda self: FineGrainedQuantize(self), - "pre_quantized": False, - }, - ) - quantizer = quantizer_cls() - - weight_mapping = [ - WeightConverter( - "self_attn.qkv_proj.weight", - [ - "self_attn.q_proj.weight", - "self_attn.k_proj.weight", - "self_attn.v_proj.weight", - ], - operations=[Chunk(dim=0), PermuteForRope(permute_layer_names=["q_proj", "k_proj"])], - ) - ] - load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) - loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config) - - self.assertEqual(loading_info.missing_keys, set()) - self.assertEqual(loading_info.unexpected_keys, set()) - self.assertEqual(loading_info.mismatched_keys, set()) - self.assertEqual(loading_info.conversion_errors, {}) - - permute_op = PermuteForRope(permute_layer_names=["q_proj", "k_proj"]) - permute_op.config = model.config - expected_q = permute_op._apply(raw_q) - expected_k = permute_op._apply(raw_k) - expected_v = raw_v # V carries no RoPE, so the permute names only q and k - - model_state = model.state_dict() - self.assertFalse(torch.allclose(raw_k, expected_k)) - torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) - - q_weight_key = "layers.0.self_attn.q_proj.weight" - scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" - self.assertIn(scale_key, model_state) - expected_dtype = torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else torch.int8 - self.assertEqual(model_state[q_weight_key].dtype, expected_dtype) - self.assertEqual(model_state[q_weight_key].shape, torch.Size((out_dim, in_dim))) - self.assertEqual(model_state[scale_key].dtype, torch.float32) - self.assertEqual( - model_state[scale_key].shape, - torch.Size((out_dim // block_size[0], in_dim // block_size[1])), - ) - - dequantized_q = FineGrainedDequantize(None)._dequantize_one( - model_state[q_weight_key], model_state[scale_key], output_dtype=torch.float32 - ) - torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) - - def test_fp8_ue8m0_quantize_dequantize_round_trip(self): - # ``scale_fmt="ue8m0"`` rounds weight_scale_inv to a power of two; the weight has to be - # quantized with that same rounded scale, or ``weight * weight_scale_inv`` at dequant - # disagrees with the scale the weight was divided by (a silent per-block error up to an octave). - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize - - if not hasattr(torch, "float8_e8m0fnu"): - self.skipTest("ue8m0 storage requires torch.float8_e8m0fnu") - - quantizer = SimpleNamespace( - quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="ue8m0") - ) - torch.manual_seed(0) - weight = torch.randn(128, 128, dtype=torch.float32) - - # no model to ask, so the op takes the block and the scale format from the config - quantized = FineGrainedQuantize(quantizer)._quantize_one("layer.weight", weight, None) - recovered = FineGrainedDequantize(quantizer)._dequantize_one( - quantized["layer.weight"], quantized["layer.weight_scale_inv"].float(), output_dtype=torch.float32 - ) - rel_err = ((recovered - weight).abs().sum() / weight.abs().sum()).item() - self.assertLess(rel_err, 5e-2) # fp8 round-trip is ~2e-2; a scale mismatch inflates it past 0.2 - - def test_fp8_float_scale_fmt_quantization_unchanged(self): - """The default `scale_fmt="float"` path stays bit-identical to its formula: the inverse - scale is `amax / MAX` and the weight is DIVIDED by it. - - Reordering that into the algebraically equal `weight * (MAX / amax)` is not bit-equal — - measured, it moves 29 of 100 random blocks' scales and 1 of their weights — so a refactor - that "simplifies" the arithmetic changes every checkpoint this path writes. Scan many - blocks, since the divergence only shows on particular fp32 maxima. - """ - from transformers.integrations.finegrained.conversions import ( - _FP8_DTYPE, - _FP8_MAX, - _FP8_MIN, - FineGrainedQuantize, - ) - - quantizer = FineGrainedQuantize( - SimpleNamespace(quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float")) - ) - for seed in range(100): - torch.manual_seed(seed) - weight = torch.randn(128, 128, dtype=torch.float32) - out = quantizer._quantize_one("layer.weight", weight, None) - ref_inv = (weight.abs().amax() / _FP8_MAX).to(torch.float32).reshape(1, 1) - ref_q = torch.clamp(weight / ref_inv, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) - self.assertTrue(torch.equal(out["layer.weight"], ref_q), f"float weight diverged at seed {seed}") - self.assertTrue( - torch.equal(out["layer.weight_scale_inv"], ref_inv), f"float scale diverged at seed {seed}" - ) - def test_scoped_renaming_does_not_leak_to_sibling_or_parent(self): """scope_prefix gates a WeightRenaming to keys under one submodel only — neither the sibling submodel nor the parent's own keys must be affected. From 9f7326c65b004d969e5106d09204fa60ef0f6e95 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 00:51:49 -0700 Subject: [PATCH 51/68] Carry the review's comment and lint fixes to the sites it did not name The review asked to shorten one four-line comment and I shortened only that one. The same cut applies to the three others this branch adds, and to the block in `sharding_utils` that two new lines had grown to four. The two `# noqa: E731` lambdas in the plan tests are plain functions now, which also settles their disagreeing parameter names. Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/sharding_utils.py | 4 +--- src/transformers/integrations/finegrained/core.py | 13 +++++-------- src/transformers/quantizers/finegrained/nvfp4.py | 7 +++---- tests/integrations/finegrained/test_core.py | 8 ++++++-- 4 files changed, 15 insertions(+), 17 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 21839e307af4..9674d21fdc46 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -87,9 +87,7 @@ def __init__(self, param: DTensor): self.param_ndim = param.ndim local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) # Axis-0 range owned by this rank (used to filter per-expert pieces) - # [_axis0_offset, _axis0_offset + _axis0_local_size) - # a 0-dim (per-tensor) parameter is replicated whole — see `shard_tensor` — so it has no - # axis 0 to take an offset on + # [_axis0_offset, _axis0_offset + _axis0_local_size); a 0-dim parameter has no axis 0 self._axis0_offset = offsets[0] if offsets else 0 self._axis0_local_size = local_shape[0] if local_shape else 1 diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 6a488ba4cb3b..65c5ddb35c2b 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -600,10 +600,8 @@ def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top class FineGrainedExperts(_FineGrainedModule, nn.Module): - # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. The default - # kind is impl-agnostic; some impls need a distinct one (megamoe needs no gradient-sync hooks - # and an EP `process_group` injection). Declared here so the quantizer does not have to know - # about impl-specific TP needs — extend when adding an impl. + # the plan kinds an `_experts_implementation` needs in place of the impl-agnostic default + # (megamoe wants no gradient-sync hooks and an EP `process_group`); extend when adding an impl _impl_tp_layer_overrides: dict[str, dict[str, str]] = { "deepgemm_megamoe": {"moe_tp_experts": "megamoe_experts", "ep_router": "megamoe_router"}, } @@ -926,10 +924,9 @@ def storage_for(module_name: str) -> dict: **storage_for(module_name), ) elif isinstance(module, nn.Linear) and hasattr(module, "n_groups"): - # Block-diagonal grouped linear (DSv4's `DeepseekV4GroupedLinear`), recognised by - # the attribute the swap needs rather than by its class NAME: a plain - # `FineGrainedLinear` would collapse the groups into one giant linear and yield - # the wrong output dim. + # block-diagonal grouped linear (DSv4), matched on the attribute the swap needs + # rather than the class name: a plain `FineGrainedLinear` collapses the groups + # and gets the output dim wrong new_module = FineGrainedGroupedLinear( in_features_per_group=module.in_features, out_features=module.out_features, diff --git a/src/transformers/quantizers/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py index d2d85ec44e02..6e147592d8e4 100644 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -94,10 +94,9 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): target_patterns="mlp.experts.down_proj_scale_inv", operations=merge, ), - # every second-level global of a layer in ONE converter: the gate|up stack's two - # calibrated halves merge into one global per expert by folding the up half's onto - # the down projection, so the weight globals and the down's input scale have to be - # decided together (`FineGrainedWeightGlobals`) + # every second-level global of a layer in ONE converter: folding the gate|up stack's + # two halves into one moves the up half's onto the down projection, so they decide + # together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ rf"{expert}gate_proj\.weight_scale_2", diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index 94cb85bfb724..eebdb5188a8d 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -730,7 +730,9 @@ def test_every_companion_shards_with_its_expert_weight(self): from transformers.distributed.tensor_parallel import _get_parameter_tp_plan plan = self._planned() - style = lambda name: _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) # noqa: E731 + + def style(name): + return _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_weight_global_scale"): self.assertEqual(style(name), "grouped_gemm", name) @@ -753,7 +755,9 @@ def test_intra_expert_entries_follow_the_projection_axis(self): "layers.*.mlp.experts": "moe_tp_experts", } plan = self._planned(raw=base) - style = lambda n: _get_parameter_tp_plan(f"layers.3.mlp.experts.{n}", plan) # noqa: E731 + + def style(name): + return _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) # the weight keeps the style the MODEL declared — `packed_colwise` is the strided split a # concatenated `[gate; up]` checkpoint needs, since the interleave into the kernels' row From f1b593dc765acbb2ddffff3a609bae2927a3b1b7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 02:50:01 -0700 Subject: [PATCH 52/68] Make tests/kernels a package so the moved fp8 test can import its fixtures `tests/kernels` had no `__init__.py`, so pytest put the directory on `sys.path` and a bare `from test_utils import ...` resolved. Moving `test_finegrained_fp8.py` to `tests/integrations`, which IS a package, took that away and CI failed to collect it. Most of `tests/` is packages already, so make this one too and import through it. Co-Authored-By: Claude Opus 5 --- tests/integrations/test_finegrained_fp8.py | 3 ++- tests/kernels/__init__.py | 0 tests/kernels/test_deepgemm.py | 3 ++- tests/kernels/test_sonicmoe.py | 3 ++- 4 files changed, 6 insertions(+), 3 deletions(-) create mode 100644 tests/kernels/__init__.py diff --git a/tests/integrations/test_finegrained_fp8.py b/tests/integrations/test_finegrained_fp8.py index 2e16b16bcead..2570aaad8ed7 100644 --- a/tests/integrations/test_finegrained_fp8.py +++ b/tests/integrations/test_finegrained_fp8.py @@ -33,7 +33,6 @@ import torch from parameterized import parameterized -from test_utils import make_fp8_experts import transformers.integrations.finegrained_fp8 as fg from transformers.integrations.finegrained_fp8 import ( @@ -43,6 +42,8 @@ ) from transformers.testing_utils import require_torch, torch_device +from ..kernels.test_utils import make_fp8_experts + def _add_one(x, *args, **kwargs): return x + 1 diff --git a/tests/kernels/__init__.py b/tests/kernels/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index aa2b3ea1403a..052f41871f2e 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -43,7 +43,6 @@ import torch from parameterized import parameterized -from test_utils import make_experts, make_fp8_experts import transformers.integrations.deepgemm as dg from transformers.integrations.deepgemm import ( @@ -58,6 +57,8 @@ torch_device, ) +from .test_utils import make_experts, make_fp8_experts + def _add_one(x, *args, **kwargs): return x + 1 diff --git a/tests/kernels/test_sonicmoe.py b/tests/kernels/test_sonicmoe.py index b4eb58db6dcc..3c6ce2567aee 100644 --- a/tests/kernels/test_sonicmoe.py +++ b/tests/kernels/test_sonicmoe.py @@ -32,12 +32,13 @@ import torch from parameterized import parameterized -from test_utils import make_experts import transformers.integrations.sonicmoe as sm from transformers.integrations.sonicmoe import sonicmoe_experts_forward from transformers.testing_utils import require_torch, torch_device +from .test_utils import make_experts + class _FakeActivationType: SWIGLU = "swiglu" From 0067189d63fbbc1ef5983362302225c1320617ee Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 03:12:59 -0700 Subject: [PATCH 53/68] Declare GLM-4v-MoE's head plan instead of deriving it in the mixin `init_parallel_plans` gave a tied head a style when the config had sharded the embedding it reads. Every model is expected to declare that itself in its `ForCausalLM` -- 81 already do -- so the hook only papered over the composites that keep their plan on `text_config`, and it had to reach for `get_text_config()`, which raises on a config with two text sub-configs. Of this branch's load-path fixtures only GLM-4v-MoE ties (the others set `tie_word_embeddings=False` or default to it), so it gets the declaration and the hook goes. The remaining composites missing one are a separate change. Co-Authored-By: Claude Opus 5 --- src/transformers/distributed/mixin.py | 15 --------------- .../models/glm4v_moe/modeling_glm4v_moe.py | 1 + .../models/glm4v_moe/modular_glm4v_moe.py | 2 ++ 3 files changed, 3 insertions(+), 15 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index f228f356e7f1..646b8ac102cc 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -80,21 +80,6 @@ def init_parallel_plans(self) -> None: if plan := getattr(module, "_fsdp_plan", None): self._fsdp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) - # A tied lm_head IS the embedding's parameter, so `embedding_rowwise` shards the tensor - # it reads while lm_head gets no style — leaving `F.linear` a plain input against a - # DTensor weight. The CONFIG flag, not `head.weight is embed.weight`: tying happens in - # `post_init`, after this. And `get_text_config` RAISES on a composite model with two - # text sub-configs (musicgen), which this would hit in every model's `post_init`. - try: - text_config = self.config.get_text_config() - except ValueError: - text_config = self.config - tied = getattr(text_config, "tie_word_embeddings", False) - if tied and "embedding_rowwise" in self._tp_plan.values(): - head = self.get_output_embeddings() - if head is not None: - self._tp_plan.setdefault(next(n for n, m in self.named_modules() if m is head), "colwise_rep") - @property def tp_plan(self) -> dict[str, str]: """The full tp plan for the model's modules.""" diff --git a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py index 9e2cfa588fd0..0d562eb8c54e 100644 --- a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py @@ -1493,6 +1493,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vMoePreTrainedModel, GenerationMixin) _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} # Reference: fix gemma3 grad acc #37208 accepts_loss_kwargs = False + _tp_plan = {"lm_head": "colwise_gather_output"} def __init__(self, config): super().__init__(config) diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 3e8c672788eb..1354adc10d25 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -350,6 +350,8 @@ class Glm4vMoeModelOutputWithPast(Qwen3VLMoeModelOutputWithPast): class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): + _tp_plan = {"lm_head": "colwise_gather_output"} + def __init__(self, config): super().__init__(config) self.num_experts = config.text_config.num_local_experts From bb24085bf9d6cca908fed98ce4c840c65137480e Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 06:36:36 -0700 Subject: [PATCH 54/68] cleanup guards and post expert norm in deepgemm --- src/transformers/integrations/deepgemm.py | 204 +++++++++++------- .../integrations/finegrained/core.py | 4 +- .../integrations/finegrained/test_forwards.py | 2 +- tests/kernels/test_deepgemm.py | 33 +-- 4 files changed, 131 insertions(+), 112 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index ac2a2a90bae6..b707164d6ba6 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -303,42 +303,6 @@ def load_deepgemm_kernel() -> DeepGEMM: # ── Scale-factor helpers ─────────────────────────────────────────────────────── -def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> None: - """Before-load guard for DeepGEMM's FP8/FP4 arch constraints on the given weight/scale dtypes: - - - FP4 (``int8``-packed) weights have no Hopper (SM90) kernel — they need Blackwell (SM100+). - - Blackwell has no float32 scale-factor path: ``_coerce_sf_for_kernel`` would silently round a - ``float32`` scale to UE8M0 and corrupt the output. UE8M0 scales load as ``float8_e8m0fnu`` (the - loader normalizes even float32-container checkpoints like dsv4-flash-base), so a ``float32`` scale - on SM100 means a genuine non-UE8M0 checkpoint. - - Uses `is_sm100()` (compile-safe via `assume_constant_result`), so the whole guard folds to a constant under - ``torch.compile``: the valid case compiles away to nothing, while an unsupported combo fails loud - rather than letting the hot path silently corrupt (unlike an ``is_compiling`` skip, which would miss a - model compiled from cold with no eager warmup). Both raise ``NotImplementedError``, which - ``fp8_linear`` treats as "DeepGEMM declined" and falls back to Triton (SM90 consuming float32 SFs - directly is fine, so those cases are no-ops). - """ - if not is_sm100(): - # SM90: DeepGEMM has no FP4 (int8-packed) kernel, but consumes float32 SFs directly. - if weight.dtype == torch.int8: - raise NotImplementedError( - "DeepGEMM's FP4 (int8-packed) path requires a Blackwell (SM100+) GPU; FP4 weights have no " - "Hopper (SM90) kernel. Use an FP8 checkpoint, or run on a Blackwell GPU." - ) - return - - # SM100: DeepGEMM has no float32 scale-factor path. - if scale.dtype == torch.float32: - raise NotImplementedError( - "DeepGEMM has no float32 scale-factor path on Blackwell (SM100): these scales are plain float32 " - "(quantization_config.scale_fmt='float'), and rounding them to UE8M0 would silently corrupt the " - "output. Use a checkpoint quantized with scale_fmt='ue8m0', or a path that consumes float32 block " - "scales directly — the FP8 linear falls back to Triton automatically; for experts use " - "`model.set_experts_implementation('grouped_mm')`." - ) - - def _ceil_to_ue8m0(sf: torch.Tensor) -> torch.Tensor: """Round each fp32 SF up to the nearest power of 2 (zero mantissa). @@ -570,6 +534,7 @@ def _combine_routed_output( def prefers_deepgemm_linear( + input: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor, *, @@ -592,6 +557,8 @@ def prefers_deepgemm_linear( # False when the model spans devices: DeepGEMM's kernels are bound to one CUDA context and # corrupt across them, and the Triton fallback is context-free. and allow_deepgemm + # DeepGEMM has no backward pass; a call that needs a gradient takes Triton, which has one + and not (torch.is_grad_enabled() and input.requires_grad) # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. and weight_scale_inv.ndim <= 2 @@ -608,48 +575,76 @@ def prefers_deepgemm_linear( ) -def deepgemm_fp8_fp4_linear( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale_inv: torch.Tensor, - bias: torch.Tensor | None = None, - block_size: tuple[int, int] | None = None, - output_dtype: torch.dtype | None = None, - activation_scale: torch.Tensor | None = None, -) -> torch.Tensor: - """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul. +def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> None: + """Before-load guard for DeepGEMM's FP8/FP4 arch constraints on the given weight/scale dtypes: + + - FP4 (``int8``-packed) weights have no Hopper (SM90) kernel — they need Blackwell (SM100+). + - Blackwell has no float32 scale-factor path: ``_coerce_sf_for_kernel`` would silently round a + ``float32`` scale to UE8M0 and corrupt the output. UE8M0 scales load as ``float8_e8m0fnu`` (the + loader normalizes even float32-container checkpoints like dsv4-flash-base), so a ``float32`` scale + on SM100 means a genuine non-UE8M0 checkpoint. - Static (per-tensor) activation quantization is rejected — DeepGEMM needs - per-row SFs. Callers should route static activations through the Triton fallback. + Uses `is_sm100()` (compile-safe via `assume_constant_result`), so the whole guard folds to a constant under + ``torch.compile``: the valid case compiles away to nothing, while an unsupported combo fails loud + rather than letting the hot path silently corrupt (unlike an ``is_compiling`` skip, which would miss a + model compiled from cold with no eager warmup). Both raise ``NotImplementedError``, which + ``fp8_linear`` treats as "DeepGEMM declined" and falls back to Triton (SM90 consuming float32 SFs + directly is fine, so those cases are no-ops). """ + if not is_sm100(): + # SM90: DeepGEMM has no FP4 (int8-packed) kernel, but consumes float32 SFs directly. + if weight.dtype == torch.int8: + raise NotImplementedError( + "DeepGEMM's FP4 (int8-packed) path requires a Blackwell (SM100+) GPU; FP4 weights have no " + "Hopper (SM90) kernel. Use an FP8 checkpoint, or run on a Blackwell GPU." + ) + return + + # SM100: DeepGEMM has no float32 scale-factor path. + if scale.dtype == torch.float32: + raise NotImplementedError( + "DeepGEMM has no float32 scale-factor path on Blackwell (SM100): these scales are plain float32 " + "(quantization_config.scale_fmt='float'), and rounding them to UE8M0 would silently corrupt the " + "output. Use a checkpoint quantized with scale_fmt='ue8m0', or a path that consumes float32 block " + "scales directly — the FP8 linear falls back to Triton automatically; for experts use " + "`model.set_experts_implementation('grouped_mm')`." + ) + + +def _assert_no_gradient(hidden_states: torch.Tensor, backend: str) -> None: + """DeepGEMM has no backward pass, so a call that needs a gradient fails rather than drop it.""" + if torch.is_grad_enabled() and hidden_states.requires_grad: + raise NotImplementedError( + f"{backend} has no backward pass; train through the Triton kernels " + "(for experts, `experts_implementation='grouped_mm'` or 'batched_mm')." + ) + + +def _assert_dynamic_activation_scale(activation_scale: torch.Tensor | None) -> None: + """DeepGEMM quantizes activations per row, so a calibrated per-tensor scale has nowhere to go.""" if activation_scale is not None: raise NotImplementedError("DeepGEMM linear does not support static activation quantization.") + + +def _assert_half_precision_input(input: torch.Tensor) -> None: if input.dtype not in (torch.bfloat16, torch.float16): raise ValueError(f"DeepGEMM linear requires FP16 or BF16 activations, got {input.dtype}") - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. - _assert_sm100_requirements(weight, weight_scale_inv) - deepgemm = load_deepgemm_kernel() - cast_kwargs = _select_fp8_cast_kwargs(weight, weight_scale_inv, block_size, is_sm100()) +def _assert_no_low_rank_adapters(self: torch.nn.Module) -> None: + """DeepGEMM runs the base weights alone, so an active adapter would be silently dropped.""" + if any(getattr(self, "low_rank_adapters", {}).values()): + raise NotImplementedError( + "DeepGEMM experts cannot apply low-rank adapters; use " + "`experts_implementation='grouped_mm'` (or 'batched_mm'), or merge the adapter." + ) - input_2d = input.view(-1, input.shape[-1]) - qinput_2d, scale_2d = deepgemm.per_token_cast_to_fp8(input_2d, **cast_kwargs) - output = torch.empty(qinput_2d.shape[0], weight.shape[0], device=input.device, dtype=input.dtype) - # Pass `(1, 1, gran_k)` for int-SF paths so the kernel uses the right K granularity - # (the default `(1, 1, 128)` mismatches FP4's gran_k=32). Float-SF leaves it None. - sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None - deepgemm.fp8_fp4_matmul( - (qinput_2d, _coerce_sf_for_kernel(scale_2d, is_sm100(), expected_mn=qinput_2d.size(0))), - (weight, _coerce_sf_for_kernel(weight_scale_inv, is_sm100(), expected_mn=weight.size(0))), - output, - recipe=sf_recipe, - ) - output = output.view(input.shape[:-1] + (weight.shape[0],)) - if bias is not None: - output.add_(bias) - return output +def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: + """Every DeepGEMM experts arm builds its intermediates and its output in bfloat16, so anything + else silently changes the dtype the caller gets back.""" + if hidden_states.dtype != torch.bfloat16: + raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") def _assert_dynamic_activations(self: torch.nn.Module) -> None: @@ -699,19 +694,27 @@ def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: ) -def _apply_post_expert_norm(self: torch.nn.Module, rows: torch.Tensor) -> torch.Tensor: - """The model's per-expert output norm, on the routed rows before the routing weights — the - same place the reference forwards apply it. A no-op for a model that declares none.""" - if not getattr(self, "has_post_expert_norm", False): - return rows - return self.post_expert_norm(rows) +def deepgemm_linear_guards(forward=None): + """State a linear's requirements on its operands, checked before any kernel work. + Usable bare (`@deepgemm_linear_guards`), like `deepgemm_experts_guards`. -def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: - """Every DeepGEMM experts arm builds its intermediates and its output in bfloat16, so anything - else silently changes the dtype the caller gets back.""" - if hidden_states.dtype != torch.bfloat16: - raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") + A static ``activation_scale`` is refused — DeepGEMM needs per-row SFs, so callers route it + through the Triton fallback — and the SM100 check fails before the hub download + JIT when the + device cannot serve these dtypes.""" + + def decorate(forward): + @functools.wraps(forward) + def guarded(input, weight, weight_scale_inv, *args, activation_scale=None, **kwargs): + _assert_half_precision_input(input) + _assert_no_gradient(input, "DeepGEMM linear") + _assert_dynamic_activation_scale(activation_scale) + _assert_sm100_requirements(weight, weight_scale_inv) + return forward(input, weight, weight_scale_inv, *args, **kwargs) + + return guarded + + return decorate(forward) if forward is not None else decorate def deepgemm_experts_guards( @@ -736,6 +739,8 @@ def decorate(forward): @functools.wraps(forward) def guarded(self, hidden_states, *args, **kwargs): _assert_bf16_hidden_states(hidden_states) + _assert_no_gradient(hidden_states, "DeepGEMM experts") + _assert_no_low_rank_adapters(self) _assert_dynamic_activations(self) if not supports_post_expert_norm: _assert_no_post_expert_norm(self) @@ -752,6 +757,37 @@ def guarded(self, hidden_states, *args, **kwargs): return decorate(forward) if forward is not None else decorate +@deepgemm_linear_guards +def deepgemm_fp8_fp4_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + bias: torch.Tensor | None = None, + block_size: tuple[int, int] | None = None, +) -> torch.Tensor: + """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul.""" + deepgemm = load_deepgemm_kernel() + cast_kwargs = _select_fp8_cast_kwargs(weight, weight_scale_inv, block_size, is_sm100()) + + input_2d = input.view(-1, input.shape[-1]) + qinput_2d, scale_2d = deepgemm.per_token_cast_to_fp8(input_2d, **cast_kwargs) + output = torch.empty(qinput_2d.shape[0], weight.shape[0], device=input.device, dtype=input.dtype) + + # Pass `(1, 1, gran_k)` for int-SF paths so the kernel uses the right K granularity + # (the default `(1, 1, 128)` mismatches FP4's gran_k=32). Float-SF leaves it None. + sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None + deepgemm.fp8_fp4_matmul( + (qinput_2d, _coerce_sf_for_kernel(scale_2d, is_sm100(), expected_mn=qinput_2d.size(0))), + (weight, _coerce_sf_for_kernel(weight_scale_inv, is_sm100(), expected_mn=weight.size(0))), + output, + recipe=sf_recipe, + ) + output = output.view(input.shape[:-1] + (weight.shape[0],)) + if bias is not None: + output.add_(bias) + return output + + @deepgemm_experts_guards def deepgemm_bf16_experts_forward( self: torch.nn.Module, @@ -806,8 +842,11 @@ def deepgemm_bf16_experts_forward( if self.has_bias: out.index_add_(0, sorted_to_padded, down_bias[expert_ids_g]) + if getattr(self, "has_post_expert_norm", False): + out = self.post_expert_norm(out) + return _combine_routed_output( - _apply_post_expert_norm(self, out), + out, sorted_weights, sentinel_mask, perm, @@ -897,8 +936,11 @@ def deepgemm_fp8_fp4_experts_forward( use_psum_layout=is_sm100(), ) + if getattr(self, "has_post_expert_norm", False): + out = self.post_expert_norm(out) + return _combine_routed_output( - _apply_post_expert_norm(self, out), + out, sorted_weights, sentinel_mask, perm, diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 65c5ddb35c2b..146f4b07495e 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -309,7 +309,8 @@ def finegrained_linear( pre-SM100 shapes it supports (FP4, UE8M0 SFs, 128×128 block FP8). Never on SM100, where Triton is tuned per shape and is the only path reading pre-swizzled scales. 2. Triton finegrained fallback everywhere else: SM100, an ``activation_scale`` - (DeepGEMM is dynamic-only), or any shape DeepGEMM declined. + (DeepGEMM is dynamic-only), a call that needs a gradient (DeepGEMM has no backward), + or any shape DeepGEMM declined. Args: input: (..., K) bf16/fp16 activations. @@ -333,6 +334,7 @@ def finegrained_linear( (``"bf16"`` = weight-only). """ if prefers_deepgemm_linear( + input, weight, weight_scale_inv, block_size=block_size, diff --git a/tests/integrations/finegrained/test_forwards.py b/tests/integrations/finegrained/test_forwards.py index b29574aee63c..c4315c45e449 100644 --- a/tests/integrations/finegrained/test_forwards.py +++ b/tests/integrations/finegrained/test_forwards.py @@ -246,7 +246,7 @@ def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): row = inter @ dequantized["down_proj"][e].t() if post_norm is not None: row = post_norm(row.to(x.dtype)).float() - out[token] += row * float(wts[token, slot]) + out[token] += row * wts[token, slot].float() return out def test_modelopt_nvfp4_experts_merge_their_globals_and_fuse_the_post_norm(self): diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 052f41871f2e..93f60841361a 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -380,37 +380,12 @@ class _Experts(torch.nn.Module): forward(_Experts(), torch.zeros(2, 4, dtype=torch.float16), None, None) self.assertIn("bfloat16", str(caught.exception)) - def test_post_expert_norm_runs_on_the_reducing_arms_and_is_refused_by_megamoe(self): - """Mega MoE fuses the routing-weighted reduce, so a per-expert output norm has nowhere to - go and is refused. The other arms reduce in `_combine_routed_output`, so the norm rides the - rows just before it — the same place the reference forwards apply it.""" + def test_megamoe_refuses_a_post_expert_norm(self): + """Mega MoE fuses the routing-weighted reduce, so a per-expert output norm has nowhere to go.""" import transformers.integrations.deepgemm as dg - applied = [] - - class _Norm(torch.nn.Module): - def forward(self, rows): - applied.append(tuple(rows.shape)) - return rows - - class _Experts(torch.nn.Module): - has_post_expert_norm = True - - def __init__(self): - super().__init__() - self.post_expert_norm = _Norm() - - module = _Experts() - rows = torch.randn(4, 8) - # the reducing arms: the norm is applied and the rows pass through - self.assertIs(dg._apply_post_expert_norm(module, rows), rows) - self.assertEqual(applied, [(4, 8)]) - - # a model declaring no norm is untouched - plain = torch.nn.Module() - self.assertIs(dg._apply_post_expert_norm(plain, rows), rows) - - # mega moe still refuses, and only mega moe + module = torch.nn.Module() + module.has_post_expert_norm = True with self.assertRaises(NotImplementedError): dg._assert_no_post_expert_norm(module) From 39eb1d192e3c1ca32513bf076adebcc5732834b0 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 06:56:30 -0700 Subject: [PATCH 55/68] not now --- src/transformers/integrations/deepgemm.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index b707164d6ba6..7f5190e6fc12 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -631,15 +631,6 @@ def _assert_half_precision_input(input: torch.Tensor) -> None: raise ValueError(f"DeepGEMM linear requires FP16 or BF16 activations, got {input.dtype}") -def _assert_no_low_rank_adapters(self: torch.nn.Module) -> None: - """DeepGEMM runs the base weights alone, so an active adapter would be silently dropped.""" - if any(getattr(self, "low_rank_adapters", {}).values()): - raise NotImplementedError( - "DeepGEMM experts cannot apply low-rank adapters; use " - "`experts_implementation='grouped_mm'` (or 'batched_mm'), or merge the adapter." - ) - - def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: """Every DeepGEMM experts arm builds its intermediates and its output in bfloat16, so anything else silently changes the dtype the caller gets back.""" @@ -740,7 +731,6 @@ def decorate(forward): def guarded(self, hidden_states, *args, **kwargs): _assert_bf16_hidden_states(hidden_states) _assert_no_gradient(hidden_states, "DeepGEMM experts") - _assert_no_low_rank_adapters(self) _assert_dynamic_activations(self) if not supports_post_expert_norm: _assert_no_post_expert_norm(self) From bf0e7ec476ccb0cd56d1f07a3113492b68a2704e Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 24 Sep 2026 08:52:00 -0700 Subject: [PATCH 56/68] minimal peft integration preparation --- .../integrations/finegrained/core.py | 165 +++++++++--------- tests/integrations/finegrained/test_core.py | 38 +++- 2 files changed, 111 insertions(+), 92 deletions(-) diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 146f4b07495e..2242f4d8dd2c 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -536,69 +536,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -def _moe_operands(kernel, module) -> dict: - """Everything the kernels' MoE forwards take from the module: the two projections with their - scales (held swizzled for dot_scaled chains), the per-expert globals and biases, - and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the - module's own GLU as a callable the kernels run on the host, so any activation works without a - kernel change.""" - up = module.gate_up_name - - act_name = module.act_fn_name - fused = act_name in kernel.get_supported_act_fns() and (module.swiglu_alpha is None or act_name == "silu") - if fused: - act_fn = act_name - else: - act_fn = module._apply_gate if module.has_gate else module.act_fn - - # A named norm is fused into the reduce; anything else runs as a callable on the routed rows. - # Fusing is off under intra-expert TP: those rows are a partial sum, so the all-reduce the - # module's wrapped forward does has to happen first (`_hf_tp_input_reduce` marks it). - post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 - if module.has_post_expert_norm: - fusable = module.post_expert_norm_name in kernel.get_supported_norms() - if not fusable or getattr(module.post_expert_norm, "_hf_tp_input_reduce", False): - post_expert_norm = module.post_expert_norm - else: - norm = module.post_expert_norm - post_expert_norm, norm_weight = module.post_expert_norm_name, norm.weight - # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves - # it None and resolves to the dtype's own - eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) - norm_eps = float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps - - return { - **{ - f"{kernel_proj}{slot}": getattr(module, f"{held_proj}{slot}") - for kernel_proj, held_proj in (("gate_up_proj", up), ("down_proj", "down_proj")) - for slot in ("", "_scale_inv", "_weight_global_scale", "_input_global_scale", "_activation_scale", "_bias") - }, - # a supported activation NAME is fused into the gate_up epilogue; any other callable - # leaves that GEMM plain and runs on the host between the two - "act_fn": act_fn, - "gate": module.has_gate, - "swiglu_alpha": module.swiglu_alpha, - "swiglu_limit": module.swiglu_limit, - "activation_format": module.activation_format, - # a supported norm NAME is folded into the routing-weighted reduce, one pass ahead of it - # computing the rows' rsqrt; any other callable is applied to the routed rows before it - "post_expert_norm": post_expert_norm, - "post_expert_norm_weight": norm_weight, - "post_expert_norm_eps": norm_eps, - } - - def finegrained_batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: """Batched (decode) experts forward: one program per routed row.""" kernel = load_finegrained_kernel() - return kernel.moe_fused_batched(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, self)) + return kernel.moe_fused_batched(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: """Grouped (prefill) experts forward: one on-device routing pass, then the chain over the expert-sorted schedule, scattering back to routed rows (EP-sentinel rows skipped).""" kernel = load_finegrained_kernel() - return kernel.moe_fused_grouped(hidden_states, top_k_index, top_k_weights, **_moe_operands(kernel, self)) + return kernel.moe_fused_grouped(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) class FineGrainedExperts(_FineGrainedModule, nn.Module): @@ -607,6 +555,9 @@ class FineGrainedExperts(_FineGrainedModule, nn.Module): _impl_tp_layer_overrides: dict[str, dict[str, str]] = { "deepgemm_megamoe": {"moe_tp_experts": "megamoe_experts", "ep_router": "megamoe_router"}, } + # a projection's operands, by the suffix they carry on the module and in the kernels' arguments: + # the weight, its scale grid, the NVFP4 weight and activation globals, a static activation scale, a bias + _projection_slots = ("", "_scale_inv", "_weight_global_scale", "_input_global_scale", "_activation_scale", "_bias") def __init__( self, @@ -727,18 +678,9 @@ def forward( continue top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) - current_state = hidden_states[token_idx] - gate_up_act_scale = ( - self.gate_up_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None - ) - proj_out = self.expert_linear( - current_state, self.gate_up_name, expert_idx, activation_scale=gate_up_act_scale - ) + proj_out = self.expert_linear(hidden_states[token_idx], self.gate_up_name, expert_idx) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - down_act_scale = ( - self.down_proj_activation_scale[expert_idx] if self.activation_scheme == "static" else None - ) - proj_out = self.expert_linear(proj_out, "down_proj", expert_idx, activation_scale=down_act_scale) + proj_out = self.expert_linear(proj_out, "down_proj", expert_idx) if self.has_post_expert_norm: proj_out = self.post_expert_norm(proj_out) routing_weights = top_k_weights[token_idx, top_k_pos, None] @@ -746,39 +688,88 @@ def forward( final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) return final_hidden_states.to(hidden_states.dtype) - def expert_linear( - self, input: torch.Tensor, proj: str, expert_idx: int, activation_scale: torch.Tensor | None = None - ) -> torch.Tensor: - """One expert's ``proj`` as a dense linear: the same weight, scale (one expert's slice of a - swizzled stack is the ``(1, ...)`` artifact ``matmul_2d`` reads), bias, NVFP4 global and - activation format the fused forwards hand the kernels.""" - weight = getattr(self, proj)[expert_idx] - bias = getattr(self, f"{proj}_bias") - bias = bias[expert_idx] if bias is not None else None - scale = getattr(self, f"{proj}_scale_inv") - weight_global_scale = getattr(self, f"{proj}_weight_global_scale") - weight_global_scale = weight_global_scale[expert_idx] if weight_global_scale is not None else None - # None under weight-only, which registers no input global at all. `% numel` picks this - # expert's, or the single one gate_up holds because its rows are quantized pre-routing. - input_global_scale = getattr(self, f"{proj}_input_global_scale") - input_global_scale = ( - input_global_scale.reshape(-1)[expert_idx % input_global_scale.numel()].reshape(1) - if input_global_scale is not None - else None + def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torch.Tensor: + """One expert's ``proj`` as a dense linear, over that expert's slice of the projection's operands.""" + weight, scale, weight_global, input_global, activation_scale, bias = ( + getattr(self, f"{proj}{slot}") for slot in self._projection_slots + ) + weight = weight[expert_idx] + # one expert's slice of a swizzled stack is the `(1, ...)` artifact `matmul_2d` reads + scale = scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx] + # None where the checkpoint fills no such slot + weight_global, activation_scale, bias = ( + operand[expert_idx] if operand is not None else None for operand in (weight_global, activation_scale, bias) ) + # gate_up holds ONE input global, its rows being quantized pre-routing; down one per expert + if input_global is not None and proj == "down_proj": + input_global = input_global[expert_idx : expert_idx + 1] + return finegrained_linear( input, weight, - scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx], + scale, self.block_size, bias=bias, activation_scale=activation_scale, allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=weight_global_scale, - input_global_scale=input_global_scale, + weight_global_scale=weight_global, + input_global_scale=input_global, activation_format=self.activation_format, ) + def _moe_operands(self, kernel) -> dict: + """Everything the kernels' MoE forwards take from the module: the two projections with their + scales (held swizzled for dot_scaled chains), the per-expert globals and biases, + and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the + module's own GLU as a callable the kernels run on the host, so any activation works without a + kernel change. + + A method rather than a function so an adapter library can extend it per module: PEFT overrides + it to add its LoRA factors, which the kernels then apply on the routed rows between the GEMMs + (`expert_linear` is the eager loop's counterpart).""" + act_name = self.act_fn_name + fused = act_name in kernel.get_supported_act_fns() and (self.swiglu_alpha is None or act_name == "silu") + if fused: + act_fn = act_name + else: + act_fn = self._apply_gate if self.has_gate else self.act_fn + + # A named norm is fused into the reduce; anything else runs as a callable on the routed rows. + # Fusing is off under intra-expert TP: those rows are a partial sum, so the all-reduce the + # module's wrapped forward does has to happen first (`_hf_tp_input_reduce` marks it). + post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 + if self.has_post_expert_norm: + fusable = self.post_expert_norm_name in kernel.get_supported_norms() + if not fusable or getattr(self.post_expert_norm, "_hf_tp_input_reduce", False): + post_expert_norm = self.post_expert_norm + else: + norm = self.post_expert_norm + post_expert_norm, norm_weight = self.post_expert_norm_name, norm.weight + # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves + # it None and resolves to the dtype's own + eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) + norm_eps = float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps + + return { + **{ + f"{kernel_proj}{slot}": getattr(self, f"{held_proj}{slot}") + for kernel_proj, held_proj in (("gate_up_proj", self.gate_up_name), ("down_proj", "down_proj")) + for slot in self._projection_slots + }, + # a supported activation NAME is fused into the gate_up epilogue; any other callable + # leaves that GEMM plain and runs on the host between the two + "act_fn": act_fn, + "gate": self.has_gate, + "swiglu_alpha": self.swiglu_alpha, + "swiglu_limit": self.swiglu_limit, + "activation_format": self.activation_format, + # a supported norm NAME is folded into the routing-weighted reduce, one pass ahead of it + # computing the rows' rsqrt; any other callable is applied to the routed rows before it + "post_expert_norm": post_expert_norm, + "post_expert_norm_weight": norm_weight, + "post_expert_norm_eps": norm_eps, + } + class FineGrainedExpertsInterface(ExpertsInterface): """Interface for registering custom FP8 experts forward functions.""" diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index eebdb5188a8d..b8a218ce7106 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -509,7 +509,7 @@ def test_static_activation_scheme_reaches_the_fused_chain(self): kernel.get_supported_act_fns.return_value = ("silu",) kernel.get_supported_norms.return_value = () for module, calibrated in ((static, True), (dynamic, False)): - operands = fg._moe_operands(kernel, module) + operands = module._moe_operands(kernel) for projection in ("gate_up_proj", "down_proj"): with self.subTest(scheme=module.activation_scheme, projection=projection): scale = operands[f"{projection}_activation_scale"] @@ -607,10 +607,11 @@ def _operands(self, *, input_reduce: bool): module.has_post_expert_norm = True module.post_expert_norm = norm module.post_expert_norm_name = "input_scaled_rms_norm" + module._projection_slots = fg.FineGrainedExperts._projection_slots kernel = mock.Mock() kernel.get_supported_act_fns.return_value = ("silu",) kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) - return fg._moe_operands(kernel, module), module + return fg.FineGrainedExperts._moe_operands(module, kernel), module def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): operands, _ = self._operands(input_reduce=False) @@ -1037,8 +1038,9 @@ class FineGrainedDeepGemmDispatchTest(unittest.TestCase): and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 scales have no swizzled layout, so the first never fires for the shape the second catches.""" - def _routed_to(self, *, sm100, scale_ndim): - """Which backend a block-FP8 linear reaches, given arch and scale layout.""" + def _routed_to(self, *, sm100, scale_ndim, trains=False): + """Which backend a block-FP8 linear reaches, given arch, scale layout, and whether the call + needs a gradient.""" kernel, _ = _fake_bundle() w = torch.randn(32, 64).to(torch.float8_e4m3fn) s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) @@ -1052,7 +1054,8 @@ def _routed_to(self, *, sm100, scale_ndim): mock.patch.object(deepgemm, "is_sm100", return_value=sm100), mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, ): - fg.finegrained_linear(torch.randn(4, 64, dtype=torch.bfloat16), w, s, block_size=[128, 128]) + x = torch.randn(4, 64, dtype=torch.bfloat16, requires_grad=trains) + fg.finegrained_linear(x, w, s, block_size=[128, 128]) return "deepgemm" if dg.called else "triton" def test_sm100_never_prefers_deepgemm(self): @@ -1061,6 +1064,31 @@ def test_sm100_never_prefers_deepgemm(self): def test_pre_sm100_still_uses_deepgemm(self): self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") + def test_deepgemm_refuses_a_call_that_needs_a_gradient(self): + """DeepGEMM has no backward pass: every entry point fails a call that needs a gradient, + before touching the kernel.""" + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size = 256, 128 + experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8") + experts._deepgemm_disabled = False + hs = torch.randn(2, 256, dtype=torch.bfloat16, requires_grad=True) + routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + w, s = torch.randn(32, 256).to(torch.float8_e4m3fn), torch.ones(1, 2) + calls = [ + ("DeepGEMM linear", lambda: deepgemm.deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128))), + ("DeepGEMM experts", lambda: deepgemm.deepgemm_fp8_fp4_experts_forward(experts, hs, *routing)), + ("DeepGEMM experts", lambda: deepgemm.deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing)), + ] + with mock.patch.object(deepgemm, "load_deepgemm_kernel") as load: + for backend, call in calls: + with self.assertRaisesRegex(NotImplementedError, f"{backend} has no backward pass"): + call() + load.assert_not_called() + + def test_a_linear_that_needs_a_gradient_prefers_triton(self): + """Where DeepGEMM would otherwise win, a call that needs a gradient still goes to triton.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2, trains=True), "triton") + def test_swizzled_scales_never_reach_deepgemm(self): """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") From 5b2b7f00197b136f7956a19cff9045ef2e4884fb Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 25 Sep 2026 01:00:51 -0700 Subject: [PATCH 57/68] addressing review comments --- src/transformers/integrations/deepgemm.py | 41 +++----- .../integrations/finegrained/core.py | 32 +------ src/transformers/integrations/moe.py | 8 +- .../quantizers/finegrained/base.py | 17 ++-- .../quantizers/finegrained/mxfp4.py | 3 +- .../quantizers/finegrained/nvfp4.py | 34 +++---- .../quantizers/quantizer_finegrained_fp8.py | 3 - .../quantizers/quantizer_nvfp4.py | 8 +- tests/integrations/finegrained/test_core.py | 93 +++++++++---------- tests/integrations/test_moe.py | 1 + tests/kernels/test_deepgemm.py | 1 + tests/kernels/test_utils.py | 1 + 12 files changed, 93 insertions(+), 149 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 7f5190e6fc12..ac03823b3779 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -36,7 +36,6 @@ import torch -from ..utils import logging from ..utils.import_utils import ( is_kernels_available, maybe_import_error, @@ -45,8 +44,6 @@ from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel -logger = logging.get_logger(__name__) - # ── Kernel loading ───────────────────────────────────────────────────────────── @@ -611,11 +608,11 @@ def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> Non ) -def _assert_no_gradient(hidden_states: torch.Tensor, backend: str) -> None: +def _assert_no_gradient(hidden_states: torch.Tensor) -> None: """DeepGEMM has no backward pass, so a call that needs a gradient fails rather than drop it.""" if torch.is_grad_enabled() and hidden_states.requires_grad: raise NotImplementedError( - f"{backend} has no backward pass; train through the Triton kernels " + "DeepGEMM does not support training; train through the Triton kernels " "(for experts, `experts_implementation='grouped_mm'` or 'batched_mm')." ) @@ -678,7 +675,7 @@ def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: """Mega MoE fuses the routing-weighted reduce into its kernel, so a model's per-expert output norm has nowhere to go — it would be dropped silently, so refuse instead. The other DeepGEMM arms reduce in `_combine_routed_output` and apply the norm on the rows just before it.""" - if getattr(self, "has_post_expert_norm", False): + if self.has_post_expert_norm: raise NotImplementedError( "DeepGEMM Mega MoE cannot apply this model's per-expert output norm; use " "`experts_implementation='deepgemm'`, 'grouped_mm' or 'batched_mm'." @@ -686,19 +683,14 @@ def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: def deepgemm_linear_guards(forward=None): - """State a linear's requirements on its operands, checked before any kernel work. - - Usable bare (`@deepgemm_linear_guards`), like `deepgemm_experts_guards`. - - A static ``activation_scale`` is refused — DeepGEMM needs per-row SFs, so callers route it - through the Triton fallback — and the SM100 check fails before the hub download + JIT when the - device cannot serve these dtypes.""" + """State a linear's requirements on its operands, checked before any kernel work; each assert + documents its own condition. Usable bare (`@deepgemm_linear_guards`), like `deepgemm_experts_guards`.""" def decorate(forward): @functools.wraps(forward) def guarded(input, weight, weight_scale_inv, *args, activation_scale=None, **kwargs): + _assert_no_gradient(input) _assert_half_precision_input(input) - _assert_no_gradient(input, "DeepGEMM linear") _assert_dynamic_activation_scale(activation_scale) _assert_sm100_requirements(weight, weight_scale_inv) return forward(input, weight, weight_scale_inv, *args, **kwargs) @@ -716,22 +708,17 @@ def deepgemm_experts_guards( supports_post_expert_norm: bool = True, stacked_gate_up: bool = False, ): - """State an experts forward's requirements on the module, checked before any kernel work. - - Usable bare (`@deepgemm_experts_guards`) for an arm that only needs the common checks. - - ``supports_post_expert_norm=False`` refuses a model that declares a per-expert output norm — - Mega MoE only, whose fused reduce has no seam for it. ``affine_scales`` refuses the SWIZZLE_32_4_4 - scales a module loaded for a triton backend holds, ``stacked_gate_up`` refuses the interleaved - rows that same load produces, and ``sm100`` fails before the hub download + JIT when the device - cannot serve these dtypes.""" + """State an experts forward's requirements on the module, checked before any kernel work; each + assert documents its own condition, and each flag enables one (``supports_post_expert_norm=False`` + the post-expert-norm refusal). Usable bare (`@deepgemm_experts_guards`) for an arm that only needs + the common checks.""" def decorate(forward): @functools.wraps(forward) def guarded(self, hidden_states, *args, **kwargs): - _assert_bf16_hidden_states(hidden_states) - _assert_no_gradient(hidden_states, "DeepGEMM experts") _assert_dynamic_activations(self) + _assert_no_gradient(hidden_states) + _assert_bf16_hidden_states(hidden_states) if not supports_post_expert_norm: _assert_no_post_expert_norm(self) if affine_scales: @@ -832,7 +819,7 @@ def deepgemm_bf16_experts_forward( if self.has_bias: out.index_add_(0, sorted_to_padded, down_bias[expert_ids_g]) - if getattr(self, "has_post_expert_norm", False): + if self.has_post_expert_norm: out = self.post_expert_norm(out) return _combine_routed_output( @@ -926,7 +913,7 @@ def deepgemm_fp8_fp4_experts_forward( use_psum_layout=is_sm100(), ) - if getattr(self, "has_post_expert_norm", False): + if self.has_post_expert_norm: out = self.post_expert_norm(out) return _combine_routed_output( diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 2242f4d8dd2c..2b1fce8152ba 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -19,9 +19,6 @@ from __future__ import annotations import functools -import importlib -import os -import sys from collections.abc import Callable from dataclasses import dataclass @@ -114,26 +111,6 @@ class FineGrained: _FINEGRAINED: FineGrained | None = None -def _import_local_finegrained(): - """A locally importable `finegrained_kernels` package takes precedence over the hub build: - `FINEGRAINED_KERNELS_PATH` (a checkout's `torch-ext` directory, or any directory containing - the package) is prepended to `sys.path`, then a plain import is attempted either way — an - installed / already-on-path package also wins. Returns the module, or `None` to fall back - to the `kernels` hub load.""" - path = os.environ.get("FINEGRAINED_KERNELS_PATH") - if path: - if not os.path.isdir(path): - raise ImportError(f"FINEGRAINED_KERNELS_PATH does not exist: {path}") - if path not in sys.path: - sys.path.insert(0, path) - try: - return importlib.import_module("finegrained_kernels") - except ImportError: - if path: - raise # an explicit local path that fails to import is a setup error, not a fallback - return None - - @torch._dynamo.allow_in_graph def _load_finegrained_kernel() -> None: """ @@ -154,11 +131,10 @@ def _load_finegrained_kernel() -> None: if _FINEGRAINED is not None: return - kernel = _import_local_finegrained() - if kernel is None: - if not is_kernels_available(): - raise ImportError(f"finegrained-fp8 kernel unavailable: {_MISSING_KERNELS_MESSAGE}") - kernel = lazy_load_kernel("finegrained-kernels") + if not is_kernels_available(): + raise ImportError(f"finegrained-fp8 kernel unavailable: {_MISSING_KERNELS_MESSAGE}") + + kernel = lazy_load_kernel("finegrained-kernels") if kernel is None: raise ImportError( "Failed to load the finegrained-kernels kernel — check that `kernels-community/finegrained-kernels` " diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index af95f980c9bd..81ba67fcc31f 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -165,9 +165,7 @@ def batched_mm_experts_forward( ) # (S, hidden_dim) # Normalize each expert application, where such a norm is defined, before it is weighted. - # `getattr`: an experts module built without the decorator has no such attribute, and this - # forward is the shared one every backend adapts onto. - if getattr(self, "has_post_expert_norm", False): + if self.has_post_expert_norm: proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) # Apply routing weights @@ -482,9 +480,7 @@ def grouped_mm_experts_forward( proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Normalize each expert application, where such a norm is defined, before it is weighted. - # `getattr`: an experts module built without the decorator has no such attribute, and this - # forward is the shared one every backend adapts onto. - if getattr(self, "has_post_expert_norm", False): + if self.has_post_expert_norm: proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) # Apply routing weights diff --git a/src/transformers/quantizers/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py index 08a4ad1019f0..03c52423f1d9 100644 --- a/src/transformers/quantizers/finegrained/base.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -59,11 +59,7 @@ class FineGrainedHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedConfig" - - def _default_activation_format(self) -> str | None: - """What to quantize activations to when the config leaves it open. `None` keeps the - kernels' weight-native choice; the MXFP4 arm overrides this to run weight-only.""" - return None + default_activation_format: str | None = None @property def supports_dequantize(self) -> bool: @@ -74,7 +70,7 @@ def supports_dequantize(self) -> bool: def _assert_dequantize_supported(self) -> None: if not self.supports_dequantize: raise NotImplementedError( - f"`dequantize=True` is not supported for {self._quant_method()} checkpoints. Load them " + f"`dequantize=True` is not supported for {self.quant_method} checkpoints. Load them " "quantized on a GPU that serves the format, or start from a bf16 checkpoint." ) @@ -184,7 +180,7 @@ def _process_model_before_weight_loading( expert_dtype = getattr(model.config.get_text_config(), "expert_dtype", None) self.quantization_config.groups = groups_with_expert_dtype(self.quantization_config.groups, expert_dtype) if self.quantization_config.activation_format is None: - self.quantization_config.activation_format = self._default_activation_format() + self.quantization_config.activation_format = self.default_activation_format if self.quantization_config.activation_format == "bf16": # Weight-only holds no activation global, so a calibrated checkpoint's `input_scale` # has no slot to load into. Not an unexpected key — one this run has no use for, which @@ -197,9 +193,9 @@ def _process_model_before_weight_loading( ) if self.pre_quantized and self.quantization_config.modules_to_convert: - if self._quant_method() != "fp8": + if self.quant_method != "fp8": logger.warning_once( - f"Embedding tables are quantized to FP8; {self._quant_method()} has no embedding path. " + f"Embedding tables are quantized to FP8; {self.quant_method} has no embedding path. " f"{self.quantization_config.modules_to_convert} will hold FP8 rows with one per-tensor scale." ) replace_with_finegrained_embedding( @@ -293,7 +289,8 @@ def get_quantize_ops(self): return FineGrainedQuantize(self) - def _quant_method(self) -> str: + @property + def quant_method(self) -> str: method = self.quantization_config.quant_method return getattr(method, "value", method) # a QuantizationMethod enum, or already its str diff --git a/src/transformers/quantizers/finegrained/mxfp4.py b/src/transformers/quantizers/finegrained/mxfp4.py index 403b15548286..30935d6b9c51 100644 --- a/src/transformers/quantizers/finegrained/mxfp4.py +++ b/src/transformers/quantizers/finegrained/mxfp4.py @@ -30,8 +30,7 @@ class FineGrainedMxfp4HfQuantizer(FineGrainedHfQuantizer): W4A16; an explicit `activation_format="mxfp4"` still gets W4A4. """ - def _default_activation_format(self) -> str: - return "bf16" + default_activation_format = "bf16" def get_weight_conversions(self): if not self.pre_quantized: diff --git a/src/transformers/quantizers/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py index 6e147592d8e4..569ba146548d 100644 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -15,13 +15,9 @@ that the other formats do not. Everything else is inherited. """ -from ...utils import logging from .base import FineGrainedHfQuantizer -logger = logging.get_logger(__name__) - - class FineGrainedNvfp4HfQuantizer(FineGrainedHfQuantizer): """NVFP4, in both key layouts modelopt ships: per expert per projection, or stacked per layer. @@ -81,43 +77,47 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): from ...core_model_loading import Concatenate, MergeModulelist, WeightConverter from ...integrations.finegrained.conversions import FineGrainedInputScales, FineGrainedWeightGlobals - merge = [MergeModulelist(dim=0)] - expert = r"mlp\.experts\..*\." converters = [ WeightConverter( - source_patterns=[rf"{expert}gate_proj\.weight_scale$", rf"{expert}up_proj\.weight_scale$"], + source_patterns=[ + r"mlp\.experts\..*\.gate_proj\.weight_scale$", + r"mlp\.experts\..*\.up_proj\.weight_scale$", + ], target_patterns="mlp.experts.gate_up_proj_scale_inv", - operations=[*merge, Concatenate(dim=1)], + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( - source_patterns=rf"{expert}down_proj\.weight_scale$", + source_patterns=r"mlp\.experts\..*\.down_proj\.weight_scale$", target_patterns="mlp.experts.down_proj_scale_inv", - operations=merge, + operations=[MergeModulelist(dim=0)], ), # every second-level global of a layer in ONE converter: folding the gate|up stack's # two halves into one moves the up half's onto the down projection, so they decide # together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ - rf"{expert}gate_proj\.weight_scale_2", - rf"{expert}up_proj\.weight_scale_2", - rf"{expert}down_proj\.weight_scale_2", - *([rf"{expert}down_proj\.input_scale"] if calibrated else []), + r"mlp\.experts\..*\.gate_proj\.weight_scale_2", + r"mlp\.experts\..*\.up_proj\.weight_scale_2", + r"mlp\.experts\..*\.down_proj\.weight_scale_2", + *([r"mlp\.experts\..*\.down_proj\.input_scale"] if calibrated else []), ], target_patterns=[ "mlp.experts.gate_up_proj_weight_global_scale", "mlp.experts.down_proj_weight_global_scale", *(["mlp.experts.down_proj_input_global_scale"] if calibrated else []), ], - operations=[*merge, FineGrainedWeightGlobals(self)], + operations=[MergeModulelist(dim=0), FineGrainedWeightGlobals(self)], ), ] if calibrated: converters.append( WeightConverter( - source_patterns=[rf"{expert}gate_proj\.input_scale", rf"{expert}up_proj\.input_scale"], + source_patterns=[ + r"mlp\.experts\..*\.gate_proj\.input_scale", + r"mlp\.experts\..*\.up_proj\.input_scale", + ], target_patterns="mlp.experts.gate_up_proj_input_global_scale", - operations=[*merge, FineGrainedInputScales(self)], + operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], ) ) return converters diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index dde5b3691bcc..df28f050349e 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -34,9 +34,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedFP8Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index 3215d28777a8..fa7161d52a45 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -16,7 +16,7 @@ import warnings from typing import TYPE_CHECKING -from ..utils import is_kernels_available, is_torch_available, logging +from ..utils import is_kernels_available, is_torch_available from ..utils.import_utils import KERNELS_MAX_VERSION, KERNELS_MIN_VERSION from .base import HfQuantizer from .quantizers_utils import get_module_from_name @@ -31,9 +31,6 @@ ) -logger = logging.get_logger(__name__) - - if is_torch_available(): import torch @@ -58,9 +55,6 @@ class NVFP4HfQuantizer(HfQuantizer): requires_calibration = False quantization_config: NVFP4Config - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, device_map, **kwargs): if self.pre_quantized: raise ValueError("Loading pre-quantized NVFP4 checkpoints is not supported yet.") diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index b8a218ce7106..c1fc0166b4c6 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -98,9 +98,6 @@ def _loaded(kernel): mock.patch.object(fg, "_FINEGRAINED", None), mock.patch.object(fg, "is_kernels_available", return_value=True), mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), - # a locally importable checkout (FINEGRAINED_KERNELS_PATH / installed package) - # takes precedence over the hub loader; tests must stay hermetic to the fake bundle - mock.patch.object(fg, "_import_local_finegrained", return_value=None), ) @@ -125,8 +122,8 @@ def test_loader_requires_every_symbol(self): del kernel.matmul_grouped # Mock auto-creates attributes; force the miss kernel.matmul_grouped = None - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, self.assertRaises(ImportError) as ctx: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3, self.assertRaises(ImportError) as ctx: load_finegrained_kernel() self.assertIn("matmul_grouped", str(ctx.exception)) @@ -136,8 +133,8 @@ def test_loader_is_compile_safe_cold(self): non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" kernel, _ = _fake_bundle() kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: torch.compiler.reset() @torch.compile(fullgraph=True) @@ -152,8 +149,8 @@ def test_loader_is_compile_safe_when_warm(self): its short-circuit at trace time — the branch that must also return None.""" kernel, _ = _fake_bundle() kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: load_finegrained_kernel() torch.compiler.reset() @@ -166,8 +163,8 @@ def run(x): def test_loader_binds_all_symbols(self): kernel, _ = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: bundle = load_finegrained_kernel() self.assertIs(bundle.matmul_2d, kernel.matmul_2d) self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) @@ -177,8 +174,8 @@ def test_loader_binds_all_symbols(self): class FineGrainedLinearMarshallingTest(unittest.TestCase): def _run(self, **linear_kwargs): kernel, rec = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): x = torch.randn(3, 5, 64, dtype=torch.bfloat16) w = torch.randn(32, 64).to(torch.float8_e4m3fn) ws = torch.randn(1, 1, dtype=torch.float32) @@ -207,8 +204,8 @@ def test_nvfp4_global_and_weight_only_format(self): def test_module_forward_threads_everything(self): kernel, rec = _fake_bundle() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): m = FineGrainedLinear( 64, 32, @@ -311,13 +308,12 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" kernel, rec = _fake_bundle() m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") - p1, p2, p3, p4 = _loaded(kernel) + p1, p2, p3 = _loaded(kernel) linear = mock.Mock(wraps=fg.finegrained_linear) with ( p1, p2, p3, - p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), mock.patch.object(fg, "finegrained_linear", linear), ): @@ -338,8 +334,8 @@ def test_eager_linear_threads_bias_global_and_activation_format(self): def test_batched_marshalling(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_batched"] hidden, top_k_index, top_k_weights = call.args @@ -358,8 +354,8 @@ def test_batched_marshalling(self): def test_grouped_marshalling(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_grouped"] self.assertIs(call.kwargs["down_proj"], m.down_proj) @@ -370,8 +366,8 @@ def test_activation_format_passes_through(self): kernel, rec = _fake_bundle() for activation_format in (None, "bf16", "mxfp8"): m = self._experts(has_gate=True, activation_format=activation_format) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_batched_mm_experts_forward(m, *self._route()) self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) @@ -381,8 +377,8 @@ def test_nvfp4_experts_thread_per_expert_globals(self): # the format table must resolve the ATTRIBUTE the forwards gate on, not just the # param allocation — a None here silently drops the global at every forward self.assertIsNotNone(m.global_scale_dtype) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_batched_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_batched"] self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) @@ -405,8 +401,8 @@ def test_weight_only_experts_hold_no_activation_global(self): would refuse a two-level path without it.""" kernel, rec = _fake_bundle() m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_grouped"] self.assertIsNone(m.gate_up_proj_input_global_scale) @@ -424,8 +420,8 @@ def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) @@ -439,7 +435,7 @@ def forward(self, x): return x m.post_expert_norm = _Recording() - with p1, p2, p3, p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): m(*self._route()) self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") @@ -462,10 +458,10 @@ def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True) m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True - p1, p2, p3, p4 = _loaded(kernel) + p1, p2, p3 = _loaded(kernel) for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): m.post_expert_norm_name = name - with p1, p2, p3, p4: + with p1, p2, p3: fg.finegrained_grouped_mm_experts_forward(m, *self._route()) call = rec.calls["moe_fused_grouped"][-1] if fused: @@ -479,8 +475,8 @@ def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): def test_biases_ride_the_kernel_chain(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True, has_bias=True) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) self.assertTrue(torch.isfinite(out).all()) (call,) = rec.calls["moe_fused_batched"] @@ -521,8 +517,8 @@ def test_unfusable_act_fn_is_passed_as_the_module_glu(self): kernel, rec = _fake_bundle() m = self._experts(has_gate=True, has_bias=True) m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: fg.finegrained_batched_mm_experts_forward(m, *self._route()) (call,) = rec.calls["moe_fused_batched"] # the kernels run the module's own GLU on the host between the two GEMMs — a new @@ -1044,12 +1040,11 @@ def _routed_to(self, *, sm100, scale_ndim, trains=False): kernel, _ = _fake_bundle() w = torch.randn(32, 64).to(torch.float8_e4m3fn) s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) - p1, p2, p3, p4 = _loaded(kernel) + p1, p2, p3 = _loaded(kernel) with ( p1, p2, p3, - p4, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), mock.patch.object(deepgemm, "is_sm100", return_value=sm100), mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, @@ -1075,13 +1070,13 @@ def test_deepgemm_refuses_a_call_that_needs_a_gradient(self): routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) w, s = torch.randn(32, 256).to(torch.float8_e4m3fn), torch.ones(1, 2) calls = [ - ("DeepGEMM linear", lambda: deepgemm.deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128))), - ("DeepGEMM experts", lambda: deepgemm.deepgemm_fp8_fp4_experts_forward(experts, hs, *routing)), - ("DeepGEMM experts", lambda: deepgemm.deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing)), + lambda: deepgemm.deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128)), + lambda: deepgemm.deepgemm_fp8_fp4_experts_forward(experts, hs, *routing), + lambda: deepgemm.deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing), ] with mock.patch.object(deepgemm, "load_deepgemm_kernel") as load: - for backend, call in calls: - with self.assertRaisesRegex(NotImplementedError, f"{backend} has no backward pass"): + for call in calls: + with self.assertRaisesRegex(NotImplementedError, "DeepGEMM does not support training"): call() load.assert_not_called() @@ -1194,8 +1189,8 @@ def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): op = FineGrainedSwizzleScales(hf_quantizer=None) grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: # the loader hands a single target under its pattern, the full name (with suffix) alongside out = op.convert( {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" @@ -1268,8 +1263,8 @@ def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) @@ -1364,8 +1359,8 @@ def test_catch_all_converters_deliver_the_full_parameter_name(self): tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) kernel = self._op_kernel() - p1, p2, p3, p4 = _loaded(kernel) - with p1, p2, p3, p4: + p1, p2, p3 = _loaded(kernel) + with p1, p2, p3: out = converter.convert(renamed, model=model) self.assertEqual(list(out), [key]) diff --git a/tests/integrations/test_moe.py b/tests/integrations/test_moe.py index c13d69fe9e9e..efef759ae22f 100644 --- a/tests/integrations/test_moe.py +++ b/tests/integrations/test_moe.py @@ -41,6 +41,7 @@ def apply_gate(gate_up): num_experts=num_experts, has_gate=True, has_bias=False, + has_post_expert_norm=False, is_transposed=False, act_fn=act_fn, _apply_gate=apply_gate, diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 93f60841361a..660e9ac92d53 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -597,6 +597,7 @@ def _make_megamoe(self, *, num_experts=4, hidden_dim=64, inter=32, swiglu_limit= requires_grad=False, ), config=types.SimpleNamespace(swiglu_limit=swiglu_limit), + has_post_expert_norm=False, ) @require_torch_greater_or_equal("2.7") # torch.float8_e8m0fnu (UE8M0) landed in 2.7 diff --git a/tests/kernels/test_utils.py b/tests/kernels/test_utils.py index dc467c451389..ac64d991ef02 100644 --- a/tests/kernels/test_utils.py +++ b/tests/kernels/test_utils.py @@ -66,6 +66,7 @@ def apply_gate(gate_up): num_experts=num_experts, has_gate=has_gate, has_bias=has_bias, + has_post_expert_norm=False, is_transposed=is_transposed, _is_expert_parallel=is_expert_parallel, act_fn=act_fn, From f3c7cc8d634aba4561ce6f8664facadaf18642c8 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 25 Sep 2026 11:21:06 -0700 Subject: [PATCH 58/68] Tiny checkpoints for every shipped layout, and the save/load fixes they found Tests - tiny random copies of the shipped checkpoints (Qwen3, Ministral-3, Mistral-Small-4, Llama NVFP4, DeepSeek-V3/V4, GPT-OSS, GLM-5.2 NVFP4, MiniMax-M3), shrunk from the real configs and layouts: clean load, logits against the dequantized weights and a golden slice, and a byte-exact save back to the checkpoint layout - test_conversions: scheme-agnostic op tests, then one dummy-model chain per scheme (load through the real quantizer, check the held tensors, save back) - kernel call-structure tests move to tests/kernels/test_finegrained.py, the sharded load paths to test_distributed.py; per-pairing forwards; activation formats a weight format can't take fail at config construction - test_reverse_loading_mapping runs for DeepSeek-V4 and MiniMax-M3 (head models) Fixes - save: a parameter the quantizer creates follows its weight's reversed converter (`quantization_param_suffixes`), replacing the swizzle-only workaround; converter outputs get the renamings after conversion; SplitModulelist squeezes only its own dim - load: a pre-quantized checkpoint's scales keep their dtype (`param_keeps_checkpoint_dtype`) - the MXFP4 weight-only default reaches the groups the modules read (GPT-OSS ran W4A4) - modelopt: `dynamic: false` is NVFP4's calibrated global, not static activations; dense NVFP4 linears; FP8 KV-cache scales ignored; save keys and storage - VLM skip lists match renamed modules (MiniMax-M3's vision tower) - DeepSeek-V4 and MiniMax-M3 mappings reverse faithfully Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/conversion_mapping.py | 27 +- src/transformers/core_model_loading.py | 54 +- .../integrations/finegrained/conversions.py | 55 +- .../integrations/finegrained/core.py | 5 + src/transformers/quantizers/base.py | 11 + .../quantizers/finegrained/base.py | 47 +- .../quantizers/finegrained/blockfp8.py | 6 +- .../quantizers/finegrained/nvfp4.py | 101 +- src/transformers/utils/quantization_config.py | 26 +- .../finegrained/test_conversions.py | 1002 +++++++++++++---- tests/integrations/finegrained/test_core.py | 951 +--------------- .../finegrained/test_distributed.py | 545 +++++++++ .../integrations/finegrained/test_forwards.py | 130 ++- tests/integrations/finegrained/test_models.py | 723 ++++-------- tests/kernels/test_finegrained.py | 577 ++++++++++ .../deepseek_v4/test_modeling_deepseek_v4.py | 17 +- .../test_modeling_minimax_m3_vl.py | 4 +- 17 files changed, 2468 insertions(+), 1813 deletions(-) create mode 100644 tests/integrations/finegrained/test_distributed.py create mode 100644 tests/kernels/test_finegrained.py diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 667da9b570a4..49e2c0818998 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -503,16 +503,21 @@ def _build_checkpoint_conversion_mapping(): ), ], "deepseek_v4": [ + # checkpoint-form renames the general `.attn.` one below would otherwise take first + WeightRenaming(source_patterns=r"\.attn\.wq_b\.", target_patterns=r".attn.q_b_proj."), # General mid-key renames acting on several keys WeightRenaming(source_patterns=r"\.attn\.", target_patterns=r".self_attn."), WeightRenaming(source_patterns=r"\.ffn\.", target_patterns=r".mlp."), WeightRenaming(source_patterns=r"\.indexer\.compressor\.", target_patterns=r"\.compressor\.indexer\."), # Leaf patterns - WeightRenaming(source_patterns=r"^embed\.weight$", target_patterns="embed_tokens.weight"), + # The checkpoint has no `model.` prefix: named here, so a save takes it back off + WeightRenaming(source_patterns=r"^layers\.", target_patterns="model.layers."), + WeightRenaming(source_patterns=r"^norm\.weight$", target_patterns="model.norm.weight"), + WeightRenaming(source_patterns=r"^embed\.weight$", target_patterns="model.embed_tokens.weight"), WeightRenaming(source_patterns=r"^head\.weight$", target_patterns="lm_head.weight"), - WeightRenaming(source_patterns=r"^hc_head_fn$", target_patterns="hc_head.hc_fn"), - WeightRenaming(source_patterns=r"^hc_head_base$", target_patterns="hc_head.hc_base"), - WeightRenaming(source_patterns=r"^hc_head_scale$", target_patterns="hc_head.hc_scale"), + WeightRenaming(source_patterns=r"^hc_head_fn$", target_patterns="model.hc_head.hc_fn"), + WeightRenaming(source_patterns=r"^hc_head_base$", target_patterns="model.hc_head.hc_base"), + WeightRenaming(source_patterns=r"^hc_head_scale$", target_patterns="model.hc_head.hc_scale"), WeightRenaming(source_patterns=r"\.attn_norm\.", target_patterns=r".input_layernorm."), WeightRenaming(source_patterns=r"\.ffn_norm\.", target_patterns=r".post_attention_layernorm."), WeightRenaming(source_patterns=r"\.hc_attn_fn$", target_patterns=r".attn_hc.fn"), @@ -527,10 +532,13 @@ def _build_checkpoint_conversion_mapping(): target_patterns=r"\.compressor\.indexer\.scorer\.weights_proj\.", ), WeightRenaming(source_patterns=r"\.indexer\.wq_b\.", target_patterns=r"\.compressor\.indexer\.q_b_proj\."), - WeightRenaming(source_patterns=r"\.norm\.", target_patterns=r"\.kv_norm\."), + # the compressors' norms only: the attention's own is `kv_norm` already, and a save reversing a + # bare `.kv_norm.` would rename it too (`[.]`: the reverse unescapes a group's `\.`) + WeightRenaming( + source_patterns=r"\.compressor\.(|indexer[.])norm\.", target_patterns=r"\.compressor\.\1kv_norm\." + ), WeightRenaming(source_patterns=r"\.ape$", target_patterns=r"\.position_bias"), WeightRenaming(source_patterns=r"\.wq_a\.", target_patterns=r".q_a_proj."), - WeightRenaming(source_patterns=r"\.self_attn\.wq_b\.", target_patterns=r"\.self_attn\.q_b_proj\."), WeightRenaming(source_patterns=r"\.wkv\.", target_patterns=r".kv_proj."), WeightRenaming(source_patterns=r"\.wgate\.", target_patterns=r".gate_proj."), WeightRenaming(source_patterns=r"\.wo_a\.", target_patterns=r".o_a_proj."), @@ -640,7 +648,12 @@ def _build_checkpoint_conversion_mapping(): source_patterns=r"^patch_merge_mlp\.linear_2\.", target_patterns="model.multi_modal_projector.merge_linear_2.", ), - WeightRenaming(source_patterns=r"\.block_sparse_moe\.", target_patterns=r"\.mlp\."), + # the MoE block only, so a save does not reverse the dense and vision `.mlp.` too (`[.]`: the + # reverse unescapes a group's `\.`) + WeightRenaming( + source_patterns=r"\.block_sparse_moe\.(experts[.]|gate[.]|shared_experts[.]|e_score_correction_bias$)", + target_patterns=r"\.mlp\.\1", + ), WeightRenaming( source_patterns=r"\.e_score_correction_bias", target_patterns=r"\.gate\.e_score_correction_bias" ), diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7d4ee840ee69..eb1e18498cb2 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -287,8 +287,9 @@ def convert( sizes = tensor.size(self.dim) targets = self.get_target_patterns(input_dict, source_pattern, target_patterns, sizes) chunks = torch.chunk(tensor, sizes, dim=self.dim) - # We squeeze each chunk here as well to make sure to give them their original shape - all_tensors.update({target: chunk.squeeze() for target, chunk in zip(targets, chunks)}) + # We squeeze each chunk here as well to make sure to give them their original shape; only the split + # dim, so a tensor that holds a size-1 dim of its own (a single column of block scales) keeps it + all_tensors.update({target: chunk.squeeze(self.dim) for target, chunk in zip(targets, chunks)}) return all_tensors def get_target_patterns( @@ -1699,6 +1700,7 @@ def convert_and_load_state_dict_in_model( if hasattr(tensor, "get_dtype") else tensor.is_floating_point() ) + or hf_quantizer.param_keeps_checkpoint_dtype(renamed_key) ) ): # if the key was renamed as it is not available in the state dict otherwise, it means that we are deserializing it, @@ -1812,6 +1814,11 @@ def convert_and_load_state_dict_in_model( return loading_info, disk_offload_index +def _is_catch_all(converter: WeightConverter, key: str) -> bool: + """Whether `converter` takes `key` by its parameter name alone, for any module that holds one.""" + return converter.compiled_sources.search(f"probe.{key.rsplit('.', 1)[-1]}") is not None + + def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]): """ Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one @@ -1844,27 +1851,52 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch converters = [entry for entry in reverse_weight_conversions if isinstance(entry, WeightConverter)] pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} - conversion_mapping: dict[str, WeightTransform] = {} + # parameters the quantizer created next to a weight (e.g. its scale when quantizing on the fly) have no + # checkpoint key, so no converter of their own: they take their weight's reversed layout ops, and stay quantized + hf_quantizer = getattr(model, "hf_quantizer", None) + quantization_param_suffixes = hf_quantizer.quantization_param_suffixes if hf_quantizer is not None else () + + # `(layer name, reversed transform, quantization-param suffix)` per realized group + conversion_mapping: dict[str, tuple[str, WeightTransform, str]] = {} state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) for original_key, tensor in state_dict: # Rename the key according to all renaming pattern and optional weight converter patterns renamed_key, source_pattern = rename_source_key(original_key, renamings, converters, reverse=True) + + source_key, suffix = original_key, "" + param_suffix = next((s for s in quantization_param_suffixes if original_key.endswith(s)), None) + own = pattern_to_converter.get(source_pattern) + # a parameter follows its weight's converter unless it has one of its own; a catch-all, which + # takes the name whatever module holds it, does not count as one + if param_suffix is not None and (own is None or _is_catch_all(own, original_key)): + weight_key = original_key.removesuffix(param_suffix) + _, weight_pattern = rename_source_key(weight_key, [], converters, reverse=True) + if weight_pattern is not None and not _is_catch_all(pattern_to_converter[weight_pattern], weight_key): + source_key, source_pattern, suffix = weight_key, weight_pattern, param_suffix + if source_pattern is not None: - new_converter = deepcopy(pattern_to_converter[source_pattern]) + # converted on its own name: the renamings apply to what the converter writes, since one written + # on the tail the converter matches would otherwise leave it nothing to expand the prefix from + converted_key, _ = rename_source_key(source_key, [], converters, reverse=True) # each target key gets its own converter instance - mapping = conversion_mapping.setdefault(renamed_key, new_converter) + _, mapping, _ = conversion_mapping.setdefault( + converted_key + suffix, (converted_key, deepcopy(pattern_to_converter[source_pattern]), suffix) + ) + mapping.add_tensor(converted_key, source_key, source_pattern, tensor) else: - mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) - source_pattern = original_key - - mapping.add_tensor(renamed_key, original_key, source_pattern, tensor) + _, mapping, _ = conversion_mapping.setdefault( + renamed_key, (renamed_key, WeightRenaming(original_key, renamed_key), "") + ) + mapping.add_tensor(renamed_key, original_key, original_key, tensor) new_state_dict = {} - for first_param_name, reversed_converter in conversion_mapping.items(): + for layer_name, reversed_transform, suffix in conversion_mapping.values(): # Apply the reverse converter - realized_value = reversed_converter.convert(first_param_name, model=model, config=model.config) + realized_value = reversed_transform.convert(layer_name, model=model, config=model.config) for target_name, param in realized_value.items(): param = param[0] if isinstance(param, list) else param + if isinstance(reversed_transform, WeightConverter): + target_name, _ = rename_source_key(target_name + suffix, renamings, [], reverse=True) new_state_dict[target_name] = param return new_state_dict diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index f46fe46f30b6..66933f33e6fa 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -111,9 +111,10 @@ def convert(self, input_dict, model=None, target_patterns=None, **kwargs): class FineGrainedScaleContainer(_FineGrainedOp): - """A block scale into the dtype its module holds. UE8M0 ships either as the exponent byte - under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast exactly); the - module records which the checkpoint used so the reverse restores it on save.""" + """A UE8M0 block scale into the ``float8_e8m0fnu`` its module holds: it ships either as the + exponent byte under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast + exactly), and the module records which so the reverse restores it on save. Every other scale + stays in the dtype the checkpoint ships (Qwen3 and Mistral BF16, DeepSeek-V3 fp32).""" def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): out = {} @@ -128,7 +129,8 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= value = as_container(value, container) else: module, held = held_scale(model, full_layer_name, key) - if held is not None and value.dtype != held.dtype: + # the one container there is: e8m0 held, shipped as its bytes or its fp32 powers of two + if held is not None and held.dtype == _get_ue8m0_dtype() and value.dtype != held.dtype: module.scale_container_dtype = value.dtype value = as_container(value, held.dtype) out[key] = value @@ -304,7 +306,8 @@ def convert(self, input_dict, model=None, target_patterns=None, **kwargs): experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if experts is not None: value = value.reshape(1).expand(experts.num_experts).contiguous() - return dict.fromkeys(target_patterns or list(input_dict), value) + # one tensor per key: a save refuses keys that share storage + return {key: value.clone() for key in target_patterns or list(input_dict)} @property def reverse_op(self) -> ConversionOps: @@ -630,45 +633,3 @@ def reverse_op(self) -> ConversionOps: # a save re-quantizes, so the checkpoint keeps its format whether the in-memory # state stayed quantized or was dequantized for compute return FineGrainedQuantize(self.hf_quantizer) - - -def keep_swizzle_reverse_for_save(model, hf_quantizer) -> None: - """Keep the swizzle's reverse reachable at save time for a model quantized ON THE FLY. - - `_weight_conversions` retains only converters that matched a checkpoint weight, and - quantizing a bf16 checkpoint CREATES the scale keys — so the converter carrying the - layout ops never matched, and saving would write the module's 5-D SWIZZLE_32_4_4 grid - where the checkpoint format is the affine one (silently: the module reads back whatever - it wrote, so only a re-save catches it). - - Adds a converter that renames nothing and carries the swizzle alone, for scales no - retained converter already covers — reversing it is the unswizzle, which reads the - affine grid off the self-describing 5-D shape. Scoped that way because a converter - appended here is matched BEFORE the retained ones (the save reverses the list), so a - broader pattern would take keys away from the real converter and skip the rest of its - chain. - """ - from ...core_model_loading import WeightConverter - - conversions = list(getattr(model, "_weight_conversions", None) or []) - covered = { - target - for conv in conversions - if isinstance(conv, WeightConverter) - for target in (getattr(conv, "target_patterns", None) or []) - } - held_swizzled = { - name.rsplit(".", 1)[-1] - for name, param in model.named_parameters() - if param.ndim == 5 and name.rsplit(".", 1)[-1].endswith("_scale_inv") - } - missing = sorted(name for name in held_swizzled if not any(name in target for target in covered)) - if not missing: - return - for name in missing: - conv = WeightConverter( - source_patterns=rf"{name}$", target_patterns=name, operations=[FineGrainedSwizzleScales(hf_quantizer)] - ) - conv._was_used = True # it describes a layout this quantizer applied, not one a checkpoint carried - conversions.append(conv) - model._weight_conversions = conversions diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 2b1fce8152ba..48e3405a7f46 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -415,7 +415,11 @@ def __init__( self.weight_scale_inv = nn.Parameter(scale, requires_grad=scale.dtype.is_floating_point) static = self.activation_scheme == "static" + calibrated = global_scale_dtype is not None and activation_format != "bf16" _set_optional_parameter(self, "activation_scale", torch.tensor(1.0, dtype=torch.float32) if static else None) + _set_optional_parameter( + self, "input_global_scale", torch.tensor(1.0, dtype=torch.float32) if calibrated else None + ) _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) def forward(self, input: torch.Tensor) -> torch.Tensor: @@ -428,6 +432,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: bias=self.bias, allow_deepgemm=not self._deepgemm_disabled, weight_global_scale=self.weight_global_scale, + input_global_scale=self.input_global_scale, activation_format=self.activation_format, ) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index fbebea2977be..b960bb496423 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -81,9 +81,13 @@ class HfQuantizer(ABC): The quantization config that defines the quantization parameters of your model that you want to quantize. requires_calibration (`bool`): Whether the quantization method requires to calibrate the model before using it. + quantization_param_suffixes (`tuple[str, ...]`): + Suffixes of the parameters quantization adds to a weight (`weight` -> `weight_scale_inv`). + A save runs each through its weight's reversed conversion, so it lands in the weight's layout. """ requires_calibration = False + quantization_param_suffixes: tuple[str, ...] = () def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): self.quantization_config = quantization_config @@ -133,6 +137,13 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** """ return False + def param_keeps_checkpoint_dtype(self, param_name: str) -> bool: + """ + Whether a pre-quantized checkpoint's `param_name` loads in the dtype it ships in rather than the model's + (e.g. quantization scales, which a checkpoint may ship narrower or wider than the model dtype). + """ + return False + def validate_environment(self, *args, **kwargs): """ This method is used to potentially check for potential conflicts with arguments that are diff --git a/src/transformers/quantizers/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py index 03c52423f1d9..4598efe41e29 100644 --- a/src/transformers/quantizers/finegrained/base.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -22,6 +22,7 @@ """ import re +from dataclasses import replace from typing import TYPE_CHECKING from ...utils import ( @@ -33,7 +34,7 @@ ) from ...utils.quantization_config import groups_with_expert_dtype from ..base import HfQuantizer -from ..quantizers_utils import get_module_from_name +from ..quantizers_utils import get_module_from_name, should_convert_module if is_torch_available(): @@ -60,6 +61,12 @@ class FineGrainedHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedConfig" default_activation_format: str | None = None + quantization_param_suffixes = ("_scale_inv",) + + def param_keeps_checkpoint_dtype(self, param_name: str) -> bool: + # a scale loads as the checkpoint ships it (Qwen3 and Mistral ship BF16, DeepSeek-V3 fp32): the + # kernels read either, and a save then writes it back unchanged + return param_name.endswith(("_scale_inv", "activation_scale", "_global_scale")) @property def supports_dequantize(self) -> bool: @@ -147,25 +154,30 @@ def param_element_size(self, model: "PreTrainedModel", param_name: str, param: " return super().param_element_size(model, param_name, param) def _normalize_modules_to_not_convert(self, model: "PreTrainedModel"): - """Rewrite the skip-list to the model's own module tree. - For models that were already released, if they have a list of modules to not quantize - we need to apply the weight renaming / weight conversion opérations to get the actual - layer name of the model in `transformers`. - """ + """Extend the skip-list, which names modules in the checkpoint's namespace, with the model's + modules it covers there. A released model renames its modules on load, so each one is named back + through the reversed renames (written on keys, so on the module's key prefix) and taken with its + sub-tree when the list covers it: a VLM's bare `vision_tower` matches no rename on its own.""" skip = self.quantization_config.modules_to_not_convert if not skip: return from ...conversion_mapping import get_model_conversion_mapping + from ...core_model_loading import WeightRenaming renamings = get_model_conversion_mapping(model) - remapped = [] - for name in skip: - renamed = name - for rename in renamings: - renamed, _ = rename.rename_source_key(renamed) - remapped.append(renamed) - self.quantization_config.modules_to_not_convert = remapped + reverse = [r.reverse_transform() for r in renamings[::-1] if isinstance(r, WeightRenaming)] + covered = [] + for module_name, _ in model.named_modules(): + if any(module_name.startswith(f"{parent}.") for parent in covered): + continue + checkpoint_name = f"{module_name}." + for rename in reverse: + checkpoint_name, _ = rename.rename_source_key(checkpoint_name) + checkpoint_name = checkpoint_name.removesuffix(".") + if checkpoint_name != module_name and not should_convert_module(checkpoint_name, skip): + covered.append(module_name) + self.quantization_config.modules_to_not_convert = [*skip, *covered] def _process_model_before_weight_loading( self, @@ -181,6 +193,13 @@ def _process_model_before_weight_loading( self.quantization_config.groups = groups_with_expert_dtype(self.quantization_config.groups, expert_dtype) if self.quantization_config.activation_format is None: self.quantization_config.activation_format = self.default_activation_format + # the modules read their group's format, so the default lands on this arm's groups too + self.quantization_config.groups = { + name: replace(group, activation_format=self.default_activation_format) + if group.activation_format is None and group.quant_method == self.quant_method + else group + for name, group in self.quantization_config.groups.items() + } if self.quantization_config.activation_format == "bf16": # Weight-only holds no activation global, so a calibrated checkpoint's `input_scale` # has no slot to load into. Not an unexpected key — one this run has no use for, which @@ -207,11 +226,9 @@ def _process_model_before_weight_loading( def _process_model_after_weight_loading(self, model, **kwargs): from ...integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device - from ...integrations.finegrained.conversions import keep_swizzle_reverse_for_save assert_modules_are_quantized(model) disable_deepgemm_on_multi_device(model) - keep_swizzle_reverse_for_save(model, self) return model diff --git a/src/transformers/quantizers/finegrained/blockfp8.py b/src/transformers/quantizers/finegrained/blockfp8.py index dbfeee6e4379..a20030237419 100644 --- a/src/transformers/quantizers/finegrained/blockfp8.py +++ b/src/transformers/quantizers/finegrained/blockfp8.py @@ -44,14 +44,14 @@ def _static_activation_conversions(self): per_expert = [ WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.input_scale", - r"mlp\.experts\..*\.up_proj\.input_scale", + "mlp.experts.*.gate_proj.input_scale", + "mlp.experts.*.up_proj.input_scale", ], target_patterns="mlp.experts.gate_up_proj_activation_scale", operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], ), WeightConverter( - source_patterns=r"mlp\.experts\..*\.down_proj\.input_scale", + source_patterns="mlp.experts.*.down_proj.input_scale", target_patterns="mlp.experts.down_proj_activation_scale", operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], ), diff --git a/src/transformers/quantizers/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py index 569ba146548d..e46872ed5fe5 100644 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -26,7 +26,7 @@ class FineGrainedNvfp4HfQuantizer(FineGrainedHfQuantizer): * the two-level globals. `weight_scale_2` has no slot in a one-level fold, so `dequantize=True` is refused here rather than tested for in the shared path. * the packed E2M1 weights, which need a uint8 -> int8 BITCAST (`copy_` would convert and - corrupt anything >= 128) and an anchored `.weight` source so it cannot swallow + corrupt anything >= 128) and a `.weight` source that cannot swallow `weight_scale` / `weight_scale_2`. """ @@ -35,37 +35,78 @@ def supports_dequantize(self) -> bool: """No: NVFP4's two-level scales do not fold into a single full-precision weight here.""" return False + def _process_model_before_weight_loading(self, model, **kwargs): + # modelopt calibrates an FP8 KV cache beside the weights; the cache here stays in the model dtype + model._keys_to_ignore_on_load_unexpected = set(model._keys_to_ignore_on_load_unexpected or []) | { + r"\.(k_proj\.k|v_proj\.v)_scale$" + } + super()._process_model_before_weight_loading(model, **kwargs) + def get_weight_conversions(self): return self._nvfp4_conversions() if self.pre_quantized else [] def update_weight_conversions(self, weight_conversions): - """The base chain, with every `.weight` source anchored and bitcast to the int8 view.""" + """The base chain, with every `.weight` source kept off modelopt's `weight_scale` / + `weight_scale_2` and bitcast to the int8 view. Quantizing on the fly too: the view passes + the full-precision source through, and its reverse saves the packed bytes as modelopt's uint8.""" from ...core_model_loading import WeightConverter from ...integrations.finegrained.conversions import FineGrainedViewPackedInt8 - if self.pre_quantized: - updated = [] - for conv in weight_conversions: - if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): - conv = WeightConverter( - source_patterns=[p + "$" if p.endswith(".weight") else p for p in conv.source_patterns], - target_patterns=conv.target_patterns, - operations=[*conv.operations, FineGrainedViewPackedInt8(self)], - ) - updated.append(conv) - weight_conversions = updated - return super().update_weight_conversions(weight_conversions) + updated = [] + for conv in weight_conversions: + if isinstance(conv, WeightConverter) and any(p.endswith(".weight") for p in conv.source_patterns): + conv = WeightConverter( + # modelopt's two scales stay out; a `weight_scale_inv` follows its weight in + source_patterns=[ + p + "(?!_scale$)(?!_scale_2$)" if p.endswith(".weight") else p for p in conv.source_patterns + ], + target_patterns=conv.target_patterns, + operations=[*conv.operations, FineGrainedViewPackedInt8(self)], + ) + updated.append(conv) + # a dense linear's packed weight has no converter of its own to take the view on; last, so + # every converter above (the experts' included) claims its keys first + dense_weights = WeightConverter( + source_patterns=r"\.weight$", target_patterns=".weight", operations=[FineGrainedViewPackedInt8(self)] + ) + return [*super().update_weight_conversions(updated), dense_weights] def _nvfp4_conversions(self): - """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`), in both layouts - a modelopt checkpoint ships them in: per expert per projection (GLM-5.2-NVFP4) or already - stacked per layer (the fused vLLM layout). A checkpoint matches one set and - the other never fires. Routed experts only — a generic `weight_scale*` rename would run - before converter collection and mangle these keys.""" + """modelopt NVFP4 scales (`weight_scale`, `weight_scale_2`, `input_scale`): the routed + experts in both layouts a modelopt checkpoint ships them in — per expert per projection + (GLM-5.2-NVFP4) or already stacked per layer (the fused vLLM layout), a checkpoint matching + one set while the other never fires — and the dense linears modelopt's default export + quantizes too.""" # weight-only keeps activations bf16, so there is no activation global for a calibrated # `input_scale` to land on — leave those keys to the unexpected-key filter calibrated = self.quantization_config.activation_format != "bf16" - return self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) + return ( + self._nvfp4_per_expert_conversions(calibrated) + + self._nvfp4_fused_conversions(calibrated) + + self._nvfp4_dense_conversions(calibrated) + ) + + def _nvfp4_dense_conversions(self, calibrated: bool): + """A dense linear's modelopt scales onto its own slots, one-to-one. Kept off the experts: + these renames run before converter collection, and would take the expert keys away from + the converters above.""" + from ...core_model_loading import WeightRenaming + + renames = [ + WeightRenaming( + source_patterns=r"^(?!.*\.experts\.)(.+)\.weight_scale$", target_patterns=r"\1.weight_scale_inv" + ), + WeightRenaming( + source_patterns=r"^(?!.*\.experts\.)(.+)\.weight_scale_2$", target_patterns=r"\1.weight_global_scale" + ), + ] + if calibrated: + renames.append( + WeightRenaming( + source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", target_patterns=r"\1.input_global_scale" + ) + ) + return renames def _nvfp4_per_expert_conversions(self, calibrated: bool): """One key per expert per projection, as GLM-5.2-NVFP4 ships them. @@ -80,14 +121,14 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): converters = [ WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.weight_scale$", - r"mlp\.experts\..*\.up_proj\.weight_scale$", + "mlp.experts.*.gate_proj.weight_scale$", + "mlp.experts.*.up_proj.weight_scale$", ], target_patterns="mlp.experts.gate_up_proj_scale_inv", operations=[MergeModulelist(dim=0), Concatenate(dim=1)], ), WeightConverter( - source_patterns=r"mlp\.experts\..*\.down_proj\.weight_scale$", + source_patterns="mlp.experts.*.down_proj.weight_scale$", target_patterns="mlp.experts.down_proj_scale_inv", operations=[MergeModulelist(dim=0)], ), @@ -96,10 +137,10 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): # together (`FineGrainedWeightGlobals`) WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.weight_scale_2", - r"mlp\.experts\..*\.up_proj\.weight_scale_2", - r"mlp\.experts\..*\.down_proj\.weight_scale_2", - *([r"mlp\.experts\..*\.down_proj\.input_scale"] if calibrated else []), + "mlp.experts.*.gate_proj.weight_scale_2", + "mlp.experts.*.up_proj.weight_scale_2", + "mlp.experts.*.down_proj.weight_scale_2", + *(["mlp.experts.*.down_proj.input_scale"] if calibrated else []), ], target_patterns=[ "mlp.experts.gate_up_proj_weight_global_scale", @@ -113,8 +154,8 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): converters.append( WeightConverter( source_patterns=[ - r"mlp\.experts\..*\.gate_proj\.input_scale", - r"mlp\.experts\..*\.up_proj\.input_scale", + "mlp.experts.*.gate_proj.input_scale", + "mlp.experts.*.up_proj.input_scale", ], target_patterns="mlp.experts.gate_up_proj_input_global_scale", operations=[MergeModulelist(dim=0), FineGrainedInputScales(self)], @@ -125,7 +166,7 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): def _nvfp4_fused_conversions(self, calibrated: bool): """Already stacked per layer, as the vLLM fused layout ships them — so no `MergeModulelist`, and the expert weights need the packed uint8 -> int8 view the - `.weight`-anchored rule gives the per-expert layout.""" + `.weight` rule gives the per-expert layout.""" from ...core_model_loading import WeightConverter from ...integrations.finegrained.conversions import ( FineGrainedInputScales, diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index eba1470b60ef..60aeea690d8b 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1721,6 +1721,19 @@ def __post_init__(self): if self.weight_block_size is not None: self.weight_block_size = tuple(self.weight_block_size) + # what each weight format's kernels quantize activations to (`None` = the weights' own format, + # `"bf16"` = weight-only); anything else would fail in the kernels at the first forward + supported = { + "fp8": (None, "fp8"), + "mxfp8": (None, "mxfp8", "mxfp4", "bf16"), + "mxfp4": (None, "mxfp8", "mxfp4", "bf16"), + "nvfp4": (None, "nvfp4", "bf16"), + }.get(self.quant_method) + if supported is not None and self.activation_format not in supported: + raise ValueError( + f"{self.quant_method} weights take activation_format in {supported}; got {self.activation_format!r}" + ) + def matches(self, module_name: str) -> bool: return any(re.search(target, module_name) for target in self.targets) @@ -1774,11 +1787,15 @@ def group_from_config_groups(spec: dict) -> FineGrainedGroup | None: return None activations = spec.get("input_activations") or {} act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) + activation_format = formats.get(act_key) if activations else None + # group-scaled activations quantize their blocks inline either way: `dynamic: false` there is + # NVFP4's calibrated second-level global (`input_scale`), which the loader carries on its own + static = not activations.get("dynamic", True) and activation_format == "fp8" return FineGrainedGroup( quant_method=quant_method, targets=targets or [".*"], - activation_format=formats.get(act_key) if activations else None, - activation_scheme="dynamic" if activations.get("dynamic", True) else "static", + activation_format=activation_format, + activation_scheme="static" if static else "dynamic", ) @@ -1906,6 +1923,11 @@ def _read_modelopt(kwargs: dict) -> tuple[str, str | None, dict | None, list | N if all(group is not None for group in mapped.values()): groups = mapped ignore = kwargs.pop("ignore", None) or kwargs.pop("exclude_modules", None) + if kwargs.pop("kv_cache_scheme", None): + logger.warning_once( + "This modelopt checkpoint calibrates an FP8 KV cache (`k_scale` / `v_scale`); the KV cache is not " + "quantized here and stays in the model's dtype, so those scales are not loaded." + ) return quant_method, activation_format, groups, ignore def to_dict(self): diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index a84d14fb9b23..e3ba9ad9a0c0 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -1,18 +1,7 @@ -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The `finegrained` conversion ops: what each one emits from a checkpoint layout, and how the -loader chains them. Pure torch, so no kernel is loaded and no GPU is needed.""" +"""The `finegrained` conversions. The ops first, which are scheme-agnostic (quantize, the NVFP4 +global scales); then one class per scheme that loads a checkpoint in that scheme's real layout +through the real quantizer onto a dummy model, checks what the modules hold, and saves it back — +with each scheme's quirks. Pure torch: no kernel is loaded and no GPU is needed.""" import re import unittest @@ -20,111 +9,149 @@ from unittest import mock import torch +from torch import nn import transformers.integrations.finegrained.core as fg from transformers import PreTrainedConfig, PreTrainedModel from transformers.core_model_loading import ( Chunk, + Concatenate, + MergeModulelist, PermuteForRope, WeightConverter, convert_and_load_state_dict_in_model, + revert_weight_conversion, ) from transformers.integrations.finegrained import FineGrainedExperts +from transformers.integrations.moe import use_experts_implementation from transformers.modeling_utils import LoadStateDictConfig from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import FineGrainedConfig +from ...kernels.test_finegrained import fake_finegrained_kernel, patch_finegrained_kernel from .test_core import _Cfg, _quantizer_for, _targets -@require_torch -class FineGrainedMxfp4ConverterTest(unittest.TestCase): - """The gpt-oss checkpoint layout through the real ConversionOps: blocks (E, N, K/32, 16) - uint8 low-nibble-first E2M1 + biased-127 exponent scales, gate|up rows INTERLEAVED — - deserialized output must dequantize identically to the reference LUT unpack.""" +E, H, I = 2, 64, 32 # experts, hidden, intermediate: one 32-group along each input - FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] - def _reference_dequant(self, blocks, scales): - # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) - lut = torch.tensor(self.FP4_VALUES) - lo = lut[(blocks & 0xF).long()] - hi = lut[(blocks >> 4).long()] - vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) - exp = (scales.long() - 127).unsqueeze(-1) - return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) +@use_experts_implementation +class _Experts(nn.Module): + """A model's unquantized experts, which the quantizer replaces with `FineGrainedExperts`.""" - def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): - from transformers.integrations.finegrained.conversions import ( - FineGrainedPackedBlocks, - FineGrainedScaleContainer, + def __init__(self, config): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.empty(E, 2 * I, H)) + self.down_proj = nn.Parameter(torch.empty(E, H, I)) + + +class _Layer(nn.Module): + def __init__(self, config): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(H, H, bias=False) + self.mlp = nn.Module() + self.mlp.experts = _Experts(config) + + +class _MoeModel(PreTrainedModel): + """One decoder layer: a dense `q_proj` and a MoE, the two module kinds a checkpoint quantizes.""" + + base_model_prefix = "model" + config: PreTrainedConfig + + def __init__(self, config): + super().__init__(config) + self.model = nn.Module() + self.model.layers = nn.ModuleList([_Layer(config)]) + self.post_init() + + def _init_weights(self, module): + pass # every weight here comes from the state dict + + +def _model_mapping(): + """The per-expert merge a Qwen3-MoE / DeepSeek model declares for its experts.""" + return [ + WeightConverter( + ["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"], + "mlp.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + "mlp.experts.*.down_proj.weight", "mlp.experts.down_proj", operations=[MergeModulelist(dim=0)] + ), + ] + + +def _load_and_save(checkpoint, quantization_config): + """`checkpoint` through the real quantizer's chain, as `from_pretrained` runs it (off SM100, so the + scales stay affine and no kernel is needed), then saved back. Returns `(model, loading_info, saved)`.""" + config = PreTrainedConfig() + config.hidden_size, config.intermediate_size, config.num_local_experts, config.hidden_act = H, I, E, "silu" + quantizer = _quantizer_for(quantization_config) + quantizer.pre_quantized = True + model = _MoeModel(config) + with mock.patch.object(fg, "is_sm100", return_value=False): + quantizer._process_model_before_weight_loading(model) + load_config = LoadStateDictConfig( + weight_mapping=quantizer.update_weight_conversions(_model_mapping()), hf_quantizer=quantizer ) + info, _ = convert_and_load_state_dict_in_model( + model, {k: v.clone() for k, v in checkpoint.items()}, load_config + ) + model.hf_quantizer = quantizer + saved = revert_weight_conversion(model, model.state_dict()) + return model, info, saved - torch.manual_seed(0) - E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows - blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) - scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) - weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ - "gate_up_proj" - ] - self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) - back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] - self.assertTrue(torch.equal(back, blocks)) +def _raw(tensor): + return tensor.reshape(-1).view(torch.uint8) if tensor.element_size() == 1 else tensor - # the scales are the container op's job: a module holding e8m0 receives the exponent bytes as is - module = fg.FineGrainedExperts.__new__(fg.FineGrainedExperts) - torch.nn.Module.__init__(module) - module.gate_up_proj_scale_inv = torch.nn.Parameter( - torch.zeros(E, N, K // 32, dtype=torch.float8_e8m0fnu), requires_grad=False + +def _per_expert(make): + """A per-expert checkpoint layout: `make(n, k)` gives one projection's `{suffix: tensor}`.""" + checkpoint = {} + for expert in range(E): + for proj, shape in (("gate_proj", (I, H)), ("up_proj", (I, H)), ("down_proj", (H, I))): + for suffix, tensor in make(*shape).items(): + checkpoint[f"model.layers.0.mlp.experts.{expert}.{proj}.{suffix}"] = tensor + for suffix, tensor in make(H, H).items(): + checkpoint[f"model.layers.0.self_attn.q_proj.{suffix}"] = tensor + return checkpoint + + +class _ChainAssertions: + def assertLoadsCleanly(self, info): + self.assertEqual( + (info.missing_keys, info.unexpected_keys, info.mismatched_keys, info.conversion_errors), + (set(), set(), set(), {}), ) - model = torch.nn.Module() - model.experts = module - scale_inv = FineGrainedScaleContainer().convert( - {"gate_up_proj_scales$": scales}, - model=model, - full_layer_name="experts.gate_up_proj_scale_inv", - target_patterns=["gate_up_proj_scale_inv"], - )["gate_up_proj_scale_inv"] - self.assertEqual(scale_inv.dtype, torch.float8_e8m0fnu) - ref = self._reference_dequant(blocks, scales) # rows pass through interleaved - # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) - lut = torch.tensor(self.FP4_VALUES) - w_u8 = weight.view(torch.uint8) - lo = lut[(w_u8 & 0xF).long()] - hi = lut[(w_u8 >> 4).long()] - vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) - exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) - got = vals * torch.pow(torch.tensor(2.0), exp) - torch.testing.assert_close(got, ref, rtol=0, atol=0) + def assertSavesBack(self, saved, checkpoint): + """The save writes the checkpoint's own keys, dtypes and bytes.""" + self.assertEqual(sorted(saved), sorted(checkpoint)) + for key, tensor in checkpoint.items(): + with self.subTest(key=key): + self.assertEqual((saved[key].dtype, saved[key].shape), (tensor.dtype, tensor.shape)) + self.assertTrue(torch.equal(_raw(saved[key]), _raw(tensor))) + + def assertHoldsInterleaved(self, experts, checkpoint, suffix, held): + """The experts hold gate|up stacked and row-interleaved: `[g0, u0, g1, u1, ...]`.""" + self.assertTrue(experts.holds_interleaved_gate_up) + for expert in range(E): + gate = checkpoint[f"model.layers.0.mlp.experts.{expert}.gate_proj.{suffix}"] + up = checkpoint[f"model.layers.0.mlp.experts.{expert}.up_proj.{suffix}"] + expected = torch.stack([_raw(gate).view(gate.shape), _raw(up).view(up.shape)], dim=1) + got = _raw(held[expert]).view(held[expert].shape) + self.assertTrue(torch.equal(got, expected.reshape(-1, gate.shape[-1])), f"expert {expert} {suffix}") @require_torch -class FineGrainedOnTheFlyQuantizeTest(unittest.TestCase): +class FineGrainedQuantizeOpTest(unittest.TestCase): """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run the kernels' own quantizers, which need a GPU), emitted in the module's layout.""" - def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): - """`dequantize=True` on a GPT-OSS MXFP4 checkpoint used to be the mxfp4 integration's job. - The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation - the unquantized experts hold — so it is compared against that integration directly.""" - from transformers.core_model_loading import Transpose - from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks - from transformers.integrations.mxfp4 import convert_moe_packed_tensors - - torch.manual_seed(0) - num_experts, rows, hidden = 2, 8, 64 - blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) - scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) - reference = convert_moe_packed_tensors(blocks.clone(), scales.clone(), dtype=torch.bfloat16) - - sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] - tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} - for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): - tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) - self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain - torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) - def test_a_static_scheme_gets_its_activation_scales_written(self): """A calibration-fed slot the checkpoint does not supply is still WRITTEN, as the identity. @@ -275,160 +302,276 @@ def test_float_scale_fmt_stays_bit_identical_to_its_formula(self): @require_torch -class FineGrainedModeloptConverterTest(unittest.TestCase): - """The ModelOpt (NVFP4) weight conversions. `MergeModulelist` is load-bearing beyond the - merge itself: `core_model_loading` stamps a source with its expert index ONLY when one is - present in the chain, and expert parallelism selects experts by that index. Without it every - rank collects all E globals and the forward asserts on the per-expert count.""" - - def _modelopt_conversions(self, **cfg_kwargs): - from transformers.quantizers.finegrained.nvfp4 import FineGrainedNvfp4HfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig +class FineGrainedLayoutOpsTest(unittest.TestCase): + """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 + scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the + scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's + `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" - cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) - return FineGrainedNvfp4HfQuantizer(cfg).get_weight_conversions() + E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks - def test_every_declared_target_is_a_slot_the_module_holds(self): - """A converter target that no module holds is a load failure, and the activation global is - exactly such a target under a weight-only run: the experts allocate one only when they - quantize activations. Both runs are checked against a real module's parameter names, since - the declarations are the only thing standing between the two.""" + def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): cfg = _Cfg() - for activation_format, holds_global in ((None, True), ("bf16", False)): - experts = FineGrainedExperts( - cfg, block_size=(4, 4), weight_format="nvfp4", activation_format=activation_format + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E + cfg._experts_implementation = impl + with ( + mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts( + cfg, block_size=(128, 128), weight_format=weight_format, activation_format=activation_format ) - slots = {name for name, _ in experts.named_parameters()} | set(experts._buffers) - self.assertEqual( - any("input_global_scale" in slot for slot in slots), - holds_global, - f"module with activation_format={activation_format!r} disagrees about the slot", + model = torch.nn.Module() + model.config = cfg + model.experts = experts + return model, experts + + def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) + self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertTrue(experts.holds_interleaved_gate_up) + # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a + # float32 grid there; GPT-OSS's config has no scale_fmt at all) + _, experts = self._experts("mxfp4", activation_format="bf16") + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + + def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): + for kw in ( + {"weight_format": "mxfp8", "sm100": False}, + {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely + {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales + {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, + ): + with self.subTest(**kw): + _, experts = self._experts(**kw) + self.assertEqual( + experts.gate_up_proj_scale_inv.shape, + (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), + ) + # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block + _, experts = self._experts("fp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) + + def test_megamoe_holds_gate_up_stacked(self): + _, experts = self._experts("fp8", impl="deepgemm_megamoe") + self.assertFalse(experts.holds_interleaved_gate_up) + + def _op_kernel(self): + kernel, _ = fake_finegrained_kernel() + kernel.swizzle_mx_scales = mock.Mock( + side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) + ) + kernel.unswizzle_mx_scales = mock.Mock( + side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) + ) + return kernel + + def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + model, experts = self._experts("mxfp8") + kernel = self._op_kernel() + op = FineGrainedSwizzleScales(hf_quantizer=None) + grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + # the loader hands a single target under its pattern, the full name (with suffix) alongside + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" ) - targets = { - t.split("experts.")[-1].rstrip("$") - for c in self._modelopt_conversions(activation_format=activation_format) - for t in _targets(c) - } - for target in targets: - self.assertIn(target, slots, f"activation_format={activation_format!r} converts onto a missing slot") - - def test_global_scale_converters_carry_an_expert_index(self): - from transformers.core_model_loading import MergeModulelist - - # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked - # tensor per layer has no per-expert sources to stamp — it is sharded like the fused - # expert weight next to it - globals_converters = [ - c - for c in self._modelopt_conversions() - if any("global_scale" in t for t in _targets(c)) - and any(re.search(p, "model.layers.0.mlp.experts.7.gate_proj.weight_scale_2") for p in c.source_patterns) - ] - self.assertTrue(globals_converters, "no global-scale converter found") - for conv in globals_converters: - self.assertTrue( - any(isinstance(op, MergeModulelist) for op in conv.operations), - f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " - "experts and every rank would collect all of them", + self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) + # the same converter's weight passes through + out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") + self.assertIs(out["mlp.experts.gate_up_proj"], weight) + # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully + out = op.convert( + { + "experts.down_proj": weight, + "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), + }, + model=model, + full_layer_name="experts.down_proj", ) - - def test_both_modelopt_scale_layouts_have_converters(self): - """modelopt ships the scales either as one tensor per expert per projection (GLM-5.2) or - as one already-stacked tensor per layer (the fused vLLM layout). Both are - converted; the block scale takes the module's layout ops either way.""" - sources = {p for c in self._modelopt_conversions() for p in c.source_patterns} - for suffix in ("weight_scale", "weight_scale_2", "input_scale"): - per_expert = f"model.layers.0.mlp.experts.7.gate_proj.{suffix}" - fused = f"model.layers.0.mlp.experts.gate_up_proj_{suffix}" - self.assertTrue(any(re.search(p, per_expert) for p in sources), f"no per-expert converter for {suffix}") - self.assertTrue(any(re.search(p, fused) for p in sources), f"no fused converter for {suffix}") - # the per-expert patterns are REGEXES: an unescaped dot also matches `_`, so - # `mlp.experts.*.up_proj.weight_scale_2` would claim the fused key as well - claimed = [p for p in sources if re.search(p, fused) and re.search(p, per_expert)] - self.assertEqual(claimed, [], f"one pattern claims BOTH layouts for {suffix}: {claimed}") - - def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): - """Two converters for one key is not additive — the later one WINS and silently drops - whatever ops the first carried. That is how a fused modelopt checkpoint lost its packed - uint8 -> int8 view (a catch-all duplicated the fused converter) and how the per-expert - patterns, which are REGEXES whose unescaped dots also match `_`, claimed the fused - globals and flattened the gate|up pair instead of folding it.""" - import re - - from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) + self.assertIs(out["experts.down_proj"], weight) + # a module that holds affine scales keeps them + affine_model, _ = self._experts("mxfp4", activation_format="bf16") + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, + model=affine_model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + self.assertIs(out["mlp.experts.gate_up_proj"], grid) + self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) + + def test_scale_container_op_round_trips_every_container(self): + """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: + the op brings both into the held e8m0 (exact cast / same bytes), records the container on + the module, and its reverse restores that container on save; weights and native scales + pass both ways.""" + from transformers.integrations.finegrained.conversions import FineGrainedScaleContainer + + op = FineGrainedScaleContainer(hf_quantizer=None) + native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) + for container, shipped in ((torch.float32, native.float()), (torch.uint8, native.view(torch.uint8))): + with self.subTest(container=container): + model, experts = self._experts("mxfp8", sm100=False) # affine e8m0 scale Parameters + out = op.convert( + {"mlp.experts.gate_up_proj": shipped}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + held = out["mlp.experts.gate_up_proj"] + self.assertEqual(held.dtype, torch.float8_e8m0fnu) + self.assertTrue(torch.equal(held.view(torch.uint8), native.view(torch.uint8))) + self.assertIs(experts.scale_container_dtype, container) + back = op.reverse_op.convert({"x": held}, model=model)["x"] + self.assertEqual(back.dtype, container) + self.assertTrue(torch.equal(back, shipped)) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + self.assertIs(op.reverse_op.convert({"w": weight}, model=model)["w"], weight) + model, experts = self._experts("mxfp8", sm100=False) + self.assertIs( + op.convert({"mlp.experts.down_proj": native}, model=model, full_layer_name="experts.down_proj_scale_inv")[ + "mlp.experts.down_proj" + ], + native, + ) + self.assertIsNone(experts.scale_container_dtype) + self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays + + def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + kernel = self._op_kernel() + reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op + artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) + affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") + self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) + kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) + out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") + self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) + + def _quantizer(self, quant_method): from transformers.utils.quantization_config import FineGrainedConfig - for quant_kwargs in ( - {"quant_method": "fp8"}, - {"quant_method": "mxfp4"}, - {"quant_method": "modelopt", "quant_algo": "NVFP4"}, - ): - quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**quant_kwargs)) - quantizer.pre_quantized = True - converters = quantizer.update_weight_conversions([]) - for key in ( - "model.layers.3.mlp.experts.gate_up_proj", - "model.layers.3.mlp.experts.down_proj", - "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", - "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", - ): - claimed = [ - c - for c in converters - for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) - if re.search(str(pattern), key) - ] - with self.subTest(quant=quant_kwargs["quant_method"], key=key.rsplit(".", 1)[1]): - self.assertLessEqual( - len(claimed), - 1, - f"{len(claimed)} converters claim this key; the last one wins and drops " - f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", - ) + return _quantizer_for(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions - def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): - """A static checkpoint calibrates one `input_scale` per quantized module, and each key - must reach the slot its module holds, in that module's shape: one value on a dense linear, - one per expert on the stacked experts, and the gate|up pair reduced to one per expert - rather than flattened to 2E. Unrouted, the module keeps its registered default of 1.0 and - every activation quantizes against the wrong scale, silently.""" - import re + def test_blocks_scales_converters_carry_the_layout_ops(self): + from transformers.core_model_loading import WeightConverter + from transformers.integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedSwizzleScales, + ) + converters = self._quantizer("mxfp4")([]) + by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} + ops = by_target["gate_up_proj"].operations + self.assertIsInstance(ops[-2], FineGrainedInterleaveGateUp) + self.assertIsInstance(ops[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["gate_up_proj_scale_inv"].operations[0], FineGrainedInterleaveGateUp) + self.assertIsInstance(by_target["down_proj"].operations[-1], FineGrainedSwizzleScales) + self.assertIsInstance(by_target["down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) + + def test_catch_all_converters_deliver_the_full_parameter_name(self): + """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax + down_proj_scale_inv) goes through the real converter plumbing: the ops receive the tensor + under the source PATTERN and must hand back the target, which the loader expands to the + full name.""" + from transformers.core_model_loading import WeightConverter, rename_source_key + + model, _ = self._experts("mxfp8") + converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] + for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): + renamed, source_pattern = rename_source_key(key, [], converters) + self.assertEqual(renamed, key) + converter = next(c for c in converters if source_pattern in c.source_patterns) + tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) + converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) + kernel = self._op_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = converter.convert(renamed, model=model) + self.assertEqual(list(out), [key]) + + def test_replacement_keeps_the_models_gate_up_convention(self): + """`replace_with_finegrained_layer` goes through `use_experts_implementation`, which stamps + the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must + survive it (it decides whether the loader interleaves).""" + from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM from transformers.utils.quantization_config import FineGrainedConfig - quantizer = _quantizer_for(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) - quantizer.pre_quantized = True - converters = quantizer.update_weight_conversions([]) + gpt_oss = GptOssConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=32, + num_local_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + sliding_window=16, + ) + qwen = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + ) + for cls, cfg, quant_method, expected in ( + (GptOssForCausalLM, gpt_oss, "mxfp4", False), + (Qwen3MoeForCausalLM, qwen, "mxfp8", True), + ): + with self.subTest(model=cls.__name__), torch.device("meta"): + model = cls(cfg) + fg.replace_with_finegrained_layer( + model, modules_to_not_convert=[], quantization_config=FineGrainedConfig(quant_method=quant_method) + ) + experts = model.model.layers[0].mlp.experts + self.assertIsInstance(experts, fg.FineGrainedExperts) + self.assertIs(experts.is_concatenated, expected) + + def test_interleave_op_follows_the_model_and_the_experts_backend(self): + """The op reads the model's experts modules: a model whose own rows already come + interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself + and keeps the stack; every other case gets the interleaved order, and the reverse restores + the stack.""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + model, experts = self._experts("mxfp8") + out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) + back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] + torch.testing.assert_close(back, stacked) + experts.is_concatenated = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) + experts.is_concatenated = True + experts.holds_interleaved_gate_up = False + self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) - def claim(key): - return [ - c - for c in converters - for pattern in (c.source_patterns if isinstance(c.source_patterns, list) else [c.source_patterns]) - if re.search(str(pattern), key) - ] - for key, expected in ( - ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), - ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.7.up_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.7.down_proj.input_scale", "mlp.experts.down_proj_activation_scale"), - ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), - ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), - ): - with self.subTest(key=key.rsplit(".", 2)[-2]): - claimed = claim(key) - self.assertEqual(len(claimed), 1, f"{len(claimed)} converters claim {key}") - # a renaming's target carries a backreference; resolve it against the key - targets = [ - re.sub(str(claimed[0].source_patterns[0]), str(t), key) if "\\1" in str(t) else str(t) - for t in claimed[0]._original_target_patterns - ] - self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") - # a dynamic checkpoint has no such key to route, and a weight-only one ignores it - dynamic = _quantizer_for(FineGrainedConfig(quant_method="fp8")) - dynamic.pre_quantized = True - self.assertEqual(dynamic.get_weight_conversions(), []) +@require_torch +class FineGrainedGlobalScaleOpsTest(unittest.TestCase): + """The NVFP4 second-level globals and calibrated input scales, in the layout the kernels index: + one per expert per projection, the gate|up pair folded, the gate_up activation global one value.""" def test_the_calibrated_expert_scales_reduce_per_expert_not_per_tensor(self): """`FineGrainedInputScales` collapses the gate_up to ONE value for the NVFP4 global — the @@ -568,6 +711,407 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) +@require_torch +class FineGrainedConversionPlanTest(unittest.TestCase): + """What every scheme's conversion plan must hold whatever the layout.""" + + def test_a_checkpoint_key_is_claimed_by_exactly_one_converter(self): + """Two converters for one key is not additive — the later one WINS and silently drops + whatever ops the first carried. That is how a fused modelopt checkpoint lost its packed + uint8 -> int8 view (a catch-all duplicated the fused converter) and how the per-expert + patterns, which are REGEXES whose unescaped dots also match `_`, claimed the fused + globals and flattened the gate|up pair instead of folding it.""" + + from transformers.quantizers.finegrained.base import FineGrainedHfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + for quant_kwargs in ( + {"quant_method": "fp8"}, + {"quant_method": "mxfp4"}, + {"quant_method": "modelopt", "quant_algo": "NVFP4"}, + ): + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(**quant_kwargs)) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + for key in ( + "model.layers.3.mlp.experts.gate_up_proj", + "model.layers.3.mlp.experts.down_proj", + "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", + "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", + ): + # the loader routes a key through each converter's compiled sources + claimed = [c for c in converters if isinstance(c, WeightConverter) and c.compiled_sources.search(key)] + with self.subTest(quant=quant_kwargs["quant_method"], key=key.rsplit(".", 1)[1]): + self.assertLessEqual( + len(claimed), + 1, + f"{len(claimed)} converters claim this key; the last one wins and drops " + f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", + ) + + +@require_torch +class FineGrainedBlockFp8ChainTest(_ChainAssertions, unittest.TestCase): + """Block-FP8 and per-tensor FP8, as DeepSeek-V3, DeepSeek-V4 and a static (Mistral) checkpoint ship them.""" + + def test_deepseek_v3_block_fp8(self): + """Per-expert e4m3 weights with fp32 `weight_scale_inv` per block, merged onto the stacked experts.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.rand(-(-n // 16), -(-k // 16)) + 0.5, + } + ) + model, info, saved = _load_and_save( + checkpoint, FineGrainedConfig(quant_method="fp8", weight_block_size=(16, 16)) + ) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + def test_deepseek_v4_ue8m0_scale_keys(self): + """DeepSeek-V4 ships its block scales as E8M0 under `.scale`; they load as `weight_scale_inv` + and save back under their own name.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "scale": torch.randint(120, 130, (-(-n // 16), -(-k // 16)), dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + } + ) + config = FineGrainedConfig(quant_method="fp8", weight_block_size=(16, 16), scale_fmt="ue8m0") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + self.assertEqual(model.model.layers[0].self_attn.q_proj.weight_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertSavesBack(saved, checkpoint) + + def test_mistral4_static_per_tensor(self): + """Mistral-Small-4 ships its experts already stacked under the module's own names: e4m3, one + BF16 weight scale per expert `(E, 1, 1)` and a calibrated BF16 activation scale per expert; + its dense linears hold one 0-dim BF16 scale of each.""" + torch.manual_seed(0) + checkpoint = { + "model.layers.0.self_attn.q_proj.weight": torch.randn(H, H).to(torch.float8_e4m3fn), + "model.layers.0.self_attn.q_proj.weight_scale_inv": (torch.rand(()) + 0.5).to(torch.bfloat16), + "model.layers.0.self_attn.q_proj.activation_scale": (torch.rand(()) + 0.5).to(torch.bfloat16), + } + for proj, (rows, cols) in (("gate_up_proj", (2 * I, H)), ("down_proj", (H, I))): + prefix = f"model.layers.0.mlp.experts.{proj}" + checkpoint[prefix] = torch.randn(E, rows, cols).to(torch.float8_e4m3fn) + checkpoint[f"{prefix}_scale_inv"] = (torch.rand(E, 1, 1) + 0.5).to(torch.bfloat16) + checkpoint[f"{prefix}_activation_scale"] = (torch.rand(E) + 0.5).to(torch.bfloat16) + config = FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + torch.testing.assert_close( + experts.gate_up_proj_activation_scale.float(), + checkpoint["model.layers.0.mlp.experts.gate_up_proj_activation_scale"].float(), + ) + self.assertSavesBack(saved, checkpoint) + + def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): + """A static checkpoint calibrates one `input_scale` per quantized module, and each key + must reach the slot its module holds, in that module's shape: one value on a dense linear, + one per expert on the stacked experts, and the gate|up pair reduced to one per expert + rather than flattened to 2E. Unrouted, the module keeps its registered default of 1.0 and + every activation quantizes against the wrong scale, silently.""" + + from transformers.utils.quantization_config import FineGrainedConfig + + quantizer = _quantizer_for(FineGrainedConfig(quant_method="fp8", activation_scheme="static")) + quantizer.pre_quantized = True + converters = quantizer.update_weight_conversions([]) + + def claim(key): + # the loader routes a key through each converter's compiled sources, not the raw patterns + return [c for c in converters if c.compiled_sources.search(key)] + + for key, expected in ( + ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), + ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.up_proj.input_scale", "mlp.experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.down_proj.input_scale", "mlp.experts.down_proj_activation_scale"), + ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), + ): + with self.subTest(key=key.rsplit(".", 2)[-2]): + claimed = claim(key) + self.assertEqual(len(claimed), 1, f"{len(claimed)} converters claim {key}") + # a renaming's target carries a backreference; resolve it against the key + targets = [ + re.sub(str(claimed[0].source_patterns[0]), str(t), key) if "\\1" in str(t) else str(t) + for t in claimed[0]._original_target_patterns + ] + self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") + # a dynamic checkpoint has no such key to route, and a weight-only one ignores it + dynamic = _quantizer_for(FineGrainedConfig(quant_method="fp8")) + dynamic.pre_quantized = True + self.assertEqual(dynamic.get_weight_conversions(), []) + + +@require_torch +class FineGrainedMxfp8ChainTest(_ChainAssertions, unittest.TestCase): + """MXFP8, as MiniMax-M3 ships it.""" + + def test_minimax_m3_e8m0_scale_bytes(self): + """MiniMax-M3-MXFP8 ships its E8M0 scales as raw `uint8` bytes: they load as the module's + `float8_e8m0fnu`, and a save restores the container the checkpoint used.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.randint(120, 130, (n, k // 32), dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, FineGrainedConfig(quant_method="mxfp8")) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + +@require_torch +class FineGrainedMxfp4ChainTest(_ChainAssertions, unittest.TestCase): + """MXFP4 as GPT-OSS ships it: blocks (E, N, K/32, 16) uint8 low-nibble-first E2M1 + biased-127 + exponent scales, gate|up rows INTERLEAVED.""" + + FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + + def _reference_dequant(self, blocks, scales): + # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) + lut = torch.tensor(self.FP4_VALUES) + lo = lut[(blocks & 0xF).long()] + hi = lut[(blocks >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) + exp = (scales.long() - 127).unsqueeze(-1) + return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) + + def test_packed_experts(self): + """Packed E2M1 per expert (two values per `int8` byte along K) with an E8M0 scale per 32, as + DeepSeek-V4 ships its experts, merged onto the stacked experts.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randint(-128, 128, (n, k // 2), dtype=torch.int8), + "weight_scale_inv": torch.randint(120, 130, (n, k // 32), dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + } + ) + model, info, saved = _load_and_save(checkpoint, FineGrainedConfig(quant_method="mxfp4")) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): + from transformers.integrations.finegrained.conversions import ( + FineGrainedPackedBlocks, + FineGrainedScaleContainer, + ) + + torch.manual_seed(0) + E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows + blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) + scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) + + weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ + "gate_up_proj" + ] + self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) + back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] + self.assertTrue(torch.equal(back, blocks)) + + # the scales are the container op's job: a module holding e8m0 receives the exponent bytes as is + module = fg.FineGrainedExperts.__new__(fg.FineGrainedExperts) + torch.nn.Module.__init__(module) + module.gate_up_proj_scale_inv = torch.nn.Parameter( + torch.zeros(E, N, K // 32, dtype=torch.float8_e8m0fnu), requires_grad=False + ) + model = torch.nn.Module() + model.experts = module + scale_inv = FineGrainedScaleContainer().convert( + {"gate_up_proj_scales$": scales}, + model=model, + full_layer_name="experts.gate_up_proj_scale_inv", + target_patterns=["gate_up_proj_scale_inv"], + )["gate_up_proj_scale_inv"] + self.assertEqual(scale_inv.dtype, torch.float8_e8m0fnu) + + ref = self._reference_dequant(blocks, scales) # rows pass through interleaved + # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) + lut = torch.tensor(self.FP4_VALUES) + w_u8 = weight.view(torch.uint8) + lo = lut[(w_u8 & 0xF).long()] + hi = lut[(w_u8 >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) + exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) + got = vals * torch.pow(torch.tensor(2.0), exp) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + def test_mxfp4_blocks_dequantize_like_the_integration_they_replace(self): + """`dequantize=True` on a GPT-OSS MXFP4 checkpoint used to be the mxfp4 integration's job. + The finegrained chain has to land the same bf16 tensor, in the (E, hidden, 2I) orientation + the unquantized experts hold — so it is compared against that integration directly.""" + from transformers.core_model_loading import Transpose + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks + from transformers.integrations.mxfp4 import convert_moe_packed_tensors + + torch.manual_seed(0) + num_experts, rows, hidden = 2, 8, 64 + blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) + scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) + reference = convert_moe_packed_tensors(blocks.clone(), scales.clone(), dtype=torch.bfloat16) + + sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] + tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} + for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): + tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) + self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain + torch.testing.assert_close(next(iter(tensors.values())).float(), reference.float(), rtol=0, atol=0) + + +@require_torch +class FineGrainedNvfp4ChainTest(_ChainAssertions, unittest.TestCase): + """NVFP4 as modelopt ships it, per expert (GLM-5.2-NVFP4), stacked per layer (vLLM) and on dense + linears (modelopt's default export). `MergeModulelist` is load-bearing beyond the merge itself: + `core_model_loading` stamps a source with its expert index ONLY when one is present in the chain, + and expert parallelism selects experts by that index.""" + + def _modelopt_checkpoint(self): + """modelopt's per-expert layout, dense linear included: packed E2M1 bytes, E4M3 block scales + per 16, fp32 second-level `weight_scale_2` and calibrated `input_scale`. Gate and up share + their globals, as modelopt's linear fusion exports them, so a save writes each back exactly.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randint(0, 256, (n, k // 2), dtype=torch.uint8), + "weight_scale": (torch.rand(n, k // 16) + 0.5).to(torch.float8_e4m3fn), + "weight_scale_2": torch.rand(()) * 1e-3 + 1e-4, + "input_scale": torch.rand(()) * 1e-2 + 1e-3, + } + ) + for expert in range(E): + prefix = f"model.layers.0.mlp.experts.{expert}" + for suffix in ("weight_scale_2", "input_scale"): + checkpoint[f"{prefix}.up_proj.{suffix}"] = checkpoint[f"{prefix}.gate_proj.{suffix}"].clone() + # the gate|up activation global is one value for the layer: its rows are quantized before routing + for expert in range(E): + for proj in ("gate_proj", "up_proj"): + checkpoint[f"model.layers.0.mlp.experts.{expert}.{proj}.input_scale"] = torch.tensor(5e-3) + return checkpoint + + def test_modelopt_per_expert_and_dense(self): + checkpoint = self._modelopt_checkpoint() + model, info, saved = _load_and_save(checkpoint, FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4")) + self.assertLoadsCleanly(info) + experts, q_proj = model.model.layers[0].mlp.experts, model.model.layers[0].self_attn.q_proj + self.assertEqual((experts.gate_up_proj.dtype, q_proj.weight.dtype), (torch.int8, torch.int8)) + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + torch.testing.assert_close( + experts.gate_up_proj_weight_global_scale, + torch.stack([checkpoint[f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale_2"] for e in range(E)]), + ) + for slot, key in ( + ("weight_global_scale", "weight_scale_2"), + ("input_global_scale", "input_scale"), + ): + torch.testing.assert_close(getattr(q_proj, slot), checkpoint[f"model.layers.0.self_attn.q_proj.{key}"]) + self.assertSavesBack(saved, checkpoint) + + def test_modelopt_fused_layout(self): + """The vLLM layout: the per-expert modelopt tensors already stacked per layer.""" + per_expert = self._modelopt_checkpoint() + checkpoint = {k: v for k, v in per_expert.items() if ".experts." not in k} + for proj, halves in (("gate_up_proj", ("gate_proj", "up_proj")), ("down_proj", ("down_proj",))): + for suffix, stacked in (("", "weight"), ("_weight_scale", "weight_scale")): + checkpoint[f"model.layers.0.mlp.experts.{proj}{suffix}"] = torch.stack( + [ + torch.cat([per_expert[f"model.layers.0.mlp.experts.{e}.{h}.{stacked}"] for h in halves]) + for e in range(E) + ] + ) + # the globals stack one per expert; the gate|up pair agrees here, so its one value covers both + # halves (a pair calibrated apart folds, which the global-scale op tests cover) + for suffix in ("weight_scale_2", "input_scale"): + checkpoint[f"model.layers.0.mlp.experts.{proj}_{suffix}"] = torch.stack( + [per_expert[f"model.layers.0.mlp.experts.{e}.{halves[0]}.{suffix}"] for e in range(E)] + ) + model, info, saved = _load_and_save(checkpoint, FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4")) + self.assertLoadsCleanly(info) + self.assertEqual(model.model.layers[0].mlp.experts.gate_up_proj.dtype, torch.int8) + self.assertSavesBack(saved, checkpoint) + + def _modelopt_conversions(self, **cfg_kwargs): + from transformers.quantizers.finegrained.nvfp4 import FineGrainedNvfp4HfQuantizer + from transformers.utils.quantization_config import FineGrainedConfig + + cfg = FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4", **cfg_kwargs) + return FineGrainedNvfp4HfQuantizer(cfg).get_weight_conversions() + + def test_every_declared_target_is_a_slot_the_module_holds(self): + """A converter target that no module holds is a load failure, and the activation global is + exactly such a target under a weight-only run: the experts allocate one only when they + quantize activations. Both runs are checked against a real module's parameter names, since + the declarations are the only thing standing between the two.""" + from transformers.integrations.finegrained import FineGrainedLinear + + cfg = _Cfg() + for activation_format, holds_global in ((None, True), ("bf16", False)): + experts = FineGrainedExperts( + cfg, block_size=(4, 4), weight_format="nvfp4", activation_format=activation_format + ) + linear = FineGrainedLinear(32, 32, weight_format="nvfp4", activation_format=activation_format) + for module in (experts, linear): + slots = {name for name, _ in module.named_parameters()} + self.assertEqual( + any("input_global_scale" in slot for slot in slots), + holds_global, + f"{type(module).__name__} with activation_format={activation_format!r} disagrees about the slot", + ) + for c in self._modelopt_conversions(activation_format=activation_format): + for target in _targets(c): + # an expert target names the experts' slot; a dense rename's, the linear's + module, slot = ( + (experts, target.split("experts.")[-1]) + if "experts." in target + else (linear, target.removeprefix("\\1.")) + ) + self.assertIn( + slot.rstrip("$"), + {name for name, _ in module.named_parameters()}, + f"activation_format={activation_format!r} converts onto a missing slot", + ) + + def test_global_scale_converters_carry_an_expert_index(self): + from transformers.core_model_loading import MergeModulelist + + # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked + # tensor per layer has no per-expert sources to stamp — it is sharded like the fused + # expert weight next to it + globals_converters = [ + c + for c in self._modelopt_conversions() + if any("global_scale" in t for t in _targets(c)) + and c.compiled_sources.search("model.layers.0.mlp.experts.7.gate_proj.weight_scale_2") + ] + self.assertTrue(globals_converters, "no global-scale converter found") + for conv in globals_converters: + self.assertTrue( + any(isinstance(op, MergeModulelist) for op in conv.operations), + f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " + "experts and every rank would collect all of them", + ) + + @require_torch class FineGrainedLoaderChainTest(unittest.TestCase): """A core conversion chain and an on-the-fly quantization in the same load: the permute has to diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index c1fc0166b4c6..664c16eff425 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -11,80 +11,24 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Marshalling tests for the multi-format ``finegrained`` integration: the kernel bundle is -mocked, so these pin exactly what the integration passes to `kernels-community/finegrained-kernels` -(the As-positional / no-block_size / expert_start / b_global_scale contract) without a GPU.""" +"""The `finegrained` integration's own pieces, no kernel involved: the modules and what replaces what, +the quantization config and its groups, environment validation, the parallel plans, and the frozen +legacy FP8 shim. The kernel call contract lives in `tests/kernels/test_finegrained.py`, the conversion +ops and chains in `test_conversions.py`.""" import re import unittest -import warnings from contextlib import ExitStack, contextmanager -from dataclasses import dataclass from unittest import mock import torch +from parameterized import parameterized -import transformers.integrations.deepgemm as deepgemm import transformers.integrations.finegrained.core as fg -from transformers.integrations.finegrained import ( - FineGrainedExperts, - FineGrainedLinear, - finegrained_linear, - load_finegrained_kernel, -) +from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear from transformers.testing_utils import require_torch -@dataclass -class _Call: - args: tuple - kwargs: dict - - -class _Recorder: - """Records every call and returns a correctly-shaped zero tensor for matmuls.""" - - def __init__(self): - self.calls = {} - - def _op(self, name): - def op(*args, **kwargs): - self.calls.setdefault(name, []).append(_Call(args, kwargs)) - if name == "swizzle_mx_scales": - return args[0] - if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) - hidden, down = args[0], kwargs["down_proj"] - return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) - a, b = args[0], args[1] - n = b.shape[-2] - gather = kwargs.get("gather_idx") - rows = gather.shape[0] if gather is not None else a.shape[0] - return a.new_zeros(rows, n, dtype=torch.bfloat16) - - return op - - -def _fake_bundle(): - rec = _Recorder() - kernel = mock.Mock() - for name in ( - "matmul_2d", - "matmul_batched", - "matmul_grouped", - "moe_fused_batched", - "moe_fused_grouped", - "swizzle_mx_scales", - "unswizzle_mx_scales", - "mxfp8_act_quant", - "mxfp4_act_quant", - "nvfp4_act_quant", - ): - setattr(kernel, name, rec._op(name)) - kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") - kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") - return kernel, rec - - class _Cfg: hidden_size = 64 num_local_experts = 4 @@ -93,14 +37,6 @@ class _Cfg: _experts_implementation = None -def _loaded(kernel): - return ( - mock.patch.object(fg, "_FINEGRAINED", None), - mock.patch.object(fg, "is_kernels_available", return_value=True), - mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), - ) - - def _quantizer_for(config): """The quantizer a `from_pretrained` would pick for this config. Since the arms were split per scheme, constructing the base directly only exercises what every scheme shares.""" @@ -116,113 +52,9 @@ def _targets(conv): @require_torch -class FineGrainedLoaderTest(unittest.TestCase): - def test_loader_requires_every_symbol(self): - kernel, _ = _fake_bundle() - del kernel.matmul_grouped - # Mock auto-creates attributes; force the miss - kernel.matmul_grouped = None - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3, self.assertRaises(ImportError) as ctx: - load_finegrained_kernel() - self.assertIn("matmul_grouped", str(ctx.exception)) - - def test_loader_is_compile_safe_cold(self): - """Cold path: the compiled call is first to load, so the opaque loader node runs its full - body under compile and must return None, never the bundle (`torch.* op returned - non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" - kernel, _ = _fake_bundle() - kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return load_finegrained_kernel().matmul_2d(x) - - out = run(torch.zeros(3)) - self.assertTrue(torch.equal(out, torch.ones(3))) - - def test_loader_is_compile_safe_when_warm(self): - """Warm path, which is the production order: eager warm-up, then compile. The loader hits - its short-circuit at trace time — the branch that must also return None.""" - kernel, _ = _fake_bundle() - kernel.matmul_2d = lambda x, *a, **k: x + 1 - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - load_finegrained_kernel() - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return load_finegrained_kernel().matmul_2d(x) - - out = run(torch.zeros(3)) - self.assertTrue(torch.equal(out, torch.ones(3))) - - def test_loader_binds_all_symbols(self): - kernel, _ = _fake_bundle() - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - bundle = load_finegrained_kernel() - self.assertIs(bundle.matmul_2d, kernel.matmul_2d) - self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) - - -@require_torch -class FineGrainedLinearMarshallingTest(unittest.TestCase): - def _run(self, **linear_kwargs): - kernel, rec = _fake_bundle() - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - x = torch.randn(3, 5, 64, dtype=torch.bfloat16) - w = torch.randn(32, 64).to(torch.float8_e4m3fn) - ws = torch.randn(1, 1, dtype=torch.float32) - out = finegrained_linear(x, w, ws, **linear_kwargs) - return out, rec.calls["matmul_2d"][-1] - - def test_scale_rides_positionally_and_block_size_dies(self): - out, call = self._run(block_size=[128, 128]) - # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here - self.assertEqual(call.args[0].shape, (15, 64)) - self.assertIsNone(call.args[2]) - self.assertEqual(call.args[3].dtype, torch.float32) - self.assertNotIn("block_size", call.kwargs) - self.assertEqual(out.shape, (3, 5, 32)) - - def test_static_activation_scale_is_As(self): - scale = torch.tensor(0.5) - _, call = self._run(activation_scale=scale) - self.assertIs(call.args[2], scale) - - def test_nvfp4_global_and_weight_only_format(self): - g = torch.tensor(2.0) - _, call = self._run(weight_global_scale=g, activation_format="bf16") - self.assertIs(call.kwargs["b_global_scale"], g) - self.assertEqual(call.kwargs["activation_format"], "bf16") - - def test_module_forward_threads_everything(self): - kernel, rec = _fake_bundle() - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - m = FineGrainedLinear( - 64, - 32, - block_size=None, - weight_format="nvfp4", - activation_format="bf16", - ) - m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) - out = m(torch.randn(2, 64, dtype=torch.bfloat16)) - call = rec.calls["matmul_2d"][-1] - self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) - self.assertEqual(call.kwargs["activation_format"], "bf16") - self.assertEqual(out.shape, (2, 32)) +class FineGrainedModulesTest(unittest.TestCase): + """The finegrained modules themselves: what they hold and what replaces what.""" - -@require_torch -class FineGrainedGroupedLinearSwapTest(unittest.TestCase): def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): """`replace_with_finegrained_layer` picks `FineGrainedGroupedLinear` for a block-diagonal linear — a plain one would collapse the groups into a single giant matmul and return the @@ -252,11 +84,8 @@ def __init__(self): self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) - -class FineGrainedEmbeddingTest(unittest.TestCase): - """A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, - rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`.""" - + # A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, + # rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`. def test_lookup_rescales_the_gathered_rows(self): torch.manual_seed(0) table = torch.randn(64, 16) @@ -281,262 +110,12 @@ def test_replacement_follows_the_patterns_and_the_skip_list(self): self.assertIs(type(model.embed_tokens), torch.nn.Embedding) self.assertIs(type(model.other), torch.nn.Embedding) - -@require_torch -class FineGrainedExpertsMarshallingTest(unittest.TestCase): - def _experts(self, **kw): - cfg = _Cfg() - weight_format = kw.pop("weight_format", "fp8") - m = FineGrainedExperts(cfg, block_size=(4, 4), weight_format=weight_format, **kw) - for name, p in m.named_parameters(): - if p is None: - continue - if p.dtype == torch.float8_e4m3fn: - p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) - elif p.dtype.is_floating_point: - p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) - return m - - def _route(self, tokens=6): - hs = torch.randn(tokens, 64, dtype=torch.bfloat16) - idx = torch.randint(0, 4, (tokens, 2)) - wts = torch.rand(tokens, 2) - return hs, idx, wts - - def test_eager_linear_threads_bias_global_and_activation_format(self): - """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: - the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" - kernel, rec = _fake_bundle() - m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") - p1, p2, p3 = _loaded(kernel) - linear = mock.Mock(wraps=fg.finegrained_linear) - with ( - p1, - p2, - p3, - mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), - mock.patch.object(fg, "finegrained_linear", linear), - ): - m(*self._route()) - self.assertTrue(linear.call_args_list) - for call in linear.call_args_list: - self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) - # this expert's NVFP4 weight global: one value for the down, and (1, 2) for the - # gate|up stack, whose halves modelopt calibrates separately - self.assertIn(tuple(call.kwargs["weight_global_scale"].shape), {(), (1, 2)}) - self.assertEqual(call.kwargs["activation_format"], "bf16") - # weight-only: the activations stay bf16, so no activation global is quantized against - self.assertIsNone(call.kwargs["input_global_scale"]) - for call in rec.calls["matmul_2d"]: - self.assertIsNotNone(call.kwargs["b_global_scale"]) - self.assertEqual(call.kwargs["activation_format"], "bf16") - - def test_batched_marshalling(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - hidden, top_k_index, top_k_weights = call.args - self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers - self.assertEqual(top_k_index.shape, (6, 2)) - self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) - # one scale tensor per projection: the affine Parameter, or the swizzled cache when the - # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) - self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) - self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) - self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name - self.assertIs(call.kwargs["gate"], True) - self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format - self.assertEqual(out.shape, (6, 64)) - - def test_grouped_marshalling(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_grouped"] - self.assertIs(call.kwargs["down_proj"], m.down_proj) - self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) - self.assertNotIn("moe_fused_batched", rec.calls) - - def test_activation_format_passes_through(self): - kernel, rec = _fake_bundle() - for activation_format in (None, "bf16", "mxfp8"): - m = self._experts(has_gate=True, activation_format=activation_format) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) - - def test_nvfp4_experts_thread_per_expert_globals(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, weight_format="nvfp4") - # the format table must resolve the ATTRIBUTE the forwards gate on, not just the - # param allocation — a None here silently drops the global at every forward - self.assertIsNotNone(m.global_scale_dtype) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) - self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) - # one weight global per expert on both projections: a gate|up stack calibrated per half - # is merged to that at load, so nothing downstream carries the pair - self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) - self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) - # the calibrated activation globals (the checkpoint's `input_scale`): one value for the - # gate_up, whose rows are quantized once before routing, and one per expert for the down, - # whose rows the gate_up epilogue requantizes per expert - self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) - self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) - self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) - self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) - - def test_weight_only_experts_hold_no_activation_global(self): - """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation - quant for a calibrated global to normalize: the module never allocates one, and the ops - would refuse a two-level path without it.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_grouped"] - self.assertIsNone(m.gate_up_proj_input_global_scale) - self.assertIsNone(m.down_proj_input_global_scale) - self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) - self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) - # the WEIGHT globals are the format's own second level and stay - self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) - - def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): - """A model whose experts norm the down output before the routing weights: - the swap carries the norm over, the kernel chain runs it on the routed rows, and the - eager loop applies it per expert application — the same place the reference forwards do. - A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) - - class _Recording(torch.nn.Module): - def __init__(self): - super().__init__() - self.rows = 0 - - def forward(self, x): - self.rows += x.shape[0] - return x - - m.post_expert_norm = _Recording() - with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): - m(*self._route()) - self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") - - def test_a_post_expert_norm_needs_nothing_but_the_norm(self): - """A class that declares a post-expert norm only has to HOLD it: every backend calls - `post_expert_norm` directly, so there is no per-class hook to define alongside it.""" - from transformers.integrations.moe import use_experts_implementation - - @use_experts_implementation(has_post_expert_norm=True) - class _JustTheNorm(torch.nn.Module): - def forward(self, hidden_states, top_k_index, top_k_weights): - return hidden_states - - self.assertFalse(hasattr(_JustTheNorm, "_apply_post_norm")) - - def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): - """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels - when they implement it, with the weight and epsilon the fused form needs, so the chain - folds it into its reduce instead of calling the module.""" - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True) - m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True - p1, p2, p3 = _loaded(kernel) - for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): - m.post_expert_norm_name = name - with p1, p2, p3: - fg.finegrained_grouped_mm_experts_forward(m, *self._route()) - call = rec.calls["moe_fused_grouped"][-1] - if fused: - self.assertEqual(call.kwargs["post_expert_norm"], name) - self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) - self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) - else: - self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) - self.assertIsNone(call.kwargs["post_expert_norm_weight"]) - - def test_biases_ride_the_kernel_chain(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, has_bias=True) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) - self.assertTrue(torch.isfinite(out).all()) - (call,) = rec.calls["moe_fused_batched"] - # the kernels add gate_up's bias before the gated split and down's before the - # routing-weight multiply; nothing is added host-side - self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) - self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) - - def test_static_activation_scheme_reaches_the_fused_chain(self): - """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a - MoE calibrates each expert separately, so the module holds one per expert. Both fused - chains take them per projection and the kernels apply each tile's own expert scale in - register — a host pre-quant could not, since top-k routes one row to several experts - whose scales differ. A dynamic module holds no such slot and passes `None`.""" - static, dynamic = ( - FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) - for scheme in ("static", "dynamic") - ) - held = static.gate_up_proj_activation_scale - self.assertIsNotNone(held, "the calibrated scale the chain needs") - self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") - # a dynamic module holds the slot empty, like every other optional parameter - self.assertIsNone(dynamic.gate_up_proj_activation_scale) - - kernel = mock.Mock() - kernel.get_supported_act_fns.return_value = ("silu",) - kernel.get_supported_norms.return_value = () - for module, calibrated in ((static, True), (dynamic, False)): - operands = module._moe_operands(kernel) - for projection in ("gate_up_proj", "down_proj"): - with self.subTest(scheme=module.activation_scheme, projection=projection): - scale = operands[f"{projection}_activation_scale"] - self.assertEqual(scale is not None, calibrated) - if calibrated: - self.assertEqual(scale.numel(), _Cfg.num_local_experts) - - def test_unfusable_act_fn_is_passed_as_the_module_glu(self): - kernel, rec = _fake_bundle() - m = self._experts(has_gate=True, has_bias=True) - m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - fg.finegrained_batched_mm_experts_forward(m, *self._route()) - (call,) = rec.calls["moe_fused_batched"] - # the kernels run the module's own GLU on the host between the two GEMMs — a new - # activation never waits for a kernel release; the bias still rides the GEMM - self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality - self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) - - -@require_torch -class FineGrainedScaleDtypeTest(unittest.TestCase): - """A scale's dtype is the format's, never the ambient default. - - `from_pretrained` sets the default dtype to the checkpoint's for the duration of model - construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it - as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches - that, whatever the format decides each scale should be. - """ - + # A scale's dtype is the format's, never the ambient default. + # + # `from_pretrained` sets the default dtype to the checkpoint's for the duration of model + # construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it + # as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches + # that, whatever the format decides each scale should be. def _dtypes(self, build, default): previous = torch.get_default_dtype() torch.set_default_dtype(default) @@ -576,73 +155,13 @@ def test_experts_scales_ignore_the_default_dtype(self): ) -class FineGrainedFusedNormGateTest(unittest.TestCase): - """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. - - Fusing runs the norm inside the launch that produced the rows, so an all-reduce that has to - land between the two has nowhere to go — and under intra-expert TP those rows are a partial - sum over the sharded intermediate, so the fused norm would normalize a fraction of the value. - The wrapped module does the reduce in its own forward, so the fallback is the correct path. - """ - - class _Norm(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(4)) - self.eps = 1e-6 - - def _operands(self, *, input_reduce: bool): - norm = self._Norm() - if input_reduce: - norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks - - module = mock.Mock() - module.has_gate = True - module.act_fn_name = "silu" - module.swiglu_alpha = None - module.has_post_expert_norm = True - module.post_expert_norm = norm - module.post_expert_norm_name = "input_scaled_rms_norm" - module._projection_slots = fg.FineGrainedExperts._projection_slots - kernel = mock.Mock() - kernel.get_supported_act_fns.return_value = ("silu",) - kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) - return fg.FineGrainedExperts._moe_operands(module, kernel), module - - def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): - operands, _ = self._operands(input_reduce=False) - self.assertEqual(operands["post_expert_norm"], "input_scaled_rms_norm") - self.assertIsNotNone(operands["post_expert_norm_weight"], "the fused form needs the norm's weight") - - def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): - operands, module = self._operands(input_reduce=True) - self.assertIs( - operands["post_expert_norm"], - module.post_expert_norm, - "fused past a collective: the kernel would normalize this rank's partial sum", - ) - self.assertIsNone(operands["post_expert_norm_weight"]) - - @require_torch class FrozenFp8ShimTest(unittest.TestCase): - def test_frozen_module_warns_and_is_self_contained(self): - import importlib - + def test_frozen_module_is_self_contained(self): import transformers.integrations.finegrained_fp8 as frozen from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer from transformers.utils.quantization_config import FineGrainedFP8Config - # the whole FILE is deprecated, so importing anything out of it is what warns. Safe to - # do at module scope only because `_LazyModule` never executes these unless asked: a - # plain `import transformers` does not reach them. - with warnings.catch_warnings(record=True) as raised: - warnings.simplefilter("always") - importlib.reload(frozen) - self.assertTrue( - [w for w in raised if issubclass(w.category, DeprecationWarning) and "frozen" in str(w.message)], - [str(w.message) for w in raised], - ) # the quantizer is still constructible — frozen means no new recipes, not removed FineGrainedFP8HfQuantizer(FineGrainedFP8Config()) # distinct machinery: the frozen classes are not the live ones @@ -918,9 +437,34 @@ def test_a_real_moe_models_experts_gain_their_companions(self): self.assertIn(f"layers.*.mlp.experts.{companion}", plan, companion) -class FineGrainedGroupsTest(unittest.TestCase): - """A quant config names a format per module SUBSET, because a checkpoint can be more than one.""" +@require_torch +class FineGrainedConfigTest(unittest.TestCase): + """The quantization config: formats per module group, and the skip-lists checkpoints ship.""" + + @parameterized.expand( + [ + (weights, activations, activations in supported) + for weights, supported in ( + ("fp8", (None, "fp8")), + ("mxfp8", (None, "mxfp8", "mxfp4", "bf16")), + ("mxfp4", (None, "mxfp8", "mxfp4", "bf16")), + ("nvfp4", (None, "nvfp4", "bf16")), + ) + for activations in (None, "bf16", "fp8", "mxfp8", "mxfp4", "nvfp4") + ] + ) + def test_an_activation_format_the_weights_cannot_take_fails_at_construction(self, weights, activations, ok): + """The kernels quantize activations only to what each weight format's MMA takes, and refuse + the rest at the first forward; the config refuses them up front.""" + from transformers.utils.quantization_config import FineGrainedConfig + + if ok: + FineGrainedConfig(quant_method=weights, activation_format=activations) + else: + with self.assertRaises(ValueError): + FineGrainedConfig(quant_method=weights, activation_format=activations) + # A quant config names a format per module SUBSET, because a checkpoint can be more than one. def _config(self, **kwargs): from transformers.utils.quantization_config import FineGrainedConfig @@ -994,7 +538,8 @@ def test_a_producers_config_groups_are_normalized(self): ) group = config.group_for("model.layers.0.self_attn.q_proj") self.assertEqual((group.quant_method, group.activation_format), ("nvfp4", "nvfp4")) - self.assertEqual(group.activation_scheme, "static") # `dynamic: False` + # NVFP4's `dynamic: False` is its calibrated activation global, not a static activation scale + self.assertEqual(group.activation_scheme, "dynamic") def test_a_format_we_do_not_serve_falls_back_rather_than_guessing(self): """Half a translation would quantize modules by a rule the producer did not write.""" @@ -1007,8 +552,6 @@ def test_a_format_we_do_not_serve_falls_back_rather_than_guessing(self): ) self.assertEqual(list(config.groups), ["default"]) - -class SubtreePatternTest(unittest.TestCase): def test_a_checkpoints_glob_skip_list_is_bounded(self): """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" @@ -1025,405 +568,3 @@ def test_a_checkpoints_glob_skip_list_is_bounded(self): self.assertFalse(should_convert_module(name, skip), name) for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): self.assertTrue(should_convert_module(name, skip), name) - - -@require_torch -class FineGrainedDeepGemmDispatchTest(unittest.TestCase): - """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled - scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine - and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 - scales have no swizzled layout, so the first never fires for the shape the second catches.""" - - def _routed_to(self, *, sm100, scale_ndim, trains=False): - """Which backend a block-FP8 linear reaches, given arch, scale layout, and whether the call - needs a gradient.""" - kernel, _ = _fake_bundle() - w = torch.randn(32, 64).to(torch.float8_e4m3fn) - s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) - p1, p2, p3 = _loaded(kernel) - with ( - p1, - p2, - p3, - mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), - mock.patch.object(deepgemm, "is_sm100", return_value=sm100), - mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, - ): - x = torch.randn(4, 64, dtype=torch.bfloat16, requires_grad=trains) - fg.finegrained_linear(x, w, s, block_size=[128, 128]) - return "deepgemm" if dg.called else "triton" - - def test_sm100_never_prefers_deepgemm(self): - self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") - - def test_pre_sm100_still_uses_deepgemm(self): - self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") - - def test_deepgemm_refuses_a_call_that_needs_a_gradient(self): - """DeepGEMM has no backward pass: every entry point fails a call that needs a gradient, - before touching the kernel.""" - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size = 256, 128 - experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8") - experts._deepgemm_disabled = False - hs = torch.randn(2, 256, dtype=torch.bfloat16, requires_grad=True) - routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) - w, s = torch.randn(32, 256).to(torch.float8_e4m3fn), torch.ones(1, 2) - calls = [ - lambda: deepgemm.deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128)), - lambda: deepgemm.deepgemm_fp8_fp4_experts_forward(experts, hs, *routing), - lambda: deepgemm.deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing), - ] - with mock.patch.object(deepgemm, "load_deepgemm_kernel") as load: - for call in calls: - with self.assertRaisesRegex(NotImplementedError, "DeepGEMM does not support training"): - call() - load.assert_not_called() - - def test_a_linear_that_needs_a_gradient_prefers_triton(self): - """Where DeepGEMM would otherwise win, a call that needs a gradient still goes to triton.""" - self.assertEqual(self._routed_to(sm100=False, scale_ndim=2, trains=True), "triton") - - def test_swizzled_scales_never_reach_deepgemm(self): - """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" - self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") - - def test_deepgemm_experts_refuse_swizzled_scales(self): - """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM - experts backend afterwards must fail loudly rather than read the permuted buffer as affine.""" - from transformers.integrations.deepgemm import ( - deepgemm_fp8_fp4_experts_forward, - deepgemm_fp8_fp4_megamoe_experts_forward, - ) - - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" - with ( - mock.patch.object(fg, "is_sm100", return_value=True), - mock.patch("torch.cuda.is_available", return_value=True), - ): - experts = fg.FineGrainedExperts(cfg, weight_format="mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) - hs = torch.randn(2, 256, dtype=torch.bfloat16) - idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) - for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): - with self.assertRaisesRegex(RuntimeError, "held swizzled"): - forward(experts, hs, idx, wts) - - -@require_torch -class FineGrainedScaleLayoutTest(unittest.TestCase): - """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 - scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the - scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's - `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" - - E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks - - def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): - cfg = _Cfg() - cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E - cfg._experts_implementation = impl - with ( - mock.patch.object(fg, "is_sm100", return_value=sm100), - mock.patch("torch.cuda.is_available", return_value=True), - ): - experts = fg.FineGrainedExperts( - cfg, block_size=(128, 128), weight_format=weight_format, activation_format=activation_format - ) - model = torch.nn.Module() - model.config = cfg - model.experts = experts - return model, experts - - def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): - _, experts = self._experts("mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) - self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) - self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) - self.assertTrue(experts.holds_interleaved_gate_up) - # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a - # float32 grid there; GPT-OSS's config has no scale_fmt at all) - _, experts = self._experts("mxfp4", activation_format="bf16") - self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) - - def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): - for kw in ( - {"weight_format": "mxfp8", "sm100": False}, - {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely - {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales - {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, - ): - with self.subTest(**kw): - _, experts = self._experts(**kw) - self.assertEqual( - experts.gate_up_proj_scale_inv.shape, - (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), - ) - # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block - _, experts = self._experts("fp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) - with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): - _, experts = self._experts("mxfp8") - self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) - - def test_megamoe_holds_gate_up_stacked(self): - _, experts = self._experts("fp8", impl="deepgemm_megamoe") - self.assertFalse(experts.holds_interleaved_gate_up) - - def _op_kernel(self): - kernel, _ = _fake_bundle() - kernel.swizzle_mx_scales = mock.Mock( - side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) - ) - kernel.unswizzle_mx_scales = mock.Mock( - side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) - ) - return kernel - - def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): - from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales - - model, experts = self._experts("mxfp8") - kernel = self._op_kernel() - op = FineGrainedSwizzleScales(hf_quantizer=None) - grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) - weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - # the loader hands a single target under its pattern, the full name (with suffix) alongside - out = op.convert( - {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" - ) - self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) - # the same converter's weight passes through - out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") - self.assertIs(out["mlp.experts.gate_up_proj"], weight) - # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully - out = op.convert( - { - "experts.down_proj": weight, - "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), - }, - model=model, - full_layer_name="experts.down_proj", - ) - self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) - self.assertIs(out["experts.down_proj"], weight) - # a module that holds affine scales keeps them - affine_model, _ = self._experts("mxfp4", activation_format="bf16") - out = op.convert( - {"mlp.experts.gate_up_proj": grid}, - model=affine_model, - full_layer_name="experts.gate_up_proj_scale_inv", - ) - self.assertIs(out["mlp.experts.gate_up_proj"], grid) - self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) - - def test_scale_container_op_round_trips_every_container(self): - """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: - the op brings both into the held e8m0 (exact cast / same bytes), records the container on - the module, and its reverse restores that container on save; weights and native scales - pass both ways.""" - from transformers.integrations.finegrained.conversions import FineGrainedScaleContainer - - op = FineGrainedScaleContainer(hf_quantizer=None) - native = torch.pow(2.0, torch.randint(-8, 8, (4, 256, 8)).float()).to(torch.float8_e8m0fnu) - for container, shipped in ((torch.float32, native.float()), (torch.uint8, native.view(torch.uint8))): - with self.subTest(container=container): - model, experts = self._experts("mxfp8", sm100=False) # affine e8m0 scale Parameters - out = op.convert( - {"mlp.experts.gate_up_proj": shipped}, - model=model, - full_layer_name="experts.gate_up_proj_scale_inv", - ) - held = out["mlp.experts.gate_up_proj"] - self.assertEqual(held.dtype, torch.float8_e8m0fnu) - self.assertTrue(torch.equal(held.view(torch.uint8), native.view(torch.uint8))) - self.assertIs(experts.scale_container_dtype, container) - back = op.reverse_op.convert({"x": held}, model=model)["x"] - self.assertEqual(back.dtype, container) - self.assertTrue(torch.equal(back, shipped)) - weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) - self.assertIs(op.reverse_op.convert({"w": weight}, model=model)["w"], weight) - model, experts = self._experts("mxfp8", sm100=False) - self.assertIs( - op.convert({"mlp.experts.down_proj": native}, model=model, full_layer_name="experts.down_proj_scale_inv")[ - "mlp.experts.down_proj" - ], - native, - ) - self.assertIsNone(experts.scale_container_dtype) - self.assertIs(op.reverse_op.convert({"x": native}, model=model)["x"], native) # nothing recorded: native stays - - def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): - from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales - - kernel = self._op_kernel() - reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op - artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) - affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") - self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) - kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) - out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") - self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) - - def _quantizer(self, quant_method): - from transformers.utils.quantization_config import FineGrainedConfig - - return _quantizer_for(FineGrainedConfig(quant_method=quant_method)).update_weight_conversions - - def _arch_converters(self): - from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter - - return [ - WeightConverter( - source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.up_proj.weight"], - target_patterns="mlp.experts.gate_up_proj", - operations=[MergeModulelist(dim=0), Concatenate(dim=1)], - ), - WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight", - target_patterns="mlp.experts.down_proj", - operations=[MergeModulelist(dim=0)], - ), - WeightConverter( - source_patterns=["mlp.fc1.weight", "mlp.fc2.weight"], - target_patterns="mlp.fc.weight", - operations=[Concatenate(dim=0)], - ), - ] - - def test_quantizer_attaches_the_layout_ops_to_expert_converters(self): - from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained.conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, - FineGrainedSwizzleScales, - ) - - converters = self._quantizer("mxfp8")(self._arch_converters()) - by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} - gate_up_ops = by_target["mlp.experts.gate_up_proj"].operations - # stacked [gate; up] -> [g0, u0, ...], the scale into the held dtype, then the swizzle packs - # the final row order (it reads 1-byte scales) - self.assertIsInstance(gate_up_ops[-3], FineGrainedInterleaveGateUp) - self.assertFalse(gate_up_ops[-3].inverse) - self.assertIsInstance(gate_up_ops[-2], FineGrainedScaleContainer) - self.assertIsInstance(gate_up_ops[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["mlp.experts.down_proj"].operations[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["weight_scale_inv"].operations[0], FineGrainedScaleContainer) # dense linears - self.assertEqual(len(by_target["mlp.fc.weight"].operations), 1) # dense converters untouched - # keys arriving under the fused names get the same layout - self.assertIn("experts.gate_up_proj_bias", by_target) # patterns are stored anchor-stripped - self.assertIsInstance(by_target["experts.down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) - # ...and saving reverses in the opposite order: unswizzle first, then de-interleave - reverse_ops = by_target["mlp.experts.gate_up_proj"].reverse_transform().operations - self.assertIsInstance(reverse_ops[0], FineGrainedSwizzleScales) - self.assertTrue(reverse_ops[0].inverse) - self.assertIsInstance(reverse_ops[2], FineGrainedInterleaveGateUp) - self.assertTrue(reverse_ops[2].inverse) - - def test_blocks_scales_converters_carry_the_layout_ops(self): - from transformers.core_model_loading import WeightConverter - from transformers.integrations.finegrained.conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedSwizzleScales, - ) - - converters = self._quantizer("mxfp4")([]) - by_target = {_targets(c)[0]: c for c in converters if isinstance(c, WeightConverter)} - ops = by_target["gate_up_proj"].operations - self.assertIsInstance(ops[-2], FineGrainedInterleaveGateUp) - self.assertIsInstance(ops[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["gate_up_proj_scale_inv"].operations[0], FineGrainedInterleaveGateUp) - self.assertIsInstance(by_target["down_proj"].operations[-1], FineGrainedSwizzleScales) - self.assertIsInstance(by_target["down_proj_scale_inv"].operations[-1], FineGrainedSwizzleScales) - - def test_catch_all_converters_deliver_the_full_parameter_name(self): - """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax - down_proj_scale_inv) goes through the real converter plumbing: the ops receive the tensor - under the source PATTERN and must hand back the target, which the loader expands to the - full name.""" - from transformers.core_model_loading import WeightConverter, rename_source_key - - model, _ = self._experts("mxfp8") - converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] - for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): - renamed, source_pattern = rename_source_key(key, [], converters) - self.assertEqual(renamed, key) - converter = next(c for c in converters if source_pattern in c.source_patterns) - tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) - converter.add_tensor(renamed, key, source_pattern, lambda t=tensor: t) - kernel = self._op_kernel() - p1, p2, p3 = _loaded(kernel) - with p1, p2, p3: - out = converter.convert(renamed, model=model) - self.assertEqual(list(out), [key]) - - def test_replacement_keeps_the_models_gate_up_convention(self): - """`replace_with_finegrained_layer` goes through `use_experts_implementation`, which stamps - the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must - survive it (it decides whether the loader interleaves).""" - from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM - from transformers.utils.quantization_config import FineGrainedConfig - - gpt_oss = GptOssConfig( - vocab_size=64, - hidden_size=64, - intermediate_size=32, - num_local_experts=2, - num_experts_per_tok=1, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - sliding_window=16, - ) - qwen = Qwen3MoeConfig( - vocab_size=64, - hidden_size=64, - intermediate_size=64, - moe_intermediate_size=32, - num_experts=2, - num_experts_per_tok=1, - num_hidden_layers=1, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - ) - for cls, cfg, quant_method, expected in ( - (GptOssForCausalLM, gpt_oss, "mxfp4", False), - (Qwen3MoeForCausalLM, qwen, "mxfp8", True), - ): - with self.subTest(model=cls.__name__), torch.device("meta"): - model = cls(cfg) - fg.replace_with_finegrained_layer( - model, modules_to_not_convert=[], quantization_config=FineGrainedConfig(quant_method=quant_method) - ) - experts = model.model.layers[0].mlp.experts - self.assertIsInstance(experts, fg.FineGrainedExperts) - self.assertIs(experts.is_concatenated, expected) - - def test_interleave_op_follows_the_model_and_the_experts_backend(self): - """The op reads the model's experts modules: a model whose own rows already come - interleaved (`is_concatenated=False`, GPT-OSS) keeps them, Mega MoE packs gate|up itself - and keeps the stack; every other case gets the interleaved order, and the reverse restores - the stack.""" - from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp - - stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] - op = FineGrainedInterleaveGateUp(hf_quantizer=None) - model, experts = self._experts("mxfp8") - out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] - torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) - back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] - torch.testing.assert_close(back, stacked) - experts.is_concatenated = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) - experts.is_concatenated = True - experts.holds_interleaved_gate_up = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) diff --git a/tests/integrations/finegrained/test_distributed.py b/tests/integrations/finegrained/test_distributed.py new file mode 100644 index 000000000000..88f77b32b297 --- /dev/null +++ b/tests/integrations/finegrained/test_distributed.py @@ -0,0 +1,545 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The sharded load paths: expert parallelism and intra-expert tensor parallelism must both reproduce +what `device_map` computes, for every format a shipped model quantizes with, and a checkpoint quantized +on the fly must save and reload to the same bytes.""" + +import os +import re +import socket +import subprocess +import tempfile + +import torch + +from transformers.testing_utils import ( + TestCasePlus, + cleanup, + require_torch_multi_accelerator, + torch_device, +) + +from .test_models import _checkpoint_bytes + + +_SHARDING_WORKER = """ +import importlib, os, sys, torch +from transformers.distributed import DistributedConfig + +model_dir, out_dir, modes, cls_path = sys.argv[1], sys.argv[2], sys.argv[3].split(","), sys.argv[4] +module_name, cls_name = cls_path.split(":") +model_cls = getattr(importlib.import_module(module_name), cls_name) +world = int(os.environ["WORLD_SIZE"]) +ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + +# both modes in ONE launch: they share the mesh, so the process group is created once and only +# the plan differs -- a second torchrun would pay another interpreter + CUDA start +for mode in modes: + model = model_cls.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=(mode == "ep")), + ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, "gate_up_proj", None) + if weight is None: + weight = experts.up_proj + local = weight.to_local() if hasattr(weight, "to_local") else weight + with torch.no_grad(): + logits = model(ids.to(model.device)).logits.float().cpu() + if int(os.environ["RANK"]) == 0: + torch.save({"logits": logits, "expert_local": tuple(local.shape)}, + os.path.join(out_dir, mode + ".pt")) + del model + torch.accelerator.empty_cache() +""" + + +def _checkpoint_expert_shape(model_dir): + """The UNSHARDED `(experts, rows)` of the stacked gate|up projection the module holds. + + Read from the checkpoint, which ships the experts either already fused (`experts.gate_up_proj`) + or one per expert (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`, whose rows are gate and + up separately). Taking the baseline from whichever leg ran first would compare one mode to + another, and a mode that placed nothing would still look sharded. + """ + import glob + import json + import os + + from safetensors import safe_open + + shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + index = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as fh: + weight_map = json.load(fh)["weight_map"] + shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) + + experts: set[int] = set() + rows = 0 + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shape = None + if re.search(r"\.experts\.gate_up_proj$", key): + shape = handle.get_slice(key).get_shape() + if len(shape) == 3: + return (shape[0], shape[1]) + if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): + experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) + rows = max(rows, handle.get_slice(key).get_shape()[0]) + if not experts or not rows: + raise AssertionError(f"no expert projection in {model_dir} to size the shard against") + return (len(experts), 2 * rows) # the module stacks gate and up into one row extent + + +@require_torch_multi_accelerator +class FineGrainedLoadPathEquivalenceTest(TestCasePlus): + """Expert parallelism and intra-expert tensor parallelism must both reproduce what + `device_map` gives, which is the baseline because it places modules across devices without + splitting a single tensor — no process group, no DTensor, no collectives. So it computes the + unsharded answer, on the multi-GPU path people actually deploy. + + Each model is built in the format it actually ships, because the conversion path differs by + format and by checkpoint layout — a per-expert FP8 checkpoint and a fused NVFP4 one reach the + experts through different converters, and bugs have hidden in exactly that gap. + + Three properties make this able to fail: + * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the + shard it already holds, so they come out correctly sized whatever the plan says and a + plan that shards no scale at all still produces the right answer. + * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty + (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is skipped EXPLICITLY rather + than passing vacuously. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + """ + + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}` — each model + in the format it actually ships, with the quantization config that format arrives under, + not just its name: block-FP8 carries a `weight_block_size`, and an NVFP4 checkpoint comes + from modelopt under `quant_algo`, which is remapped on construction. Dims are multiples + of the 128 block so a 2-way split stays block-aligned, and the expert count divides the + mesh.""" + from transformers import ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + DeepseekV4Config, + DeepseekV4ForCausalLM, + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + Glm4vMoeTextConfig, + Glm4vMoeVisionConfig, + GptOssConfig, + GptOssForCausalLM, + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLConfig, + Mistral4Config, + Mistral4ForCausalLM, + ) + from transformers.utils.quantization_config import FineGrainedConfig + + return { + "deepseek_v3-fp8": ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "n_group": 1, + "topk_group": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "q_lora_rank": None, + "kv_lora_rank": 32, + "qk_nope_head_dim": 32, + "qk_rope_head_dim": 16, + "v_head_dim": 32, + }, + # DeepSeek-V3 ships block-FP8: 128x128 weight blocks, activations quantized + # per token at run time + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), + ), + # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a + # calibrated activation scale per quantized module, which for a MoE is one per expert. + # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense + # counterpart of the same export. + "mistral4-fp8_tensor_static": ( + Mistral4Config, + Mistral4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static"), + ), + # MULTIMODAL: the experts' plans live on `text_config`, not on the config the + # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer + # one adds no companion while the weights still shard, which is how a real + # multimodal MoE ends up with whole scales against sharded weights. + "glm4v_moe-nvfp4": ( + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + { + "text_config": Glm4vMoeTextConfig( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, + rope_theta=10000, + tie_word_embeddings=True, + bos_token_id=0, + eos_token_id=0, + pad_token_id=0, + n_routed_experts=4, + n_shared_experts=1, + n_group=1, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ), + "vision_config": Glm4vMoeVisionConfig( + depth=2, + num_heads=4, + hidden_size=64, + out_hidden_size=256, + intermediate_size=64, + patch_size=14, + spatial_merge_size=1, + temporal_patch_size=2, + ), + }, + # the GLM NVFP4 checkpoints are modelopt exports: `quant_algo` names the format + # and `FineGrainedConfig` remaps it to nvfp4 at construction + FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4"), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases — and + # an EP plan that already names those biases, so the companion rules meet entries + # the model wrote itself + "gpt_oss-mxfp4": ( + GptOssConfig, + GptOssForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + }, + # GPT-OSS ships weight-only: raw bf16 activations against packed fp4 weights + FineGrainedConfig(quant_method="mxfp4", activation_format="bf16"), + ), + # MIXED precision, and the only entry whose expert format is not the quantization + # config's: `expert_dtype` is a model-config side-channel that makes the EXPERTS + # mxfp4 while the dense and attention paths stay block-FP8 — with scales in UE8M0 + # containers rather than fp32, the other `scale_fmt` + "deepseek_v4-fp4_experts+fp8_dense": ( + DeepseekV4Config, + DeepseekV4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "expert_dtype": "fp4", + }, + FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second MULTIMODAL nesting, in the group-32 MX format. The sub-config shape + # follows this model's own tester; only the MoE dims are raised to a multiple of + # the 128 block so a 2-way split stays block-aligned. + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + # 512 so the 2-way split leaves 256: the experts' down projection + # contracts over this, and a sharded 128 has no 128-wide swizzled tile + "intermediate_size": 512, + "dense_intermediate_size": 256, + "shared_intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "rotary_dim": 32, + "vocab_size": 64, + "max_position_embeddings": 32, + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "moe_layer_freq": [0, 1], + "layer_types": ["full_attention", "minimax_m3_sparse"], + "tie_word_embeddings": False, + "index_n_heads": 2, + "index_head_dim": 16, + "index_block_size": 8, + "index_topk_blocks": 4, + "index_local_blocks": 1, + }, + "vision_config": { + # 256 so the 2-way SPLIT is still 128-aligned: an MXFP8 weight with + # pre-swizzled scales is read in 128-wide K tiles, and a sharded 128 dim + # leaves 64 — which has none to offer, and the tuner has no config at all + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 14, + "patch_size": 14, + "temporal_patch_size": 2, + "spatial_merge_size": 1, + }, + "image_token_index": 4, + "video_token_index": 5, + "projector_hidden_size": 256, + "pad_token_id": 2, + }, + FineGrainedConfig(quant_method="mxfp8"), + ), + } + + @classmethod + def setUpClass(cls): + cls._tmp = tempfile.TemporaryDirectory() + # the kernels autotune per shape from a cold cache, which dwarfs everything else here + cls._env = { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), + } + os.environ.update(cls._env) + + @classmethod + def tearDownClass(cls): + cls._tmp.cleanup() + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + @staticmethod + def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("", 0)) + return sock.getsockname()[1] + + def _fixture(self, label): + """A PRE-QUANTIZED checkpoint of one tiny model, in the format that model ships.""" + + config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] + config = config_cls(**kwargs) + + bf16_dir = os.path.join(self._tmp.name, label, "bf16") + quant_dir = os.path.join(self._tmp.name, label, "quantized") + torch.manual_seed(0) + # BF16 is what every model here ships, and the dtype decides which kernel arms run: a + # float32 checkpoint (torch's default for a freshly built model, which `dtype="auto"` + # then faithfully reloads) sends the weight-only formats down arms `tl.dot_scaled` + # cannot serve at all, so the suite would exercise a dtype nobody deploys and leave the + # real one uncovered. Saving in bf16 makes `"auto"` mean bf16 for every load below. + model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) + quantized = model_cls.from_pretrained( + bf16_dir, + dtype="auto", + attn_implementation="eager", + quantization_config=quantization_config, + device_map=f"{torch_device}:0", + ) + quantized.save_pretrained(quant_dir, safe_serialization=True) + + # SAVE must restore the checkpoint's own layout: the layout ops (gate|up interleave, + # scale container, swizzle) each have a reverse, and a quantized model reloaded and + # written again has to land on the same bytes. A broken reverse writes a + # corrupt checkpoint silently — the module still reads back whatever it wrote. + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") + round_trip = os.path.join(self._tmp.name, label, "round_trip") + reloaded.save_pretrained(round_trip, safe_serialization=True) + saved, resaved = _checkpoint_bytes(quant_dir), _checkpoint_bytes(round_trip) + mismatched = sorted(key for key in saved.keys() | resaved.keys() if saved.get(key) != resaved.get(key)) + self.assertEqual(mismatched, [], f"{label}: save did not round-trip") + return quant_dir, config, model_cls + + def _logits(self, model_dir, model_cls, **load_kwargs): + model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + logits = model(ids).logits.float().cpu() + return logits + + def _rounding_floor(self, model_dir, model_cls, reference): + """How far this model's logits move when every sharded block is perturbed by BF16 rounding. + + Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different + order, so the first sharded op differs by an ULP. How far that travels is a property of + the MODEL, not of the sharding — a chain of MoE layers can amplify it a hundredfold while + a dense stack barely moves. Injecting the same magnitude on ONE device measures that + amplification directly, giving each model a budget its own conditioning earns. + """ + + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" + ).eval() + # Perturb EVERY block a sharded reduction passes through, not one: TP re-orders the sum + # in each of them, so the rounding accumulates down the stack instead of cancelling. + # Perturbing a single site with random noise measured LESS movement than sharding caused + # even at 14x the magnitude, which is what accumulation looks like from the wrong model. + blocks = [ + m + for n, m in model.named_modules() + if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 + ] + generator = torch.Generator(device=model.device).manual_seed(0) + + def perturb(module, args, output): + tensor = output[0] if isinstance(output, tuple) else output + if not torch.is_tensor(tensor): + return output + ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() + noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp + tensor = tensor + noise + return (tensor,) + output[1:] if isinstance(output, tuple) else tensor + + handles = [b.register_forward_hook(perturb) for b in blocks] + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + perturbed = model(ids).logits.float().cpu() + for handle in handles: + handle.remove() + return (perturbed - reference).abs().max().item() + + def _sharded(self, model_dir, model_cls, modes): + """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" + + script = os.path.join(self._tmp.name, "sharding_worker.py") + with open(script, "w", encoding="utf-8") as fh: + fh.write(_SHARDING_WORKER) + out = os.path.join(self._tmp.name, "out") + os.makedirs(out, exist_ok=True) + subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={self._free_port()}", + script, + model_dir, + out, + ",".join(modes), + # the model's own class: `AutoModelForCausalLM` cannot resolve a multimodal + # `ForConditionalGeneration` from its config + f"{model_cls.__module__}:{model_cls.__name__}", + ], + check=True, + env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, + ) + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} + + def test_every_load_path_agrees(self): + for label in self._model_table(): + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # the baseline, spread across both devices: `max_memory` forces a real split, + # since `auto` alone fits this whole model on one GPU and would quietly compare + # the sharded legs against a single-device load + reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) + noise_floor = self._rounding_floor(model_dir, model_cls, reference) + + # Only the modes this model's OWN plans shard experts under. A mode whose plan + # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places + # nothing, so running it would pass without being evidence of anything. A + # multimodal model keeps these on a sub-config. + owners = [config] + [ + c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) + ] + modes = [ + mode + for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) + if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) + ] + self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") + whole = _checkpoint_expert_shape(model_dir) + for mode, payload in self._sharded(model_dir, model_cls, modes).items(): + with self.subTest(model=label, mode=mode): + local = payload["expert_local"] + # against the CHECKPOINT's own shape, never another mode's shard: seeding + # this from the first leg made the second leg check EP against TP, so a + # mode that placed nothing still looked sharded (glm4v's TP left its + # experts replicated and passed). + self.assertLess( + local[0] * local[1], + whole[0] * whole[1], + f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " + f"this leg cannot catch a bad axis", + ) + # Sharding reorders reductions, so the legs differ by BF16 rounding + # whatever the axes — and how much that shows at the logits is the + # model's own business (the NVFP4 fixture amplifies ~100x, the dense + # ones barely). Hence each model's measured floor, not a fixed + # tolerance. A wrong shard axis lands orders of magnitude above it. + budget = max(2e-2, noise_floor) + gap = (payload["logits"] - reference).abs().max().item() + self.assertLessEqual( + gap, + budget, + f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " + f"model's own rounding budget {budget:.4g} — that is a sharding bug, " + f"not reduction order", + ) + if noise_floor <= 2e-2: + # only meaningful where rounding does NOT already move the argmax + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) diff --git a/tests/integrations/finegrained/test_forwards.py b/tests/integrations/finegrained/test_forwards.py index c4315c45e449..25008522d003 100644 --- a/tests/integrations/finegrained/test_forwards.py +++ b/tests/integrations/finegrained/test_forwards.py @@ -15,6 +15,7 @@ what those assert is passed to the kernels, these assert comes back out of them. Needs an accelerator, since nothing here is mocked.""" +import os import unittest from unittest import mock @@ -23,7 +24,7 @@ import transformers.integrations.finegrained.core as fg from transformers.integrations.finegrained import ( FineGrainedExperts, - finegrained_linear, + FineGrainedLinear, load_finegrained_kernel, ) from transformers.testing_utils import ( @@ -86,6 +87,22 @@ def _unpack_e2m1(packed): return pairs.reshape(*packed.shape[:-1], packed.shape[-1] * 2) +# `(weight format, activation format, scale format, linear floor, experts floor)`: every pairing the +# kernels serve (`None` is the weights' own format, `"bf16"` weight-only), against the dequantized weights. +# A floor is the activation quantization's own rounding, measured, with headroom; weight-only has nothing +# else that rounds, so a mis-read scale or a wrong arm lands far above it. +PAIRINGS = [ + ("fp8", None, "float", 0.05, 0.08), # W8A8 + ("fp8", None, "ue8m0", 0.05, 0.08), # W8A8, power-of-two block scales + ("mxfp8", None, "float", 0.05, 0.09), + ("mxfp4", None, "float", 0.2, 0.35), # W4A4 + ("mxfp4", "mxfp8", "float", 0.05, 0.09), # W4A8 + ("mxfp4", "bf16", "float", 0.01, 0.02), # W4A16 + ("nvfp4", None, "float", 0.17, 0.25), # W4A4 + ("nvfp4", "bf16", "float", 0.01, 0.02), # W4A16 +] + + @require_kernels @require_torch_accelerator class FineGrainedForwardTest(unittest.TestCase): @@ -98,33 +115,85 @@ class FineGrainedForwardTest(unittest.TestCase): @classmethod def setUpClass(cls): + # the kernels autotune per shape from a cold cache, which dwarfs everything else here + os.environ["FINEGRAINED_AUTOTUNE_TRIALS"] = "1" load_finegrained_kernel() - def _block_fp8_weight(self, N, K, *, ue8m0): - """(N, K) E4M3 + its (N/128, K/128) inv-scale grid, and the exact dequantized values the - kernel will see — so the reference is the quantization floor, not the pre-quant weight.""" - w = torch.randn(N, K, device=torch_device, dtype=torch.float32) - blocks = w.reshape(N // 128, 128, K // 128, 128) - amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12) - inv = amax / 448.0 - if ue8m0: # power-of-two scales: the tcgen05 dot_scaled format - inv = torch.pow(2.0, torch.ceil(torch.log2(inv))) - q = (blocks / inv).to(torch.float8_e4m3fn) - deq = (q.float() * inv).reshape(N, K) - return q.reshape(N, K), inv.reshape(N // 128, K // 128).contiguous(), deq - - def test_linear_forward_matches_the_quantization_floor(self): - """A real quantized linear against its own dequantized weight. Loose vs bf16 would pass - even if the kernel silently dropped a K-block; comparing to the floor does not.""" - M, N, K = 256, 512, 384 - for ue8m0 in (False, True): - with self.subTest(ue8m0=ue8m0): - w, s, deq = self._block_fp8_weight(N, K, ue8m0=ue8m0) - x = torch.randn(M, K, device=torch_device, dtype=torch.bfloat16) - out = finegrained_linear(x, w, s, block_size=[128, 128]) - ref = x.float() @ deq.t() - rel = ((out.float() - ref).norm() / ref.norm()).item() - self.assertLess(rel, 5e-2, f"ue8m0={ue8m0}: {rel:.2e} vs the dequantized weight") + def _quantize_into(self, module, key, value, model): + """`value` quantized by the kernels' own quantizers into `module`'s held layout.""" + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + for name, tensor in FineGrainedQuantize(hf_quantizer=None).convert({key: value}, model=model).items(): + setattr(module, name.rsplit(".", 1)[-1], torch.nn.Parameter(tensor, requires_grad=False)) + + def _dequantized(self, module, proj, scale, global_scale, rows, experts=None): + """The values the kernels multiply by: the held weight times its (unswizzled) scales and global.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + + scale = getattr(module, scale) + if scale.ndim == 5: + scale = load_finegrained_kernel().unswizzle_mx_scales(scale, rows, scale.shape[2] * 4, num_experts=experts) + weight = FineGrainedDequantize(None)._dequantize_one( + getattr(module, proj), scale.float(), output_dtype=torch.float32 + ) + global_scale = getattr(module, global_scale, None) + if global_scale is not None: + weight = weight * (global_scale.reshape(-1, 1, 1) if experts else global_scale) + return weight + + def _relative(self, out, reference): + return ((out.float() - reference).norm() / reference.norm()).item() + + def test_every_pairing_forwards_within_its_floor(self): + """Each weight × activation pairing, on the linear and the three experts forwards.""" + torch.manual_seed(0) + hidden, intermediate, num_experts = 512, 256, 4 + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = hidden, intermediate, num_experts + cfg._experts_implementation = "grouped_mm" + forwards = ( + ("grouped", fg.finegrained_grouped_mm_experts_forward), + ("batched", fg.finegrained_batched_mm_experts_forward), + ("eager", lambda module, *args: module(*args)), + ) + for weight_format, activation_format, scale_fmt, linear_floor, experts_floor in PAIRINGS: + storage = {"block_size": (128, 128), "scale_fmt": scale_fmt} if weight_format == "fp8" else {} + with self.subTest(weights=weight_format, activations=activation_format, scales=scale_fmt, module="linear"): + linear = FineGrainedLinear( + hidden, hidden, weight_format=weight_format, activation_format=activation_format, **storage + ).to(torch_device) + model = torch.nn.Module() + model.proj = linear + self._quantize_into(linear, "proj.weight", self._weight(hidden, hidden), model) + x = torch.randn(64, hidden, device=torch_device, dtype=torch.bfloat16) + reference = ( + x.float() + @ self._dequantized(linear, "weight", "weight_scale_inv", "weight_global_scale", hidden).t() + ) + self.assertLess(self._relative(linear(x), reference), linear_floor) + + experts = FineGrainedExperts( + cfg, weight_format=weight_format, activation_format=activation_format, **storage + ).to(torch_device) + model = torch.nn.Module() + model.experts = experts + dequantized = {} + for proj, rows, cols in (("gate_up_proj", 2 * intermediate, hidden), ("down_proj", hidden, intermediate)): + self._quantize_into(experts, f"experts.{proj}", self._weight(num_experts, rows, cols), model) + dequantized[proj] = self._dequantized( + experts, proj, f"{proj}_scale_inv", f"{proj}_weight_global_scale", rows, experts=num_experts + ) + x = torch.randn(8, hidden, device=torch_device, dtype=torch.bfloat16) + idx = torch.randint(0, num_experts, (8, 2), device=torch_device, dtype=torch.long) + wts = torch.rand(8, 2, device=torch_device, dtype=torch.bfloat16) + reference = self._moe_reference(dequantized, x, idx, wts, interleaved=experts.holds_interleaved_gate_up) + for name, forward in forwards: + with self.subTest(weights=weight_format, activations=activation_format, scales=scale_fmt, module=name): + self.assertLess(self._relative(forward(experts, x, idx, wts), reference), experts_floor) + + @staticmethod + def _weight(*shape): + return (torch.randn(*shape, device=torch_device) * 0.02).to(torch.bfloat16) def test_on_the_fly_group_formats_round_trip_within_their_floor(self): """MXFP8 / MXFP4 / NVFP4 on-the-fly quantization through the kernels' quantizers, emitted in @@ -234,15 +303,16 @@ def _nvfp4_experts_modelopt_layout(self, cfg, activation_format=None, magnitudes setattr(experts, target.split(".")[-1], torch.nn.Parameter(value, requires_grad=False)) return experts, dequantized, single_global - def _moe_reference(self, dequantized, x, idx, wts, post_norm=None): - """The experts chain in torch over the dequantized weights: gate|up (interleaved columns), - SwiGLU, down, the optional per-expert output norm, then the routing-weighted sum.""" + def _moe_reference(self, dequantized, x, idx, wts, post_norm=None, interleaved=True): + """The experts chain in torch over the dequantized weights: gate|up (interleaved columns, or + stacked halves), SwiGLU, down, the optional per-expert output norm, then the routing-weighted sum.""" out = torch.zeros_like(x, dtype=torch.float32) for token in range(x.shape[0]): for slot in range(idx.shape[1]): e = int(idx[token, slot]) pre = x[token].float() @ dequantized["gate_up_proj"][e].t() - inter = torch.nn.functional.silu(pre[0::2]) * pre[1::2] + gate, up = (pre[0::2], pre[1::2]) if interleaved else pre.chunk(2) + inter = torch.nn.functional.silu(gate) * up row = inter @ dequantized["down_proj"][e].t() if post_norm is not None: row = post_norm(row.to(x.dtype)).float() diff --git a/tests/integrations/finegrained/test_models.py b/tests/integrations/finegrained/test_models.py index b8d3d315ae96..e41585fe7c6e 100644 --- a/tests/integrations/finegrained/test_models.py +++ b/tests/integrations/finegrained/test_models.py @@ -11,551 +11,240 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""The sharded load paths on real checkpoints: expert parallelism and intra-expert tensor -parallelism must both reproduce what `device_map` computes, for every format a shipped model -quantizes with.""" +"""The shipped checkpoint layouts end to end: tiny random copies of the real checkpoints must load +cleanly, compute what their dequantized weights compute and the logits recorded for them, and save back +to their own layout.""" import os import re -import socket -import subprocess -import tempfile import torch +from parameterized import parameterized +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText from transformers.testing_utils import ( + Expectations, TestCasePlus, - backend_empty_cache, - require_torch_multi_accelerator, + cleanup, + require_kernels, + require_torch_accelerator, + slow, torch_device, ) -def _checkpoint_shapes(path): - """`{key: shape}` of a saved checkpoint, for comparing one save against another.""" +def _checkpoint_bytes(path): + """`{key: (dtype, shape, raw bytes)}` of a saved checkpoint, for comparing one save against another.""" import glob from safetensors import safe_open - shapes = {} + tensors = {} for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))): with safe_open(shard, framework="pt") as handle: for key in handle.keys(): - shapes[key] = tuple(handle.get_slice(key).get_shape()) - return shapes - - -_SHARDING_WORKER = """ -import importlib, os, sys, torch -from transformers.distributed import DistributedConfig - -model_dir, out_dir, modes, cls_path = sys.argv[1], sys.argv[2], sys.argv[3].split(","), sys.argv[4] -module_name, cls_name = cls_path.split(":") -model_cls = getattr(importlib.import_module(module_name), cls_name) -world = int(os.environ["WORLD_SIZE"]) -ids = torch.arange(16, dtype=torch.long).unsqueeze(0) - -# both modes in ONE launch: they share the mesh, so the process group is created once and only -# the plan differs -- a second torchrun would pay another interpreter + CUDA start -for mode in modes: - model = model_cls.from_pretrained( - model_dir, - dtype="auto", - attn_implementation="eager", - distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=(mode == "ep")), - ).eval() - experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) - weight = getattr(experts, "gate_up_proj", None) - if weight is None: - weight = experts.up_proj - local = weight.to_local() if hasattr(weight, "to_local") else weight - with torch.no_grad(): - logits = model(ids.to(model.device)).logits.float().cpu() - if int(os.environ["RANK"]) == 0: - torch.save({"logits": logits, "expert_local": tuple(local.shape)}, - os.path.join(out_dir, mode + ".pt")) - del model - torch.accelerator.empty_cache() -""" - - -def _checkpoint_expert_shape(model_dir): - """The UNSHARDED `(experts, rows)` of the stacked gate|up projection the module holds. - - Read from the checkpoint, which ships the experts either already fused (`experts.gate_up_proj`) - or one per expert (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`, whose rows are gate and - up separately). Taking the baseline from whichever leg ran first would compare one mode to - another, and a mode that placed nothing would still look sharded. - """ - import glob - import json - import os - - from safetensors import safe_open - - shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) - index = os.path.join(model_dir, "model.safetensors.index.json") - if os.path.exists(index): - with open(index, encoding="utf-8") as fh: - weight_map = json.load(fh)["weight_map"] - shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) - - experts: set[int] = set() - rows = 0 - for shard in shards: - with safe_open(shard, framework="pt") as handle: - for key in handle.keys(): - shape = None - if re.search(r"\.experts\.gate_up_proj$", key): - shape = handle.get_slice(key).get_shape() - if len(shape) == 3: - return (shape[0], shape[1]) - if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): - experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) - rows = max(rows, handle.get_slice(key).get_shape()[0]) - if not experts or not rows: - raise AssertionError(f"no expert projection in {model_dir} to size the shard against") - return (len(experts), 2 * rows) # the module stacks gate and up into one row extent - - -@require_torch_multi_accelerator -class FineGrainedLoadPathEquivalenceTest(TestCasePlus): - """Expert parallelism and intra-expert tensor parallelism must both reproduce what - `device_map` gives, which is the baseline because it places modules across devices without - splitting a single tensor — no process group, no DTensor, no collectives. So it computes the - unsharded answer, on the multi-GPU path people actually deploy. - - Each model is built in the format it actually ships, because the conversion path differs by - format and by checkpoint layout — a per-expert FP8 checkpoint and a fused NVFP4 one reach the - experts through different converters, and bugs have hidden in exactly that gap. - - Three properties make this able to fail: - * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the - shard it already holds, so they come out correctly sized whatever the plan says and a - plan that shards no scale at all still produces the right answer. - * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty - (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is skipped EXPLICITLY rather - than passing vacuously. - * every leg reports the local expert shard, so a leg that placed nothing fails loudly. - """ + tensor = handle.get_tensor(key) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + tensors[key] = (tensor.dtype, tuple(tensor.shape), raw) + return tensors + + +# Tiny random copies of the shipped checkpoints, shrunk from the real configs and written in their layouts. +# Each repo's `reference/` subfolder holds the same weights dequantized to a plain bf16 model; the budget is +# each format's own rounding against it, measured, with headroom. +TINY_CHECKPOINTS = { + "IlyasMoutawwakil/tiny-random-Qwen3-FP8": 0.1, + "IlyasMoutawwakil/tiny-random-Mistral3-FP8-static": 0.1, + "IlyasMoutawwakil/tiny-random-Mistral4-FP8-static": 0.15, + "IlyasMoutawwakil/tiny-random-Llama-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-DeepseekV3-FP8": 0.25, + "IlyasMoutawwakil/tiny-random-DeepseekV4-Flash": 0.25, + "IlyasMoutawwakil/tiny-random-GptOss-MXFP4": 0.05, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": 0.6, + "IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": 0.3, +} + + +# The logits `[0, :3, :8]` each checkpoint gives on `torch.arange(3, 35)`, per device, as the model integration +# tests record them: bit-stable across tuner caches, so a tight tolerance holds. +EXPECTED_SLICES = { + "IlyasMoutawwakil/tiny-random-Qwen3-FP8": Expectations( + { + ("cuda", (10, 0)): [ + [0.6094, -0.6055, -0.3262, 0.1396, -0.1270, -0.0420, 0.3750, -0.1592], + [0.1689, -0.4961, 0.1621, 0.0752, -0.1895, 0.0688, 0.3320, -0.0234], + [0.1069, -0.3027, 0.1050, 0.1553, 0.1699, -0.1572, 0.1699, 0.2598], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Mistral3-FP8-static": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1992, -0.4141, -0.1650, -0.1826, 0.1309, 0.4375, -0.2832, 0.2480], + [-0.1367, -0.2422, -0.1138, -0.3418, 0.2910, 0.3965, -0.0293, -0.0330], + [-0.1436, -0.1055, -0.3320, -0.3203, 0.3633, 0.2480, 0.1523, -0.1621], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Mistral4-FP8-static": Expectations( + { + ("cuda", (10, 0)): [ + [-0.0850, 0.1816, -0.2344, 0.5547, -0.0747, 0.6875, -0.2949, 0.1982], + [-0.0056, -0.0065, 0.0530, 0.0996, 0.1436, 0.2217, -0.3281, 0.1846], + [0.0055, 0.2559, -0.1045, -0.0031, 0.0608, 0.0449, 0.0962, 0.1416], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Llama-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.7617, -0.5938, -0.2832, 0.2559, -0.1021, 0.0082, 0.3418, -0.1699], + [0.6523, -0.7148, 0.2930, 0.3750, -0.5039, 0.0835, 0.2793, -0.0153], + [0.3457, -0.4785, 0.1924, 0.4102, -0.2715, -0.0835, 0.2070, 0.1602], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-DeepseekV3-FP8": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1260, -0.2051, -0.1836, -0.6523, 0.4395, -0.0771, -0.7461, 0.3613], + [-0.1025, -0.0640, -0.0132, -0.6602, 0.3672, -0.1270, -0.9023, -0.0786], + [-0.0874, 0.0552, 0.0442, -0.4824, 0.4375, -0.4121, -0.4980, -0.1836], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-DeepseekV4-Flash": Expectations( + { + ("cuda", (10, 0)): [ + [-0.3984, 0.3633, 0.1348, -0.0476, -0.1118, -0.0518, -0.0952, -0.4590], + [-0.4336, 0.1836, 0.0049, -0.1514, -0.0449, -0.2383, -0.1113, -0.4766], + [-0.2559, 0.0693, -0.0664, -0.2207, -0.2500, -0.1177, -0.1562, -0.2949], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GptOss-MXFP4": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1309, 0.2061, 0.2559, 0.5156, -0.2715, 0.5469, 0.4004, -0.0854], + [-0.2490, 0.6406, 0.1289, 0.4277, -0.3809, 0.3125, 0.3516, -0.0811], + [-0.2334, 0.1367, 0.2930, 0.6875, -0.1836, 0.4844, 0.1953, -0.2363], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.1709, -0.4199, -0.4199, -0.2402, -0.3125, 0.6016, 0.3008, 0.1914], + [-0.2051, -0.6328, -0.5391, 0.3184, 0.1299, 0.2773, 0.4824, 0.0938], + [0.0447, -0.5078, -0.3672, -0.1992, 0.1553, 0.6523, 0.4141, -0.0038], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": Expectations( + { + ("cuda", (10, 0)): [ + [0.2266, -0.2793, -0.3906, -0.1357, 0.0170, 0.4980, 0.1035, 0.2041], + [-0.0649, -0.2188, -0.5547, 0.2041, 0.1201, 0.3750, 0.4355, 0.1099], + [0.0908, -0.4492, -0.6094, -0.1836, 0.3203, 0.4453, 0.4199, 0.1172], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": Expectations( + { + ("cuda", (10, 0)): [ + [0.2334, -0.0771, -0.3633, 0.1670, -0.2832, 0.1611, 0.4609, 0.0574], + [0.4199, -0.1011, -0.2031, -0.2334, -0.2832, 0.1182, 0.5586, 0.1201], + [0.1836, -0.0981, -0.6836, 0.2969, 0.0684, 0.0942, 0.5625, -0.0635], + ], + } + ), # fmt: skip +} - @staticmethod - def _model_table(): - """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}` — each model - in the format it actually ships, with the quantization config that format arrives under, - not just its name: block-FP8 carries a `weight_block_size`, and an NVFP4 checkpoint comes - from modelopt under `quant_algo`, which is remapped on construction. Dims are multiples - of the 128 block so a 2-way split stays block-aligned, and the expert count divides the - mesh.""" - from transformers import ( - DeepseekV3Config, - DeepseekV3ForCausalLM, - DeepseekV4Config, - DeepseekV4ForCausalLM, - Glm4vMoeConfig, - Glm4vMoeForConditionalGeneration, - Glm4vMoeTextConfig, - Glm4vMoeVisionConfig, - GptOssConfig, - GptOssForCausalLM, - MiniMaxM3SparseForConditionalGeneration, - MiniMaxM3VLConfig, - Mistral4Config, - Mistral4ForCausalLM, - ) - from transformers.utils.quantization_config import FineGrainedConfig - return { - "deepseek_v3-fp8": ( - DeepseekV3Config, - DeepseekV3ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "n_group": 1, - "topk_group": 1, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - "q_lora_rank": None, - "kv_lora_rank": 32, - "qk_nope_head_dim": 32, - "qk_rope_head_dim": 16, - "v_head_dim": 32, - }, - # DeepSeek-V3 ships block-FP8: 128x128 weight blocks, activations quantized - # per token at run time - FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128)), - ), - # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a - # calibrated activation scale per quantized module, which for a MoE is one per expert. - # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense - # counterpart of the same export. - "mistral4-fp8_tensor_static": ( - Mistral4Config, - Mistral4ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - }, - FineGrainedConfig(quant_method="fp8", weight_block_size=None, activation_scheme="static"), - ), - # MULTIMODAL: the experts' plans live on `text_config`, not on the config the - # quantizer is handed — whose `base_model_ep_plan` is None. Reading only the outer - # one adds no companion while the weights still shard, which is how a real - # multimodal MoE ends up with whole scales against sharded weights. - "glm4v_moe-nvfp4": ( - Glm4vMoeConfig, - Glm4vMoeForConditionalGeneration, - { - "text_config": Glm4vMoeTextConfig( - vocab_size=64, - hidden_size=256, - intermediate_size=256, - moe_intermediate_size=256, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=32, - rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, - rope_theta=10000, - tie_word_embeddings=True, - bos_token_id=0, - eos_token_id=0, - pad_token_id=0, - n_routed_experts=4, - n_shared_experts=1, - n_group=1, - topk_group=1, - num_experts_per_tok=2, - first_k_dense_replace=0, - ), - "vision_config": Glm4vMoeVisionConfig( - depth=2, - num_heads=4, - hidden_size=64, - out_hidden_size=256, - intermediate_size=64, - patch_size=14, - spatial_merge_size=1, - temporal_patch_size=2, - ), - }, - # the GLM NVFP4 checkpoints are modelopt exports: `quant_algo` names the format - # and `FineGrainedConfig` remaps it to nvfp4 at construction - FineGrainedConfig(quant_method="modelopt", quant_algo="NVFP4"), - ), - # interleaved rows (`is_concatenated=False`), transposed, with expert biases — and - # an EP plan that already names those biases, so the companion rules meet entries - # the model wrote itself - "gpt_oss-mxfp4": ( - GptOssConfig, - GptOssForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "num_local_experts": 4, - "num_experts_per_tok": 2, - "max_position_embeddings": 32, - }, - # GPT-OSS ships weight-only: raw bf16 activations against packed fp4 weights - FineGrainedConfig(quant_method="mxfp4", activation_format="bf16"), - ), - # MIXED precision, and the only entry whose expert format is not the quantization - # config's: `expert_dtype` is a model-config side-channel that makes the EXPERTS - # mxfp4 while the dense and attention paths stay block-FP8 — with scales in UE8M0 - # containers rather than fp32, the other `scale_fmt` - "deepseek_v4-fp4_experts+fp8_dense": ( - DeepseekV4Config, - DeepseekV4ForCausalLM, - { - "vocab_size": 64, - "hidden_size": 256, - "intermediate_size": 256, - "moe_intermediate_size": 256, - "num_hidden_layers": 1, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "n_routed_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "first_k_dense_replace": 0, - "max_position_embeddings": 32, - "expert_dtype": "fp4", - }, - FineGrainedConfig(quant_method="fp8", weight_block_size=(128, 128), scale_fmt="ue8m0"), - ), - # a second MULTIMODAL nesting, in the group-32 MX format. The sub-config shape - # follows this model's own tester; only the MoE dims are raised to a multiple of - # the 128 block so a 2-way split stays block-aligned. - "minimax_m3_vl-mxfp8": ( - MiniMaxM3VLConfig, - MiniMaxM3SparseForConditionalGeneration, - { - "text_config": { - "hidden_size": 256, - # 512 so the 2-way split leaves 256: the experts' down projection - # contracts over this, and a sharded 128 has no 128-wide swizzled tile - "intermediate_size": 512, - "dense_intermediate_size": 256, - "shared_intermediate_size": 256, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "head_dim": 64, - "rotary_dim": 32, - "vocab_size": 64, - "max_position_embeddings": 32, - "bos_token_id": 0, - "eos_token_id": 1, - "pad_token_id": 2, - "num_local_experts": 4, - "num_experts_per_tok": 2, - "n_shared_experts": 1, - "moe_layer_freq": [0, 1], - "layer_types": ["full_attention", "minimax_m3_sparse"], - "tie_word_embeddings": False, - "index_n_heads": 2, - "index_head_dim": 16, - "index_block_size": 8, - "index_topk_blocks": 4, - "index_local_blocks": 1, - }, - "vision_config": { - # 256 so the 2-way SPLIT is still 128-aligned: an MXFP8 weight with - # pre-swizzled scales is read in 128-wide K tiles, and a sharded 128 dim - # leaves 64 — which has none to offer, and the tuner has no config at all - "hidden_size": 256, - "intermediate_size": 256, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_channels": 3, - "image_size": 14, - "patch_size": 14, - "temporal_patch_size": 2, - "spatial_merge_size": 1, - }, - "image_token_index": 4, - "video_token_index": 5, - "projector_hidden_size": 256, - "pad_token_id": 2, - }, - FineGrainedConfig(quant_method="mxfp8"), - ), - } +@slow +@require_kernels +@require_torch_accelerator +class FineGrainedTinyCheckpointTest(TestCasePlus): + """Each shipped layout, from the hub: it must load with every key accounted for, compute what + its dequantized weights compute within the format's own rounding, and save back to its own + layout. The checkpoints come from the formats' own producers (modelopt, OpenAI's and DeepSeek's + quantizers), so a converter is checked against what the format really is, not against what + this code believes it to be.""" @classmethod def setUpClass(cls): - cls._tmp = tempfile.TemporaryDirectory() # the kernels autotune per shape from a cold cache, which dwarfs everything else here - cls._env = { - "FINEGRAINED_AUTOTUNE_TRIALS": "1", - "TRITON_CACHE_DIR": os.path.join(cls._tmp.name, "triton"), - } - os.environ.update(cls._env) + os.environ["FINEGRAINED_AUTOTUNE_TRIALS"] = "1" - @classmethod - def tearDownClass(cls): - cls._tmp.cleanup() - - @staticmethod - def _free_port() -> int: - with socket.socket() as sock: - sock.bind(("", 0)) - return sock.getsockname()[1] + def setUp(self): + super().setUp() + # below that the quantizer dequantizes instead, which is not what is under test + if torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9): + self.skipTest("fine-grained quantized compute needs compute capability >= 8.9") - def _fixture(self, label): - """A PRE-QUANTIZED checkpoint of one tiny model, in the format that model ships.""" + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) - config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] - config = config_cls(**kwargs) - - bf16_dir = os.path.join(self._tmp.name, label, "bf16") - quant_dir = os.path.join(self._tmp.name, label, "quantized") - torch.manual_seed(0) - # BF16 is what every model here ships, and the dtype decides which kernel arms run: a - # float32 checkpoint (torch's default for a freshly built model, which `dtype="auto"` - # then faithfully reloads) sends the weight-only formats down arms `tl.dot_scaled` - # cannot serve at all, so the suite would exercise a dtype nobody deploys and leave the - # real one uncovered. Saving in bf16 makes `"auto"` mean bf16 for every load below. - model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) - quantized = model_cls.from_pretrained( - bf16_dir, - dtype="auto", - attn_implementation="eager", - quantization_config=quantization_config, - device_map=f"{torch_device}:0", - ) - quantized.save_pretrained(quant_dir, safe_serialization=True) - del quantized - backend_empty_cache(torch_device) - - # SAVE must restore the checkpoint's own layout: the layout ops (gate|up interleave, - # scale container, swizzle) each have a reverse, and a quantized model reloaded and - # written again has to land on the same keys and shapes. A broken reverse writes a - # corrupt checkpoint silently — the module still reads back whatever it wrote. - reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") - round_trip = os.path.join(self._tmp.name, label, "round_trip") - reloaded.save_pretrained(round_trip, safe_serialization=True) - del reloaded - backend_empty_cache(torch_device) - self.assertEqual( - _checkpoint_shapes(quant_dir), _checkpoint_shapes(round_trip), f"{label}: save did not round-trip" - ) - return quant_dir, config, model_cls + @staticmethod + def _load(repo, **kwargs): + config = AutoConfig.from_pretrained(repo, **kwargs) + model_cls = AutoModelForImageTextToText if "vision_config" in config.to_dict() else AutoModelForCausalLM + return model_cls.from_pretrained(repo, device_map=torch_device, output_loading_info=True, **kwargs) - def _logits(self, model_dir, model_cls, **load_kwargs): - model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() - ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + @staticmethod + def _unbuilt_mtp_keys(model, keys): + """The multi-token-prediction layer the checkpoint ships and the model does not build: it + sits at index `num_hidden_layers`, which the models' own ignore rule only names at the + full-size index.""" + mtp = model.config.get_text_config().num_hidden_layers + return {key for key in keys if re.search(rf"(^|\.)layers\.{mtp}\.", key)} + + def _not_loaded(self, model, keys): + """Checkpoint keys the model does not load, so a save does not write them back: the unbuilt MTP + layer, and what the load declares it ignores (an FP8 KV cache's scales).""" + ignored = model._keys_to_ignore_on_load_unexpected or [] + return self._unbuilt_mtp_keys(model, keys) | {k for k in keys if any(re.search(p, k) for p in ignored)} + + @parameterized.expand(TINY_CHECKPOINTS.items()) + def test_logits_match_the_dequantized_reference(self, repo, budget): + model, info = self._load(repo) + unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) + self.assertEqual((set(info["missing_keys"]), unexpected, set(info["mismatched_keys"])), (set(), set(), set())) + + reference, _ = self._load(repo, subfolder="reference", dtype=torch.bfloat16) + ids = torch.arange(3, 35, dtype=torch.long, device=model.device).unsqueeze(0) with torch.no_grad(): - logits = model(ids).logits.float().cpu() - del model - backend_empty_cache(torch_device) - return logits - - def _rounding_floor(self, model_dir, model_cls, reference): - """How far this model's logits move when every sharded block is perturbed by BF16 rounding. - - Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different - order, so the first sharded op differs by an ULP. How far that travels is a property of - the MODEL, not of the sharding — a chain of MoE layers can amplify it a hundredfold while - a dense stack barely moves. Injecting the same magnitude on ONE device measures that - amplification directly, giving each model a budget its own conditioning earns. - """ - - model = model_cls.from_pretrained( - model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" - ).eval() - # Perturb EVERY block a sharded reduction passes through, not one: TP re-orders the sum - # in each of them, so the rounding accumulates down the stack instead of cancelling. - # Perturbing a single site with random noise measured LESS movement than sharding caused - # even at 14x the magnitude, which is what accumulation looks like from the wrong model. - blocks = [ - m - for n, m in model.named_modules() - if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 - ] - generator = torch.Generator(device=model.device).manual_seed(0) - - def perturb(module, args, output): - tensor = output[0] if isinstance(output, tuple) else output - if not torch.is_tensor(tensor): - return output - ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() - noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp - tensor = tensor + noise - return (tensor,) + output[1:] if isinstance(output, tuple) else tensor - - handles = [b.register_forward_hook(perturb) for b in blocks] - ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + got, want = model(ids).logits.float(), reference(ids).logits.float() + error = ((got - want).norm() / want.norm()).item() + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + + @parameterized.expand(TINY_CHECKPOINTS) + def test_logits_match_the_expected_slice(self, repo): + model, _ = self._load(repo) + ids = torch.arange(3, 35, dtype=torch.long, device=model.device).unsqueeze(0) with torch.no_grad(): - perturbed = model(ids).logits.float().cpu() - for handle in handles: - handle.remove() - del model - backend_empty_cache(torch_device) - return (perturbed - reference).abs().max().item() - - def _sharded(self, model_dir, model_cls, modes): - """`{mode: payload}` from one 2-rank `torchrun`, through the real EP / TP load path.""" - - script = os.path.join(self._tmp.name, "sharding_worker.py") - with open(script, "w", encoding="utf-8") as fh: - fh.write(_SHARDING_WORKER) - out = os.path.join(self._tmp.name, "out") - os.makedirs(out, exist_ok=True) - subprocess.run( - [ - "torchrun", - "--nproc_per_node=2", - f"--master_port={self._free_port()}", - script, - model_dir, - out, - ",".join(modes), - # the model's own class: `AutoModelForCausalLM` cannot resolve a multimodal - # `ForConditionalGeneration` from its config - f"{model_cls.__module__}:{model_cls.__name__}", - ], - check=True, - env={**os.environ, **self._env, "TOKENIZERS_PARALLELISM": "false"}, - ) - return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} - - def test_every_load_path_agrees(self): - for label in self._model_table(): - with self.subTest(model=label): - model_dir, config, model_cls = self._fixture(label) - # the baseline, spread across both devices: `max_memory` forces a real split, - # since `auto` alone fits this whole model on one GPU and would quietly compare - # the sharded legs against a single-device load - reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) - noise_floor = self._rounding_floor(model_dir, model_cls, reference) - - # Only the modes this model's OWN plans shard experts under. A mode whose plan - # has no expert entry (GPT-OSS and DeepSeek-V4 ship no TP plan at all) places - # nothing, so running it would pass without being evidence of anything. A - # multimodal model keeps these on a sub-config. - owners = [config] + [ - c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) - ] - modes = [ - mode - for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) - if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) - ] - self.assertTrue(modes, f"{label}: no plan shards experts, so nothing is under test") - whole = _checkpoint_expert_shape(model_dir) - for mode, payload in self._sharded(model_dir, model_cls, modes).items(): - with self.subTest(model=label, mode=mode): - local = payload["expert_local"] - # against the CHECKPOINT's own shape, never another mode's shard: seeding - # this from the first leg made the second leg check EP against TP, so a - # mode that placed nothing still looked sharded (glm4v's TP left its - # experts replicated and passed). - self.assertLess( - local[0] * local[1], - whole[0] * whole[1], - f"{label}/{mode}: experts were not sharded ({local} of {whole}) — " - f"this leg cannot catch a bad axis", - ) - # Sharding reorders reductions, so the legs differ by BF16 rounding - # whatever the axes — and how much that shows at the logits is the - # model's own business (the NVFP4 fixture amplifies ~100x, the dense - # ones barely). Hence each model's measured floor, not a fixed - # tolerance. A wrong shard axis lands orders of magnitude above it. - budget = max(2e-2, noise_floor) - gap = (payload["logits"] - reference).abs().max().item() - self.assertLessEqual( - gap, - budget, - f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " - f"model's own rounding budget {budget:.4g} — that is a sharding bug, " - f"not reduction order", - ) - if noise_floor <= 2e-2: - # only meaningful where rounding does NOT already move the argmax - self.assertTrue( - torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), - f"{label}/{mode}: argmax diverged from the unsharded model", - ) + logits = model(ids).logits[0, :3, :8].float().cpu() + expected = torch.tensor(EXPECTED_SLICES[repo].get_expectation()) + torch.testing.assert_close(logits, expected, rtol=1e-2, atol=5e-3) + + @parameterized.expand(TINY_CHECKPOINTS) + def test_save_writes_the_checkpoint_layout_back(self, repo): + """A pre-quantized model saved again lands on its checkpoint's keys, dtypes and bytes. The + one exception is by design: NVFP4's gate and up `input_scale` share one activation global, + which a save writes back to both.""" + from huggingface_hub import snapshot_download + + model, _ = self._load(repo) + saved = os.path.join(self.get_auto_remove_tmp_dir(), "saved") + model.save_pretrained(saved) + source = _checkpoint_bytes(snapshot_download(repo, allow_patterns=["*.safetensors", "*.json"])) + source = {key: value for key, value in source.items() if key not in self._not_loaded(model, source)} + written = _checkpoint_bytes(saved) + self.assertEqual(sorted(written), sorted(source)) + for key, (dtype, shape, raw) in written.items(): + with self.subTest(key=key): + self.assertEqual((dtype, shape), source[key][:2]) + if not re.search(r"\.(gate|up)_proj\.input_scale$", key): + self.assertEqual(raw, source[key][2]) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py new file mode 100644 index 000000000000..987c776ff9ae --- /dev/null +++ b/tests/kernels/test_finegrained.py @@ -0,0 +1,577 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""What the `finegrained` integration hands `kernels-community/finegrained-kernels`: the kernel bundle is +mocked, so these pin the call contract (As-positional / no-block_size / expert_start / b_global_scale), the +loader and the DeepGEMM dispatch without a GPU.""" + +import unittest +from dataclasses import dataclass +from unittest import mock + +import torch + +import transformers.integrations.deepgemm as deepgemm +import transformers.integrations.finegrained.core as fg +from transformers.integrations.finegrained import ( + FineGrainedExperts, + FineGrainedLinear, + finegrained_linear, + load_finegrained_kernel, +) +from transformers.testing_utils import require_torch + +from ..integrations.finegrained.test_core import _Cfg + + +@dataclass +class _Call: + args: tuple + kwargs: dict + + +class _Recorder: + """Records every call and returns a correctly-shaped zero tensor for matmuls.""" + + def __init__(self): + self.calls = {} + + def _op(self, name): + def op(*args, **kwargs): + self.calls.setdefault(name, []).append(_Call(args, kwargs)) + if name == "swizzle_mx_scales": + return args[0] + if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) + hidden, down = args[0], kwargs["down_proj"] + return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) + a, b = args[0], args[1] + n = b.shape[-2] + gather = kwargs.get("gather_idx") + rows = gather.shape[0] if gather is not None else a.shape[0] + return a.new_zeros(rows, n, dtype=torch.bfloat16) + + return op + + +def fake_finegrained_kernel(): + """A `finegrained-kernels` stand-in whose ops record their calls and return correctly shaped zeros.""" + rec = _Recorder() + kernel = mock.Mock() + for name in ( + "matmul_2d", + "matmul_batched", + "matmul_grouped", + "moe_fused_batched", + "moe_fused_grouped", + "swizzle_mx_scales", + "unswizzle_mx_scales", + "mxfp8_act_quant", + "mxfp4_act_quant", + "nvfp4_act_quant", + ): + setattr(kernel, name, rec._op(name)) + kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") + kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") + return kernel, rec + + +def patch_finegrained_kernel(kernel): + """Patches that make the finegrained integration load `kernel` as its bundle.""" + return ( + mock.patch.object(fg, "_FINEGRAINED", None), + mock.patch.object(fg, "is_kernels_available", return_value=True), + mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), + ) + + +@require_torch +class FineGrainedLoaderTest(unittest.TestCase): + def test_loader_requires_every_symbol(self): + kernel, _ = fake_finegrained_kernel() + del kernel.matmul_grouped + # Mock auto-creates attributes; force the miss + kernel.matmul_grouped = None + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, self.assertRaises(ImportError) as ctx: + load_finegrained_kernel() + self.assertIn("matmul_grouped", str(ctx.exception)) + + def test_loader_is_compile_safe_cold(self): + """Cold path: the compiled call is first to load, so the opaque loader node runs its full + body under compile and must return None, never the bundle (`torch.* op returned + non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_is_compile_safe_when_warm(self): + """Warm path, which is the production order: eager warm-up, then compile. The loader hits + its short-circuit at trace time — the branch that must also return None.""" + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + load_finegrained_kernel() + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_binds_all_symbols(self): + kernel, _ = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + bundle = load_finegrained_kernel() + self.assertIs(bundle.matmul_2d, kernel.matmul_2d) + self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) + + +@require_torch +class FineGrainedLinearMarshallingTest(unittest.TestCase): + def _run(self, **linear_kwargs): + kernel, rec = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + x = torch.randn(3, 5, 64, dtype=torch.bfloat16) + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + ws = torch.randn(1, 1, dtype=torch.float32) + out = finegrained_linear(x, w, ws, **linear_kwargs) + return out, rec.calls["matmul_2d"][-1] + + def test_scale_rides_positionally_and_block_size_dies(self): + out, call = self._run(block_size=[128, 128]) + # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here + self.assertEqual(call.args[0].shape, (15, 64)) + self.assertIsNone(call.args[2]) + self.assertEqual(call.args[3].dtype, torch.float32) + self.assertNotIn("block_size", call.kwargs) + self.assertEqual(out.shape, (3, 5, 32)) + + def test_static_activation_scale_is_As(self): + scale = torch.tensor(0.5) + _, call = self._run(activation_scale=scale) + self.assertIs(call.args[2], scale) + + def test_nvfp4_global_and_weight_only_format(self): + g = torch.tensor(2.0) + _, call = self._run(weight_global_scale=g, activation_format="bf16") + self.assertIs(call.kwargs["b_global_scale"], g) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_module_forward_threads_everything(self): + kernel, rec = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m = FineGrainedLinear( + 64, + 32, + block_size=None, + weight_format="nvfp4", + activation_format="bf16", + ) + m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) + out = m(torch.randn(2, 64, dtype=torch.bfloat16)) + call = rec.calls["matmul_2d"][-1] + self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) + self.assertEqual(call.kwargs["activation_format"], "bf16") + self.assertEqual(out.shape, (2, 32)) + + +@require_torch +class FineGrainedExpertsMarshallingTest(unittest.TestCase): + def _experts(self, **kw): + cfg = _Cfg() + weight_format = kw.pop("weight_format", "fp8") + m = FineGrainedExperts(cfg, block_size=(4, 4), weight_format=weight_format, **kw) + for name, p in m.named_parameters(): + if p is None: + continue + if p.dtype == torch.float8_e4m3fn: + p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) + elif p.dtype.is_floating_point: + p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) + return m + + def _route(self, tokens=6): + hs = torch.randn(tokens, 64, dtype=torch.bfloat16) + idx = torch.randint(0, 4, (tokens, 2)) + wts = torch.rand(tokens, 2) + return hs, idx, wts + + def test_eager_linear_threads_bias_global_and_activation_format(self): + """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: + the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") + p1, p2, p3 = patch_finegrained_kernel(kernel) + linear = mock.Mock(wraps=fg.finegrained_linear) + with ( + p1, + p2, + p3, + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), + mock.patch.object(fg, "finegrained_linear", linear), + ): + m(*self._route()) + self.assertTrue(linear.call_args_list) + for call in linear.call_args_list: + self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) + # this expert's NVFP4 weight global: one value for the down, and (1, 2) for the + # gate|up stack, whose halves modelopt calibrates separately + self.assertIn(tuple(call.kwargs["weight_global_scale"].shape), {(), (1, 2)}) + self.assertEqual(call.kwargs["activation_format"], "bf16") + # weight-only: the activations stay bf16, so no activation global is quantized against + self.assertIsNone(call.kwargs["input_global_scale"]) + for call in rec.calls["matmul_2d"]: + self.assertIsNotNone(call.kwargs["b_global_scale"]) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_batched_marshalling(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + hidden, top_k_index, top_k_weights = call.args + self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers + self.assertEqual(top_k_index.shape, (6, 2)) + self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) + # one scale tensor per projection: the affine Parameter, or the swizzled cache when the + # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) + self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) + self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) + self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name + self.assertIs(call.kwargs["gate"], True) + self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format + self.assertEqual(out.shape, (6, 64)) + + def test_grouped_marshalling(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIs(call.kwargs["down_proj"], m.down_proj) + self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) + self.assertNotIn("moe_fused_batched", rec.calls) + + def test_activation_format_passes_through(self): + kernel, rec = fake_finegrained_kernel() + for activation_format in (None, "bf16", "mxfp8"): + m = self._experts(has_gate=True, activation_format=activation_format) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], activation_format) + + def test_nvfp4_experts_thread_per_expert_globals(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, weight_format="nvfp4") + # the format table must resolve the ATTRIBUTE the forwards gate on, not just the + # param allocation — a None here silently drops the global at every forward + self.assertIsNotNone(m.global_scale_dtype) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) + self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) + # one weight global per expert on both projections: a gate|up stack calibrated per half + # is merged to that at load, so nothing downstream carries the pair + self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) + # the calibrated activation globals (the checkpoint's `input_scale`): one value for the + # gate_up, whose rows are quantized once before routing, and one per expert for the down, + # whose rows the gate_up epilogue requantizes per expert + self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) + self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) + self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) + self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) + + def test_weight_only_experts_hold_no_activation_global(self): + """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation + quant for a calibrated global to normalize: the module never allocates one, and the ops + would refuse a two-level path without it.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIsNone(m.gate_up_proj_input_global_scale) + self.assertIsNone(m.down_proj_input_global_scale) + self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) + self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) + # the WEIGHT globals are the format's own second level and stay + self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) + + def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): + """A model whose experts norm the down output before the routing weights: + the swap carries the norm over, the kernel chain runs it on the routed rows, and the + eager loop applies it per expert application — the same place the reference forwards do. + A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) + + class _Recording(torch.nn.Module): + def __init__(self): + super().__init__() + self.rows = 0 + + def forward(self, x): + self.rows += x.shape[0] + return x + + m.post_expert_norm = _Recording() + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m(*self._route()) + self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") + + def test_a_post_expert_norm_needs_nothing_but_the_norm(self): + """A class that declares a post-expert norm only has to HOLD it: every backend calls + `post_expert_norm` directly, so there is no per-class hook to define alongside it.""" + from transformers.integrations.moe import use_experts_implementation + + @use_experts_implementation(has_post_expert_norm=True) + class _JustTheNorm(torch.nn.Module): + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states + + self.assertFalse(hasattr(_JustTheNorm, "_apply_post_norm")) + + def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): + """The norm follows `act_fn`'s shape: the name the model declared goes to the kernels + when they implement it, with the weight and epsilon the fused form needs, so the chain + folds it into its reduce instead of calling the module.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True + p1, p2, p3 = patch_finegrained_kernel(kernel) + for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): + m.post_expert_norm_name = name + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + call = rec.calls["moe_fused_grouped"][-1] + if fused: + self.assertEqual(call.kwargs["post_expert_norm"], name) + self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) + self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) + else: + self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) + self.assertIsNone(call.kwargs["post_expert_norm_weight"]) + + def test_biases_ride_the_kernel_chain(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, has_bias=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertTrue(torch.isfinite(out).all()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels add gate_up's bias before the gated split and down's before the + # routing-weight multiply; nothing is added host-side + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) + + def test_static_activation_scheme_reaches_the_fused_chain(self): + """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a + MoE calibrates each expert separately, so the module holds one per expert. Both fused + chains take them per projection and the kernels apply each tile's own expert scale in + register — a host pre-quant could not, since top-k routes one row to several experts + whose scales differ. A dynamic module holds no such slot and passes `None`.""" + static, dynamic = ( + FineGrainedExperts(_Cfg(), block_size=(128, 128), activation_scheme=scheme) + for scheme in ("static", "dynamic") + ) + held = static.gate_up_proj_activation_scale + self.assertIsNotNone(held, "the calibrated scale the chain needs") + self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") + # a dynamic module holds the slot empty, like every other optional parameter + self.assertIsNone(dynamic.gate_up_proj_activation_scale) + + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = () + for module, calibrated in ((static, True), (dynamic, False)): + operands = module._moe_operands(kernel) + for projection in ("gate_up_proj", "down_proj"): + with self.subTest(scheme=module.activation_scheme, projection=projection): + scale = operands[f"{projection}_activation_scale"] + self.assertEqual(scale is not None, calibrated) + if calibrated: + self.assertEqual(scale.numel(), _Cfg.num_local_experts) + + def test_unfusable_act_fn_is_passed_as_the_module_glu(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, has_bias=True) + m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels run the module's own GLU on the host between the two GEMMs — a new + # activation never waits for a kernel release; the bias still rides the GEMM + self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + + +@require_torch +class FineGrainedDeepGemmDispatchTest(unittest.TestCase): + """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled + scale is not readable as row-major, so DeepGEMM would consume a permuted buffer as affine + and silently return garbage) and an SM100 perf one. They cover different cases — block-FP8 + scales have no swizzled layout, so the first never fires for the shape the second catches.""" + + def _routed_to(self, *, sm100, scale_ndim, trains=False): + """Which backend a block-FP8 linear reaches, given arch, scale layout, and whether the call + needs a gradient.""" + kernel, _ = fake_finegrained_kernel() + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with ( + p1, + p2, + p3, + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=True), + mock.patch.object(deepgemm, "is_sm100", return_value=sm100), + mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as dg, + ): + x = torch.randn(4, 64, dtype=torch.bfloat16, requires_grad=trains) + fg.finegrained_linear(x, w, s, block_size=[128, 128]) + return "deepgemm" if dg.called else "triton" + + def test_sm100_never_prefers_deepgemm(self): + self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") + + def test_pre_sm100_still_uses_deepgemm(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") + + def test_deepgemm_refuses_a_call_that_needs_a_gradient(self): + """DeepGEMM has no backward pass: every entry point fails a call that needs a gradient, + before touching the kernel.""" + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size = 256, 128 + experts = fg.FineGrainedExperts(cfg, block_size=(128, 128), weight_format="fp8") + experts._deepgemm_disabled = False + hs = torch.randn(2, 256, dtype=torch.bfloat16, requires_grad=True) + routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + w, s = torch.randn(32, 256).to(torch.float8_e4m3fn), torch.ones(1, 2) + calls = [ + lambda: deepgemm.deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128)), + lambda: deepgemm.deepgemm_fp8_fp4_experts_forward(experts, hs, *routing), + lambda: deepgemm.deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing), + ] + with mock.patch.object(deepgemm, "load_deepgemm_kernel") as load: + for call in calls: + with self.assertRaises(NotImplementedError): + call() + load.assert_not_called() + + def test_a_linear_that_needs_a_gradient_prefers_triton(self): + """Where DeepGEMM would otherwise win, a call that needs a gradient still goes to triton.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2, trains=True), "triton") + + def test_swizzled_scales_never_reach_deepgemm(self): + """Correctness gate, and it must hold on any arch — not just where the perf gate does.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + + def test_deepgemm_experts_refuse_swizzled_scales(self): + """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM + experts backend afterwards must fail loudly rather than read the permuted buffer as affine.""" + from transformers.integrations.deepgemm import ( + deepgemm_fp8_fp4_experts_forward, + deepgemm_fp8_fp4_megamoe_experts_forward, + ) + + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" + with ( + mock.patch.object(fg, "is_sm100", return_value=True), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts(cfg, weight_format="mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) + hs = torch.randn(2, 256, dtype=torch.bfloat16) + idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.assertRaises(RuntimeError): + forward(experts, hs, idx, wts) + + +class FineGrainedFusedNormGateTest(unittest.TestCase): + """A norm the kernels can fuse must NOT be fused when a forward collective sits on it. + + Fusing runs the norm inside the launch that produced the rows, so an all-reduce that has to + land between the two has nowhere to go — and under intra-expert TP those rows are a partial + sum over the sharded intermediate, so the fused norm would normalize a fraction of the value. + The wrapped module does the reduce in its own forward, so the fallback is the correct path. + """ + + class _Norm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(4)) + self.eps = 1e-6 + + def _operands(self, *, input_reduce: bool): + norm = self._Norm() + if input_reduce: + norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks + + module = mock.Mock() + module.has_gate = True + module.act_fn_name = "silu" + module.swiglu_alpha = None + module.has_post_expert_norm = True + module.post_expert_norm = norm + module.post_expert_norm_name = "input_scaled_rms_norm" + module._projection_slots = fg.FineGrainedExperts._projection_slots + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) + return fg.FineGrainedExperts._moe_operands(module, kernel), module + + def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): + operands, _ = self._operands(input_reduce=False) + self.assertEqual(operands["post_expert_norm"], "input_scaled_rms_norm") + self.assertIsNotNone(operands["post_expert_norm_weight"], "the fused form needs the norm's weight") + + def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): + operands, module = self._operands(input_reduce=True) + self.assertIs( + operands["post_expert_norm"], + module.post_expert_norm, + "fused past a collective: the kernel would normalize this rank's partial sum", + ) + self.assertIsNone(operands["post_expert_norm_weight"]) diff --git a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py index ca13e05a7813..8e75dc4fcde6 100644 --- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py +++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py @@ -185,22 +185,9 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l self.assertEqual(keys.shape[3], head_dim) self.assertEqual(keys.shape, values.shape) - @unittest.skip( - reason=( - "V4's conversion mapping is two-pass: a structural prefix rename " - "(``layers.X.attn.`` → ``model.layers.X.self_attn.``) runs first, then specific in-prefix " - "renames operate on the already-prefixed HF-form keys (``model.layers.X.self_attn.compressor.norm.`` " - "→ ``...compressor.kv_norm.``). This split is load-bearing for save / load round-tripping — " - "any single-pass ordering loses information in either direction (the general prefix rule " - "and a specific in-prefix rule both want to match the same upstream key, and one of the " - "two directions ends up with the general rule stealing the match). The base " - "``test_reverse_loading_mapping`` checks every source pattern against the *upstream-form* " - "serialized keys, so the Pass 2 patterns (written in HF form) inherently can't satisfy " - "that invariant. The actual round-trip is exercised by ``test_save_load``." - ) - ) def test_reverse_loading_mapping(self): - pass + # the base model has no `model.` prefix for the checkpoint's prefix renames to reverse onto + super().test_reverse_loading_mapping(skip_base_model=True) @unittest.skip( reason=( diff --git a/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py b/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py index 1d220e55500d..55b5f3e3854c 100644 --- a/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py +++ b/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py @@ -258,9 +258,9 @@ def test_indexer_selects_blocks_per_head(self): def test_config(self): self.config_tester.run_common_tests() - @unittest.skip(reason="IDK exactly why, can be adressed later") def test_reverse_loading_mapping(self): - pass + # the base model has no `model.` prefix for the `^language_model\.model\.` rename to reverse onto, as in llava + super().test_reverse_loading_mapping(skip_base_model=True) @unittest.skip( reason=( From 6a7c92e3c6cccb86651922bba417298b25253edc Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 25 Sep 2026 11:28:16 -0700 Subject: [PATCH 59/68] Keep pre-quantized scale dtypes through the dtype plan The quantizer's `dtype_plan` joins the model's in `from_pretrained` (`None` keeps the checkpoint's dtype), replacing the `param_keeps_checkpoint_dtype` hook in the loader. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/core_model_loading.py | 1 - src/transformers/modeling_utils.py | 2 ++ src/transformers/quantizers/base.py | 11 ++++------- src/transformers/quantizers/finegrained/base.py | 8 +++----- tests/integrations/finegrained/test_conversions.py | 4 +++- 5 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index eb1e18498cb2..24d5d960f018 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1700,7 +1700,6 @@ def convert_and_load_state_dict_in_model( if hasattr(tensor, "get_dtype") else tensor.is_floating_point() ) - or hf_quantizer.param_keeps_checkpoint_dtype(renamed_key) ) ): # if the key was renamed as it is not available in the state dict otherwise, it means that we are deserializing it, diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 8cf6f591e02d..d142aaa7783c 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4284,6 +4284,8 @@ def from_pretrained( # Create the dtype_plan to potentially use the `keep_in_fp32` flags (this needs to be called on the already # instantiated model, as the flags can be modified by instances sometimes) dtype_plan = model._get_dtype_plan(dtype) + if hf_quantizer is not None: + dtype_plan.update(hf_quantizer.dtype_plan) # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index b960bb496423..72a5c7513636 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -84,10 +84,14 @@ class HfQuantizer(ABC): quantization_param_suffixes (`tuple[str, ...]`): Suffixes of the parameters quantization adds to a weight (`weight` -> `weight_scale_inv`). A save runs each through its weight's reversed conversion, so it lands in the weight's layout. + dtype_plan (`dict`): + Entries added to the model's dtype plan: a parameter pattern and the dtype it loads in, `None` + keeping the checkpoint's (e.g. quantization scales, which a checkpoint may ship in any float dtype). """ requires_calibration = False quantization_param_suffixes: tuple[str, ...] = () + dtype_plan: dict = {} def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): self.quantization_config = quantization_config @@ -137,13 +141,6 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** """ return False - def param_keeps_checkpoint_dtype(self, param_name: str) -> bool: - """ - Whether a pre-quantized checkpoint's `param_name` loads in the dtype it ships in rather than the model's - (e.g. quantization scales, which a checkpoint may ship narrower or wider than the model dtype). - """ - return False - def validate_environment(self, *args, **kwargs): """ This method is used to potentially check for potential conflicts with arguments that are diff --git a/src/transformers/quantizers/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py index 4598efe41e29..1935c8653fa9 100644 --- a/src/transformers/quantizers/finegrained/base.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -62,11 +62,9 @@ class FineGrainedHfQuantizer(HfQuantizer): quantization_config: "FineGrainedConfig" default_activation_format: str | None = None quantization_param_suffixes = ("_scale_inv",) - - def param_keeps_checkpoint_dtype(self, param_name: str) -> bool: - # a scale loads as the checkpoint ships it (Qwen3 and Mistral ship BF16, DeepSeek-V3 fp32): the - # kernels read either, and a save then writes it back unchanged - return param_name.endswith(("_scale_inv", "activation_scale", "_global_scale")) + # a scale loads as the checkpoint ships it (Qwen3 and Mistral ship BF16, DeepSeek-V3 fp32): the kernels + # read either, and a save then writes it back unchanged + dtype_plan = dict.fromkeys(("_scale_inv", "activation_scale", "_global_scale")) @property def supports_dequantize(self) -> bool: diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index e3ba9ad9a0c0..58e8dc92aaef 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -95,7 +95,9 @@ def _load_and_save(checkpoint, quantization_config): with mock.patch.object(fg, "is_sm100", return_value=False): quantizer._process_model_before_weight_loading(model) load_config = LoadStateDictConfig( - weight_mapping=quantizer.update_weight_conversions(_model_mapping()), hf_quantizer=quantizer + weight_mapping=quantizer.update_weight_conversions(_model_mapping()), + hf_quantizer=quantizer, + dtype_plan=quantizer.dtype_plan, ) info, _ = convert_and_load_state_dict_in_model( model, {k: v.clone() for k, v in checkpoint.items()}, load_config From 2d1f0bc7c7326f7017fa5bde516900a7a39cd425 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:04:30 +0200 Subject: [PATCH 60/68] Update src/transformers/distributed/sharding_utils.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/distributed/sharding_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 9674d21fdc46..9d94f03794bc 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -87,7 +87,7 @@ def __init__(self, param: DTensor): self.param_ndim = param.ndim local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) # Axis-0 range owned by this rank (used to filter per-expert pieces) - # [_axis0_offset, _axis0_offset + _axis0_local_size); a 0-dim parameter has no axis 0 + # [_axis0_offset, _axis0_offset + _axis0_local_size); a 0-dim parameter (a quantization scale) has no axis 0 self._axis0_offset = offsets[0] if offsets else 0 self._axis0_local_size = local_shape[0] if local_shape else 1 From f37d1d465ce77da962d14df8cd39c991f9ee97ec Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:04:55 +0200 Subject: [PATCH 61/68] Update src/transformers/distributed/sharding_utils.py Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/distributed/sharding_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 9d94f03794bc..826388a178b9 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -112,8 +112,7 @@ def shard_tensor( # Dense path if tensor_idx is None: - # nothing to slice: a 0-dim tensor has no axis, and no placement names one - # (`source[...]`, since a lazy safetensors slice rejects `source[()]`) + # on no dim, return early parameter is replicated if not source_shape or not dim_placements: return source[...].to(device=device, dtype=dtype) From 5ea2a3ea3d132c15960f55d6036c2ba8e05dca0d Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 30 Sep 2026 05:44:57 -0700 Subject: [PATCH 62/68] adapters --- src/transformers/integrations/deepgemm.py | 8 +-- .../integrations/finegrained/core.py | 70 +++++++++++-------- 2 files changed, 45 insertions(+), 33 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index ac03823b3779..0f2b17813ccb 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -31,7 +31,7 @@ import os import re import shutil -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass import torch @@ -541,6 +541,7 @@ def prefers_deepgemm_linear( input_global_scale: torch.Tensor | None, activation_format: str | None, allow_deepgemm: bool, + adapters: Sequence[tuple[torch.Tensor, torch.Tensor, float]] = (), ) -> bool: """Whether a dense fine-grained linear should go to DeepGEMM rather than Triton. @@ -554,16 +555,15 @@ def prefers_deepgemm_linear( # False when the model spans devices: DeepGEMM's kernels are bound to one CUDA context and # corrupt across them, and the Triton fallback is context-free. and allow_deepgemm - # DeepGEMM has no backward pass; a call that needs a gradient takes Triton, which has one and not (torch.is_grad_enabled() and input.requires_grad) # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. and weight_scale_inv.ndim <= 2 - and activation_scale is None - # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format and weight_global_scale is None and input_global_scale is None and activation_format is None + and activation_scale is None + and not adapters and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" # last: an environment probe (arch, `CUDA_HOME/bin/nvcc`, nvcc version) rather than a diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 48e3405a7f46..1c894f1f264e 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -19,7 +19,7 @@ from __future__ import annotations import functools -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass import torch @@ -277,6 +277,7 @@ def finegrained_linear( weight_global_scale: torch.Tensor | None = None, input_global_scale: torch.Tensor | None = None, activation_format: str | None = None, + adapters: Sequence[tuple[torch.Tensor, torch.Tensor, float]] = (), ) -> torch.Tensor: """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. @@ -308,6 +309,8 @@ def finegrained_linear( ``input_scale`` (None = quantize against the block scales alone). activation_format: the activation format where the weights leave it open (``"bf16"`` = weight-only). + adapters: low-rank adapters `(A (r, K), B (N, r), scaling)` on the frozen weight, whose term the kernel + adds in its epilogue. """ if prefers_deepgemm_linear( input, @@ -319,6 +322,7 @@ def finegrained_linear( input_global_scale=input_global_scale, activation_format=activation_format, allow_deepgemm=allow_deepgemm, + adapters=adapters, ): try: return deepgemm_fp8_fp4_linear( @@ -348,6 +352,7 @@ def finegrained_linear( output_dtype=input.dtype, a_global_scale=input_global_scale, b_global_scale=weight_global_scale, + adapters=adapters, ) output = output.reshape(*original_shape[:-1], output.shape[-1]) if bias is not None: @@ -423,18 +428,22 @@ def __init__( _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) def forward(self, input: torch.Tensor) -> torch.Tensor: - return finegrained_linear( - input, - self.weight, - self.weight_scale_inv, - block_size=self.block_size, - activation_scale=self.activation_scale, - bias=self.bias, - allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=self.weight_global_scale, - input_global_scale=self.input_global_scale, - activation_format=self.activation_format, - ) + return finegrained_linear(input, **self._linear_operands(input)) + + def _linear_operands(self, input: torch.Tensor) -> dict: + """Everything `finegrained_linear` takes from the module for `input`. A method so an adapter library can + extend it per module: PEFT adds its LoRA factors as `adapters`, whose term the kernel adds in its epilogue.""" + return { + "weight": self.weight, + "weight_scale_inv": self.weight_scale_inv, + "block_size": self.block_size, + "activation_scale": self.activation_scale, + "bias": self.bias, + "allow_deepgemm": not self._deepgemm_disabled, + "weight_global_scale": self.weight_global_scale, + "input_global_scale": self.input_global_scale, + "activation_format": self.activation_format, + } class FineGrainedEmbedding(nn.Embedding): @@ -671,10 +680,15 @@ def forward( def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torch.Tensor: """One expert's ``proj`` as a dense linear, over that expert's slice of the projection's operands.""" + return finegrained_linear(input, **self._linear_operands(input, proj, expert_idx)) + + def _linear_operands(self, input: torch.Tensor, proj: str, expert_idx: int) -> dict: + """Everything `finegrained_linear` takes for one expert's ``proj`` on `input`: that expert's slice of the + projection's operands. A method so an adapter library can extend it, as `_moe_operands` for the fused + forwards: PEFT adds its LoRA factors as `adapters`.""" weight, scale, weight_global, input_global, activation_scale, bias = ( getattr(self, f"{proj}{slot}") for slot in self._projection_slots ) - weight = weight[expert_idx] # one expert's slice of a swizzled stack is the `(1, ...)` artifact `matmul_2d` reads scale = scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx] # None where the checkpoint fills no such slot @@ -684,19 +698,17 @@ def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torc # gate_up holds ONE input global, its rows being quantized pre-routing; down one per expert if input_global is not None and proj == "down_proj": input_global = input_global[expert_idx : expert_idx + 1] - - return finegrained_linear( - input, - weight, - scale, - self.block_size, - bias=bias, - activation_scale=activation_scale, - allow_deepgemm=not self._deepgemm_disabled, - weight_global_scale=weight_global, - input_global_scale=input_global, - activation_format=self.activation_format, - ) + return { + "weight": weight[expert_idx], + "weight_scale_inv": scale, + "block_size": self.block_size, + "bias": bias, + "activation_scale": activation_scale, + "allow_deepgemm": not self._deepgemm_disabled, + "weight_global_scale": weight_global, + "input_global_scale": input_global, + "activation_format": self.activation_format, + } def _moe_operands(self, kernel) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their @@ -706,8 +718,8 @@ def _moe_operands(self, kernel) -> dict: kernel change. A method rather than a function so an adapter library can extend it per module: PEFT overrides - it to add its LoRA factors, which the kernels then apply on the routed rows between the GEMMs - (`expert_linear` is the eager loop's counterpart).""" + it to add its LoRA factors, which the kernels then add in each GEMM's epilogue (`expert_linear` + is the eager loop's counterpart).""" act_name = self.act_fn_name fused = act_name in kernel.get_supported_act_fns() and (self.swiglu_alpha is None or act_name == "silu") if fused: From c766206b32decc6125d631eaed6a6895e4d20e0f Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 2 Oct 2026 06:31:32 -0700 Subject: [PATCH 63/68] Fine-grained layout ops: follow the experts module each tensor fills A mixed checkpoint keeps some layers' experts unquantized; the gate|up interleave and the input-scale split read the model's first FineGrainedExperts and applied its layout to every layer. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../integrations/finegrained/conversions.py | 25 +++++++++++--- .../finegrained/test_conversions.py | 34 ++++++++++++++++--- 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 66933f33e6fa..787dab8bb961 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -64,6 +64,21 @@ def held_scale(model, full_layer_name, key): return module, getattr(module, param_name, None) +def finegrained_experts(model, full_layer_name): + """The ``FineGrainedExperts`` a converter output fills, else ``None`` (a mixed checkpoint keeps some + layers' experts unquantized). Saving names the key past the module (``experts.*.gate_proj.weight``), + so the deepest module the name resolves to is the one.""" + name = full_layer_name or "" + while model is not None and name: + name = name.rpartition(".")[0] + try: + module = model.get_submodule(name) + except AttributeError: + continue + return module if isinstance(module, FineGrainedExperts) else None + return None + + def as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: """The same bytes under a same-width dtype (uint8 <-> e8m0), an exact numeric cast otherwise.""" return scale.view(dtype) if scale.element_size() == dtype.itemsize else scale.to(dtype) @@ -92,9 +107,9 @@ class FineGrainedInterleaveGateUp(_FineGrainedOp): core ``Interleave`` along dim 1, over the weight, scale grid and bias alike. Runs on load and save wherever the checkpoint's row order differs from the one the experts hold.""" - def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): input_dict = keyed_by_target(input_dict, target_patterns) - experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) if model else None + experts = finegrained_experts(model, full_layer_name) if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): return input_dict interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved @@ -299,11 +314,11 @@ class FineGrainedInputScalesSplit(_FineGrainedOp): are gone and each projection is written the value that covers it. A collapsed global re-expands per expert, the shape the checkpoint holds.""" - def convert(self, input_dict, model=None, target_patterns=None, **kwargs): + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): value = next(iter(input_dict.values())) value = value[0] if isinstance(value, list) else value - if value.numel() == 1 and model is not None: - experts = next((m for m in model.modules() if isinstance(m, FineGrainedExperts)), None) + if value.numel() == 1: + experts = finegrained_experts(model, full_layer_name) if experts is not None: value = value.reshape(1).expand(experts.num_experts).contiguous() # one tensor per key: a save refuses keys that share storage diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index 58e8dc92aaef..a5f3502f229c 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -559,15 +559,41 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] op = FineGrainedInterleaveGateUp(hf_quantizer=None) model, experts = self._experts("mxfp8") - out = op.convert({"mlp.experts.gate_up_proj": stacked}, model=model)["mlp.experts.gate_up_proj"] + key = "experts.gate_up_proj" + out = op.convert({key: stacked}, model=model, full_layer_name=key)[key] torch.testing.assert_close(out, stacked[:, [0, 3, 1, 4, 2, 5]]) - back = op.reverse_op.convert({"mlp.experts.gate_up_proj": out}, model=model)["mlp.experts.gate_up_proj"] + back = op.reverse_op.convert({key: out}, model=model, full_layer_name=key)[key] torch.testing.assert_close(back, stacked) experts.is_concatenated = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) experts.is_concatenated = True experts.holds_interleaved_gate_up = False - self.assertIs(op.convert({"x": stacked}, model=model)["x"], stacked) + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) + + def test_interleave_op_follows_the_layer_it_fills(self): + """A mixed checkpoint keeps some layers' experts in full precision: their rows stay stacked + while a quantized layer's interleave, on load (model key) and on save (the checkpoint's + per-expert key).""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] + interleaved = stacked[:, [0, 3, 1, 4, 2, 5]] + full_precision = torch.nn.Module() + full_precision.gate_up_proj = torch.nn.Parameter(torch.empty(2, 6, 4)) + full_precision.down_proj = torch.nn.Parameter(torch.empty(2, 4, 3)) + model = torch.nn.Module() + model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()]) + model.layers[0].experts = full_precision + model.layers[1].experts = self._experts("mxfp8")[1] + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + + for layer, expected in ((0, stacked), (1, interleaved)): + key = f"layers.{layer}.experts.gate_up_proj" + out = op.convert({key: stacked}, model=model, full_layer_name=key)[key] + torch.testing.assert_close(out, expected) + saved_key = f"layers.{layer}.experts.*.gate_proj.weight" + back = op.reverse_op.convert({saved_key: out}, model=model, full_layer_name=saved_key)[saved_key] + torch.testing.assert_close(back, stacked) @require_torch From 7d44b81e9bda3ddab11e346c513fb4e96db391c1 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 2 Oct 2026 06:31:51 -0700 Subject: [PATCH 64/68] Revert "adapters" This reverts commit 5ea2a3ea3d132c15960f55d6036c2ba8e05dca0d. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/integrations/deepgemm.py | 8 +-- .../integrations/finegrained/core.py | 70 ++++++++----------- 2 files changed, 33 insertions(+), 45 deletions(-) diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 0f2b17813ccb..ac03823b3779 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -31,7 +31,7 @@ import os import re import shutil -from collections.abc import Callable, Sequence +from collections.abc import Callable from dataclasses import dataclass import torch @@ -541,7 +541,6 @@ def prefers_deepgemm_linear( input_global_scale: torch.Tensor | None, activation_format: str | None, allow_deepgemm: bool, - adapters: Sequence[tuple[torch.Tensor, torch.Tensor, float]] = (), ) -> bool: """Whether a dense fine-grained linear should go to DeepGEMM rather than Triton. @@ -555,15 +554,16 @@ def prefers_deepgemm_linear( # False when the model spans devices: DeepGEMM's kernels are bound to one CUDA context and # corrupt across them, and the Triton fallback is context-free. and allow_deepgemm + # DeepGEMM has no backward pass; a call that needs a gradient takes Triton, which has one and not (torch.is_grad_enabled() and input.requires_grad) # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. and weight_scale_inv.ndim <= 2 + and activation_scale is None + # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format and weight_global_scale is None and input_global_scale is None and activation_format is None - and activation_scale is None - and not adapters and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" # last: an environment probe (arch, `CUDA_HOME/bin/nvcc`, nvcc version) rather than a diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 1c894f1f264e..48e3405a7f46 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -19,7 +19,7 @@ from __future__ import annotations import functools -from collections.abc import Callable, Sequence +from collections.abc import Callable from dataclasses import dataclass import torch @@ -277,7 +277,6 @@ def finegrained_linear( weight_global_scale: torch.Tensor | None = None, input_global_scale: torch.Tensor | None = None, activation_format: str | None = None, - adapters: Sequence[tuple[torch.Tensor, torch.Tensor, float]] = (), ) -> torch.Tensor: """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. @@ -309,8 +308,6 @@ def finegrained_linear( ``input_scale`` (None = quantize against the block scales alone). activation_format: the activation format where the weights leave it open (``"bf16"`` = weight-only). - adapters: low-rank adapters `(A (r, K), B (N, r), scaling)` on the frozen weight, whose term the kernel - adds in its epilogue. """ if prefers_deepgemm_linear( input, @@ -322,7 +319,6 @@ def finegrained_linear( input_global_scale=input_global_scale, activation_format=activation_format, allow_deepgemm=allow_deepgemm, - adapters=adapters, ): try: return deepgemm_fp8_fp4_linear( @@ -352,7 +348,6 @@ def finegrained_linear( output_dtype=input.dtype, a_global_scale=input_global_scale, b_global_scale=weight_global_scale, - adapters=adapters, ) output = output.reshape(*original_shape[:-1], output.shape[-1]) if bias is not None: @@ -428,22 +423,18 @@ def __init__( _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) def forward(self, input: torch.Tensor) -> torch.Tensor: - return finegrained_linear(input, **self._linear_operands(input)) - - def _linear_operands(self, input: torch.Tensor) -> dict: - """Everything `finegrained_linear` takes from the module for `input`. A method so an adapter library can - extend it per module: PEFT adds its LoRA factors as `adapters`, whose term the kernel adds in its epilogue.""" - return { - "weight": self.weight, - "weight_scale_inv": self.weight_scale_inv, - "block_size": self.block_size, - "activation_scale": self.activation_scale, - "bias": self.bias, - "allow_deepgemm": not self._deepgemm_disabled, - "weight_global_scale": self.weight_global_scale, - "input_global_scale": self.input_global_scale, - "activation_format": self.activation_format, - } + return finegrained_linear( + input, + self.weight, + self.weight_scale_inv, + block_size=self.block_size, + activation_scale=self.activation_scale, + bias=self.bias, + allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=self.weight_global_scale, + input_global_scale=self.input_global_scale, + activation_format=self.activation_format, + ) class FineGrainedEmbedding(nn.Embedding): @@ -680,15 +671,10 @@ def forward( def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torch.Tensor: """One expert's ``proj`` as a dense linear, over that expert's slice of the projection's operands.""" - return finegrained_linear(input, **self._linear_operands(input, proj, expert_idx)) - - def _linear_operands(self, input: torch.Tensor, proj: str, expert_idx: int) -> dict: - """Everything `finegrained_linear` takes for one expert's ``proj`` on `input`: that expert's slice of the - projection's operands. A method so an adapter library can extend it, as `_moe_operands` for the fused - forwards: PEFT adds its LoRA factors as `adapters`.""" weight, scale, weight_global, input_global, activation_scale, bias = ( getattr(self, f"{proj}{slot}") for slot in self._projection_slots ) + weight = weight[expert_idx] # one expert's slice of a swizzled stack is the `(1, ...)` artifact `matmul_2d` reads scale = scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx] # None where the checkpoint fills no such slot @@ -698,17 +684,19 @@ def _linear_operands(self, input: torch.Tensor, proj: str, expert_idx: int) -> d # gate_up holds ONE input global, its rows being quantized pre-routing; down one per expert if input_global is not None and proj == "down_proj": input_global = input_global[expert_idx : expert_idx + 1] - return { - "weight": weight[expert_idx], - "weight_scale_inv": scale, - "block_size": self.block_size, - "bias": bias, - "activation_scale": activation_scale, - "allow_deepgemm": not self._deepgemm_disabled, - "weight_global_scale": weight_global, - "input_global_scale": input_global, - "activation_format": self.activation_format, - } + + return finegrained_linear( + input, + weight, + scale, + self.block_size, + bias=bias, + activation_scale=activation_scale, + allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=weight_global, + input_global_scale=input_global, + activation_format=self.activation_format, + ) def _moe_operands(self, kernel) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their @@ -718,8 +706,8 @@ def _moe_operands(self, kernel) -> dict: kernel change. A method rather than a function so an adapter library can extend it per module: PEFT overrides - it to add its LoRA factors, which the kernels then add in each GEMM's epilogue (`expert_linear` - is the eager loop's counterpart).""" + it to add its LoRA factors, which the kernels then apply on the routed rows between the GEMMs + (`expert_linear` is the eager loop's counterpart).""" act_name = self.act_fn_name fused = act_name in kernel.get_supported_act_fns() and (self.swiglu_alpha is None or act_name == "silu") if fused: From eb9c09fe334afa9f55d9a2591f13590e647b67b1 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 2 Oct 2026 06:43:02 -0700 Subject: [PATCH 65/68] trimming --- src/transformers/integrations/finegrained/conversions.py | 5 ++--- tests/integrations/finegrained/test_conversions.py | 4 +--- tests/integrations/finegrained/test_core.py | 9 +++------ 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 787dab8bb961..84e694890e63 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -65,9 +65,8 @@ def held_scale(model, full_layer_name, key): def finegrained_experts(model, full_layer_name): - """The ``FineGrainedExperts`` a converter output fills, else ``None`` (a mixed checkpoint keeps some - layers' experts unquantized). Saving names the key past the module (``experts.*.gate_proj.weight``), - so the deepest module the name resolves to is the one.""" + """The ``FineGrainedExperts`` the tensor at ``full_layer_name`` fills, else ``None``. On save the name + runs past the module (``experts.*.gate_proj.weight``).""" name = full_layer_name or "" while model is not None and name: name = name.rpartition(".")[0] diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index a5f3502f229c..5594921bb33d 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -571,9 +571,7 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) def test_interleave_op_follows_the_layer_it_fills(self): - """A mixed checkpoint keeps some layers' experts in full precision: their rows stay stacked - while a quantized layer's interleave, on load (model key) and on save (the checkpoint's - per-expert key).""" + """Unquantized experts in a mixed checkpoint keep their stacked rows, on load and on save.""" from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index 664c16eff425..cb50160e4c21 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -110,12 +110,9 @@ def test_replacement_follows_the_patterns_and_the_skip_list(self): self.assertIs(type(model.embed_tokens), torch.nn.Embedding) self.assertIs(type(model.other), torch.nn.Embedding) - # A scale's dtype is the format's, never the ambient default. - # - # `from_pretrained` sets the default dtype to the checkpoint's for the duration of model - # construction, so a scale allocated as `torch.ones(n)` comes out bf16 and the kernels read it - # as fp32 — NaN logits, no error. Building under both defaults and comparing is what catches - # that, whatever the format decides each scale should be. + # `from_pretrained` sets the default dtype to the checkpoint's during construction, so a scale + # allocated as `torch.ones(n)` comes out bf16 and the kernels read it as fp32 (NaN logits, no + # error). Building under both defaults catches it. def _dtypes(self, build, default): previous = torch.get_default_dtype() torch.set_default_dtype(default) From 77f135cd15e32d922b4975c247ed2c334dc72cce Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Mon, 5 Oct 2026 06:07:44 -0700 Subject: [PATCH 66/68] Fine-grained NVFP4: exact per-expert gate_up activation globals The gate_up holds one activation global per expert, as calibrated, and the kernels quantize each routed row against its own expert's. TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL (`global` or `rank`) shares their max instead, so the rows quantize once per token: faster at prefill, at some block-scale precision for experts below the max. The NVFP4 quantizer registers it after loading; `_projection_operands` hands the shared global to both forwards. EP plans shard the gate_up's activation globals with its other per-expert companions. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../integrations/finegrained/__init__.py | 2 + .../integrations/finegrained/conversions.py | 7 +-- .../integrations/finegrained/core.py | 57 +++++++++++++++---- .../quantizers/finegrained/base.py | 10 +++- .../quantizers/finegrained/nvfp4.py | 8 +++ .../finegrained/test_conversions.py | 37 ++++++++++-- tests/integrations/finegrained/test_core.py | 10 ++-- 7 files changed, 100 insertions(+), 31 deletions(-) diff --git a/src/transformers/integrations/finegrained/__init__.py b/src/transformers/integrations/finegrained/__init__.py index eb08ca99af47..8c93719211e3 100644 --- a/src/transformers/integrations/finegrained/__init__.py +++ b/src/transformers/integrations/finegrained/__init__.py @@ -25,6 +25,7 @@ disable_deepgemm_on_multi_device, finegrained_linear, load_finegrained_kernel, + register_shared_gate_up_input_globals, replace_with_finegrained_embedding, replace_with_finegrained_layer, ) @@ -38,6 +39,7 @@ "FineGrainedLinear", "assert_modules_are_quantized", "disable_deepgemm_on_multi_device", + "register_shared_gate_up_input_globals", "finegrained_linear", "load_finegrained_kernel", "replace_with_finegrained_embedding", diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 84e694890e63..a7a7739cf052 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -287,9 +287,7 @@ def _assert_no_gate_up_bias(model) -> None: class FineGrainedInputScales(_FineGrainedOp): """A calibrated checkpoint's ``input_scale`` in the layout the module holds: one value per quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max since - both halves read the same rows. The NVFP4 gate_up global collapses to ONE value — its rows are - the hidden states, quantized once before routing. A static activation scale IS the - quantization scale, with no block level to absorb an inflated one, so it stays per-expert.""" + both halves read the same rows.""" def convert(self, input_dict, full_layer_name=None, **kwargs): values = [] @@ -299,8 +297,7 @@ def convert(self, input_dict, full_layer_name=None, **kwargs): stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) # a scale is a magnitude, so the reduce is over absolute values per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1) - one_value = full_layer_name.endswith("gate_up_proj_input_global_scale") # the NVFP4 global only - return {full_layer_name: (per_expert.amax().reshape(1) if one_value else per_expert).contiguous()} + return {full_layer_name: per_expert.contiguous()} @property def reverse_op(self) -> ConversionOps: diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/core.py index 48e3405a7f46..0ccad12ff8bb 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/core.py @@ -19,6 +19,7 @@ from __future__ import annotations import functools +import os from collections.abc import Callable from dataclasses import dataclass @@ -586,19 +587,19 @@ def __init__( self.gate_up_name = "gate_up_proj" if self.has_gate else "up_proj" up_rows = (2 if self.has_gate else 1) * self.intermediate_dim storage = (format_spec, scale_dtype, scale_group, swizzled) - # gate_up takes ONE activation global — its rows are the pre-routing hidden states, - # quantized once before routing; down's rows belong to an expert each - gate_up_globals, down_globals = 1, self.num_experts self._register_projection( - self.gate_up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, gate_up_globals, storage + self.gate_up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, self.num_experts, storage ) - self._register_projection("down_proj", self.hidden_dim, self.intermediate_dim, 1, down_globals, storage) + self._register_projection("down_proj", self.hidden_dim, self.intermediate_dim, 1, self.num_experts, storage) # the model's per-expert output norm, filled by the swap self.post_expert_norm = None self.has_post_expert_norm = False self.post_expert_norm_name = None + # optimization knob, set by `register_shared_gate_up_input_globals`: one gate_up activation global + self.register_buffer("shared_gate_up_input_global", None, persistent=False) + def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, storage): """One expert projection: the packed weight and its scale grid, then the slots a checkpoint may or may not fill — a bias, the NVFP4 second-level globals, a calibrated @@ -669,11 +670,16 @@ def forward( final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) return final_hidden_states.to(hidden_states.dtype) + def _projection_operands(self, proj: str) -> dict[str, torch.Tensor | None]: + """``proj``'s tensors by slot, with the gate_up's shared activation global where one is set.""" + operands = {slot: getattr(self, f"{proj}{slot}") for slot in self._projection_slots} + if proj == self.gate_up_name and self.shared_gate_up_input_global is not None: + operands["_input_global_scale"] = self.shared_gate_up_input_global + return operands + def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torch.Tensor: """One expert's ``proj`` as a dense linear, over that expert's slice of the projection's operands.""" - weight, scale, weight_global, input_global, activation_scale, bias = ( - getattr(self, f"{proj}{slot}") for slot in self._projection_slots - ) + weight, scale, weight_global, input_global, activation_scale, bias = self._projection_operands(proj).values() weight = weight[expert_idx] # one expert's slice of a swizzled stack is the `(1, ...)` artifact `matmul_2d` reads scale = scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx] @@ -681,8 +687,8 @@ def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torc weight_global, activation_scale, bias = ( operand[expert_idx] if operand is not None else None for operand in (weight_global, activation_scale, bias) ) - # gate_up holds ONE input global, its rows being quantized pre-routing; down one per expert - if input_global is not None and proj == "down_proj": + # one per expert, or one shared by all + if input_global is not None and input_global.numel() > 1: input_global = input_global[expert_idx : expert_idx + 1] return finegrained_linear( @@ -733,9 +739,9 @@ def _moe_operands(self, kernel) -> dict: return { **{ - f"{kernel_proj}{slot}": getattr(self, f"{held_proj}{slot}") + f"{kernel_proj}{slot}": operand for kernel_proj, held_proj in (("gate_up_proj", self.gate_up_name), ("down_proj", "down_proj")) - for slot in self._projection_slots + for slot, operand in self._projection_operands(held_proj).items() }, # a supported activation NAME is fused into the gate_up epilogue; any other callable # leaves that GEMM plain and runs on the host between the two @@ -788,6 +794,33 @@ def assert_modules_are_quantized(model: nn.Module) -> None: ) +def register_shared_gate_up_input_globals(model: nn.Module) -> None: + """Share the gate_up's NVFP4 activation global across experts, as the max of their calibrated ones, so + its rows quantize once per token before routing and the GEMM gathers them. Per-expert globals, the + calibration exactly, quantize per routed row (top_k times the rows), which the GEMM then reads + materialized: 3-11% slower on a GLM NVFP4 MoE layer at prefill, the most at small batches. Decode is + unaffected. A max never clips, but experts below it lose block-scale precision. + + ``TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL`` turns it on and picks the experts the max + runs over under expert parallelism: ``global`` all of them, so the numerics do not depend on the EP + size, or ``rank`` this rank's, a tighter global. Unset keeps every expert's own.""" + scope = os.environ.get("TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL") + if scope is None: + return + if scope not in ("global", "rank"): + raise ValueError( + f"TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL takes `global` or `rank`, got {scope!r}" + ) + for module in model.modules(): + if not isinstance(module, FineGrainedExperts): + continue + input_global = getattr(module, f"{module.gate_up_name}_input_global_scale") + if input_global is not None: + if hasattr(input_global, "full_tensor"): # sharded over experts + input_global = input_global.full_tensor() if scope == "global" else input_global.to_local() + module.shared_gate_up_input_global = input_global.detach().amax().reshape(1) + + def disable_deepgemm_on_multi_device(model: nn.Module) -> None: """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` diff --git a/src/transformers/quantizers/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py index 1935c8653fa9..6e5377f1ba1e 100644 --- a/src/transformers/quantizers/finegrained/base.py +++ b/src/transformers/quantizers/finegrained/base.py @@ -254,9 +254,13 @@ def update_tp_plan(self, config): if projection not in ("gate_up_proj", "up_proj", "down_proj"): continue if style == "grouped_gemm": - companions = ["_scale_inv", "_bias", "_weight_global_scale", "_activation_scale"] - if projection == "down_proj": - companions.append("_input_global_scale") + companions = [ + "_bias", + "_scale_inv", + "_weight_global_scale", + "_input_global_scale", + "_activation_scale", + ] for suffix in companions: updated_plan.setdefault(f"{key}{suffix}", style) continue diff --git a/src/transformers/quantizers/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py index e46872ed5fe5..c5205914224a 100644 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ b/src/transformers/quantizers/finegrained/nvfp4.py @@ -28,6 +28,7 @@ class FineGrainedNvfp4HfQuantizer(FineGrainedHfQuantizer): * the packed E2M1 weights, which need a uint8 -> int8 BITCAST (`copy_` would convert and corrupt anything >= 128) and a `.weight` source that cannot swallow `weight_scale` / `weight_scale_2`. + * the activation globals, whose gate_up ones may be shared across experts after loading. """ @property @@ -42,6 +43,13 @@ def _process_model_before_weight_loading(self, model, **kwargs): } super()._process_model_before_weight_loading(model, **kwargs) + def _process_model_after_weight_loading(self, model, **kwargs): + from ...integrations.finegrained import register_shared_gate_up_input_globals + + model = super()._process_model_after_weight_loading(model, **kwargs) + register_shared_gate_up_input_globals(model) + return model + def get_weight_conversions(self): return self._nvfp4_conversions() if self.pre_quantized else [] diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index 5594921bb33d..3043380dff10 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -3,6 +3,7 @@ through the real quantizer onto a dummy model, checks what the modules hold, and saves it back — with each scheme's quirks. Pure torch: no kernel is loaded and no GPU is needed.""" +import os import re import unittest from types import SimpleNamespace @@ -570,6 +571,34 @@ def test_interleave_op_follows_the_model_and_the_experts_backend(self): experts.holds_interleaved_gate_up = False self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) + def test_shared_gate_up_input_global_follows_the_env_var(self): + """`TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL` shares the gate_up's activation global + as the max over experts; the down and the per-expert parameters keep the checkpoint's values.""" + from transformers.integrations.finegrained.core import register_shared_gate_up_input_globals + + for env, expected in ( + ({}, None), + ({"TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL": "global"}, 4.0), + ): + model, experts = self._experts("nvfp4") + with torch.no_grad(): + experts.gate_up_proj_input_global_scale.copy_(torch.tensor([1.0, 4.0, 2.0, 3.0])) + with mock.patch.dict(os.environ, env): + register_shared_gate_up_input_globals(model) + gate_up_global = experts._projection_operands("gate_up_proj")["_input_global_scale"] + if expected is None: + self.assertIsNone(experts.shared_gate_up_input_global) + self.assertIs(gate_up_global, experts.gate_up_proj_input_global_scale) + else: + torch.testing.assert_close(gate_up_global, torch.tensor([expected])) + down_global = experts._projection_operands("down_proj")["_input_global_scale"] + self.assertIs(down_global, experts.down_proj_input_global_scale) + self.assertEqual(experts.gate_up_proj_input_global_scale.numel(), 4) + model, _ = self._experts("nvfp4") + with mock.patch.dict(os.environ, {"TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL": "node"}): + with self.assertRaises(ValueError): + register_shared_gate_up_input_globals(model) + def test_interleave_op_follows_the_layer_it_fills(self): """Unquantized experts in a mixed checkpoint keep their stacked rows, on load and on save.""" from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp @@ -716,10 +745,8 @@ def test_the_fused_layout_ships_both_halves_in_one_tensor(self): torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) torch.testing.assert_close(out[targets[1]], torch.tensor([5.0 * 3.0, 6.0 * 2.0])) - def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_down(self): - """The gate_up quantizes the hidden states once, BEFORE routing, so its calibrated - `input_scale` reduces to one value; the down's rows are per expert, so its stays per - expert — the requant epilogue normalizes each row by its own expert's value.""" + def test_activation_global_is_per_expert_for_both_projections(self): + """Each expert keeps its calibrated `input_scale`, the gate|up pair reducing to their max.""" from transformers.integrations.finegrained.conversions import FineGrainedInputScales up = FineGrainedInputScales().convert( @@ -729,7 +756,7 @@ def test_activation_global_is_one_value_for_the_gate_up_and_per_expert_for_the_d }, full_layer_name="mlp.experts.gate_up_proj_input_global_scale", )["mlp.experts.gate_up_proj_input_global_scale"] - torch.testing.assert_close(up, torch.tensor([3.0])) + torch.testing.assert_close(up, torch.tensor([2.0, 3.0])) down = FineGrainedInputScales().convert( {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, full_layer_name="mlp.experts.down_proj_input_global_scale", diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py index cb50160e4c21..d86c1c4a3bca 100644 --- a/tests/integrations/finegrained/test_core.py +++ b/tests/integrations/finegrained/test_core.py @@ -247,12 +247,10 @@ def test_every_companion_shards_with_its_expert_weight(self): def style(name): return _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) - for name in ("gate_up_proj", "gate_up_proj_scale_inv", "gate_up_proj_weight_global_scale"): - self.assertEqual(style(name), "grouped_gemm", name) - for name in ("down_proj_scale_inv", "down_proj_bias", "down_proj_input_global_scale"): - self.assertEqual(style(name), "grouped_gemm", name) - # one per-tensor value for the pre-routing hidden states: replicated, so no entry of its own - self.assertEqual(style("gate_up_proj_input_global_scale"), "moe_tp_experts") + for proj in ("gate_up_proj", "down_proj"): + for name in (proj, f"{proj}_scale_inv", f"{proj}_weight_global_scale", f"{proj}_input_global_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + self.assertEqual(style("down_proj_bias"), "grouped_gemm") self.assertEqual(plan["layers.*.self_attn.q_proj"], "colwise") def test_intra_expert_entries_follow_the_projection_axis(self): From 4e8aa7e7d1a64bbd4385b4880702d284930be785 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 15 Oct 2026 07:13:26 -0700 Subject: [PATCH 67/68] Fine-grained tests: per-expert gate_up activation globals, and the operands helper on the mock Co-Authored-By: Claude Opus 5.5 --- tests/kernels/test_finegrained.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 987c776ff9ae..1f6ec20f7478 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -305,12 +305,12 @@ def test_nvfp4_experts_thread_per_expert_globals(self): # is merged to that at load, so nothing downstream carries the pair self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) - # the calibrated activation globals (the checkpoint's `input_scale`): one value for the - # gate_up, whose rows are quantized once before routing, and one per expert for the down, - # whose rows the gate_up epilogue requantizes per expert + # the calibrated activation globals (the checkpoint's `input_scale`), one per expert on both + # projections: the gate_up's exact ones unless `register_shared_gate_up_input_globals` set one + # shared value, and the down's, whose rows the gate_up epilogue requantizes per expert self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_input_global_scale) self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_input_global_scale) - self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (1,)) + self.assertEqual(tuple(m.gate_up_proj_input_global_scale.shape), (m.num_experts,)) self.assertEqual(tuple(m.down_proj_input_global_scale.shape), (m.num_experts,)) def test_weight_only_experts_hold_no_activation_global(self): @@ -557,6 +557,9 @@ def _operands(self, *, input_reduce: bool): module.post_expert_norm = norm module.post_expert_norm_name = "input_scaled_rms_norm" module._projection_slots = fg.FineGrainedExperts._projection_slots + module.gate_up_name = "gate_up_proj" + module.shared_gate_up_input_global = None + module._projection_operands = lambda proj: fg.FineGrainedExperts._projection_operands(module, proj) kernel = mock.Mock() kernel.get_supported_act_fns.return_value = ("silu",) kernel.get_supported_norms.return_value = ("input_scaled_rms_norm",) From 585946e8dbfc114a5f45e62b41e35a468e9e85b9 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 15 Oct 2026 11:50:07 -0700 Subject: [PATCH 68/68] test muse sparknvfp4 after merge --- src/transformers/models/glm4v_moe/modular_glm4v_moe.py | 2 -- tests/integrations/finegrained/test_models.py | 10 ++++++++++ tests/integrations/test_moe.py | 1 - 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index ef4bcba4f730..9dfc57662939 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -351,8 +351,6 @@ class Glm4vMoeModelOutputWithPast(Qwen3VLMoeModelOutputWithPast): class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): - _tp_plan = {"lm_head": "colwise_gather_output"} - def __init__(self, config): super().__init__(config) self.num_experts = config.text_config.num_local_experts diff --git a/tests/integrations/finegrained/test_models.py b/tests/integrations/finegrained/test_models.py index e41585fe7c6e..e09856af3257 100644 --- a/tests/integrations/finegrained/test_models.py +++ b/tests/integrations/finegrained/test_models.py @@ -62,6 +62,7 @@ def _checkpoint_bytes(path): "IlyasMoutawwakil/tiny-random-GptOss-MXFP4": 0.05, "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": 0.25, "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": 0.6, + "IlyasMoutawwakil/tiny-random-Muse-Spark-NVFP4": 0.3, "IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": 0.3, } @@ -150,6 +151,15 @@ def _checkpoint_bytes(path): ], } ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Muse-Spark-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.0869, 0.0566, -0.0145, 0.0112, 0.0103, 0.0012, -0.0679, -0.0064], + [0.0283, -0.0557, 0.0192, 0.0322, -0.0295, -0.0117, 0.0173, 0.0267], + [-0.0304, -0.0079, 0.0510, 0.0081, 0.0280, 0.0067, 0.0635, 0.0201], + ], + } + ), # fmt: skip "IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": Expectations( { ("cuda", (10, 0)): [ diff --git a/tests/integrations/test_moe.py b/tests/integrations/test_moe.py index 470d0f95e673..efef759ae22f 100644 --- a/tests/integrations/test_moe.py +++ b/tests/integrations/test_moe.py @@ -50,7 +50,6 @@ def apply_gate(gate_up): down_proj=torch.randn(num_experts, hidden, inter, device=torch_device), down_proj_bias=None, _is_expert_parallel=is_expert_parallel, - has_post_expert_norm=False, )