From 64708a2472b0bac19b9a300ad959c9e67cd062b7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 13 Aug 2026 12:54:31 -0700 Subject: [PATCH 01/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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/91] 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 8dc70c4d3b8208436d3ea92ccb32d12e7c39e90b Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 7 Oct 2026 01:37:27 -0700 Subject: [PATCH 67/91] Fine-grained quantization: one quantizer per format, and FineGrainedConfig to mix them Four standalone format quantizers, each with its own config and usable alone, pre-quantized or on the fly: BlockFp8 (`quant_method="fp8"`), Mxfp8, Mxfp4 and NVFP4 (the last two rewritten in place). Their shared machinery is `BlockScaledQuantizer` (`quantizers/block_scaled_utils.py`) and `BlockScaledConfig`. `FineGrainedConfig(dtype_plan={pattern: format config})` and `FineGrainedHfQuantizer` host several of them in one model. - The deprecated `FineGrainedFP8Config` translates to the config it spells (BlockFp8, Mxfp8, or a FineGrainedConfig); DeepSeek-V4's `expert_dtype="fp4"` reads as block-FP8 over MXFP4 experts. - modelopt and compressed-tensors checkpoints in served formats read as these configs: one format as its quantizer, several as a FineGrainedConfig (`read_modelopt`, `read_compressed_tensors`). - The modules live in `integrations/finegrained/layers.py` (was `core.py`); conversions are per-parameter ops from each format quantizer. - A quantizer quantizes what its plan names and nothing else; the automatic dequantize fallback stays. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/source/en/experts_interface.md | 8 +- docs/source/en/main_classes/quantization.md | 24 +- docs/source/en/model_doc/deepseek_v32.md | 4 +- docs/source/en/model_doc/deepseek_v4.md | 35 ++ docs/source/en/model_doc/minimax_m3_vl.md | 4 +- .../source/en/quantization/finegrained_fp8.md | 8 +- docs/source/en/quantization/mxfp4.md | 7 - docs/source/en/quantization/nvfp4.md | 7 - src/transformers/__init__.py | 9 +- src/transformers/conversion_mapping.py | 2 + src/transformers/core_model_loading.py | 91 ++- src/transformers/integrations/__init__.py | 4 - .../integrations/finegrained/__init__.py | 8 +- .../integrations/finegrained/conversions.py | 131 ++-- .../finegrained/{core.py => layers.py} | 339 +++++----- .../integrations/finegrained_fp8.py | 5 +- src/transformers/integrations/moe.py | 7 +- src/transformers/integrations/mxfp4.py | 5 +- src/transformers/integrations/nvfp4.py | 5 +- .../convert_ministral3_weights_to_hf.py | 9 +- .../mistral4/convert_mistral4_weight_to_hf.py | 6 +- src/transformers/quantizers/auto.py | 106 +++- src/transformers/quantizers/base.py | 9 +- .../quantizers/block_scaled_utils.py | 475 ++++++++++++++ .../quantizers/finegrained/__init__.py | 31 - .../quantizers/finegrained/base.py | 475 -------------- .../quantizers/finegrained/blockfp8.py | 76 --- .../quantizers/finegrained/mxfp4.py | 82 --- .../quantizers/finegrained/nvfp4.py | 229 ------- .../quantizers/quantizer_blockfp8.py | 53 ++ .../quantizers/quantizer_finegrained.py | 63 ++ .../quantizers/quantizer_finegrained_fp8.py | 5 +- .../quantizers/quantizer_mxfp4.py | 333 ++-------- .../mxfp8.py => quantizer_mxfp8.py} | 19 +- .../quantizers/quantizer_nvfp4.py | 231 +++---- .../quantizers/quantizers_utils.py | 66 ++ src/transformers/utils/quantization_config.py | 579 ++++++++++-------- 37 files changed, 1615 insertions(+), 1935 deletions(-) rename src/transformers/integrations/finegrained/{core.py => layers.py} (80%) create mode 100644 src/transformers/quantizers/block_scaled_utils.py delete mode 100644 src/transformers/quantizers/finegrained/__init__.py delete mode 100644 src/transformers/quantizers/finegrained/base.py delete mode 100644 src/transformers/quantizers/finegrained/blockfp8.py delete mode 100644 src/transformers/quantizers/finegrained/mxfp4.py delete mode 100644 src/transformers/quantizers/finegrained/nvfp4.py create mode 100644 src/transformers/quantizers/quantizer_blockfp8.py create mode 100644 src/transformers/quantizers/quantizer_finegrained.py rename src/transformers/quantizers/{finegrained/mxfp8.py => quantizer_mxfp8.py} (55%) diff --git a/docs/source/en/experts_interface.md b/docs/source/en/experts_interface.md index ab673c3d802a..cb40eb96b6d7 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 [`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. +DeepSeek-style checkpoints are usually pre-quantized and carry their own quantization config, so you don't need to pass a [`BlockFp8Config`]. 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 [`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. +The main reason to pass a [`BlockFp8Config`] 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, FineGrainedConfig +from transformers import AutoModelForCausalLM, BlockFp8Config model = AutoModelForCausalLM.from_pretrained( "deepseek-ai/DeepSeek-V3", - quantization_config=FineGrainedConfig(dequantize=True), + quantization_config=BlockFp8Config(dequantize=True), experts_implementation="deepgemm", ) ``` diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index d9da976a992c..ae5c655e2b23 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -98,21 +98,29 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] SpQRConfig -## FineGrainedConfig +## BlockScaledConfig -> [!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]] BlockScaledConfig -[[autodoc]] FineGrainedConfig +## BlockFp8Config + +[[autodoc]] BlockFp8Config + +## Mxfp8Config -## FineGrainedGroup +[[autodoc]] Mxfp8Config -[[autodoc]] FineGrainedGroup +## FineGrainedConfig + +[[autodoc]] FineGrainedConfig ## FineGrainedFP8Config +> [!WARNING] +> `FineGrainedFP8Config` is frozen and will be removed in a future release. A config passed to `from_pretrained` +> loads as the config it spells: [`Mxfp8Config`] for a `[1, 32]` block, [`BlockFp8Config`] otherwise, and a +> [`FineGrainedConfig`] of block-FP8 over MXFP4 experts for a model that declares `expert_dtype="fp4"` (DeepSeek-V4). + [[autodoc]] FineGrainedFP8Config ## QuarkConfig diff --git a/docs/source/en/model_doc/deepseek_v32.md b/docs/source/en/model_doc/deepseek_v32.md index c7cfdaf1c84a..6722253256c9 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 FineGrainedConfig, AutoModelForCausalLM, AutoTokenizer +from transformers import BlockFp8Config, AutoModelForCausalLM, AutoTokenizer import torch model_name = "deepseek-ai/DeepSeek-V3.2-Exp" -quantization_config = FineGrainedConfig( +quantization_config = BlockFp8Config( 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/deepseek_v4.md b/docs/source/en/model_doc/deepseek_v4.md index a87c733721ac..efc9f6e10019 100644 --- a/docs/source/en/model_doc/deepseek_v4.md +++ b/docs/source/en/model_doc/deepseek_v4.md @@ -29,6 +29,41 @@ share the same architecture; they differ only in width / depth / expert count an > [!TIP] > Set `use_kernels=True` in [`~PreTrainedModel.from_pretrained`] to replace supported layers with optimized kernels from the Hub. Refer to [Loading kernels](../kernel_doc/loading_kernels) to learn more. +## Usage examples + +DeepSeek-V4 checkpoints store block-FP8 weights with MXFP4 routed experts (`expert_dtype="fp4"` in the model config). +They load as-is, with no quantization config: + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_name = "deepseek-ai/DeepSeek-V4-Flash" +model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto", device_map="auto") +tokenizer = AutoTokenizer.from_pretrained(model_name) + +inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) +outputs = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(outputs[0], skip_special_tokens=True)) +``` + +The checkpoint mixes two formats, so it loads through [`FineGrainedConfig`]. The config below is the one built +automatically; pass your own to change a format, for example `Mxfp4Config(activation_format="bf16")` to run the experts +weight-only, or `dequantize=True` to load everything in bf16: + +```python +from transformers import AutoModelForCausalLM, BlockFp8Config, FineGrainedConfig, Mxfp4Config + +quantization_config = FineGrainedConfig( + dtype_plan={ + r"\.experts($|\.)": Mxfp4Config(activation_format="mxfp4"), + ".*": BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), + } +) +model = AutoModelForCausalLM.from_pretrained( + "deepseek-ai/DeepSeek-V4-Flash", dtype="auto", device_map="auto", quantization_config=quantization_config +) +``` + ## Architecture (paper §2) ### Hybrid attention (§2.3) diff --git a/docs/source/en/model_doc/minimax_m3_vl.md b/docs/source/en/model_doc/minimax_m3_vl.md index a90c61ab4f26..322b7c1269e7 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, - FineGrainedConfig, + Mxfp8Config, ) 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=FineGrainedConfig(dequantize=True), + quantization_config=Mxfp8Config(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 afb55a2a7414..e438d984c5c1 100644 --- a/docs/source/en/quantization/finegrained_fp8.md +++ b/docs/source/en/quantization/finegrained_fp8.md @@ -17,11 +17,9 @@ 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. +> `FineGrainedFP8Config` is frozen and will be removed in a future release; use [`BlockFp8Config`] or +> [`Mxfp8Config`]. A `FineGrainedFP8Config` passed to `from_pretrained` loads as the config it spells (MXFP8 for a +> `[1, 32]` block, block-FP8 otherwise, with DeepSeek-V4's MXFP4 experts), and existing checkpoints load unchanged. Fine-grained FP8 quantization quantizes the weights and activations to fp8. diff --git a/docs/source/en/quantization/mxfp4.md b/docs/source/en/quantization/mxfp4.md index 09ac3c73a56d..f96833c0cc09 100644 --- a/docs/source/en/quantization/mxfp4.md +++ b/docs/source/en/quantization/mxfp4.md @@ -16,13 +16,6 @@ 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 1ea263317c15..42a2d31cfa5a 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -16,13 +16,6 @@ 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/__init__.py b/src/transformers/__init__.py index 8189220c43a9..3c5724f6f340 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -256,12 +256,13 @@ "AwqConfig", "BitNetQuantConfig", "BitsAndBytesConfig", + "BlockFp8Config", + "BlockScaledConfig", "CompressedTensorsConfig", "EetqConfig", "FbgemmFp8Config", "FineGrainedConfig", "FineGrainedFP8Config", - "FineGrainedGroup", "FourOverSixConfig", "FPQuantConfig", "GemmaQuantizationConfig", @@ -271,6 +272,7 @@ "HqqConfig", "MetalConfig", "Mxfp4Config", + "Mxfp8Config", "NVFP4Config", "QuantoConfig", "QuarkConfig", @@ -789,12 +791,13 @@ from .utils.quantization_config import AwqConfig as AwqConfig from .utils.quantization_config import BitNetQuantConfig as BitNetQuantConfig from .utils.quantization_config import BitsAndBytesConfig as BitsAndBytesConfig + from .utils.quantization_config import BlockFp8Config as BlockFp8Config + from .utils.quantization_config import BlockScaledConfig as BlockScaledConfig 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 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 @@ -803,6 +806,8 @@ from .utils.quantization_config import HiggsConfig as HiggsConfig from .utils.quantization_config import HqqConfig as HqqConfig from .utils.quantization_config import MetalConfig as MetalConfig + from .utils.quantization_config import Mxfp4Config as Mxfp4Config + from .utils.quantization_config import Mxfp8Config as Mxfp8Config from .utils.quantization_config import NVFP4Config as NVFP4Config from .utils.quantization_config import QuantoConfig as QuantoConfig from .utils.quantization_config import QuarkConfig as QuarkConfig diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 49e2c0818998..ddc523d2c153 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -503,6 +503,8 @@ def _build_checkpoint_conversion_mapping(): ), ], "deepseek_v4": [ + # the per-block weight scales ship as `.scale`; first, while `hc_*_scale` keys still carry no dot + WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv"), # 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 diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 24d5d960f018..57c6b554e9fb 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1144,13 +1144,8 @@ def reverse_transform(self) -> WeightTransform: 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. - """ + """The classes known to be able to use m:n. A function so the finegrained integration, which imports this module, + is imported at call time.""" from .integrations.finegrained.conversions import FineGrainedWeightGlobals return ( @@ -1648,13 +1643,12 @@ def convert_and_load_state_dict_in_model( renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)] converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)] + pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} param_name_to_load: dict[str, WeightRenaming | WeightConverter] = {} if dtype_plan != {}: dtype_policy_alt, dtype_policy_by_group_name, _ = build_glob_alternation(list(dtype_plan.keys())) - pattern_to_converter = {k: converter for converter in converters for k in converter.source_patterns} - state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) for original_key, tensor in state_dict: # 1. Rename the key according to all renaming and weight conversion patterns. @@ -1670,14 +1664,19 @@ def convert_and_load_state_dict_in_model( # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: empty_param = meta_model_state_dict.get(renamed_key) - # If we enter here, we have a WeightConverter operation to perform - if source_pattern is not None: - new_converter = deepcopy(pattern_to_converter[source_pattern]) - # each target key gets its own converter instance - mapping = param_name_to_load.setdefault(renamed_key, new_converter) - # Otherwise, only potential renaming - else: - mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) + # each target key gets its own transform, with the ops the quantizer gives the module it fills + if renamed_key not in param_name_to_load: + param_ops = hf_quantizer.get_param_ops(model, renamed_key) if hf_quantizer is not None else [] + if source_pattern is not None: + mapping = deepcopy(pattern_to_converter[source_pattern]) + mapping.operations = [*mapping.operations, *param_ops] + elif param_ops: + mapping = WeightConverter(original_key, renamed_key, operations=param_ops) + else: + mapping = WeightRenaming(original_key, renamed_key) + param_name_to_load[renamed_key] = mapping + mapping = param_name_to_load[renamed_key] + if source_pattern is None: source_pattern = original_key # 3. Handle dtype casting @@ -1813,9 +1812,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 _reversed_param_ops(hf_quantizer, model, param_name: str) -> list[ConversionOps]: + """The quantizer's ops for the module `param_name` fills, undone first on save since they ran last on load.""" + if hf_quantizer is None: + return [] + return [op.reverse_op for op in hf_quantizer.get_param_ops(model, param_name)[::-1]] def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]): @@ -1850,52 +1851,44 @@ 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} - # 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]] = {} + # `(layer name, reversed transform)` per realized group + conversion_mapping: dict[str, tuple[str, WeightTransform]] = {} 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: # 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) + converted_key, _ = rename_source_key(original_key, [], converters, reverse=True) # each target key gets its own converter instance - _, 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) + if converted_key not in conversion_mapping: + converter = deepcopy(pattern_to_converter[source_pattern]) + converter.operations = [*_reversed_param_ops(hf_quantizer, model, original_key), *converter.operations] + conversion_mapping[converted_key] = (converted_key, converter) + _, mapping = conversion_mapping[converted_key] + mapping.add_tensor(converted_key, original_key, source_pattern, tensor) else: - _, mapping, _ = conversion_mapping.setdefault( - renamed_key, (renamed_key, WeightRenaming(original_key, renamed_key), "") - ) - mapping.add_tensor(renamed_key, original_key, original_key, tensor) + if renamed_key not in conversion_mapping: + # with ops, a converter of them alone on the model's own name, which the renamings then apply to + reverse_ops = _reversed_param_ops(hf_quantizer, model, original_key) + conversion_mapping[renamed_key] = ( + (original_key, WeightConverter(original_key, original_key, operations=reverse_ops)) + if reverse_ops + else (renamed_key, WeightRenaming(original_key, renamed_key)) + ) + layer_name, mapping = conversion_mapping[renamed_key] + mapping.add_tensor(layer_name, original_key, original_key, tensor) new_state_dict = {} - for layer_name, reversed_transform, suffix in conversion_mapping.values(): + for layer_name, reversed_transform in conversion_mapping.values(): # Apply the reverse converter 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) + target_name, _ = rename_source_key(target_name, renamings, [], reverse=True) new_state_dict[target_name] = param return new_state_dict diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 5e21b782694f..bad766e16787 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -54,8 +54,6 @@ "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"], @@ -221,8 +219,6 @@ 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 diff --git a/src/transformers/integrations/finegrained/__init__.py b/src/transformers/integrations/finegrained/__init__.py index 8c93719211e3..ffb0093164d2 100644 --- a/src/transformers/integrations/finegrained/__init__.py +++ b/src/transformers/integrations/finegrained/__init__.py @@ -15,8 +15,9 @@ `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 ( +from .layers import ( ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + SUPPORTED_FUSED_NORMS, FineGrainedEmbedding, FineGrainedExperts, FineGrainedGroupedLinear, @@ -26,13 +27,12 @@ finegrained_linear, load_finegrained_kernel, register_shared_gate_up_input_globals, - replace_with_finegrained_embedding, - replace_with_finegrained_layer, ) __all__ = [ "ALL_FINEGRAINED_EXPERTS_FUNCTIONS", + "SUPPORTED_FUSED_NORMS", "FineGrainedEmbedding", "FineGrainedExperts", "FineGrainedGroupedLinear", @@ -42,6 +42,4 @@ "register_shared_gate_up_input_globals", "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 index a7a7739cf052..9270291906b3 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.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. -"""Loader/saver conversion ops for the fine-grained family. +"""Loader/saver conversion ops for the block-scaled formats. 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 @@ -26,7 +26,7 @@ from ...core_model_loading import ConversionOps, Interleave from ...quantizers.quantizers_utils import get_module_from_name -from .core import ( +from .layers import ( _FP8_DTYPE, _FP8_MAX, _FP8_MIN, @@ -64,8 +64,8 @@ def held_scale(model, full_layer_name, key): return module, getattr(module, param_name, None) -def finegrained_experts(model, full_layer_name): - """The ``FineGrainedExperts`` the tensor at ``full_layer_name`` fills, else ``None``. On save the name +def finegrained_module(model, full_layer_name): + """The finegrained module 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: @@ -74,10 +74,16 @@ def finegrained_experts(model, full_layer_name): module = model.get_submodule(name) except AttributeError: continue - return module if isinstance(module, FineGrainedExperts) else None + return module if isinstance(module, _FineGrainedModule) else None return None +def finegrained_experts(model, full_layer_name): + """The ``FineGrainedExperts`` the tensor at ``full_layer_name`` fills, else ``None``.""" + module = finegrained_module(model, full_layer_name) + return module if isinstance(module, FineGrainedExperts) else 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) @@ -89,16 +95,18 @@ class _FineGrainedOp(ConversionOps): 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. + class — merge/split, quantize/dequantize — override `reverse_op` and ignore the flag. A layout op + built for one module carries it, since a save may name the tensor as the checkpoint does. """ - def __init__(self, hf_quantizer=None, inverse: bool = False): + def __init__(self, hf_quantizer=None, inverse: bool = False, module=None): self.hf_quantizer = hf_quantizer self.inverse = inverse + self.module = module @property def reverse_op(self) -> ConversionOps: - return type(self)(self.hf_quantizer, inverse=not self.inverse) + return type(self)(self.hf_quantizer, inverse=not self.inverse, module=self.module) class FineGrainedInterleaveGateUp(_FineGrainedOp): @@ -135,10 +143,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: - # 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) + # the container its module recorded at load; e8m0 is only ever a scale, so the dtype + # alone picks the tensors to restore + module = self.module or finegrained_module(model, full_layer_name) + container = getattr(module, "scale_container_dtype", None) if container is not None and value.dtype == _get_ue8m0_dtype(): value = as_container(value, container) else: @@ -211,6 +219,17 @@ def convert(self, input_dict, target_patterns=None, **kwargs): return out +class FineGrainedReciprocal(_FineGrainedOp): + """compressed-tensors stores NVFP4's globals as ``448 * 6 / amax``, the reciprocal of the multiplier the kernels + (and modelopt) take. Its own inverse, so a save writes them back as shipped.""" + + def convert(self, input_dict, target_patterns=None, **kwargs): + return { + key: [1 / tensor for tensor in value] if isinstance(value, list) else 1 / value + for key, value in keyed_by_target(input_dict, target_patterns).items() + } + + 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 @@ -234,7 +253,7 @@ def role(key: str) -> tuple[str, str]: `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" + level = "input" if "input_scale" in key or "activation_scale" in key else "weight" return ("down" if down else "gate_up"), level def convert(self, input_dict, target_patterns=None, model=None, **kwargs): @@ -350,18 +369,12 @@ def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> @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.""" + """An expert stack in the ``(E, rows, K)`` the module's slot and the kernels take. A model may store its + experts transposed (GPT-OSS: ``(E, H, 2I)``), which the module it replaced declares: quantizing before the + swap packs the wrong axis, and a square stack (GPT-OSS's down projection) cannot tell from its shape.""" 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 + return value.transpose(1, 2).contiguous() if getattr(module, "ships_transposed", False) else value @staticmethod def _weight_holder(model, key: str): @@ -443,12 +456,12 @@ def _identity_activation_scales(module, prefix: str, scale_name: str | None, dev # `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) + slot = f"{proj}activation_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 @@ -528,6 +541,19 @@ def _scale_pattern_for(self, weight_pattern: str) -> str: scale = base + "_scale_inv" return scale + "$" if anchored else scale + @property + def _reciprocal_globals(self) -> bool: + """compressed-tensors ships NVFP4's globals as their reciprocal (see `FineGrainedReciprocal`).""" + plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} + return any(config.checkpoint_format == "compressed-tensors" for config in plan.values()) + + @staticmethod + def _global_pattern_for(weight_pattern: str) -> str: + """NVFP4's second-level global beside `weight_pattern`: `x.weight` -> `x.weight_global_scale`.""" + anchored = weight_pattern.endswith("$") + base = weight_pattern[:-1] if anchored else weight_pattern + return base + "_global_scale" + ("$" if anchored else "") + # 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) @@ -541,11 +567,16 @@ def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: 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 + self, + quantized: torch.Tensor, + scales: torch.Tensor, + output_dtype: torch.dtype | None = None, + global_scale: torch.Tensor | 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): + # modelopt ships its packed bytes as uint8 + if quantized.dtype in (torch.int8, torch.uint8) or (fp4_dtype is not None and quantized.dtype == fp4_dtype): quantized_fp32 = self._unpack_fp4(quantized) else: quantized_fp32 = quantized.to(torch.float32) @@ -558,9 +589,17 @@ def _dequantize_one( # 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 # 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): + # `(5, 56)` grid against a 128x128 block); only guess when no block the plan names reproduces it + plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} + configured = next( + ( + block + for block in {tuple(config.weight_block_size) for config in plan.values() if config.weight_block_size} + if (_cdiv(rows, block[0]), _cdiv(cols, block[1])) == (scale_rows, scale_cols) + ), + None, + ) + if configured: block_m, block_n = configured elif rows % scale_rows or cols % scale_cols: raise ValueError( @@ -586,8 +625,11 @@ def _dequantize_one( 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) - out = (blocked * per_block).reshape(*original_shape[:-2], rows + pad_m, cols + pad_n) - return out[..., :rows, :cols].to(output_dtype) + out = (blocked * per_block).reshape(*original_shape[:-2], rows + pad_m, cols + pad_n)[..., :rows, :cols] + if global_scale is not None: + global_scale = global_scale.to(torch.float32) + out = out / global_scale if self._reciprocal_globals else out * global_scale + return out.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: @@ -604,7 +646,7 @@ def convert( **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``) + # The dense converter (``["weight$", "weight_scale_inv", "weight_global_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` @@ -615,11 +657,16 @@ def convert( 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)} + global_scale = input_dict.get("weight_global_scale") + global_scale = global_scale[0] if isinstance(global_scale, list) else global_scale + dequantized = self._dequantize_one(quantized, scales, output_dtype, global_scale) + return {target_key: dequantized} 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} + consumed = { + pattern(key) for key in input_dict for pattern in (self._scale_pattern_for, self._global_pattern_for) + } result: dict[str, list[torch.Tensor] | torch.Tensor] = {} for key, value in input_dict.items(): if "activation_scale" in key or key in consumed: @@ -636,7 +683,11 @@ def convert( 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)] + global_scales = input_dict.get(self._global_pattern_for(key), [None] * len(weights)) + global_scales = global_scales if isinstance(global_scales, list) else [global_scales] + result[key] = [ + self._dequantize_one(w, s, output_dtype, g) for w, s, g in zip(weights, scales, global_scales) + ] return result @property diff --git a/src/transformers/integrations/finegrained/core.py b/src/transformers/integrations/finegrained/layers.py similarity index 80% rename from src/transformers/integrations/finegrained/core.py rename to src/transformers/integrations/finegrained/layers.py index 0ccad12ff8bb..0cf01d79722d 100644 --- a/src/transformers/integrations/finegrained/core.py +++ b/src/transformers/integrations/finegrained/layers.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. -"""Quantized modules and forwards for the fine-grained family — block-FP8, MXFP8, MXFP4, NVFP4 +"""Quantized modules and forwards for the block-scaled formats — 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`. """ @@ -27,9 +27,9 @@ 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 ...utils.quantization_config import BlockFp8Config, BlockScaledConfig from ..deepgemm import ( deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, @@ -356,6 +356,11 @@ def finegrained_linear( return output +# the kernels' fused form of a model's post-expert norm, by the norm's exact class: a subclass may compute +# something else, so a model whose own norm one of them implements registers that class here +SUPPORTED_FUSED_NORMS: dict[type[nn.Module], str] = {nn.RMSNorm: "rms_norm"} + + class _FineGrainedModule: """What every finegrained module (dense linear or experts) carries besides its parameters.""" @@ -364,34 +369,54 @@ class _FineGrainedModule: # read back by its reverse so a save restores the checkpoint's container scale_container_dtype: torch.dtype | None = None + # 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 + _hf_quantized_needs_local_tp = True + # 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 + quantization_config: BlockScaledConfig + + @property + def weight_format(self) -> str: + return self.quantization_config.method.value + + @property + def activation_format(self) -> str | None: + return self.quantization_config.activation_format + + @property + def activation_scheme(self) -> str: + return self.quantization_config.activation_scheme + + @property + def block_size(self) -> tuple[int, int] | None: + return self.quantization_config.weight_block_size + class FineGrainedLinear(_FineGrainedModule, nn.Linear): def __init__( self, in_features: int, out_features: int, - block_size: tuple[int, int] | None = None, - activation_scheme: str = "dynamic", - scale_fmt: str = "float", + *, + quantization_config: BlockScaledConfig | None = None, has_bias: bool = False, - weight_format: str = "fp8", - activation_format: str | None = None, ): super().__init__(in_features, out_features) + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) 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 + activation_format = self.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 - format_spec, sf_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + format_spec, sf_dtype, scale_group = resolve_weight_format( + self.weight_format, self.quantization_config.scale_fmt, self.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 @@ -415,25 +440,38 @@ 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) + # the checkpoint's calibrated activation scale: a static FP8 scale, or NVFP4's activation global + calibrated = self.activation_scheme == "static" or ( + global_scale_dtype is not None and activation_format != "bf16" + ) _set_optional_parameter( - self, "input_global_scale", torch.tensor(1.0, dtype=torch.float32) if calibrated else None + self, "activation_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) + @classmethod + def from_module( + cls, module: nn.Linear, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedLinear: + return cls( + module.in_features, + module.out_features, + quantization_config=quantization_config, + has_bias=module.bias is not None, + ) + def forward(self, input: torch.Tensor) -> torch.Tensor: + two_level = self.weight_global_scale is not None return finegrained_linear( input, self.weight, self.weight_scale_inv, block_size=self.block_size, - activation_scale=self.activation_scale, + activation_scale=None if two_level else 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, + input_global_scale=self.activation_scale if two_level else None, activation_format=self.activation_format, ) @@ -444,12 +482,30 @@ class FineGrainedEmbedding(nn.Embedding): ~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): + _hf_quantized_needs_local_tp = True + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int | None = None, + *, + quantization_config: BlockScaledConfig | None = None, + ): super().__init__(num_embeddings, embedding_dim, padding_idx=padding_idx) + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) 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) + @classmethod + def from_module( + cls, module: nn.Embedding, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedEmbedding: + return cls( + module.num_embeddings, module.embedding_dim, module.padding_idx, quantization_config=quantization_config + ) + def forward(self, input: torch.Tensor) -> torch.Tensor: rows = super().forward(input) return rows.to(self.weight_scale.dtype) * self.weight_scale @@ -470,25 +526,27 @@ def __init__( 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", + *, + quantization_config: BlockScaledConfig | None = None, has_bias: bool = False, - weight_format: str = "fp8", - activation_format: str | None = None, ): 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, - weight_format=weight_format, - activation_format=activation_format, + in_features_per_group, out_features, quantization_config=quantization_config, has_bias=has_bias ) self.n_groups = n_groups + @classmethod + def from_module( + cls, module: nn.Linear, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedGroupedLinear: + return cls( + in_features_per_group=module.in_features, + out_features=module.out_features, + n_groups=module.n_groups, + quantization_config=quantization_config, + has_bias=module.bias is not None, + ) + def forward(self, x: torch.Tensor) -> torch.Tensor: input_shape = x.shape[:-2] hidden_dim = x.shape[-1] @@ -537,70 +595,86 @@ 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") + # a projection's tensors, by the suffix they carry on the module: the weight, its scale grid, the NVFP4 + # weight global, the calibrated activation scale, a bias + _projection_slots = ("", "_scale_inv", "_weight_global_scale", "_activation_scale", "_bias") def __init__( self, config, - block_size: tuple[int, int] | None = None, - activation_scheme: str = "dynamic", - scale_fmt: str = "float", - has_bias: bool = False, + *, + quantization_config: BlockScaledConfig | None = None, has_gate: bool = True, - weight_format: str = "fp8", - activation_format: str | None = None, + has_bias: bool = False, is_concatenated: bool = True, ): super().__init__() self.config = config - self.has_bias = has_bias - self.has_gate = has_gate - self.block_size = block_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.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.act_fn_name = _first_attr(config, "hidden_activation", "hidden_act") 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] - self.weight_format = weight_format - format_spec, scale_dtype, scale_group = resolve_weight_format(weight_format, scale_fmt, block_size) + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) + activation_format = self.activation_format + format_spec, scale_dtype, scale_group = resolve_weight_format( + self.weight_format, self.quantization_config.scale_fmt, self.block_size + ) self.global_scale_dtype = format_spec.global_scale_dtype - # The layouts: what the CHECKPOINT ships (`is_concatenated`) against what the module now + # The layouts: what the CHECKPOINT ships (`is_concatenated`: stacked ``[gate; up]`` or, like + # GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``) 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.has_gate = has_gate + self.has_bias = has_bias 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) - 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) - self._register_projection( - 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, self.num_experts, storage) + self._register_projection(self.gate_up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, storage) + self._register_projection("down_proj", self.hidden_dim, self.intermediate_dim, 1, 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 + self.ships_transposed = False # 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): + @classmethod + def from_module( + cls, module: nn.Module, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedExperts: + """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. The model's per-expert + output norm is the model's own module, so it is taken from the experts this replaces. + """ + 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=cls, experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, **flags + ) + config = getattr(module, "config", model.config.get_text_config()) + experts = experts_class(config=config, quantization_config=quantization_config, **flags) + experts.post_expert_norm, experts.has_post_expert_norm = module.post_expert_norm, module.has_post_expert_norm + # a full-precision stack quantized on load arrives in the layout of the module this replaces + experts.ships_transposed = getattr(module, "is_transposed", False) + return experts + + def _register_projection(self, proj, rows, in_dim, min_scale_out, 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.""" @@ -616,16 +690,14 @@ def _register_projection(self, proj, rows, in_dim, min_scale_out, input_globals, 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`. + # down to that). The calibrated activation scale is the checkpoint's `input_scale`: a static FP8 scale, + # or NVFP4's activation global. two_level = self.global_scale_dtype is not None - calibrated = two_level and self.activation_format != "bf16" - static = self.activation_scheme == "static" + calibrated = self.activation_scheme == "static" or (two_level and self.activation_format != "bf16") 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 + activation = torch.ones(self.num_experts, dtype=torch.float32) if calibrated 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: @@ -671,15 +743,24 @@ def forward( 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.""" + """``proj``'s tensors by the kernels' slot names: the calibrated activation scale goes in as NVFP4's + activation global (the gate_up's shared one where it is set) or as a static scale.""" operands = {slot: getattr(self, f"{proj}{slot}") for slot in self._projection_slots} + activation = operands.pop("_activation_scale") 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 + activation = self.shared_gate_up_input_global + two_level = self.global_scale_dtype is not None + operands["_input_global_scale"] = activation if two_level else None + operands["_activation_scale"] = None if two_level else activation 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 = self._projection_operands(proj).values() + operands = self._projection_operands(proj) + weight, scale, weight_global, input_global, activation_scale, bias = ( + operands[slot] + for slot in ("", "_scale_inv", "_weight_global_scale", "_input_global_scale", "_activation_scale", "_bias") + ) 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] @@ -726,12 +807,13 @@ def _moe_operands(self, kernel) -> dict: # 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() + norm_name = SUPPORTED_FUSED_NORMS.get(type(self.post_expert_norm)) + fusable = 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 + post_expert_norm, norm_weight = 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) @@ -814,7 +896,7 @@ def register_shared_gate_up_input_globals(model: nn.Module) -> None: for module in model.modules(): if not isinstance(module, FineGrainedExperts): continue - input_global = getattr(module, f"{module.gate_up_name}_input_global_scale") + input_global = getattr(module, f"{module.gate_up_name}_activation_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() @@ -847,110 +929,3 @@ def disable_deepgemm_on_multi_device(model: nn.Module) -> None: "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_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 _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: - # 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 - 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 - 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 - - 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(): - if not should_convert_module(module_name, modules_to_not_convert): - continue - - new_module = None - with torch.device("meta"): - if module_name.endswith(".experts"): - new_module = _quantized_experts(module, model, storage_for(module_name)) - elif type(module) is nn.Linear: - new_module = FineGrainedLinear( - in_features=module.in_features, - out_features=module.out_features, - has_bias=module.bias is not None, - **storage_for(module_name), - ) - elif isinstance(module, nn.Linear) and hasattr(module, "n_groups"): - # 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, - n_groups=module.n_groups, - has_bias=module.bias is not None, - **storage_for(module_name), - ) - if new_module is not None: - # 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 - - if not has_been_replaced: - logger.warning( - "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 diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 5cc3bcdae74d..d783e035219b 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -40,9 +40,8 @@ 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.", + "`transformers.integrations.finegrained_fp8` is frozen and will be removed in a future release; " + "`BlockFp8Config` and `Mxfp8Config` load through `transformers.integrations.finegrained`.", DeprecationWarning, stacklevel=2, ) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 81ba67fcc31f..59a45655f323 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -550,7 +550,6 @@ def use_experts_implementation( is_transposed: bool = False, has_bias: bool = False, has_gate: bool = True, - has_post_expert_norm: bool = False, ) -> type[torch.nn.Module]: """Decorator to modify experts class to support different experts implementations. @@ -569,9 +568,6 @@ 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. - has_post_expert_norm (`bool`, *optional*, defaults to `False`): - 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. @@ -589,7 +585,8 @@ def __init__(self, config, *args, **kwargs): self.has_bias = has_bias self.is_transposed = is_transposed self.is_concatenated = is_concatenated - self.has_post_expert_norm = has_post_expert_norm + self.post_expert_norm = getattr(self, "post_expert_norm", None) + self.has_post_expert_norm = self.post_expert_norm is not None self._is_expert_parallel = False @wraps(original_forward) diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index b8a01bcae243..24be86f65e5d 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -28,9 +28,8 @@ 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.", + "`transformers.integrations.mxfp4` is frozen and will be removed in a future release; " + "`Mxfp4Config` now loads through `transformers.integrations.finegrained`.", DeprecationWarning, stacklevel=2, ) diff --git a/src/transformers/integrations/nvfp4.py b/src/transformers/integrations/nvfp4.py index cb74485ea95b..daa10c4bb43a 100644 --- a/src/transformers/integrations/nvfp4.py +++ b/src/transformers/integrations/nvfp4.py @@ -30,9 +30,8 @@ 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.", + "`transformers.integrations.nvfp4` is frozen and will be removed in a future release; " + "`NVFP4Config` now loads through `transformers.integrations.finegrained`.", DeprecationWarning, stacklevel=2, ) 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 f270a0c31c9f..59084138d78a 100644 --- a/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py +++ b/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py @@ -29,9 +29,8 @@ PixtralProcessor, PixtralVisionConfig, ) -from transformers.integrations.finegrained import replace_with_finegrained_layer from transformers.integrations.mistral import convert_tekken_tokenizer -from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig # fmt: off @@ -256,9 +255,9 @@ def convert_and_write_model(input_dir: str, output_dir: str, max_position_embedd # swap nn.Linear for the quantized module before loading if hasattr(model.config, "quantization_config"): - model = replace_with_finegrained_layer( - model, model.config.quantization_config.modules_to_not_convert, model.config.quantization_config - ) + qconfig = model.config.quantization_config + quantizer = AutoHfQuantizer.from_config(qconfig, pre_quantized=True) + quantizer.replace_modules(model, modules_to_not_convert=qconfig.modules_to_not_convert) model.load_state_dict(full_state_dict, strict=True, assign=True) model.save_pretrained(output_dir) 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 4997f4f633b4..0fcd6e7c3549 100644 --- a/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py +++ b/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py @@ -37,10 +37,9 @@ MergeModulelist, WeightRenaming, ) -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 +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig _FP8_DTYPE = torch.float8_e4m3fn @@ -536,7 +535,8 @@ def convert_and_write_model( if output_fp8 and hasattr(model.config, "quantization_config"): qconfig = model.config.quantization_config - model = replace_with_finegrained_layer(model, qconfig.modules_to_not_convert, qconfig) + quantizer = AutoHfQuantizer.from_config(qconfig, pre_quantized=True) + quantizer.replace_modules(model, modules_to_not_convert=qconfig.modules_to_not_convert) model.load_state_dict(full_state_dict, strict=True, assign=True) model.save_pretrained(str(output_dir)) diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 48047595d1aa..d6ef42c90b89 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -22,6 +22,8 @@ AwqConfig, BitNetQuantConfig, BitsAndBytesConfig, + BlockFp8Config, + BlockScaledConfig, CompressedTensorsConfig, EetqConfig, FbgemmFp8Config, @@ -36,6 +38,8 @@ HqqConfig, MetalConfig, Mxfp4Config, + Mxfp8Config, + NVFP4Config, QuantizationConfigMixin, QuantizationMethod, QuantoConfig, @@ -44,23 +48,21 @@ SpQRConfig, TorchAoConfig, VptqConfig, + read_compressed_tensors, + read_modelopt, ) 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 from .quantizer_bitnet import BitNetHfQuantizer +from .quantizer_blockfp8 import BlockFp8HfQuantizer from .quantizer_bnb_4bit import Bnb4BitHfQuantizer from .quantizer_bnb_8bit import Bnb8BitHfQuantizer from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer from .quantizer_eetq import EetqHfQuantizer from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer +from .quantizer_finegrained import FineGrainedHfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer from .quantizer_gemma import GemmaQuantizer @@ -69,12 +71,16 @@ from .quantizer_higgs import HiggsHfQuantizer from .quantizer_hqq import HqqHfQuantizer from .quantizer_metal import MetalHfQuantizer +from .quantizer_mxfp4 import Mxfp4HfQuantizer +from .quantizer_mxfp8 import Mxfp8HfQuantizer +from .quantizer_nvfp4 import NVFP4HfQuantizer from .quantizer_quanto import QuantoHfQuantizer from .quantizer_quark import QuarkHfQuantizer from .quantizer_sinq import SinqHfQuantizer from .quantizer_spqr import SpQRHfQuantizer from .quantizer_torchao import TorchAoHfQuantizer from .quantizer_vptq import VptqHfQuantizer +from .quantizers_utils import EXPERTS_PATTERN AUTO_QUANTIZER_MAPPING = { @@ -99,13 +105,12 @@ "gguf": GgufHfQuantizer, "metal": MetalHfQuantizer, "auto-round": AutoRoundQuantizer, - # 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": FineGrainedMxfp8HfQuantizer, - "mxfp4": FineGrainedMxfp4HfQuantizer, - "nvfp4": FineGrainedNvfp4HfQuantizer, - "modelopt": FineGrainedNvfp4HfQuantizer, + # fine-grained: one quantizer per weight format, and one hosting several on the parts of a model + "fp8": BlockFp8HfQuantizer, + "mxfp8": Mxfp8HfQuantizer, + "mxfp4": Mxfp4HfQuantizer, + "nvfp4": NVFP4HfQuantizer, + "finegrained": FineGrainedHfQuantizer, "sinq": SinqHfQuantizer, "gemma": GemmaQuantizer, } @@ -132,11 +137,11 @@ "spqr": SpQRConfig, "metal": MetalConfig, "auto-round": AutoRoundConfig, - # the finegrained formats (see AUTO_QUANTIZER_MAPPING) - "fp8": FineGrainedConfig, - "mxfp8": FineGrainedConfig, - "mxfp4": FineGrainedConfig, - "nvfp4": FineGrainedConfig, + "fp8": BlockFp8Config, + "mxfp8": Mxfp8Config, + "mxfp4": Mxfp4Config, + "nvfp4": NVFP4Config, + "finegrained": FineGrainedConfig, "modelopt": FineGrainedConfig, "sinq": SinqConfig, "gemma": GemmaQuantizationConfig, @@ -148,10 +153,9 @@ AutoRoundConfig, FbgemmFp8Config, CompressedTensorsConfig, - Mxfp4Config, MetalConfig, FineGrainedConfig, - FineGrainedFP8Config, + BlockScaledConfig, ) logger = logging.get_logger(__name__) @@ -175,6 +179,14 @@ def from_dict(cls, quantization_config_dict: dict): "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized" ) + # a modelopt export reads as its one format's config, or a `FineGrainedConfig` when it mixes them + if quant_method == QuantizationMethod.MODELOPT: + return read_modelopt(quantization_config_dict) + # so does a compressed-tensors checkpoint in formats these kernels serve + if quant_method == QuantizationMethod.COMPRESSED_TENSORS: + if (config := read_compressed_tensors(quantization_config_dict)) is not None: + return config + if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING: raise ValueError( f"Unknown quantization type, got {quant_method} - supported types are:" @@ -209,6 +221,7 @@ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwar # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = AutoQuantizationConfig.from_dict(quantization_config) + quantization_config = _from_legacy_fp8(quantization_config) quant_method = quantization_config.quant_method @@ -257,8 +270,8 @@ def merge_quantization_configs( if isinstance(quantization_config, dict): # Convert the config based on the type of quantization_config_from_args (e.g., AutoRoundConfig), which takes priority before automatic configuration dispatch. - if isinstance(quantization_config_from_args, AutoRoundConfig): - quantization_config = AutoRoundConfig.from_dict(quantization_config) + if isinstance(quantization_config_from_args, (AutoRoundConfig, CompressedTensorsConfig)): + quantization_config = type(quantization_config_from_args).from_dict(quantization_config) else: quantization_config = AutoQuantizationConfig.from_dict(quantization_config) @@ -282,7 +295,7 @@ def merge_quantization_configs( 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, FineGrainedConfig, FineGrainedFP8Config) + quantization_config, (MetalConfig, FineGrainedConfig, BlockScaledConfig) ): warnings.warn(warning_msg) else: @@ -343,6 +356,45 @@ def register_quantizer_fn(cls): return register_quantizer_fn +def _from_legacy_fp8(quantization_config): + """A deprecated `FineGrainedFP8Config` as the config it spells: MXFP8 for its `[1, 32]` block with UE8M0 scales, + block-FP8 otherwise (its FP4 experts come from the model config, see `finegrained_quantization_config`). Any other + config unchanged.""" + if not isinstance(quantization_config, FineGrainedFP8Config): + return quantization_config + options = { + "activation_scheme": quantization_config.activation_scheme, + "modules_to_not_convert": quantization_config.modules_to_not_convert, + "dequantize": quantization_config.dequantize, + } + block_size = quantization_config.weight_block_size + if block_size is not None and tuple(block_size) == (1, 32): + return Mxfp8Config(**options) + return BlockFp8Config( + weight_block_size=block_size, + scale_fmt=quantization_config.scale_fmt, + modules_to_convert=quantization_config.modules_to_convert, + **options, + ) + + +def finegrained_quantization_config(quantization_config, model_config): + """The config a checkpoint really describes, with its model config in hand: a deprecated `FineGrainedFP8Config` as + the config it spells, and DeepSeek-V4, which ships `quant_method="fp8"` and declares MXFP4 experts on the MODEL + config (`expert_dtype`), as block-FP8 over W4A4 MXFP4 experts. Any other config unchanged.""" + quantization_config = _from_legacy_fp8(quantization_config) + if not isinstance(quantization_config, BlockFp8Config): + return quantization_config + if getattr(model_config.get_text_config(), "expert_dtype", None) != "fp4": + return quantization_config + return FineGrainedConfig( + dtype_plan={EXPERTS_PATTERN: Mxfp4Config(activation_format="mxfp4"), ".*": quantization_config}, + modules_to_not_convert=quantization_config.modules_to_not_convert, + modules_to_convert=quantization_config.modules_to_convert, + dequantize=quantization_config.dequantize, + ) + + def get_hf_quantizer(config, quantization_config, device_map, weights_only, user_agent, gguf_file=None): if gguf_file is not None: if quantization_config is None: @@ -363,9 +415,15 @@ def get_hf_quantizer(config, quantization_config, device_map, weights_only, user pre_quantized = False if pre_quantized or quantization_config is not None: + # both read as what they describe for this model before they are compared + quantization_config = finegrained_quantization_config(quantization_config, config) if pre_quantized: + # a passed `AutoRoundConfig` or `CompressedTensorsConfig` reads the checkpoint as its own (the merge does) + routed = isinstance(quantization_config, (AutoRoundConfig, CompressedTensorsConfig)) + if isinstance(quantization_params_from_config, dict) and not routed: + quantization_params_from_config = AutoQuantizationConfig.from_dict(quantization_params_from_config) config.quantization_config = AutoHfQuantizer.merge_quantization_configs( - quantization_params_from_config, quantization_config + finegrained_quantization_config(quantization_params_from_config, config), quantization_config ) else: config.quantization_config = quantization_config diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 72a5c7513636..819a535a36ab 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -81,16 +81,12 @@ 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. 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): @@ -302,6 +298,11 @@ def get_quantize_ops(self): def get_weight_conversions(self): return [] + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: + """The ops a tensor loading into `param_name` runs after its converter's own, for the layout the module + holding it keeps; reversed, the same ops write it back on save. Default: none.""" + return [] + def update_weight_conversions(self, weight_conversions): """Give the quantizer a chance to rewrite the weight conversion pipeline. diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py new file mode 100644 index 000000000000..7c1e4def094b --- /dev/null +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -0,0 +1,475 @@ +# 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 block-scaled weight formats' quantizers share: block-FP8, MXFP8, MXFP4 and NVFP4 are each a quantizer with +its own config, served by the modules in `integrations/finegrained`, and each quantizes a whole model alone; +`FineGrainedHfQuantizer` hosts several of them, each on the modules its `dtype_plan` gives it. +""" + +import re +from typing import TYPE_CHECKING + +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 ( + embedding_pattern, + get_module_from_name, + plan_entry_for, + replace_with_quantized_modules, + should_convert_module, +) + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..utils.quantization_config import BlockScaledConfig + +logger = logging.get_logger(__name__) + + +class BlockScaledQuantizer(HfQuantizer): + """The quantizer of a block-scaled weight format (block-FP8, MXFP8, MXFP4, NVFP4), written over `plan()`: the + modules it covers and the format config each loads in. Alone, a format covers every module with its own config; + `FineGrainedHfQuantizer` covers each part of a model with the format quantizer its plan entry names.""" + + requires_calibration = False + # 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")) + # checkpoint keys this format's producers ship beside the weights that no module holds + unexpected_keys: tuple[str, ...] = () + # narrowed once the modules are built + holds_experts = True + + def plan(self) -> dict[str, "BlockScaledConfig"]: + """`{module regex: format config}` for the modules this quantizer covers.""" + return {".*": self.quantization_config} + + def members(self) -> list["BlockScaledQuantizer"]: + """The quantizer of each format config in `plan()`.""" + return [self] + + def member_of(self, module) -> "BlockScaledQuantizer | None": + """The quantizer whose plan entry `module` was built from; `None` for a module left as it was.""" + config = getattr(module, "quantization_config", None) + return next((member for member in self.members() if member.quantization_config is config), None) + + # what a format brings + + def format_conversions(self, dequantize: bool) -> list: + """The converters this format's producers' key layouts need on top of the model's plan. Under `dequantize`, + what folding this format's keys needs beyond the shared `weight` + `weight_scale_inv` chain.""" + return [] + + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: + """The ops a tensor loading into `param_name` runs, for the layout this format's modules hold it in: + an expert's stacked gate|up rows interleaved (weight, scale grid and bias share the row axis), a scale + in the container dtype its module holds, then in the swizzled layout. Each op has a reverse, so a save + restores the checkpoint layout.""" + from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear + from ..integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedScaleContainer, + FineGrainedSwizzleScales, + ) + + module, name = get_module_from_name(model, param_name) + if isinstance(module, FineGrainedExperts) and re.fullmatch( + r"(gate_up_proj|down_proj)(_scale_inv|_bias)?", name + ): + ops = [FineGrainedInterleaveGateUp(self)] if name.startswith("gate_up_proj") else [] + if not name.endswith("_bias"): + ops += [FineGrainedScaleContainer(self, module=module), FineGrainedSwizzleScales(self)] + return ops + if isinstance(module, FineGrainedLinear) and name == "weight_scale_inv": + return [FineGrainedScaleContainer(self, module=module)] + return [] + + def _process_model_after_weight_loading(self, model, **kwargs): + from ..integrations.finegrained import assert_modules_are_quantized + + assert_modules_are_quantized(model) + return model + + def _activation_scale_conversions(self, projections=("gate_up_proj", "up_proj", "down_proj"), ops=()) -> list: + """A calibrated checkpoint's `input_scale` onto the `activation_scale` its module holds: renamed for every + module first, then one value per expert of each of `projections`, the gate|up pair reduced to the one both + halves read (`up_proj` is an ungated stack's); `ops` run on each value before that reduction. Both expert layouts + are covered, a checkpoint matching one while the other never fires; neither when this format holds no + experts.""" + from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming + from ..integrations.finegrained.conversions import FineGrainedInputScales + + rename = WeightRenaming(source_patterns=r"^(.+)\.input_scale$", target_patterns=r"\1.activation_scale") + if not self.holds_experts: + return [rename] + per_expert = { + "gate_up_proj": ["mlp.experts.*.gate_proj.activation_scale", "mlp.experts.*.up_proj.activation_scale"], + "down_proj": ["mlp.experts.*.down_proj.activation_scale"], + } + return [ + rename, + *( + WeightConverter( + source_patterns=per_expert[proj], + target_patterns=f"mlp.experts.{proj}_activation_scale", + operations=[MergeModulelist(dim=0), *ops, FineGrainedInputScales(self)], + ) + for proj in projections + if proj in per_expert + ), + *( + WeightConverter( + source_patterns=rf"experts\.{proj}_input_scale$", + target_patterns=f"experts.{proj}_activation_scale", + operations=[*ops, FineGrainedInputScales(self)], + ) + for proj in projections + ), + ] + + # what every format shares + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Loading a fine-grained 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( + "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 fine-grained 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( + "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 + + device_map = kwargs.get("device_map") + if device_map is None: + logger.warning_once( + "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 ( + 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 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: + # `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) + 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"): + """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) + 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, model: "PreTrainedModel", **kwargs): + self._normalize_modules_to_not_convert(model) + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules + ) + # an embedding table worth quantizing is too big to expand, so it stays FP8 under `dequantize=True` + # too; quantizing one on the fly has no path + plan = self.plan() + embeddings = { + embedding_pattern(name) for name in getattr(self.quantization_config, "modules_to_convert", None) or [] + } + if not self.pre_quantized: + plan = {pattern: config for pattern, config in plan.items() if pattern not in embeddings} + elif self.quantization_config.dequantize: + plan = {pattern: config for pattern, config in plan.items() if pattern in embeddings} + replaced = self.replace_modules(model, plan, self.modules_to_not_convert) + from ..integrations.finegrained import FineGrainedExperts + + experts = [module for module in model.modules() if isinstance(module, FineGrainedExperts)] + ignored = set(model._keys_to_ignore_on_load_unexpected or []) + for member in self.members(): + # a format's expert converters exist only where its experts do + member.holds_experts = any(module.quantization_config is member.quantization_config for module in experts) + ignored.update(member.unexpected_keys) + # with no module holding an activation scale (weight-only, dynamic activations), a calibrated + # `input_scale` is one this run has no use for, not an unexpected key: a later run may read it + if not any(name.endswith("activation_scale") for name, _ in model.named_parameters()): + ignored.add(r"input(_global)?_scale$") + model._keys_to_ignore_on_load_unexpected = ignored + if not replaced and not self.quantization_config.dequantize: + logger.warning( + "You are loading a fine-grained quantized model but no linear or experts modules were " + "found to convert. Please double check your model architecture." + ) + + def replace_modules(self, model, plan=None, modules_to_not_convert=None) -> bool: + """Swap the modules `plan` (this quantizer's own by default) covers for their quantized classes.""" + from ..integrations.finegrained import ( + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedGroupedLinear, + FineGrainedLinear, + ) + + module_classes = { + "linear": FineGrainedLinear, + "grouped_linear": FineGrainedGroupedLinear, + "experts": FineGrainedExperts, + "embedding": FineGrainedEmbedding, + } + plan = self.plan() if plan is None else plan + return replace_with_quantized_modules(model, plan, modules_to_not_convert, module_classes) + + 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 + + impl = getattr(config, "_experts_implementation", None) + layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) + 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 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] + if projection not in ("gate_up_proj", "up_proj", "down_proj"): + continue + if style == "grouped_gemm": + companions = ["_bias", "_scale_inv", "_weight_global_scale", "_activation_scale"] + for suffix in companions: + updated_plan.setdefault(f"{key}{suffix}", style) + continue + + # ...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 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)`. + entry = plan_entry_for(self.plan(), key.rpartition(".")[0]) + blocked = entry is not None and _has_scale_grid(entry) + 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"): + # 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) + updated_plan.setdefault(f"{key}_bias", rows) + + if updated_plan != base_plan: + setattr(plan_owner, 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.conversions import FineGrainedQuantize + + return FineGrainedQuantize(self) + + def get_weight_conversions(self): + """The converters every format's producer key layout needs, on top of the model's plan.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained.conversions import FineGrainedDequantize + + if not self.pre_quantized: + return [] + dequantize = self.quantization_config.dequantize + conversions, claimed = [], {} + if dequantize: + # every quantized key folded back to a full-precision weight from the `weight` + `weight_scale_inv` pair, + # times NVFP4's `weight_global_scale` where there is one; anchored `weight$` so the scale keys land in the + # scale slots, and the activation scales collected only to be dropped + conversions.append( + WeightConverter( + source_patterns=["weight$", "weight_scale_inv", "weight_global_scale", "activation_scale"], + target_patterns="weight", + operations=[FineGrainedDequantize(self)], + ) + ) + # once each: formats reading the same producer keys emit the same converter; two that read a key + # differently cannot both load it, so the plan is refused rather than one picked + for member in self.members(): + for conversion in member.format_conversions(dequantize): + owner = next((claimed[p] for p in conversion.source_patterns if p in claimed), None) + if owner is None: + conversions.append(conversion) + claimed.update(dict.fromkeys(conversion.source_patterns, conversion)) + elif _signature(owner) != _signature(conversion): + raise ValueError( + f"The plan's formats read {conversion.source_patterns} differently, so they cannot be " + "loaded together." + ) + return conversions + + def update_weight_conversions(self, weight_conversions): + """Rewrite the model's conversion plan for this checkpoint, then append :meth:`get_weight_conversions`. + + Each of the model's converters carries a weight's block-scale grid along with it, both ways: the merge + that stacks the experts' weights stacks their grids too, and nothing else rides it (a global or an + activation scale has a converter of its own). Under ``dequantize=True`` a :class:`FineGrainedDequantize` + runs first in every converter with a ``.weight`` source, with the scales collected alongside, so each + per-expert pair folds to full precision before the merge collapses the per-expert structure. The + layout each module holds its tensors in comes per module instead, from :meth:`get_param_ops`.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained.conversions import FineGrainedDequantize + + dequantize = self.pre_quantized and self.quantization_config.dequantize + updated = [] + for conv in weight_conversions: + weights = [p for p in conv.source_patterns if p.endswith(".weight")] + if isinstance(conv, WeightConverter) and weights and dequantize: + conv = _rebuilt( + conv, + [p + "$" for p in weights] + + [p.removesuffix(".weight") + ".weight_scale_inv$" for p in weights] + + [p + "_global_scale$" for p in weights] + + [p for p in conv.source_patterns if not p.endswith(".weight")], + operations=[FineGrainedDequantize(self), *conv.operations], + ) + elif isinstance(conv, WeightConverter): + conv = _rebuilt( + conv, [_with_scale_grid(p) if p.endswith(".weight") else p for p in conv.source_patterns] + ) + updated.append(conv) + return updated + self.get_weight_conversions() + + +# a parameter and its block-scale grid, and none of its other companions (`_weight_global_scale`, `_activation_scale`) +_WITH_SCALE_GRID = r"(?=$|_scale_inv$)" + + +def _with_scale_grid(pattern: str) -> str: + return pattern if pattern.endswith((_WITH_SCALE_GRID, "$")) else pattern + _WITH_SCALE_GRID + + +def _rebuilt(conv, source_patterns: list[str], operations: list | None = None): + """`conv` over `source_patterns`, its targets taking only a parameter and its scale grid on save (a target is + a source pattern there, and a parameter's name prefixes all its companions'), and its scope kept.""" + from ..core_model_loading import WeightConverter + + params = (".weight", "gate_up_proj", "up_proj", "down_proj") + rebuilt = WeightConverter( + source_patterns=source_patterns, + target_patterns=[_with_scale_grid(p) if p.endswith(params) else p for p in conv._original_target_patterns], + operations=conv.operations if operations is None else operations, + ) + rebuilt.scope_prefix, rebuilt.base_model_prefix = conv.scope_prefix, conv.base_model_prefix + return rebuilt + + +def _has_scale_grid(config) -> bool: + """Whether a weight in `config`'s format carries a scale grid, which shards along with it under tensor parallelism, + rather than one scale for the whole tensor.""" + from ..integrations.finegrained.layers import resolve_weight_format + + return resolve_weight_format(config.method.value, config.scale_fmt, config.weight_block_size)[2] is not None + + +def _signature(conversion) -> tuple: + """What a converter does: which keys it reads, where they go and through which ops.""" + operations = tuple(type(op).__name__ for op in getattr(conversion, "operations", ())) + return type(conversion), tuple(conversion.source_patterns), tuple(conversion.target_patterns), operations diff --git a/src/transformers/quantizers/finegrained/__init__.py b/src/transformers/quantizers/finegrained/__init__.py deleted file mode 100644 index b7574e4b876a..000000000000 --- a/src/transformers/quantizers/finegrained/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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/finegrained/base.py b/src/transformers/quantizers/finegrained/base.py deleted file mode 100644 index 6e5377f1ba1e..000000000000 --- a/src/transformers/quantizers/finegrained/base.py +++ /dev/null @@ -1,475 +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. -"""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 -from dataclasses import replace -from typing import TYPE_CHECKING - -from ...utils import ( - is_accelerate_available, - is_torch_available, - is_torch_distributed_available, - is_torch_xpu_available, - logging, -) -from ...utils.quantization_config import groups_with_expert_dtype -from ..base import HfQuantizer -from ..quantizers_utils import get_module_from_name, should_convert_module - - -if is_torch_available(): - import torch - -if TYPE_CHECKING: - from ...modeling_utils import PreTrainedModel - from ...utils.quantization_config import FineGrainedConfig - -logger = logging.get_logger(__name__) - - -class FineGrainedHfQuantizer(HfQuantizer): - """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" - default_activation_format: str | None = None - quantization_param_suffixes = ("_scale_inv",) - # 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: - """Whether `dequantize=True` can fold this format's scales into a full-precision weight. - 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(): - raise ImportError("Loading a fine-grained 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( - "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 fine-grained 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): - self._assert_dequantize_supported() - logger.warning_once( - "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 - - device_map = kwargs.get("device_map") - if device_map is None: - logger.warning_once( - "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 ( - 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 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: - # `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) - 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"): - """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) - 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, - model: "PreTrainedModel", - **kwargs, - ): - 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 - 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 - # 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 - # 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 - ) - - 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 - ) - 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): - from ...integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device - - assert_modules_are_quantized(model) - disable_deepgemm_on_multi_device(model) - - return model - - 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 - - impl = getattr(config, "_experts_implementation", None) - layer_overrides = FineGrainedExperts._impl_tp_layer_overrides.get(impl, {}) - 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 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] - if projection not in ("gate_up_proj", "up_proj", "down_proj"): - continue - if style == "grouped_gemm": - companions = [ - "_bias", - "_scale_inv", - "_weight_global_scale", - "_input_global_scale", - "_activation_scale", - ] - for suffix in companions: - updated_plan.setdefault(f"{key}{suffix}", style) - continue - - # ...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 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"): - # 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) - updated_plan.setdefault(f"{key}_bias", rows) - - if updated_plan != base_plan: - setattr(plan_owner, 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.conversions import FineGrainedQuantize - - return FineGrainedQuantize(self) - - @property - def quant_method(self) -> str: - 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.""" - if not self.pre_quantized: - return [] - if self.quantization_config.dequantize: - return self._dequantize_conversions() - return [] - - 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 - - # 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 update_weight_conversions(self, weight_conversions): - """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.conversions import FineGrainedDequantize - - scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") - weight_conversions = [scale_rename, *weight_conversions] - - 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 = [] - for conv in weight_conversions: - # 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 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], - ) - updated.append(conv) - return anchor(updated + 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 - (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.conversions import ( - FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, - FineGrainedSwizzleScales, - ) - - 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 - - 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 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]) - for target in expert_targets: - fused = target.rpartition("experts.")[2].removesuffix("$") - 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 - 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"): - # 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) - 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=ops, - ) - ) - return updated diff --git a/src/transformers/quantizers/finegrained/blockfp8.py b/src/transformers/quantizers/finegrained/blockfp8.py deleted file mode 100644 index a20030237419..000000000000 --- a/src/transformers/quantizers/finegrained/blockfp8.py +++ /dev/null @@ -1,76 +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. -"""The BLOCKFP8 arm of the fine-grained quantizer.""" - -from .base import FineGrainedHfQuantizer - - -class FineGrainedBlockFp8HfQuantizer(FineGrainedHfQuantizer): - """Block-FP8 and per-tensor FP8 — the only formats that take a CALIBRATED activation scale. - - 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): - 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, 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( - source_patterns=[ - "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="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)], - ) - # `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 = [ - WeightRenaming( - source_patterns=r"^(?!.*\.experts\.)(.+)\.input_scale$", - target_patterns=r"\1.activation_scale", - ) - ] - return per_expert + fused + dense diff --git a/src/transformers/quantizers/finegrained/mxfp4.py b/src/transformers/quantizers/finegrained/mxfp4.py deleted file mode 100644 index 30935d6b9c51..000000000000 --- a/src/transformers/quantizers/finegrained/mxfp4.py +++ /dev/null @@ -1,82 +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. -"""The MXFP4 arm of the fine-grained quantizer: GPT-OSS's packed `_blocks` / `_scales` layout, -and the weight-only default that layout implies. -""" - -from .base 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. - """ - - default_activation_format = "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/finegrained/nvfp4.py b/src/transformers/quantizers/finegrained/nvfp4.py deleted file mode 100644 index c5205914224a..000000000000 --- a/src/transformers/quantizers/finegrained/nvfp4.py +++ /dev/null @@ -1,229 +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. -"""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 .base import FineGrainedHfQuantizer - - -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 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 - 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 _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 [] - - def update_weight_conversions(self, weight_conversions): - """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 - - 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`): 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) - + 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. - - `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 - - 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)], - ), - 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: 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=[ - "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: - converters.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), FineGrainedInputScales(self)], - ) - ) - return converters - - 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` 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/quantizers/quantizer_blockfp8.py b/src/transformers/quantizers/quantizer_blockfp8.py new file mode 100644 index 000000000000..5c950d880b97 --- /dev/null +++ b/src/transformers/quantizers/quantizer_blockfp8.py @@ -0,0 +1,53 @@ +# 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. +"""Block-FP8: `float8_e4m3fn` weights with fp32 or UE8M0 scales over blocks (or one per tensor), W8A8.""" + +from typing import TYPE_CHECKING + +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import embedding_pattern + + +if TYPE_CHECKING: + from ..utils.quantization_config import BlockFp8Config + + +class BlockFp8HfQuantizer(BlockScaledQuantizer): + """Block-FP8 and per-tensor FP8: the only format that takes a static activation scale, and the only one that + quantizes embedding tables (`modules_to_convert`).""" + + quantization_config: "BlockFp8Config" + + def plan(self): + embeddings = self.quantization_config.modules_to_convert or [] + return {**dict.fromkeys(map(embedding_pattern, embeddings), self.quantization_config), **super().plan()} + + def _process_model_after_weight_loading(self, model, **kwargs): + from ..integrations.finegrained import disable_deepgemm_on_multi_device + + disable_deepgemm_on_multi_device(model) + return super()._process_model_after_weight_loading(model, **kwargs) + + def format_conversions(self, dequantize): + from ..core_model_loading import WeightRenaming + + conversions = [] + if self.quantization_config.checkpoint_format == "compressed-tensors": + # its block scale ships as `weight_scale`, and rides its weight's converter once renamed + conversions.append( + WeightRenaming(source_patterns=r"^(.+)\.weight_scale$", target_patterns=r"\1.weight_scale_inv") + ) + if not dequantize and self.quantization_config.activation_scheme == "static": + conversions += self._activation_scale_conversions() + return conversions diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py new file mode 100644 index 000000000000..dd5e8285c195 --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -0,0 +1,63 @@ +# 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 fine-grained quantizer: a model whose modules are in different block-scaled formats, each part loaded by its +own format's quantizer.""" + +from typing import TYPE_CHECKING + +from ..utils.quantization_config import BlockFp8Config +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import embedding_pattern, get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..utils.quantization_config import FineGrainedConfig + + +class FineGrainedHfQuantizer(BlockScaledQuantizer): + """Hosts one quantizer per entry of the config's `dtype_plan`, each the one its format config names in + `AUTO_QUANTIZER_MAPPING`: what the formats share runs once over the whole plan, and each module's tensors take + the ops of the quantizer its entry belongs to. DeepSeek-V4 is MXFP4 experts over block-FP8 linears; a modelopt + export can mix NVFP4 and MXFP8.""" + + quantization_config: "FineGrainedConfig" + + def __init__(self, quantization_config: "FineGrainedConfig", **kwargs): + super().__init__(quantization_config, **kwargs) + from .auto import AUTO_QUANTIZER_MAPPING + + # the embedding tables `modules_to_convert` names hold FP8 rows with one per-tensor scale + self._embedding_config = BlockFp8Config(weight_block_size=None) + configs = {id(config): config for config in self.plan().values()}.values() + self._members = [AUTO_QUANTIZER_MAPPING[config.method.value](config) for config in configs] + + def plan(self): + embeddings = self.quantization_config.modules_to_convert or [] + return { + **dict.fromkeys(map(embedding_pattern, embeddings), self._embedding_config), + **self.quantization_config.dtype_plan, + } + + def members(self): + return self._members + + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: + member = self.member_of(get_module_from_name(model, param_name)[0]) + return member.get_param_ops(model, param_name) if member is not None else [] + + def _process_model_after_weight_loading(self, model, **kwargs): + for member in self.members(): + member._process_model_after_weight_loading(model, **kwargs) + return model diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index df28f050349e..7b5004804ed4 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -7,9 +7,8 @@ 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.", + "`transformers.quantizers.quantizer_finegrained_fp8` is frozen and will be removed in a future release; " + "use `BlockFp8Config` or `Mxfp8Config`, which a passed `FineGrainedFP8Config` loads as.", DeprecationWarning, stacklevel=2, ) diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 14222c668fcd..99e32040bc17 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.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. @@ -11,317 +11,66 @@ # 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 +"""MXFP4: E2M1 weights two per byte with a UE8M0 scale per 32 — GPT-OSS's experts, which ship each projection as +`{proj}_blocks` + `{proj}_scales`, and DeepSeek-V4's, which ship a plain weight/scale pair.""" + from typing import TYPE_CHECKING -from .base import HfQuantizer +from .block_scaled_utils import BlockScaledQuantizer if TYPE_CHECKING: - from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import Mxfp4Config -from ..utils import ( - is_accelerate_available, - is_kernels_available, - is_torch_available, - is_triton_available, - logging, -) -from ..utils.import_utils import KERNELS_MAX_VERSION, KERNELS_MIN_VERSION -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 - - from ..core_model_loading import WeightConverter - -logger = logging.get_logger(__name__) -triton_kernels_hub = None +class Mxfp4HfQuantizer(BlockScaledQuantizer): + """MXFP4 weights, W4A16 (the config's default, what GPT-OSS runs), W4A8 or W4A4 (DeepSeek-V4's experts). -class Mxfp4HfQuantizer(HfQuantizer): - """ - FP4 quantization using fbgemm kernels + One thing is MXFP4's alone: GPT-OSS names each projection `{proj}_blocks` + `{proj}_scales` rather than a + weight/scale pair. """ - requires_calibration = False quantization_config: "Mxfp4Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - self.triton_kernels_hub = None - - def _lazy_import_kernels(self): - """Lazy import and initialize kernels only when needed""" - if self.triton_kernels_hub is None: - try: - from ..integrations.hub_kernels import get_kernel - - self.triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels", version=1) - except ImportError: - raise ImportError("kernels package is required for MXFP4 quantization") - return self.triton_kernels_hub - - def validate_environment(self, *args, **kwargs): - if not is_torch_available(): - raise ImportError( - "Using mxfp4 quantization requires torch" - "Please install the latest version of torch ( pip install --upgrade torch )" - ) - - if self.quantization_config.dequantize: - return + def _process_model_after_weight_loading(self, model, **kwargs): + from ..integrations.finegrained import disable_deepgemm_on_multi_device - if not is_accelerate_available(): - raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`") + disable_deepgemm_on_multi_device(model) + return super()._process_model_after_weight_loading(model, **kwargs) - device = torch.accelerator.current_accelerator() or torch.device("cpu") - if device.type not in ["cuda", "xpu", "cpu"]: - if self.pre_quantized: - logger.warning_once( - f"Using MXFP4 quantized models requires model on cuda/xpu/cpu, but found {device}, we will default to dequantizing the model to bf16. To use mxfp4, please disable the current accelerator." - ) - self.quantization_config.dequantize = True - return - else: - raise RuntimeError( - f"Quantizing a model using MXFP4 requires model on cuda/xpu/cpu, but found {device}. To use mxfp4, please disable the current accelerator." - ) + def format_conversions(self, dequantize): + return self._dequantize_conversions() if dequantize else self._mxfp4_conversions() - if torch.xpu.is_available(): - is_device_supported_mxfp4 = True - triton_available = is_triton_available("3.5.0") - kernels_installed = is_kernels_available() - elif torch.cuda.is_available(): - compute_capability = torch.cuda.get_device_capability() - is_device_supported_mxfp4 = compute_capability >= (7, 5) - triton_available = is_triton_available("3.4.0") - kernels_installed = is_kernels_available() - elif device.type == "cpu": - is_device_supported_mxfp4 = True - triton_available = is_triton_available("3.5.0") - kernels_installed = is_kernels_available() - else: - is_device_supported_mxfp4 = False - triton_available = False - kernels_installed = False - - if self.pre_quantized: - if not is_device_supported_mxfp4: - logger.warning_once( - "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 " - "(e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series). " - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - - if not triton_available: - logger.warning_once( - "MXFP4 quantization requires Triton: CUDA requires Triton >= 3.4.0, " - "XPU/CPU requires Triton >= 3.5.0. Please install triton: `pip install triton`. " - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - - if not kernels_installed: - logger.warning_once( - "MXFP4 quantization requires the `kernels` package: " - f"Please install a compatible version ({KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), " - f"e.g. `pip install kernels=={KERNELS_MIN_VERSION}`" - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - elif not is_device_supported_mxfp4: - raise ValueError( - "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 " - "(e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) or CPU" - ) - elif not triton_available: - raise ValueError( - "MXFP4 quantization requires Triton: CUDA requires Triton >= 3.4.0, " - "XPU/CPU requires Triton >= 3.5.0. Please install triton: `pip install triton`" - ) - elif not kernels_installed: - raise ValueError("MXFP4 quantization requires the `kernels` package: `pip install kernels>=0.12.0`") - - if not self.pre_quantized: - self._lazy_import_kernels() - - device_map = kwargs.get("device_map") - if device_map is not None and isinstance(device_map, dict): - if not self.pre_quantized and "disk" in device_map.values(): - raise ValueError( - "You are attempting to load an FP4 model with a device_map that contains a disk device." - "This is not supported when the model is quantized on the fly. " - "Please use a quantized checkpoint or remove the disk device from the device_map." - ) - - def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: - from ..integrations import Mxfp4GptOssExperts - - module, tensor_name = get_module_from_name(model, param_name) - if isinstance(module, Mxfp4GptOssExperts): - if tensor_name in ["down_proj_bias", "gate_up_proj_bias"]: - return False - return True - return False - - def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): - # clean cache due to triton ops - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif torch.xpu.is_available(): - torch.xpu.empty_cache() - - def _process_model_before_weight_loading( - self, - model: "PreTrainedModel", - use_kernels: bool = False, - **kwargs, - ): - from ..integrations import replace_with_mxfp4_linear - - # if we are using kernels, we can't use the quantized model, since the forward pass is different and needs special handling - # only CPU and XPU kernels can work with pre-quantized models - device = torch.accelerator.current_accelerator() or torch.device("cpu") - if use_kernels and device.type not in ["cpu", "xpu"]: - logger.warning_once( - "You are using full precision kernels, we will dequantize the model to bf16. " - "To use the quantized model with quantization kernels, please set use_kernels=False" - ) - self.quantization_config.dequantize = True + 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 - if not use_kernels and device.type in ["cpu"]: - logger.warning_once( - "MXFP4 inference on CPU requires use_kernels=True, but use_kernels is disabled. " - "We will dequantize the model to bf16. To run MXFP4 natively on CPU, please set use_kernels=True." + return [ + WeightConverter( + source_patterns=[rf"{proj}_blocks$", rf"{proj}_scales$"], + target_patterns=f"{proj}$", + operations=[FineGrainedPackedBlocks(self), FineGrainedDequantize(self), Transpose(1, 2)], ) - self.quantization_config.dequantize = True - - 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_mxfp4_linear( - model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config - ) - - def update_tp_plan(self, config): - if "GptOssConfig" in config.__class__.__name__: - if getattr(config, "base_model_tp_plan", None) is not None: - config.base_model_tp_plan.update( - { - "layers.*.mlp.experts.gate_up_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.gate_up_proj_scales": "grouped_gemm", - "layers.*.mlp.experts.down_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.down_proj_scales": "grouped_gemm", - } - ) - return config - - def update_ep_plan(self, config): - if "GptOssConfig" in config.__class__.__name__: - if getattr(config, "base_model_ep_plan", None) is not None: - config.base_model_ep_plan.update( - { - "layers.*.mlp.experts.gate_up_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.gate_up_proj_scales": "grouped_gemm", - "layers.*.mlp.experts.down_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.down_proj_scales": "grouped_gemm", - } - ) - return config - - def get_state_dict_and_metadata(self, model): - from ..integrations import Mxfp4GptOssExperts - - state_dict = model.state_dict() - num_local_experts = getattr(model.config, "num_local_experts", 32) - hidden_size = getattr(model.config, "hidden_size", 2880) - - for name, module in model.named_modules(): - if not ( - isinstance(module, Mxfp4GptOssExperts) - and hasattr(module, "gate_up_proj") - and hasattr(module, "down_proj") - ): - continue - - for proj in ("gate_up_proj", "down_proj"): - triton_tensor = getattr(module, proj) - precision_config = getattr(module, f"{proj}_precision_config") - - blocks = triton_tensor.storage.layout.unswizzle_data(triton_tensor.storage.data).transpose(-1, -2) - if proj == "gate_up_proj": - blocks = blocks.reshape(num_local_experts, -1, 90, 16) - else: - blocks = blocks.reshape(num_local_experts, hidden_size, 90, -1) - - scales = precision_config.weight_scale.storage.layout.unswizzle_data( - precision_config.weight_scale.storage.data - ).transpose(-1, -2) - - state_dict[f"{name}.{proj}_blocks"] = blocks - state_dict[f"{name}.{proj}_scales"] = scales - - metadata = {} - return state_dict, metadata - - def is_serializable(self): - return True - - @property - def is_trainable(self) -> bool: - logger.warning_once( - "MXFP4 quantization don't support training, please consider dequantizing the model first by passing quantization_config=Mxfp4Config(dequantize=True) to .from_pretrained()" - ) - return False - - def get_quantize_ops(self): - from ..integrations.mxfp4 import Mxfp4Quantize - - return Mxfp4Quantize(self) + for proj in ("gate_up_proj", "down_proj") + ] - def get_weight_conversions(self): - from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize + def _mxfp4_conversions(self): + """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed weight and + the exponent bytes are the scale grid. Bare patterns, with no `experts.` scope; the layout each lands + in is the module's, from `get_param_ops`.""" + from ..core_model_loading import WeightConverter, WeightRenaming + from ..integrations.finegrained.conversions import FineGrainedPackedBlocks - if self.pre_quantized and self.quantization_config.dequantize: - return [ + conversions = [] + for proj in ("gate_up_proj", "down_proj"): + conversions += [ WeightConverter( - source_patterns=["down_proj_blocks", "down_proj_scales"], - target_patterns=r"down_proj$", - operations=[Mxfp4Dequantize(self)], - ), - WeightConverter( - source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], - target_patterns=["gate_up_proj$"], - operations=[Mxfp4Dequantize(self)], + source_patterns=rf"{proj}_blocks$", + target_patterns=f"{proj}$", + operations=[FineGrainedPackedBlocks(self)], ), + WeightRenaming(source_patterns=rf"{proj}_scales$", target_patterns=f"{proj}_scale_inv"), ] - - return [ - WeightConverter( - source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], - target_patterns=r"gate_up_proj$", - operations=[Mxfp4Deserialize(self)], - ), - WeightConverter( - source_patterns=["down_proj_blocks", "down_proj_scales"], - target_patterns=r"down_proj$", - operations=[Mxfp4Deserialize(self)], - ), - ] + return conversions diff --git a/src/transformers/quantizers/finegrained/mxfp8.py b/src/transformers/quantizers/quantizer_mxfp8.py similarity index 55% rename from src/transformers/quantizers/finegrained/mxfp8.py rename to src/transformers/quantizers/quantizer_mxfp8.py index 1456fc72c14d..a57bf3ff7adc 100644 --- a/src/transformers/quantizers/finegrained/mxfp8.py +++ b/src/transformers/quantizers/quantizer_mxfp8.py @@ -11,12 +11,19 @@ # 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.""" +"""MXFP8: `float8_e4m3fn` weights with a UE8M0 scale per 32. It ships the plain `weight` / `weight_scale_inv` +pair, so it needs nothing beyond what every fine-grained format runs.""" -from .base import FineGrainedHfQuantizer +from typing import TYPE_CHECKING +from .block_scaled_utils import BlockScaledQuantizer -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.""" + +if TYPE_CHECKING: + from ..utils.quantization_config import Mxfp8Config + + +class Mxfp8HfQuantizer(BlockScaledQuantizer): + """MXFP8 weights, W8A8, W8A4 or weight-only.""" + + quantization_config: "Mxfp8Config" diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index fa7161d52a45..752b8fbf6486 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -11,122 +11,147 @@ # 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 +"""NVFP4: E2M1 weights two per byte with an E4M3 scale per 16 and an fp32 global, in the layouts NVIDIA modelopt +exports (per expert per projection, or stacked per layer); quantizes a full-precision checkpoint on the fly too.""" -import warnings from typing import TYPE_CHECKING -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 .block_scaled_utils import BlockScaledQuantizer 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, -) - - -if is_torch_available(): - import torch - if TYPE_CHECKING: - from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import NVFP4Config -def _as_cuda_device(device) -> torch.device | None: - if isinstance(device, int): - return torch.device("cuda", device) - try: - device = torch.device(device) - except (TypeError, RuntimeError): - return None - return device if device.type == "cuda" else None +class NVFP4HfQuantizer(BlockScaledQuantizer): + """NVFP4, in both key layouts modelopt ships: per expert per projection, or stacked per layer.""" + + quantization_config: "NVFP4Config" + # modelopt calibrates an FP8 KV cache beside the weights; the cache here stays in the model dtype + unexpected_keys = (r"\.(k_proj\.k|v_proj\.v)_scale$",) + + def _process_model_after_weight_loading(self, model, **kwargs): + from ..integrations.finegrained import register_shared_gate_up_input_globals + + register_shared_gate_up_input_globals(model) + return super()._process_model_after_weight_loading(model, **kwargs) + + def format_conversions(self, dequantize): + # weight-only keeps activations bf16, so there is no activation global for `input_scale` to land on; nor + # does a dequantized model, which leaves those keys to the unexpected-key filter + calibrated = not dequantize and self.quantization_config.activation_format != "bf16" + conversions = self._nvfp4_renames(calibrated) + if calibrated: + # the down projection's is read with its globals + conversions += self._activation_scale_conversions(projections=("gate_up_proj",), ops=self._reciprocal_ops) + if self.holds_experts and not dequantize: + conversions += self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) + return conversions + + def get_param_ops(self, model, param_name): + """The packed weights first take the int8 view the modules hold: the loader keeps a pre-quantized + non-float tensor in the checkpoint's dtype, modelopt's uint8, and the reverse saves it back as that.""" + from ..integrations.finegrained import FineGrainedLinear + from ..integrations.finegrained.conversions import FineGrainedViewPackedInt8 + + ops = super().get_param_ops(model, param_name) + module, name = get_module_from_name(model, param_name) + if name in ("weight", "gate_up_proj", "up_proj", "down_proj"): + ops = [FineGrainedViewPackedInt8(self), *ops] + # an experts module's globals are inverted in their converters, before the per-expert reductions + elif isinstance(module, FineGrainedLinear) and name in ("weight_global_scale", "activation_scale"): + ops = [*self._reciprocal_ops, *ops] + return ops + @property + def _reciprocal_ops(self) -> list: + """compressed-tensors ships the globals as their reciprocal; modelopt as the multiplier the kernels take.""" + from ..integrations.finegrained.conversions import FineGrainedReciprocal -class NVFP4HfQuantizer(HfQuantizer): - """Quantize eligible linear weights to NVFP4 while loading a full-precision checkpoint.""" - - requires_calibration = False - quantization_config: NVFP4Config + return ( + [FineGrainedReciprocal(self)] if self.quantization_config.checkpoint_format == "compressed-tensors" else [] + ) - def validate_environment(self, device_map, **kwargs): - if self.pre_quantized: - raise ValueError("Loading pre-quantized NVFP4 checkpoints is not supported yet.") - if not is_kernels_available(): - raise ImportError( - "NVFP4 quantization requires the `kernels` package. " - f"Install a compatible version ({KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), " - f"for example with `pip install kernels=={KERNELS_MIN_VERSION}`." - ) - if not torch.cuda.is_available(): - raise RuntimeError("NVFP4 quantization requires a CUDA GPU.") - - if device_map is None: - raise ValueError("NVFP4 quantization requires a CUDA `device_map`.") - devices = set(device_map.values()) if isinstance(device_map, dict) else {device_map} - cuda_devices = {_as_cuda_device(device) for device in devices} - if None in cuda_devices: - raise ValueError("NVFP4 quantization does not support CPU or disk offload.") - if len(cuda_devices) != 1: - raise ValueError( - "NVFP4 quantization currently supports one CUDA device. Tensor parallelism and multi-device " - "`device_map` configurations are not supported yet." + def _nvfp4_renames(self, calibrated: bool): + """modelopt's scales onto the slots every module holds them in, experts included: renamed before converter + collection, a block scale then rides its weight's converter as `weight_scale_inv` (the merge that stacks + the experts' weights stacks their grids too), and the globals reach the converters below by slot name.""" + from ..core_model_loading import WeightRenaming + + renames = [ + WeightRenaming(source_patterns=r"^(.+)\.weight_scale$", target_patterns=r"\1.weight_scale_inv"), + WeightRenaming(source_patterns=r"^(.+)\.weight_scale_2$", target_patterns=r"\1.weight_global_scale"), + ] + if self.quantization_config.checkpoint_format == "compressed-tensors": + # its packed values ship as `weight_packed` and its activation global as `input_global_scale`; its weight + # global already carries the slot's name + renames.append(WeightRenaming(source_patterns=r"^(.+)\.weight_packed$", target_patterns=r"\1.weight")) + if calibrated: + renames.append( + WeightRenaming( + source_patterns=r"^(.+)\.input_global_scale$", target_patterns=r"\1.activation_scale" + ) + ) + return renames + + 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 MergeModulelist, WeightConverter + from ..integrations.finegrained.conversions import FineGrainedWeightGlobals + + converters = [ + # 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=[ + "mlp.experts.*.gate_proj.weight_global_scale", + "mlp.experts.*.up_proj.weight_global_scale", + "mlp.experts.*.down_proj.weight_global_scale", + *(["mlp.experts.*.down_proj.activation_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_activation_scale"] if calibrated else []), + ], + operations=[MergeModulelist(dim=0), *self._reciprocal_ops, FineGrainedWeightGlobals(self)], + ), + ] + return converters + + def _nvfp4_fused_conversions(self, calibrated: bool): + """Already stacked per layer, as the vLLM fused layout ships them — so no `MergeModulelist`: + the scales only take their slot's name, and the module's ops give every tensor its layout.""" + from ..core_model_loading import WeightConverter, WeightRenaming + from ..integrations.finegrained.conversions import FineGrainedWeightGlobals + + converters = [ + WeightRenaming( + source_patterns=rf"experts\.{proj}_weight_scale$", target_patterns=f"experts.{proj}_scale_inv" ) - self.quantization_device = cuda_devices.pop() - - major, minor = torch.cuda.get_device_capability(self.quantization_device) - if major < 10: - raise RuntimeError(f"NVFP4 requires Blackwell (sm100+) block-scaled tensor cores; found sm{major}{minor}.") - - def update_tp_plan(self, config): - distributed_config = getattr(config, "distributed_config", None) - tp_size = getattr(distributed_config, "tp_size", None) - if tp_size is not None and tp_size > 1: - raise ValueError( - "NVFP4 quantization does not support tensor parallelism yet because its scale metadata does not have " - "a sharding plan." + 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_activation_scale"] if calibrated else []), + ], + operations=[*self._reciprocal_ops, FineGrainedWeightGlobals(self)], ) - return config - - def param_needs_quantization(self, model: PreTrainedModel, param_name: str, **kwargs) -> bool: - from ..integrations.nvfp4 import NVFP4Linear - - module, tensor_name = get_module_from_name(model, param_name) - return isinstance(module, NVFP4Linear) and tensor_name == "weight" - - def param_element_size(self, model: PreTrainedModel, param_name: str, param: torch.Tensor) -> float: - if self.param_needs_quantization(model, param_name): - # NVFP4 packs two 4-bit weights per byte; this estimate excludes the smaller scale metadata. - return 0.5 - return super().param_element_size(model, param_name, param) - - def _process_model_before_weight_loading(self, model: PreTrainedModel, **kwargs): - from ..integrations.nvfp4 import replace_with_nvfp4_linear - - self.modules_to_not_convert = self.get_modules_to_not_convert( - model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules ) - replace_with_nvfp4_linear(model, modules_to_not_convert=self.modules_to_not_convert) - - def get_quantize_ops(self): - from ..integrations.nvfp4 import NVFP4Quantize - - return NVFP4Quantize(self.quantization_device) - - @property - def is_serializable(self) -> bool: - return False - - @property - def is_trainable(self) -> bool: - return False - - @property - def is_compileable(self) -> bool: - return True + return converters diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 42036ff4416f..7b95f484f8d6 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -23,6 +23,72 @@ def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: return module, tensor_name +def module_kind(module_name: str, module) -> str | None: + """The kind of layer a quantized-module registry is keyed by: `"experts"`, `"grouped_linear"` (a + block-diagonal linear, matched on the attribute a swap needs rather than its class name), `"linear"` or + `"embedding"`; `None` for anything a registry does not replace.""" + from torch import nn + + if module_name.endswith(".experts"): + return "experts" + if isinstance(module, nn.Linear) and hasattr(module, "n_groups"): + return "grouped_linear" + if type(module) is nn.Linear: + return "linear" + if type(module) is nn.Embedding: + return "embedding" + return None + + +EXPERTS_PATTERN = r"\.experts($|\.)" + + +def embedding_pattern(name: str) -> str: + """The `dtype_plan` key of a `modules_to_convert` entry, which names an embedding table by suffix.""" + return rf"{re.escape(name)}$" + + +def plan_entry_for(dtype_plan: dict, module_name: str, catch_all: bool = True): + """The format config `dtype_plan` gives `module_name`, `None` if none covers it. + + A targeted entry 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 entries claiming the same module is ambiguous, and says so. Without + `catch_all`, only an entry that names the module counts. + """ + claimed = [pattern for pattern in dtype_plan if pattern != ".*" and re.search(pattern, module_name)] + if len(claimed) > 1: + raise ValueError(f"{module_name!r} is claimed by more than one dtype_plan entry: {claimed}") + if claimed: + return dtype_plan[claimed[0]] + return dtype_plan.get(".*") if catch_all else None + + +def replace_with_quantized_modules( + model, dtype_plan: dict, modules_to_not_convert: list[str] | None, module_classes: dict[str, type] +) -> bool: + """Swap every module `dtype_plan` gives a config for the class `module_classes` builds that kind of layer as, built + on `meta` by its `from_module(module, config, model)`. A module the plan gives no config, of a kind with no class, + or named by `modules_to_not_convert` is kept. Returns whether any module was replaced.""" + import torch + + replaced = False + for module_name, module in model.named_modules(): + kind = module_kind(module_name, module) + if kind is None or not should_convert_module(module_name, modules_to_not_convert): + continue + # the catch-all covers linears and experts; an embedding is quantized only where an entry names it + config = plan_entry_for(dtype_plan, module_name, catch_all=kind != "embedding") + quantized_class = module_classes.get(kind) if config is not None else None + if quantized_class is None: + continue + with torch.device("meta"): + quantized = quantized_class.from_module(module, config, model) + model.set_submodule(module_name, quantized) + replaced = True + return replaced + + 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 60aeea690d8b..dc67a505c6a5 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 asdict, dataclass, field, replace +from dataclasses import dataclass from enum import Enum from typing import Any, Optional, Union @@ -63,6 +63,8 @@ class QuantizationMethod(str, Enum): AUTOROUND = "auto-round" MXFP4 = "mxfp4" NVFP4 = "nvfp4" + FINEGRAINED = "finegrained" + MODELOPT = "modelopt" MXFP8 = "mxfp8" METAL = "metal" FOUR_OVER_SIX = "fouroversix" @@ -1701,65 +1703,200 @@ def subtree_patterns(globs: list[str]) -> list[str]: return patterns -@dataclass -class FineGrainedGroup: - """One set of modules and the format they are quantized in. +class BlockScaledConfig(QuantizationConfigMixin): + """What every block-scaled weight format's config holds: the options of the format itself, which the quantized + modules read, and the loading options. One such config quantizes a whole model; a [`FineGrainedConfig`] maps + several onto parts of one. - 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. + Args: + activation_format (`str`, *optional*): + What activations are quantized to: `None` for the weights' own format, `"bf16"` for weight-only, or + another format the weights' kernels serve. + activation_scheme (`str`, *optional*, defaults to `"dynamic"`): + `"dynamic"` (inline per-token/per-block) or `"static"` (a calibrated per-tensor scale from the checkpoint). + weight_block_size (`tuple[int, int]`, *optional*): + The block of the weight scale grid, for formats whose granularity is not fixed by the format. + scale_fmt (`str`, *optional*, defaults to `"float"`): + Storage dtype of block scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`). + dequantize (`bool`, *optional*, defaults to `False`): + Whether to dequantize the model to full precision while loading. + modules_to_not_convert (`list`, *optional*): + Modules to keep in full precision. `ignored_layers` is accepted as an alias. + checkpoint_format (`str`, *optional*): + The producer whose key layout the checkpoint ships when it is not this format's own: + `"compressed-tensors"`. """ - quant_method: str - 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 - 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) - - # 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: + # what `quant_method` this config is, and the activation formats its weights' kernels serve + method: QuantizationMethod + activation_formats: tuple[str | None, ...] = (None,) + + def __init__( + self, + activation_format: str | None = None, + activation_scheme: str = "dynamic", + weight_block_size: tuple[int, int] | None = None, + scale_fmt: str = "float", + dequantize: bool = False, + modules_to_not_convert: list | None = None, + checkpoint_format: str | None = None, + **kwargs, + ): + self.quant_method = self.method + self.checkpoint_format = checkpoint_format + kwargs.pop("quant_method", None) + if modules_to_not_convert is None and "ignored_layers" in kwargs: + modules_to_not_convert = kwargs.pop("ignored_layers") + self.activation_format = activation_format + self.activation_scheme = activation_scheme + self.weight_block_size = tuple(weight_block_size) if weight_block_size is not None else None + self.scale_fmt = scale_fmt + self.dequantize = dequantize + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + self.activation_scheme = self.activation_scheme.lower() + if self.activation_format not in self.activation_formats: raise ValueError( - f"{self.quant_method} weights take activation_format in {supported}; got {self.activation_format!r}" + f"{self.method.value} weights take activation_format in {self.activation_formats}; " + f"got {self.activation_format!r}" ) + if self.activation_scheme not in ("dynamic", "static"): + raise ValueError(f"activation_scheme must be 'dynamic' or 'static'; got {self.activation_scheme!r}") + 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.checkpoint_format not in (None, "compressed-tensors"): + raise ValueError(f"checkpoint_format must be None or 'compressed-tensors'; got {self.checkpoint_format!r}") + if self.weight_block_size is not None and ( + len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0 + ): + raise ValueError(f"weight_block_size must be two positive integers; got {self.weight_block_size}") - def matches(self, module_name: str) -> bool: - return any(re.search(target, module_name) for target in self.targets) + def get_loading_attributes(self): + # a skip-list left unset keeps the checkpoint's + attributes = {"dequantize": self.dequantize} + if self.modules_to_not_convert is not None: + attributes["modules_to_not_convert"] = self.modules_to_not_convert + return attributes -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. +class BlockFp8Config(BlockScaledConfig): + """Block-FP8 weights: `float8_e4m3fn` values with fp32 or UE8M0 scales over `weight_block_size` blocks (one + per-tensor scale without a block), as DeepSeek-V3 and Qwen3 ship them; W8A8 with dynamic or calibrated activations. - 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. + Args: + weight_block_size (`tuple[int, int]`, *optional*, defaults to `(128, 128)`): + The block each weight scale covers; `None` for one scale per tensor. + scale_fmt (`str`, *optional*, defaults to `"float"`): + Storage dtype of the scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`, DeepSeek-V4). + activation_scheme (`str`, *optional*, defaults to `"dynamic"`): + `"dynamic"` or `"static"` (a calibrated per-tensor activation scale from the checkpoint). + modules_to_convert (`list`, *optional*): + Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.FP8 + activation_formats = (None, "fp8") + + def __init__( + self, + weight_block_size: tuple[int, int] | None = (128, 128), + scale_fmt: str = "float", + activation_scheme: str = "dynamic", + modules_to_convert: list | None = None, + **kwargs, + ): + self.modules_to_convert = modules_to_convert + super().__init__( + weight_block_size=weight_block_size, scale_fmt=scale_fmt, activation_scheme=activation_scheme, **kwargs + ) + + +class Mxfp8Config(BlockScaledConfig): + """MXFP8 weights: `float8_e4m3fn` values with a UE8M0 scale per 32, as MiniMax-M3 ships them. + + Args: + activation_format (`str`, *optional*): + `None` for MXFP8 activations, `"mxfp4"` for W8A4, or `"bf16"` for weight-only. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.MXFP8 + activation_formats = (None, "mxfp8", "mxfp4", "bf16") + + def __init__(self, activation_format: str | None = None, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) + + +class Mxfp4Config(BlockScaledConfig): + """MXFP4 weights: E2M1 values two per byte with a UE8M0 scale per 32, as GPT-OSS and DeepSeek-V4's experts ship + them. + + Args: + activation_format (`str`, *optional*, defaults to `"bf16"`): + `"bf16"` for weight-only (W4A16, what GPT-OSS runs), `"mxfp8"` for W4A8, or `"mxfp4"` (or `None`) for W4A4. + modules_to_not_convert (`list`, *optional*): + Modules to keep in full precision. + dequantize (`bool`, *optional*, defaults to `False`): + Whether to dequantize the model to full precision while loading. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.MXFP4 + activation_formats = (None, "mxfp8", "mxfp4", "bf16") + + def __init__( + self, + modules_to_not_convert: list | None = None, + dequantize: bool = False, + activation_format: str | None = "bf16", + **kwargs, + ): + super().__init__( + activation_format=activation_format, + dequantize=dequantize, + modules_to_not_convert=modules_to_not_convert, + **kwargs, + ) + + +class NVFP4Config(BlockScaledConfig): + """NVFP4 weights: E2M1 values two per byte with an E4M3 scale per 16 and an fp32 global, as NVIDIA modelopt + exports them; quantizes a full-precision checkpoint on the fly as well. + + Args: + modules_to_not_convert (`list[str]`, *optional*): + Modules to keep in full precision. + activation_format (`str`, *optional*): + `None` for NVFP4 activations (W4A4, with the checkpoint's calibrated activation global) or `"bf16"` for + weight-only. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. """ - 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, - } + method = QuantizationMethod.NVFP4 + activation_formats = (None, "nvfp4", "bf16") -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. + def __init__( + self, modules_to_not_convert: list[str] | None = None, activation_format: str | None = None, **kwargs + ): + super().__init__(activation_format=activation_format, modules_to_not_convert=modules_to_not_convert, **kwargs) + + +# the block-scaled formats by their `quant_method`, for a plan entry read back from `config.json` +BLOCK_SCALED_CONFIGS = { + config.method.value: config for config in (BlockFp8Config, Mxfp8Config, Mxfp4Config, NVFP4Config) +} + + +def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: + """One `config_groups` entry as `dtype_plan` entries, one per target, 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. `group_size` is what separates the two 4-bit formats: nvfp4's block scale is E4M3 over 16 values, mxfp4's is E8M0 @@ -1773,9 +1910,9 @@ def group_from_config_groups(spec: dict) -> FineGrainedGroup | None: (8, "float", None): "fp8", } weights = spec.get("weights") or {} - key = (weights.get("num_bits"), weights.get("type"), weights.get("group_size")) - quant_method = formats.get(key) - if quant_method is None: + quant_method = formats.get((weights.get("num_bits"), weights.get("type"), weights.get("group_size"))) + # one FP8 scale per output channel has no block the kernels tile + if quant_method is None or weights.get("strategy") == "channel": return None targets = [] for target in spec.get("targets") or []: @@ -1787,195 +1924,149 @@ 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=activation_format, - activation_scheme="static" if static else "dynamic", - ) + # no activation quantization is weight-only, not the weights' native activation format + activation_format = formats.get(act_key) if activations else "bf16" + options = {"activation_format": None if activation_format == quant_method else activation_format} + if quant_method == "fp8": + # 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" + options.update( + activation_scheme="static" if static else "dynamic", weight_block_size=weights.get("block_structure") + ) + config = BLOCK_SCALED_CONFIGS[quant_method](**options) + return dict.fromkeys(targets or [".*"], config) + + +def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: + """The config an NVIDIA modelopt export describes: its one format's, or a [`FineGrainedConfig`] when its + `config_groups` cover different modules with different formats. + + `quant_algo` names what it exported; modelopt also covers FP8, INT4 AWQ, W4A8 and MXFP4, which are refused + rather than silently mis-read. The calibrated `input_scale` TENSORS are the loader's business, not the config's. + """ + config_dict = dict(config_dict) + quant_algo = config_dict.pop("quant_algo", None) + if quant_algo != "NVFP4": + raise ValueError(f"modelopt checkpoints are supported for quant_algo in ['NVFP4']; got {quant_algo!r}") + logger.info("modelopt checkpoint exported with quant_algo='NVFP4'.") + if config_dict.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." + ) + ignore = config_dict.pop("ignore", None) or config_dict.pop("exclude_modules", None) + modules_to_not_convert = subtree_patterns(ignore) if ignore is not None else None + dtype_plan = {".*": NVFP4Config(activation_format="nvfp4")} + # the same thing described parametrically, taken only when every group maps, so a multi-group export is not + # flattened to one format + if groups := config_dict.pop("config_groups", None): + mapped = [plan_from_config_groups(spec) for spec in groups.values()] + if all(entries is not None for entries in mapped): + dtype_plan = {pattern: config for entries in mapped for pattern, config in entries.items()} + return _config_for_plan(dtype_plan, modules_to_not_convert) + + +def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None: + """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, MX, + per-channel, sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's + `checkpoint_format` names the layout its tensors ship in, which a save writes back.""" + if config_dict.get("sparsity_config") or config_dict.get("transform_config"): + return None + if config_dict.get("quantization_status", "compressed") != "compressed": + return None + groups = config_dict.get("config_groups") or {} + served = ("float-quantized", "nvfp4-pack-quantized") + if not groups or any(spec.get("format", config_dict.get("format")) not in served for spec in groups.values()): + return None + mapped = [plan_from_config_groups(spec) for spec in groups.values()] + if any(entries is None for entries in mapped): + return None + dtype_plan = {pattern: config for entries in mapped for pattern, config in entries.items()} + for config in dtype_plan.values(): + config.checkpoint_format = "compressed-tensors" + if config_dict.get("kv_cache_scheme"): + logger.warning_once( + "This compressed-tensors checkpoint calibrates a quantized KV cache; the KV cache is not quantized here " + "and stays in the model's dtype, so those scales are not loaded." + ) + # entries are module names, or `re:`-prefixed patterns; a bare name is that module and its sub-tree, not a prefix + ignore = config_dict.get("ignore") or [] + skip = [entry[3:] for entry in ignore if entry.startswith("re:")] + skip += subtree_patterns([entry for entry in ignore if not entry.startswith("re:")]) + return _config_for_plan(dtype_plan, skip or None) + + +def _config_for_plan(dtype_plan: dict, modules_to_not_convert: list | None) -> QuantizationConfigMixin: + """The one format config of a plan that covers everything with it, else the `FineGrainedConfig` of the plan.""" + if list(dtype_plan) == [".*"]: + (config,) = dtype_plan.values() + config.modules_to_not_convert = modules_to_not_convert + return config + return FineGrainedConfig(dtype_plan=dtype_plan, modules_to_not_convert=modules_to_not_convert) -@dataclass class FineGrainedConfig(QuantizationConfigMixin): """ - Configuration for the fine-grained quantization family served by the - `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` | + Configuration for a model whose modules are quantized in different fine-grained formats: which format covers + which modules, each as its own format config ([`BlockFp8Config`], [`Mxfp8Config`], [`Mxfp4Config`], + [`NVFP4Config`]), and the options that hold for the whole model. A model in one format takes that format's config + alone. + + | 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"`): - 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 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. + dtype_plan (`dict[str, BlockScaledConfig]`, *optional*): + A module regex and the format config the modules it matches load in. A targeted entry wins over the + `".*"` catch-all; modules no entry covers stay in full precision. modules_to_not_convert (`list`, *optional*): - A list of module names that should not be converted during quantization. + Modules to keep in full precision. 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). - 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. - 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. + Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. + dequantize (`bool`, *optional*, defaults to `False`): + Whether to dequantize the model to full precision while loading. """ def __init__( self, - activation_scheme: str = "dynamic", - weight_block_size: tuple[int, int] = (128, 128), - dequantize: bool = False, + dtype_plan: dict[str, BlockScaledConfig] | None = None, modules_to_not_convert: list | None = None, modules_to_convert: list | None = None, - scale_fmt: str = "float", - activation_format: str | None = None, - groups: dict[str, FineGrainedGroup] | None = None, + dequantize: bool = False, **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") - # "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: - self.quant_method, algo_activation_format, algo_groups, ignore = self._read_modelopt(kwargs) - self.activation_format = activation_format or algo_activation_format - groups = groups or algo_groups - if modules_to_not_convert is None and ignore is not None: - modules_to_not_convert = subtree_patterns(ignore) + self.quant_method = QuantizationMethod.FINEGRAINED + # from `config.json` an entry is a plain dict naming its format + self.dtype_plan = { + pattern: config + if isinstance(config, BlockScaledConfig) + else BLOCK_SCALED_CONFIGS[config["quant_method"]].from_dict(config) + for pattern, config in (dtype_plan or {}).items() + } self.modules_to_not_convert = modules_to_not_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 - # 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() - - @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) - 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): - """`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 != [".*"] 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 == [".*"]: - return group - raise KeyError(f"no group covers {module_name!r}; groups: {list(self.groups)}") - - 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}") - 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 to_dict(self) -> dict[str, Any]: + return { + "quant_method": self.quant_method, + "dtype_plan": {pattern: config.to_dict() for pattern, config in self.dtype_plan.items()}, + "modules_to_not_convert": self.modules_to_not_convert, + "modules_to_convert": self.modules_to_convert, + "dequantize": self.dequantize, + } def get_loading_attributes(self): - return {"dequantize": self.dequantize, "modules_to_not_convert": self.modules_to_not_convert} + # a skip-list left unset keeps the checkpoint's + attributes = {"dequantize": self.dequantize} + if self.modules_to_not_convert is not None: + attributes["modules_to_not_convert"] = self.modules_to_not_convert + return attributes class FineGrainedFP8Config(QuantizationConfigMixin): @@ -2080,41 +2171,6 @@ def __init__( self.quant_method = QuantizationMethod.QUARK -@dataclass -class Mxfp4Config(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using mxfp4 quantization. - - Args: - modules_to_not_convert (`list`, *optional*, default to `None`): - The list of modules to not quantize, useful for quantizing models that explicitly require to have - some modules left in their original precision. - dequantize (`bool`, *optional*, default to `False`): - Whether we dequantize the model to bf16 precision or not - """ - - def __init__( - self, - modules_to_not_convert: list | None = None, - dequantize: bool = False, - **kwargs, - ): - self.quant_method = QuantizationMethod.MXFP4 - self.modules_to_not_convert = modules_to_not_convert - self.dequantize = dequantize - - def get_loading_attributes(self): - return {"dequantize": self.dequantize} - - def to_dict(self) -> dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Returns: - `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. - """ - return {"quant_method": self.quant_method, "modules_to_not_convert": self.modules_to_not_convert} - - class MetalConfig(QuantizationConfigMixin): """ Configuration class for Metal affine quantization targeting Apple Silicon (MPS) devices. @@ -2357,25 +2413,6 @@ def __init__( self.modules_to_not_convert = modules_to_not_convert -class NVFP4Config(QuantizationConfigMixin): - """Configuration for on-the-fly NVFP4 weight quantization. - - Args: - modules_to_not_convert (`list[str]`, *optional*): - Module-name patterns that should remain in their original precision. - kwargs (`dict[str, Any]`, *optional*): - Additional values are ignored and reported through the Transformers logger. - """ - - def __init__(self, modules_to_not_convert: list[str] | None = None, **kwargs): - self.quant_method = QuantizationMethod.NVFP4 - self.modules_to_not_convert = modules_to_not_convert - if kwargs: - logger.info( - f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__.__name__}." - ) - - @dataclass class GgufConfig(QuantizationConfigMixin): """Load a quantized GGUF checkpoint with its weights left in GGUF blocks. From 0b3eba2fabc410e2172010fbfad1639ee3ffa625 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 7 Oct 2026 07:46:01 -0700 Subject: [PATCH 68/91] Block-scaled quantization: config review fixes, scoped checkpoint renames, shared stack names - A passed config lends a pre-quantized checkpoint only the activation format it names; NVFP4 weight-only checkpoints refuse an activation format they ship no calibrated globals for. - `FineGrainedFP8Config` passed explicitly loads through its frozen integration; a checkpoint's `fp8` config takes the block-FP8 path. - DeepSeek-V4's `.scale` block-scale rename applies only to pre-quantized loads (`deepseek_v4_quantized`), so other models' own `.scale` parameters are untouched. - Experts' TP plan resolves every layer's format for the glob it shards, refusing a plan that mixes block-scale grids and per-tensor scales across layers. - Gated experts classes cached per flag combination (picklable); per-forward assertions and the gate|up fold checks share one walk. - Expert stack names and companions defined once (`FUSED_STACKS`, `STACK_COMPANIONS`). - Tests: shared block-scaled tester (format specs, tiny checkpoints, load-path equivalence, mixed plans), reusing the repo's port helper and the real conversion mappings. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/source/en/_toctree.yml | 10 +- docs/source/en/experts_interface.md | 2 +- docs/source/en/main_classes/quantization.md | 22 +- docs/source/en/model_doc/deepseek_v32.md | 16 +- docs/source/en/quantization/blockfp8.md | 121 ++ .../en/quantization/compressed_tensors.md | 8 + docs/source/en/quantization/concept_guide.md | 2 +- docs/source/en/quantization/contribute.md | 2 +- docs/source/en/quantization/finegrained.md | 68 + .../source/en/quantization/finegrained_fp8.md | 82 +- docs/source/en/quantization/mxfp4.md | 90 +- docs/source/en/quantization/mxfp8.md | 72 ++ docs/source/en/quantization/nvfp4.md | 72 +- docs/source/en/quantization/overview.md | 7 +- docs/source/en/quantization/quark.md | 4 +- src/transformers/conversion_mapping.py | 11 +- src/transformers/core_model_loading.py | 75 +- src/transformers/integrations/deepgemm.py | 50 +- .../integrations/finegrained/__init__.py | 2 +- .../integrations/finegrained/conversions.py | 145 ++- .../integrations/finegrained/layers.py | 121 +- src/transformers/quantizers/auto.py | 163 ++- .../quantizers/block_scaled_utils.py | 116 +- .../quantizers/quantizer_blockfp8.py | 12 +- .../quantizers/quantizer_finegrained.py | 13 +- .../quantizers/quantizer_finegrained_fp8.py | 2 +- .../quantizers/quantizer_mxfp4.py | 6 - .../quantizers/quantizer_nvfp4.py | 69 +- .../quantizers/quantizers_utils.py | 5 + src/transformers/utils/quantization_config.py | 276 ++-- .../finegrained/test_conversions.py | 1142 ++++------------- tests/integrations/finegrained/test_core.py | 565 -------- .../finegrained/test_distributed.py | 545 -------- .../integrations/finegrained/test_forwards.py | 442 ------- tests/integrations/finegrained/test_layers.py | 500 ++++++++ tests/integrations/finegrained/test_models.py | 250 ---- tests/integrations/finegrained/utils.py | 258 ++++ tests/integrations/test_finegrained_fp8.py | 349 ----- tests/kernels/test_deepgemm.py | 153 +++ tests/kernels/test_finegrained.py | 289 +---- .../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 | 22 +- tests/quantization/block_scaled_tester.py | 1031 +++++++++++++++ tests/quantization/blockfp8/__init__.py | 0 tests/quantization/blockfp8/test_blockfp8.py | 529 ++++++++ tests/quantization/finegrained/__init__.py | 0 .../finegrained/test_finegrained.py | 598 +++++++++ .../quantization/finegrained_fp8/test_fp8.py | 545 +------- tests/quantization/mxfp4/test_mxfp4.py | 664 +++------- tests/quantization/mxfp8/__init__.py | 0 tests/quantization/mxfp8/test_mxfp8.py | 70 + tests/quantization/nvfp4/test_nvfp4.py | 586 +++++++-- 54 files changed, 5258 insertions(+), 4942 deletions(-) create mode 100644 docs/source/en/quantization/blockfp8.md create mode 100644 docs/source/en/quantization/finegrained.md create mode 100644 docs/source/en/quantization/mxfp8.md delete mode 100644 tests/integrations/finegrained/test_core.py delete mode 100644 tests/integrations/finegrained/test_distributed.py delete mode 100644 tests/integrations/finegrained/test_forwards.py create mode 100644 tests/integrations/finegrained/test_layers.py delete mode 100644 tests/integrations/finegrained/test_models.py create mode 100644 tests/integrations/finegrained/utils.py delete mode 100644 tests/integrations/test_finegrained_fp8.py create mode 100644 tests/quantization/block_scaled_tester.py create mode 100644 tests/quantization/blockfp8/__init__.py create mode 100644 tests/quantization/blockfp8/test_blockfp8.py create mode 100644 tests/quantization/finegrained/__init__.py create mode 100644 tests/quantization/finegrained/test_finegrained.py create mode 100644 tests/quantization/mxfp8/__init__.py create mode 100644 tests/quantization/mxfp8/test_mxfp8.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index e069c7966028..21e4b4dbfd6e 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -281,18 +281,24 @@ title: torchao title: Quantize on the fly - sections: + - local: quantization/blockfp8 + title: Block FP8 - local: quantization/fbgemm_fp8 title: FBGEMM - - local: quantization/finegrained_fp8 - title: Fine-grained FP8 - local: quantization/fouroversix title: Four Over Six - local: quantization/fp_quant title: FP-Quant - local: quantization/mxfp4 title: MXFP4 + - local: quantization/mxfp8 + title: MXFP8 - local: quantization/nvfp4 title: NVFP4 + - local: quantization/finegrained + title: Fine-grained + - local: quantization/finegrained_fp8 + title: Fine-grained FP8 (deprecated) title: FP4 and FP8 formats - sections: - local: quantization/auto_round diff --git a/docs/source/en/experts_interface.md b/docs/source/en/experts_interface.md index cb40eb96b6d7..8411290aade3 100644 --- a/docs/source/en/experts_interface.md +++ b/docs/source/en/experts_interface.md @@ -95,7 +95,7 @@ model = AutoModelForCausalLM.from_pretrained( ## DeepGEMM -The `"deepgemm"` backend routes expert matmuls through the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels distributed by [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). It works with unquantized `bfloat16` experts and with FP8/FP4-quantized experts loaded through [Fine-grained FP8](./quantization/finegrained_fp8). +The `"deepgemm"` backend routes expert matmuls through the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels distributed by [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). It works with unquantized `bfloat16` experts, with [Block FP8](./quantization/blockfp8) experts (128x128 blocks), and with [MXFP4](./quantization/mxfp4) experts run with FP8 activations (DeepSeek-V4's W4A8, Blackwell only), including a [Fine-grained](./quantization/finegrained) plan of them. On quantized experts it quantizes activations to FP8 at run time whatever their `activation_format`, and adds no expert biases; use `grouped_mm` or `batched_mm` for MXFP8 or NVFP4 experts, biased experts, or another activation format. The `"deepgemm"` backend requires: diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index ae5c655e2b23..b1307ff3d002 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -70,14 +70,6 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] MetalConfig -## Mxfp4Config - -[[autodoc]] Mxfp4Config - -## NVFP4Config - -[[autodoc]] NVFP4Config - ## FbgemmFp8Config [[autodoc]] FbgemmFp8Config @@ -110,6 +102,14 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] Mxfp8Config +## Mxfp4Config + +[[autodoc]] Mxfp4Config + +## NVFP4Config + +[[autodoc]] NVFP4Config + ## FineGrainedConfig [[autodoc]] FineGrainedConfig @@ -117,9 +117,9 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui ## FineGrainedFP8Config > [!WARNING] -> `FineGrainedFP8Config` is frozen and will be removed in a future release. A config passed to `from_pretrained` -> loads as the config it spells: [`Mxfp8Config`] for a `[1, 32]` block, [`BlockFp8Config`] otherwise, and a -> [`FineGrainedConfig`] of block-FP8 over MXFP4 experts for a model that declares `expert_dtype="fp4"` (DeepSeek-V4). +> `FineGrainedFP8Config` is deprecated and will be removed in a future release. Passed to `from_pretrained`, it loads +> through its own frozen integration; use [`BlockFp8Config`], or [`Mxfp8Config`] for a `[1, 32]` block with UE8M0 +> scales (see [Fine-grained FP8](../quantization/finegrained_fp8)). [[autodoc]] FineGrainedFP8Config diff --git a/docs/source/en/model_doc/deepseek_v32.md b/docs/source/en/model_doc/deepseek_v32.md index 6722253256c9..2bfac5f953a8 100644 --- a/docs/source/en/model_doc/deepseek_v32.md +++ b/docs/source/en/model_doc/deepseek_v32.md @@ -55,23 +55,13 @@ In DeepSeek-V3.2 **every layer runs its own indexer** — there is no cross-laye ## Usage examples -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: +DeepSeek-V3.2-Exp is distributed as an FP8 checkpoint, loaded from its own quantization config. The router and the indexer's `weights_proj` stay in full precision: ```python -from transformers import BlockFp8Config, AutoModelForCausalLM, AutoTokenizer -import torch +from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "deepseek-ai/DeepSeek-V3.2-Exp" -quantization_config = BlockFp8Config( - modules_to_not_convert=["model.layers.*.mlp.gate.*", "*.self_attn.indexer.weights_proj.*"], - weight_block_size=(128, 128), -) -model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype="auto", - device_map="auto", - quantization_config=quantization_config, -) +model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) diff --git a/docs/source/en/quantization/blockfp8.md b/docs/source/en/quantization/blockfp8.md new file mode 100644 index 000000000000..a6dd9d8776d3 --- /dev/null +++ b/docs/source/en/quantization/blockfp8.md @@ -0,0 +1,121 @@ + + +# Block FP8 + +Block FP8 stores weights in 8-bit floating point (`float8_e4m3fn`) with one scale per 2D block of the weight matrix +(`weight_block_size=(128, 128)` by default), or one scale per tensor (`weight_block_size=None`). Activations are +quantized to FP8 too, per token and block at run time (`activation_scheme="dynamic"`), or against a calibrated scale +shipped with the checkpoint (`activation_scheme="static"`). + +It is the format of [DeepSeek-V3](https://hf.co/papers/2412.19437) and DeepSeek-R1, Qwen3-FP8, and Mistral's static +FP8 checkpoints. [`BlockFp8Config`] loads these checkpoints as they ship and quantizes full-precision checkpoints on the +fly. Checkpoints declare it as `quant_method="fp8"`. + +
+ +
+ +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer (Ada, Hopper, Blackwell) or an Intel XPU, +> and the [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. + +```bash +pip install --upgrade accelerate kernels +``` + +## Load a pre-quantized checkpoint + +A pre-quantized checkpoint carries its quantization config, so no config is needed. + +```py +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_id = "Qwen/Qwen3-8B-FP8" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") + +tokenizer = AutoTokenizer.from_pretrained(model_id) +inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) +output = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(output[0], skip_special_tokens=True)) +``` + +Checkpoints exported by [compressed-tensors](./compressed_tensors) in block FP8 (`format="float-quantized"`) load as +[`BlockFp8Config`] as well. + +## Quantize on the fly + +Pass a [`BlockFp8Config`] to [`~PreTrainedModel.from_pretrained`] to quantize a full-precision checkpoint as it loads. +Set `dtype="auto"` so the unquantized modules keep the checkpoint's dtype. + +```py +from transformers import AutoModelForCausalLM, BlockFp8Config + +quantization_config = BlockFp8Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "meta-llama/Llama-3.1-8B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +``` + +[`~PreTrainedModel.save_pretrained`] writes the quantized model back in the same layout, and +[`~PreTrainedModel.from_pretrained`] reloads it. + +```py +model.save_pretrained("llama-3.1-8b-fp8") +model = AutoModelForCausalLM.from_pretrained("llama-3.1-8b-fp8", device_map="auto") +``` + +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"fp8"` | W8A8 | FP8, per token and block at run time, or against a calibrated per-tensor scale with `activation_scheme="static"` | + +## Options + +- `weight_block_size` is the `(rows, columns)` block each weight scale covers; `None` gives one scale per tensor. +- `activation_scheme="static"` quantizes activations against the calibrated scale a checkpoint ships (its + `input_scale`), instead of deriving one per token at run time. +- `scale_fmt="ue8m0"` holds the block scales as powers of two (`float8_e8m0fnu`), as DeepSeek-V4 ships them. +- `modules_to_not_convert` keeps the modules it names in full precision; `modules_to_convert` names embedding tables to + hold in FP8 with one per-tensor scale (a table too large to keep in full precision). +- `dequantize=True` loads a pre-quantized checkpoint dequantized to full precision. + +## DeepGEMM fast path + +On Hopper GPUs, a dense FP8 linear with `weight_block_size=(128, 128)` and dynamic activations dispatches to the +[DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels from +[kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). On Blackwell the Triton kernels are +faster and are always used. + +DeepGEMM JIT-compiles its kernels with the system `nvcc`, so a full CUDA toolkit must be installed (12.3 or later on +Hopper). Transformers finds the toolkit by checking `CUDA_HOME`, then `CUDA_PATH`, then `nvcc` on your `PATH`, then +`/usr/local/cuda`. If DeepGEMM cannot load, Transformers logs a warning once and uses the Triton kernels. Set +`TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to always use them. + +For MoE experts, DeepGEMM is opt-in: pass `experts_implementation="deepgemm"` (or `"deepgemm_megamoe"` on Blackwell) +to [`~PreTrainedModel.from_pretrained`]. The DeepGEMM experts kernels on Blackwell read only UE8M0 scales; with +`float32` scales, use the default `grouped_mm` and `batched_mm` experts implementations. See the +[Experts backends](../experts_interface) guide for every option. + +## Related formats + +- [MXFP8](./mxfp8) holds FP8 weights with one power-of-two scale per 32 values of a row. +- [Fine-grained](./finegrained) mixes block FP8 with the 4-bit formats in one model, as DeepSeek-V4 ships it. diff --git a/docs/source/en/quantization/compressed_tensors.md b/docs/source/en/quantization/compressed_tensors.md index cdb397c16555..18f985872e13 100644 --- a/docs/source/en/quantization/compressed_tensors.md +++ b/docs/source/en/quantization/compressed_tensors.md @@ -49,6 +49,14 @@ pip install -e . Search using the compressed-tensors [tag](https://huggingface.co/models?other=compressed-tensors) to find a compatible model on the Hugging Face Hub. +> [!NOTE] +> A checkpoint in a format the [fine-grained](./finegrained) kernels serve loads through them, with +> `checkpoint_format="compressed-tensors"`: `float-quantized` FP8 with per-tensor or block weight scales as +> [`BlockFp8Config`], `nvfp4-pack-quantized` as [`NVFP4Config`], and one whose config groups mix formats as a +> [`FineGrainedConfig`]. Integer, per-channel FP8, sparse and transformed checkpoints stay with compressed-tensors, as +> do `naive-quantized` ones such as the examples on this page. Pass a [`CompressedTensorsConfig`] to +> [`~PreTrainedModel.from_pretrained`] to load any checkpoint with compressed-tensors. + Pre-quantized models can be loaded directly. To quantize a model into the compressed-tensors format, see [llm-compressor](https://github.com/vllm-project/llm-compressor). Alternatively, models can be created independently and serialized with a compressed-tensors config. ```python diff --git a/docs/source/en/quantization/concept_guide.md b/docs/source/en/quantization/concept_guide.md index 47c3b037c867..c9e429bdd89b 100644 --- a/docs/source/en/quantization/concept_guide.md +++ b/docs/source/en/quantization/concept_guide.md @@ -110,7 +110,7 @@ FP8 is used in the *A8W8* quantization scheme, which quantizes both activations While int8 has broad support, efficient FP8 computation requires specific hardware capabilities found in newer GPUs like NVIDIA H100/H200/B100, AMD Instinct MI300 series, and Intel XPU. Without native hardware acceleration, the benefits of FP8 might not be fully realized. -Transformers supports FP8 through specific backends like [FBGEMM](./fbgemm_fp8), [FineGrainedFP8](./finegrained_fp8), and [compressed-tensors](./compressed_tensors). These backends handle the underlying FP8 conversion and computation when the appropriate hardware and configurations are used. +Transformers supports FP8 through specific backends like [FBGEMM](./fbgemm_fp8), [Block FP8](./blockfp8), [MXFP8](./mxfp8), and [compressed-tensors](./compressed_tensors). These backends handle the underlying FP8 conversion and computation when the appropriate hardware and configurations are used. ## Granularity diff --git a/docs/source/en/quantization/contribute.md b/docs/source/en/quantization/contribute.md index 190f7692231b..fd5d540470f9 100644 --- a/docs/source/en/quantization/contribute.md +++ b/docs/source/en/quantization/contribute.md @@ -46,7 +46,7 @@ Some quantization methods may require "pre-quantizing" the model through data ca ## Create new HFQuantizer class -0. The best starting point would be to have a look at another quantization method such as Finegrained Fp8. You will have to update or create three files in total: the [config file](https://github.com/huggingface/transformers/blob/main/src/transformers/utils/quantization_config.py), the [integration file](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/finegrained_fp8.py) and the [quantizer file](https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_finegrained_fp8.py). +0. The best starting point would be to have a look at another quantization method such as Block FP8. You will have to update or create three files in total: the [config file](https://github.com/huggingface/transformers/blob/main/src/transformers/utils/quantization_config.py), the [integration file](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/finegrained/layers.py) and the [quantizer file](https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_blockfp8.py). 1. Create a new quantization config class inside [src/transformers/utils/quantization_config.py](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/utils/quantization_config.py). Add the new quantization config to the [\_import_structure](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/__init__.py#L1088) inside Transformers' [src/transformers/\_\_init\_\_.py](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/__init__.py) file. diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md new file mode 100644 index 000000000000..43548e0e943d --- /dev/null +++ b/docs/source/en/quantization/finegrained.md @@ -0,0 +1,68 @@ + + +# Fine-grained + +[Block FP8](./blockfp8), [MXFP8](./mxfp8), [MXFP4](./mxfp4) and [NVFP4](./nvfp4) each have their own config and +quantize a whole model on their own. [`FineGrainedConfig`] mixes them in one model: its `dtype_plan` maps module name +patterns to the format config each matching module is quantized in. A typical plan holds the MoE experts, which carry +most of the weights, in a 4-bit format and every other linear in 8 bits. + +```py +from transformers import AutoModelForCausalLM, BlockFp8Config, FineGrainedConfig, Mxfp4Config + +quantization_config = FineGrainedConfig( + dtype_plan={ + r"\.experts($|\.)": Mxfp4Config(), + ".*": BlockFp8Config(), + }, + modules_to_not_convert=["lm_head"], +) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mixed") +``` + +The saved model records its plan, so it reloads with no config. + +> [!TIP] +> The requirements are those of the formats in the plan: a CUDA GPU with compute capability 8.9 or newer or an Intel +> XPU, and the [kernels](https://github.com/huggingface/kernels) package. + +## How a plan resolves + +- A pattern names the modules it matches with `re.search`. A targeted entry wins over the `".*"` catch-all, whatever + order the plan is written in, since a plan saved to `config.json` loses its order. Two targeted entries matching one + module raise an error. +- A module no entry matches, or that `modules_to_not_convert` names, stays in full precision. +- `modules_to_convert` names embedding tables to hold in FP8 with one per-tensor scale, whatever the plan's formats. +- `dequantize=True` loads every entry dequantized to full precision. + +## Checkpoints that mix formats + +These load as a [`FineGrainedConfig`] without one being passed: + +- DeepSeek-V4 ships block-FP8 weights with MXFP4 routed experts. Its checkpoint declares `quant_method="fp8"` and its + model config `expert_dtype="fp4"`, which together read as a plan of MXFP4 experts over block FP8 (DeepSeek-V4.1 + declares `expert_dtype` in its quantization config instead, and reads the same way). See the + [DeepSeek-V4](../model_doc/deepseek_v4) page for the plan written out. +- A ModelOpt export labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats + reads as a plan with one entry per group target. Other `quant_algo` values, such as `"MIXED_PRECISION"`, raise an + error. + +To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass +`FineGrainedConfig(dequantize=True)`. diff --git a/docs/source/en/quantization/finegrained_fp8.md b/docs/source/en/quantization/finegrained_fp8.md index e438d984c5c1..11ca965e841e 100644 --- a/docs/source/en/quantization/finegrained_fp8.md +++ b/docs/source/en/quantization/finegrained_fp8.md @@ -14,78 +14,20 @@ rendered properly in your Markdown viewer. --> -# Fine-grained FP8 +# Fine-grained FP8 (deprecated) > [!WARNING] -> `FineGrainedFP8Config` is frozen and will be removed in a future release; use [`BlockFp8Config`] or -> [`Mxfp8Config`]. A `FineGrainedFP8Config` passed to `from_pretrained` loads as the config it spells (MXFP8 for a -> `[1, 32]` block, block-FP8 otherwise, with DeepSeek-V4's MXFP4 experts), and existing checkpoints load unchanged. +> `FineGrainedFP8Config` is deprecated and will be removed in a future release. Existing code keeps working: a +> `FineGrainedFP8Config` passed to `from_pretrained` loads through its own frozen integration. A checkpoint declaring +> `quant_method="fp8"`, loaded without one, loads as [`BlockFp8Config`]. -Fine-grained FP8 quantization quantizes the weights and activations to fp8. +New code uses the config that describes the checkpoint's format: -- The weights are quantized to 8-bits for each 2D block (`weight_block_size=(128, 128)`). -- The activations are quantized to 8-bits for each group per token. The group value matches the weights in the input channel (128 by default). +| `FineGrainedFP8Config` | Use | Guide | +|---|---|---| +| `weight_block_size=[1, 32]` with `scale_fmt="ue8m0"` | [`Mxfp8Config`] | [MXFP8](./mxfp8) | +| any other `weight_block_size` | [`BlockFp8Config`] | [Block FP8](./blockfp8) | +| on a model whose config declares `expert_dtype="fp4"` (DeepSeek-V4) | [`BlockFp8Config`], which loads as a [`FineGrainedConfig`] of MXFP4 experts over block FP8 | [Fine-grained](./finegrained) | -FP8 quantization enables support for [DeepSeek-V3](https://hf.co/papers/2412.19437) and DeepSeek-R1. - -
- -
- -> [!TIP] -> You need a GPU with Compute Capability>=9 (H100), and install a PyTorch version compatible with the CUDA version of your GPU. - -Install Accelerate and upgrade to the latest version of PyTorch. - -```bash -pip install --upgrade accelerate torch -``` - -Create a [`FineGrainedFP8Config`] class and pass it to [`~PreTrainedModel.from_pretrained`] to quantize it. The weights are loaded in full precision (`torch.float32`) by default regardless of the actual data type the weights are stored in. Set `dtype="auto"` to load the weights in the data type defined in a models `config.json` file to automatically load the most memory-optimal data type. - -```py -from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer - -model_name = "meta-llama/Meta-Llama-3-8B" -quantization_config = FineGrainedFP8Config() -quantized_model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto", device_map="auto", quantization_config=quantization_config) - -tokenizer = AutoTokenizer.from_pretrained(model_name) -input_text = "What are we having for dinner?" -input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device.type) - -output = quantized_model.generate(**input_ids, max_new_tokens=10) -print(tokenizer.decode(output[0], skip_special_tokens=True)) -``` - -Use [`~PreTrainedModel.save_pretrained`] to save the quantized model and reload it with [`~PreTrainedModel.from_pretrained`]. - -```py -quant_path = "/path/to/save/quantized/model" -quantized_model.save_pretrained(quant_path) -model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto") -``` - -## DeepGEMM fast path - -On Hopper (SM90+) and Blackwell (SM100+) GPUs, every FP8 linear automatically dispatches to the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels from [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm) when `weight_block_size=(128, 128)` and `activation_scheme="dynamic"`. DeepGEMM is 3-6x faster than the Triton fallback. Install or upgrade the [kernels](https://github.com/huggingface/kernels) package to enable it. - -```bash -pip install -U kernels -``` - -DeepGEMM JIT-compiles its kernels with the system `nvcc`, so a full CUDA toolkit must be installed. A runtime-only install (CUDA libraries without `nvcc`) won't work. The minimum toolkit version depends on the hardware: 12.3 or later on Hopper (SM90), 12.9 or later on Blackwell (SM100). - -Transformers finds the toolkit by checking `CUDA_HOME`, then `CUDA_PATH`, then `nvcc` on your `PATH`, then `/usr/local/cuda`. Set `CUDA_HOME` to the toolkit root if `nvcc` isn't discovered. - -If the kernel cannot load (missing `kernels`, unsupported GPU, no CUDA toolkit, or an `nvcc` older than the required version), Transformers logs a warning once and falls back to the Triton finegrained-fp8 kernel. Static activation quantization always stays on the Triton path. - -To force the Triton fallback even when DeepGEMM is available, set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`. This only affects the FP8 linear dispatch and leaves the `"deepgemm"` experts backend untouched, which you switch with [`~PreTrainedModel.set_experts_implementation`]. - -For MoE experts, the DeepGEMM path is opt-in. Pass `experts_implementation="deepgemm"` (or `"deepgemm_megamoe"` on Blackwell) at load time to route the expert matmuls through DeepGEMM. See the [Experts backends](../experts_interface) guide for the full set of options. - -## UE8M0 scale format - -DeepSeek V4-style checkpoints store FP8 weight scales in the packed `float8_e8m0fnu` format instead of `float32`. These checkpoints are pre-quantized and set `scale_fmt="ue8m0"` in their quantization config. Both the DeepGEMM and Triton kernels read UE8M0 scales, so these checkpoints run on either path. - -On Blackwell (SM100+), the DeepGEMM experts kernels only supports UE8M0 scales. A checkpoint with plain `float32` scales (`scale_fmt="float"`) raises a `ValueError`. Use a `scale_fmt="ue8m0"` checkpoint, or run the experts with `grouped_mm` or `batched_mm`, which support `float32` scales directly. Hopper (SM90+) supports `float32` scales on the DeepGEMM path without conversion. See the [Experts backends](../experts_interface) guide for the experts backend options. +The fields keep their names and meanings, so `FineGrainedFP8Config(weight_block_size=(128, 128), +activation_scheme="static")` becomes `BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static")`. diff --git a/docs/source/en/quantization/mxfp4.md b/docs/source/en/quantization/mxfp4.md index f96833c0cc09..456b584ab50c 100644 --- a/docs/source/en/quantization/mxfp4.md +++ b/docs/source/en/quantization/mxfp4.md @@ -16,62 +16,70 @@ rendered properly in your Markdown viewer. # MXFP4 -Note: MXFP4 quantization currently only works for OpenAI GPT-OSS 120b and 20b. +MXFP4 is the OCP microscaling 4-bit format: weights in 4-bit floating point (E2M1, two values per byte) with one +power-of-two scale (UE8M0) per group of 32 values along a row. It cuts weight memory about 4x against bf16, so +GPT-OSS-120B fits on a single 80GB GPU and GPT-OSS-20B needs about 16GB. + +[`Mxfp4Config`] loads MXFP4 checkpoints, such as OpenAI's GPT-OSS, as they ship and quantizes full-precision +checkpoints on the fly. Checkpoints declare it as `quant_method="mxfp4"`. DeepSeek-V4 ships its routed experts in MXFP4 +beside block-FP8 weights, which loads through [Fine-grained](./finegrained). + +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. + +```bash +pip install --upgrade accelerate kernels +``` -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. +## Load a pre-quantized checkpoint -To use MXPF4, make sure your hardware meets the following requirements. +A checkpoint whose quantization config says `"quant_method": "mxfp4"` loads in MXFP4 with no config. -- Install Accelerate, kernels, and Triton ≥ 3.4. Only manually install Triton ≥ 3.4 if you're using PyTorch 2.7 because it is already supported in PyTorch 2.8. -- NVIDIA GPU Compute Capability ≥ 7.5 which includes Tesla GPUs and newer. Use [get_device_capability](https://docs.pytorch.org/docs/stable/generated/torch.cuda.get_device_capability.html) to check Compute Capability. +```py +from transformers import AutoModelForCausalLM, AutoTokenizer -```python -from torch import cuda -cuda.get_device_capability() +model_id = "openai/gpt-oss-20b" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") -# (7, 5) +tokenizer = AutoTokenizer.from_pretrained(model_id) +inputs = tokenizer("Once upon a time", return_tensors="pt").to(model.device) +output = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -Check a model's quantization config as shown below to see if it supports MXFP4. If `'quant_method': 'mxfp4'`, then the model automatically uses MXFP4. +Pass `Mxfp4Config(dequantize=True)` to load the weights dequantized to bf16 instead, for example to fine-tune them. + +## Quantize on the fly ```py -from transformers import GptOssConfig - -model_id = "openai/gpt-oss-120b" -cfg = GptOssConfig.from_pretrained(model_id) -print(cfg.quantization_config) - -# Example output: -# { -# 'modules_to_not_convert': [ -# 'model.layers.*.self_attn', -# 'model.layers.*.mlp.router', -# 'model.embed_tokens', -# 'lm_head' -# ], -# 'quant_method': 'mxfp4' -# } -``` +from transformers import AutoModelForCausalLM, Mxfp4Config -## MXFP4 kernels +quantization_config = Mxfp4Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mxfp4") +``` -Transformers automatically pulls the MXFP4-aware Triton kernels from the community repository when you load a model that needs them. The kernels are stored in your local cache and used during the forward pass. +## Activation formats -MXFP4 kernels are used by default, if available and supported, and does not require any code changes. +`activation_format` picks what the activations are quantized to before each matmul. -You can use [hf cache scan](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache#scan-your-cache) to verify the kernels are downloaded. +| `activation_format` | Scheme | Activations | +|---|---|---| +| `"bf16"` (default) | W4A16 (weight-only) | kept in full precision, as GPT-OSS runs | +| `"mxfp8"` | W4A8 | MXFP8 | +| `None` or `"mxfp4"` | W4A4 | MXFP4 | -```shell -hf cache scan -``` +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. -```shell -REPO ID REPO TYPE SIZE ON DISK --------------------------------- --------- ------------ -kernels-community/triton_kernels model 536.2K -openai/gpt-oss-20b model 13.8G -``` +The 4-bit activation scheme uses the 4-bit block-scaled MMA of Blackwell GPUs. ## Resources -Learn more about MXFP4 quantization and how blockwise scaling works in this [blog post](https://huggingface.co/blog/faster-transformers#mxfp4-quantization). +Learn more about MXFP4 and how blockwise scaling works in this +[blog post](https://huggingface.co/blog/faster-transformers#mxfp4-quantization). diff --git a/docs/source/en/quantization/mxfp8.md b/docs/source/en/quantization/mxfp8.md new file mode 100644 index 000000000000..7f61ae5486ed --- /dev/null +++ b/docs/source/en/quantization/mxfp8.md @@ -0,0 +1,72 @@ + + +# MXFP8 + +MXFP8 is the OCP microscaling FP8 format: weights in 8-bit floating point (`float8_e4m3fn`) with one power-of-two +scale (UE8M0) per group of 32 values along a row. [`Mxfp8Config`] loads MXFP8 checkpoints, such as `MiniMaxAI/MiniMax-M3-MXFP8`, as +they ship and quantizes full-precision checkpoints on the fly. Checkpoints declare it as `quant_method="mxfp8"`. + +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package. Without one, a pre-quantized checkpoint is dequantized to +> full precision as it loads. On Blackwell, the group scales feed the tensor cores' block-scaled MMA directly. + +```bash +pip install --upgrade accelerate kernels +``` + +## Load a pre-quantized checkpoint + +```py +from transformers import AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("MiniMaxAI/MiniMax-M3-MXFP8", dtype="auto", device_map="auto") +``` + +## Quantize on the fly + +```py +from transformers import AutoModelForCausalLM, Mxfp8Config + +quantization_config = Mxfp8Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mxfp8") +``` + +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"mxfp8"` | W8A8 | MXFP8 | +| `"mxfp4"` | W8A4 | MXFP4 | +| `"bf16"` | W8A16 (weight-only) | kept in full precision | + +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. + +Other options are shared with the other formats: `modules_to_not_convert` keeps modules in full precision, and +`dequantize=True` loads a pre-quantized checkpoint dequantized. + +## Related formats + +- [Block FP8](./blockfp8) holds FP8 weights with one scale per 2D block or per tensor. +- [MXFP4](./mxfp4) is the 4-bit microscaling format. +- [Fine-grained](./finegrained) mixes MXFP8 with other formats in one model. diff --git a/docs/source/en/quantization/nvfp4.md b/docs/source/en/quantization/nvfp4.md index 42a2d31cfa5a..7fd709255916 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -16,38 +16,35 @@ rendered properly in your Markdown viewer. # NVFP4 -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 -not replaced. +NVFP4 is NVIDIA's 4-bit floating-point format: weights in E2M1 (two values per byte) with two levels of scale, an FP8 +(E4M3) scale per group of 16 values and one FP32 global scale per tensor. [`NVFP4Config`] loads NVFP4 checkpoints as +they ship and quantizes full-precision checkpoints on the fly. Checkpoints declare it as `quant_method="nvfp4"`. > [!TIP] -> NVFP4 requires a Blackwell GPU with compute capability 10.0 or newer, a compatible CUDA-enabled PyTorch build, and -> the [kernels](https://github.com/huggingface/kernels) package. - -Install Accelerate and a compatible version of `kernels`. +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. The 4-bit activation scheme uses the NVFP4 +> MMA of Blackwell GPUs. ```bash pip install --upgrade accelerate kernels ``` -Pass [`NVFP4Config`] to [`~PreTrainedModel.from_pretrained`] with a single CUDA device. Weights are quantized as they -are loaded, so the source checkpoint should contain floating-point weights. - -```py -import torch +## Load a pre-quantized checkpoint -from transformers import AutoModelForCausalLM, AutoTokenizer, NVFP4Config +The NVFP4 exports of [NVIDIA ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) (`quant_method="modelopt"`, +`quant_algo="NVFP4"`) and of [compressed-tensors](./compressed_tensors) (`format="nvfp4-pack-quantized"`) load as +[`NVFP4Config`], with their calibrated activation scales, and save back in their own layout. A ModelOpt export +labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats loads as a +[`FineGrainedConfig`] (see [Fine-grained](./finegrained)); other `quant_algo` values, such as `"MIXED_PRECISION"`, raise +an error. +```py +from transformers import AutoModelForCausalLM, AutoTokenizer -model_id = "meta-llama/Llama-3.2-1B" -quantization_config = NVFP4Config() -model = AutoModelForCausalLM.from_pretrained( - model_id, - dtype=torch.bfloat16, - device_map="cuda", - quantization_config=quantization_config, -) +model_id = "nvidia/Llama-3.1-8B-Instruct-NVFP4" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_id) inputs = tokenizer("NVFP4 is", return_tensors="pt").to(model.device) @@ -55,19 +52,30 @@ output = model.generate(**inputs, max_new_tokens=20) print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -Use `modules_to_not_convert` to keep selected modules in their original precision. +Pass `NVFP4Config(dequantize=True)` to load the weights dequantized to bf16 instead. + +## Quantize on the fly ```py -quantization_config = NVFP4Config(modules_to_not_convert=["vision", "lm_head"]) +from transformers import AutoModelForCausalLM, NVFP4Config + +quantization_config = NVFP4Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "meta-llama/Llama-3.1-8B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("llama-3.1-8b-nvfp4") ``` -NVFP4 linear modules support `torch.compile`. The first compiled invocation includes graph compilation time, so warm up -the model before measuring generation throughput. +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"nvfp4"` | W4A4 | NVFP4, against the checkpoint's calibrated global scale when it ships one | +| `"bf16"` | W4A16 (weight-only) | kept in full precision | -## Current limitations +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. -- Only one CUDA device is supported. Tensor parallelism and multi-device `device_map` configurations are rejected until - the sharding behavior of the NVFP4 scale metadata is defined. -- CPU and disk offload are not supported. -- Pre-quantized NVFP4 checkpoints are not supported. -- NVFP4 models cannot currently be serialized with [`~PreTrainedModel.save_pretrained`] or trained. +`modules_to_not_convert` keeps the modules it names in full precision. diff --git a/docs/source/en/quantization/overview.md b/docs/source/en/quantization/overview.md index e833c334a768..055ade5aec0d 100644 --- a/docs/source/en/quantization/overview.md +++ b/docs/source/en/quantization/overview.md @@ -37,13 +37,16 @@ Use the Space below to help you pick a quantization method depending on your har | [HIGGS](./higgs) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 2/4 | 🔴 | 🟢 | 🟢 | https://github.com/HanGuo97/flute | | [HQQ](./hqq) | 🟢 | 🟢 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 1/8 | 🟢 | 🔴 | 🟢 | https://github.com/mobiusml/hqq/ | | [Metal](./metal) | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 🔴 | 🔴 | 2/4/8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/mlx-quantization-metal-kernels) | -| [NVFP4](./nvfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 4 | 🔴 | 🔴 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/nvfp4-gemm) | +| [Block FP8](./blockfp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [Fine-grained](./finegrained) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4/8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [MXFP4](./mxfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [MXFP8](./mxfp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [NVFP4](./nvfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | | [optimum-quanto](./quanto) | 🟢 | 🟢 | 🟢 | 🔴 | 🟢 | 🟢 | 🟢 | 2/4/8 | 🔴 | 🔴 | 🟢 | https://github.com/huggingface/optimum-quanto | | [SINQ](./sinq) | 🟢 | 🟢 | 🟢 | 🟡 | 🟡 | 🟡 | 🟡 | 2/3/4/6/8 | 🔴 | 🟢 | 🟢 | https://github.com/huawei-csl/SINQ | | [FBGEMM_FP8](./fbgemm_fp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🔴 | 8 | 🔴 | 🟢 | 🟢 | https://github.com/pytorch/FBGEMM | | [torchao](./torchao) | 🟢 | 🟢 | 🟢 | 🔴 | 🟡 | 🟢 | | 4/8 | | 🟢🔴 | 🟢 | https://github.com/pytorch/ao | | [VPTQ](./vptq) | 🔴 | 🔴 | 🟢 | 🟡 | 🔴 | 🔴 | 🟢 | 1/8 | 🔴 | 🟢 | 🟢 | https://github.com/microsoft/VPTQ | -| [FINEGRAINED_FP8](./finegrained_fp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🔴 | 8 | 🔴 | 🟢 | 🟢 | Built-in | | [SpQR](./spqr) | 🔴 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 3 | 🔴 | 🟢 | 🟢 | https://github.com/Vahe1994/SpQR/ | | [Quark](./quark) | 🔴 | 🟢 | 🟢 | 🟢 | 🟢 | 🟢 | ? | 2/4/6/8/9/16 | 🔴 | 🔴 | 🟢 | https://quark.docs.amd.com/latest/ | diff --git a/docs/source/en/quantization/quark.md b/docs/source/en/quantization/quark.md index 5688a6328c7e..1c5180fffa74 100644 --- a/docs/source/en/quantization/quark.md +++ b/docs/source/en/quantization/quark.md @@ -22,7 +22,7 @@ The PyTorch support integrated through 🤗 Transformers primarily targets AMD C Users interested in Quark can refer to its [documentation](https://quark.docs.amd.com/latest/) to get started quantizing models and using them in supported open-source libraries! -Although Quark has its own checkpoint / [configuration format](https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test/blob/main/config.json#L26), the library also supports producing models with a serialization layout compliant with other quantization/runtime implementations ([AutoAWQ](https://huggingface.co/docs/transformers/quantization/awq), [native fp8 in 🤗 Transformers](https://huggingface.co/docs/transformers/quantization/finegrained_fp8)). +Although Quark has its own checkpoint / [configuration format](https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test/blob/main/config.json#L26), the library also supports producing models with a serialization layout compliant with other quantization/runtime implementations ([AutoAWQ](https://huggingface.co/docs/transformers/quantization/awq), [native fp8 in 🤗 Transformers](https://huggingface.co/docs/transformers/quantization/blockfp8)). To be able to load Quark quantized models in Transformers, the library first needs to be installed: @@ -51,7 +51,7 @@ The table below shows a few features supported by Quark: Public models using Quark native serialization can be found at https://huggingface.co/models?other=quark. -Although Quark also supports [models using `quant_method="fp8"`](https://huggingface.co/models?other=fp8) and [models using `quant_method="awq"`](https://huggingface.co/models?other=awq), Transformers loads these models rather through [AutoAWQ](https://huggingface.co/docs/transformers/quantization/awq) or uses the [native fp8 support in 🤗 Transformers](https://huggingface.co/docs/transformers/quantization/finegrained_fp8). +Although Quark also supports [models using `quant_method="fp8"`](https://huggingface.co/models?other=fp8) and [models using `quant_method="awq"`](https://huggingface.co/models?other=awq), Transformers loads these models rather through [AutoAWQ](https://huggingface.co/docs/transformers/quantization/awq) or uses the [native fp8 support in 🤗 Transformers](https://huggingface.co/docs/transformers/quantization/blockfp8). ## Using Quark models in Transformers diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index ddc523d2c153..60b540f548c0 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -502,9 +502,12 @@ def _build_checkpoint_conversion_mapping(): ], ), ], - "deepseek_v4": [ - # the per-block weight scales ship as `.scale`; first, while `hc_*_scale` keys still carry no dot + # a pre-quantized checkpoint ships a weight's block scales as `.scale`; ahead of the `hc_*_scale` renames, which + # give the model's own `.scale` parameters their dot + "deepseek_v4_quantized": [ WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv"), + ], + "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 @@ -2049,6 +2052,10 @@ def get_model_conversion_mapping( if model_type is not None: conversions = get_checkpoint_conversion_mapping(model_type) + if hf_quantizer is not None and hf_quantizer.pre_quantized and model_type is not None: + if (quantized := get_checkpoint_conversion_mapping(f"{model_type}_quantized")) is not None: + conversions = quantized + (conversions or []) + if conversions is None: continue diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 57c6b554e9fb..e727c9dd55dc 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -932,6 +932,36 @@ def rename_source_key(self, source_key: str) -> tuple[str, str | None]: renamed_key = prefix_dot + renamed_key return renamed_key, source_pattern_that_matched + def quantize( + self, + collected_tensors: dict[str, Any], + full_layer_name: str, + model=None, + config=None, + hf_quantizer=None, + loading_info: LoadStateDictInfo | None = None, + ) -> dict[str, Any]: + """`quantization_operation` over the tensors under their full names, then the quantizer's ops for the + module they fill: those lay out what the quantization emits (a gate|up interleave would otherwise mix + the halves inside one quantization block), so a quantized target takes them here and not in its chain.""" + kwargs = { + "source_patterns": self.source_patterns, + "full_layer_name": full_layer_name, + "model": model, + "config": config, + "missing_keys": loading_info.missing_keys if loading_info else None, + } + context = (len(collected_tensors), full_layer_name) + with log_conversion_errors(full_layer_name, loading_info, context, self.quantization_operation): + collected_tensors = self.quantization_operation.convert( + collected_tensors, target_patterns=self.target_patterns, **kwargs + ) + for op in hf_quantizer.get_param_ops(model, full_layer_name): + with log_conversion_errors(full_layer_name, loading_info, context, op): + # every tensor is under its full name by now + collected_tensors = op.convert(collected_tensors, target_patterns=list(collected_tensors), **kwargs) + return collected_tensors + def reverse_transform(self) -> WeightTransform: """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations.""" # TODO: check this and relax when quantizer have `reverse_op` @@ -1014,18 +1044,7 @@ def convert( collected_tensors = {target_key: collected_tensors[self.source_patterns[0]]} if hf_quantizer is not None and self.quantization_operation is not None: - with log_conversion_errors( - layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation - ): - collected_tensors = self.quantization_operation.convert( - collected_tensors, - source_patterns=self.source_patterns, - target_patterns=self.target_patterns, - full_layer_name=target_key, - model=model, - config=config, - missing_keys=loading_info.missing_keys if loading_info else None, - ) + collected_tensors = self.quantize(collected_tensors, target_key, model, config, hf_quantizer, loading_info) return collected_tensors @@ -1218,18 +1237,7 @@ def convert( pass if hf_quantizer is not None and self.quantization_operation is not None: - with log_conversion_errors( - layer_name, loading_info, (len(collected_tensors), layer_name), self.quantization_operation - ): - collected_tensors = self.quantization_operation.convert( - collected_tensors, - source_patterns=self.source_patterns, - target_patterns=self.target_patterns, - full_layer_name=layer_name, - config=config, - model=model, - missing_keys=loading_info.missing_keys if loading_info else None, - ) + collected_tensors = self.quantize(collected_tensors, layer_name, model, config, hf_quantizer, loading_info) return collected_tensors @@ -1664,9 +1672,19 @@ def convert_and_load_state_dict_in_model( # 2. finally, collect the tensor into the proper converter if renamed_key in meta_model_state_dict: empty_param = meta_model_state_dict.get(renamed_key) - # each target key gets its own transform, with the ops the quantizer gives the module it fills + needs_quantization = ( + hf_quantizer + and not hf_quantizer.pre_quantized + and hf_quantizer.param_needs_quantization(model, renamed_key) + ) + # each target key gets its own transform, with the ops the quantizer gives the module it fills; a + # target quantized on load runs them after its quantization instead (`WeightTransform.quantize`) if renamed_key not in param_name_to_load: - param_ops = hf_quantizer.get_param_ops(model, renamed_key) if hf_quantizer is not None else [] + param_ops = ( + hf_quantizer.get_param_ops(model, renamed_key) + if hf_quantizer is not None and not needs_quantization + else [] + ) if source_pattern is not None: mapping = deepcopy(pattern_to_converter[source_pattern]) mapping.operations = [*mapping.operations, *param_ops] @@ -1680,11 +1698,6 @@ def convert_and_load_state_dict_in_model( source_pattern = original_key # 3. Handle dtype casting - needs_quantization = ( - hf_quantizer - and not hf_quantizer.pre_quantized - and hf_quantizer.param_needs_quantization(model, renamed_key) - ) if needs_quantization: mapping.quantization_operation = hf_quantizer.get_quantize_ops() diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index ac03823b3779..e00b2819028c 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -36,6 +36,7 @@ import torch +from ..quantizers.quantizers_utils import FUSED_STACKS from ..utils.import_utils import ( is_kernels_available, maybe_import_error, @@ -155,7 +156,7 @@ def is_deepgemm_loadable(raise_error: bool = False) -> bool: `@assume_constant_result` makes dynamo evaluate this once at trace time and inline the bool (its probe is untraceable — an `lru_cache`'d import check + the `get_device_capability` pybind + nvcc lookup), so - the check can sit in a compiled dispatch gate (e.g. `fp8_linear`) without breaking the graph. + the check can sit in a compiled dispatch gate (e.g. `prefers_deepgemm_linear`) without breaking the graph. FP4 / Mega MoE additionally need Blackwell (SM100); the relevant forwards enforce that via `_assert_sm100_requirements` / `is_sm100()`, since this env check can't see which dtype path the @@ -585,7 +586,7 @@ def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> Non ``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 + ``finegrained_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(): @@ -603,8 +604,8 @@ def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> Non "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')`." + "scales directly — the finegrained linear falls back to Triton automatically; for experts load with " + "`experts_implementation='grouped_mm'`." ) @@ -660,17 +661,37 @@ def _assert_affine_scales(module: torch.nn.Module) -> None: 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.""" + """The quantized DeepGEMM arms take the stacked ``[gate; up]`` rows: Mega MoE's + ``transform_weights_for_mega_moe`` does its own gate/up interleave, and the grouped GEMM reads a + block-scale grid row-major, whose blocks an interleaved stack splits between the halves. 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 " + "DeepGEMM experts need gate|up stacked, but this module was loaded interleaved 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_block_scales_only(module: torch.nn.Module) -> None: + """The quantized DeepGEMM arms read the weights and their UE8M0/float32 block scales, and quantize the + activations to FP8 themselves. An expert bias, NVFP4's second-level globals and E4M3 block scales, or + weight-only (bf16) activations have nowhere to go, so they are refused rather than dropped.""" + scales = (getattr(module, f"{proj}_scale_inv", None) for proj in FUSED_STACKS) + if getattr(module, "has_bias", False): + raise NotImplementedError("DeepGEMM quantized experts do not apply expert biases; use 'grouped_mm'.") + if getattr(module, "global_scale_dtype", None) is not None or any( + scale is not None and scale.dtype == torch.float8_e4m3fn for scale in scales + ): + raise NotImplementedError("DeepGEMM serves no NVFP4 experts (E4M3 block scales, globals); use 'grouped_mm'.") + if getattr(module, "activation_format", None) == "bf16": + raise NotImplementedError( + "DeepGEMM quantizes the activations, so it cannot serve weight-only experts " + "(activation_format='bf16'); use 'grouped_mm'." + ) + + 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 @@ -707,6 +728,7 @@ def deepgemm_experts_guards( sm100: bool = False, supports_post_expert_norm: bool = True, stacked_gate_up: bool = False, + block_scales_only: bool = False, ): """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`` @@ -725,6 +747,8 @@ def guarded(self, hidden_states, *args, **kwargs): _assert_affine_scales(self) if stacked_gate_up: _assert_stacked_gate_up(self) + if block_scales_only: + _assert_block_scales_only(self) if sm100: _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) return forward(self, hidden_states, *args, **kwargs) @@ -835,7 +859,7 @@ def deepgemm_bf16_experts_forward( ) -@deepgemm_experts_guards(affine_scales=True, sm100=True) +@deepgemm_experts_guards(affine_scales=True, sm100=True, stacked_gate_up=True, block_scales_only=True) def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -930,7 +954,7 @@ def deepgemm_fp8_fp4_experts_forward( def setup_megamoe_weights(module: torch.nn.Module) -> None: - """One-shot pack + permute of an FP8Experts module's L1/L2 weights into the + """One-shot pack + permute of a `FineGrainedExperts` module's L1/L2 weights into the Mega MoE UTCCP layout. Called lazily on the first megamoe forward; idempotent via the caller's ``_megamoe_transformed`` flag. @@ -986,7 +1010,9 @@ 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, supports_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, block_scales_only=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 index ffb0093164d2..08160a809afd 100644 --- a/src/transformers/integrations/finegrained/__init__.py +++ b/src/transformers/integrations/finegrained/__init__.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. -"""The fine-grained quantization integration: `core` holds the modules and their forwards, +"""The fine-grained quantization integration: `layers` 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.""" diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 9270291906b3..294ac394bb90 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -25,7 +25,7 @@ import torch from ...core_model_loading import ConversionOps, Interleave -from ...quantizers.quantizers_utils import get_module_from_name +from ...quantizers.quantizers_utils import FUSED_STACKS, get_module_from_name from .layers import ( _FP8_DTYPE, _FP8_MAX, @@ -40,6 +40,16 @@ ) +def single_tensor(value): + """A converter's tensor, which the loader hands in a list: one source's is a one-element list. Several are + sources this op does not merge, so they are refused rather than cut down to the first.""" + if not isinstance(value, list): + return value + if len(value) != 1: + raise ValueError(f"expected a single tensor, got a list of {len(value)}") + return value[0] + + 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 @@ -47,7 +57,7 @@ def keyed_by_target(input_dict: dict, target_patterns) -> dict: 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 {target_patterns[0]: single_tensor(value)} return input_dict @@ -110,22 +120,24 @@ def reverse_op(self) -> ConversionOps: 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.""" + """Expert gate|up rows between the checkpoint's order and the one the experts hold: stacked + ``[gate; up]`` or the kernels' ``[g0, u0, g1, u1, ...]`` — the core ``Interleave`` along dim 1, + over the weight, scale grid and bias alike. Runs on load and save wherever the two differ.""" 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 = finegrained_experts(model, full_layer_name) - if experts is None or not (experts.is_concatenated and experts.holds_interleaved_gate_up): + # `is_concatenated` is the checkpoint's stacked order + if experts is None or not experts.has_gate or experts.is_concatenated != experts.holds_interleaved_gate_up: return input_dict - interleave = Interleave(dim=1, inverse=not self.inverse) # inverse=True is stacked -> interleaved + # inverse=True is stacked -> interleaved + interleave = Interleave(dim=1, inverse=experts.holds_interleaved_gate_up != self.inverse) out = {} for key, value in input_dict.items(): - probe = value[0] if isinstance(value, list) and len(value) == 1 else value + value = single_tensor(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: + if getattr(value, "ndim", 0) < 2 or value.shape[1] < 2: out[key] = value else: out[key] = interleave.convert({key: value}, None, [key])[key] @@ -141,7 +153,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(): - value = value[0] if isinstance(value, list) else value + value = single_tensor(value) if self.inverse: # the container its module recorded at load; e8m0 is only ever a scale, so the dtype # alone picks the tensors to restore @@ -167,7 +179,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(): - value = value[0] if isinstance(value, list) else value + value = single_tensor(value) out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) return out @@ -195,7 +207,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(): - value = value[0] if isinstance(value, list) else value + value = single_tensor(value) if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) out[key] = value elif self.inverse: @@ -214,7 +226,7 @@ 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 + value = single_tensor(value) out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value return out @@ -279,9 +291,10 @@ def convert(self, input_dict, target_patterns=None, model=None, **kwargs): # 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].chunk(2, dim=1) gate, up = gate.squeeze(1), up.squeeze(1) + if not torch.equal(gate, up): + self._assert_halves_foldable(model) globals_[stack] = gate ratio = (up / gate)[:, None] for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): @@ -290,17 +303,23 @@ def convert(self, input_dict, target_patterns=None, model=None, **kwargs): 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() - ): + def _assert_halves_foldable(model) -> None: + """Merging the halves' globals moves the up half's past the GLU. A gate|up bias is added AFTER the global, so + its up rows would have to be divided by the same ratio, in the bias converter the layout ops own; no + checkpoint pairs the two. And the GLU must be linear in `up`: a clamped SwiGLU (`swiglu_limit`) clips the + rescaled values, and SwiGLU-OAI (`swiglu_alpha`) multiplies by `up + 1`.""" + experts = [module for module in model.modules() if isinstance(module, FineGrainedExperts)] if model else [] + if any(module.has_bias for module in experts): 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." ) + if any(module.swiglu_limit is not None or module.swiglu_alpha is not None for module in experts): + raise NotImplementedError( + "this checkpoint calibrates the gate|up halves separately, and this model's GLU " + "(`swiglu_limit` / `swiglu_alpha`) is not linear in the up half, so the halves' globals " + "cannot be merged into one per expert." + ) class FineGrainedInputScales(_FineGrainedOp): @@ -311,7 +330,9 @@ class FineGrainedInputScales(_FineGrainedOp): 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 + # unwrapped, a fused `(E,)` scale in the loader's list reads as `(1, E)` and reduces across experts + if isinstance(value, list): + value = value[0] if len(value) == 1 else torch.stack(value, dim=0) 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 the reduce is over absolute values @@ -330,8 +351,7 @@ class FineGrainedInputScalesSplit(_FineGrainedOp): re-expands per expert, the shape the checkpoint holds.""" 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 + value = single_tensor(next(iter(input_dict.values()))) if value.numel() == 1: experts = finegrained_experts(model, full_layer_name) if experts is not None: @@ -346,7 +366,9 @@ 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 that module's layout. + scale (and the NVFP4 global) in its held dtype. The weight is read as its full stack, so a block + covers rows of one projection; the loader then runs the module's layout ops (gate|up interleave, + scale swizzle) on what this emits. 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 @@ -356,7 +378,7 @@ class FineGrainedQuantize(_FineGrainedOp): 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 + tensor = single_tensor(value) 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 @@ -372,7 +394,7 @@ def _as_expert_rows(module, name: str, value: torch.Tensor) -> torch.Tensor: """An expert stack in the ``(E, rows, K)`` the module's slot and the kernels take. A model may store its experts transposed (GPT-OSS: ``(E, H, 2I)``), which the module it replaced declares: quantizing before the swap packs the wrong axis, and a square stack (GPT-OSS's down projection) cannot tell from its shape.""" - if module is None or value.ndim != 3 or name not in ("gate_up_proj", "up_proj", "down_proj"): + if module is None or value.ndim != 3 or name not in FUSED_STACKS: return value return value.transpose(1, 2).contiguous() if getattr(module, "ships_transposed", False) else value @@ -386,7 +408,7 @@ def _weight_holder(model, key: str): return None if name == "weight": return module, "weight_scale_inv" - if name in ("gate_up_proj", "up_proj", "down_proj"): + if name in FUSED_STACKS: return module, f"{name}_scale_inv" return None @@ -431,8 +453,6 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc ) weight, scale, global_scale = self._quantize_group(value, format_spec) 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: # the experts name their global `_weight_global_scale`, a dense linear just @@ -547,12 +567,26 @@ def _reciprocal_globals(self) -> bool: plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} return any(config.checkpoint_format == "compressed-tensors" for config in plan.values()) + # the weights a converter hands one at a time: a dense `weight`, or an experts module's fused stack + _SINGLE_WEIGHTS = ("weight", *FUSED_STACKS) + + @staticmethod + def _name(pattern: str) -> str: + """The name a source pattern matches: `(? `gate_up_proj`.""" + return pattern.removeprefix(r"(? str | None: + """The source pattern in `input_dict` matching the same name as `pattern`, if the converter collected one.""" + return next((key for key in input_dict if self._name(key) == self._name(pattern)), None) + @staticmethod def _global_pattern_for(weight_pattern: str) -> str: - """NVFP4's second-level global beside `weight_pattern`: `x.weight` -> `x.weight_global_scale`.""" + """NVFP4's second-level global beside `weight_pattern`: `x.weight` -> `x.weight_global_scale`, and a fused + stack's `gate_up_proj` -> `gate_up_proj_weight_global_scale`.""" anchored = weight_pattern.endswith("$") base = weight_pattern[:-1] if anchored else weight_pattern - return base + "_global_scale" + ("$" if anchored else "") + suffix = "_global_scale" if base.endswith("weight") else "_weight_global_scale" + return base + suffix + ("$" if anchored else "") # 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) @@ -583,11 +617,8 @@ def _dequantize_one( 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 + # a per-tensor scale has no grid: one block over the whole matrix + scale_rows, scale_cols = scales.shape[-2:] if scales.ndim >= 2 else (1, 1) # 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 block the plan names reproduces it plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} @@ -628,6 +659,10 @@ def _dequantize_one( out = (blocked * per_block).reshape(*original_shape[:-2], rows + pad_m, cols + pad_n)[..., :rows, :cols] if global_scale is not None: global_scale = global_scale.to(torch.float32) + if out.ndim > 2 and global_scale.ndim: + # a stack's globals, one per expert (E,) or one per stacked gate|up half (E, 2), over its rows + global_scale = global_scale.reshape(global_scale.shape[0], -1) + global_scale = global_scale.repeat_interleave(rows // global_scale.shape[1], dim=1)[..., None] out = out / global_scale if self._reciprocal_globals else out * global_scale return out.to(output_dtype) @@ -646,22 +681,24 @@ def convert( **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", "weight_global_scale", ...] -> weight``) - # hands one weight, with or without a scale (RMSNorm weights match ``weight$`` too). - if "weight$" in input_dict: + # The dense converter (``["weight$", companions...] -> weight``) and a fused stack's (``["gate_up_proj$", its + # companions...] -> gate_up_proj``) each hand one weight, with or without a scale (RMSNorm weights match + # ``weight$`` too). + weight_pattern = next( + (p for p in input_dict if p.endswith("$") and self._name(p) in self._SINGLE_WEIGHTS), None + ) + if weight_pattern is not None: # 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 - global_scale = input_dict.get("weight_global_scale") - global_scale = global_scale[0] if isinstance(global_scale, list) else global_scale - dequantized = self._dequantize_one(quantized, scales, output_dtype, global_scale) - return {target_key: dequantized} - return {target_key: quantized} + # (direct invocation) key it by the weight's own name + target_key = full_layer_name if full_layer_name is not None else self._name(weight_pattern) + quantized = single_tensor(input_dict[weight_pattern]) + scale, global_scale = ( + single_tensor(input_dict.get(self._companion(input_dict, companion(weight_pattern)))) + for companion in (self._scale_pattern_for, self._global_pattern_for) + ) + if scale is None: + return {target_key: quantized} + return {target_key: self._dequantize_one(quantized, scale, output_dtype, global_scale)} # Generic chain path: dequantize every weight pattern that has a sibling scale. consumed = { @@ -692,6 +729,6 @@ def convert( @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 + # a dequantized model saves in full precision: none of its weights fills a finegrained module, + # so the quantize passes each through return FineGrainedQuantize(self.hf_quantizer) diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index 0cf01d79722d..9874b044efcb 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -13,13 +13,14 @@ # limitations under the License. """Quantized modules and forwards for the block-scaled formats — 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`. +moved onto these modules by the ops in `conversions`. """ from __future__ import annotations import functools import os +import sys from collections.abc import Callable from dataclasses import dataclass @@ -27,6 +28,7 @@ import torch.nn as nn from ...activations import ACT2FN +from ...quantizers.quantizers_utils import FUSED_STACKS, STACK_COMPANIONS from ...utils import logging from ...utils.import_utils import is_kernels_available from ...utils.quantization_config import BlockFp8Config, BlockScaledConfig @@ -115,7 +117,7 @@ class FineGrained: @torch._dynamo.allow_in_graph def _load_finegrained_kernel() -> None: """ - Load the finegrained-fp8 Triton kernel once into the `_FINEGRAINED` module global. + Load the finegrained-kernels 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 @@ -133,7 +135,7 @@ def _load_finegrained_kernel() -> None: return if not is_kernels_available(): - raise ImportError(f"finegrained-fp8 kernel unavailable: {_MISSING_KERNELS_MESSAGE}") + raise ImportError(f"finegrained-kernels kernel unavailable: {_MISSING_KERNELS_MESSAGE}") kernel = lazy_load_kernel("finegrained-kernels") if kernel is None: @@ -253,6 +255,10 @@ def _alloc_expert_proj( ) +# the experts backends that read gate|up stacked and block scales affine +_DEEPGEMM_EXPERTS_IMPLEMENTATIONS = ("deepgemm", "deepgemm_megamoe") + + 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, @@ -260,7 +266,7 @@ def _holds_swizzled_scales(config, format_spec: WeightFormat, activation_format: 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") + getattr(config, "_experts_implementation", None) not in _DEEPGEMM_EXPERTS_IMPLEMENTATIONS and format_spec.scale_group is not None and activation_format != "bf16" and is_sm100() @@ -496,7 +502,7 @@ def __init__( self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) 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) + self.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) @classmethod def from_module( @@ -551,7 +557,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: input_shape = x.shape[:-2] hidden_dim = x.shape[-1] - w = self.weight.view(self.n_groups, -1, hidden_dim) + # the stored K, which packed FP4 halves + w = self.weight.view(self.n_groups, -1, self.weight.shape[-1]) 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)) @@ -560,14 +567,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # (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 + two_level = self.weight_global_scale is not None kernel = load_finegrained_kernel() y = kernel.matmul_grouped( x, w, - None, + None if two_level else self.activation_scale, scale_inv, expert_start=expert_start, activation_format=self.activation_format, + a_global_scale=self.activation_scale if two_level else None, b_global_scale=self.weight_global_scale, ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) @@ -576,8 +585,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y +def _assert_interleaved_gate_up(module: nn.Module) -> None: + """The kernels' fused chains read gate|up interleaved. A module loaded for a DeepGEMM backend holds it + stacked; switched to a triton backend after load, the GLU would pair the wrong rows.""" + if module.has_gate and not module.holds_interleaved_gate_up: + raise NotImplementedError( + "The triton experts backends need gate|up interleaved, but this module was loaded stacked for a " + "DeepGEMM backend. Switching backends after load is not supported — pass `experts_implementation=...` " + "to `from_pretrained` so the model loads for it." + ) + + 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.""" + _assert_interleaved_gate_up(self) kernel = load_finegrained_kernel() return kernel.moe_fused_batched(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) @@ -585,6 +606,7 @@ def finegrained_batched_mm_experts_forward(self, hidden_states, top_k_index, top 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).""" + _assert_interleaved_gate_up(self) kernel = load_finegrained_kernel() return kernel.moe_fused_grouped(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) @@ -593,11 +615,15 @@ class FineGrainedExperts(_FineGrainedModule, nn.Module): # 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"}, + "deepgemm_megamoe": { + "moe_tp_experts": "megamoe_experts", + "ep_dispatch_experts": "megamoe_experts", + "ep_router": "megamoe_router", + }, } # a projection's tensors, by the suffix they carry on the module: the weight, its scale grid, the NVFP4 # weight global, the calibrated activation scale, a bias - _projection_slots = ("", "_scale_inv", "_weight_global_scale", "_activation_scale", "_bias") + _projection_slots = ("", *STACK_COMPANIONS, "_bias") def __init__( self, @@ -627,14 +653,16 @@ def __init__( self.global_scale_dtype = format_spec.global_scale_dtype # The layouts: what the CHECKPOINT ships (`is_concatenated`: stacked ``[gate; up]`` or, like - # GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``) against what the module now + # GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``) against what the module # 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. + # disagree. Frozen at load rather than derived on read, since it describes the bytes. The + # triton backends and the eager loop share the interleaved rows; the DeepGEMM ones read + # gate|up stacked (Mega MoE packs it itself, the grouped GEMM reads a block grid row-major). self.has_gate = has_gate self.has_bias = has_bias self.is_concatenated = is_concatenated impl = getattr(config, "_experts_implementation", None) - self.holds_interleaved_gate_up = self.has_gate and impl != "deepgemm_megamoe" + self.holds_interleaved_gate_up = self.has_gate and impl not in _DEEPGEMM_EXPERTS_IMPLEMENTATIONS swizzled = _holds_swizzled_scales(config, format_spec, activation_format) 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 @@ -664,9 +692,7 @@ def from_module( "has_bias": getattr(module, "has_bias", False), "is_concatenated": getattr(module, "is_concatenated", True), } - experts_class = use_experts_implementation( - experts_class=cls, experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, **flags - ) + experts_class = _experts_class_with(cls, **flags) config = getattr(module, "config", model.config.get_text_config()) experts = experts_class(config=config, quantization_config=quantization_config, **flags) experts.post_expert_norm, experts.has_post_expert_norm = module.post_expert_norm, module.has_post_expert_norm @@ -701,17 +727,18 @@ def _register_projection(self, proj, rows, in_dim, min_scale_out, storage): _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 - 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). + # interleaved rows give alternating output columns, the split the fused epilogue does; stacked rows two halves + if self.holds_interleaved_gate_up: + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] + else: + gate, up = gate_up.chunk(2, dim=-1) + if self.swiglu_limit is not None: gate = gate.clamp(max=self.swiglu_limit) up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + if self.swiglu_alpha is not None: + # SwiGLU-OAI gate (same math as the model's non-quantized experts) glu = gate * torch.sigmoid(gate * self.swiglu_alpha) return (up + 1.0) * glu - 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( @@ -772,18 +799,35 @@ def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torc if input_global is not None and input_global.numel() > 1: input_global = input_global[expert_idx : expert_idx + 1] - return finegrained_linear( - input, - weight, - scale, - self.block_size, - bias=bias, + linear = functools.partial( + finegrained_linear, + block_size=self.block_size, 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, ) + # A scale spanning several interleaved rows belongs to one half (the loader interleaves the grid per + # block), so a row-major read would pair half the rows with the other half's scale: each half runs on its + # own rows and blocks, its output columns interleaved back. + if ( + proj == self.gate_up_name + and self.holds_interleaved_gate_up + and scale.ndim == 2 + and scale.shape[0] < weight.shape[0] + ): + halves = [ + linear( + input, + weight[half::2].contiguous(), + scale[half::2].contiguous(), + bias=bias[half::2] if bias is not None else None, + ) + for half in (0, 1) + ] + return torch.stack(halves, dim=-1).flatten(-2) + return linear(input, weight, scale, bias=bias) def _moe_operands(self, kernel) -> dict: """Everything the kernels' MoE forwards take from the module: the two projections with their @@ -840,8 +884,25 @@ def _moe_operands(self, kernel) -> dict: } +@functools.cache +def _experts_class_with(cls, *, has_gate: bool, has_bias: bool, is_concatenated: bool) -> type: + """`cls` with the experts dispatch and these flags: `use_experts_implementation` wraps the class it is handed in + place, so each flag set wraps a subclass of its own, once, and `cls` itself stays as defined. The subclass is + registered on `cls`'s module under its own name, so a pickled model finds it.""" + qualname = f"{cls.__qualname__}_gate{int(has_gate)}_bias{int(has_bias)}_concatenated{int(is_concatenated)}" + subclass = type(cls.__name__, (cls,), {"__module__": cls.__module__, "__qualname__": qualname}) + setattr(sys.modules[cls.__module__], qualname, subclass) + return use_experts_implementation( + experts_class=subclass, + experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + has_gate=has_gate, + has_bias=has_bias, + is_concatenated=is_concatenated, + ) + + class FineGrainedExpertsInterface(ExpertsInterface): - """Interface for registering custom FP8 experts forward functions.""" + """Interface for registering custom finegrained experts forward functions.""" _global_mapping = { "deepgemm": deepgemm_fp8_fp4_experts_forward, @@ -866,7 +927,7 @@ def assert_modules_are_quantized(model: nn.Module) -> None: for name, module in model.named_modules(): if not isinstance(module, _FineGrainedModule): continue - for attr in ("weight", "gate_up_proj", "up_proj", "down_proj"): + for attr in ("weight", *FUSED_STACKS): weight = getattr(module, attr, None) if weight is not None and weight.element_size() > 1: raise ValueError( diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index d6ef42c90b89..3df674489103 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -12,6 +12,8 @@ # 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 copy +import json import warnings from ..models.auto.configuration_auto import AutoConfig @@ -142,11 +144,17 @@ "mxfp4": Mxfp4Config, "nvfp4": NVFP4Config, "finegrained": FineGrainedConfig, + # a modelopt export reads through `read_modelopt` (see `AutoQuantizationConfig.from_dict`) as its format's config + # or a `FineGrainedConfig`; listed so the method counts as supported "modelopt": FineGrainedConfig, "sinq": SinqConfig, "gemma": GemmaQuantizationConfig, } +# a passed config of these types reads a pre-quantized checkpoint's config as its own; the deprecated +# `FineGrainedFP8Config` so keeps loading through its frozen integration +CHECKPOINT_READING_CONFIG_TYPES = (AutoRoundConfig, CompressedTensorsConfig, FineGrainedFP8Config) + LOADING_ATTRIBUTES_CONFIG_TYPES = ( GPTQConfig, AwqConfig, @@ -156,6 +164,7 @@ MetalConfig, FineGrainedConfig, BlockScaledConfig, + FineGrainedFP8Config, ) logger = logging.get_logger(__name__) @@ -221,7 +230,11 @@ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwar # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = AutoQuantizationConfig.from_dict(quantization_config) - quantization_config = _from_legacy_fp8(quantization_config) + # the deprecated config keeps its own, frozen integration + if isinstance(quantization_config, FineGrainedFP8Config): + from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + + return FineGrainedFP8HfQuantizer(quantization_config, **kwargs) quant_method = quantization_config.quant_method @@ -270,7 +283,7 @@ def merge_quantization_configs( if isinstance(quantization_config, dict): # Convert the config based on the type of quantization_config_from_args (e.g., AutoRoundConfig), which takes priority before automatic configuration dispatch. - if isinstance(quantization_config_from_args, (AutoRoundConfig, CompressedTensorsConfig)): + if isinstance(quantization_config_from_args, CHECKPOINT_READING_CONFIG_TYPES): quantization_config = type(quantization_config_from_args).from_dict(quantization_config) else: quantization_config = AutoQuantizationConfig.from_dict(quantization_config) @@ -284,22 +297,33 @@ def merge_quantization_configs( "Please make sure to pass the same quantization config class to `from_pretrained` with different loading attributes." ) + ignored = [] if isinstance(quantization_config, LOADING_ATTRIBUTES_CONFIG_TYPES) and isinstance( quantization_config_from_args, LOADING_ATTRIBUTES_CONFIG_TYPES ): loading_attr_dict = quantization_config_from_args.get_loading_attributes() + if isinstance(quantization_config_from_args, (BlockScaledConfig, FineGrainedConfig)): + ignored = _ignored_options(quantization_config, quantization_config_from_args, loading_attr_dict) for attr, val in loading_attr_dict.items(): + if attr == "dtype_plan": + val = _with_activation_formats(quantization_config.dtype_plan, val) + elif attr == "activation_format": + _check_calibrated(quantization_config, val) setattr(quantization_config, attr, val) 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 ignored: + warning_msg += f" Ignored: {ignored}." - if warning_msg != "" and not isinstance( - quantization_config, (MetalConfig, FineGrainedConfig, BlockScaledConfig) - ): + # the block-scaled formats take a passed config for its loading attributes, so that alone is not worth a warning + lends_only = ( + isinstance(quantization_config, (MetalConfig, FineGrainedConfig, BlockScaledConfig, FineGrainedFP8Config)) + and not ignored + ) + if warning_msg != "" and not lends_only: warnings.warn(warning_msg) else: - # in the case of mxfp4, we don't want to print the warning message, bit confusing for users logger.info(warning_msg) return quantization_config @@ -356,43 +380,103 @@ def register_quantizer_fn(cls): return register_quantizer_fn -def _from_legacy_fp8(quantization_config): - """A deprecated `FineGrainedFP8Config` as the config it spells: MXFP8 for its `[1, 32]` block with UE8M0 scales, - block-FP8 otherwise (its FP4 experts come from the model config, see `finegrained_quantization_config`). Any other - config unchanged.""" - if not isinstance(quantization_config, FineGrainedFP8Config): - return quantization_config - options = { - "activation_scheme": quantization_config.activation_scheme, - "modules_to_not_convert": quantization_config.modules_to_not_convert, - "dequantize": quantization_config.dequantize, - } - block_size = quantization_config.weight_block_size - if block_size is not None and tuple(block_size) == (1, 32): - return Mxfp8Config(**options) - return BlockFp8Config( - weight_block_size=block_size, - scale_fmt=quantization_config.scale_fmt, - modules_to_convert=quantization_config.modules_to_convert, - **options, +def _as_json(config) -> dict: + """`config` as `config.json` holds it, so a tuple and the list it serializes to compare equal.""" + return json.loads(config.to_json_string(use_diff=False)) + + +def _set_options(config) -> dict: + """The options `config` sets away from its class's defaults.""" + defaults = _as_json(type(config)()) + return {name: value for name, value in _as_json(config).items() if defaults.get(name) != value} + + +def _ignored_options(quantization_config, quantization_config_from_args, loading_attributes: dict) -> list[str]: + """The options a config passed for a pre-quantized checkpoint sets away from its defaults, differently from the + checkpoint's own, and that are not loading attributes, so the checkpoint's stand; a passed plan's, per entry.""" + own, passed = _as_json(quantization_config), _set_options(quantization_config_from_args) + ignored = [ + name + for name, value in passed.items() + if name not in loading_attributes and name != "quant_method" and own.get(name) != value + ] + for pattern, entry in loading_attributes.get("dtype_plan", {}).items(): + own_entry = quantization_config.dtype_plan.get(pattern) + if type(own_entry) is not type(entry): + ignored.append(f"dtype_plan[{pattern!r}]") + continue + own_options = _as_json(own_entry) + ignored += [ + f"dtype_plan[{pattern!r}].{name}" + for name, value in _set_options(entry).items() + if name != "activation_format" and own_options.get(name) != value + ] + return ignored + + +def _with_activation_formats(dtype_plan: dict, passed_plan: dict) -> dict: + """`dtype_plan` with the activation format of each `passed_plan` entry naming the same format under the same + pattern: the kernels the same weights run on.""" + merged = dict(dtype_plan) + for pattern, entry in passed_plan.items(): + if type(merged.get(pattern)) is type(entry) and entry._names_activation_format: + _check_calibrated(merged[pattern], entry.activation_format) + merged[pattern] = copy.copy(merged[pattern]) + merged[pattern].activation_format = entry.activation_format + return merged + + +def _check_calibrated(checkpoint_config, activation_format): + """NVFP4 activations quantize against the checkpoint's calibrated activation global, which a weight-only + NVFP4 checkpoint does not ship: such a checkpoint can only run weight-only.""" + if ( + isinstance(checkpoint_config, NVFP4Config) + and checkpoint_config.activation_format == "bf16" + and activation_format != "bf16" + ): + raise ValueError( + "This NVFP4 checkpoint is weight-only: it ships no calibrated activation globals (`input_scale`), which " + f"activation_format={activation_format!r} quantizes against. Load it with activation_format='bf16'." + ) + + +def _with_mxfp4_experts(block_fp8_config): + """Block-FP8 over W4A4 MXFP4 experts: what a checkpoint declaring `expert_dtype="fp4"` ships. The loading options + are the plan's.""" + dense = BlockFp8Config( + weight_block_size=block_fp8_config.weight_block_size, + scale_fmt=block_fp8_config.scale_fmt, + activation_scheme=block_fp8_config.activation_scheme, + activation_format=block_fp8_config.activation_format, + checkpoint_format=block_fp8_config.checkpoint_format, + ) + return FineGrainedConfig( + dtype_plan={EXPERTS_PATTERN: Mxfp4Config(activation_format=None), ".*": dense}, + modules_to_not_convert=block_fp8_config.modules_to_not_convert, + modules_to_convert=block_fp8_config.modules_to_convert, + dequantize=block_fp8_config.dequantize, + ) + + +def _declared_expert_dtype(model_config): + """`expert_dtype` where a checkpoint declares it: on its model config (DeepSeek-V4), or in the quantization config + it ships (DeepSeek-V4.1), read from the raw dict since the parsed config has no such field.""" + text_config = model_config.get_text_config() + shipped = getattr(model_config, "quantization_config", None) or getattr(text_config, "quantization_config", None) + return getattr(text_config, "expert_dtype", None) or ( + shipped.get("expert_dtype") if isinstance(shipped, dict) else None ) def finegrained_quantization_config(quantization_config, model_config): - """The config a checkpoint really describes, with its model config in hand: a deprecated `FineGrainedFP8Config` as - the config it spells, and DeepSeek-V4, which ships `quant_method="fp8"` and declares MXFP4 experts on the MODEL - config (`expert_dtype`), as block-FP8 over W4A4 MXFP4 experts. Any other config unchanged.""" - quantization_config = _from_legacy_fp8(quantization_config) + """The config a checkpoint really describes, with its model config in hand: block-FP8 whose checkpoint declares + FP4 experts (`expert_dtype`, DeepSeek-V4 and V4.1) as block-FP8 over W4A4 MXFP4 experts. Any other config + unchanged.""" if not isinstance(quantization_config, BlockFp8Config): return quantization_config - if getattr(model_config.get_text_config(), "expert_dtype", None) != "fp4": + if _declared_expert_dtype(model_config) != "fp4": return quantization_config - return FineGrainedConfig( - dtype_plan={EXPERTS_PATTERN: Mxfp4Config(activation_format="mxfp4"), ".*": quantization_config}, - modules_to_not_convert=quantization_config.modules_to_not_convert, - modules_to_convert=quantization_config.modules_to_convert, - dequantize=quantization_config.dequantize, - ) + return _with_mxfp4_experts(quantization_config) def get_hf_quantizer(config, quantization_config, device_map, weights_only, user_agent, gguf_file=None): @@ -415,11 +499,12 @@ def get_hf_quantizer(config, quantization_config, device_map, weights_only, user pre_quantized = False if pre_quantized or quantization_config is not None: - # both read as what they describe for this model before they are compared + # read as what it describes for this model, the model config's `expert_dtype` included: on the fly too, as + # the checkpoint it writes declares it quantization_config = finegrained_quantization_config(quantization_config, config) if pre_quantized: - # a passed `AutoRoundConfig` or `CompressedTensorsConfig` reads the checkpoint as its own (the merge does) - routed = isinstance(quantization_config, (AutoRoundConfig, CompressedTensorsConfig)) + # a passed config of `CHECKPOINT_READING_CONFIG_TYPES` reads the checkpoint as its own (the merge does) + routed = isinstance(quantization_config, CHECKPOINT_READING_CONFIG_TYPES) if isinstance(quantization_params_from_config, dict) and not routed: quantization_params_from_config = AutoQuantizationConfig.from_dict(quantization_params_from_config) config.quantization_config = AutoHfQuantizer.merge_quantization_configs( diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index 7c1e4def094b..52a39367aaa1 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -28,6 +28,8 @@ ) from .base import HfQuantizer from .quantizers_utils import ( + FUSED_STACKS, + STACK_COMPANIONS, embedding_pattern, get_module_from_name, plan_entry_for, @@ -57,8 +59,15 @@ class BlockScaledQuantizer(HfQuantizer): dtype_plan = dict.fromkeys(("_scale_inv", "activation_scale", "_global_scale")) # checkpoint keys this format's producers ship beside the weights that no module holds unexpected_keys: tuple[str, ...] = () - # narrowed once the modules are built + # narrowed once the modules are built: whether this format holds experts, and the up projection they stack + # (`gate_up_proj` gated, `up_proj` not) holds_experts = True + up_projections: tuple[str, ...] = ("gate_up_proj",) + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + # the embedding tables the plan quantizes, by name suffix: each holds its own per-tensor `weight_scale` + self.embeddings = list(getattr(quantization_config, "modules_to_convert", None) or []) def plan(self) -> dict[str, "BlockScaledConfig"]: """`{module regex: format config}` for the modules this quantizer covers.""" @@ -94,7 +103,7 @@ def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: module, name = get_module_from_name(model, param_name) if isinstance(module, FineGrainedExperts) and re.fullmatch( - r"(gate_up_proj|down_proj)(_scale_inv|_bias)?", name + r"(gate_up_proj|up_proj|down_proj)(_scale_inv|_bias)?", name ): ops = [FineGrainedInterleaveGateUp(self)] if name.startswith("gate_up_proj") else [] if not name.endswith("_bias"): @@ -104,38 +113,54 @@ def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: return [FineGrainedScaleContainer(self, module=module)] return [] + def finish_modules(self, model): + """What this format does to its loaded modules beyond what every format does.""" + def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import assert_modules_are_quantized + from ..integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device + for member in self.members(): + member.finish_modules(model) + disable_deepgemm_on_multi_device(model) assert_modules_are_quantized(model) return model - def _activation_scale_conversions(self, projections=("gate_up_proj", "up_proj", "down_proj"), ops=()) -> list: + @staticmethod + def _per_expert_sources(proj: str, slot: str) -> list[str]: + """The per-expert checkpoint keys of `slot` that stack into the experts' `proj`, under any parent + (`mlp.experts`, `mixer.experts`, `feed_forward.experts`): the gate|up stack reads both halves.""" + halves = ("gate_proj", "up_proj") if proj == "gate_up_proj" else (proj,) + return [f"experts.*.{half}.{slot}" for half in halves] + + def _block_scale_rename(self): + """A producer's `weight_scale` block scale onto the `weight_scale_inv` slot every module holds it in; the + embedding tables the plan quantizes hold their per-tensor scale under that name, and keep it.""" + from ..core_model_loading import WeightRenaming + + kept = "".join(rf"(? list: """A calibrated checkpoint's `input_scale` onto the `activation_scale` its module holds: renamed for every - module first, then one value per expert of each of `projections`, the gate|up pair reduced to the one both - halves read (`up_proj` is an ungated stack's); `ops` run on each value before that reduction. Both expert layouts - are covered, a checkpoint matching one while the other never fires; neither when this format holds no - experts.""" + module first, then one value per expert of each of `projections` the experts hold, the gate|up pair reduced + to the one both halves read; `ops` run on each value before that reduction. Both expert layouts are covered, a + checkpoint matching one while the other never fires; neither when this format holds no experts.""" from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming from ..integrations.finegrained.conversions import FineGrainedInputScales rename = WeightRenaming(source_patterns=r"^(.+)\.input_scale$", target_patterns=r"\1.activation_scale") if not self.holds_experts: return [rename] - per_expert = { - "gate_up_proj": ["mlp.experts.*.gate_proj.activation_scale", "mlp.experts.*.up_proj.activation_scale"], - "down_proj": ["mlp.experts.*.down_proj.activation_scale"], - } + held = [proj for proj in projections if proj in (*self.up_projections, "down_proj")] return [ rename, *( WeightConverter( - source_patterns=per_expert[proj], - target_patterns=f"mlp.experts.{proj}_activation_scale", + source_patterns=self._per_expert_sources(proj, "activation_scale"), + target_patterns=f"experts.{proj}_activation_scale", operations=[MergeModulelist(dim=0), *ops, FineGrainedInputScales(self)], ) - for proj in projections - if proj in per_expert + for proj in held ), *( WeightConverter( @@ -143,7 +168,7 @@ def _activation_scale_conversions(self, projections=("gate_up_proj", "up_proj", target_patterns=f"experts.{proj}_activation_scale", operations=[*ops, FineGrainedInputScales(self)], ) - for proj in projections + for proj in held ), ] @@ -214,7 +239,7 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** 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 + # one byte per held element: an FP8 value, or two FP4 values packed into a weight already half as wide return 1 return super().param_element_size(model, param_name, param) @@ -265,8 +290,10 @@ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwarg experts = [module for module in model.modules() if isinstance(module, FineGrainedExperts)] ignored = set(model._keys_to_ignore_on_load_unexpected or []) for member in self.members(): - # a format's expert converters exist only where its experts do - member.holds_experts = any(module.quantization_config is member.quantization_config for module in experts) + # a format's expert converters exist only where its experts do, for the stacks they hold + held = [module for module in experts if module.quantization_config is member.quantization_config] + member.holds_experts = bool(held) + member.up_projections = tuple(sorted({module.gate_up_name for module in held})) ignored.update(member.unexpected_keys) # with no module holding an activation scale (weight-only, dynamic activations), a calibrated # `input_scale` is one this run has no use for, not an unexpected key: a later run may read it @@ -297,10 +324,22 @@ def replace_modules(self, model, plan=None, modules_to_not_convert=None) -> bool plan = self.plan() if plan is None else plan return replace_with_quantized_modules(model, plan, modules_to_not_convert, module_classes) + def _experts_hold_scale_grid(self, glob: str, config) -> bool: + """Whether the experts a TP glob (`layers.*.mlp.experts`) names hold a scale grid, which splits with its + weight. The glob's one style serves every layer, so its layers must agree.""" + layers = range(getattr(config, "num_hidden_layers", None) or 1) if "*" in glob else (0,) + entries = (plan_entry_for(self.plan(), glob.replace("*", str(layer))) for layer in layers) + grids = {_has_scale_grid(entry) for entry in entries if entry is not None} + if len(grids) > 1: + raise ValueError( + f"Tensor parallelism shards every `{glob}` alike, but this plan gives some of them a block-scale " + "grid and others a per-tensor scale. Give those experts one format." + ) + return True in grids + 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. + # Per-impl rewrite of the experts parallel-layer kind, over the plan the model declares. 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) @@ -318,11 +357,10 @@ def update_tp_plan(self, config): 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"): + if projection not in FUSED_STACKS: continue if style == "grouped_gemm": - companions = ["_bias", "_scale_inv", "_weight_global_scale", "_activation_scale"] - for suffix in companions: + for suffix in ("_bias", *STACK_COMPANIONS): updated_plan.setdefault(f"{key}{suffix}", style) continue @@ -336,8 +374,7 @@ def update_tp_plan(self, config): # 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)`. - entry = plan_entry_for(self.plan(), key.rpartition(".")[0]) - blocked = entry is not None and _has_scale_grid(entry) + blocked = self._experts_hold_scale_grid(key.rpartition(".")[0], plan_owner) if projection == "down_proj" and style in ("rowwise", "packed_rowwise"): if blocked: updated_plan.setdefault(f"{key}_scale_inv", "moe_experts_rowwise") @@ -381,15 +418,25 @@ def get_weight_conversions(self): conversions, claimed = [], {} if dequantize: # every quantized key folded back to a full-precision weight from the `weight` + `weight_scale_inv` pair, - # times NVFP4's `weight_global_scale` where there is one; anchored `weight$` so the scale keys land in the - # scale slots, and the activation scales collected only to be dropped + # times NVFP4's `weight_global_scale` where there is one, and the activation scales collected only to be + # dropped; the companions start a name (`(? str: @@ -448,14 +497,15 @@ def _with_scale_grid(pattern: str) -> str: def _rebuilt(conv, source_patterns: list[str], operations: list | None = None): """`conv` over `source_patterns`, its targets taking only a parameter and its scale grid on save (a target is - a source pattern there, and a parameter's name prefixes all its companions'), and its scope kept.""" + a source pattern there, and a parameter's name prefixes all its companions'), and its scope and placement kept.""" from ..core_model_loading import WeightConverter - params = (".weight", "gate_up_proj", "up_proj", "down_proj") + params = (".weight", *FUSED_STACKS) rebuilt = WeightConverter( source_patterns=source_patterns, target_patterns=[_with_scale_grid(p) if p.endswith(params) else p for p in conv._original_target_patterns], operations=conv.operations if operations is None else operations, + force_cpu=conv.force_cpu, ) rebuilt.scope_prefix, rebuilt.base_model_prefix = conv.scope_prefix, conv.base_model_prefix return rebuilt diff --git a/src/transformers/quantizers/quantizer_blockfp8.py b/src/transformers/quantizers/quantizer_blockfp8.py index 5c950d880b97..d13746c40caa 100644 --- a/src/transformers/quantizers/quantizer_blockfp8.py +++ b/src/transformers/quantizers/quantizer_blockfp8.py @@ -33,21 +33,11 @@ def plan(self): embeddings = self.quantization_config.modules_to_convert or [] return {**dict.fromkeys(map(embedding_pattern, embeddings), self.quantization_config), **super().plan()} - def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import disable_deepgemm_on_multi_device - - disable_deepgemm_on_multi_device(model) - return super()._process_model_after_weight_loading(model, **kwargs) - def format_conversions(self, dequantize): - from ..core_model_loading import WeightRenaming - conversions = [] if self.quantization_config.checkpoint_format == "compressed-tensors": # its block scale ships as `weight_scale`, and rides its weight's converter once renamed - conversions.append( - WeightRenaming(source_patterns=r"^(.+)\.weight_scale$", target_patterns=r"\1.weight_scale_inv") - ) + conversions.append(self._block_scale_rename()) if not dequantize and self.quantization_config.activation_scheme == "static": conversions += self._activation_scale_conversions() return conversions diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index dd5e8285c195..4b7d3dfee206 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -41,12 +41,14 @@ def __init__(self, quantization_config: "FineGrainedConfig", **kwargs): # the embedding tables `modules_to_convert` names hold FP8 rows with one per-tensor scale self._embedding_config = BlockFp8Config(weight_block_size=None) configs = {id(config): config for config in self.plan().values()}.values() - self._members = [AUTO_QUANTIZER_MAPPING[config.method.value](config) for config in configs] + self._members = [AUTO_QUANTIZER_MAPPING[config.method.value](config, **kwargs) for config in configs] + # the plan's embedding tables are the host's to name + for member in self._members: + member.embeddings = self.embeddings def plan(self): - embeddings = self.quantization_config.modules_to_convert or [] return { - **dict.fromkeys(map(embedding_pattern, embeddings), self._embedding_config), + **dict.fromkeys(map(embedding_pattern, self.embeddings), self._embedding_config), **self.quantization_config.dtype_plan, } @@ -56,8 +58,3 @@ def members(self): def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: member = self.member_of(get_module_from_name(model, param_name)[0]) return member.get_param_ops(model, param_name) if member is not None else [] - - def _process_model_after_weight_loading(self, model, **kwargs): - for member in self.members(): - member._process_model_after_weight_loading(model, **kwargs) - return model diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 7b5004804ed4..bb893834637d 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -8,7 +8,7 @@ warnings.warn( "`transformers.quantizers.quantizer_finegrained_fp8` is frozen and will be removed in a future release; " - "use `BlockFp8Config` or `Mxfp8Config`, which a passed `FineGrainedFP8Config` loads as.", + "a passed `FineGrainedFP8Config` still loads through it. Use `BlockFp8Config` or `Mxfp8Config`.", DeprecationWarning, stacklevel=2, ) diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 99e32040bc17..87c7196f5e65 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -32,12 +32,6 @@ class Mxfp4HfQuantizer(BlockScaledQuantizer): quantization_config: "Mxfp4Config" - def _process_model_after_weight_loading(self, model, **kwargs): - from ..integrations.finegrained import disable_deepgemm_on_multi_device - - disable_deepgemm_on_multi_device(model) - return super()._process_model_after_weight_loading(model, **kwargs) - def format_conversions(self, dequantize): return self._dequantize_conversions() if dequantize else self._mxfp4_conversions() diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index 752b8fbf6486..a7f355eb80c5 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING from .block_scaled_utils import BlockScaledQuantizer -from .quantizers_utils import get_module_from_name +from .quantizers_utils import FUSED_STACKS, get_module_from_name if TYPE_CHECKING: @@ -31,11 +31,10 @@ class NVFP4HfQuantizer(BlockScaledQuantizer): # modelopt calibrates an FP8 KV cache beside the weights; the cache here stays in the model dtype unexpected_keys = (r"\.(k_proj\.k|v_proj\.v)_scale$",) - def _process_model_after_weight_loading(self, model, **kwargs): + def finish_modules(self, model): from ..integrations.finegrained import register_shared_gate_up_input_globals register_shared_gate_up_input_globals(model) - return super()._process_model_after_weight_loading(model, **kwargs) def format_conversions(self, dequantize): # weight-only keeps activations bf16, so there is no activation global for `input_scale` to land on; nor @@ -44,8 +43,16 @@ def format_conversions(self, dequantize): conversions = self._nvfp4_renames(calibrated) if calibrated: # the down projection's is read with its globals - conversions += self._activation_scale_conversions(projections=("gate_up_proj",), ops=self._reciprocal_ops) - if self.holds_experts and not dequantize: + conversions += self._activation_scale_conversions( + projections=self.up_projections, ops=self._reciprocal_ops + ) + if dequantize: + # a stacked layout's scales reach its stack's dequantize converter under the slot names it collects; the + # experts are left as they are, so every stack a model may hold is covered + conversions += self._nvfp4_fused_renames( + FUSED_STACKS, (("_weight_scale", "_scale_inv"), ("_weight_scale_2", "_weight_global_scale")) + ) + elif self.holds_experts: conversions += self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) return conversions @@ -57,7 +64,7 @@ def get_param_ops(self, model, param_name): ops = super().get_param_ops(model, param_name) module, name = get_module_from_name(model, param_name) - if name in ("weight", "gate_up_proj", "up_proj", "down_proj"): + if name in ("weight", *FUSED_STACKS): ops = [FineGrainedViewPackedInt8(self), *ops] # an experts module's globals are inverted in their converters, before the per-expert reductions elif isinstance(module, FineGrainedLinear) and name in ("weight_global_scale", "activation_scale"): @@ -80,7 +87,7 @@ def _nvfp4_renames(self, calibrated: bool): from ..core_model_loading import WeightRenaming renames = [ - WeightRenaming(source_patterns=r"^(.+)\.weight_scale$", target_patterns=r"\1.weight_scale_inv"), + self._block_scale_rename(), WeightRenaming(source_patterns=r"^(.+)\.weight_scale_2$", target_patterns=r"\1.weight_global_scale"), ] if self.quantization_config.checkpoint_format == "compressed-tensors": @@ -105,53 +112,59 @@ def _nvfp4_per_expert_conversions(self, calibrated: bool): from ..core_model_loading import MergeModulelist, WeightConverter from ..integrations.finegrained.conversions import FineGrainedWeightGlobals - converters = [ + return [ # 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=[ - "mlp.experts.*.gate_proj.weight_global_scale", - "mlp.experts.*.up_proj.weight_global_scale", - "mlp.experts.*.down_proj.weight_global_scale", - *(["mlp.experts.*.down_proj.activation_scale"] if calibrated else []), + *self._per_expert_sources(up, "weight_global_scale"), + *self._per_expert_sources("down_proj", "weight_global_scale"), + *(self._per_expert_sources("down_proj", "activation_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_activation_scale"] if calibrated else []), + f"experts.{up}_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_activation_scale"] if calibrated else []), ], operations=[MergeModulelist(dim=0), *self._reciprocal_ops, FineGrainedWeightGlobals(self)], - ), + ) + for up in self.up_projections + ] + + @staticmethod + def _nvfp4_fused_renames(projections, slots) -> list: + """The stacked layout's `{proj}{shipped}` scales onto the `{proj}{slot}` slots, for each `(shipped, slot)`.""" + from ..core_model_loading import WeightRenaming + + return [ + WeightRenaming(source_patterns=rf"experts\.{proj}{shipped}$", target_patterns=f"experts.{proj}{slot}") + for proj in projections + for shipped, slot in slots ] - return converters def _nvfp4_fused_conversions(self, calibrated: bool): """Already stacked per layer, as the vLLM fused layout ships them — so no `MergeModulelist`: the scales only take their slot's name, and the module's ops give every tensor its layout.""" - from ..core_model_loading import WeightConverter, WeightRenaming + from ..core_model_loading import WeightConverter from ..integrations.finegrained.conversions import FineGrainedWeightGlobals - converters = [ - WeightRenaming( - source_patterns=rf"experts\.{proj}_weight_scale$", target_patterns=f"experts.{proj}_scale_inv" - ) - for proj in ("gate_up_proj", "down_proj") - ] - converters.append( + converters = self._nvfp4_fused_renames((*self.up_projections, "down_proj"), (("_weight_scale", "_scale_inv"),)) + converters += [ # 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$", + rf"experts\.{up}_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", + f"experts.{up}_weight_global_scale", "experts.down_proj_weight_global_scale", *(["experts.down_proj_activation_scale"] if calibrated else []), ], operations=[*self._reciprocal_ops, FineGrainedWeightGlobals(self)], ) - ) + for up in self.up_projections + ] return converters diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 7b95f484f8d6..0236837b408b 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -16,6 +16,11 @@ from typing import Any +# an experts module's fused projection stacks, and what each holds beside it +FUSED_STACKS = ("gate_up_proj", "up_proj", "down_proj") +STACK_COMPANIONS = ("_scale_inv", "_weight_global_scale", "_activation_scale") + + def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: if "." in tensor_name: module_name, tensor_name = tensor_name.rsplit(".", 1) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index dc67a505c6a5..a546b006be73 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -19,6 +19,7 @@ import json import os import re +import warnings from dataclasses import dataclass from enum import Enum from typing import Any, Optional, Union @@ -1703,21 +1704,40 @@ def subtree_patterns(globs: list[str]) -> list[str]: return patterns +# the block and scale MXFP8 and MXFP4 fix +_MX_FORMAT_OPTIONS = {"weight_block_size": (1, 32), "scale_fmt": "ue8m0"} + + +class _FormatDefault: + """The `activation_format` a format config takes when none is passed.""" + + def __repr__(self): + return "the format's default" + + +_FORMAT_DEFAULT = _FormatDefault() + + +def _loading_attributes(config, **attributes) -> dict[str, Any]: + """What a config passed to `from_pretrained` lends a pre-quantized checkpoint's own config: `dequantize`, the + skip-list when it sets one (left unset, the checkpoint's stays), and `attributes`.""" + loading = {"dequantize": config.dequantize} + if config.modules_to_not_convert is not None: + loading["modules_to_not_convert"] = config.modules_to_not_convert + return {**loading, **attributes} + + class BlockScaledConfig(QuantizationConfigMixin): """What every block-scaled weight format's config holds: the options of the format itself, which the quantized modules read, and the loading options. One such config quantizes a whole model; a [`FineGrainedConfig`] maps - several onto parts of one. + several onto parts of one. It is built as one of its formats: [`BlockFp8Config`], [`Mxfp8Config`], + [`Mxfp4Config`] or [`NVFP4Config`]. Args: activation_format (`str`, *optional*): - What activations are quantized to: `None` for the weights' own format, `"bf16"` for weight-only, or - another format the weights' kernels serve. - activation_scheme (`str`, *optional*, defaults to `"dynamic"`): - `"dynamic"` (inline per-token/per-block) or `"static"` (a calibrated per-tensor scale from the checkpoint). - weight_block_size (`tuple[int, int]`, *optional*): - The block of the weight scale grid, for formats whose granularity is not fixed by the format. - scale_fmt (`str`, *optional*, defaults to `"float"`): - Storage dtype of block scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`). + What activations are quantized to: `None` for the weights' own format (which its name spells too), + `"bf16"` for weight-only, or another format the weights' kernels serve. On a pre-quantized checkpoint, a + passed config's activation format replaces the checkpoint's: it picks the kernels the same weights run on. dequantize (`bool`, *optional*, defaults to `False`): Whether to dequantize the model to full precision while loading. modules_to_not_convert (`list`, *optional*): @@ -1727,58 +1747,86 @@ class BlockScaledConfig(QuantizationConfigMixin): `"compressed-tensors"`. """ - # what `quant_method` this config is, and the activation formats its weights' kernels serve + # what `quant_method` this config is, the activation formats its weights' kernels serve, and the one it takes method: QuantizationMethod activation_formats: tuple[str | None, ...] = (None,) + default_activation_format: str | None = None + # fixed by the MX and NVFP4 formats; block-FP8 sets them + weight_block_size: tuple[int, int] | None = None + scale_fmt: str = "float" + activation_scheme: str = "dynamic" + # the values those fields take in this format, which a checkpoint's `config.json` may spell out + format_options: dict[str, Any] = {} def __init__( self, - activation_format: str | None = None, - activation_scheme: str = "dynamic", - weight_block_size: tuple[int, int] | None = None, - scale_fmt: str = "float", + *, + activation_format: str | None = _FORMAT_DEFAULT, dequantize: bool = False, modules_to_not_convert: list | None = None, checkpoint_format: str | None = None, **kwargs, ): + if not hasattr(self, "method"): + raise TypeError( + f"{type(self).__name__} is what the block-scaled format configs share; build one of them: " + f"{', '.join(config.__name__ for config in BLOCK_SCALED_CONFIGS.values())}" + ) self.quant_method = self.method - self.checkpoint_format = checkpoint_format kwargs.pop("quant_method", None) if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") - self.activation_format = activation_format - self.activation_scheme = activation_scheme - self.weight_block_size = tuple(weight_block_size) if weight_block_size is not None else None - self.scale_fmt = scale_fmt + # a `config.json` written before the format fixed these carries their defaults + fixed = { + name: kwargs.pop(name) + for name in ("weight_block_size", "scale_fmt", "activation_scheme") + if name in kwargs + } + changed = { + name: value + for name, value in fixed.items() + # a block size read back from `config.json` is a list + if (tuple(value) if isinstance(value, list) else value) + not in (getattr(BlockScaledConfig, name), self.format_options.get(name)) + } + if changed: + raise ValueError( + f"{self.method.value} weights fix {sorted(changed)} by their format; got {changed}. Only " + "`BlockFp8Config` sets them." + ) + if kwargs: + logger.info(f"Unused kwargs: {list(kwargs)}. These kwargs are not used in {type(self).__name__}.") + self.checkpoint_format = checkpoint_format + # a passed config lends its activation format to a pre-quantized checkpoint only where it names one + self._names_activation_format = activation_format is not _FORMAT_DEFAULT + self.activation_format = ( + self.default_activation_format if activation_format is _FORMAT_DEFAULT else activation_format + ) self.dequantize = dequantize self.modules_to_not_convert = modules_to_not_convert self.post_init() def post_init(self): - self.activation_scheme = self.activation_scheme.lower() if self.activation_format not in self.activation_formats: raise ValueError( f"{self.method.value} weights take activation_format in {self.activation_formats}; " f"got {self.activation_format!r}" ) - if self.activation_scheme not in ("dynamic", "static"): - raise ValueError(f"activation_scheme must be 'dynamic' or 'static'; got {self.activation_scheme!r}") - if self.scale_fmt not in ("float", "ue8m0"): - raise ValueError(f"scale_fmt must be 'float' or 'ue8m0'; got {self.scale_fmt!r}") + # one spelling for the weights' own activation format + if self.activation_format == self.method.value: + self.activation_format = None if self.checkpoint_format not in (None, "compressed-tensors"): raise ValueError(f"checkpoint_format must be None or 'compressed-tensors'; got {self.checkpoint_format!r}") - if self.weight_block_size is not None and ( - len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0 - ): - raise ValueError(f"weight_block_size must be two positive integers; got {self.weight_block_size}") def get_loading_attributes(self): - # a skip-list left unset keeps the checkpoint's - attributes = {"dequantize": self.dequantize} - if self.modules_to_not_convert is not None: - attributes["modules_to_not_convert"] = self.modules_to_not_convert - return attributes + named = {"activation_format": self.activation_format} if self._names_activation_format else {} + return _loading_attributes(self, **named) + + def to_dict(self) -> dict[str, Any]: + return {name: value for name, value in super().to_dict().items() if not name.startswith("_")} + + def __iter__(self): + yield from self.to_dict().items() class BlockFp8Config(BlockScaledConfig): @@ -1791,7 +1839,8 @@ class BlockFp8Config(BlockScaledConfig): scale_fmt (`str`, *optional*, defaults to `"float"`): Storage dtype of the scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`, DeepSeek-V4). activation_scheme (`str`, *optional*, defaults to `"dynamic"`): - `"dynamic"` or `"static"` (a calibrated per-tensor activation scale from the checkpoint). + `"dynamic"` (quantized inline, per 128-wide group) or `"static"` (a calibrated per-tensor activation + scale from the checkpoint). modules_to_convert (`list`, *optional*): Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. kwargs (`dict[str, Any]`, *optional*): @@ -1803,16 +1852,30 @@ class BlockFp8Config(BlockScaledConfig): def __init__( self, + *, weight_block_size: tuple[int, int] | None = (128, 128), scale_fmt: str = "float", activation_scheme: str = "dynamic", modules_to_convert: list | None = None, **kwargs, ): + self.weight_block_size = tuple(weight_block_size) if weight_block_size is not None else None + self.scale_fmt = scale_fmt + self.activation_scheme = activation_scheme self.modules_to_convert = modules_to_convert - super().__init__( - weight_block_size=weight_block_size, scale_fmt=scale_fmt, activation_scheme=activation_scheme, **kwargs - ) + super().__init__(**kwargs) + + def post_init(self): + super().post_init() + self.activation_scheme = self.activation_scheme.lower() + if self.activation_scheme not in ("dynamic", "static"): + raise ValueError(f"activation_scheme must be 'dynamic' or 'static'; got {self.activation_scheme!r}") + 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.weight_block_size is not None and ( + len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0 + ): + raise ValueError(f"weight_block_size must be two positive integers; got {self.weight_block_size}") class Mxfp8Config(BlockScaledConfig): @@ -1820,15 +1883,16 @@ class Mxfp8Config(BlockScaledConfig): Args: activation_format (`str`, *optional*): - `None` for MXFP8 activations, `"mxfp4"` for W8A4, or `"bf16"` for weight-only. + `None` (or `"mxfp8"`) for MXFP8 activations, `"mxfp4"` for W8A4, or `"bf16"` for weight-only. kwargs (`dict[str, Any]`, *optional*): The options every block-scaled format takes, see [`BlockScaledConfig`]. """ method = QuantizationMethod.MXFP8 activation_formats = (None, "mxfp8", "mxfp4", "bf16") + format_options = _MX_FORMAT_OPTIONS - def __init__(self, activation_format: str | None = None, **kwargs): + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): super().__init__(activation_format=activation_format, **kwargs) @@ -1838,31 +1902,18 @@ class Mxfp4Config(BlockScaledConfig): Args: activation_format (`str`, *optional*, defaults to `"bf16"`): - `"bf16"` for weight-only (W4A16, what GPT-OSS runs), `"mxfp8"` for W4A8, or `"mxfp4"` (or `None`) for W4A4. - modules_to_not_convert (`list`, *optional*): - Modules to keep in full precision. - dequantize (`bool`, *optional*, defaults to `False`): - Whether to dequantize the model to full precision while loading. + `"bf16"` for weight-only (W4A16, what GPT-OSS runs), `"mxfp8"` for W4A8, or `None` (or `"mxfp4"`) for W4A4. kwargs (`dict[str, Any]`, *optional*): The options every block-scaled format takes, see [`BlockScaledConfig`]. """ method = QuantizationMethod.MXFP4 activation_formats = (None, "mxfp8", "mxfp4", "bf16") + default_activation_format = "bf16" + format_options = _MX_FORMAT_OPTIONS - def __init__( - self, - modules_to_not_convert: list | None = None, - dequantize: bool = False, - activation_format: str | None = "bf16", - **kwargs, - ): - super().__init__( - activation_format=activation_format, - dequantize=dequantize, - modules_to_not_convert=modules_to_not_convert, - **kwargs, - ) + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) class NVFP4Config(BlockScaledConfig): @@ -1870,22 +1921,19 @@ class NVFP4Config(BlockScaledConfig): exports them; quantizes a full-precision checkpoint on the fly as well. Args: - modules_to_not_convert (`list[str]`, *optional*): - Modules to keep in full precision. activation_format (`str`, *optional*): - `None` for NVFP4 activations (W4A4, with the checkpoint's calibrated activation global) or `"bf16"` for - weight-only. + `None` (or `"nvfp4"`) for NVFP4 activations (W4A4, with the checkpoint's calibrated activation global) or + `"bf16"` for weight-only. kwargs (`dict[str, Any]`, *optional*): The options every block-scaled format takes, see [`BlockScaledConfig`]. """ method = QuantizationMethod.NVFP4 activation_formats = (None, "nvfp4", "bf16") + format_options = {"weight_block_size": (1, 16)} - def __init__( - self, modules_to_not_convert: list[str] | None = None, activation_format: str | None = None, **kwargs - ): - super().__init__(activation_format=activation_format, modules_to_not_convert=modules_to_not_convert, **kwargs) + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) # the block-scaled formats by their `quant_method`, for a plan entry read back from `config.json` @@ -1901,7 +1949,8 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: 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. + i.e. every linear), and an `re:` prefix introduces a pattern. Activations are read the same way, and only where + the weights' kernels quantize to them. """ formats = { (4, "float", 16): "nvfp4", @@ -1914,6 +1963,11 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: # one FP8 scale per output channel has no block the kernels tile if quant_method is None or weights.get("strategy") == "channel": return None + # MX block scales are E8M0 (some producers store its exponent bytes as uint8): group-32 values over another + # scale dtype are not an MX format + scale_dtype = str(weights.get("scale_dtype") or "e8m0") + if quant_method in ("mxfp8", "mxfp4") and not any(name in scale_dtype for name in ("e8m0", "uint8")): + return None targets = [] for target in spec.get("targets") or []: if target.startswith("re:"): @@ -1924,17 +1978,21 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: return None activations = spec.get("input_activations") or {} act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) - # no activation quantization is weight-only, not the weights' native activation format - activation_format = formats.get(act_key) if activations else "bf16" - options = {"activation_format": None if activation_format == quant_method else activation_format} + # block-FP8's own activations are FP8 per token, or per 128-wide group as its kernels quantize them; + # no activation quantization is weight-only + activation_format = {**formats, (8, "float", 128): "fp8"}.get(act_key) if activations else "bf16" + config_class = BLOCK_SCALED_CONFIGS[quant_method] + if activation_format is None or activation_format not in config_class.activation_formats: + return None + options = {"activation_format": activation_format} if quant_method == "fp8": - # 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" + # only a per-tensor scale is calibrated: `dynamic: false` on group-scaled activations is NVFP4's + # second-level global (`input_scale`), which the loader carries on its own + static = not activations.get("dynamic", True) and act_key == (8, "float", None) options.update( activation_scheme="static" if static else "dynamic", weight_block_size=weights.get("block_structure") ) - config = BLOCK_SCALED_CONFIGS[quant_method](**options) + config = config_class(**options) return dict.fromkeys(targets or [".*"], config) @@ -1957,7 +2015,7 @@ def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: ) ignore = config_dict.pop("ignore", None) or config_dict.pop("exclude_modules", None) modules_to_not_convert = subtree_patterns(ignore) if ignore is not None else None - dtype_plan = {".*": NVFP4Config(activation_format="nvfp4")} + dtype_plan = {".*": NVFP4Config()} # the same thing described parametrically, taken only when every group maps, so a multi-group export is not # flattened to one format if groups := config_dict.pop("config_groups", None): @@ -2023,7 +2081,10 @@ class FineGrainedConfig(QuantizationConfigMixin): Args: dtype_plan (`dict[str, BlockScaledConfig]`, *optional*): A module regex and the format config the modules it matches load in. A targeted entry wins over the - `".*"` catch-all; modules no entry covers stay in full precision. + `".*"` catch-all; modules no entry covers stay in full precision. An entry takes no loading options + (`dequantize`, `modules_to_not_convert`, `modules_to_convert`): those are this config's. On a + pre-quantized checkpoint, a passed plan's entries replace the activation formats of the checkpoint's + entries under the same pattern. modules_to_not_convert (`list`, *optional*): Modules to keep in full precision. modules_to_convert (`list`, *optional*): @@ -2032,8 +2093,12 @@ class FineGrainedConfig(QuantizationConfigMixin): Whether to dequantize the model to full precision while loading. """ + # what an entry of the plan does not hold: they are the whole model's + _entry_loading_options = ("dequantize", "modules_to_not_convert", "modules_to_convert") + def __init__( self, + *, dtype_plan: dict[str, BlockScaledConfig] | None = None, modules_to_not_convert: list | None = None, modules_to_convert: list | None = None, @@ -2041,43 +2106,66 @@ def __init__( **kwargs, ): self.quant_method = QuantizationMethod.FINEGRAINED - # from `config.json` an entry is a plain dict naming its format - self.dtype_plan = { - pattern: config - if isinstance(config, BlockScaledConfig) - else BLOCK_SCALED_CONFIGS[config["quant_method"]].from_dict(config) - for pattern, config in (dtype_plan or {}).items() - } + kwargs.pop("quant_method", None) + if kwargs: + logger.info(f"Unused kwargs: {list(kwargs)}. These kwargs are not used in {type(self).__name__}.") + self.dtype_plan = {pattern: self._plan_entry(pattern, entry) for pattern, entry in (dtype_plan or {}).items()} self.modules_to_not_convert = modules_to_not_convert self.modules_to_convert = modules_to_convert self.dequantize = dequantize + def _plan_entry(self, pattern: str, entry) -> BlockScaledConfig: + """`entry` as its format config: from `config.json`, a plain dict naming its format.""" + if isinstance(entry, dict): + method = entry.get("quant_method") + if method not in BLOCK_SCALED_CONFIGS: + raise ValueError( + f"dtype_plan entry {pattern!r} has quant_method {method!r}; an entry is one of " + f"{sorted(BLOCK_SCALED_CONFIGS)}" + ) + entry = BLOCK_SCALED_CONFIGS[method].from_dict(entry) + if not isinstance(entry, BlockScaledConfig): + raise TypeError( + f"dtype_plan entry {pattern!r} is a {type(entry).__name__}; an entry is one of " + f"{', '.join(config.__name__ for config in BLOCK_SCALED_CONFIGS.values())}" + ) + loading = [name for name in self._entry_loading_options if getattr(entry, name, None)] + if loading: + raise ValueError( + f"dtype_plan entry {pattern!r} sets {loading}, which hold for the whole model: set them on the " + "FineGrainedConfig" + ) + return entry + def to_dict(self) -> dict[str, Any]: return { "quant_method": self.quant_method, - "dtype_plan": {pattern: config.to_dict() for pattern, config in self.dtype_plan.items()}, + "dtype_plan": { + pattern: { + name: value for name, value in config.to_dict().items() if name not in self._entry_loading_options + } + for pattern, config in self.dtype_plan.items() + }, "modules_to_not_convert": self.modules_to_not_convert, "modules_to_convert": self.modules_to_convert, "dequantize": self.dequantize, } def get_loading_attributes(self): - # a skip-list left unset keeps the checkpoint's - attributes = {"dequantize": self.dequantize} - if self.modules_to_not_convert is not None: - attributes["modules_to_not_convert"] = self.modules_to_not_convert - return attributes + return _loading_attributes(self, **({"dtype_plan": self.dtype_plan} if self.dtype_plan else {})) class FineGrainedFP8Config(QuantizationConfigMixin): """ - FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. + Deprecated: passed to `from_pretrained`, it loads through its own frozen integration. Use [`BlockFp8Config`], or + [`Mxfp8Config`] for a `[1, 32]` block with UE8M0 scales. Args: activation_scheme (`str`, *optional*, defaults to `"dynamic"`): - The scheme used for activation, the defaults and only support scheme for now is "dynamic". + `"dynamic"` (quantized inline) or `"static"` (a calibrated per-tensor activation scale from 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 block each weight scale covers; `None` for one scale per tensor. dequantize (`bool`, *optional*, defaults to `False`): Whether to dequantize the model during loading. modules_to_not_convert (`list`, *optional*): @@ -2099,12 +2187,16 @@ def __init__( scale_fmt: str = "float", **kwargs, ): + warnings.warn( + "`FineGrainedFP8Config` is deprecated and will be removed in a future release: use `BlockFp8Config`, or " + "`Mxfp8Config` for a `[1, 32]` block with UE8M0 scales.", + FutureWarning, + ) 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 diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index 3043380dff10..a03da65ce018 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -1,307 +1,69 @@ -"""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.""" +# 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 a load and a save run whatever the format: the layout ops (gate|up +interleave, scale container, swizzle) and the ops each expert tensor takes from its module, and a core chain +composed with on-the-fly quantization. Each format's checkpoint layouts are tested in its own folder. Pure +torch: no kernel is loaded and no GPU is needed.""" -import os -import re import unittest from types import SimpleNamespace from unittest import mock import torch -from torch import nn -import transformers.integrations.finegrained.core as fg +import transformers.integrations.finegrained.layers 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 - - -E, H, I = 2, 64, 32 # experts, hidden, intermediate: one 32-group along each input - +from transformers.utils.quantization_config import ( + BlockFp8Config, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) -@use_experts_implementation -class _Experts(nn.Module): - """A model's unquantized experts, which the quantizer replaces with `FineGrainedExperts`.""" +from .utils import ( + E, + H, + I, + _Cfg, + _format_config, + _load_and_save, + _model_mapping, + _MoeModel, + _quantizer_for, + _replace, + fake_finegrained_kernel, + patch_finegrained_kernel, +) - 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)) +@use_experts_implementation(is_transposed=True) +class _TransposedExperts(torch.nn.Module): + """GPT-OSS's layout: the stacked experts stored `(E, in, out)`.""" -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, - 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 - ) - model.hf_quantizer = quantizer - saved = revert_weight_conversion(model, model.state_dict()) - return model, info, saved - - -def _raw(tensor): - return tensor.reshape(-1).view(torch.uint8) if tensor.element_size() == 1 else tensor - - -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(), {}), - ) - - 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 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_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}" - ) + self.gate_up_proj = torch.nn.Parameter(torch.empty(E, H, 2 * I)) + self.down_proj = torch.nn.Parameter(torch.empty(E, I, H)) @require_torch @@ -322,7 +84,13 @@ def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm1 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 + cfg, + quantization_config=_format_config( + weight_format, + activation_format=activation_format, + # the group formats fix their own block + **({"weight_block_size": (128, 128)} if weight_format == "fp8" else {}), + ), ) model = torch.nn.Module() model.config = cfg @@ -364,6 +132,11 @@ def test_megamoe_holds_gate_up_stacked(self): _, experts = self._experts("fp8", impl="deepgemm_megamoe") self.assertFalse(experts.holds_interleaved_gate_up) + def test_deepgemm_holds_gate_up_stacked(self): + """DeepGEMM's grouped GEMM reads a block grid row-major, which interleaved rows do not follow.""" + _, experts = self._experts("fp8", impl="deepgemm") + self.assertFalse(experts.holds_interleaved_gate_up) + def _op_kernel(self): kernel, _ = fake_finegrained_kernel() kernel.swizzle_mx_scales = mock.Mock( @@ -434,7 +207,9 @@ def test_scale_container_op_round_trips_every_container(self): 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"] + back = op.reverse_op.convert( + {"x": held}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" + )["x"] self.assertEqual(back.dtype, container) self.assertTrue(torch.equal(back, shipped)) weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) @@ -464,54 +239,54 @@ def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): 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 test_blocks_scales_converters_carry_the_layout_ops(self): - from transformers.core_model_loading import WeightConverter + def test_expert_tensors_take_their_layout_from_their_module(self): + """The layout an expert tensor lands in comes from the module it fills, whatever converter + brought it: gate|up interleaved (weight, scale grid and bias), then the scale container and + swizzle; a bias only interleaves.""" from transformers.integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, + FineGrainedScaleContainer, 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) + model, experts = self._experts("mxfp4") + experts.quantization_config = Mxfp4Config(activation_format=None) # what the replacement walk records + quantizer = _quantizer_for(Mxfp4Config()) + layout = [FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales] + for param, expected in ( + ("experts.gate_up_proj", layout), + ("experts.gate_up_proj_scale_inv", layout), + ("experts.gate_up_proj_bias", [FineGrainedInterleaveGateUp]), + ("experts.down_proj_scale_inv", [FineGrainedScaleContainer, FineGrainedSwizzleScales]), + ("experts.down_proj_bias", []), + ): + ops = quantizer.get_param_ops(model, param) + self.assertEqual([type(op) for op in ops], expected, param) - def test_catch_all_converters_deliver_the_full_parameter_name(self): + def test_a_key_under_its_fused_name_takes_its_modules_ops_by_its_full_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 + down_proj_scale_inv) has no converter of its own: the loader runs its module's ops in a + one-to-one converter, whose ops receive the tensor under the key and hand it back by name.""" + from transformers.core_model_loading import WeightConverter - model, _ = self._experts("mxfp8") - converters = [c for c in self._quantizer("mxfp8")([]) if isinstance(c, WeightConverter)] + model, experts = self._experts("mxfp8") + experts.quantization_config = Mxfp8Config() # what the replacement walk records + quantizer = _quantizer_for(Mxfp8Config()) 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) + converter = WeightConverter(key, key, operations=quantizer.get_param_ops(model, key)) 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) + converter.add_tensor(key, key, key, 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) + out = converter.convert(key, 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 + """`FineGrainedExperts.from_module` 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, @@ -543,64 +318,16 @@ def test_replacement_keeps_the_models_gate_up_convention(self): ): 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) - ) + _replace(model, _format_config(quant_method, flat=True)) 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") - 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({key: out}, model=model, full_layer_name=key)[key] - torch.testing.assert_close(back, stacked) - experts.is_concatenated = False - 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({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.""" + """The op reads the experts module a key fills: quantized experts get the interleaved order and the + reverse restores the stack, while a model whose rows already come interleaved (`is_concatenated=False`, + GPT-OSS), Mega MoE (which packs gate|up itself) and full-precision experts of a mixed checkpoint keep + 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] @@ -611,7 +338,7 @@ def test_interleave_op_follows_the_layer_it_fills(self): 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] + model.layers[1].experts = quantized = self._experts("mxfp8")[1] op = FineGrainedInterleaveGateUp(hf_quantizer=None) for layer, expected in ((0, stacked), (1, interleaved)): @@ -622,553 +349,185 @@ def test_interleave_op_follows_the_layer_it_fills(self): back = op.reverse_op.convert({saved_key: out}, model=model, full_layer_name=saved_key)[saved_key] torch.testing.assert_close(back, stacked) + key = "layers.1.experts.gate_up_proj" + quantized.is_concatenated = False + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) + quantized.is_concatenated, quantized.holds_interleaved_gate_up = True, False + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) -@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 - 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 + def test_interleave_op_stacks_an_interleaved_checkpoint_for_stacked_experts(self): + """An interleaved checkpoint (GPT-OSS) loaded for a backend that holds gate|up stacked (DeepGEMM) is + de-interleaved on load and interleaved back on save.""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp - 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, + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) + interleaved = stacked[:, [0, 3, 1, 4, 2, 5]] + model, experts = self._experts("mxfp8", impl="deepgemm") + experts.is_concatenated = False + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + key = "experts.gate_up_proj" + loaded = op.convert({key: interleaved}, model=model, full_layer_name=key)[key] + torch.testing.assert_close(loaded, stacked) + torch.testing.assert_close( + op.reverse_op.convert({key: loaded}, model=model, full_layer_name=key)[key], interleaved ) - 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, + def test_ops_refuse_a_list_of_several_tensors(self): + """The loader hands a single-source tensor in a one-element list; an op that unwraps it must not + silently keep the first of several.""" + from transformers.integrations.finegrained.conversions import ( + FineGrainedInputScalesSplit, + FineGrainedPackedBlocks, + FineGrainedQuantize, + FineGrainedScaleContainer, + FineGrainedViewPackedInt8, ) - 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.""" + tensor = torch.zeros(2, 4, 32, 16, dtype=torch.uint8) + for op in ( + FineGrainedScaleContainer(None), + FineGrainedPackedBlocks(None), + FineGrainedViewPackedInt8(None), + FineGrainedInputScalesSplit(None), + FineGrainedQuantize(None), + ): + with self.subTest(op=type(op).__name__), self.assertRaises(ValueError): + op.convert({"experts.down_proj": [tensor, tensor]}) + # a single target, re-keyed by the op before anything reads it + with self.assertRaises(ValueError): + FineGrainedViewPackedInt8(None).convert({"x": [tensor, tensor]}, target_patterns=["experts.down_proj"]) + + def test_per_half_globals_are_not_folded_past_a_glu_nonlinear_in_up(self): + """Folding the up half's NVFP4 global onto the down projection assumes the GLU is linear in `up`; a + clamped SwiGLU or SwiGLU-OAI's `(up + 1)` is not, so differing halves are refused there. Halves that + agree have nothing to fold.""" 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_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( - { - "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([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", - )["mlp.experts.down_proj_input_global_scale"] - 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]}", + targets = ["experts.gate_up_proj_weight_global_scale", "experts.down_proj_weight_global_scale"] + for attribute in ("swiglu_limit", "swiglu_alpha"): + config = _Cfg() + setattr(config, attribute, 7.0) + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts(config, quantization_config=NVFP4Config()) + with self.subTest(attribute=attribute): + with self.assertRaises(NotImplementedError): + 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, + model=model, ) + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 1.0], [2.0, 2.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, + model=model, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([5.0, 6.0])) @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.""" +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_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 _quantize_on_load(self, model, checkpoint, weight_mapping): + """`checkpoint`'s full-precision weights quantized on load into 32x32-block FP8, off SM100.""" + quantizer = _quantizer_for(BlockFp8Config(weight_block_size=(32, 32))) + quantizer.pre_quantized = False + 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(weight_mapping), + hf_quantizer=quantizer, + dtype_plan=quantizer.dtype_plan, + ) + info, _ = convert_and_load_state_dict_in_model(model, checkpoint, load_config) + self.assertEqual(info.conversion_errors, {}) + return model.model.layers[0].mlp.experts - 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_dequantizing_a_fused_nvfp4_stack_applies_each_experts_globals(self): + """The vLLM layout ships each stack's globals per expert, `(E,)`, and a gate|up stack's per half, + `(E, 2)`: the stack's dequantize scales each expert's rows, and each half's, by its own.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize - 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))): + checkpoint, expected = {}, {} + for proj, rows, cols, halves in (("gate_up_proj", 2 * I, H, 2), ("down_proj", H, I, 1)): + codes = torch.randint(0, 256, (E, rows, cols // 2), dtype=torch.uint8) + blocks = (torch.rand(E, rows, cols // 16) + 0.5).to(torch.float8_e4m3fn) + globals_ = torch.pow(2.0, torch.randint(-8, 0, (E, halves)).float()) 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), + checkpoint |= { + prefix: codes, + f"{prefix}_weight_scale": blocks, + f"{prefix}_weight_scale_2": globals_.squeeze(1) if halves == 1 else globals_, } - ) - 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) + half_rows = rows // halves + expected[proj] = torch.stack( + [ + torch.cat( + [ + FineGrainedDequantize(None)._dequantize_one( + codes[e, h * half_rows : (h + 1) * half_rows], + blocks[e, h * half_rows : (h + 1) * half_rows], + torch.float32, + globals_[e, h], + ) + for h in range(halves) + ] + ) + for e in range(E) + ] + ) + model, info, _ = _load_and_save(checkpoint, NVFP4Config(dequantize=True)) + self.assertEqual(info.conversion_errors, {}) 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) + for proj, reference in expected.items(): + torch.testing.assert_close(getattr(experts, proj).float(), reference, msg=proj) - 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_on_the_fly_quantization_blocks_each_projection_whole(self): + """The gate|up interleave runs on the quantized weight and its grid, so every 32-row block holds rows of + ONE projection, scaled by its own magnitude: a gate four orders below its up half keeps its values. + Both stack layouts: a per-expert checkpoint concatenated by the model's converter, and one stored + transposed `(E, H, 2I)`, which the quantization reads as rows.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize - 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"), + gate, up, down = torch.randn(E, I, H) * 1e-5, torch.randn(E, I, H), torch.randn(E, H, I) + dense = {"model.layers.0.self_attn.q_proj.weight": torch.randn(H, H)} + per_expert = { + f"model.layers.0.mlp.experts.{expert}.{proj}.weight": tensor[expert] + for expert in range(E) + for proj, tensor in (("gate_proj", gate), ("up_proj", up), ("down_proj", down)) + } + transposed = { + "model.layers.0.mlp.experts.gate_up_proj": torch.cat([gate, up], dim=1).transpose(1, 2).contiguous(), + "model.layers.0.mlp.experts.down_proj": down.transpose(1, 2).contiguous(), + } + config = PreTrainedConfig() + config.hidden_size, config.intermediate_size, config.num_local_experts, config.hidden_act = H, I, E, "silu" + for layout, checkpoint, mapping in ( + ("concatenated", per_expert, _model_mapping()), + ("transposed", transposed, []), ): - 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", + with self.subTest(layout=layout): + model = _MoeModel(config) + if layout == "transposed": + model.model.layers[0].mlp.experts = _TransposedExperts(config) + experts = self._quantize_on_load(model, {**dense, **checkpoint}, mapping) + self.assertTrue(experts.holds_interleaved_gate_up) + for half, expected in ((0, gate), (1, up)): + held = FineGrainedDequantize(None)._dequantize_one( + experts.gate_up_proj[:, half::2], experts.gate_up_proj_scale_inv[:, half::2], torch.float32 ) - - 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 - run on the full-precision tensor, and only the finegrained module of the three gets a scale.""" + error = ((held - expected).norm() / expected.norm()).item() + self.assertLess(error, 0.1, f"half {half}") def test_a_chunked_rope_permute_quantizes_only_the_finegrained_target(self): from transformers.integrations.finegrained import FineGrainedLinear @@ -1188,7 +547,9 @@ def __init__(self): 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.q_proj = FineGrainedLinear( + in_dim, out_dim, quantization_config=BlockFp8Config(weight_block_size=block_size) + ) self.k_proj = RopeProjector() self.v_proj = RopeProjector() @@ -1239,6 +600,7 @@ def _init_weights(self, module): ), "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), "get_quantize_ops": lambda self: FineGrainedQuantize(self), + "get_param_ops": lambda self, _model, _param_name: [], "pre_quantized": False, }, ) diff --git a/tests/integrations/finegrained/test_core.py b/tests/integrations/finegrained/test_core.py deleted file mode 100644 index d86c1c4a3bca..000000000000 --- a/tests/integrations/finegrained/test_core.py +++ /dev/null @@ -1,565 +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. -"""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 -from contextlib import ExitStack, contextmanager -from unittest import mock - -import torch -from parameterized import parameterized - -import transformers.integrations.finegrained.core as fg -from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear -from transformers.testing_utils import require_torch - - -class _Cfg: - hidden_size = 64 - num_local_experts = 4 - intermediate_size = 32 - hidden_act = "silu" - _experts_implementation = 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 FineGrainedModulesTest(unittest.TestCase): - """The finegrained modules themselves: what they hold and what replaces what.""" - - 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) - - # 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) - - # `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) - 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 - ) - ) - - -@require_torch -class FrozenFp8ShimTest(unittest.TestCase): - 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 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() - - def style(name): - return _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) - - 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): - """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) - - 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 - # 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) - - -@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 - - 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")) - # 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.""" - 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"]) - - 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) diff --git a/tests/integrations/finegrained/test_distributed.py b/tests/integrations/finegrained/test_distributed.py deleted file mode 100644 index 88f77b32b297..000000000000 --- a/tests/integrations/finegrained/test_distributed.py +++ /dev/null @@ -1,545 +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. -"""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 deleted file mode 100644 index 25008522d003..000000000000 --- a/tests/integrations/finegrained/test_forwards.py +++ /dev/null @@ -1,442 +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. -"""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 os -import unittest -from unittest import mock - -import torch - -import transformers.integrations.finegrained.core as fg -from transformers.integrations.finegrained import ( - FineGrainedExperts, - FineGrainedLinear, - 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) - - -# `(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): - """`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): - # the kernels autotune per shape from a cold cache, which dwarfs everything else here - os.environ["FINEGRAINED_AUTOTUNE_TRIALS"] = "1" - load_finegrained_kernel() - - 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 - 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, 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() - 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() - out[token] += row * wts[token, slot].float() - 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_layers.py b/tests/integrations/finegrained/test_layers.py new file mode 100644 index 000000000000..8e7872c37fa4 --- /dev/null +++ b/tests/integrations/finegrained/test_layers.py @@ -0,0 +1,500 @@ +# 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` integration's modules, no kernel involved: what replaces what, what the modules hold and +hand the kernels, and the parallel plans a quantized model needs. The kernel call contract lives in +`tests/kernels/test_finegrained.py`, the conversion ops in `test_conversions.py`.""" + +import pickle +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +import transformers.integrations.finegrained.layers as fg +from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import BlockFp8Config, NVFP4Config + +from .utils import ( + _Cfg, + _format_config, + _quantizer_for, + _replace, + fake_finegrained_kernel, + patch_finegrained_kernel, +) + + +def _block_scaled_matmul_2d(A, B, As, Bs, *, output_dtype=None, **kwargs): + """`matmul_2d`'s contract on a block-scaled weight: the grid read row-major, its block inferred from the shapes.""" + block_n, block_k = B.shape[0] // Bs.shape[0], B.shape[1] // Bs.shape[1] + weight = B.float() * Bs.float().repeat_interleave(block_n, 0).repeat_interleave(block_k, 1) + return (A.float() @ weight.T).to(output_dtype) + + +@require_torch +class FineGrainedModulesTest(unittest.TestCase): + """The finegrained modules themselves: what they hold and what replaces what.""" + + def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): + """`replace_with_quantized_modules` 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 transformers.integrations.finegrained import FineGrainedGroupedLinear + + class OddlyNamedBlockDiagonal(torch.nn.Linear): + def __init__(self): + super().__init__(8, 16, bias=False) + self.n_groups = 2 + + model = torch.nn.Module() + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + model.grouped = OddlyNamedBlockDiagonal() + model.plain = torch.nn.Linear(8, 16, bias=False) + + with torch.device("meta"): + _replace(model, BlockFp8Config()) + self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) + self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) + + # 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_replacing_experts_leaves_the_experts_class_as_defined(self): + """Each layer's flags come with its own class, shared by the layers that agree on them; a load modifies + neither `FineGrainedExperts` nor what a later direct construction gets.""" + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + num_experts=2, + num_experts_per_tok=1, + ) + with torch.device("meta"): + model = Qwen3MoeForCausalLM(config) + _replace(model, BlockFp8Config()) + experts = {type(layer.mlp.experts) for layer in model.model.layers} + self.assertEqual(len(experts), 1) + self.assertTrue(issubclass(experts.pop(), FineGrainedExperts)) + self.assertFalse(hasattr(FineGrainedExperts.forward, "__wrapped__")) + direct = FineGrainedExperts(_Cfg(), quantization_config=_format_config("fp8"), has_bias=True) + self.assertTrue(direct.has_bias) + # a replaced experts class is found by name, so a pickled model loads back + experts = model.model.layers[0].mlp.experts + self.assertIs(pickle.loads(pickle.dumps(type(experts))), type(experts)) + + 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_a_skip_list_names_modules_by_pattern(self): + """`should_convert_module`: everything converts by default, a skip-list entry names its modules.""" + from transformers.quantizers.quantizers_utils import should_convert_module + + self.assertTrue(should_convert_module("model", None)) + self.assertTrue(should_convert_module("model", [])) + patterns = ["model.layers.*.self_attn", "lm_head"] + self.assertFalse(should_convert_module("lm_head", patterns)) + self.assertTrue(should_convert_module("experts", patterns)) + + def test_an_empty_plan_replaces_nothing(self): + """A dequantized load walks an empty plan; it must not fall back to the quantizer's own.""" + model = torch.nn.Sequential(torch.nn.Linear(256, 256)) + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + self.assertFalse(_quantizer_for(BlockFp8Config()).replace_modules(model, plan={})) + self.assertIsInstance(model[0], torch.nn.Linear) + + def test_replacement_follows_the_patterns_and_the_skip_list(self): + """A table is quantized where `modules_to_convert` names it, never through the plan's catch-all.""" + + 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) + _replace(model, BlockFp8Config(modules_to_convert=["ngram_embedding", "other"]), ["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) + + # `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) + 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( + quantization_config=_format_config( + f, **({"weight_block_size": (4, 4), "activation_scheme": s} if f == "fp8" else {}) + ), + in_features=64, + out_features=32, + ) + ) + + def test_embedding_scale_ignores_the_default_dtype(self): + self._assert_default_dtype_invariant(lambda: fg.FineGrainedEmbedding(64, 16)) + + 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(), + quantization_config=_format_config( + f, **({"weight_block_size": (4, 4), "activation_scheme": s} if f == "fp8" else {}) + ), + ) + ) + + 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(), quantization_config=BlockFp8Config(weight_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_eager_gated_block_scales_pair_each_half_with_its_own_blocks(self): + """The loader interleaves a gate|up block grid per BLOCK, so one scale covers rows of one half while + the rows alternate between halves: the eager loop runs each half on its own rows and blocks rather + than hand `matmul_2d` the stack as a row-major grid. Halves of very different magnitudes, as real + gate and up projections have, so a scale read from the wrong half shows.""" + experts = FineGrainedExperts(_Cfg(), quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + experts._deepgemm_disabled = True + num_experts, hidden, inter = _Cfg.num_local_experts, _Cfg.hidden_size, _Cfg.intermediate_size + torch.manual_seed(0) + codes = torch.randn(num_experts, 2, inter, hidden).to(torch.float8_e4m3fn) # [gate, up] + exponents = torch.tensor([-8.0, 0.0]).reshape(1, 2, 1, 1) + torch.randint(-2, 3, (num_experts, 2, 1, 2)) + grid = torch.pow(2.0, exponents) # (experts, half, row blocks, column blocks) + # interleaved: rows [g0, u0, g1, u1, ...], grid rows [gate block, up block] + experts.gate_up_proj = torch.nn.Parameter(codes.transpose(1, 2).reshape(num_experts, 2 * inter, hidden)) + experts.gate_up_proj_scale_inv = torch.nn.Parameter(grid.transpose(1, 2).reshape(num_experts, 2, 2)) + + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = _block_scaled_matmul_2d + x = torch.randn(3, hidden) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = experts.expert_linear(x, "gate_up_proj", 1) + halves = [ + x @ (codes[1, half].float() * grid[1, half].repeat_interleave(32, 0).repeat_interleave(32, 1)).T + for half in (0, 1) + ] + torch.testing.assert_close(out, torch.stack(halves, dim=-1).flatten(-2)) + + def test_stacked_gate_up_is_gated_by_halves_and_kept_from_the_triton_chains(self): + """The DeepGEMM backends load gate|up stacked: the host GLU splits the halves, and the triton chains, + which read the rows interleaved, refuse such a module rather than pair the wrong rows.""" + config = _Cfg() + config._experts_implementation = "deepgemm" + experts = FineGrainedExperts(config, quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + self.assertFalse(experts.holds_interleaved_gate_up) + gate, up = torch.randn(3, _Cfg.intermediate_size), torch.randn(3, _Cfg.intermediate_size) + torch.testing.assert_close( + experts._apply_gate(torch.cat([gate, up], dim=-1)), torch.nn.functional.silu(gate) * up + ) + + hidden = torch.randn(2, _Cfg.hidden_size, dtype=torch.bfloat16) + routing = torch.zeros(2, 1, dtype=torch.long), torch.ones(2, 1) + with mock.patch.object(fg, "load_finegrained_kernel") as load: + for forward in (fg.finegrained_grouped_mm_experts_forward, fg.finegrained_batched_mm_experts_forward): + with ( + self.subTest(forward=forward.__name__), + self.assertRaisesRegex(NotImplementedError, "interleaved"), + ): + forward(experts, hidden, *routing) + load.assert_not_called() + + def test_swiglu_alpha_without_a_limit_runs_unclamped(self): + config = _Cfg() + config.swiglu_alpha = 1.702 + experts = FineGrainedExperts(config, quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + gate_up = torch.randn(3, 2 * _Cfg.intermediate_size) + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] + torch.testing.assert_close(experts._apply_gate(gate_up), (up + 1) * gate * torch.sigmoid(gate * 1.702)) + + def test_grouped_linear_hands_over_its_stored_weight_and_activation_scales(self): + """The groups split the weight as stored (packed FP4 halves K), and the calibrated activation scale goes + where `FineGrainedLinear` hands it: NVFP4's as the activation global, a static FP8 one as `As`.""" + from transformers.integrations.finegrained import FineGrainedGroupedLinear + + for config, packed in ( + (NVFP4Config(), True), + (BlockFp8Config(weight_block_size=(16, 16), activation_scheme="static"), False), + ): + with self.subTest(weight_format=config.method.value): + linear = FineGrainedGroupedLinear(64, 32, 2, quantization_config=config) + kernel, recorder = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + linear(torch.randn(3, 2, 64, dtype=torch.bfloat16)) + (call,) = recorder.calls["matmul_grouped"] + self.assertEqual(call.args[1].shape, (2, 16, 32 if packed else 64)) + two_level = linear.weight_global_scale is not None + self.assertIs(call.args[2], None if two_level else linear.activation_scale) + self.assertIs(call.kwargs["a_global_scale"], linear.activation_scale if two_level else None) + + +@require_torch +class FineGrainedFusedNormTest(unittest.TestCase): + """A norm the kernels can fuse is not fused when a forward collective sits on its input: under intra-expert + TP the rows are a partial sum over the sharded intermediate, which the wrapped module all-reduces in its own + forward.""" + + def _operands(self, *, input_reduce): + experts = FineGrainedExperts(_Cfg(), quantization_config=_format_config("fp8", weight_block_size=(4, 4))) + experts.post_expert_norm, experts.has_post_expert_norm = torch.nn.RMSNorm(_Cfg.hidden_size), True + if input_reduce: + experts.post_expert_norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = ("rms_norm",) + with mock.patch.dict(fg.SUPPORTED_FUSED_NORMS, {torch.nn.RMSNorm: "rms_norm"}): + return experts._moe_operands(kernel), experts + + def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): + operands, experts = self._operands(input_reduce=False) + self.assertEqual(operands["post_expert_norm"], "rms_norm") + self.assertIs(operands["post_expert_norm_weight"], experts.post_expert_norm.weight) + + def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): + operands, experts = self._operands(input_reduce=True) + self.assertIs(operands["post_expert_norm"], experts.post_expert_norm) + self.assertIsNone(operands["post_expert_norm_weight"]) + + +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 transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + quantizer = BlockFp8HfQuantizer(BlockFp8Config()) + config = SimpleNamespace(base_model_tp_plan=dict(raw or self.BASE_EP), _experts_implementation=impl) + # the config is updated in place and handed back, as the loader takes it + self.assertIs(quantizer.update_tp_plan(config), 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() + + def style(name): + return _get_parameter_tp_plan(f"layers.3.mlp.experts.{name}", plan) + + for proj in ("gate_up_proj", "down_proj"): + for name in (proj, f"{proj}_scale_inv", f"{proj}_weight_global_scale", f"{proj}_activation_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + for name in (f"{proj}_bias",): + self.assertEqual(style(name), "grouped_gemm", name) + 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 + follows its weight while the per-expert globals stay replicated.""" + 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) + + 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 + # 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; an expert-axis style + would shard a dim they do not have.""" + 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: rewriting only the config the quantizer + is handed leaves the scales whole against sharded weights.""" + from transformers import Glm4vMoeConfig + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + 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 {}) + BlockFp8HfQuantizer(BlockFp8Config()).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", "_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_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") + # Mega MoE dispatches across ranks itself, in place of the EP dispatch a plan may default to + dispatched = self._planned("deepgemm_megamoe", raw={"layers.*.mlp.experts": "ep_dispatch_experts"}) + self.assertEqual(dispatched["layers.*.mlp.experts"], "megamoe_experts") + + def test_a_real_models_plan_keeps_its_entries_and_gains_the_expert_companions(self): + """`update_tp_plan` adds companions and never drops what the model ships. Real configs, since a branch + keyed on the config class is invisible to a `SimpleNamespace`.""" + from transformers import Qwen3Config, Qwen3MoeConfig, Qwen3VLMoeConfig + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + 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") + BlockFp8HfQuantizer(BlockFp8Config()).update_tp_plan(config) + for key, style in shipped.items(): + self.assertEqual(config.base_model_tp_plan.get(key), style, key) + for companion in ("gate_up_proj_scale_inv", "down_proj_scale_inv", "gate_up_proj_bias"): + self.assertIn(f"layers.*.mlp.experts.{companion}", config.base_model_tp_plan, companion) + + # a multimodal model keeps its plan on the sub-config, and nothing is invented on the outer one + multimodal = Qwen3VLMoeConfig() + BlockFp8HfQuantizer(BlockFp8Config()).update_tp_plan(multimodal) + self.assertFalse(getattr(multimodal, "base_model_tp_plan", None) or {}) + self.assertTrue(multimodal.text_config.base_model_tp_plan) diff --git a/tests/integrations/finegrained/test_models.py b/tests/integrations/finegrained/test_models.py deleted file mode 100644 index e41585fe7c6e..000000000000 --- a/tests/integrations/finegrained/test_models.py +++ /dev/null @@ -1,250 +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. -"""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 torch -from parameterized import parameterized - -from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText -from transformers.testing_utils import ( - Expectations, - TestCasePlus, - cleanup, - require_kernels, - require_torch_accelerator, - slow, - torch_device, -) - - -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 - - 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(): - 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 -} - - -@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): - # the kernels autotune per shape from a cold cache, which dwarfs everything else here - os.environ["FINEGRAINED_AUTOTUNE_TRIALS"] = "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 tearDown(self): - super().tearDown() - cleanup(torch_device, gc_collect=True) - - @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) - - @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(): - 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(): - 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/integrations/finegrained/utils.py b/tests/integrations/finegrained/utils.py new file mode 100644 index 000000000000..f31c0cfc6f0c --- /dev/null +++ b/tests/integrations/finegrained/utils.py @@ -0,0 +1,258 @@ +# 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. +"""Helpers shared by the fine-grained quantization tests: the format configs, the quantizer a load picks, the +checkpoint-config builders, the mocked kernel bundle, and the tiny MoE checkpoint the conversion chains load +and save back.""" + +from dataclasses import dataclass +from types import SimpleNamespace +from unittest import mock + +import torch +from torch import nn + +import transformers.integrations.finegrained.layers as fg +from transformers import PreTrainedConfig, PreTrainedModel +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import ( + convert_and_load_state_dict_in_model, + revert_weight_conversion, +) +from transformers.integrations.moe import use_experts_implementation +from transformers.modeling_utils import LoadStateDictConfig +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.utils.quantization_config import BLOCK_SCALED_CONFIGS + + +def _format_config(weight_format, flat=False, **options): + """`weight_format`'s config. Built for a module directly, per-tensor FP8 and W4A4 MXFP4 unless asked otherwise; + `flat`, the defaults a config loading a whole checkpoint takes.""" + defaults = ( + {} + if flat + else {"fp8": {"weight_block_size": None}, "mxfp4": {"activation_format": None}}.get(weight_format, {}) + ) + return BLOCK_SCALED_CONFIGS[weight_format](**{**defaults, **options}) + + +class _Cfg: + hidden_size = 64 + num_local_experts = 4 + intermediate_size = 32 + hidden_act = "silu" + _experts_implementation = None + + +def _quantizer_for(config): + """The quantizer a `from_pretrained` picks for this config.""" + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING + + method = getattr(config.quant_method, "value", config.quant_method) + return AUTO_QUANTIZER_MAPPING[method](config) + + +def _replace(model, config, modules_to_not_convert=None): + """The module swap a load with `config` makes, through the quantizer it picks.""" + return _quantizer_for(config).replace_modules(model, modules_to_not_convert=modules_to_not_convert) + + +def _targets(conv): + t = conv.target_patterns + return t if isinstance(t, (list, tuple)) else [t] + + +def _modelopt(**fields): + return AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4", **fields}) + + +def _model_config(**text_fields): + text = SimpleNamespace(**text_fields) + return SimpleNamespace(get_text_config=lambda: text) + + +def _compressed_tensors(weights, activations=None, fmt="nvfp4-pack-quantized", **fields): + group = { + "weights": weights, + "targets": ["Linear"], + **({"input_activations": activations} if activations else {}), + } + return {"quant_method": "compressed-tensors", "format": fmt, "config_groups": {"group_0": group}, **fields} + + +@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), + ) + + +E, H, I = 2, 64, 32 # experts, hidden, intermediate: one 32-group along each input + + +@use_experts_implementation +class _Experts(nn.Module): + """A model's unquantized experts, which the quantizer replaces with `FineGrainedExperts`.""" + + 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 get_checkpoint_conversion_mapping("qwen2_moe") + + +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, + 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 + ) + model.hf_quantizer = quantizer + saved = revert_weight_conversion(model, model.state_dict()) + return model, info, saved + + +def _raw(tensor): + return tensor.reshape(-1).view(torch.uint8) if tensor.element_size() == 1 else tensor + + +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(), {}), + ) + + 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}") diff --git a/tests/integrations/test_finegrained_fp8.py b/tests/integrations/test_finegrained_fp8.py deleted file mode 100644 index 2570aaad8ed7..000000000000 --- a/tests/integrations/test_finegrained_fp8.py +++ /dev/null @@ -1,349 +0,0 @@ -# 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. -"""Integration tests for the finegrained-fp8 Triton kernel bindings. - -Two layers, both mocking only what needs the `kernels-community/finegrained-fp8` hub build: - -* `FinegrainedFp8LoaderTest` mocks `is_kernels_available` and `lazy_load_kernel` to drive - `load_finegrained_fp8_kernel` with no hub download — asserting it gates correctly (raises when a - precondition is unmet, returns the bundle otherwise) and stays torch-compile safe. -* `FinegrainedFp8ForwardTest` mocks only the loaded kernel bundle (`matmul` / `batched_matmul` / - `grouped_matmul`) with capturing fakes that return correctly-shaped tensors, and runs the real - `finegrained_fp8_linear`, `fp8_batched_mm_experts_forward` and `fp8_grouped_mm_experts_forward` so - their repeat-interleave / routing-flatten / sort-and-histogram / gating / sentinel-mask / reshape-sum - glue executes for real (device-agnostic, so it runs on CPU); it then asserts the tensors handed to the kernel are what the - kernel expects (flattened routing, unclamped sentinels, weight/scale pairing, ...) and - the value-exact output the surrounding marshalling produces. -""" - -import contextlib -import unittest -from unittest import mock - -import torch -from parameterized import parameterized - -import transformers.integrations.finegrained_fp8 as fg -from transformers.integrations.finegrained_fp8 import ( - finegrained_fp8_linear, - fp8_batched_mm_experts_forward, - fp8_grouped_mm_experts_forward, -) -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 - - -# ── Fake kernels (a "good" one exposing every symbol the loader resolves, plus a variant missing one) ── - - -class _FakeFinegrainedKernel: - matmul_2d = staticmethod(_add_one) - matmul_batched = staticmethod(_add_one) - matmul_grouped = staticmethod(_add_one) - - -class _FinegrainedKernelMissingSymbol: - matmul_2d = staticmethod(_add_one) # missing matmul_batched / matmul_grouped - - -@require_torch -class FinegrainedFp8LoaderTest(unittest.TestCase): - def setUp(self): - fg._FINEGRAINED_FP8 = None - self.addCleanup(setattr, fg, "_FINEGRAINED_FP8", None) - - def _env(self, *, kernels_available=True, kernel=_FakeFinegrainedKernel): - stack = contextlib.ExitStack() - stack.enter_context(mock.patch.object(fg, "is_kernels_available", return_value=kernels_available)) - stack.enter_context(mock.patch.object(fg, "lazy_load_kernel", return_value=kernel)) - return stack - - def test_loads_when_environment_is_valid(self): - with self._env(): - bundle = fg.load_finegrained_fp8_kernel() - self.assertIsInstance(bundle, fg.FineGrainedFP8) - self.assertIs(bundle.matmul, _add_one) - - @parameterized.expand( - [ - ("no_kernels", {"kernels_available": False}, "kernel unavailable"), - ("kernel_load_fails", {"kernel": None}, "Failed to load the finegrained-fp8 kernel"), - ("missing_symbols", {"kernel": _FinegrainedKernelMissingSymbol}, "missing required symbols"), - ] - ) - def test_loader_raises(self, _name, env_kwargs, pattern): - with self._env(**env_kwargs), self.assertRaisesRegex(ImportError, pattern): - fg.load_finegrained_fp8_kernel() - - def test_loader_is_compile_safe(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 (`Unsupported: torch.* op returned - # non-Tensor`). Default (inductor) backend. The finegrained loader has no arch/CUDA gate, so - # nothing here fakes the device — torch.compile sees the real one; no GPU required. - with self._env(): - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return fg.load_finegrained_fp8_kernel().matmul(x) - - out = run(torch.zeros(3, device=torch_device)) - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - - def test_loader_is_compile_safe_when_warm(self): - # Warm path (production order: eager warmup, then compile). The loader hits its short-circuit at - # trace time — the branch that must also return None, not the already-loaded bundle. - with self._env(): - fg.load_finegrained_fp8_kernel() - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return fg.load_finegrained_fp8_kernel().matmul(x) - - out = run(torch.zeros(3, device=torch_device)) - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - - -@require_torch -class FinegrainedFp8ForwardTest(unittest.TestCase): - """Drives the real finegrained-fp8 forwards with only the loaded kernel bundle mocked.""" - - def setUp(self): - fg._FINEGRAINED_FP8 = None - self.addCleanup(setattr, fg, "_FINEGRAINED_FP8", None) - - @contextlib.contextmanager - def _mocked_kernel(self): - # Each fake records every call and returns a rightly-shaped, deterministic tensor: - # `matmul` fills 3.0 (so the bias add is value-checkable); the batched/grouped experts fakes - # return ones (so the routing-weight * mask * reshape-sum reduction is analytically exact and - # independent of the — non-stable on CUDA — expert sort). - calls = {"matmul": [], "batched_matmul": [], "grouped_matmul": []} - - def fake_matmul(input, weight, weight_scale_inv, block_size, output_dtype, *, activation_scale=None): - calls["matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale_inv": weight_scale_inv, - "block_size": block_size, - "output_dtype": output_dtype, - "activation_scale": activation_scale, - } - ) - return torch.full((input.shape[0], weight.shape[0]), 3.0, dtype=output_dtype, device=input.device) - - def fake_batched(input, weight, weight_scale, *, block_size, expert_ids): - calls["batched_matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale": weight_scale, - "block_size": block_size, - "expert_ids": expert_ids, - } - ) - return torch.ones(input.shape[0], weight.shape[1], dtype=torch.float32, device=input.device) - - def fake_grouped(input, weight, weight_scale, *, offsets, tokens_per_expert, block_size): - calls["grouped_matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale": weight_scale, - "offsets": offsets, - "tokens_per_expert": tokens_per_expert, - "block_size": block_size, - } - ) - return torch.ones(input.shape[0], weight.shape[1], dtype=torch.float32, device=input.device) - - bundle = fg.FineGrainedFP8(matmul=fake_matmul, batched_matmul=fake_batched, grouped_matmul=fake_grouped) - with mock.patch.object(fg, "load_finegrained_fp8_kernel", return_value=bundle): - yield calls - - @staticmethod - def _expected_experts_output(hidden_states, top_k_index, top_k_weights, num_experts): - # Mirror the shared tail of both experts forwards with the ones-returning down projection: - # weighted_out = ones * routing_weight, sentinel rows zeroed, then per-token reduction. - num_tokens, top_k = top_k_index.shape - hidden = hidden_states.size(-1) - sw = top_k_weights.reshape(-1).to(torch.float32) - weighted = torch.ones(sw.shape[0], hidden, device=hidden_states.device) * sw.unsqueeze(-1) - sentinel = (top_k_index.reshape(-1) >= num_experts).unsqueeze(-1) - weighted = weighted.masked_fill(sentinel, 0.0) - return weighted.view(num_tokens, top_k, hidden).sum(dim=1).to(hidden_states.dtype) - - # ── finegrained_fp8_linear ──────────────────────────────────────────────────────────────────── - - def test_linear_marshals_args_and_defaults_output_dtype(self): - input = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - weight = torch.randn(16, 8, device=torch_device).to(torch.float8_e4m3fn) - weight_scale_inv = torch.randn(1, 1, dtype=torch.float32, device=torch_device) - block_size = [128, 128] - with self._mocked_kernel() as calls: - out = finegrained_fp8_linear(input, weight, weight_scale_inv, block_size) - call = calls["matmul"][0] - # input / weight / scale pass straight through, positionally. - self.assertIs(call["input"], input) - self.assertIs(call["weight"], weight) - self.assertIs(call["weight_scale_inv"], weight_scale_inv) - # output_dtype defaults to input.dtype when the caller leaves it None, and no activation scale. - self.assertEqual(call["output_dtype"], torch.bfloat16) - self.assertIsNone(call["activation_scale"]) - self.assertEqual(out.dtype, torch.bfloat16) - - def test_linear_adds_bias_in_place(self): - input = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - weight = torch.randn(16, 8, device=torch_device).to(torch.float8_e4m3fn) - weight_scale_inv = torch.randn(1, 1, dtype=torch.float32, device=torch_device) - bias = torch.randn(16, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(): - out = finegrained_fp8_linear(input, weight, weight_scale_inv, [128, 128], bias=bias) - # The kernel returned a 3.0-filled (3, 16) tensor; bias is broadcast-added onto it. - self.assertTrue(torch.equal(out, torch.full((3, 16), 3.0, dtype=torch.bfloat16, device=torch_device) + bias)) - - # ── fp8_batched_mm_experts_forward ────────────────────────────────────────────────────────────── - - def test_batched_mm_kernel_inputs_and_output(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - up, down = calls["batched_matmul"] - # Up projection: each token replicated top_k times (S = 6), routing flattened row-major. - self.assertTrue(torch.equal(up["input"], hidden_states.repeat_interleave(2, dim=0))) - self.assertTrue(torch.equal(up["expert_ids"], top_k_index.reshape(-1))) - # Weight/scale pairing passes straight through (no copies). - self.assertIs(up["weight"], experts.gate_up_proj) - self.assertIs(up["weight_scale"], experts.gate_up_proj_scale_inv) - # Down projection sees the gated activation: 2*inter (32) collapsed to inter (16). - self.assertEqual(down["input"].shape, (6, 16)) - self.assertIs(down["weight"], experts.down_proj) - self.assertIs(down["weight_scale"], experts.down_proj_scale_inv) - self.assertTrue(torch.equal(down["expert_ids"], top_k_index.reshape(-1))) - - # Value-exact per-token reduction, recast to the input dtype. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_batched_mm_non_gated_uses_up_proj_and_activation(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, has_gate=False) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - up, down = calls["batched_matmul"] - self.assertIs(up["weight"], experts.up_proj) - self.assertIs(up["weight_scale"], experts.up_proj_scale_inv) - # Non-gated: act_fn keeps the inter dim (16), no 2*inter chunk. - self.assertEqual(down["input"].shape, (6, 16)) - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_batched_mm_passes_sentinel_expert_ids_unclamped(self): - # EP sentinels (expert_ids >= num_experts) reach the kernel unclamped; the post-mask zeroes the - # matching output rows before the per-token reduction (the kernel leaves them uninitialized). - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # 4 == num_experts -> sentinel - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - self.assertTrue(torch.equal(calls["batched_matmul"][0]["expert_ids"], top_k_index.reshape(-1))) - self.assertEqual(int(calls["batched_matmul"][0]["expert_ids"].max()), 4) - # Sentinel token-expert pairs contribute 0 to the reduction -> each token keeps only its non-sentinel weight. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - self.assertTrue(torch.equal(out, top_k_weights[:, :1].to(torch.float32).expand(3, 8).to(torch.bfloat16))) - - def test_batched_mm_rejects_static_activation_scheme(self): - # Static activation quant needs a per-tensor activation scale the batched kernel can't consume; - # this guards a genuine unsupported-config path, not a trivial type/device check. - experts = make_fp8_experts(activation_scheme="static") - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(), self.assertRaisesRegex(NotImplementedError, "activation_scheme='static'"): - fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - # ── fp8_grouped_mm_experts_forward ──────────────────────────────────────────────────────────── - - def test_grouped_mm_kernel_inputs_and_output(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - up, down = calls["grouped_matmul"] - # S = num_tokens * top_k selected pairs, gathered by the expert-sort permutation. - self.assertEqual(up["input"].shape, (6, 8)) - self.assertIs(up["weight"], experts.gate_up_proj) - self.assertIs(up["weight_scale"], experts.gate_up_proj_scale_inv) - # Down projection sees the gated activation: 2*inter (32) collapsed to inter (16). - self.assertEqual(down["input"].shape, (6, 16)) - self.assertIs(down["weight"], experts.down_proj) - self.assertIs(down["weight_scale"], experts.down_proj_scale_inv) - - # offsets / tokens_per_expert are the per-expert histogram over the sorted expert ids. - # CPU histc requires float input, CUDA requires int (matches the source's dispatch). - expert_ids_g, _ = torch.sort(top_k_index.reshape(-1)) - histc_input = expert_ids_g.float() if expert_ids_g.device.type == "cpu" else expert_ids_g.int() - expected_tpe = torch.histc(histc_input, bins=experts.num_experts, min=0, max=experts.num_experts - 1) - expected_offsets = torch.cumsum(expected_tpe, dim=0, dtype=torch.int32) - self.assertTrue(torch.equal(up["tokens_per_expert"], expected_tpe)) - self.assertTrue(torch.equal(up["offsets"], expected_offsets)) - # Both projections share the same offsets/tokens_per_expert tensors. - self.assertIs(down["offsets"], up["offsets"]) - self.assertIs(down["tokens_per_expert"], up["tokens_per_expert"]) - - # Output restored to original token order (inv_perm) then reduced; value-exact. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_grouped_mm_sentinels_dropped_from_histogram(self): - # Sentinels are left unclamped so the sort pushes them to the tail and histc(max=num_experts-1) - # drops them from tokens_per_expert -> no wasted GEMM rows; the post-mask zeroes their output. - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # three sentinels (== num_experts) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - tpe = calls["grouped_matmul"][0]["tokens_per_expert"] - # Only experts 0,1,2 got one token each; the 3 sentinels are absent from the histogram. - self.assertTrue(torch.equal(tpe, torch.tensor([1.0, 1.0, 1.0, 0.0], device=torch_device))) - self.assertEqual(int(tpe.sum()), 3) - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_grouped_mm_rejects_static_activation_scheme(self): - experts = make_fp8_experts(activation_scheme="static") - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(), self.assertRaisesRegex(NotImplementedError, "activation_scheme='static'"): - fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 660e9ac92d53..fb7963b72afb 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -45,18 +45,24 @@ from parameterized import parameterized import transformers.integrations.deepgemm as dg +import transformers.integrations.finegrained.layers as fg from transformers.integrations.deepgemm import ( deepgemm_bf16_experts_forward, deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, deepgemm_fp8_fp4_megamoe_experts_forward, ) +from transformers.integrations.finegrained import FineGrainedLinear from transformers.testing_utils import ( require_torch, + require_torch_gpu, require_torch_greater_or_equal, + require_torch_multi_gpu, torch_device, ) +from transformers.utils.quantization_config import BlockFp8Config, Mxfp8Config +from ..integrations.finegrained.utils import _Cfg, fake_finegrained_kernel, patch_finegrained_kernel from .test_utils import make_experts, make_fp8_experts @@ -676,3 +682,150 @@ def test_megamoe_requires_fp4_weights_and_process_group(self): @staticmethod def _megamoe_flag_set(module): return getattr(module, "_megamoe_transformed", False) + + +@require_torch +class DeepGemmFineGrainedDispatchTest(unittest.TestCase): + """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled scale is not + readable as row-major) and an SM100 perf one. 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.dict(os.environ, {"TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR": "0"}), + mock.patch.object(dg, "is_deepgemm_loadable", return_value=True), + mock.patch.object(dg, "is_sm100", return_value=sm100), + mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as deepgemm_linear, + ): + x = torch.randn(4, 64, dtype=torch.bfloat16, requires_grad=trains) + fg.finegrained_linear(x, w, s, block_size=[128, 128]) + return "deepgemm" if deepgemm_linear.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_a_linear_that_needs_a_gradient_prefers_triton(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2, trains=True), "triton") + + def test_swizzled_scales_never_reach_deepgemm(self): + """The correctness gate holds on any arch, not just where the perf gate does.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + + 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, quantization_config=BlockFp8Config(weight_block_size=(128, 128))) + 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_fp8_fp4_linear(hs, w, s, block_size=(128, 128)), + lambda: deepgemm_fp8_fp4_experts_forward(experts, hs, *routing), + lambda: deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing), + ] + with mock.patch.object(dg, "load_deepgemm_kernel") as load: + for call in calls: + with self.assertRaises(NotImplementedError): + call() + load.assert_not_called() + + 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 fails loudly rather than reading the permuted buffer as affine.""" + 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, quantization_config=Mxfp8Config()) + 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 DeepGemmMultiDeviceTest(unittest.TestCase): + """`disable_deepgemm_on_multi_device` flags the quantized modules by the devices they occupy: DeepGEMM's + kernels are bound to one CUDA context and corrupt across devices, but a model on one device keeps DeepGEMM + even with other GPUs visible.""" + + @staticmethod + def _linear(device): + return FineGrainedLinear(256, 256, quantization_config=BlockFp8Config()).to(device) + + @require_torch_multi_gpu + def test_multi_device_disables_deepgemm(self): + model = torch.nn.Module() + model.a = self._linear("cuda:0") + model.b = self._linear("cuda:1") + fg.disable_deepgemm_on_multi_device(model) + self.assertTrue(model.a._deepgemm_disabled) + self.assertTrue(model.b._deepgemm_disabled) + + @require_torch_gpu + def test_single_device_keeps_deepgemm(self): + model = torch.nn.Module() + model.a = self._linear("cuda:0") + model.b = self._linear("cuda:0") + fg.disable_deepgemm_on_multi_device(model) + self.assertFalse(model.a._deepgemm_disabled) + self.assertFalse(model.b._deepgemm_disabled) + + +@require_torch +class DeepGemmQuantizedExpertsOperandsTest(unittest.TestCase): + """The quantized DeepGEMM arms read the stacked weights and their block scales and nothing else: whatever + else a module carries is refused before any kernel work rather than dropped.""" + + def _experts(self, quantization_config, impl="deepgemm", has_bias=False): + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, impl + experts = fg.FineGrainedExperts(cfg, quantization_config=quantization_config, has_bias=has_bias) + # `from_module` wraps the class's `__init__` to stamp the flags of the module it last replaced + experts.has_bias = has_bias + return experts + + def _assert_refused(self, experts, message): + hidden = torch.randn(2, 256, dtype=torch.bfloat16) + routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + with mock.patch.object(dg, "load_deepgemm_kernel") as load: + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.subTest(forward=forward.__name__), self.assertRaisesRegex(NotImplementedError, message): + forward(experts, hidden, *routing) + load.assert_not_called() + + def test_an_interleaved_gate_up_is_refused(self): + """A module loaded for a triton backend holds gate|up interleaved, its block grid interleaved per block, + which a row-major read pairs with the wrong half.""" + experts = self._experts(BlockFp8Config(weight_block_size=(128, 128)), impl="grouped_mm") + self.assertTrue(experts.holds_interleaved_gate_up) + self._assert_refused(experts, "stacked") + + def test_operands_the_kernels_would_drop_are_refused(self): + from transformers.utils.quantization_config import Mxfp4Config, NVFP4Config + + for experts, message in ( + (self._experts(BlockFp8Config(weight_block_size=(128, 128)), has_bias=True), "bias"), + (self._experts(NVFP4Config()), "NVFP4"), + (self._experts(Mxfp4Config(activation_format="bf16")), "weight-only"), + ): + with self.subTest(message=message): + self._assert_refused(experts, message) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py index 987c776ff9ae..7a4bf22216fa 100644 --- a/tests/kernels/test_finegrained.py +++ b/tests/kernels/test_finegrained.py @@ -12,17 +12,16 @@ # 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.""" +mocked, so these pin the loader and the call contract (As-positional / no-block_size / b_global_scale) 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 +import transformers.integrations.finegrained.layers as fg from transformers.integrations.finegrained import ( FineGrainedExperts, FineGrainedLinear, @@ -30,68 +29,9 @@ load_finegrained_kernel, ) from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import NVFP4Config -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), - ) +from ..integrations.finegrained.utils import _Cfg, _format_config, fake_finegrained_kernel, patch_finegrained_kernel @require_torch @@ -186,11 +126,7 @@ def test_module_forward_threads_everything(self): 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", + 64, 32, quantization_config=NVFP4Config(weight_block_size=None, activation_format="bf16") ) m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) out = m(torch.randn(2, 64, dtype=torch.bfloat16)) @@ -205,7 +141,9 @@ 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) + block = {"weight_block_size": (4, 4)} if weight_format == "fp8" else {} + quantized_type = _format_config(weight_format, activation_format=kw.pop("activation_format", None), **block) + m = FineGrainedExperts(cfg, quantization_config=quantized_type, **kw) for name, p in m.named_parameters(): if p is None: continue @@ -239,9 +177,7 @@ 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) - # 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["weight_global_scale"].shape, ()) # this expert's NVFP4 weight global 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"]) @@ -282,12 +218,13 @@ def test_grouped_marshalling(self): 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) + # the format's own activations read as `None` + for activation_format, handed in ((None, None), ("bf16", "bf16"), ("nvfp4", None)): + m = self._experts(has_gate=True, weight_format="nvfp4", 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) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], handed) def test_nvfp4_experts_thread_per_expert_globals(self): kernel, rec = fake_finegrained_kernel() @@ -305,13 +242,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 - 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,)) + # the calibrated activation globals (the checkpoint's `input_scale`): one per expert on both + # projections, each routed row quantized against its own expert's + self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_activation_scale) + self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_activation_scale) + self.assertEqual(tuple(m.gate_up_proj_activation_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_activation_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 @@ -323,8 +259,8 @@ def test_weight_only_experts_hold_no_activation_global(self): 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(m.gate_up_proj_activation_scale) + self.assertIsNone(m.down_proj_activation_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 @@ -357,29 +293,16 @@ 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_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.""" + """The norm follows `act_fn`'s shape: the name its class maps to 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: + with p1, p2, p3, mock.patch.dict(fg.SUPPORTED_FUSED_NORMS, {torch.nn.RMSNorm: name}): fg.finegrained_grouped_mm_experts_forward(m, *self._route()) call = rec.calls["moe_fused_grouped"][-1] if fused: @@ -403,34 +326,6 @@ 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_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) @@ -443,135 +338,3 @@ def test_unfusable_act_fn_is_passed_as_the_module_glu(self): # 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 8e75dc4fcde6..6dd8f9a3ad37 100644 --- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py +++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py @@ -448,12 +448,12 @@ def main() -> int: tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) inputs = tokenizer(PROMPT, return_tensors="pt", add_special_tokens=ADD_SPECIAL_TOKENS).to(model.device) - # FP8 only: Triton multi-expert kernels generate wrong tokens after DeepGEMM has run on + # FP8 experts only: Triton multi-expert kernels generate wrong tokens after DeepGEMM has run on # the same module (root cause unknown — per-row kernel outputs measure bit-perfect, yet # end-to-end generation degrades; not reproducible with DeepGEMM disabled). Force the - # Triton fallback in `fp8_linear` for batched_mm / grouped_mm so attention/MLP linears - # stay Triton-only too. FP4 is unaffected and keeps DeepGEMM for performance. - is_fp8 = getattr(model.config, "expert_dtype", "fp8") != "fp4" + # Triton dense-linear dispatch for batched_mm / grouped_mm so attention/MLP linears + # stay Triton-only too. FP4 experts are unaffected and keep DeepGEMM for performance. + is_fp8 = next((m.weight_format for m in model.modules() if hasattr(m, "holds_interleaved_gate_up")), None) == "fp8" failed = [] for dispatch in RUNTIME_DISPATCHES: @@ -528,7 +528,7 @@ def main() -> int: inputs = tokenizer(PROMPT, return_tensors="pt", add_special_tokens=ADD_SPECIAL_TOKENS).to(model.device) expected_logits_shape = inputs.input_ids.shape - is_fp8 = getattr(model.config, "expert_dtype", "fp8") != "fp4" + is_fp8 = next((m.weight_format for m in model.modules() if hasattr(m, "holds_interleaved_gate_up")), None) == "fp8" failed = [] for dispatch in RUNTIME_DISPATCHES: 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 c5335f780b41..868fa6eb2d4a 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, - FineGrainedConfig, + BlockFp8Config, 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 = FineGrainedConfig( + quantization_config = BlockFp8Config( 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 fd6c1b0c2518..fb8f5818ea93 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -28,7 +28,6 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, - FineGrainedConfig, is_torch_available, ) from transformers.testing_utils import ( @@ -45,6 +44,7 @@ slow, torch_device, ) +from transformers.utils.quantization_config import Mxfp4Config from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester @@ -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=FineGrainedConfig(quant_method="mxfp4", dequantize=True), + quantization_config=Mxfp4Config(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 e4d899ec3d74..c5fe9b32aa90 100644 --- a/tests/models/qwen4_exp/test_modeling_qwen4_exp.py +++ b/tests/models/qwen4_exp/test_modeling_qwen4_exp.py @@ -22,6 +22,7 @@ from transformers import is_torch_available from transformers.testing_utils import require_torch, torch_device +from transformers.utils.quantization_config import NVFP4Config from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester from ...test_modeling_common import floats_tensor @@ -46,8 +47,9 @@ ) from transformers.distributed.fsdp import verify_fsdp_plan from transformers.integrations.finegrained import FineGrainedEmbedding - from transformers.quantizers.finegrained import FineGrainedHfQuantizer - from transformers.utils.quantization_config import FineGrainedConfig + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig class Qwen4ExpTextModelTester(CausalLMModelTester): @@ -204,8 +206,8 @@ def test_finegrained_embedding_conversion(self): with torch.device("meta"): model = Qwen4ExpForCausalLM(config) - quantizer = FineGrainedHfQuantizer( - FineGrainedConfig(modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[]) + quantizer = BlockFp8HfQuantizer( + BlockFp8Config(modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[]) ) quantizer.pre_quantized = True quantizer.preprocess_model(model) @@ -215,24 +217,22 @@ def test_finegrained_embedding_conversion(self): 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.""" + def test_a_table_beside_another_format_is_held_in_fp8(self): + """An embedding table is FP8 whatever its model's other modules are: beside NVFP4 linears, the plan's + `modules_to_convert` gives it a block-FP8 entry of its own.""" config = self.model_tester.get_config() with torch.device("meta"): model = Qwen4ExpForCausalLM(config) quantizer = FineGrainedHfQuantizer( FineGrainedConfig( - quant_method="nvfp4", + dtype_plan={".*": NVFP4Config()}, 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) + quantizer.preprocess_model(model) 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): diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py new file mode 100644 index 000000000000..dab496bc5e3d --- /dev/null +++ b/tests/quantization/block_scaled_tester.py @@ -0,0 +1,1031 @@ +# 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 tests every block-scaled format runs, driven by its `FormatSpec`: a format's test file instantiates the +testers with its spec and adds only what is its own (its checkpoint layouts, its producers' quirks). + +- `BlockScaledFormatTesterMixin`: the config, the quantizer a load picks, the activation formats it takes, the + no-accelerator fallbacks, one converter per checkpoint key (CPU); every forward against the format's + dequantized weights and the on-the-fly quantization round trip (GPU, the real kernels). +- `TinyCheckpointTesterMixin` (slow): the format's pushed tiny checkpoints load with every key accounted for, + compute what their dequantized weights compute, match the recorded logits, and save back to their layout. + `FINEGRAINED` runs it on the checkpoints that mix formats. +- `LoadPathEquivalenceTesterMixin` (slow, 2 GPUs): expert and tensor parallelism reproduce `device_map`.""" + +import os +import re +import subprocess +import tempfile +from collections.abc import Callable +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from unittest import mock + +import torch + +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText +from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key +from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig +from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer +from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer +from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer +from transformers.quantizers.quantizer_mxfp8 import Mxfp8HfQuantizer +from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer +from transformers.testing_utils import ( + Expectations, + cleanup, + get_device_properties, + get_torch_dist_unique_port, + require_accelerate, + require_kernels, + require_torch_accelerator, + require_torch_multi_accelerator, + slow, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, + FineGrainedConfig, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) + +from ..integrations.finegrained.utils import _Cfg, _model_mapping, _quantizer_for + + +def _cheap_tuning(): + """One autotune trial per shape: a cold tune dwarfs everything a test runs.""" + return mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}) + + +@dataclass(frozen=True) +class FormatSpec: + """What the shared tests need to know about one format; no test hard-codes it. + + `quantize_kwargs` configure the on-the-fly quantization whose layout is checked (`None`: the format's + quantize op is tested on its own). `checkpoint_keys` are keys of the format's shipped layouts, each of which + at most one converter may claim; `dequantize_keys` are a quantized linear's keys, each of which a + dequantizing load must route to a converter. `tiny_checkpoints` maps each pushed tiny + checkpoint to its logits budget against its `reference/` bf16 copy; `distributed_models` names the + `LoadPathEquivalenceTesterMixin` fixtures built in this format.""" + + name: str + config_cls: type + quantizer_cls: type + checkpoint_config: Callable + activation_formats: tuple = () + weight_bits: int = 16 + two_level: bool = False + quantize_kwargs: dict | None = None + checkpoint_keys: tuple = () + dequantize_keys: tuple = () + tiny_checkpoints: dict = field(default_factory=dict) + distributed_models: tuple = () + + +BLOCK_FP8 = FormatSpec( + name="fp8", + config_cls=BlockFp8Config, + quantizer_cls=BlockFp8HfQuantizer, + activation_formats=(None, "fp8"), + checkpoint_config=BlockFp8Config, + weight_bits=8, + checkpoint_keys=( + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.gate_up_proj_scale_inv", + "model.layers.3.self_attn.q_proj.weight_scale_inv", + ), + dequantize_keys=("model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.weight_scale_inv"), + 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-DeepseekV3-FP8": 0.25, + }, + distributed_models=("deepseek_v3-fp8", "mistral4-fp8_tensor_static"), +) +MXFP8 = FormatSpec( + name="mxfp8", + config_cls=Mxfp8Config, + quantizer_cls=Mxfp8HfQuantizer, + activation_formats=(None, "mxfp8", "mxfp4", "bf16"), + checkpoint_config=Mxfp8Config, + weight_bits=8, + quantize_kwargs={}, + checkpoint_keys=( + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.down_proj_scale_inv", + ), + dequantize_keys=("model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.weight_scale_inv"), + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": 0.3}, + distributed_models=("minimax_m3_vl-mxfp8",), +) +MXFP4 = FormatSpec( + name="mxfp4", + config_cls=Mxfp4Config, + quantizer_cls=Mxfp4HfQuantizer, + activation_formats=(None, "mxfp8", "mxfp4", "bf16"), + checkpoint_config=Mxfp4Config, + weight_bits=4, + quantize_kwargs={"activation_format": None}, + checkpoint_keys=( + "model.layers.3.mlp.experts.gate_up_proj_blocks", + "model.layers.3.mlp.experts.gate_up_proj_scales", + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + ), + # GPT-OSS's blocks + scales pair + dequantize_keys=( + "model.layers.0.mlp.experts.gate_up_proj_blocks", + "model.layers.0.mlp.experts.gate_up_proj_scales", + ), + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-GptOss-MXFP4": 0.05}, + distributed_models=("gpt_oss-mxfp4",), +) +NVFP4 = FormatSpec( + name="nvfp4", + config_cls=NVFP4Config, + quantizer_cls=NVFP4HfQuantizer, + activation_formats=(None, "nvfp4", "bf16"), + # what a modelopt export reads as + checkpoint_config=lambda: AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4"}), + weight_bits=4, + two_level=True, + quantize_kwargs={}, + checkpoint_keys=( + "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", + ), + dequantize_keys=( + "model.layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.down_proj.weight_scale", + "model.layers.0.mlp.down_proj.weight_scale_2", + ), + tiny_checkpoints={ + "IlyasMoutawwakil/tiny-random-Llama-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": 0.6, + }, + distributed_models=("glm4v_moe-nvfp4",), +) +FORMATS = (BLOCK_FP8, MXFP8, MXFP4, NVFP4) + +# Plans mixing the formats, as DeepSeek-V4 ships one (block-FP8 over MXFP4 experts). +FINEGRAINED = FormatSpec( + name="finegrained", + config_cls=FineGrainedConfig, + quantizer_cls=FineGrainedHfQuantizer, + checkpoint_config=FineGrainedConfig, + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-DeepseekV4-Flash": 0.25}, + distributed_models=("deepseek_v4-fp4_experts+fp8_dense",), +) + + +@contextmanager +def no_accelerator(): + with ExitStack() as stack: + stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) + if hasattr(torch, "xpu"): + stack.enter_context(mock.patch("torch.xpu.is_available", return_value=False)) + stack.enter_context( + mock.patch("transformers.quantizers.block_scaled_utils.is_torch_xpu_available", return_value=False) + ) + yield + + +def _require_kernel_capability(test): + # below it the quantizer dequantizes instead, which is not what is under test + if torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9): + test.skipTest("fine-grained quantized compute needs compute capability >= 8.9") + + +class BlockScaledFormatTesterMixin: + """The contract every format keeps. A format's test file sets `spec` and inherits these.""" + + spec: FormatSpec + + def test_config_round_trips_through_a_dict(self): + config = self.spec.config_cls(modules_to_not_convert=["lm_head"]) + read = self.spec.config_cls.from_dict(config.to_dict()) + self.assertEqual(read.to_dict(), config.to_dict()) + self.assertEqual(read.quant_method.value, self.spec.name) + + def test_from_pretrained_picks_the_format_quantizer(self): + quantizer = AutoHfQuantizer.from_config(self.spec.config_cls()) + self.assertIsInstance(quantizer, self.spec.quantizer_cls) + + def test_the_config_takes_only_the_activation_formats_its_kernels_serve(self): + """The kernels quantize activations only to what the weight format's MMA takes, and refuse the rest at + the first forward; the config refuses them up front.""" + for activation_format in {None} | {f for spec in FORMATS for f in spec.activation_formats}: + with self.subTest(activation_format=activation_format): + if activation_format in self.spec.activation_formats: + self.spec.config_cls(activation_format=activation_format) + else: + with self.assertRaises(ValueError): + self.spec.config_cls(activation_format=activation_format) + + @require_accelerate + def test_no_accelerator_refuses_to_quantize(self): + # quantizing a bf16 checkpoint needs the hardware; there is nothing to fall back to + quantizer = _quantizer_for(self.spec.config_cls()) + quantizer.pre_quantized = False + with no_accelerator(), self.assertRaises(RuntimeError): + quantizer.validate_environment() + + @require_accelerate + def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + with no_accelerator(): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + @require_accelerate + def test_a_gpu_below_compute_capability_8_9_dequantizes_a_quantized_checkpoint(self): + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + with ( + mock.patch("torch.cuda.is_available", return_value=True), + mock.patch("torch.cuda.get_device_capability", return_value=(8, 0)), + ): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + def test_a_checkpoint_key_is_claimed_by_at_most_one_converter(self): + """Two converters claiming one key is not additive: the later one wins and drops the ops the first carried. + Keys are routed as the loader routes them, renamed first.""" + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions(_model_mapping()) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + for key in self.spec.checkpoint_keys: + renamed, _ = rename_source_key(key, renamings, converters) + claimed = [c for c in converters if c.compiled_sources.search(renamed)] + with self.subTest(key=key): + 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_dequantize_claims_every_key_of_a_quantized_linear(self): + """A dequantizing load swaps in no fine-grained module, so no scale slot exists: every key of the pair + (and the global, for a two-level format) must reach a converter.""" + quantizer = _quantizer_for(self.spec.config_cls(dequantize=True)) + quantizer.pre_quantized = True + conversions = quantizer.get_weight_conversions() + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + for key in self.spec.dequantize_keys: + renamed, _ = rename_source_key(key, renamings, converters) + with self.subTest(key=key): + self.assertTrue( + any(c.compiled_sources.search(renamed) for c in converters), f"nothing dequantizes {key}" + ) + + def test_module_sizes_count_the_quantized_weights(self): + """The device map is sized before any weight loads, from the replaced modules' dtypes and shapes.""" + from transformers import Qwen3MoeConfig + from transformers.integrations.accelerate import compute_module_sizes + from transformers.modeling_utils import expand_device_map, get_total_byte_count + + config = Qwen3MoeConfig( + 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=2, + head_dim=64, + num_experts=4, + num_experts_per_tok=2, + ) + with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) + sizes, _ = compute_module_sizes(model, only_modules=False) + quantizer = AutoHfQuantizer.from_config(self.spec.config_cls(), pre_quantized=False) + quantizer.preprocess_model(model=model, config=model.config) + quantized_sizes, _ = compute_module_sizes(model, quantizer, only_modules=False) + keys = [name for name, _ in model.named_parameters()] + [name for name, _ in model.named_buffers()] + (total,) = get_total_byte_count(model, expand_device_map({"": torch_device}, keys), quantizer).values() + + weights = [ + f"{name}.{param}" + for name, module in model.named_modules() + if isinstance(module, (FineGrainedLinear, FineGrainedExperts)) + for param in (("weight",) if isinstance(module, FineGrainedLinear) else ("gate_up_proj", "down_proj")) + ] + self.assertTrue(weights) + for weight in weights: + self.assertEqual(quantized_sizes[weight], sizes[weight] * self.spec.weight_bits // 16, weight) + self.assertEqual(total, quantized_sizes[""]) + + @require_kernels + @require_torch_accelerator + def test_on_the_fly_quantization_lands_in_the_module_layout(self): + """The kernels' quantizers then the module's layout ops fill the held layout (swizzled where it is held + so), with the second-level global only for a two-level format.""" + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + if self.spec.quantize_kwargs is None: + self.skipTest(f"{self.spec.name}'s quantize op is tested on its own") + _require_kernel_capability(self) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + experts = FineGrainedExperts(cfg, quantization_config=self.spec.config_cls(**self.spec.quantize_kwargs)) + experts = experts.to(torch_device) + model = torch.nn.Module() + model.experts = experts + original = torch.randn(4, 512, 512, device=torch_device, dtype=torch.bfloat16) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + for name in ("gate_up_proj", "gate_up_proj_scale_inv"): + for op in _quantizer_for(self.spec.config_cls()).get_param_ops(model, f"experts.{name}"): + key = f"experts.{name}" + out.update(op.convert({key: out[key]}, model=model, full_layer_name=key)) + for name in ("gate_up_proj", "gate_up_proj_scale_inv"): + held = getattr(experts, name) + self.assertEqual((out[f"experts.{name}"].dtype, out[f"experts.{name}"].shape), (held.dtype, held.shape)) + global_key = "experts.gate_up_proj_weight_global_scale" + if self.spec.two_level: + self.assertEqual(out[global_key].shape, experts.gate_up_proj_weight_global_scale.shape) + else: + self.assertNotIn(global_key, out) + + +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 + + 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(): + tensor = handle.get_tensor(key) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + tensors[key] = (tensor.dtype, tuple(tensor.shape), raw) + return tensors + + +# The logits `[0, :3, :8]` each checkpoint gives on `torch.arange(3, 35)`, per device: 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.3086, -0.2773, -0.4883, -0.1797, -0.0073, 0.4531, 0.2158, 0.2754], + [0.0295, -0.2598, -0.5586, 0.2432, -0.0106, 0.3496, 0.4805, 0.1279], + [0.1533, -0.4551, -0.5742, -0.1973, 0.2139, 0.4141, 0.3574, 0.1338], + ], + } + ), # 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 +} + + +class TinyCheckpointTesterMixin: + """A format's pushed tiny checkpoints, random copies of the shipped ones written in their producers' + layouts (modelopt, OpenAI's and DeepSeek's quantizers): each loads with every key accounted for, computes + what its dequantized weights compute within the format's rounding, matches its recorded logits, and saves + back to its own layout. So a converter is checked against what the format really is, not against what this + code believes it to be. A format's test file sets `spec`.""" + + spec: FormatSpec + + def setUp(self): + super().setUp() + self.enterContext(_cheap_tuning()) + _require_kernel_capability(self) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + @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) + + @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)} + + def _reference_error(self, model, repo): + """The logits' error against the dequantized reference, and the logits.""" + 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(): + got, want = model(ids).logits.float(), reference(ids).logits.float() + return ((got - want).norm() / want.norm()).item(), got + + @staticmethod + def _recorded_slice(repo): + """The slice recorded on this exact device, if any: another device's kernels round differently.""" + return dict(EXPECTED_SLICES[repo].unpacked()).get(get_device_properties()) + + @slow + @require_kernels + @require_torch_accelerator + def test_logits_match_the_dequantized_reference_and_the_recorded_slice(self): + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + 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()) + ) + error, logits = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + expected = self._recorded_slice(repo) + if expected is not None: + torch.testing.assert_close(logits[0, :3, :8].cpu(), torch.tensor(expected), rtol=1e-2, atol=5e-3) + + @slow + @require_kernels + @require_torch_accelerator + def test_dequantized_load_matches_the_reference(self): + """`dequantize=True` folds each weight against its scales (and its global, for a two-level format).""" + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + model, info = self._load(repo, quantization_config=self.spec.config_cls(dequantize=True)) + unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) + self.assertEqual((set(info["missing_keys"]), unexpected), (set(), set())) + # nothing stays quantized; a model keeps its own fp32 parameters (DeepSeek-V4's hyper-connections) + quantized = {torch.float8_e4m3fn, torch.float8_e8m0fnu, torch.int8, torch.uint8} + self.assertEqual([name for name, param in model.named_parameters() if param.dtype in quantized], []) + error, _ = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + + @slow + @require_kernels + @require_torch_accelerator + def test_save_writes_the_checkpoint_layout_back(self): + """A pre-quantized model saved again lands on its checkpoint's keys, dtypes and bytes.""" + from huggingface_hub import snapshot_download + + for repo in self.spec.tiny_checkpoints: + with self.subTest(repo=repo): + 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(repo=repo, key=key): + self.assertEqual((dtype, shape), source[key][:2]) + self.assertEqual(raw, source[key][2]) + + +_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 and differ only in the plan +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, experts.gate_up_name) + 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 fused (`experts.gate_up_proj`) or one per expert + (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`).""" + import glob + import json + + 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 + + +class LoadPathEquivalenceTesterMixin: + """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. + + 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 not run. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + + A format's test file sets `spec`; the fixtures are its `distributed_models`.""" + + spec: FormatSpec + + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}`: each model in the format and + under the quantization config it ships with. 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, + ) + + 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, + }, + BlockFp8Config(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, + }, + BlockFp8Config(weight_block_size=None, activation_scheme="static"), + ), + # multimodal: the experts' plans live on `text_config`, not on the config the quantizer is handed + "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, + ), + }, + NVFP4.checkpoint_config(), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases its EP plan already names + "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, + }, + Mxfp4Config(activation_format="bf16"), + ), + # mixed: `expert_dtype` makes the experts MXFP4 over block-FP8 with UE8M0 scales + "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", + }, + BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second multimodal nesting; the MoE dims are raised so a 2-way split stays 128-aligned + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + # the 2-way split leaves 256: a swizzled tile is 128 wide along the down's K + "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": { + # the 2-way split leaves 128, one swizzled K tile + "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, + }, + Mxfp8Config(), + ), + } + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) + cls._env = {"FINEGRAINED_AUTOTUNE_TRIALS": "1", "TRITON_CACHE_DIR": os.path.join(cls._tmp, "triton")} + cls.enterClassContext(mock.patch.dict(os.environ, cls._env)) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + 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, label, "bf16") + quant_dir = os.path.join(self._tmp, label, "quantized") + torch.manual_seed(0) + # saved in bf16, the dtype every model here ships, so `dtype="auto"` reloads bf16 + 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) + + # a reload saved again lands on the same bytes: every layout op's reverse restores the checkpoint + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") + round_trip = os.path.join(self._tmp, 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. How far + that travels is a property of the model (a chain of MoE layers can amplify it a hundredfold), so the + same magnitude injected on one device gives each model the budget its own conditioning earns. + """ + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" + ).eval() + # every block a sharded reduction passes through: the rounding accumulates down the stack + 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, "sharding_worker.py") + with open(script, "w", encoding="utf-8") as fh: + fh.write(_SHARDING_WORKER) + out = os.path.join(self._tmp, "out") + os.makedirs(out, exist_ok=True) + subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + f"--master_port={get_torch_dist_unique_port()}", + script, + model_dir, + out, + ",".join(modes), + f"{model_cls.__module__}:{model_cls.__name__}", + ], + check=True, + env={**os.environ, "TOKENIZERS_PARALLELISM": "false"}, + ) + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} + + @slow + @require_kernels + @require_torch_multi_accelerator + def test_every_load_path_agrees(self): + for label in self.spec.distributed_models: + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # `max_memory` forces the baseline across both devices; `auto` alone fits it on one + 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 whose plans shard experts (a multimodal model keeps them 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 shape, never another mode's shard + 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", + ) + # a wrong shard axis lands orders of magnitude above the model's rounding floor + 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: # rounding alone does not 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/quantization/blockfp8/__init__.py b/tests/quantization/blockfp8/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/blockfp8/test_blockfp8.py b/tests/quantization/blockfp8/test_blockfp8.py new file mode 100644 index 000000000000..aac173eb4a4e --- /dev/null +++ b/tests/quantization/blockfp8/test_blockfp8.py @@ -0,0 +1,529 @@ +# 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. +"""Block-FP8 (`BlockFp8Config`, `quant_method="fp8"`): the config, the quantize op and the checkpoint +layouts it loads and saves back, and model-level loading.""" + +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch +from parameterized import parameterized + +import transformers.integrations.finegrained.layers as fg +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.integrations.finegrained import FineGrainedExperts +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + get_device_properties, + require_accelerate, + require_torch, + require_torch_accelerator, + require_torch_multi_accelerator, + slow, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, +) + +from ...integrations.finegrained.utils import ( + E, + H, + I, + _Cfg, + _ChainAssertions, + _compressed_tensors, + _load_and_save, + _per_expert, + _quantizer_for, +) +from ...test_memory_cleanup_mixin import MemoryCleanupMixin +from ..block_scaled_tester import ( + BLOCK_FP8, + BlockScaledFormatTesterMixin, + LoadPathEquivalenceTesterMixin, + TinyCheckpointTesterMixin, + no_accelerator, +) + + +@require_torch +class BlockFp8ConfigTest(unittest.TestCase): + def test_a_passed_config_overrides_the_checkpoints_loading_attributes(self): + from transformers.quantizers.auto import AutoHfQuantizer + + passed = BlockFp8Config(dequantize=True, modules_to_not_convert=["lm_head", "model.layers.0"]) + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "fp8", "modules_to_not_convert": ["lm_head"]}, passed + ) + self.assertEqual((merged.dequantize, merged.modules_to_not_convert), (True, ["lm_head", "model.layers.0"])) + + def test_a_compressed_tensors_block_fp8_checkpoint_reads_as_block_fp8(self): + block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} + fp8 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 128, "dynamic": True} + config = AutoQuantizationConfig.from_dict(_compressed_tensors(block, fp8, fmt="float-quantized")) + self.assertEqual((type(config), config.weight_block_size), (BlockFp8Config, (128, 128))) + + +@require_torch +class BlockFp8QuantizeOpTest(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_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, quantization_config=BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static") + ) + model.proj = fg.FineGrainedLinear( + quantization_config=BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static"), + in_features=256, + out_features=256, + ) + 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, quantization_config=BlockFp8Config(weight_block_size=(128, 128)), 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 = _quantizer_for(BlockFp8Config(weight_block_size=(128, 128))) + 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, quantization_config=BlockFp8Config(weight_block_size=(128, 128), 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_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 BlockFp8ChainTest(_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, BlockFp8Config(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_ue8m0_block_scales(self): + """Block scales shipped as E8M0 load as `float8_e8m0fnu` and save back unchanged. (DeepSeek-V4's ship under + `.scale`, which its own conversion mapping renames.)""" + 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 // 16), -(-k // 16)), dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + } + ) + config = BlockFp8Config(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 = BlockFp8Config(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.core_model_loading import WeightConverter, WeightRenaming, rename_source_key + + quantizer = _quantizer_for(BlockFp8Config(activation_scheme="static")) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions([]) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + + 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", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.up_proj.input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.down_proj.input_scale", "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]): + # as the loader routes it: every renaming, then the one converter that matches + renamed, pattern = rename_source_key(key, renamings, converters) + targets = ( + [renamed] + if pattern is None + else [t for c in converters if pattern in c.source_patterns for t in c.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(BlockFp8Config()) + dynamic.pre_quantized = True + self.assertEqual(dynamic.get_weight_conversions(), []) + + +@slow +@require_accelerate +@require_torch_accelerator +@unittest.skipIf( + get_device_properties()[0] == "cuda" + and (get_device_properties()[1] < 8 or (get_device_properties()[1] == 8 and get_device_properties()[2] < 9)), + "Skipping BlockFp8QuantizerTest because it is not supported on GPU with capability < 8.9", +) +class BlockFp8QuantizerTest(MemoryCleanupMixin, unittest.TestCase): + model_name = "meta-llama/Llama-3.2-1B" + quantized_model_name = "hf-internal-testing/Llama-3.2-1B-Instruct-fp8" + input_text = "Once upon a time" + max_new_tokens = 10 + EXPECTED_OUTPUTS = { + "Once upon a time, there was a little girl who loved to play", + "Once upon a time, there was a man who was very rich.", + } + EXPECTED_DEQUANTIZED_OUTPUT = "Once upon a time, in a small village nestled in the rolling hills" + device_map = torch_device + offload_device_map = { + "model.embed_tokens": 0, + "model.layers.0": 0, + "model.layers.1": 0, + "model.layers.2": 0, + "model.layers.3": 0, + "model.layers.4": 0, + "model.layers.5": 0, + "model.layers.6": 0, + "model.layers.7": "cpu", + "model.layers.8": "cpu", + "model.layers.9": "cpu", + "model.layers.10": "cpu", + "model.layers.11": "cpu", + "model.layers.12": "cpu", + "model.layers.13": "cpu", + "model.layers.14": "cpu", + "model.layers.15": "cpu", + "model.rotary_emb": "cpu", + "model.norm": "cpu", + "lm_head": 0, + } + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.quantization_config = BlockFp8Config() + cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name) + cls.quantized_model = AutoModelForCausalLM.from_pretrained( + cls.model_name, device_map=cls.device_map, quantization_config=cls.quantization_config + ) + + @parameterized.expand( + [ + "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM", + "hf-internal-testing/tiny-random-MixtralForCausalLM", + ] + ) + def test_moe_conversion(self, model_id): + quantization_config = BlockFp8Config(weight_block_size=(32, 32)) + model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config) + self.assertTrue(any(isinstance(module, FineGrainedExperts) for module in model.modules())) + + def test_quantized_model(self): + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = self.quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertIn(output_tokens, self.EXPECTED_OUTPUTS) + + def test_dequantized_model(self): + quantization_config = BlockFp8Config(dequantize=True) + dequantized_model = AutoModelForCausalLM.from_pretrained( + self.quantized_model_name, device_map=self.device_map, quantization_config=quantization_config + ) + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + output = dequantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertEqual(output_tokens, self.EXPECTED_DEQUANTIZED_OUTPUT) + del dequantized_model + + def test_dequantize_when_no_accelerator(self): + with no_accelerator(): + dequantized_model = AutoModelForCausalLM.from_pretrained(self.quantized_model_name, device_map="cpu") + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to("cpu") + output = dequantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertEqual(output_tokens, self.EXPECTED_DEQUANTIZED_OUTPUT) + del dequantized_model + + def test_save_pretrained(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + + model = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.device_map) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + def test_weight_and_weight_scale_inv(self): + weight = self.quantized_model.model.layers[0].self_attn.q_proj.weight + weight_scale_inv = self.quantized_model.model.layers[0].self_attn.q_proj.weight_scale_inv + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(weight_scale_inv.dtype, torch.float32) + self.assertEqual(weight.shape, (weight_scale_inv.shape[0] * 128, weight_scale_inv.shape[1] * 128)) + + def test_block_size(self): + self.assertEqual(self.quantized_model.config.quantization_config.weight_block_size, (128, 128)) + quantization_config = BlockFp8Config(weight_block_size=(32, 32)) + quantized_model = AutoModelForCausalLM.from_pretrained( + self.model_name, device_map=self.device_map, quantization_config=quantization_config + ) + self.assertEqual(quantized_model.config.quantization_config.weight_block_size, (32, 32)) + + @require_torch_multi_accelerator + def test_quantized_model_multi_accelerators(self): + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + quantization_config = BlockFp8Config() + # need to empty cache or set max_memory, otherwise we will use the reserved memory that was not allocated when computing max-memory + # this will lead to put the entire model to device 0. + quantized_model = AutoModelForCausalLM.from_pretrained( + self.model_name, + device_map="auto", + quantization_config=quantization_config, + max_memory={0: "1GB", 1: "10GB"}, + ) + self.assertTrue(set(quantized_model.hf_device_map.values()) == {0, 1}) + + output = quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + @require_torch_multi_accelerator + def test_save_pretrained_multi_accelerators(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + # need to empty cache or set max_memory, otherwise we will use the reserved memory that was not allocated when computing max-memory + # this will lead to put the entire model to device 0. + model = AutoModelForCausalLM.from_pretrained( + tmpdirname, device_map="auto", max_memory={0: "1GB", 1: "10GB"} + ) + self.assertTrue(set(model.hf_device_map.values()) == {0, 1}) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + def test_quantized_model_offload(self): + with self.assertRaisesRegex( + 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 + ) + + def test_save_pretrained_offload(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + quantized_model = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.offload_device_map) + output = quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + @parameterized.expand(["eager", "batched_mm", "grouped_mm", "deepgemm"]) + def test_quantized_moe_forward(self, experts_implementation): + """MiniMax-M2 routes indices past `top_k`; every experts implementation has to run them.""" + if experts_implementation == "deepgemm" and torch_device != "cuda": + self.skipTest("deepgemm is only supported on CUDA") + + model = AutoModelForCausalLM.from_pretrained( + "hf-internal-testing/MiniMax-M2-Tiny-FP8", # single layer version + experts_implementation=experts_implementation, + device_map=self.device_map, + ) + assert model.config._experts_implementation == experts_implementation + + tokenizer = AutoTokenizer.from_pretrained("MiniMaxAI/MiniMax-M2") + messages = [ + {"role": "user", "content": [{"type": "text", "text": "What is your favourite condiment?"}]}, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!", + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "Do you have mayonnaise recipes?"}]}, + ] + model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to( + self.device_map + ) + + # Only caring about this not crashing + _ = model.generate(**model_inputs, max_new_tokens=24) + + +class BlockFp8Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = BLOCK_FP8 + + +class BlockFp8TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = BLOCK_FP8 + + +class BlockFp8LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = BLOCK_FP8 diff --git a/tests/quantization/finegrained/__init__.py b/tests/quantization/finegrained/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py new file mode 100644 index 000000000000..06faae1cae36 --- /dev/null +++ b/tests/quantization/finegrained/test_finegrained.py @@ -0,0 +1,598 @@ +# 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 quantizer (`FineGrainedConfig`): a plan mixing the format quantizers, the checkpoints that +read into one, and the conversions it builds. Each format's own coverage lives in its folder.""" + +import os +import tempfile +import unittest +import warnings +from types import SimpleNamespace +from unittest import mock + +import torch +from parameterized import parameterized +from torch import nn + +import transformers.integrations.finegrained.layers as fg +from transformers import AutoModelForCausalLM, PreTrainedConfig, PreTrainedModel +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + cleanup, + require_kernels, + require_torch, + require_torch_accelerator, + slow, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, + BlockScaledConfig, + FineGrainedConfig, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) + +from ...integrations.finegrained.utils import _Cfg, _model_config, _modelopt, _quantizer_for, _targets +from ..block_scaled_tester import ( + FINEGRAINED, + FORMATS, + LoadPathEquivalenceTesterMixin, + TinyCheckpointTesterMixin, + _cheap_tuning, +) + + +@require_torch +class FineGrainedPlanTest(unittest.TestCase): + """`FineGrainedConfig`: a plan of formats per module group, how a checkpoint that mixes formats reads into one, + and what each plan entry loads.""" + + def test_a_legacy_expert_dtype_becomes_the_two_entries_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.quantizers.auto import finegrained_quantization_config + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = finegrained_quantization_config(BlockFp8Config(scale_fmt="ue8m0"), _model_config(expert_dtype="fp4")) + self.assertIsInstance(config, FineGrainedConfig) + 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(plan_entry_for(config.dtype_plan, module).method.value, expected, module) + # without one, block-FP8 stays what it is + plain = BlockFp8Config() + self.assertIs(finegrained_quantization_config(plain, _model_config()), plain) + + def test_deepseek_v4_takes_one_config_whether_its_checkpoint_is_from_the_hub_or_a_save(self): + """Both sides are read for the model before they are compared, so the class a user passes does not depend on + whether the checkpoint still says `fp8` (the Hub) or already `finegrained` (a `save_pretrained`).""" + from transformers.quantizers.auto import AutoHfQuantizer, finegrained_quantization_config + + model_config = _model_config(expert_dtype="fp4") + hub = {"quant_method": "fp8", "scale_fmt": "ue8m0", "weight_block_size": [128, 128]} + saved = finegrained_quantization_config(AutoQuantizationConfig.from_dict(hub), model_config).to_dict() + for checkpoint in (hub, saved): + resolved = finegrained_quantization_config(AutoQuantizationConfig.from_dict(checkpoint), model_config) + passed = finegrained_quantization_config(BlockFp8Config(dequantize=True), model_config) + self.assertTrue(AutoHfQuantizer.merge_quantization_configs(resolved, passed).dequantize) + + def test_an_expert_dtype_in_the_quantization_config_reads_as_the_same_plan(self): + """DeepSeek-V4.1 declares its FP4 experts in the quantization config instead of the model config: the same + plan, block-FP8 at the checkpoint's block over W4A4 MXFP4 experts, whichever config the load is handed.""" + from transformers.quantizers.auto import finegrained_quantization_config + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + shipped = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [32, 32], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4", + } + model_config = SimpleNamespace(quantization_config=shipped, get_text_config=lambda: SimpleNamespace()) + for passed in (AutoQuantizationConfig.from_dict(shipped), BlockFp8Config(dequantize=True)): + with self.subTest(passed=type(passed).__name__): + config = finegrained_quantization_config(passed, model_config) + self.assertIsInstance(config, FineGrainedConfig) + experts, rest = config.dtype_plan[EXPERTS_PATTERN], config.dtype_plan[".*"] + # W4A4: the weights' own activation format + self.assertEqual((type(experts), experts.activation_format), (Mxfp4Config, None)) + self.assertEqual( + (type(rest), rest.weight_block_size, rest.scale_fmt), + (BlockFp8Config, passed.weight_block_size, passed.scale_fmt), + ) + # the loading options are the plan's, not an entry's + self.assertEqual((config.dequantize, rest.dequantize), (passed.dequantize, False)) + + def test_a_mixed_plans_dequantize_reads_every_entrys_block(self): + """One dequantize chain serves the whole plan, so a grid that rounds up finds its block whichever entry names it.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.quantizers_utils import FUSED_STACKS + + config = FineGrainedConfig( + dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config(weight_block_size=(128, 128))}, + modules_to_convert=["embed_tokens"], + dequantize=True, + ) + quantizer = FineGrainedHfQuantizer(config, pre_quantized=True) + dequantize = [ + op + for c in quantizer.get_weight_conversions() + for op in getattr(c, "operations", ()) + if isinstance(op, FineGrainedDequantize) + ] + # the dense weights' converter and one per fused expert stack + self.assertEqual(len(dequantize), 1 + len(FUSED_STACKS)) + self.assertTrue(all(op.hf_quantizer is quantizer for op in dequantize)) + + 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 entries were not preferred outright.""" + import json + + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = FineGrainedConfig( + dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config(weight_block_size=(128, 128))} + ) + back = FineGrainedConfig.from_dict(json.loads(config.to_json_string())) + self.assertIsInstance(plan_entry_for(back.dtype_plan, "model.layers.0.mlp.experts"), NVFP4Config) + self.assertIsInstance(plan_entry_for(back.dtype_plan, "model.layers.0.self_attn.q_proj"), BlockFp8Config) + # a list survives JSON where a tuple does not, and the block size is compared as a tuple + self.assertEqual(back.dtype_plan[".*"].weight_block_size, (128, 128)) + + def test_two_targeted_entries_claiming_one_module_is_an_error(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + + plan = {r"\.experts": BlockFp8Config(weight_block_size=None), r"mlp\.": NVFP4Config()} + with self.assertRaises(ValueError): + plan_entry_for(plan, "model.layers.0.mlp.experts") + + def test_an_export_mixing_formats_is_a_finegrained_config(self): + config = _modelopt( + config_groups={ + "group_0": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}, "targets": ["re:.*mlp.*"]}, + "group_1": {"weights": {"num_bits": 8, "type": "float", "group_size": 32}, "targets": ["re:.*attn.*"]}, + }, + ) + self.assertIsInstance(config, FineGrainedConfig) + self.assertEqual( + {pattern: type(entry) for pattern, entry in config.dtype_plan.items()}, + {".*mlp.*": NVFP4Config, ".*attn.*": Mxfp8Config}, + ) + # no `input_activations` is weight-only, whatever the weights' own activation format + self.assertEqual(config.dtype_plan[".*mlp.*"].activation_format, "bf16") + + def test_expert_dtype_reaches_a_config_quantizing_on_the_fly(self): + """The checkpoint an on-the-fly quantization writes declares the model config's `expert_dtype`, and is read + back with it: the experts are quantized to MXFP4 then, so the save loads as what it holds.""" + from transformers.quantizers.auto import get_hf_quantizer + from transformers.quantizers.block_scaled_utils import BlockScaledQuantizer + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + model_config = PreTrainedConfig() + model_config.expert_dtype = "fp4" + with mock.patch.object(BlockScaledQuantizer, "validate_environment"): + quantizer, model_config, _ = get_hf_quantizer(model_config, BlockFp8Config(), None, True, {}) + self.assertIsInstance(quantizer.quantization_config, FineGrainedConfig) + self.assertIsInstance(quantizer.quantization_config.dtype_plan[EXPERTS_PATTERN], Mxfp4Config) + self.assertFalse(quantizer.pre_quantized) + + def test_a_passed_configs_activation_format_reaches_a_prequantized_checkpoint(self): + """The activation format picks the kernel arm the same weights run on, so a passed config sets it; the + checkpoint's format options stay its own, and a passed one that differs from them is warned about.""" + from transformers.quantizers.auto import AutoHfQuantizer, finegrained_quantization_config + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + with warnings.catch_warnings(): + warnings.simplefilter("error") + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "mxfp4"}, Mxfp4Config(activation_format="mxfp8") + ) + self.assertEqual((type(merged), merged.activation_format), (Mxfp4Config, "mxfp8")) + + # a plan's entries lend theirs to the checkpoint's own entries + model_config = _model_config(expert_dtype="fp4") + hub = {"quant_method": "fp8", "scale_fmt": "ue8m0", "weight_block_size": [128, 128]} + passed = FineGrainedConfig( + dtype_plan={ + EXPERTS_PATTERN: Mxfp4Config(activation_format="bf16"), + ".*": BlockFp8Config(weight_block_size=(64, 64)), + } + ) + with self.assertWarns(UserWarning): + merged = AutoHfQuantizer.merge_quantization_configs( + finegrained_quantization_config(AutoQuantizationConfig.from_dict(hub), model_config), passed + ) + experts, rest = merged.dtype_plan[EXPERTS_PATTERN], merged.dtype_plan[".*"] + self.assertEqual(experts.activation_format, "bf16") + self.assertEqual((rest.weight_block_size, rest.scale_fmt), ((128, 128), "ue8m0")) + + def test_a_passed_config_lends_only_the_activation_format_it_names(self): + """A config passed for its skip-list keeps the checkpoint's activation format; one naming a format NVFP4 + cannot run on a weight-only checkpoint (no calibrated activation globals) is refused.""" + from transformers.quantizers.auto import AutoHfQuantizer + + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "nvfp4", "activation_format": "bf16"}, NVFP4Config(modules_to_not_convert=["lm_head"]) + ) + self.assertEqual((merged.activation_format, merged.modules_to_not_convert), ("bf16", ["lm_head"])) + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "mxfp4", "activation_format": "mxfp8"}, Mxfp4Config(modules_to_not_convert=["lm_head"]) + ) + self.assertEqual(merged.activation_format, "mxfp8") + with self.assertRaises(ValueError): + AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "nvfp4", "activation_format": "bf16"}, NVFP4Config(activation_format=None) + ) + self.assertFalse([name for name in NVFP4Config().to_dict() if name.startswith("_")]) + + def test_the_formats_own_activation_name_is_native(self): + """One spelling for a format's own activations: `None`.""" + for config_class, own in ( + (BlockFp8Config, "fp8"), + (Mxfp8Config, "mxfp8"), + (Mxfp4Config, "mxfp4"), + (NVFP4Config, "nvfp4"), + ): + with self.subTest(config_class=config_class.__name__): + self.assertIsNone(config_class(activation_format=own).activation_format) + self.assertIsNone(_modelopt().activation_format) + + def test_options_the_format_fixes_are_refused_and_not_written(self): + """MXFP8, MXFP4 and NVFP4 fix their block, scale dtype and activation scheme; only block-FP8 sets them. A + `config.json` spelling the format's own values (MiniMax-M3's `weight_block_size: [1, 32]`) or the old + defaults reads back.""" + for config_class in (Mxfp8Config, Mxfp4Config, NVFP4Config): + for option in ({"weight_block_size": (128, 128)}, {"activation_scheme": "static"}): + with self.subTest(config_class=config_class.__name__, option=option): + with self.assertRaises(ValueError): + config_class(**option) + for written in ( + {"weight_block_size": None, "scale_fmt": "float", "activation_scheme": "dynamic"}, + {"weight_block_size": list(config_class.format_options["weight_block_size"])}, + ): + config = config_class.from_dict({"quant_method": config_class.method.value, **written}) + self.assertFalse(set(written) & set(config.to_dict())) + self.assertEqual(BlockFp8Config(scale_fmt="ue8m0").to_dict()["scale_fmt"], "ue8m0") + + def test_a_format_configs_options_are_keyword_only_and_unknown_ones_reported(self): + with self.assertRaises(TypeError): + BlockFp8Config((64, 64)) + with self.assertLogs("transformers.utils.quantization_config", level="INFO") as logs: + Mxfp8Config(activaton_format="bf16") + self.assertIn("activaton_format", "".join(logs.output)) + + def test_the_base_config_is_not_a_format(self): + with self.assertRaises(TypeError): + BlockScaledConfig() + + def test_a_plan_entry_is_a_format_config_without_loading_options(self): + """The loading options hold for the whole model, so an entry naming one is refused rather than ignored.""" + for entry in ( + BlockFp8Config(dequantize=True), + NVFP4Config(modules_to_not_convert=["lm_head"]), + BlockFp8Config(modules_to_convert=["embed_tokens"]), + {"quant_method": "mxfp8", "dequantize": True}, + ): + with self.subTest(entry=entry): + with self.assertRaises(ValueError): + FineGrainedConfig(dtype_plan={".*": entry}) + with self.assertRaises(ValueError): + FineGrainedConfig(dtype_plan={".*": {"quant_method": "awq"}}) + with self.assertRaises(TypeError): + FineGrainedConfig(dtype_plan={".*": "fp8"}) + + def test_the_plans_members_load_as_the_host_does(self): + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + config = FineGrainedConfig(dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config()}) + quantizer = FineGrainedHfQuantizer(config, pre_quantized=False) + self.assertEqual([member.pre_quantized for member in quantizer.members()], [False, False]) + + def test_a_loaded_model_is_finished_once_whatever_its_formats(self): + """Every format checks its modules and keeps DeepGEMM off a model spread over devices, once per load.""" + import transformers.integrations.finegrained as finegrained + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + for config in ( + FineGrainedConfig(dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config()}), + Mxfp8Config(), + ): + with self.subTest(config=type(config).__name__): + quantizer = ( + FineGrainedHfQuantizer(config) if isinstance(config, FineGrainedConfig) else _quantizer_for(config) + ) + with ( + mock.patch.object(finegrained, "assert_modules_are_quantized") as checked, + mock.patch.object(finegrained, "disable_deepgemm_on_multi_device") as deepgemm, + ): + quantizer._process_model_after_weight_loading(nn.Module()) + self.assertEqual((checked.call_count, deepgemm.call_count), (1, 1)) + + def test_the_tensor_parallel_plan_resolves_the_experts_entry_for_the_tp_glob(self): + """Under intra-expert TP, whether the expert scale grid splits with its weight is read off the experts' + plan entries for the modules the TP glob (`layers.*.mlp.experts`) names. The glob's one style serves every + layer, so a plan giving some of them a grid and others a per-tensor scale is refused.""" + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + def tp_plan(dtype_plan): + config = SimpleNamespace( + base_model_tp_plan={ + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + }, + num_hidden_layers=2, + _experts_implementation=None, + ) + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(dtype_plan=dtype_plan)) + return quantizer.update_tp_plan(config).base_model_tp_plan + + per_tensor = BlockFp8Config(weight_block_size=None) + plan = tp_plan({r"layers\.\d+\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + self.assertEqual(plan.get("layers.*.mlp.experts.gate_up_proj_scale_inv"), "moe_experts_colwise") + self.assertEqual(plan.get("layers.*.mlp.experts.down_proj_scale_inv"), "moe_experts_rowwise") + with self.assertRaises(ValueError): + tp_plan({r"layers\.1\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + + +@require_torch +class FineGrainedPlanConversionTest(unittest.TestCase): + """The conversions a plan of several formats builds.""" + + def test_input_scale_reaches_each_formats_own_slot(self): + """One checkpoint key, two formats: `input_scale` is a static FP8 module's activation scale and an + NVFP4 module's activation global, held in one slot either way. Only the format holding the experts + reads their keys, and two formats that both would refuse the plan.""" + from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + config = FineGrainedConfig( + dtype_plan={ + EXPERTS_PATTERN: NVFP4Config(activation_format="nvfp4"), + ".*": BlockFp8Config(weight_block_size=None, activation_scheme="static"), + }, + ) + quantizer = FineGrainedHfQuantizer(config) + quantizer.pre_quantized = True + fp8 = next(member for member in quantizer.members() if member.quantization_config is config.dtype_plan[".*"]) + fp8.holds_experts = False # what the replacement walk leaves it + conversions = quantizer.update_weight_conversions([]) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + + renamed, _ = rename_source_key("model.layers.0.self_attn.q_proj.input_scale", renamings, converters) + self.assertEqual(renamed, "model.layers.0.self_attn.q_proj.activation_scale") + _, pattern = rename_source_key("model.layers.0.mlp.experts.3.down_proj.input_scale", renamings, converters) + (converter,) = [c for c in converters if pattern in c.source_patterns] + self.assertIn("experts.down_proj_activation_scale", _targets(converter)) + self.assertIn("experts.down_proj_weight_global_scale", _targets(converter)) + + fp8.holds_experts = True + with self.assertRaises(ValueError): + quantizer.update_weight_conversions([]) + + @staticmethod + def _split(conversions): + from transformers.core_model_loading import WeightConverter, WeightRenaming + + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + return renamings, [c for c in conversions if isinstance(c, WeightConverter)] + + def test_a_rewritten_model_converter_keeps_its_placement(self): + """Qwen4-Exp concatenates its ~95 GiB n-gram table on CPU (`force_cpu`); the rewrite that adds the scales + keeps it there, both ways the quantizer rewrites a converter.""" + from transformers.core_model_loading import MergeModulelist, WeightConverter + + for dequantize in (False, True): + with self.subTest(dequantize=dequantize): + model_converter = WeightConverter( + "ngram_embeddings.*.weight", "ngram_embedding.weight", [MergeModulelist(dim=0)], force_cpu=True + ) + quantizer = _quantizer_for(BlockFp8Config(dequantize=dequantize)) + (rewritten, *_) = quantizer.update_weight_conversions([model_converter]) + self.assertTrue(rewritten.force_cpu) + + def test_the_fused_nvfp4_layout_reaches_the_dequantize_converter(self): + """vLLM's stacked layout names its scales `{proj}_weight_scale` / `{proj}_weight_scale_2`: under `dequantize` + they take the slot names the fused stack's dequantize converter collects, so the packed weight is folded.""" + from transformers.core_model_loading import rename_source_key + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + + quantizer = _quantizer_for(NVFP4Config(dequantize=True)) + renamings, converters = self._split(quantizer.update_weight_conversions([])) + prefix = "model.layers.0.mlp.experts" + for proj in ("gate_up_proj", "down_proj"): + for shipped, slot in (("_weight_scale", "_scale_inv"), ("_weight_scale_2", "_weight_global_scale")): + with self.subTest(key=f"{proj}{shipped}"): + renamed, pattern = rename_source_key(f"{prefix}.{proj}{shipped}", renamings, converters) + self.assertEqual(renamed, f"{prefix}.{proj}") + (converter,) = [c for c in converters if pattern in c.source_patterns] + self.assertTrue(any(isinstance(op, FineGrainedDequantize) for op in converter.operations)) + self.assertIn(rf"(? quantized_model_size[""] * 1.5 - - @parameterized.expand(["eager", "batched_mm", "grouped_mm", "deepgemm"]) - def test_quantized_moe_forward(self, experts_implementation): - """ - Checks implicitly if the moe implementation is correct, i.e. it does not crash for cases - where the indices go over `top_k` as shown within the Minimax M2 model - """ - # deepgemm only has CUDA kernels, skip on other devices - if experts_implementation == "deepgemm" and torch_device != "cuda": - self.skipTest("deepgemm is only supported on CUDA") - - model = AutoModelForCausalLM.from_pretrained( - "hf-internal-testing/MiniMax-M2-Tiny-FP8", # single layer version - experts_implementation=experts_implementation, - device_map=self.device_map, - ) - assert model.config._experts_implementation == experts_implementation - - tokenizer = AutoTokenizer.from_pretrained("MiniMaxAI/MiniMax-M2") - messages = [ - {"role": "user", "content": [{"type": "text", "text": "What is your favourite condiment?"}]}, - { - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!", - } - ], - }, - {"role": "user", "content": [{"type": "text", "text": "Do you have mayonnaise recipes?"}]}, - ] - model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to( - self.device_map - ) - - # Only caring about this not crashing - _ = model.generate(**model_inputs, max_new_tokens=24) - - -@require_torch_accelerator -@unittest.skipIf( - get_device_properties()[0] == "cuda" - and (get_device_properties()[1] < 8 or (get_device_properties()[1] == 8 and get_device_properties()[2] < 9)), - "Skipping FP8LinearTest because it is not supported on GPU with capability < 8.9", -) -class FP8LinearTest(unittest.TestCase): - device = torch_device - - def test_linear_preserves_shape(self): - """ - Test that FP8Linear preserves shape when in_features == out_features. - """ - from transformers.integrations import FP8Linear - - linear = FP8Linear(256, 256, block_size=(128, 128)).to(self.device) - x = torch.rand((1, 5, 256)).to(self.device) - - x_ = linear(x) - self.assertEqual(x_.shape, x.shape) - - def test_linear_with_diff_feature_size_preserves_shape(self): - """ - Test that FP8Linear generates the correct shape when in_features != out_features. - """ - from transformers.integrations import FP8Linear - - linear = FP8Linear(128, 256, block_size=(128, 128)).to(self.device) - x = torch.rand((1, 5, 128)).to(self.device) - - x_ = linear(x) - self.assertEqual(x_.shape, (1, 5, 256)) - - -class FP8EmbeddingTest(unittest.TestCase): - """ - Some checkpoints quantize an embedding *table* rather than a linear (Qwen4-Exp's n-gram/PLE - table). - """ - - def test_lookup_applies_the_per_tensor_scale(self): - """Without the rescale the rows come back as raw FP8 and blow up the next matmul.""" - from transformers.integrations import FP8Embedding - - table = torch.randn(64, 16, dtype=torch.bfloat16) - scale = table.abs().max().float() / torch.finfo(torch.float8_e4m3fn).max - embedding = FP8Embedding(64, 16) - embedding.weight = torch.nn.Parameter((table.float() / 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.randint(0, 64, (2, 5)) - rows, expected = embedding(ids), table[ids] - - self.assertEqual(rows.dtype, expected.dtype) - error = (rows.float() - expected.float()).abs().max() / expected.float().abs().max() - self.assertLess(error, 0.1) - - def test_patterns_match_by_suffix_and_honour_the_skip_list(self): - """Suffixes match whether or not the text model is nested under a multimodal wrapper.""" - from transformers.integrations import FP8Embedding, replace_with_fp8_embedding - - model = torch.nn.ModuleDict( - {n: torch.nn.ModuleDict({"table": torch.nn.Embedding(8, 4)}) for n in ("layers", "language_model")} - ) - model.other = torch.nn.Embedding(8, 4) - replace_with_fp8_embedding(model, ["table"], modules_to_not_convert=["language_model.table"]) - - self.assertIsInstance(model["layers"]["table"], FP8Embedding) # matched by suffix - self.assertIs(type(model["language_model"]["table"]), torch.nn.Embedding) # skip-listed - self.assertIs(type(model.other), torch.nn.Embedding) # no pattern match - - -class FP8DeepGEMMMultiDeviceTest(unittest.TestCase): - """`_disable_deepgemm_on_multi_device` must flag FP8 modules based on the devices they actually - occupy — DeepGEMM's kernels are bound to a single CUDA context and corrupt across devices, but a - model that fits on one device must keep DeepGEMM even when other GPUs are visible (no overshoot). - """ - - @staticmethod - def _fp8_module(device): - from transformers.integrations import FP8Linear - - return FP8Linear(256, 256, block_size=(128, 128)).to(device) - - @require_torch_multi_gpu - def test_multi_device_disables_deepgemm(self): - from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device - - model = torch.nn.Module() - model.a = self._fp8_module("cuda:0") - model.b = self._fp8_module("cuda:1") - _disable_deepgemm_on_multi_device(model) - self.assertTrue(model.a._deepgemm_disabled) - self.assertTrue(model.b._deepgemm_disabled) - - @require_torch_gpu - def test_single_device_keeps_deepgemm(self): - from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device - - model = torch.nn.Module() - model.a = self._fp8_module("cuda:0") - model.b = self._fp8_module("cuda:0") - _disable_deepgemm_on_multi_device(model) - self.assertFalse(model.a._deepgemm_disabled) - self.assertFalse(model.b._deepgemm_disabled) +class FineGrainedFP8LegacyTest(unittest.TestCase): + def test_building_one_is_deprecated(self): + with self.assertWarns(FutureWarning): + FineGrainedFP8Config() + + def test_a_passed_deprecated_config_loads_through_its_frozen_integration(self): + """On the fly and over an `fp8` checkpoint alike: the checkpoint's config is read as the passed one's class.""" + from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + + for checkpoint in (None, {"quant_method": "fp8", "weight_block_size": [128, 128]}): + with self.subTest(pre_quantized=checkpoint is not None): + quantizer = _quantizer(_legacy(dequantize=True), checkpoint) + self.assertIsInstance(quantizer, FineGrainedFP8HfQuantizer) + self.assertIsInstance(quantizer.quantization_config, FineGrainedFP8Config) + self.assertTrue(quantizer.quantization_config.dequantize) + + def test_an_fp8_checkpoint_loads_as_block_fp8(self): + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + quantizer = _quantizer(None, {"quant_method": "fp8", "weight_block_size": [128, 128]}) + self.assertIsInstance(quantizer, BlockFp8HfQuantizer) + self.assertIsInstance(quantizer.quantization_config, BlockFp8Config) diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index e9dbfbf2e202..b6450364f7c1 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -12,399 +12,101 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc +"""MXFP4 (`Mxfp4Config`): GPT-OSS's blocks layout and the packed per-expert one it loads and saves back, the +generate path, and the published GPT-OSS checkpoints.""" + import tempfile import unittest -from contextlib import ExitStack, contextmanager -from unittest.mock import patch -from transformers import AutoTokenizer, GptOssForCausalLM, Mxfp4Config +import torch + +import transformers.integrations.finegrained.layers as fg +from transformers import AutoTokenizer, GptOssConfig, GptOssForCausalLM, Mxfp4Config from transformers.testing_utils import ( + TestCasePlus, + cleanup, require_kernels, require_torch, + require_torch_accelerator, require_torch_accelerator_memory, require_torch_large_accelerator, - require_triton, slow, torch_device, ) -from transformers.utils import ( - is_torch_available, -) - - -if is_torch_available(): - import torch - - -if torch.cuda.is_available(): - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.4.0") -elif hasattr(torch, "xpu") and torch.xpu.is_available(): - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0") -elif torch_device == "cpu": - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0") -else: - REQUIRE_TRITON_MXFP4 = unittest.skip("test requires CUDA or XPU") - - -def _empty_accelerator_cache(): - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - torch.xpu.empty_cache() - - -@contextmanager -def _patch_no_accelerator(): - with ExitStack() as stack: - stack.enter_context(patch("torch.cuda.is_available", return_value=False)) - if hasattr(torch, "xpu"): - stack.enter_context(patch("torch.xpu.is_available", return_value=False)) - stack.enter_context(patch("torch.accelerator.current_accelerator", return_value=None)) - yield - - -class Mxfp4ConfigTest(unittest.TestCase): - def test_basic_config_creation(self): - """Test basic configuration creation with default values""" - config = Mxfp4Config() - self.assertEqual(config.quant_method.value, "mxfp4") - self.assertIsNone(config.modules_to_not_convert) - self.assertFalse(config.dequantize) - - def test_config_with_modules_to_not_convert(self): - """Test configuration with modules to not convert""" - modules = ["model.layers.*.self_attn", "lm_head"] - config = Mxfp4Config(modules_to_not_convert=modules) - self.assertEqual(config.modules_to_not_convert, modules) - - def test_config_with_dequantize(self): - """Test configuration with dequantize enabled""" - config = Mxfp4Config(dequantize=True) - self.assertTrue(config.dequantize) - - def test_get_loading_attributes(self): - """Test get_loading_attributes method""" - config = Mxfp4Config(dequantize=True) - attrs = config.get_loading_attributes() - self.assertEqual(attrs["dequantize"], True) - - def test_to_dict(self): - """Test configuration serialization to dict""" - config = Mxfp4Config(modules_to_not_convert=["lm_head"], dequantize=True) - config_dict = config.to_dict() - self.assertEqual(config_dict["quant_method"], "mxfp4") - self.assertEqual(config_dict["modules_to_not_convert"], ["lm_head"]) - # we don't keep dequantize in config_dict - self.assertTrue("dequantize" not in config_dict) - - def test_from_dict(self): - """Test configuration creation from dict""" - config_dict = {"quant_method": "mxfp4", "modules_to_not_convert": ["lm_head"], "dequantize": True} - config = Mxfp4Config.from_dict(config_dict) - self.assertEqual(config.modules_to_not_convert, ["lm_head"]) - self.assertTrue(config.dequantize) - - -class Mxfp4QuantizerTest(unittest.TestCase): - """Test the Mxfp4HfQuantizer class""" - - def setUp(self): - gc.collect() - _empty_accelerator_cache() - from transformers.utils.logging import warning_once - - warning_once.cache_clear() - - def test_quantizer_validation_no_torch(self): - """Test quantizer validation when torch is not available""" - with patch("transformers.quantizers.quantizer_mxfp4.is_torch_available", return_value=False): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - with self.assertRaises(ImportError): - quantizer.validate_environment() - - def test_quantizer_validation_no_accelerator(self): - """Test quantizer validation when CUDA/XPU is not available""" - with ( - _patch_no_accelerator(), - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.Mxfp4HfQuantizer._lazy_import_kernels"), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - # CPU already supported MXFP4 - quantizer.validate_environment() - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_quantizer_validation_low_compute_capability(self): - """Test quantizer validation with CUDA low compute capability or supported XPU""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with self.assertRaises(ValueError): - quantizer.validate_environment() - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_quantizer_validation_low_compute_capability_with_prequantized(self): - """Test pre-quantized validation with CUDA low compute capability or supported XPU""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - # Should automatically set dequantize=True and warn - quantizer.validate_environment() - self.assertTrue(quantizer.quantization_config.dequantize) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_quantizer_validation_low_compute_capability_with_dequantize(self): - """Test quantizer validation with dequantize enabled""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config(dequantize=True) - quantizer = Mxfp4HfQuantizer(config) - - # Should not raise error with dequantize=True - try: - quantizer.validate_environment() - except ValueError as e: - if "compute capability" in str(e): - self.fail("Should not raise compute capability error when dequantize=True") - - def test_quantizer_validation_order_dequantize_before_accelerator_check(self): - """Test that dequantize check happens before CUDA/XPU availability check""" - # Mock torch.cuda.is_available - with ( - _patch_no_accelerator(), - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.Mxfp4HfQuantizer._lazy_import_kernels"), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Test with dequantize=True - should pass even without CUDA/XPU and accelerate - config = Mxfp4Config(dequantize=True) - quantizer = Mxfp4HfQuantizer(config) - - # This should not raise any error because dequantize check comes first - quantizer.validate_environment() - - # Test with dequantize=False - should still fail due to missing CUDA/XPU - config = Mxfp4Config(dequantize=False) - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - # CPU already supported MXFP4 - quantizer.validate_environment() - - def test_quantizer_validation_missing_triton(self): - """Test quantizer validation when triton is not available""" - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - with self.assertRaises(ValueError): - quantizer.validate_environment() - - def test_quantizer_validation_missing_triton_pre_quantized_no_dequantize(self): - """Test quantizer validation when triton is not available but model is pre-quantized and dequantize is False""" - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - # Should automatically set dequantize=True and warn - quantizer.validate_environment() - self.assertTrue(quantizer.quantization_config.dequantize) - - def test_is_trainable(self): - """Test trainability""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - # MXFP4 is not trainable - self.assertFalse(quantizer.is_trainable) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_warning_distinguishes_triton_from_kernels(self): - """When only one dependency is missing, warning should mention it specifically.""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Missing kernels only -> warning should mention kernels - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - self.assertLogs("transformers", level="WARNING") as cm, - ): - quantizer.validate_environment() - - warning_text = " ".join(cm.output) - self.assertIn("kernels", warning_text.lower()) - self.assertTrue(quantizer.quantization_config.dequantize) - - # Missing triton only -> warning should mention triton - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - self.assertLogs("transformers", level="WARNING") as cm, - ): - quantizer.validate_environment() - - warning_text = " ".join(cm.output) - self.assertIn("triton", warning_text.lower()) - self.assertTrue(quantizer.quantization_config.dequantize) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_error_distinguishes_triton_from_kernels(self): - """When quantizing without a dependency, ValueError should mention it specifically.""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Missing kernels only -> error should mention kernels - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - with self.assertRaises(ValueError) as ctx: - quantizer.validate_environment() - - self.assertIn("kernels", str(ctx.exception).lower()) - - # Missing triton only -> error should mention triton - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - ): - with self.assertRaises(ValueError) as ctx: - quantizer.validate_environment() - self.assertIn("triton", str(ctx.exception).lower()) - - -class Mxfp4IntegrationTest(unittest.TestCase): - """Test mxfp4 integration functions""" - - def test_should_convert_module(self): - """Test module conversion decision logic""" - from transformers.quantizers.quantizers_utils import should_convert_module - - # Should convert by default - self.assertTrue(should_convert_module("model", None)) - self.assertTrue(should_convert_module("model", [])) - - # Should not convert if in exclusion list - patterns = ["model.layers.*.self_attn", "lm_head"] - self.assertFalse(should_convert_module("lm_head", patterns)) - self.assertTrue(should_convert_module("experts", patterns)) - - @require_torch - def test_convert_moe_packed_tensors(self): - """Test unpacking of quantized tensors""" - from transformers.integrations.mxfp4 import convert_moe_packed_tensors - - # Create dummy packed tensors - blocks = torch.randint(0, 255, (2, 4, 8, 16), dtype=torch.uint8) - scales = torch.randint(100, 150, (2, 4, 8), dtype=torch.uint8) - - result = convert_moe_packed_tensors(blocks, scales, dtype=torch.bfloat16) - self.assertEqual(result.shape, (2, 8 * 16 * 2, 4)) - self.assertEqual(result.dtype, torch.bfloat16) - - @REQUIRE_TRITON_MXFP4 - @require_kernels - @require_torch - def test_quantize_to_mxfp4(self): - """Test quantization function""" - from transformers.integrations.mxfp4 import quantize_to_mxfp4 - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) +from ...integrations.finegrained.utils import _ChainAssertions, _load_and_save, _per_expert, _quantizer_for +from ..block_scaled_tester import ( + MXFP4, + BlockScaledFormatTesterMixin, + LoadPathEquivalenceTesterMixin, + TinyCheckpointTesterMixin, +) - # Create dummy weight tensor - device = torch_device - w = torch.randn(32, 64, 128, dtype=torch.bfloat16, device=torch.device(device)) - quantized_w, w_scale = quantize_to_mxfp4(w, quantizer._lazy_import_kernels()) +@require_torch_accelerator +@require_kernels +class Mxfp4GenerateTest(unittest.TestCase): + def test_generate_matches_forward_logits(self): + """`generate` switches the experts to the batched kernels, while a plain forward runs the grouped ones: on + a tiny GPT-OSS with biased experts, quantized on the fly (W4A16), every decoding step's logits must match + a forward over the same prefix. Large biases, so a kernel that drops one cannot pass.""" + torch.manual_seed(0) + config = GptOssConfig( + vocab_size=512, + hidden_size=256, + intermediate_size=256, + head_dim=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + ) + model = GptOssForCausalLM(config).to(torch.bfloat16) + with torch.no_grad(): + for layer in model.model.layers: + layer.mlp.experts.gate_up_proj_bias.normal_() + layer.mlp.experts.down_proj_bias.normal_() + with tempfile.TemporaryDirectory() as tmp: + model.save_pretrained(tmp) + quantized = GptOssForCausalLM.from_pretrained( + tmp, quantization_config=Mxfp4Config(), dtype=torch.bfloat16, device_map=torch_device + ) - # Check that shapes are reasonable - self.assertEqual(quantized_w.dtype, torch.uint8) + prompt = torch.tensor([[1, 2, 3, 4, 5]], device=torch_device) + for use_cache in (False, True): + with self.subTest(use_cache=use_cache), torch.no_grad(): + out = quantized.generate( + prompt, + max_new_tokens=6, + do_sample=False, + use_cache=use_cache, + output_logits=True, + return_dict_in_generate=True, + ) + for step, logits in enumerate(out.logits): + prefix = out.sequences[:, : prompt.shape[1] + step] + expected = quantized(prefix).logits[:, -1].float() + torch.testing.assert_close(logits.float(), expected, atol=5e-2, rtol=5e-2) -@require_torch @require_torch_large_accelerator -@REQUIRE_TRITON_MXFP4 @require_kernels @slow class Mxfp4ModelTest(unittest.TestCase): - """Test mxfp4 with actual models (requires specific model and hardware)""" + """The published GPT-OSS checkpoint, quantized and dequantized, and saved back.""" - # These should be paths to real OpenAI MoE models for proper testing model_name = "openai/gpt-oss-20b" - input_text = "Once upon a time" - - # Expected outputs for generation tests - EXPECTED_OUTPUTS = set() - EXPECTED_OUTPUTS.add("Once upon a time, in a small town, there lived a young") - - def setUp(self): - gc.collect() - _empty_accelerator_cache() + EXPECTED_OUTPUTS = {"Once upon a time, in a small town, there lived a young"} + # what the published checkpoint keeps in full precision + MODULES_TO_NOT_CONVERT = ["model.layers.*.self_attn", "model.layers.*.mlp.router", "model.embed_tokens", "lm_head"] def tearDown(self): - gc.collect() - _empty_accelerator_cache() + cleanup(torch_device, gc_collect=True) def check_inference_correctness_quantized(self, model, tokenizer): - # Check that inference pass works on the model encoded_input = tokenizer(self.input_text, return_tensors="pt").to(model.device) - - # Set pad token if not set - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - with torch.no_grad(): output_sequences = model.generate( **encoded_input, @@ -419,8 +121,6 @@ def check_inference_correctness_quantized(self, model, tokenizer): self.assertIn(generated_text, self.EXPECTED_OUTPUTS) def test_gpt_oss_model_loading_quantized_with_device_map(self): - """Test loading OpenAI MoE model with mxfp4 quantization and device_map""" - model = GptOssForCausalLM.from_pretrained( self.model_name, dtype=torch.bfloat16, @@ -432,39 +132,18 @@ def test_gpt_oss_model_loading_quantized_with_device_map(self): # ~39 GiB for the dequantized bfloat16 weights. @require_torch_accelerator_memory(memory=48) def test_gpt_oss_model_loading_dequantized_with_device_map(self): - """Test loading OpenAI MoE model with mxfp4 dequantization and device_map""" - - quantization_config = Mxfp4Config(dequantize=True) - - # Test that config is properly set up - self.assertTrue(quantization_config.dequantize) - model = GptOssForCausalLM.from_pretrained( self.model_name, - quantization_config=quantization_config, + quantization_config=Mxfp4Config(dequantize=True), dtype=torch.bfloat16, device_map="auto", ) tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.check_inference_correctness_quantized(model, tokenizer) - def test_model_device_map_validation(self): - """Test device map validation""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - # Test with CPU in device map (CPU already support mxfp4) - quantizer.validate_environment(device_map={"": "cpu"}) - # Holds the packed (~13 GiB) and the dequantized (~39 GiB) model at the same time, by design. @require_torch_accelerator_memory(memory=56) def test_memory_footprint_comparison(self): - """Test memory footprint differences between quantized and unquantized models""" - - # Expected: quantized < dequantized < unquantized memory usage quantization_config = Mxfp4Config(dequantize=True) quantized_model = GptOssForCausalLM.from_pretrained( self.model_name, @@ -484,32 +163,26 @@ def test_memory_footprint_comparison(self): # Keeps the packed model alive while reloading the saved copy, dequantized (~13 + ~39 GiB). @require_torch_accelerator_memory(memory=56) def test_save_mxfp4(self): - """Test saving quantized OpenAI MoE model with device_map""" - model = GptOssForCausalLM.from_pretrained( self.model_name, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) tokenizer = AutoTokenizer.from_pretrained(self.model_name) with tempfile.TemporaryDirectory() as tmp: - # Save the model in mxfp4 format model.save_pretrained(tmp) - _empty_accelerator_cache() - gc.collect() - # test quantized model + cleanup(torch_device, gc_collect=True) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) - # test dequantized model loaded_model = GptOssForCausalLM.from_pretrained( tmp, quantization_config=Mxfp4Config(dequantize=True), - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) @@ -517,83 +190,162 @@ def test_save_mxfp4(self): # Quantizes from a bfloat16 checkpoint, so it holds the bfloat16 source and the mxfp4 result at once. @require_torch_accelerator_memory(memory=56) def test_save_mxfp4_non_quantized(self): - """Test saving dequantized OpenAI MoE model with mxfp4 quantization and device_map""" non_quantized_model_name = "hf-internal-testing/gpt-oss-20b-bf16" tokenizer = AutoTokenizer.from_pretrained(non_quantized_model_name) loaded_model = GptOssForCausalLM.from_pretrained( non_quantized_model_name, - quantization_config=Mxfp4Config(), - torch_dtype=torch.bfloat16, + quantization_config=Mxfp4Config(modules_to_not_convert=self.MODULES_TO_NOT_CONVERT), + dtype=torch.bfloat16, device_map="auto", ) - # save the quantized model with tempfile.TemporaryDirectory() as tmp: loaded_model.save_pretrained(tmp) - _empty_accelerator_cache() - gc.collect() - # load it back to check with everything works as expected + cleanup(torch_device, gc_collect=True) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - quantization_config=Mxfp4Config(dequantized=True), - torch_dtype=torch.bfloat16, + quantization_config=Mxfp4Config(dequantize=True), + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) - def test_compute_module_sizes(self): - r""" - Test if we compute the right module sizes needed to generate the device map. - Also test if we get the right values for `total_byte_count` in `caching_allocator_warmup`. - """ - from transformers import AutoConfig, AutoModelForCausalLM - from transformers.integrations import Mxfp4GptOssExperts - from transformers.integrations.accelerate import compute_module_sizes - from transformers.modeling_utils import expand_device_map, get_total_byte_count - from transformers.quantizers import AutoHfQuantizer - - # we need to preprocess the model like that because device_map calculation happens before we load the weights inside the model. - # For normal wieghts, it's fine but for quantized weights, the tensors dtype might change during loading. - with torch.device("meta"): - config = AutoConfig.from_pretrained(self.model_name) - model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) - model_size, _ = compute_module_sizes(model, only_modules=False) - - expected_keys = [name for name, _ in model.named_parameters()] + [ - name for name, _ in model.named_buffers() - ] - expanded_device_map = expand_device_map({"": torch_device}, expected_keys) - total_byte_count = list(get_total_byte_count(model, expanded_device_map).values())[0] - - # testing prequantized = False should be enough, the shape should be the same whether it is pre-quantized or not - hf_quantizer = AutoHfQuantizer.from_config(Mxfp4Config(), pre_quantized=False) - hf_quantizer.preprocess_model(model=model, config=model.config) - quantized_model_size, _ = compute_module_sizes(model, hf_quantizer, only_modules=False) - - expected_keys = [name for name, _ in model.named_parameters()] + [ - name for name, _ in model.named_buffers() - ] - expanded_device_map = expand_device_map({"": torch_device}, expected_keys) - quantized_total_byte_count = list(get_total_byte_count(model, expanded_device_map, hf_quantizer).values())[ - 0 - ] - for name, module in model.named_modules(): - if isinstance(module, Mxfp4GptOssExperts): - # from 16 bits to 4 bits - assert int(model_size[f"{name}.gate_up_proj"] // 4) == int( - quantized_model_size[f"{name}.gate_up_proj"] - ) - assert int(model_size[f"{name}.down_proj"] // 4) == int(quantized_model_size[f"{name}.down_proj"]) - # check that we get the same value, as we use `compute_module_sizes` in `get_total_byte_count` - assert total_byte_count == model_size[""] - assert quantized_total_byte_count == quantized_model_size[""] +@require_torch +class Mxfp4ChainTest(_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, Mxfp4Config()) + 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_blocks_dequantize_to_the_unquantized_experts_orientation(self): + """`dequantize=True` lands the blocks as bf16 in the `(E, hidden, 2I)` orientation the unquantized + GPT-OSS experts hold.""" + from transformers.core_model_loading import Transpose + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks + + 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 = self._reference_dequant(blocks, scales).transpose(1, 2) + + 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, rtol=0, atol=0) + + def test_blocks_and_scales_take_only_the_ops_their_layout_needs(self): + """GPT-OSS's blocks unpack to the packed weight and its scales only rename; the module's own layout ops + follow.""" + from transformers.core_model_loading import WeightConverter, WeightRenaming + from transformers.integrations.finegrained.conversions import FineGrainedPackedBlocks + + quantizer = _quantizer_for(Mxfp4Config()) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions([]) + blocks = next( + c + for c in conversions + if isinstance(c, WeightConverter) and c.compiled_sources.search("gate_up_proj_blocks") + ) + self.assertEqual([type(op) for op in blocks.operations], [FineGrainedPackedBlocks]) + self.assertTrue( + any( + isinstance(c, WeightRenaming) and c.compiled_sources.search("experts.gate_up_proj_scales") + for c in conversions + ) + ) + + +class Mxfp4Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = MXFP4 + + +class Mxfp4TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = MXFP4 + - # we should at least have 3 times memory reduction in total for this model - assert model_size[""] > quantized_model_size[""] * 3 +class Mxfp4LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = MXFP4 diff --git a/tests/quantization/mxfp8/__init__.py b/tests/quantization/mxfp8/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/mxfp8/test_mxfp8.py b/tests/quantization/mxfp8/test_mxfp8.py new file mode 100644 index 000000000000..c6fa0307a7e9 --- /dev/null +++ b/tests/quantization/mxfp8/test_mxfp8.py @@ -0,0 +1,70 @@ +# 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. +"""MXFP8 (`Mxfp8Config`): the config and the checkpoint layouts it loads and saves back.""" + +import unittest + +import torch + +from transformers.testing_utils import ( + TestCasePlus, + require_torch, +) +from transformers.utils.quantization_config import Mxfp8Config + +from ...integrations.finegrained.utils import ( + _ChainAssertions, + _load_and_save, + _per_expert, +) +from ..block_scaled_tester import ( + MXFP8, + BlockScaledFormatTesterMixin, + LoadPathEquivalenceTesterMixin, + TinyCheckpointTesterMixin, +) + + +@require_torch +class Mxfp8ChainTest(_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, Mxfp8Config()) + 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) + + +class Mxfp8Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = MXFP8 + + +class Mxfp8TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = MXFP8 + + +class Mxfp8LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = MXFP8 diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index 5847e352d7c2..181ee1987ff9 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -11,157 +11,467 @@ # 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 unittest -from types import SimpleNamespace -from unittest.mock import patch - -from transformers import NVFP4Config -from transformers.integrations.nvfp4 import NVFP4Linear, NVFP4Quantize, replace_with_nvfp4_linear -from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer -from transformers.utils import is_torch_available -from transformers.utils.quantization_config import QuantizationMethod - - -if is_torch_available(): - import torch - from torch import nn - - -class _PackedWeight: - def __init__(self, **kwargs): - self.__dict__.update(kwargs) +"""NVFP4 (`NVFP4Config`): the two-level global-scale ops, the modelopt and compressed-tensors checkpoint layouts +it loads and saves back, and how those exports read.""" - -def _fake_kernel(): - def pack(weight, device): - out_features, in_features = weight.shape - return SimpleNamespace( - qweight=torch.empty(out_features, in_features // 2, dtype=torch.uint8, device=device), - sf=torch.empty(out_features, in_features // 64, dtype=torch.int32, device=device), - sf_rowmajor=torch.empty(out_features, in_features // 16, dtype=torch.uint8, device=device), - global_scale=torch.ones(1, dtype=torch.float32, device=device), +import os +import unittest +from unittest import mock + +import torch + +from transformers.integrations.finegrained import FineGrainedExperts +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + require_torch, +) +from transformers.utils.quantization_config import NVFP4Config + +from ...integrations.finegrained.utils import ( + E, + _Cfg, + _ChainAssertions, + _compressed_tensors, + _load_and_save, + _modelopt, + _per_expert, + _quantizer_for, + _targets, +) +from ..block_scaled_tester import ( + NVFP4, + BlockScaledFormatTesterMixin, + LoadPathEquivalenceTesterMixin, + TinyCheckpointTesterMixin, +) + + +@require_torch +class Nvfp4GlobalScaleOpsTest(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.""" + + 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.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_activation_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_activation_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, ) - - return SimpleNamespace( - PackedWeight=_PackedWeight, - gemm=lambda packed_weight, input: input.new_zeros((*input.shape[:-1], packed_weight.n)), - pack=pack, - swizzled_sf_shape=lambda rows, columns: (rows, columns // 64), - ) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4ConfigTest(unittest.TestCase): - def test_config_uses_quantization_method_enum(self): - config = NVFP4Config(modules_to_not_convert=["lm_head"]) - - self.assertEqual(config.quant_method, QuantizationMethod.NVFP4) - self.assertEqual(config.to_dict()["modules_to_not_convert"], ["lm_head"]) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4IntegrationTest(unittest.TestCase): - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_replaces_only_eligible_linears(self, _): - model = nn.ModuleDict( + 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_activation_scale", + ] + out = FineGrainedWeightGlobals().convert( { - "convert": nn.Linear(64, 32, bias=False), - "skip": nn.Linear(64, 32, bias=False), - "biased": nn.Linear(64, 32, bias=True), - "unaligned": nn.Linear(63, 32, bias=False), - } + "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, ) - - replace_with_nvfp4_linear(model, modules_to_not_convert=["skip"]) - - self.assertIsInstance(model["convert"], NVFP4Linear) - self.assertIsInstance(model["skip"], nn.Linear) - self.assertIsInstance(model["biased"], nn.Linear) - self.assertIsInstance(model["unaligned"], nn.Linear) - self.assertEqual(model["convert"].weight_sf.dtype, torch.int32) - self.assertEqual(model["convert"].weight_sf_rowmajor.dtype, torch.uint8) - self.assertEqual(model["convert"].weight_global_scale.dtype, torch.float32) - - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_linear_forward_uses_normal_function_entry_point(self, _): - layer = NVFP4Linear(64, 32) - output = layer(torch.randn(2, 3, 64)) - - self.assertEqual(output.shape, (2, 3, 32)) - - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_packed_tensors_are_registered_in_state_dict(self, _): - layer = NVFP4Linear(64, 32) - - for buffer_name in ("weight", "weight_sf", "weight_sf_rowmajor", "weight_global_scale"): - self.assertIsInstance(getattr(layer, buffer_name), nn.Buffer) - self.assertEqual( - set(layer.state_dict()), - {"weight", "weight_sf", "weight_sf_rowmajor", "weight_global_scale"}, + 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, ) - self.assertNotIn("nvfp4_kernel", layer.state_dict()) + 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])) - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_conversion_produces_all_registered_buffers(self, _): - conversion = NVFP4Quantize(torch.device("cpu")) + 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 - output = conversion.convert({"model.proj.weight": torch.randn(32, 64)}) - - self.assertEqual( - set(output), + up = FineGrainedInputScales().convert( { - "model.proj.weight", - "model.proj.weight_sf", - "model.proj.weight_sf_rowmajor", - "model.proj.weight_global_scale", + "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_activation_scale", + )["mlp.experts.gate_up_proj_activation_scale"] + 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_activation_scale", + )["mlp.experts.down_proj_activation_scale"] + torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) + + def test_shared_gate_up_input_global_follows_the_env_var(self): + """`TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL=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.layers import register_shared_gate_up_input_globals + + variable = "TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL" + for value, expected in ((None, None), ("global", 4.0)): + model = torch.nn.Module() + model.experts = experts = FineGrainedExperts(_Cfg(), quantization_config=NVFP4Config()) + with torch.no_grad(): + experts.gate_up_proj_activation_scale.copy_(torch.tensor([1.0, 4.0, 2.0, 3.0])) + with mock.patch.dict(os.environ): + os.environ.pop(variable, None) + if value is not None: + os.environ[variable] = value + 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_activation_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_activation_scale) + self.assertEqual(experts.gate_up_proj_activation_scale.numel(), 4) + with mock.patch.dict(os.environ, {variable: "node"}), self.assertRaises(ValueError): + register_shared_gate_up_input_globals(model) + + +@require_torch +class Nvfp4ChainTest(_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() + return checkpoint + + def test_modelopt_per_expert_and_dense(self): + checkpoint = self._modelopt_checkpoint() + model, info, saved = _load_and_save( + checkpoint, AutoQuantizationConfig.from_dict({"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"), + ("activation_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_compressed_tensors_per_expert_and_dense(self): + """compressed-tensors names the packed values `weight_packed` and the activation global `input_global_scale`, + and stores both globals as the reciprocal of modelopt's: it loads into the same slots, and saves back as it + shipped. Globals that are powers of two keep the reciprocal exact.""" + modelopt = self._modelopt_checkpoint() + checkpoint = {} + for key, tensor in modelopt.items(): + stem, _, suffix = key.rpartition(".") + if suffix in ("weight_scale_2", "input_scale"): + modelopt[key] = tensor = torch.exp2(torch.round(torch.log2(tensor))) + name = "weight_global_scale" if suffix == "weight_scale_2" else "input_global_scale" + checkpoint[f"{stem}.{name}"] = 1 / tensor + else: + checkpoint[f"{stem}.weight_packed" if suffix == "weight" and tensor.dtype == torch.uint8 else key] = ( + tensor + ) + config = NVFP4Config(activation_format="nvfp4", checkpoint_format="compressed-tensors") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + experts, q_proj = model.model.layers[0].mlp.experts, model.model.layers[0].self_attn.q_proj + torch.testing.assert_close( + experts.gate_up_proj_weight_global_scale, + torch.stack([modelopt[f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale_2"] for e in range(E)]), + ) + torch.testing.assert_close( + q_proj.weight_global_scale, modelopt["model.layers.0.self_attn.q_proj.weight_scale_2"] + ) + torch.testing.assert_close(q_proj.activation_scale, modelopt["model.layers.0.self_attn.q_proj.input_scale"]) + 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, AutoQuantizationConfig.from_dict({"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): + # a pure modelopt export reads as this config + config = NVFP4Config(**{"activation_format": "nvfp4", **cfg_kwargs}) + return _quantizer_for(config).format_conversions(dequantize=False) + + 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, + quantization_config=NVFP4Config(activation_format=activation_format), + ) + linear = FineGrainedLinear(32, 32, quantization_config=NVFP4Config(activation_format=activation_format)) + for module in (experts, linear): + slots = {name for name, _ in module.named_parameters()} + self.assertEqual( + any("activation_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): + # 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 + from transformers.core_model_loading import MergeModulelist, WeightConverter, WeightRenaming, rename_source_key + + conversions = self._modelopt_conversions() + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + _, pattern = rename_source_key("model.layers.0.mlp.experts.7.gate_proj.weight_scale_2", renamings, converters) + globals_converters = [c for c in converters if pattern 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", + ) + + +@require_torch +class Nvfp4CheckpointReaderTest(unittest.TestCase): + """modelopt and compressed-tensors NVFP4 exports read as `NVFP4Config`, skip-lists bounded to the modules they name.""" + + 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, so the export is one format.""" + config = _modelopt( + 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"], + } + }, + ) + self.assertIsInstance(config, NVFP4Config) + self.assertIsNone(config.activation_format) + # NVFP4's `dynamic: False` is its calibrated activation global, not a static activation scale + self.assertEqual(config.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.""" + config = _modelopt( + config_groups={ + "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} + }, + ) + self.assertIsInstance(config, NVFP4Config) + + def test_a_compressed_tensors_ignore_entry_is_that_module_not_a_prefix(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + skip = AutoQuantizationConfig.from_dict( + _compressed_tensors(nvfp4, ignore=["model.layers.0.mlp.gate", "re:.*router$"]) + ).modules_to_not_convert + self.assertFalse(should_convert_module("model.layers.0.mlp.gate", skip)) + self.assertFalse(should_convert_module("model.layers.3.mlp.router", skip)) + self.assertTrue(should_convert_module("model.layers.0.mlp.gate_proj", skip)) + + 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 + + # as a modelopt checkpoint ships them, through the config that normalizes them + skip = AutoQuantizationConfig.from_dict( + { + "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) + + def test_a_compressed_tensors_nvfp4_checkpoint_reads_as_nvfp4(self): + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + config = AutoQuantizationConfig.from_dict(_compressed_tensors(nvfp4, nvfp4, ignore=["lm_head"])) + self.assertIsInstance(config, NVFP4Config) + self.assertEqual(config.checkpoint_format, "compressed-tensors") + + def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): + """Integer, per-channel FP8 and sparse checkpoints are compressed-tensors' own.""" + from transformers.utils.quantization_config import read_compressed_tensors + + int4 = {"num_bits": 4, "type": "int", "group_size": 128, "strategy": "group"} + channel = {"num_bits": 8, "type": "float", "strategy": "channel"} + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + for checkpoint in ( + _compressed_tensors(int4, fmt="pack-quantized"), + _compressed_tensors(channel, fmt="float-quantized"), + _compressed_tensors(nvfp4, sparsity_config={"format": "sparse-24-bitmask"}), + ): + self.assertIsNone(read_compressed_tensors(checkpoint)) + + def test_activations_the_weights_kernels_do_not_serve_stay_with_compressed_tensors(self): + """An activation scheme is read only where the weights' kernels serve it: NVFP4 weights with FP8 per-token + activations, weight-only block-FP8 and FP8 weights with int8 activations stay compressed-tensors' own, as do + group-32 FP8 weights whose scales are not E8M0 (MXFP8's).""" + from transformers.utils.quantization_config import Mxfp8Config, read_compressed_tensors + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} + fp8_token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} + int8_token = {"num_bits": 8, "type": "int", "strategy": "token", "dynamic": True} + group_32 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32} + for name, checkpoint in ( + ("nvfp4 x fp8", _compressed_tensors(nvfp4, fp8_token)), + ("block-fp8 weight-only", _compressed_tensors(block, fmt="float-quantized")), + ("block-fp8 x int8", _compressed_tensors(block, int8_token, fmt="float-quantized")), + ( + "group-32 fp32 scales", + _compressed_tensors({**group_32, "scale_dtype": "torch.float32"}, group_32, fmt="float-quantized"), + ), + ): + with self.subTest(name): + self.assertIsNone(read_compressed_tensors(checkpoint)) + e8m0 = {**group_32, "scale_dtype": "torch.float8_e8m0fnu"} + self.assertIsInstance( + read_compressed_tensors(_compressed_tensors(e8m0, group_32, fmt="float-quantized")), Mxfp8Config ) - self.assertEqual(output["model.proj.weight"].shape, (32, 32)) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4QuantizerValidationTest(unittest.TestCase): - def setUp(self): - self.quantizer = NVFP4HfQuantizer(NVFP4Config(), pre_quantized=False) - - def test_rejects_pre_quantized_checkpoints(self): - quantizer = NVFP4HfQuantizer(NVFP4Config(), pre_quantized=True) - - with self.assertRaisesRegex(ValueError, "pre-quantized NVFP4"): - quantizer.validate_environment(device_map={"": "cuda"}) - - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - @patch("torch.cuda.get_device_capability", return_value=(10, 0)) - def test_accepts_one_blackwell_cuda_device(self, _, __, ___): - self.quantizer.validate_environment(device_map={"": "cuda:1"}) - - self.assertEqual(self.quantizer.quantization_device, torch.device("cuda:1")) - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - def test_rejects_multi_device_map_until_tp_metadata_is_supported(self, _, __): - with self.assertRaisesRegex(ValueError, "Tensor parallelism"): - self.quantizer.validate_environment(device_map={"model.layers.0": 0, "model.layers.1": 1}) - def test_rejects_native_tp_until_scale_metadata_is_supported(self): - config = SimpleNamespace(distributed_config=SimpleNamespace(tp_size=2)) +class Nvfp4Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = NVFP4 - with self.assertRaisesRegex(ValueError, "scale metadata does not have a sharding plan"): - self.quantizer.update_tp_plan(config) - def test_update_tp_plan_keeps_non_tp_configs(self): - configs = ( - SimpleNamespace(), - SimpleNamespace(distributed_config=SimpleNamespace(tp_size=1)), - ) +class Nvfp4TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = NVFP4 - for config in configs: - with self.subTest(config=config): - self.assertIs(self.quantizer.update_tp_plan(config), config) - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - def test_rejects_cpu_offload(self, _, __): - with self.assertRaisesRegex(ValueError, "CPU or disk offload"): - self.quantizer.validate_environment(device_map={"model": "cuda", "lm_head": "cpu"}) +class Nvfp4LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = NVFP4 From d89e0c4bf56660c5d40fad1c5a5f88b25f6efef5 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 7 Oct 2026 09:08:54 -0700 Subject: [PATCH 69/91] Block-scaled quantization: name per-tensor scales `replicate` in the parallel plans A per-tensor scale is 0-dim, with no axis to split. Rather than every dense style special-casing 0-dim parameters, the quantizer names each one a TP/EP plan would shard along with its module `replicate` (a new no-op style), which the parameter's own entry makes take precedence over the module's. Also borrows huggingface/transformers#49383: no FSDP wrap for expert-parallel token dispatch when fsdp_size == 1 (FSDP refuses quantized modules: mixed dtypes, 0-dim scales). Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/distributed/mixin.py | 10 +++---- .../distributed/sharding_utils.py | 14 ++++------ .../distributed/tensor_parallel.py | 26 +++++++------------ .../quantizers/block_scaled_utils.py | 17 ++++++++++++ .../finegrained/test_finegrained.py | 21 +++++++++++++++ .../utils/test_distributed_sharding_utils.py | 12 --------- 6 files changed, 57 insertions(+), 43 deletions(-) diff --git a/src/transformers/distributed/mixin.py b/src/transformers/distributed/mixin.py index 29bef1debfe1..cb02ae43ab6a 100644 --- a/src/transformers/distributed/mixin.py +++ b/src/transformers/distributed/mixin.py @@ -223,7 +223,7 @@ def maybe_distribute_model( # we will have to slice the batch here in order to avoid computing tp_size times the same batch. model = apply_expert_parallelism(model, ep_mesh, tp_mesh, ep_plan) - if distributed_config.fsdp_size > 1 or "ep_dispatch_experts" in ep_plan.values(): + if distributed_config.fsdp_size > 1: model = apply_fully_sharded_data_parallelism(model, mesh_manager) return model @@ -286,11 +286,9 @@ def gather_sharded_state_dict_for_save( if distributed_config is None: return state_dict - if distributed_config.fsdp_size > 1 or ( - distributed_config.ep_size > 1 and "ep_dispatch_experts" in self.ep_plan.values() - ): - # Also covers the 2-D (fsdp, tp) mesh and token dispatch: every parameter is FSDP-managed, and the - # full state dict is only materialized on rank 0. + if distributed_config.fsdp_size > 1: + # Also covers the 2-D (fsdp, tp) mesh and token dispatch with expert replicas: every parameter is + # FSDP-managed, and the full state dict is only materialized on rank 0. if not _is_torch_distributed_initialized(): raise ValueError( "Saving an FSDP-wrapped model requires torch.distributed to be initialized. " diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 826388a178b9..083fbe5f915f 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -87,9 +87,9 @@ 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 (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 + # [_axis0_offset, _axis0_offset + _axis0_local_size) + self._axis0_offset = offsets[0] + self._axis0_local_size = local_shape[0] def shard_tensor( self, source: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None @@ -112,8 +112,7 @@ def shard_tensor( # Dense path if tensor_idx is None: - # on no dim, return early parameter is replicated - if not source_shape or not dim_placements: + if 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. @@ -139,6 +138,7 @@ 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) @@ -162,10 +162,6 @@ def shard_tensor( 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) - # `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/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 330c49a0bcaa..41fbca973ae1 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -203,13 +203,7 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - # 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) + placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -280,10 +274,7 @@ def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: return - # 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): + if isinstance(module, torch.nn.Embedding): placement = Shard(0) else: # bias is replicated (added after the row-reduce); weight shards on input dim (-1) @@ -354,6 +345,11 @@ def transform_output_post_forward(self, module, output, mesh): return output +class ReplicateParallel(TensorParallelLayer): + """A parameter every rank holds whole, such as a per-tensor quantization scale: named in a plan, it takes + precedence over the style its module's entry would give it.""" + + class ReplicatedWithGradAllReduce(TensorParallelLayer): """Replicated parameter whose gradient is partial. @@ -489,10 +485,7 @@ 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. - # 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: + if meta.ndim == 1: placement = Shard(shard_dim) else: placement = _StridedShard(dim=shard_dim, split_factor=self.split_factor) @@ -533,7 +526,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, @@ -959,6 +952,7 @@ class ParallelInterface(GeneralInterface): "moe_tp_experts": MoeExpertsParallel(), "megamoe_experts": MoeTensorParalellMegaMoeExperts(), "moe_identity_expert": MoeIdentityParallel(), + "replicate": ReplicateParallel(), "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), "mla_kv_a_proj": MlaKvAProjParallel(), "all_reduce": AllReduceParallel(), diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index 52a39367aaa1..4f46aee225f5 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -300,12 +300,29 @@ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwarg if not any(name.endswith("activation_scale") for name, _ in model.named_parameters()): ignored.add(r"input(_global)?_scale$") model._keys_to_ignore_on_load_unexpected = ignored + self._replicate_per_tensor_scales(model) if not replaced and not self.quantization_config.dequantize: logger.warning( "You are loading a fine-grained quantized model but no linear or experts modules were " "found to convert. Please double check your model architecture." ) + @staticmethod + def _replicate_per_tensor_scales(model: "PreTrainedModel"): + """A per-tensor scale is 0-dim, with no axis to split: each one a parallel plan would shard along with its + module is named `replicate`, which takes precedence over the module's entry.""" + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard + + for plan in (model._tp_plan, model._ep_plan): + if plan: + plan.update( + { + replace_layer_number_by_wildcard(name): "replicate" + for name, param in model.named_parameters() + if param.ndim == 0 and _get_parameter_plan(name, plan) is not None + } + ) + def replace_modules(self, model, plan=None, modules_to_not_convert=None) -> bool: """Swap the modules `plan` (this quantizer's own by default) covers for their quantized classes.""" from ..integrations.finegrained import ( diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py index 06faae1cae36..3b63cc4c5b64 100644 --- a/tests/quantization/finegrained/test_finegrained.py +++ b/tests/quantization/finegrained/test_finegrained.py @@ -356,6 +356,27 @@ def tp_plan(dtype_plan): with self.assertRaises(ValueError): tp_plan({r"layers\.1\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + def test_parallel_plans_replicate_the_per_tensor_scales(self): + """A per-tensor scale is 0-dim, with no axis to split: the quantizer names each one a plan would shard with + its module `replicate`, which the parameter's own entry makes take precedence. A scale grid follows its + module's style.""" + from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES + from transformers.integrations.finegrained import FineGrainedLinear + from transformers.quantizers.block_scaled_utils import BlockScaledQuantizer + + model = nn.Module() + model.q_proj = FineGrainedLinear(64, 64, quantization_config=NVFP4Config()) + model.o_proj = FineGrainedLinear(64, 64, quantization_config=NVFP4Config()) + model._tp_plan, model._ep_plan = {"q_proj": "colwise"}, {} + BlockScaledQuantizer._replicate_per_tensor_scales(model) + self.assertEqual( + model._tp_plan, + {"q_proj": "colwise", "q_proj.weight_global_scale": "replicate", "q_proj.activation_scale": "replicate"}, + ) + scale = model.q_proj.weight_global_scale + ALL_PARALLEL_STYLES["replicate"].shard_param(model.q_proj, "weight_global_scale", mesh=None) + self.assertIs(model.q_proj.weight_global_scale, scale) + @require_torch class FineGrainedPlanConversionTest(unittest.TestCase): diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index b6a6d4d26093..35e897e3e219 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -294,18 +294,6 @@ 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) From ae746d55d99547bccc1de7e85dc7fe5a565800fa Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Wed, 7 Oct 2026 09:32:51 -0700 Subject: [PATCH 70/91] Borrow the expert-parallel dispatch wait for raw-pointer kernels From the `ep-dispatch-wait-collective` branch: modules flagged `_hf_quantized_needs_local_tp` get the dispatch all-to-all's received rows waited on, so the fine-grained kernels never read a pending collective's null pointer. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/distributed/tensor_parallel.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 41fbca973ae1..5eb725199aa3 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -38,7 +38,7 @@ if is_torch_distributed_available(): import torch.distributed as dist - from torch.distributed._functional_collectives import all_to_all_single + from torch.distributed._functional_collectives import all_to_all_single, wait_tensor from torch.distributed.tensor import DTensor, Partial, Replicate, Shard, distribute_tensor from torch.distributed.tensor.placement_types import _StridedShard @@ -896,6 +896,10 @@ def ep_forward(hidden_states, top_k_index, top_k_weights): tokens, expert_ids, order, send_sizes, recv_sizes = self._dispatch_tokens( hidden_states, top_k_index, module.num_experts, ep_group, ep_size ) + # kernels reading raw pointers (the local tensors above) get the received rows themselves: the + # all-to-all's pending result holds a null pointer until a torch op waits on it + if getattr(module, "_hf_quantized_needs_local_tp", False): + tokens = wait_tensor(tokens) expert_output = self._run_local_experts(experts_forward, tokens, expert_ids) output = self._combine_tokens( expert_output, top_k_weights, order, send_sizes, recv_sizes, ep_group From e8886aba3a794358d28449a5de5b51e34acca87c Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 15 Oct 2026 23:13:36 -0700 Subject: [PATCH 71/91] Add Muse Spark NVFP4 checkpoint test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/quantization/block_scaled_tester.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index a537fa6ed069..37e47e86af59 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -178,6 +178,7 @@ class FormatSpec: "IlyasMoutawwakil/tiny-random-Llama-NVFP4": 0.25, "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": 0.25, "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": 0.6, + "IlyasMoutawwakil/tiny-random-Muse-Spark-NVFP4": 0.3, }, distributed_models=("glm4v_moe-nvfp4",), ) @@ -470,6 +471,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)): [ From 30b036fdde3a764bf960dea3e5e9706662cf93cb Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 15 Oct 2026 23:19:02 -0700 Subject: [PATCH 72/91] Apply repository consistency fixes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../integrations/finegrained/layers.py | 9 +++---- src/transformers/utils/quantization_config.py | 26 +++++-------------- tests/integrations/test_moe.py | 1 - 3 files changed, 9 insertions(+), 27 deletions(-) diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index 9874b044efcb..81f31fa1375f 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -652,12 +652,9 @@ def __init__( ) self.global_scale_dtype = format_spec.global_scale_dtype - # The layouts: what the CHECKPOINT ships (`is_concatenated`: stacked ``[gate; up]`` or, like - # GPT-OSS, already the kernels' interleaved ``[g0, u0, ...]``) against what the module - # 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. The - # triton backends and the eager loop share the interleaved rows; the DeepGEMM ones read - # gate|up stacked (Mega MoE packs it itself, the grouped GEMM reads a block grid row-major). + # Checkpoint and module layouts may differ: conversion ops normalize stacked or interleaved + # gate/up rows and swizzle scales at load. The triton and eager paths use interleaved rows; + # DeepGEMM paths use stacked rows. self.has_gate = has_gate self.has_bias = has_bias self.is_concatenated = is_concatenated diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index ef7ed2df2865..6c329ee92126 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1734,17 +1734,10 @@ class BlockScaledConfig(QuantizationConfigMixin): [`Mxfp4Config`] or [`NVFP4Config`]. Args: - activation_format (`str`, *optional*): - What activations are quantized to: `None` for the weights' own format (which its name spells too), - `"bf16"` for weight-only, or another format the weights' kernels serve. On a pre-quantized checkpoint, a - passed config's activation format replaces the checkpoint's: it picks the kernels the same weights run on. - dequantize (`bool`, *optional*, defaults to `False`): - Whether to dequantize the model to full precision while loading. - modules_to_not_convert (`list`, *optional*): - Modules to keep in full precision. `ignored_layers` is accepted as an alias. - checkpoint_format (`str`, *optional*): - The producer whose key layout the checkpoint ships when it is not this format's own: - `"compressed-tensors"`. + activation_format (`str | None`, *optional*, defaults to `the format's default`): + dequantize (`bool`, *optional*, defaults to `False`): + modules_to_not_convert (`list | None`, *optional*): + checkpoint_format (`str | None`, *optional*): """ # what `quant_method` this config is, the activation formats its weights' kernels serve, and the one it takes @@ -1882,10 +1875,7 @@ class Mxfp8Config(BlockScaledConfig): """MXFP8 weights: `float8_e4m3fn` values with a UE8M0 scale per 32, as MiniMax-M3 ships them. Args: - activation_format (`str`, *optional*): - `None` (or `"mxfp8"`) for MXFP8 activations, `"mxfp4"` for W8A4, or `"bf16"` for weight-only. - kwargs (`dict[str, Any]`, *optional*): - The options every block-scaled format takes, see [`BlockScaledConfig`]. + activation_format (`str | None`, *optional*, defaults to `the format's default`): """ method = QuantizationMethod.MXFP8 @@ -1921,11 +1911,7 @@ class NVFP4Config(BlockScaledConfig): exports them; quantizes a full-precision checkpoint on the fly as well. Args: - activation_format (`str`, *optional*): - `None` (or `"nvfp4"`) for NVFP4 activations (W4A4, with the checkpoint's calibrated activation global) or - `"bf16"` for weight-only. - kwargs (`dict[str, Any]`, *optional*): - The options every block-scaled format takes, see [`BlockScaledConfig`]. + activation_format (`str | None`, *optional*, defaults to `the format's default`): """ method = QuantizationMethod.NVFP4 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, ) From 7a7e5fd2ae5777a395fdbdc067f8eb0d582788ba Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 03:35:24 -0700 Subject: [PATCH 73/91] Limit kernel correctness tests to one autotuning trial --- docs/source/en/quantization/finegrained.md | 7 ++++++ tests/kernels/conftest.py | 27 ++++++++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 tests/kernels/conftest.py diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 43548e0e943d..5bcf5a974b5b 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -66,3 +66,10 @@ These load as a [`FineGrainedConfig`] without one being passed: To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass `FineGrainedConfig(dequantize=True)`. + +## Kernel autotuning + +Fine-grained kernels tune each new input shape and cache the chosen configuration. For correctness tests or quick +experiments, set `FINEGRAINED_AUTOTUNE_TRIALS=1` before loading the kernel bundle to benchmark only one configuration +per shape. Leave it unset when measuring performance so the tuner can search its full budget. See the +[kernel autotuning documentation](https://huggingface.co/kernels-community/finegrained-kernels#autotuning). diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py new file mode 100644 index 000000000000..24be14502c36 --- /dev/null +++ b/tests/kernels/conftest.py @@ -0,0 +1,27 @@ +# 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. + +import os +from unittest import mock + +import pytest + + +@pytest.fixture(scope="session", autouse=True) +def cheap_finegrained_tuning(): + # Correctness tests need one valid configuration per shape. Set this before the + # kernel bundle loads: its autotuners read the trial budget at import time. + # Subprocesses inherit it too, including distributed workers. + with mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}): + yield From 9dba82da477efcbdcfa0be4f982251e05dcefc49 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 03:45:59 -0700 Subject: [PATCH 74/91] Clarify autotuning trial budget and remaining refinement --- docs/source/en/quantization/finegrained.md | 5 +++-- tests/kernels/conftest.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 5bcf5a974b5b..121d908b3590 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -70,6 +70,7 @@ To change one format of such a checkpoint, pass the plan you want; to dequantize ## Kernel autotuning Fine-grained kernels tune each new input shape and cache the chosen configuration. For correctness tests or quick -experiments, set `FINEGRAINED_AUTOTUNE_TRIALS=1` before loading the kernel bundle to benchmark only one configuration -per shape. Leave it unset when measuring performance so the tuner can search its full budget. See the +experiments, set `FINEGRAINED_AUTOTUNE_TRIALS=1` before loading the kernel bundle to reduce the Bayesian search +budget. The tuner still probes initial configurations and refines its choice, so this does not cap the total number of +benchmarks at one. Leave it unset when measuring performance so the tuner can search its full budget. See the [kernel autotuning documentation](https://huggingface.co/kernels-community/finegrained-kernels#autotuning). diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py index 24be14502c36..31861ce9245c 100644 --- a/tests/kernels/conftest.py +++ b/tests/kernels/conftest.py @@ -20,7 +20,7 @@ @pytest.fixture(scope="session", autouse=True) def cheap_finegrained_tuning(): - # Correctness tests need one valid configuration per shape. Set this before the + # Use the minimum Bayesian search budget for correctness tests. Set this before the # kernel bundle loads: its autotuners read the trial budget at import time. # Subprocesses inherit it too, including distributed workers. with mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}): From 363e4198228d10d7840089406eef93318f3b8da9 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 04:21:32 -0700 Subject: [PATCH 75/91] Restore fused dense scale containers and align MiniMax dtype expectations --- .../integrations/finegrained/conversions.py | 12 +++++++++ .../finegrained/test_conversions.py | 26 +++++++++++++++++++ tests/quantization/block_scaled_tester.py | 17 ++++++++---- 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 028c6d4746bd..59393e23bb57 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -132,6 +132,18 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= if not self.hf_quantizer.pre_quantized: return input_dict module = finegrained_module(model, full_layer_name) + if module is None and self.inverse: + # Saving names the first checkpoint output (e.g. gate_proj), while the + # model holds the fused input (gate_up_proj). Recover that module before + # reversing its layouts, which runs before the projection is split. + for target in target_patterns or []: + for source in kwargs.get("source_patterns", []): + candidate = (full_layer_name or "").replace(target, source.replace(r"\.", "."), 1) + module = finegrained_module(model, candidate) + if module is not None: + break + if module is not None: + break if module is None: return input_dict module_name = next(name for name, candidate in model.named_modules() if candidate is module) diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index 57b8057601d0..a59b10008762 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -26,6 +26,7 @@ from transformers import PreTrainedConfig, PreTrainedModel from transformers.core_model_loading import ( Chunk, + Concatenate, PermuteForRope, WeightConverter, convert_and_load_state_dict_in_model, @@ -187,6 +188,31 @@ def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): self.assertIs(out["mlp.experts.gate_up_proj"], grid) self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) + def test_fused_dense_scales_save_back_as_exponent_bytes(self): + from transformers.integrations.finegrained.conversions import FineGrainedDeserialize + + config = Mxfp8Config() + model = torch.nn.Module() + model.config = PreTrainedConfig() + model.mlp = torch.nn.Module() + model.mlp.gate_up_proj = fg.FineGrainedLinear(64, 64, quantization_config=config) + model.mlp.gate_up_proj.scale_container_dtype = torch.uint8 + scales = torch.arange(120, 128, dtype=torch.uint8).repeat(16, 1).reshape(64, 2) + model._weight_conversions = [ + WeightConverter( + [r"\.gate_proj.weight_scale_inv", r"\.up_proj.weight_scale_inv"], + r"\.gate_up_proj.weight_scale_inv", + [Concatenate(dim=0), FineGrainedDeserialize(_quantizer_for(config))], + ) + ] + saved = revert_weight_conversion( + model, {"mlp.gate_up_proj.weight_scale_inv": scales.view(torch.float8_e8m0fnu)} + ) + for name, expected in zip(("gate_proj", "up_proj"), scales.chunk(2)): + actual = saved[f"mlp.{name}.weight_scale_inv"] + self.assertEqual(actual.dtype, torch.uint8) + self.assertTrue(torch.equal(actual, expected)) + def test_scale_container_op_keeps_exponent_bytes_and_normalizes_float32(self): """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: both load into held e8m0. Exponent bytes save back as uint8, while float32 containers diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index 848c39eead6a..d9dde7eb6654 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -483,9 +483,9 @@ def _checkpoint_bytes(path): "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], + [0.2637, -0.0815, -0.3418, 0.1553, -0.2715, 0.1562, 0.4590, 0.0635], + [0.4355, -0.1016, -0.2002, -0.2412, -0.2930, 0.1187, 0.5664, 0.1104], + [0.1865, -0.0996, -0.6875, 0.2988, 0.0635, 0.1035, 0.5586, -0.0703], ], } ), # fmt: skip @@ -597,14 +597,21 @@ def test_save_writes_the_checkpoint_layout_back(self): source_dtype, source_shape, source_raw = source[key] self.assertEqual(shape, source_shape) if dtype != source_dtype: - # Loading honors the module's initialized FP32 norm dtype. A UE8M0 - # scale shipped in a float32 container saves as native UE8M0. + # Loading honors initialized dtypes, including FP32 norms and BF16 + # MiniMax routers/patch embeddings. Float32 UE8M0 containers save natively. self.assertTrue( ( dtype == torch.float32 and key.endswith("norm.weight") and source_dtype == torch.bfloat16 ) + or ( + dtype == torch.bfloat16 + and source_dtype == torch.float32 + and key.endswith( + ("gate.weight", "e_score_correction_bias", "patch_embedding.weight") + ) + ) or ( dtype == torch.float8_e8m0fnu and source_dtype == torch.float32 From 5b2cb429929be3cb2c6c98ccbe1351a71090b537 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 04:56:29 -0700 Subject: [PATCH 76/91] Use temporary masked EP and test compiled static-cache decoding --- docs/source/en/quantization/finegrained.md | 6 + .../quantizers/block_scaled_utils.py | 13 ++ tests/quantization/block_scaled_tester.py | 125 ++++++++++-------- .../finegrained/test_finegrained.py | 29 ++++ tests/quantization/mxfp4/test_mxfp4.py | 4 +- 5 files changed, 123 insertions(+), 54 deletions(-) diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 121d908b3590..04bd9e79e5e1 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -67,6 +67,12 @@ These load as a [`FineGrainedConfig`] without one being passed: To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass `FineGrainedConfig(dequantize=True)`. +## Expert parallelism + +Until the EP implementation PR is merged, the default fine-grained expert-parallel path temporarily keeps tokens +replicated across ranks, masks routes to non-local experts, and sums their outputs with all-reduce. Set `tp_size` and +`ep_size` to the same value in `DistributedConfig`. DeepGEMM MegaMoE uses its own routing and communication. + ## Kernel autotuning Fine-grained kernels tune each new input shape and cache the chosen configuration. For correctness tests or quick diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index abb76edbc7d0..b277d38de36b 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -353,6 +353,19 @@ def update_parallel_plan(self, model, plan): getattr(model.config, "_experts_implementation", None), {} ) plan.update({key: overrides.get(style, style) for key, style in plan.items()}) + # Temporary until the EP implementation PR is merged: use replicated routing, + # local expert masking, and all-reduce to keep token shapes fixed across ranks. + if "ep_dispatch_experts" in plan.values(): + modules = {replace_layer_number_by_wildcard(name) for name, _ in model.named_modules()} + for name, style in list(plan.items()): + if style != "ep_dispatch_experts": + continue + parent = name.rpartition(".")[0] + router = next((f"{parent}.{key}" for key in ("gate", "router") if f"{parent}.{key}" in modules), None) + if router is None: + raise ValueError(f"Cannot find a gate or router for masked expert parallelism at `{name}`.") + plan[name] = "moe_tp_experts" + plan[router] = "ep_router" grids = {} for name, module in model.named_modules(): if not isinstance(module, FineGrainedExperts): diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index d9dde7eb6654..89fd5aac8669 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -24,7 +24,6 @@ import os import re -import subprocess import tempfile from collections.abc import Callable from contextlib import ExitStack, contextmanager @@ -33,7 +32,7 @@ import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, CompileConfig from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig @@ -46,11 +45,11 @@ Expectations, cleanup, get_device_properties, - get_torch_dist_unique_port, require_accelerate, require_kernels, require_torch_accelerator, require_torch_multi_accelerator, + set_model_for_less_flaky_test, slow, torch_device, ) @@ -63,10 +62,11 @@ ) from ..integrations.finegrained.utils import _Cfg, _conversion_model, _model_mapping, _quantizer_for +from ..test_tensor_parallel_mixin import _init_distributed def _cheap_tuning(): - """One autotune trial per shape: a cold tune dwarfs everything a test runs.""" + """Use the minimum Bayesian search budget for correctness tests.""" return mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}) @@ -624,35 +624,67 @@ def test_save_writes_the_checkpoint_layout_back(self): self.assertEqual(raw, source_raw) -_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 and differ only in the plan -for mode in modes: - model = model_cls.from_pretrained( - model_dir, - dtype="auto", - attn_implementation="eager", - distributed_config=DistributedConfig(tp_size=world, ep_size=world if mode == "ep" else None), - ).eval() - experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) - weight = getattr(experts, experts.gate_up_name) - 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 _check_static_cache_decode(model, ids): + from torch._dynamo.backends.registry import lookup_backend + + generation_kwargs = dict( + input_ids=ids.to(model.device), + attention_mask=torch.ones_like(ids, device=model.device), + max_new_tokens=3, + do_sample=False, + eos_token_id=None, + pad_token_id=0, + cache_implementation="static", + output_logits=True, + return_dict_in_generate=True, + ) + # Match the generation compile tests' norm stabilization for tiny random models. + set_model_for_less_flaky_test(model) + eager = model.generate(**generation_kwargs, disable_compile=True) + graphs = [] + + def compile_backend(graph, inputs, **kwargs): + graphs.append(graph) + return lookup_backend("inductor")(graph, inputs, **kwargs) + + compiled = model.generate( + **generation_kwargs, + compile_config=CompileConfig(backend=compile_backend, fullgraph=True, dynamic=False, mode="default"), + ) + assert graphs, "Static-cache generation did not compile the decode forward" + torch.testing.assert_close(compiled.sequences, eager.sequences) + for actual, expected in zip(compiled.logits, eager.logits): + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + +def _check_sharded_load_paths(rank, model_dir, model_cls, modes, out_dir): + from transformers.distributed import DistributedConfig + + world = torch.distributed.get_world_size() + ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + for mode in modes: + model = model_cls.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, ep_size=world if mode == "ep" else None), + ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, experts.gate_up_name) + local = weight.to_local() if hasattr(weight, "to_local") else weight + compiled_decode = False + with torch.no_grad(): + logits = model(ids.to(model.device)).logits.float().cpu() + if model._can_compile_fullgraph: + _check_static_cache_decode(model, ids) + compiled_decode = True + if rank == 0: + torch.save( + {"logits": logits, "expert_local": tuple(local.shape), "compiled_decode": compiled_decode}, + os.path.join(out_dir, mode + ".pt"), + ) + del model + torch.accelerator.empty_cache() def _checkpoint_expert_shape(model_dir): @@ -918,7 +950,10 @@ def _model_table(): def setUpClass(cls): super().setUpClass() cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) - cls._env = {"FINEGRAINED_AUTOTUNE_TRIALS": "1", "TRITON_CACHE_DIR": os.path.join(cls._tmp, "triton")} + cls._env = { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton")), + } cls.enterClassContext(mock.patch.dict(os.environ, cls._env)) def tearDown(self): @@ -996,26 +1031,10 @@ def perturb(module, args, output): 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, "sharding_worker.py") - with open(script, "w", encoding="utf-8") as fh: - fh.write(_SHARDING_WORKER) + """`{mode: payload}` from two ranks, through the real EP / TP load path.""" out = os.path.join(self._tmp, "out") os.makedirs(out, exist_ok=True) - subprocess.run( - [ - "torchrun", - "--nproc_per_node=2", - f"--master_port={get_torch_dist_unique_port()}", - script, - model_dir, - out, - ",".join(modes), - f"{model_cls.__module__}:{model_cls.__name__}", - ], - check=True, - env={**os.environ, "TOKENIZERS_PARALLELISM": "false"}, - ) + _init_distributed(tp=2)(_check_sharded_load_paths)(model_dir, model_cls, modes, out) return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} @slow @@ -1042,6 +1061,8 @@ def test_every_load_path_agrees(self): 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): + if model_cls._can_compile_fullgraph: + self.assertTrue(payload["compiled_decode"]) local = payload["expert_local"] # against the checkpoint's shape, never another mode's shard self.assertLess( diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py index 35c70d3462dd..ece764b0dd6f 100644 --- a/tests/quantization/finegrained/test_finegrained.py +++ b/tests/quantization/finegrained/test_finegrained.py @@ -380,6 +380,35 @@ def tp_plan(dtype_plan): with self.assertRaises(ValueError): tp_plan({r"layers\.1\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + def test_ep_plans_use_masked_routing_for_every_format(self): + configs = [spec.config_cls() for spec in FORMATS] + configs.append(FineGrainedConfig(dtype_plan={r"\.experts($|\.)": Mxfp4Config(), ".*": BlockFp8Config()})) + for config in configs: + for router in ("gate", "router"): + with self.subTest(config=type(config).__name__, router=router): + model = nn.Module() + model.config = SimpleNamespace(_experts_implementation=None) + model.layers = nn.ModuleList([nn.Module()]) + model.layers[0].mlp = nn.Module() + setattr(model.layers[0].mlp, router, nn.Module()) + model.layers[0].mlp.experts = nn.Module() + plan = { + "layers.*.mlp.experts": "ep_dispatch_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + quantizer = _quantizer_for(config) + quantizer.update_parallel_plan(model, plan) + expected = { + "layers.*.mlp.experts": "moe_tp_experts", + f"layers.*.mlp.{router}": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + self.assertEqual(plan, expected) + quantizer.update_parallel_plan(model, plan) + self.assertEqual(plan, expected) + def test_parallel_plans_replicate_the_per_tensor_scales(self): """A per-tensor scale is 0-dim, with no axis to split: the quantizer names each one a plan would shard with its module `replicate`, which the parameter's own entry makes take precedence. A scale grid follows its diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index 41b9dad1c7ef..7eee6cc291c6 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -79,12 +79,12 @@ def test_generate_matches_forward_logits(self): tmp, quantization_config=Mxfp4Config(), dtype=torch.bfloat16, device_map=torch_device ) - prompt = torch.tensor([[1, 2, 3, 4, 5]], device=torch_device) + prompt = torch.arange(1, 17, device=torch_device).unsqueeze(0) for use_cache in (False, True): with self.subTest(use_cache=use_cache), torch.no_grad(): out = quantized.generate( prompt, - max_new_tokens=6, + max_new_tokens=3, do_sample=False, use_cache=use_cache, output_logits=True, From 5d234d3d77ef11dae400ffdd818c62a722973ba7 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 04:57:46 -0700 Subject: [PATCH 77/91] Apply minimal tuning to standalone generation tests and fix style --- tests/quantization/block_scaled_tester.py | 22 +++++++++++----------- tests/quantization/mxfp4/test_mxfp4.py | 4 ++++ 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index 89fd5aac8669..9a01406cc5be 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -627,17 +627,17 @@ def test_save_writes_the_checkpoint_layout_back(self): def _check_static_cache_decode(model, ids): from torch._dynamo.backends.registry import lookup_backend - generation_kwargs = dict( - input_ids=ids.to(model.device), - attention_mask=torch.ones_like(ids, device=model.device), - max_new_tokens=3, - do_sample=False, - eos_token_id=None, - pad_token_id=0, - cache_implementation="static", - output_logits=True, - return_dict_in_generate=True, - ) + generation_kwargs = { + "input_ids": ids.to(model.device), + "attention_mask": torch.ones_like(ids, device=model.device), + "max_new_tokens": 3, + "do_sample": False, + "eos_token_id": None, + "pad_token_id": 0, + "cache_implementation": "static", + "output_logits": True, + "return_dict_in_generate": True, + } # Match the generation compile tests' norm stabilization for tiny random models. set_model_for_less_flaky_test(model) eager = model.generate(**generation_kwargs, disable_compile=True) diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index 7eee6cc291c6..f395effee93b 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -46,12 +46,16 @@ BlockScaledFormatTesterMixin, LoadPathEquivalenceTesterMixin, TinyCheckpointTesterMixin, + _cheap_tuning, ) @require_torch_accelerator @require_kernels class Mxfp4GenerateTest(unittest.TestCase): + def setUp(self): + self.enterContext(_cheap_tuning()) + def test_generate_matches_forward_logits(self): """`generate` switches the experts to the batched kernels, while a plain forward runs the grouped ones: on a tiny GPT-OSS with biased experts, quantized on the fly (W4A16), every decoding step's logits must match From 040d4a9c50f5ebd1fc4b6374851642f4470b64ca Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 05:00:13 -0700 Subject: [PATCH 78/91] Reuse tuning context in distributed quantization tests --- tests/quantization/block_scaled_tester.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index 9a01406cc5be..fc07f1a069f3 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -950,11 +950,12 @@ def _model_table(): def setUpClass(cls): super().setUpClass() cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) - cls._env = { - "FINEGRAINED_AUTOTUNE_TRIALS": "1", - "TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton")), - } - cls.enterClassContext(mock.patch.dict(os.environ, cls._env)) + cls.enterClassContext(_cheap_tuning()) + cls.enterClassContext( + mock.patch.dict( + os.environ, {"TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton"))} + ) + ) def tearDown(self): super().tearDown() From aea13cac2c84951dd9eb623e2293d4361ffb5598 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 05:15:36 -0700 Subject: [PATCH 79/91] Simplify finegrained conversions and separate distributed test coverage --- docs/source/en/quantization/blockfp8.md | 2 +- src/transformers/core_model_loading.py | 42 +- src/transformers/integrations/deepgemm.py | 6 +- .../integrations/finegrained/conversions.py | 54 +- .../integrations/finegrained/layers.py | 57 +- src/transformers/integrations/moe.py | 10 +- src/transformers/utils/quantization_config.py | 2 +- tests/integrations/finegrained/test_layers.py | 61 ++ tests/kernels/test_deepgemm.py | 7 + .../block_scaled_checkpoint_tester.py | 285 +++++++ .../block_scaled_parallel_tester.py | 513 ++++++++++++ tests/quantization/block_scaled_tester.py | 742 +----------------- tests/quantization/blockfp8/test_blockfp8.py | 4 +- .../finegrained/test_finegrained.py | 4 +- tests/quantization/mxfp4/test_mxfp4.py | 4 +- tests/quantization/mxfp8/test_mxfp8.py | 4 +- tests/quantization/nvfp4/test_nvfp4.py | 4 +- tests/utils/test_core_model_loading.py | 7 + 18 files changed, 971 insertions(+), 837 deletions(-) create mode 100644 tests/quantization/block_scaled_checkpoint_tester.py create mode 100644 tests/quantization/block_scaled_parallel_tester.py diff --git a/docs/source/en/quantization/blockfp8.md b/docs/source/en/quantization/blockfp8.md index 9a25379fd036..68b081208f82 100644 --- a/docs/source/en/quantization/blockfp8.md +++ b/docs/source/en/quantization/blockfp8.md @@ -93,7 +93,7 @@ model = AutoModelForCausalLM.from_pretrained("llama-3.1-8b-fp8", device_map="aut - `weight_block_size` is the `(rows, columns)` block each weight scale covers; `None` gives one scale per tensor. - `activation_scheme="static"` quantizes activations against the calibrated scale a checkpoint ships (its `input_scale`), instead of deriving one per token at run time. -- `scale_fmt="ue8m0"` holds the block scales as powers of two (`float8_e8m0fnu`), as DeepSeek-V4 ships them. +- `scale_fmt="ue8m0"` holds the block scales as powers of two (`float8_e8m0fnu`, with a `uint8` exponent fallback when the dtype is unavailable), as DeepSeek-V4 ships them. Float32 scale containers, such as DeepSeek-V4-Base's, load and save as native UE8M0. - `modules_to_not_convert` keeps the modules it names in full precision; `modules_to_convert` names embedding tables to hold in FP8 with one per-tensor scale (a table too large to keep in full precision). diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index d31c5f6d0861..f0c2bf3ab6e6 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,6 +81,8 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" + supports_many_to_many = False + def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -647,6 +649,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 @@ -709,6 +713,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 @@ -1180,18 +1186,6 @@ 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 so the finegrained integration, which imports this module, - is imported at call time.""" - from .integrations.finegrained.conversions import FineGrainedWeightGlobals - - return ( - ErnieFuseAndSplitTextVisionExperts, - ErnieSplitAndDecoupleTextVisionExperts, - FineGrainedWeightGlobals, - ) - - class WeightConverter(WeightTransform): __slots__ = ("operations", "force_cpu") @@ -1208,7 +1202,7 @@ def __init__( 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): + 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." ) @@ -1863,30 +1857,24 @@ 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} - # `(layer name, reversed transform)` per realized group - conversion_mapping: dict[str, tuple[str, WeightTransform]] = {} + conversion_mapping: dict[str, WeightTransform] = {} 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) + converted_key, source_pattern = rename_source_key(original_key, [], converters, reverse=True) if source_pattern is not None: - # 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(original_key, [], converters, reverse=True) - # each target key gets its own converter instance + # Apply checkpoint renames after conversion, while the converter can still resolve the model's names. if converted_key not in conversion_mapping: - converter = deepcopy(pattern_to_converter[source_pattern]) - conversion_mapping[converted_key] = (converted_key, converter) - _, mapping = conversion_mapping[converted_key] + conversion_mapping[converted_key] = deepcopy(pattern_to_converter[source_pattern]) + mapping = conversion_mapping[converted_key] mapping.add_tensor(converted_key, original_key, source_pattern, tensor) else: - mapping = conversion_mapping.setdefault( - renamed_key, (renamed_key, WeightRenaming(original_key, renamed_key)) - )[1] + renamed_key, _ = rename_source_key(original_key, renamings, [], reverse=True) + mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) mapping.add_tensor(renamed_key, original_key, original_key, tensor) new_state_dict = {} - for layer_name, reversed_transform in conversion_mapping.values(): + for layer_name, reversed_transform in conversion_mapping.items(): # Apply the reverse converter realized_value = reversed_transform.convert(layer_name, model=model, config=model.config) for target_name, param in realized_value.items(): diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index ed29617603af..1396a9206017 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -342,14 +342,14 @@ def _coerce_sf_for_kernel(sf: torch.Tensor, is_sm100: bool, expected_mn: int | N DeepGEMM kernel branch is the only UE8M0 path on SM100; for `gran_mn > 1` the kernel only handles FP32 SFs and would otherwise reject our INT SF here. """ - if sf.dtype == torch.float8_e8m0fnu: + if sf.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)): if expected_mn is not None and sf.size(-2) < expected_mn: gran_mn = expected_mn // sf.size(-2) sf = sf.repeat_interleave(gran_mn, dim=-2) if is_sm100: sf = sf.contiguous().view(torch.int32) else: - sf = sf.float() + sf = (sf.float() - 127).exp2() if sf.dtype == torch.uint8 else sf.float() elif sf.dtype == torch.float32 and is_sm100: sf = _ceil_to_ue8m0(sf) @@ -397,7 +397,7 @@ def _select_fp8_cast_kwargs( block_size = tuple(block_size) if block_size not in ((128, 128), (1, 128)): raise ValueError(f"DeepGEMM requires `block_size` ∈ {{(128, 128), (1, 128)}}, got {block_size}.") - if weight_scale_inv.dtype == torch.float8_e8m0fnu and is_sm100: + if weight_scale_inv.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) and is_sm100: return {"use_ue8m0": True, "gran_k": 128, "use_packed_ue8m0": True} return {"use_ue8m0": False, "gran_k": 128} diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 59393e23bb57..e9f454823087 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -35,8 +35,8 @@ _cdiv, _FineGrainedModule, _get_ue8m0_dtype, + get_weight_format, load_finegrained_kernel, - weight_formats, ) @@ -75,8 +75,10 @@ def held_scale(model, full_layer_name, key): def finegrained_module(model, full_layer_name): - """The finegrained module the tensor at ``full_layer_name`` fills, else ``None``. On save the name - runs past the module (``experts.*.gate_proj.weight``).""" + """The name and finegrained module owning a tensor, or ``(None, None)``. + + On save the name can run past the module (``experts.*.gate_proj.weight``). + """ name = full_layer_name or "" while model is not None and name: name = name.rpartition(".")[0] @@ -84,19 +86,26 @@ def finegrained_module(model, full_layer_name): module = model.get_submodule(name) except AttributeError: continue - return module if isinstance(module, _FineGrainedModule) else None - return None + return (name, module) if isinstance(module, _FineGrainedModule) else (None, None) + return None, None def finegrained_experts(model, full_layer_name): """The ``FineGrainedExperts`` the tensor at ``full_layer_name`` fills, else ``None``.""" - module = finegrained_module(model, full_layer_name) + _, module = finegrained_module(model, full_layer_name) return module if isinstance(module, FineGrainedExperts) else 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) + """Reinterpret UE8M0 bytes or convert its power-of-two values between containers.""" + if scale.element_size() == dtype.itemsize: + return scale.view(dtype) + if dtype == torch.uint8: + # UE8M0 is a biased exponent, never a numeric float-to-integer cast. + return (scale.float().log2().ceil() + 127).clamp(0, 255).masked_fill(torch.isnan(scale), 255).to(torch.uint8) + if scale.dtype == torch.uint8: + return (scale.float() - 127).exp2().masked_fill(scale == 255, float("nan")).to(dtype) + return scale.to(dtype) class _FineGrainedOp(ConversionOps): @@ -131,7 +140,7 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= input_dict = keyed_by_target(input_dict, target_patterns) if not self.hf_quantizer.pre_quantized: return input_dict - module = finegrained_module(model, full_layer_name) + module_name, module = finegrained_module(model, full_layer_name) if module is None and self.inverse: # Saving names the first checkpoint output (e.g. gate_proj), while the # model holds the fused input (gate_up_proj). Recover that module before @@ -139,14 +148,13 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= for target in target_patterns or []: for source in kwargs.get("source_patterns", []): candidate = (full_layer_name or "").replace(target, source.replace(r"\.", "."), 1) - module = finegrained_module(model, candidate) + module_name, module = finegrained_module(model, candidate) if module is not None: break if module is not None: break if module is None: return input_dict - module_name = next(name for name, candidate in model.named_modules() if candidate is module) result = {} for key, value in input_dict.items(): name = key.rpartition(".")[2] @@ -189,7 +197,7 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= class FineGrainedScaleContainer(_FineGrainedOp): - """A UE8M0 block scale into the ``float8_e8m0fnu`` its module holds: it ships either as the + """A UE8M0 block scale into the native dtype or exponent bytes its module holds: it ships either as the exponent byte under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast exactly). Exponent-byte checkpoints save back as uint8; float32 containers save as native UE8M0. Every other scale @@ -202,14 +210,20 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= if self.inverse: # the container its module recorded at load; e8m0 is only ever a scale, so the dtype # alone picks the tensors to restore - module = self.module or finegrained_module(model, full_layer_name) + module = self.module + if module is None: + _, module = finegrained_module(model, full_layer_name) container = getattr(module, "scale_container_dtype", None) - if container is not None and value.dtype == _get_ue8m0_dtype(): + if container is not None and value.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)): value = as_container(value, container) else: module, held = held_scale(model, full_layer_name, key) # 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: + if ( + held is not None + and held.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) + and value.dtype != held.dtype + ): if value.dtype == torch.uint8: module.scale_container_dtype = torch.uint8 value = as_container(value, held.dtype) @@ -300,6 +314,8 @@ 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.""" + supports_many_to_many = True + @staticmethod def role(key: str) -> tuple[str, str]: """What a global-scale key holds, for a checkpoint key and a module target alike: which @@ -471,7 +487,7 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc module, scale_name = holder if holder is not None else (None, None) value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) 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 + format_spec = get_weight_format(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) @@ -486,7 +502,7 @@ 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.dtype == _get_ue8m0_dtype() + held.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) if held is not None else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" ) @@ -558,7 +574,9 @@ def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool 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 when the shape tiled - return quantized[..., :rows, :cols].contiguous(), inv_scale.to(_get_ue8m0_dtype() if ue8m0 else torch.float32) + return quantized[..., :rows, :cols].contiguous(), as_container( + inv_scale, _get_ue8m0_dtype() + ) if ue8m0 else inv_scale @staticmethod def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index 81f31fa1375f..bc60cf2887d7 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -53,17 +53,8 @@ @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 + """Use native UE8M0 when available, otherwise store its biased exponent bytes.""" + return getattr(torch, "float8_e8m0fnu", torch.uint8) def _first_attr(obj, *names, raise_error: bool = True): @@ -180,25 +171,27 @@ class WeightFormat: global_scale_dtype: torch.dtype | None = None -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 - "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( +def get_weight_format(weight_format: str) -> WeightFormat: + """Resolve only the requested format, so FP8 and NVFP4 do not require UE8M0 support.""" + if weight_format == "fp8": + return WeightFormat(weight_dtype=_FP8_DTYPE) + if weight_format in ("mxfp8", "mxfp4"): + packed = weight_format == "mxfp4" + return WeightFormat( + weight_dtype=torch.int8 if packed else _FP8_DTYPE, + values_per_byte=2 if packed else 1, + scale_dtype=_get_ue8m0_dtype(), + scale_group=(1, 32), + ) + if weight_format == "nvfp4": + return WeightFormat( weight_dtype=torch.int8, values_per_byte=2, scale_dtype=_FP8_DTYPE, scale_group=(1, 16), global_scale_dtype=torch.float32, - ), - } + ) + raise ValueError(f"unknown weight_format {weight_format!r}; expected fp8, mxfp8, mxfp4 or nvfp4") def resolve_weight_format( @@ -208,10 +201,7 @@ 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() - if weight_format not in formats: - raise ValueError(f"unknown weight_format {weight_format!r}; expected one of {sorted(formats)}") - format_spec = formats[weight_format] + format_spec = get_weight_format(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 @@ -370,9 +360,7 @@ def finegrained_linear( 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 + # Preserve uint8 exponent containers on save; float32 containers normalize to native UE8M0. scale_container_dtype: torch.dtype | None = None # the kernels take raw pointers, so the TP layer must hand local shards and not DTensors, whatever @@ -690,9 +678,10 @@ def from_module( "is_concatenated": getattr(module, "is_concatenated", True), } experts_class = _experts_class_with(cls, **flags) - config = getattr(module, "config", model.config.get_text_config()) + config = getattr(module, "config", None) or model.config.get_text_config() experts = experts_class(config=config, quantization_config=quantization_config, **flags) - experts.post_expert_norm, experts.has_post_expert_norm = module.post_expert_norm, module.has_post_expert_norm + experts.post_expert_norm = getattr(module, "post_expert_norm", None) + experts.has_post_expert_norm = getattr(module, "has_post_expert_norm", False) # a full-precision stack quantized on load arrives in the layout of the module this replaces experts.ships_transposed = getattr(module, "is_transposed", False) return experts diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 4ab8138db4fc..2e006c459dd5 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -164,9 +164,9 @@ 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. + # Apply optional post expert norm if self.has_post_expert_norm: - proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) + proj_out = self.post_expert_norm(proj_out) # Apply routing weights if sentinel_mask is not None: @@ -484,9 +484,9 @@ 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. + # Apply optional post expert norm if self.has_post_expert_norm: - proj_out = self.post_expert_norm(proj_out) # (S, hidden_dim) + proj_out = self.post_expert_norm(proj_out) # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) @@ -594,8 +594,6 @@ def __init__(self, config, *args, **kwargs): self.has_post_expert_norm = has_post_expert_norm self.is_transposed = is_transposed self.is_concatenated = is_concatenated - self.post_expert_norm = getattr(self, "post_expert_norm", None) - self.has_post_expert_norm = self.post_expert_norm is not None self._is_expert_parallel = False @wraps(original_forward) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 0d6668a41bbf..ffff4be62680 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -2024,7 +2024,7 @@ def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None: - """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, MX, + """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, per-channel, sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's `checkpoint_format` names the layout its tensors ship in, which a save writes back.""" if config_dict.get("sparsity_config") or config_dict.get("transform_config"): diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py index 581b9f03cb04..0b12d8c4348b 100644 --- a/tests/integrations/finegrained/test_layers.py +++ b/tests/integrations/finegrained/test_layers.py @@ -48,6 +48,67 @@ def _block_scaled_matmul_2d(A, B, As, Bs, *, output_dtype=None, **kwargs): class FineGrainedModulesTest(unittest.TestCase): """The finegrained modules themselves: what they hold and what replaces what.""" + def test_non_mx_formats_do_not_require_ue8m0(self): + from transformers.integrations.finegrained import conversions + + with ( + mock.patch.object(fg, "_get_ue8m0_dtype", side_effect=RuntimeError("UE8M0 unavailable")), + mock.patch.object(conversions, "_get_ue8m0_dtype", side_effect=RuntimeError("UE8M0 unavailable")), + ): + for config in (BlockFp8Config(), NVFP4Config()): + with self.subTest(format=config.method): + linear = FineGrainedLinear(64, 64, quantization_config=config) + self.assertEqual(linear.weight.element_size(), 1) + model = torch.nn.Module() + config = BlockFp8Config() + model.linear = FineGrainedLinear(64, 64, quantization_config=config) + quantizer = _quantizer_for(config) + quantizer.pre_quantized = False + result = quantizer.get_quantize_ops().convert({"linear.weight": torch.randn(64, 64)}, model=model) + self.assertEqual(result["linear.weight"].dtype, torch.float8_e4m3fn) + self.assertEqual(result["linear.weight_scale_inv"].dtype, torch.float32) + + def test_ue8m0_falls_back_to_exponent_bytes(self): + from transformers.integrations.finegrained.conversions import FineGrainedQuantize, as_container + from transformers.utils.quantization_config import Mxfp8Config + + fg._get_ue8m0_dtype.cache_clear() + try: + with mock.patch.object(torch, "float8_e8m0fnu", create=True): + del torch.float8_e8m0fnu + self.assertEqual(fg._get_ue8m0_dtype(), torch.uint8) + linear = FineGrainedLinear(64, 64, quantization_config=Mxfp8Config()) + self.assertEqual(linear.weight_scale_inv.dtype, torch.uint8) + scales = torch.tensor([0.5, 1.0, 2.0, 4.0]) + encoded = as_container(scales, torch.uint8) + torch.testing.assert_close(encoded, torch.tensor([126, 127, 128, 129], dtype=torch.uint8)) + torch.testing.assert_close(as_container(encoded, torch.float32), scales) + _, scale = FineGrainedQuantize._quantize_block_fp8(torch.full((2, 2), 448.0), (2, 2), True) + torch.testing.assert_close(scale, torch.tensor([[127]], dtype=torch.uint8)) + finally: + fg._get_ue8m0_dtype.cache_clear() + + def test_replacement_respects_a_disabled_post_expert_norm(self): + from transformers.integrations.moe import use_experts_implementation + + @use_experts_implementation(has_post_expert_norm=False) + class Experts(torch.nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.post_expert_norm = torch.nn.RMSNorm(config.hidden_size) + + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states + + config = _Cfg() + source = Experts(config) + self.assertFalse(source.has_post_expert_norm) + model = SimpleNamespace(config=SimpleNamespace(get_text_config=lambda: config)) + quantized = FineGrainedExperts.from_module(source, BlockFp8Config(), model) + self.assertIs(quantized.post_expert_norm, source.post_expert_norm) + self.assertFalse(quantized.has_post_expert_norm) + def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): """`replace_with_quantized_modules` picks `FineGrainedGroupedLinear` for a block-diagonal linear — a plain one would collapse the groups into a single giant matmul and return the diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index 09ce53cfa09b..2316a6b475c1 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -96,6 +96,13 @@ class _DeepGemmKernelMissingSymbol(_FakeDeepGemmKernel): @require_torch class DeepGemmLoaderTest(unittest.TestCase): + def test_uint8_scale_containers_use_exponents_on_both_architectures(self): + scales = torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8).repeat(2, 1) + hopper = dg._coerce_sf_for_kernel(scales, is_sm100=False) + torch.testing.assert_close(hopper, torch.tensor([[0.5, 1.0, 2.0, 4.0]]).repeat(2, 1)) + blackwell = dg._coerce_sf_for_kernel(scales, is_sm100=True) + torch.testing.assert_close(blackwell, scales.view(torch.int32)) + def setUp(self): dg._DEEPGEMM = None self.addCleanup(setattr, dg, "_DEEPGEMM", None) diff --git a/tests/quantization/block_scaled_checkpoint_tester.py b/tests/quantization/block_scaled_checkpoint_tester.py new file mode 100644 index 000000000000..730b2193db37 --- /dev/null +++ b/tests/quantization/block_scaled_checkpoint_tester.py @@ -0,0 +1,285 @@ +# 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. +"""Accuracy and serialization checks against independently generated tiny checkpoints.""" + +import os +import re + +import torch + +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText +from transformers.testing_utils import ( + Expectations, + cleanup, + get_device_properties, + require_kernels, + require_torch_accelerator, + slow, + torch_device, +) + +from .block_scaled_tester import FormatSpec, _cheap_tuning, _require_kernel_capability + + +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 + + 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(): + tensor = handle.get_tensor(key) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + tensors[key] = (tensor.dtype, tuple(tensor.shape), raw) + return tensors + + +# The logits `[0, :3, :8]` each checkpoint gives on `torch.arange(3, 35)`, per device: 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.3086, -0.2773, -0.4883, -0.1797, -0.0073, 0.4531, 0.2158, 0.2754], + [0.0295, -0.2598, -0.5586, 0.2432, -0.0106, 0.3496, 0.4805, 0.1279], + [0.1533, -0.4551, -0.5742, -0.1973, 0.2139, 0.4141, 0.3574, 0.1338], + ], + } + ), # 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)): [ + [0.2637, -0.0815, -0.3418, 0.1553, -0.2715, 0.1562, 0.4590, 0.0635], + [0.4355, -0.1016, -0.2002, -0.2412, -0.2930, 0.1187, 0.5664, 0.1104], + [0.1865, -0.0996, -0.6875, 0.2988, 0.0635, 0.1035, 0.5586, -0.0703], + ], + } + ), # fmt: skip +} + + +class TinyCheckpointTesterMixin: + """A format's pushed tiny checkpoints, random copies of the shipped ones written in their producers' + layouts (modelopt, OpenAI's and DeepSeek's quantizers): each loads with every key accounted for, computes + what its dequantized weights compute within the format's rounding, matches its recorded logits, and saves + back to its own layout. So a converter is checked against what the format really is, not against what this + code believes it to be. A format's test file sets `spec`.""" + + spec: FormatSpec + + def setUp(self): + super().setUp() + self.enterContext(_cheap_tuning()) + _require_kernel_capability(self) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + @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) + + @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)} + + def _reference_error(self, model, repo): + """The logits' error against the dequantized reference, and the logits.""" + 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(): + got, want = model(ids).logits.float(), reference(ids).logits.float() + return ((got - want).norm() / want.norm()).item(), got + + @staticmethod + def _recorded_slice(repo): + """The slice recorded on this exact device, if any: another device's kernels round differently.""" + return dict(EXPECTED_SLICES[repo].unpacked()).get(get_device_properties()) + + @slow + @require_kernels + @require_torch_accelerator + def test_logits_match_the_dequantized_reference_and_the_recorded_slice(self): + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + 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()) + ) + error, logits = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + expected = self._recorded_slice(repo) + if expected is not None: + torch.testing.assert_close(logits[0, :3, :8].cpu(), torch.tensor(expected), rtol=1e-2, atol=5e-3) + + @slow + @require_kernels + @require_torch_accelerator + def test_dequantized_load_matches_the_reference(self): + """`dequantize=True` folds each weight against its scales (and its global, for a two-level format).""" + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + model, info = self._load(repo, quantization_config=self.spec.config_cls(dequantize=True)) + unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) + self.assertEqual((set(info["missing_keys"]), unexpected), (set(), set())) + # nothing stays quantized; a model keeps its own fp32 parameters (DeepSeek-V4's hyper-connections) + quantized = {torch.float8_e4m3fn, torch.float8_e8m0fnu, torch.int8, torch.uint8} + self.assertEqual([name for name, param in model.named_parameters() if param.dtype in quantized], []) + error, _ = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + + @slow + @require_kernels + @require_torch_accelerator + def test_save_writes_the_checkpoint_layout_back(self): + """A pre-quantized model saved again lands on its checkpoint's keys, dtypes and bytes.""" + from huggingface_hub import snapshot_download + + for repo in self.spec.tiny_checkpoints: + with self.subTest(repo=repo): + 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(repo=repo, key=key): + source_dtype, source_shape, source_raw = source[key] + self.assertEqual(shape, source_shape) + if dtype != source_dtype: + # Loading honors initialized dtypes, including FP32 norms and BF16 + # MiniMax routers/patch embeddings. Float32 UE8M0 containers save natively. + self.assertTrue( + ( + dtype == torch.float32 + and key.endswith("norm.weight") + and source_dtype == torch.bfloat16 + ) + or ( + dtype == torch.bfloat16 + and source_dtype == torch.float32 + and key.endswith( + ("gate.weight", "e_score_correction_bias", "patch_embedding.weight") + ) + ) + or ( + dtype == torch.float8_e8m0fnu + and source_dtype == torch.float32 + and key.endswith((".scale", "_scale_inv")) + ) + ) + expected = torch.frombuffer(bytearray(source_raw), dtype=source_dtype).to(dtype) + self.assertEqual(raw, expected.view(torch.uint8).numpy().tobytes()) + else: + self.assertEqual(raw, source_raw) diff --git a/tests/quantization/block_scaled_parallel_tester.py b/tests/quantization/block_scaled_parallel_tester.py new file mode 100644 index 000000000000..08c46f098fec --- /dev/null +++ b/tests/quantization/block_scaled_parallel_tester.py @@ -0,0 +1,513 @@ +# 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. +"""Distributed loading, sharding, and compiled static-cache decode checks.""" + +import os +import re +import tempfile +from unittest import mock + +import torch + +from transformers import CompileConfig +from transformers.testing_utils import ( + cleanup, + require_kernels, + require_torch_multi_accelerator, + set_model_for_less_flaky_test, + slow, + torch_device, +) +from transformers.utils.quantization_config import BlockFp8Config, Mxfp4Config, Mxfp8Config + +from ..test_tensor_parallel_mixin import _init_distributed +from .block_scaled_checkpoint_tester import _checkpoint_bytes +from .block_scaled_tester import NVFP4, FormatSpec, _cheap_tuning + + +def _check_static_cache_decode(model, ids): + from torch._dynamo.backends.registry import lookup_backend + + generation_kwargs = { + "input_ids": ids.to(model.device), + "attention_mask": torch.ones_like(ids, device=model.device), + "max_new_tokens": 3, + "do_sample": False, + "eos_token_id": None, + "pad_token_id": 0, + "cache_implementation": "static", + "output_logits": True, + "return_dict_in_generate": True, + } + # Match the generation compile tests' norm stabilization for tiny random models. + set_model_for_less_flaky_test(model) + eager = model.generate(**generation_kwargs, disable_compile=True) + graphs = [] + + def compile_backend(graph, inputs, **kwargs): + graphs.append(graph) + return lookup_backend("inductor")(graph, inputs, **kwargs) + + compiled = model.generate( + **generation_kwargs, + compile_config=CompileConfig(backend=compile_backend, fullgraph=True, dynamic=False, mode="default"), + ) + assert graphs, "Static-cache generation did not compile the decode forward" + torch.testing.assert_close(compiled.sequences, eager.sequences) + for actual, expected in zip(compiled.logits, eager.logits): + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + +def _check_sharded_load_paths(rank, model_dir, model_cls, modes, out_dir, compile_decode): + from transformers.distributed import DistributedConfig + + world = torch.distributed.get_world_size() + ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + for mode in modes: + model = model_cls.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, ep_size=world if mode == "ep" else None), + ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, experts.gate_up_name) + local = weight.to_local() if hasattr(weight, "to_local") else weight + payload = {"expert_local": tuple(local.shape)} + with torch.no_grad(): + if compile_decode: + _check_static_cache_decode(model, ids) + else: + payload["logits"] = model(ids.to(model.device)).logits.float().cpu() + if rank == 0: + torch.save(payload, 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 fused (`experts.gate_up_proj`) or one per expert + (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`).""" + import glob + import json + + 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 + + +class LoadPathEquivalenceTesterMixin: + """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. + + 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 not run. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + + A format's test file sets `spec`; the fixtures are its `distributed_models`.""" + + spec: FormatSpec + + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}`: each model in the format and + under the quantization config it ships with. 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, + ) + + 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, + }, + BlockFp8Config(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, + }, + BlockFp8Config(weight_block_size=None, activation_scheme="static"), + ), + # multimodal: the experts' plans live on `text_config`, not on the config the quantizer is handed + "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, + ), + }, + NVFP4.checkpoint_config(), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases its EP plan already names + "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, + }, + Mxfp4Config(activation_format="bf16"), + ), + # mixed: `expert_dtype` makes the experts MXFP4 over block-FP8 with UE8M0 scales + "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", + }, + BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second multimodal nesting; the MoE dims are raised so a 2-way split stays 128-aligned + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + # the 2-way split leaves 256: a swizzled tile is 128 wide along the down's K + "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": { + # the 2-way split leaves 128, one swizzled K tile + "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, + }, + Mxfp8Config(), + ), + } + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) + cls._fixtures = {} + cls.enterClassContext(_cheap_tuning()) + cls.enterClassContext( + mock.patch.dict( + os.environ, {"TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton"))} + ) + ) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + def _fixture(self, label): + """A pre-quantized checkpoint of one tiny model, in the format that model ships.""" + if label in self._fixtures: + return self._fixtures[label] + config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] + config = config_cls(**kwargs) + + bf16_dir = os.path.join(self._tmp, label, "bf16") + quant_dir = os.path.join(self._tmp, label, "quantized") + torch.manual_seed(0) + # saved in bf16, the dtype every model here ships, so `dtype="auto"` reloads bf16 + 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) + + # a reload saved again lands on the same bytes: every layout op's reverse restores the checkpoint + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") + round_trip = os.path.join(self._tmp, 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") + self._fixtures[label] = quant_dir, config, model_cls + return self._fixtures[label] + + 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. How far + that travels is a property of the model (a chain of MoE layers can amplify it a hundredfold), so the + same magnitude injected on one device gives each model the budget its own conditioning earns. + """ + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" + ).eval() + # every block a sharded reduction passes through: the rounding accumulates down the stack + 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, compile_decode=False): + """`{mode: payload}` from two ranks, through the real EP / TP load path.""" + out = os.path.join(self._tmp, "out") + os.makedirs(out, exist_ok=True) + _init_distributed(tp=2)(_check_sharded_load_paths)(model_dir, model_cls, modes, out, compile_decode) + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} + + @slow + @require_kernels + @require_torch_multi_accelerator + def test_every_load_path_agrees(self): + for label in self.spec.distributed_models: + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # `max_memory` forces the baseline across both devices; `auto` alone fits it on one + 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) + + modes = self._parallel_modes(config) + self.assertTrue(modes, f"{label}: no plan shards experts") + 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): + expected = (whole[0] // 2, whole[1]) if mode == "ep" else (whole[0], whole[1] // 2) + self.assertEqual(payload["expert_local"][:2], expected, f"{label}/{mode}: wrong shard axis") + # a wrong shard axis lands orders of magnitude above the model's rounding floor + 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: # rounding alone does not move the argmax + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) + + @staticmethod + def _parallel_modes(config): + owners = [config] + [ + c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) + ] + return [ + 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 {})) + ] + + @slow + @require_kernels + @require_torch_multi_accelerator + def test_static_cache_decode_compiles(self): + supported = [ + label for label in self.spec.distributed_models if self._model_table()[label][1]._can_compile_fullgraph + ] + if not supported: + self.skipTest("These model architectures do not support fullgraph compilation") + for label in supported: + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + modes = self._parallel_modes(config) + self.assertTrue(modes, f"{label}: no plan shards experts") + self._sharded(model_dir, model_cls, modes, compile_decode=True) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index fc07f1a069f3..c2016691ed77 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -11,20 +11,13 @@ # 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 tests every block-scaled format runs, driven by its `FormatSpec`: a format's test file instantiates the -testers with its spec and adds only what is its own (its checkpoint layouts, its producers' quirks). +"""Shared format specifications and fast quantizer contract tests. -- `BlockScaledFormatTesterMixin`: the config, the quantizer a load picks, the activation formats it takes, the - no-accelerator fallbacks, one converter per checkpoint key (CPU); every forward against the format's - dequantized weights and the on-the-fly quantization round trip (GPU, the real kernels). -- `TinyCheckpointTesterMixin` (slow): the format's pushed tiny checkpoints load with every key accounted for, - compute what their dequantized weights compute, match the recorded logits, and save back to their layout. - `FINEGRAINED` runs it on the checkpoints that mix formats. -- `LoadPathEquivalenceTesterMixin` (slow, 2 GPUs): expert and tensor parallelism reproduce `device_map`.""" +Real-checkpoint tests live in `block_scaled_checkpoint_tester.py`; distributed +loading and compiled decode tests live in `block_scaled_parallel_tester.py`. +""" import os -import re -import tempfile from collections.abc import Callable from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field @@ -32,7 +25,7 @@ import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, CompileConfig +from transformers import AutoModelForCausalLM from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig @@ -42,15 +35,9 @@ from transformers.quantizers.quantizer_mxfp8 import Mxfp8HfQuantizer from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer from transformers.testing_utils import ( - Expectations, - cleanup, - get_device_properties, require_accelerate, require_kernels, require_torch_accelerator, - require_torch_multi_accelerator, - set_model_for_less_flaky_test, - slow, torch_device, ) from transformers.utils.quantization_config import ( @@ -62,7 +49,6 @@ ) from ..integrations.finegrained.utils import _Cfg, _conversion_model, _model_mapping, _quantizer_for -from ..test_tensor_parallel_mixin import _init_distributed def _cheap_tuning(): @@ -369,721 +355,3 @@ def test_on_the_fly_quantization_lands_in_the_module_layout(self): self.assertEqual(out[global_key].shape, experts.gate_up_proj_weight_global_scale.shape) else: self.assertNotIn(global_key, out) - - -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 - - 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(): - tensor = handle.get_tensor(key) - raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() - tensors[key] = (tensor.dtype, tuple(tensor.shape), raw) - return tensors - - -# The logits `[0, :3, :8]` each checkpoint gives on `torch.arange(3, 35)`, per device: 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.3086, -0.2773, -0.4883, -0.1797, -0.0073, 0.4531, 0.2158, 0.2754], - [0.0295, -0.2598, -0.5586, 0.2432, -0.0106, 0.3496, 0.4805, 0.1279], - [0.1533, -0.4551, -0.5742, -0.1973, 0.2139, 0.4141, 0.3574, 0.1338], - ], - } - ), # 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)): [ - [0.2637, -0.0815, -0.3418, 0.1553, -0.2715, 0.1562, 0.4590, 0.0635], - [0.4355, -0.1016, -0.2002, -0.2412, -0.2930, 0.1187, 0.5664, 0.1104], - [0.1865, -0.0996, -0.6875, 0.2988, 0.0635, 0.1035, 0.5586, -0.0703], - ], - } - ), # fmt: skip -} - - -class TinyCheckpointTesterMixin: - """A format's pushed tiny checkpoints, random copies of the shipped ones written in their producers' - layouts (modelopt, OpenAI's and DeepSeek's quantizers): each loads with every key accounted for, computes - what its dequantized weights compute within the format's rounding, matches its recorded logits, and saves - back to its own layout. So a converter is checked against what the format really is, not against what this - code believes it to be. A format's test file sets `spec`.""" - - spec: FormatSpec - - def setUp(self): - super().setUp() - self.enterContext(_cheap_tuning()) - _require_kernel_capability(self) - - def tearDown(self): - super().tearDown() - cleanup(torch_device, gc_collect=True) - - @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) - - @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)} - - def _reference_error(self, model, repo): - """The logits' error against the dequantized reference, and the logits.""" - 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(): - got, want = model(ids).logits.float(), reference(ids).logits.float() - return ((got - want).norm() / want.norm()).item(), got - - @staticmethod - def _recorded_slice(repo): - """The slice recorded on this exact device, if any: another device's kernels round differently.""" - return dict(EXPECTED_SLICES[repo].unpacked()).get(get_device_properties()) - - @slow - @require_kernels - @require_torch_accelerator - def test_logits_match_the_dequantized_reference_and_the_recorded_slice(self): - for repo, budget in self.spec.tiny_checkpoints.items(): - with self.subTest(repo=repo): - 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()) - ) - error, logits = self._reference_error(model, repo) - self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") - expected = self._recorded_slice(repo) - if expected is not None: - torch.testing.assert_close(logits[0, :3, :8].cpu(), torch.tensor(expected), rtol=1e-2, atol=5e-3) - - @slow - @require_kernels - @require_torch_accelerator - def test_dequantized_load_matches_the_reference(self): - """`dequantize=True` folds each weight against its scales (and its global, for a two-level format).""" - for repo, budget in self.spec.tiny_checkpoints.items(): - with self.subTest(repo=repo): - model, info = self._load(repo, quantization_config=self.spec.config_cls(dequantize=True)) - unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) - self.assertEqual((set(info["missing_keys"]), unexpected), (set(), set())) - # nothing stays quantized; a model keeps its own fp32 parameters (DeepSeek-V4's hyper-connections) - quantized = {torch.float8_e4m3fn, torch.float8_e8m0fnu, torch.int8, torch.uint8} - self.assertEqual([name for name, param in model.named_parameters() if param.dtype in quantized], []) - error, _ = self._reference_error(model, repo) - self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") - - @slow - @require_kernels - @require_torch_accelerator - def test_save_writes_the_checkpoint_layout_back(self): - """A pre-quantized model saved again lands on its checkpoint's keys, dtypes and bytes.""" - from huggingface_hub import snapshot_download - - for repo in self.spec.tiny_checkpoints: - with self.subTest(repo=repo): - 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(repo=repo, key=key): - source_dtype, source_shape, source_raw = source[key] - self.assertEqual(shape, source_shape) - if dtype != source_dtype: - # Loading honors initialized dtypes, including FP32 norms and BF16 - # MiniMax routers/patch embeddings. Float32 UE8M0 containers save natively. - self.assertTrue( - ( - dtype == torch.float32 - and key.endswith("norm.weight") - and source_dtype == torch.bfloat16 - ) - or ( - dtype == torch.bfloat16 - and source_dtype == torch.float32 - and key.endswith( - ("gate.weight", "e_score_correction_bias", "patch_embedding.weight") - ) - ) - or ( - dtype == torch.float8_e8m0fnu - and source_dtype == torch.float32 - and key.endswith((".scale", "_scale_inv")) - ) - ) - expected = torch.frombuffer(bytearray(source_raw), dtype=source_dtype).to(dtype) - self.assertEqual(raw, expected.view(torch.uint8).numpy().tobytes()) - else: - self.assertEqual(raw, source_raw) - - -def _check_static_cache_decode(model, ids): - from torch._dynamo.backends.registry import lookup_backend - - generation_kwargs = { - "input_ids": ids.to(model.device), - "attention_mask": torch.ones_like(ids, device=model.device), - "max_new_tokens": 3, - "do_sample": False, - "eos_token_id": None, - "pad_token_id": 0, - "cache_implementation": "static", - "output_logits": True, - "return_dict_in_generate": True, - } - # Match the generation compile tests' norm stabilization for tiny random models. - set_model_for_less_flaky_test(model) - eager = model.generate(**generation_kwargs, disable_compile=True) - graphs = [] - - def compile_backend(graph, inputs, **kwargs): - graphs.append(graph) - return lookup_backend("inductor")(graph, inputs, **kwargs) - - compiled = model.generate( - **generation_kwargs, - compile_config=CompileConfig(backend=compile_backend, fullgraph=True, dynamic=False, mode="default"), - ) - assert graphs, "Static-cache generation did not compile the decode forward" - torch.testing.assert_close(compiled.sequences, eager.sequences) - for actual, expected in zip(compiled.logits, eager.logits): - torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) - - -def _check_sharded_load_paths(rank, model_dir, model_cls, modes, out_dir): - from transformers.distributed import DistributedConfig - - world = torch.distributed.get_world_size() - ids = torch.arange(16, dtype=torch.long).unsqueeze(0) - for mode in modes: - model = model_cls.from_pretrained( - model_dir, - dtype="auto", - attn_implementation="eager", - distributed_config=DistributedConfig(tp_size=world, ep_size=world if mode == "ep" else None), - ).eval() - experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) - weight = getattr(experts, experts.gate_up_name) - local = weight.to_local() if hasattr(weight, "to_local") else weight - compiled_decode = False - with torch.no_grad(): - logits = model(ids.to(model.device)).logits.float().cpu() - if model._can_compile_fullgraph: - _check_static_cache_decode(model, ids) - compiled_decode = True - if rank == 0: - torch.save( - {"logits": logits, "expert_local": tuple(local.shape), "compiled_decode": compiled_decode}, - 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 fused (`experts.gate_up_proj`) or one per expert - (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`).""" - import glob - import json - - 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 - - -class LoadPathEquivalenceTesterMixin: - """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. - - 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 not run. - * every leg reports the local expert shard, so a leg that placed nothing fails loudly. - - A format's test file sets `spec`; the fixtures are its `distributed_models`.""" - - spec: FormatSpec - - @staticmethod - def _model_table(): - """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}`: each model in the format and - under the quantization config it ships with. 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, - ) - - 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, - }, - BlockFp8Config(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, - }, - BlockFp8Config(weight_block_size=None, activation_scheme="static"), - ), - # multimodal: the experts' plans live on `text_config`, not on the config the quantizer is handed - "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, - ), - }, - NVFP4.checkpoint_config(), - ), - # interleaved rows (`is_concatenated=False`), transposed, with expert biases its EP plan already names - "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, - }, - Mxfp4Config(activation_format="bf16"), - ), - # mixed: `expert_dtype` makes the experts MXFP4 over block-FP8 with UE8M0 scales - "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", - }, - BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), - ), - # a second multimodal nesting; the MoE dims are raised so a 2-way split stays 128-aligned - "minimax_m3_vl-mxfp8": ( - MiniMaxM3VLConfig, - MiniMaxM3SparseForConditionalGeneration, - { - "text_config": { - "hidden_size": 256, - # the 2-way split leaves 256: a swizzled tile is 128 wide along the down's K - "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": { - # the 2-way split leaves 128, one swizzled K tile - "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, - }, - Mxfp8Config(), - ), - } - - @classmethod - def setUpClass(cls): - super().setUpClass() - cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) - cls.enterClassContext(_cheap_tuning()) - cls.enterClassContext( - mock.patch.dict( - os.environ, {"TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton"))} - ) - ) - - def tearDown(self): - super().tearDown() - cleanup(torch_device, gc_collect=True) - - 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, label, "bf16") - quant_dir = os.path.join(self._tmp, label, "quantized") - torch.manual_seed(0) - # saved in bf16, the dtype every model here ships, so `dtype="auto"` reloads bf16 - 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) - - # a reload saved again lands on the same bytes: every layout op's reverse restores the checkpoint - reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") - round_trip = os.path.join(self._tmp, 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. How far - that travels is a property of the model (a chain of MoE layers can amplify it a hundredfold), so the - same magnitude injected on one device gives each model the budget its own conditioning earns. - """ - model = model_cls.from_pretrained( - model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" - ).eval() - # every block a sharded reduction passes through: the rounding accumulates down the stack - 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 two ranks, through the real EP / TP load path.""" - out = os.path.join(self._tmp, "out") - os.makedirs(out, exist_ok=True) - _init_distributed(tp=2)(_check_sharded_load_paths)(model_dir, model_cls, modes, out) - return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} - - @slow - @require_kernels - @require_torch_multi_accelerator - def test_every_load_path_agrees(self): - for label in self.spec.distributed_models: - with self.subTest(model=label): - model_dir, config, model_cls = self._fixture(label) - # `max_memory` forces the baseline across both devices; `auto` alone fits it on one - 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 whose plans shard experts (a multimodal model keeps them 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): - if model_cls._can_compile_fullgraph: - self.assertTrue(payload["compiled_decode"]) - local = payload["expert_local"] - # against the checkpoint's shape, never another mode's shard - 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", - ) - # a wrong shard axis lands orders of magnitude above the model's rounding floor - 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: # rounding alone does not 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/quantization/blockfp8/test_blockfp8.py b/tests/quantization/blockfp8/test_blockfp8.py index 2628df56042b..fcad893d0b0c 100644 --- a/tests/quantization/blockfp8/test_blockfp8.py +++ b/tests/quantization/blockfp8/test_blockfp8.py @@ -53,11 +53,11 @@ _quantizer_for, ) from ...test_memory_cleanup_mixin import MemoryCleanupMixin +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin from ..block_scaled_tester import ( BLOCK_FP8, BlockScaledFormatTesterMixin, - LoadPathEquivalenceTesterMixin, - TinyCheckpointTesterMixin, no_accelerator, ) diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py index ece764b0dd6f..d6610591a5da 100644 --- a/tests/quantization/finegrained/test_finegrained.py +++ b/tests/quantization/finegrained/test_finegrained.py @@ -55,11 +55,11 @@ _quantizer_for, _targets, ) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin from ..block_scaled_tester import ( FINEGRAINED, FORMATS, - LoadPathEquivalenceTesterMixin, - TinyCheckpointTesterMixin, _cheap_tuning, ) diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index f395effee93b..51e94307958a 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -41,11 +41,11 @@ _per_expert, _quantizer_for, ) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin from ..block_scaled_tester import ( MXFP4, BlockScaledFormatTesterMixin, - LoadPathEquivalenceTesterMixin, - TinyCheckpointTesterMixin, _cheap_tuning, ) diff --git a/tests/quantization/mxfp8/test_mxfp8.py b/tests/quantization/mxfp8/test_mxfp8.py index c6fa0307a7e9..6b2d161a3221 100644 --- a/tests/quantization/mxfp8/test_mxfp8.py +++ b/tests/quantization/mxfp8/test_mxfp8.py @@ -28,11 +28,11 @@ _load_and_save, _per_expert, ) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin from ..block_scaled_tester import ( MXFP8, BlockScaledFormatTesterMixin, - LoadPathEquivalenceTesterMixin, - TinyCheckpointTesterMixin, ) diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index 181ee1987ff9..35f3b885c411 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -39,11 +39,11 @@ _quantizer_for, _targets, ) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin from ..block_scaled_tester import ( NVFP4, BlockScaledFormatTesterMixin, - LoadPathEquivalenceTesterMixin, - TinyCheckpointTesterMixin, ) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 51c9d51cfb05..c490b05395c5 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -51,6 +51,13 @@ class TestWeightGlobMatching(unittest.TestCase): + def test_many_to_many_requires_an_operation_that_supports_it(self): + with self.assertRaisesRegex(ValueError, "one to many"): + WeightConverter(["a", "b"], ["c", "d"], operations=[Concatenate(dim=0)]) + for operation in (ErnieFuseAndSplitTextVisionExperts(), ErnieFuseAndSplitTextVisionExperts().reverse_op): + with self.subTest(operation=operation): + WeightConverter(["a", "b"], ["c", "d"], operations=[operation]) + def setUp(self): self.weight_globs_digits = [ "model.layers.*.mlp.gate_up_proj.weight", From e1f3eae44b559b692ee37a9fbb6687a142dc66cf Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 06:05:54 -0700 Subject: [PATCH 80/91] Remove scale container conversions and only quantize weights --- .../integrations/finegrained/conversions.py | 63 ++---------------- .../integrations/finegrained/layers.py | 3 - .../quantizers/block_scaled_utils.py | 22 +++---- .../finegrained/test_conversions.py | 65 +++++-------------- tests/integrations/finegrained/test_layers.py | 16 ++++- tests/quantization/blockfp8/test_blockfp8.py | 6 +- .../finegrained/test_finegrained.py | 6 +- tests/quantization/mxfp4/test_mxfp4.py | 21 +----- tests/quantization/mxfp8/test_mxfp8.py | 6 +- 9 files changed, 62 insertions(+), 146 deletions(-) diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index e9f454823087..4a0021996be2 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -103,8 +103,6 @@ def as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: if dtype == torch.uint8: # UE8M0 is a biased exponent, never a numeric float-to-integer cast. return (scale.float().log2().ceil() + 127).clamp(0, 255).masked_fill(torch.isnan(scale), 255).to(torch.uint8) - if scale.dtype == torch.uint8: - return (scale.float() - 127).exp2().masked_fill(scale == 255, float("nan")).to(dtype) return scale.to(dtype) @@ -114,45 +112,31 @@ class _FineGrainedOp(ConversionOps): 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. A layout op - built for one module carries it, since a save may name the tensor as the checkpoint does. + class — merge/split, quantize/dequantize — override `reverse_op` and ignore the flag. """ - def __init__(self, hf_quantizer=None, inverse: bool = False, module=None): + def __init__(self, hf_quantizer=None, inverse: bool = False): self.hf_quantizer = hf_quantizer self.inverse = inverse - self.module = module @property def reverse_op(self) -> ConversionOps: - return type(self)(self.hf_quantizer, inverse=not self.inverse, module=self.module) + return type(self)(self.hf_quantizer, inverse=not self.inverse) class FineGrainedDeserialize(_FineGrainedOp): """Apply a module's reversible layouts after its checkpoint tensors have been merged. - On-the-fly quantization applies these layouts to its outputs instead, after quantizing the - full-precision rows. Saving reverses the layouts before splitting or renaming the tensors. + On-the-fly quantization applies weight and scale layouts to its outputs instead, after + quantizing the full-precision rows. Biases only need their row layout. Saving reverses the layouts before splitting or renaming the tensors. """ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): if not self.inverse: input_dict = keyed_by_target(input_dict, target_patterns) - if not self.hf_quantizer.pre_quantized: + if not self.hf_quantizer.pre_quantized and not (full_layer_name or "").endswith("_bias"): return input_dict module_name, module = finegrained_module(model, full_layer_name) - if module is None and self.inverse: - # Saving names the first checkpoint output (e.g. gate_proj), while the - # model holds the fused input (gate_up_proj). Recover that module before - # reversing its layouts, which runs before the projection is split. - for target in target_patterns or []: - for source in kwargs.get("source_patterns", []): - candidate = (full_layer_name or "").replace(target, source.replace(r"\.", "."), 1) - module_name, module = finegrained_module(model, candidate) - if module is not None: - break - if module is not None: - break if module is None: return input_dict result = {} @@ -196,41 +180,6 @@ def convert(self, input_dict, model=None, full_layer_name=None, target_patterns= return out -class FineGrainedScaleContainer(_FineGrainedOp): - """A UE8M0 block scale into the native dtype or exponent bytes its module holds: it ships either as the - exponent byte under ``uint8`` (reinterpreted) or as its power-of-two value in float32 (cast - exactly). Exponent-byte checkpoints save back as uint8; float32 containers save as native UE8M0. - 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 = {} - for key, value in keyed_by_target(input_dict, target_patterns).items(): - value = single_tensor(value) - if self.inverse: - # the container its module recorded at load; e8m0 is only ever a scale, so the dtype - # alone picks the tensors to restore - module = self.module - if module is None: - _, module = finegrained_module(model, full_layer_name) - container = getattr(module, "scale_container_dtype", None) - if container is not None and value.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)): - value = as_container(value, container) - else: - module, held = held_scale(model, full_layer_name, key) - # 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 in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) - and value.dtype != held.dtype - ): - if value.dtype == torch.uint8: - module.scale_container_dtype = torch.uint8 - value = as_container(value, held.dtype) - out[key] = value - return out - - 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 diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index bc60cf2887d7..e2b1223fa5e1 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -360,9 +360,6 @@ def finegrained_linear( class _FineGrainedModule: """What every finegrained module (dense linear or experts) carries besides its parameters.""" - # Preserve uint8 exponent containers on save; float32 containers normalize to native UE8M0. - scale_container_dtype: torch.dtype | None = None - # 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 _hf_quantized_needs_local_tp = True diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index b277d38de36b..f9487e7b71b4 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -97,12 +97,11 @@ def format_conversions(self, dequantize: bool) -> list: def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: """The ops a tensor loading into `param_name` runs, for the layout this format's modules hold it in: an expert's stacked gate|up rows interleaved (weight, scale grid and bias share the row axis), a scale - in the container dtype its module holds, then in the swizzled layout. Each op has a reverse, so a save + in the swizzled layout. Each op has a reverse, so a save restores the checkpoint layout.""" - from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear + from ..integrations.finegrained import FineGrainedExperts from ..integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, FineGrainedSwizzleScales, ) @@ -112,10 +111,8 @@ def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: ): ops = [FineGrainedInterleaveGateUp(self)] if name.startswith("gate_up_proj") else [] if not name.endswith("_bias"): - ops += [FineGrainedScaleContainer(self, module=module), FineGrainedSwizzleScales(self)] + ops += [FineGrainedSwizzleScales(self)] return ops - if isinstance(module, FineGrainedLinear) and name == "weight_scale_inv": - return [FineGrainedScaleContainer(self, module=module)] return [] def finish_modules(self, model): @@ -238,12 +235,11 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** 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 + if self.pre_quantized: + return False + if isinstance(module, FineGrainedExperts): + return tensor_name in FUSED_STACKS + return isinstance(module, FineGrainedLinear) and tensor_name == "weight" @property def quantized_element_size(self) -> float: @@ -254,7 +250,7 @@ def get_new_dtype(self, converter, param, dtype): if any(name.endswith(("activation_scale", "_global_scale")) for name in converter.target_patterns): return None # The initialized scale dtype is authoritative for UE8M0, including float32 containers - # (DeepSeek-V4-Base). Exponent bytes still need a bitcast in the deserializer. + # (DeepSeek-V4-Base). Integer exponent containers retain their checkpoint dtype. is_float = ( param.get_dtype().startswith(("F", "BF")) if hasattr(param, "get_dtype") else param.is_floating_point() ) diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index a59b10008762..93cda729bac0 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -188,7 +188,7 @@ def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): self.assertIs(out["mlp.experts.gate_up_proj"], grid) self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) - def test_fused_dense_scales_save_back_as_exponent_bytes(self): + def test_fused_dense_scales_save_without_changing_the_runtime_dtype(self): from transformers.integrations.finegrained.conversions import FineGrainedDeserialize config = Mxfp8Config() @@ -196,7 +196,6 @@ def test_fused_dense_scales_save_back_as_exponent_bytes(self): model.config = PreTrainedConfig() model.mlp = torch.nn.Module() model.mlp.gate_up_proj = fg.FineGrainedLinear(64, 64, quantization_config=config) - model.mlp.gate_up_proj.scale_container_dtype = torch.uint8 scales = torch.arange(120, 128, dtype=torch.uint8).repeat(16, 1).reshape(64, 2) model._weight_conversions = [ WeightConverter( @@ -210,45 +209,8 @@ def test_fused_dense_scales_save_back_as_exponent_bytes(self): ) for name, expected in zip(("gate_proj", "up_proj"), scales.chunk(2)): actual = saved[f"mlp.{name}.weight_scale_inv"] - self.assertEqual(actual.dtype, torch.uint8) - self.assertTrue(torch.equal(actual, expected)) - - def test_scale_container_op_keeps_exponent_bytes_and_normalizes_float32(self): - """dsv4-flash-base ships UE8M0 scales as float32 values, MiniMax as uint8 exponent bytes: - both load into held e8m0. Exponent bytes save back as uint8, while float32 containers - save as native UE8M0; 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, torch.uint8 if container == torch.uint8 else None) - back = op.reverse_op.convert( - {"x": held}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" - )["x"] - self.assertEqual(back.dtype, torch.uint8 if container == torch.uint8 else torch.float8_e8m0fnu) - self.assertTrue(torch.equal(back.view(torch.uint8), native.view(torch.uint8))) - 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 + self.assertEqual(actual.dtype, torch.float8_e8m0fnu) + self.assertTrue(torch.equal(actual.view(torch.uint8), expected)) def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales @@ -267,23 +229,22 @@ def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): def test_expert_tensors_take_their_layout_from_their_module(self): """The layout an expert tensor lands in comes from the module it fills, whatever converter - brought it: gate|up interleaved (weight, scale grid and bias), then the scale container and + brought it: gate|up interleaved (weight, scale grid and bias), then the scale swizzle; a bias only interleaves.""" from transformers.integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, FineGrainedSwizzleScales, ) model, experts = self._experts("mxfp4") experts.quantization_config = Mxfp4Config(activation_format=None) # what the replacement walk records quantizer = _quantizer_for(Mxfp4Config()) - layout = [FineGrainedInterleaveGateUp, FineGrainedScaleContainer, FineGrainedSwizzleScales] + layout = [FineGrainedInterleaveGateUp, FineGrainedSwizzleScales] for param, expected in ( ("experts.gate_up_proj", layout), ("experts.gate_up_proj_scale_inv", layout), ("experts.gate_up_proj_bias", [FineGrainedInterleaveGateUp]), - ("experts.down_proj_scale_inv", [FineGrainedScaleContainer, FineGrainedSwizzleScales]), + ("experts.down_proj_scale_inv", [FineGrainedSwizzleScales]), ("experts.down_proj_bias", []), ): ops = quantizer.get_param_ops(model, param) @@ -349,6 +310,18 @@ def test_replacement_keeps_the_models_gate_up_convention(self): self.assertIsInstance(experts, fg.FineGrainedExperts) self.assertIs(experts.is_concatenated, expected) + def test_on_the_fly_quantization_still_interleaves_expert_biases(self): + from transformers.integrations.finegrained.conversions import FineGrainedDeserialize + + model, experts = self._experts("fp8", sm100=False) + quantizer = _quantizer_for(experts.quantization_config) + quantizer.pre_quantized = False + bias = torch.arange(4 * 256, dtype=torch.float32).reshape(4, 256) + key = "experts.gate_up_proj_bias" + actual = FineGrainedDeserialize(quantizer).convert({key: bias}, model=model, full_layer_name=key)[key] + gate, up = bias.chunk(2, dim=1) + torch.testing.assert_close(actual, torch.stack((gate, up), dim=-1).flatten(1)) + def test_interleave_op_follows_the_layer_it_fills(self): """The op reads the experts module a key fills: quantized experts get the interleaved order and the reverse restores the stack, while a model whose rows already come interleaved (`is_concatenated=False`, @@ -405,13 +378,11 @@ def test_ops_refuse_a_list_of_several_tensors(self): FineGrainedInputScalesSplit, FineGrainedPackedBlocks, FineGrainedQuantize, - FineGrainedScaleContainer, FineGrainedViewPackedInt8, ) tensor = torch.zeros(2, 4, 32, 16, dtype=torch.uint8) for op in ( - FineGrainedScaleContainer(None), FineGrainedPackedBlocks(None), FineGrainedViewPackedInt8(None), FineGrainedInputScalesSplit(None), diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py index 0b12d8c4348b..6694457b4f59 100644 --- a/tests/integrations/finegrained/test_layers.py +++ b/tests/integrations/finegrained/test_layers.py @@ -82,12 +82,26 @@ def test_ue8m0_falls_back_to_exponent_bytes(self): scales = torch.tensor([0.5, 1.0, 2.0, 4.0]) encoded = as_container(scales, torch.uint8) torch.testing.assert_close(encoded, torch.tensor([126, 127, 128, 129], dtype=torch.uint8)) - torch.testing.assert_close(as_container(encoded, torch.float32), scales) _, scale = FineGrainedQuantize._quantize_block_fp8(torch.full((2, 2), 448.0), (2, 2), True) torch.testing.assert_close(scale, torch.tensor([[127]], dtype=torch.uint8)) finally: fg._get_ue8m0_dtype.cache_clear() + def test_only_weights_need_on_the_fly_quantization(self): + config = BlockFp8Config() + quantizer = _quantizer_for(config) + quantizer.pre_quantized = False + model = torch.nn.Module() + model.linear = FineGrainedLinear(64, 64, has_bias=True, quantization_config=config) + model.experts = FineGrainedExperts(_Cfg(), quantization_config=config, has_bias=True) + for module_name, module in model.named_children(): + weights = {"weight"} if module_name == "linear" else {"gate_up_proj", "down_proj"} + for name, _ in module.named_parameters(recurse=False): + with self.subTest(parameter=f"{module_name}.{name}"): + self.assertEqual( + quantizer.param_needs_quantization(model, f"{module_name}.{name}"), name in weights + ) + def test_replacement_respects_a_disabled_post_expert_norm(self): from transformers.integrations.moe import use_experts_implementation diff --git a/tests/quantization/blockfp8/test_blockfp8.py b/tests/quantization/blockfp8/test_blockfp8.py index fcad893d0b0c..1254003114c0 100644 --- a/tests/quantization/blockfp8/test_blockfp8.py +++ b/tests/quantization/blockfp8/test_blockfp8.py @@ -299,7 +299,11 @@ def test_mistral4_static_per_tensor(self): experts.gate_up_proj_activation_scale.float(), checkpoint["model.layers.0.mlp.experts.gate_up_proj_activation_scale"].float(), ) - self.assertSavesBack(saved, checkpoint) + # The loader uses the initialized FP32 dtype for a dense block scale. + expected = dict(checkpoint) + key = "model.layers.0.self_attn.q_proj.weight_scale_inv" + expected[key] = expected[key].float() + self.assertSavesBack(saved, expected) 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 diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py index d6610591a5da..a87221cf4064 100644 --- a/tests/quantization/finegrained/test_finegrained.py +++ b/tests/quantization/finegrained/test_finegrained.py @@ -511,11 +511,11 @@ def test_the_fused_nvfp4_layout_reaches_the_dequantize_converter(self): def test_experts_under_any_parent_and_ungated_stacks_get_their_converters(self): """Nemotron-H holds its experts under `mixer.experts`, ungated (`up_proj`): the per-expert scales reach that - stack, and its scales take the container and swizzle ops without the gate|up interleave.""" + stack, and its scales take the swizzle op without the gate|up interleave.""" from transformers.core_model_loading import rename_source_key from transformers.integrations.finegrained.conversions import ( FineGrainedInterleaveGateUp, - FineGrainedScaleContainer, + FineGrainedSwizzleScales, ) from transformers.integrations.moe import use_experts_implementation @@ -557,7 +557,7 @@ def _init_weights(self, module): self.assertEqual(renamed, f"layers.0.mixer.experts.up_proj_{slot}") ops = quantizer.get_param_ops(model, "layers.0.mixer.experts.up_proj_scale_inv") - self.assertTrue(any(isinstance(op, FineGrainedScaleContainer) for op in ops)) + self.assertTrue(any(isinstance(op, FineGrainedSwizzleScales) for op in ops)) self.assertFalse(any(isinstance(op, FineGrainedInterleaveGateUp) for op in ops)) def test_an_embedding_tables_own_scale_keeps_its_name(self): diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index 51e94307958a..556ca11ac2f0 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -20,7 +20,6 @@ import torch -import transformers.integrations.finegrained.layers as fg from transformers import AutoTokenizer, GptOssConfig, GptOssForCausalLM, Mxfp4Config from transformers.testing_utils import ( TestCasePlus, @@ -59,7 +58,7 @@ def setUp(self): def test_generate_matches_forward_logits(self): """`generate` switches the experts to the batched kernels, while a plain forward runs the grouped ones: on a tiny GPT-OSS with biased experts, quantized on the fly (W4A16), every decoding step's logits must match - a forward over the same prefix. Large biases, so a kernel that drops one cannot pass.""" + a forward over the same prefix. Large biases make differences between the two kernels' bias handling visible.""" torch.manual_seed(0) config = GptOssConfig( vocab_size=512, @@ -265,7 +264,6 @@ def test_packed_experts(self): 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) @@ -280,21 +278,8 @@ def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): 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) + # Exponent bytes remain uint8; only the packed weights need a layout conversion. + scale_inv = scales ref = self._reference_dequant(blocks, scales) # rows pass through interleaved # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) diff --git a/tests/quantization/mxfp8/test_mxfp8.py b/tests/quantization/mxfp8/test_mxfp8.py index 6b2d161a3221..b061fda2f63e 100644 --- a/tests/quantization/mxfp8/test_mxfp8.py +++ b/tests/quantization/mxfp8/test_mxfp8.py @@ -41,8 +41,8 @@ class Mxfp8ChainTest(_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.""" + """MiniMax-M3-MXFP8 ships its E8M0 scales as raw `uint8` bytes: their bytes reach the kernels + unchanged, and saving preserves the runtime representation.""" torch.manual_seed(0) checkpoint = _per_expert( lambda n, k: { @@ -53,7 +53,7 @@ def test_minimax_m3_e8m0_scale_bytes(self): model, info, saved = _load_and_save(checkpoint, Mxfp8Config()) self.assertLoadsCleanly(info) experts = model.model.layers[0].mlp.experts - self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.uint8) self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) self.assertSavesBack(saved, checkpoint) From 6624d7ad805a54fae6a1ac2508080a925833f3b6 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 06:58:45 -0700 Subject: [PATCH 81/91] Fix finegrained loading regressions and map ModelOpt mixed precision --- docs/source/en/quantization/blockfp8.md | 1 + docs/source/en/quantization/finegrained.md | 31 +- docs/source/en/quantization/nvfp4.md | 4 +- .../integrations/finegrained/conversions.py | 89 +++-- .../integrations/finegrained/layers.py | 24 +- .../quantizers/block_scaled_utils.py | 58 ++- .../quantizers/quantizer_blockfp8.py | 2 +- .../quantizers/quantizer_finegrained.py | 47 ++- .../quantizers/quantizer_mxfp4.py | 20 +- .../quantizers/quantizer_mxfp8.py | 9 +- .../quantizers/quantizer_nvfp4.py | 3 +- src/transformers/utils/quantization_config.py | 148 ++++++-- tests/integrations/finegrained/test_layers.py | 53 +++ .../finegrained/test_loading_regressions.py | 334 ++++++++++++++++++ tests/quantization/block_scaled_tester.py | 6 +- .../finegrained/test_config_readers.py | 239 +++++++++++++ tests/quantization/mxfp4/test_mxfp4.py | 27 ++ tests/quantization/nvfp4/test_nvfp4.py | 109 ------ 18 files changed, 1012 insertions(+), 192 deletions(-) create mode 100644 tests/integrations/finegrained/test_loading_regressions.py create mode 100644 tests/quantization/finegrained/test_config_readers.py diff --git a/docs/source/en/quantization/blockfp8.md b/docs/source/en/quantization/blockfp8.md index 68b081208f82..d6b98b20e84d 100644 --- a/docs/source/en/quantization/blockfp8.md +++ b/docs/source/en/quantization/blockfp8.md @@ -87,6 +87,7 @@ model = AutoModelForCausalLM.from_pretrained("llama-3.1-8b-fp8", device_map="aut | `activation_format` | Scheme | Activations | |---|---|---| | `None` (default) or `"fp8"` | W8A8 | FP8, per token and block at run time, or against a calibrated per-tensor scale with `activation_scheme="static"` | +| `"bf16"` | W8A16 | Full precision; block-FP8 weights are decoded inside the matmul | ## Options diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 04bd9e79e5e1..0e76cd1bdefc 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -61,8 +61,35 @@ These load as a [`FineGrainedConfig`] without one being passed: declares `expert_dtype` in its quantization config instead, and reads the same way). See the [DeepSeek-V4](../model_doc/deepseek_v4) page for the plan written out. - A ModelOpt export labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats - reads as a plan with one entry per group target. Other `quant_algo` values, such as `"MIXED_PRECISION"`, raise an - error. + reads as a plan with one entry per group target. +- ModelOpt `quant_algo="MIXED_PRECISION"` reads the `quantized_layers` mapping. Each concrete module name covers + that module and its descendants; modules absent from the mapping stay in full precision. Supported algorithms + are `FP8` (calibrated per-tensor FP8), `FP8_PB_WO` (block-FP8 weights with full-precision activations), `NVFP4`, `MXFP8`, `MXFP4`, and the + `W4A16_NVFP4` / `W4A16_MXFP4` weight-only variants. Fields can appear directly in the quantization config or + inside its `quantization` object. Unsupported algorithms or config groups raise an error. + +For example, ModelOpt's per-module assignment can combine weight-only block-FP8 attention and NVFP4 experts: + +```json +{ + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "layers.8.attention.wq": {"quant_algo": "FP8_PB_WO"}, + "layers.9.experts": {"quant_algo": "NVFP4", "group_size": 16} + } +} +``` + +The model must support the checkpoint's architecture and parameter names. Format parsing alone does not add +support for a new architecture. Module names are resolved through the model's checkpoint key renames. An +assignment that matches no supported linear or expert module raises an error, including per-expert assignments +that cannot be represented by an architecture's fused expert module. + +AWQ smoothing (`pre_quant_scale`), SVD residuals, asymmetric zero points, sparse or transformed layouts, and +compressed scale encodings such as `weight_scale_encoding="csf"` are rejected. NVFP4 embedding tables are not +supported by the FP8 embedding integration. Weight-only block-FP8 currently runs on dense linears and eager +experts; fused expert backends require kernel support for that activation mode. To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass `FineGrainedConfig(dequantize=True)`. diff --git a/docs/source/en/quantization/nvfp4.md b/docs/source/en/quantization/nvfp4.md index 7fd709255916..88373397e25e 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -37,8 +37,8 @@ The NVFP4 exports of [NVIDIA ModelOpt](https://github.com/NVIDIA/TensorRT-Model- `quant_algo="NVFP4"`) and of [compressed-tensors](./compressed_tensors) (`format="nvfp4-pack-quantized"`) load as [`NVFP4Config`], with their calibrated activation scales, and save back in their own layout. A ModelOpt export labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats loads as a -[`FineGrainedConfig`] (see [Fine-grained](./finegrained)); other `quant_algo` values, such as `"MIXED_PRECISION"`, raise -an error. +[`FineGrainedConfig`]. ModelOpt `quant_algo="MIXED_PRECISION"` uses its explicit `quantized_layers` mapping; see +[Fine-grained](./finegrained) for supported formats and restrictions. ```py from transformers import AutoModelForCausalLM, AutoTokenizer diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index 4a0021996be2..b1fc0a52acd6 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -25,7 +25,7 @@ import torch from ...core_model_loading import ConversionOps, Interleave -from ...quantizers.quantizers_utils import FUSED_STACKS, get_module_from_name +from ...quantizers.quantizers_utils import FUSED_STACKS, get_module_from_name, plan_entry_for from .layers import ( _FP8_DTYPE, _FP8_MAX, @@ -226,6 +226,21 @@ def convert(self, input_dict, target_patterns=None, **kwargs): return out +class FineGrainedScalarScales(_FineGrainedOp): + """Give per-expert scalar scales a 1x1 grid before stacking and concatenating gate/up scales.""" + + def convert(self, input_dict, **kwargs): + def reshape(tensor): + if self.inverse: + return tensor.reshape(()) + return tensor.reshape(1, 1) + + return { + key: [reshape(tensor) for tensor in value] if isinstance(value, list) else reshape(value) + for key, value in input_dict.items() + } + + 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 @@ -580,11 +595,15 @@ def _scale_pattern_for(self, weight_pattern: str) -> str: scale = base + "_scale_inv" return scale + anchor - @property - def _reciprocal_globals(self) -> bool: - """compressed-tensors ships NVFP4's globals as their reciprocal (see `FineGrainedReciprocal`).""" - plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} - return any(config.checkpoint_format == "compressed-tensors" for config in plan.values()) + def _config_for(self, full_layer_name: str | None): + """Resolve this module's format independently of other entries in a mixed plan.""" + if self.hf_quantizer is None: + return None + config = self.hf_quantizer.quantization_config + if full_layer_name is not None: + module_name = full_layer_name.rpartition(".")[0] + config = plan_entry_for(self.hf_quantizer.plan(), module_name) + return config # the weights a converter hands one at a time: a dense `weight`, or an experts module's fused stack _SINGLE_WEIGHTS = ("weight", *FUSED_STACKS) @@ -625,6 +644,9 @@ def _dequantize_one( scales: torch.Tensor, output_dtype: torch.dtype | None = None, global_scale: torch.Tensor | None = None, + *, + reciprocal_globals: bool = False, + weight_block_size: tuple[int, int] | 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) @@ -634,23 +656,17 @@ def _dequantize_one( 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 - # a per-tensor scale has no grid: one block over the whole matrix + # A per-tensor scale has no grid: one block over the whole matrix. scale_rows, scale_cols = scales.shape[-2:] if scales.ndim >= 2 else (1, 1) - # 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 block the plan names reproduces it - plan = self.hf_quantizer.plan() if hasattr(self.hf_quantizer, "plan") else {} - configured = next( - ( - block - for block in {tuple(config.weight_block_size) for config in plan.values() if config.weight_block_size} - if (_cdiv(rows, block[0]), _cdiv(cols, block[1])) == (scale_rows, scale_cols) - ), - None, - ) - if configured: - block_m, block_n = configured + # Ceil-rounded grids do not uniquely identify a block size. Use this module's config, + # including in mixed plans; infer from the grid only when no configured block matches. + if weight_block_size is None: + weight_block_size = getattr(self._config_for(None), "weight_block_size", None) + if weight_block_size and tuple(_cdiv(size, block) for size, block in zip((rows, cols), weight_block_size)) == ( + scale_rows, + scale_cols, + ): + block_m, block_n = weight_block_size elif rows % scale_rows or cols % scale_cols: raise ValueError( f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." @@ -682,7 +698,7 @@ def _dequantize_one( # a stack's globals, one per expert (E,) or one per stacked gate|up half (E, 2), over its rows global_scale = global_scale.reshape(global_scale.shape[0], -1) global_scale = global_scale.repeat_interleave(rows // global_scale.shape[1], dim=1)[..., None] - out = out / global_scale if self._reciprocal_globals else out * global_scale + out = out / global_scale if reciprocal_globals else out * global_scale return out.to(output_dtype) def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: @@ -700,6 +716,9 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor] | torch.Tensor]: output_dtype = self._get_target_dtype(model, full_layer_name) + config = self._config_for(full_layer_name) + reciprocal_globals = getattr(config, "checkpoint_format", None) == "compressed-tensors" + weight_block_size = getattr(config, "weight_block_size", None) # The dense converter (``["weight$", companions...] -> weight``) and a fused stack's (``["gate_up_proj$", its # companions...] -> gate_up_proj``) each hand one weight, with or without a scale (RMSNorm weights match # ``weight$`` too). @@ -714,8 +733,17 @@ def convert( for companion in (self._scale_pattern_for, self._global_pattern_for) ) if scale is None: - return {target_key: quantized} - return {target_key: self._dequantize_one(quantized, scale, output_dtype, global_scale)} + return {target_key: quantized.to(output_dtype) if output_dtype is not None else quantized} + return { + target_key: self._dequantize_one( + quantized, + scale, + output_dtype, + global_scale, + reciprocal_globals=reciprocal_globals, + weight_block_size=weight_block_size, + ) + } # Generic chain path: dequantize every weight pattern that has a sibling scale. consumed = { @@ -727,6 +755,12 @@ 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: + if output_dtype is not None: + value = ( + [tensor.to(output_dtype) for tensor in value] + if isinstance(value, list) + else value.to(output_dtype) + ) result[key] = value continue weights = value if isinstance(value, list) else [value] @@ -740,7 +774,10 @@ def convert( global_scales = input_dict.get(self._global_pattern_for(key), [None] * len(weights)) global_scales = global_scales if isinstance(global_scales, list) else [global_scales] result[key] = [ - self._dequantize_one(w, s, output_dtype, g) for w, s, g in zip(weights, scales, global_scales) + self._dequantize_one( + w, s, output_dtype, g, reciprocal_globals=reciprocal_globals, weight_block_size=weight_block_size + ) + for w, s, g in zip(weights, scales, global_scales) ] return result diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index e2b1223fa5e1..a51d6b58e7f5 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -546,11 +546,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: w = self.weight.view(self.n_groups, -1, self.weight.shape[-1]) 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)) + if scale_inv.numel() == 1: + # A per-tensor checkpoint calibrates the whole stored matrix, so every + # group uses the same scale. The grouped kernel expects an expert axis. + scale_inv = scale_inv.reshape(1, 1, 1).expand(self.n_groups, 1, 1).contiguous() + else: + 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 + # Kernels pad the expert schedule to a power of two. Empty trailing groups + # point at the final row, including when n_groups itself is not a power of two. + padded_groups = 1 << (self.n_groups - 1).bit_length() + expert_start = ( + torch.arange(padded_groups + 1, device=x.device, dtype=torch.int32).clamp_max(self.n_groups) + * tokens_per_group + ) two_level = self.weight_global_scale is not None kernel = load_finegrained_kernel() @@ -791,14 +801,14 @@ def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torc input_global_scale=input_global, activation_format=self.activation_format, ) - # A scale spanning several interleaved rows belongs to one half (the loader interleaves the grid per - # block), so a row-major read would pair half the rows with the other half's scale: each half runs on its - # own rows and blocks, its output columns interleaved back. + # Separately calibrated gate/up grids interleave per block, rather than per weight row: + # run each half with its own blocks, then interleave the output columns. A single scale row + # is shared by both halves and can be read directly. if ( proj == self.gate_up_name and self.holds_interleaved_gate_up and scale.ndim == 2 - and scale.shape[0] < weight.shape[0] + and 1 < scale.shape[0] < weight.shape[0] ): halves = [ linear( diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index f9487e7b71b4..4ee0578e2b74 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -17,6 +17,7 @@ """ import re +from copy import deepcopy from typing import TYPE_CHECKING from ..utils import ( @@ -125,6 +126,12 @@ def _process_model_after_weight_loading(self, model, **kwargs): member.finish_modules(model) disable_deepgemm_on_multi_device(model) assert_modules_are_quantized(model) + if self.pre_quantized and self.quantization_config.dequantize: + # The loading plan also consumed scales. Saving full-precision weights must reverse only + # the original model transforms, preserving expert split/merge and user key mappings. + model._weight_conversions = [ + original for original, loaded in self._dequantized_save_conversions if loaded.was_used() + ] if not self.pre_quantized and getattr(model, "_weight_conversions", None) is not None: # Quantization emitted scales that were absent on disk. Give them the same # reverse conversions as their weights when saving. @@ -201,11 +208,18 @@ def validate_environment(self, *args, **kwargs): 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): + # W4A16 decodes packed FP4 weights before a regular dot product; it does not + # need the FP8 instructions used by the other activation/weight formats. + weight_only_mxfp4 = all( + config.quant_method == "mxfp4" and config.activation_format == "bf16" + for config in self.plan().values() + ) + minimum_capability = (7, 5) if weight_only_mxfp4 else (8, 9) + if compute_capability < minimum_capability: logger.warning_once( - "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" + "This fine-grained quantization plan requires a GPU with compute capability >= " + f"{minimum_capability[0]}.{minimum_capability[1]}, actual = `{major}.{minor}`. " + "We will default to dequantizing the model to bf16." ) self.quantization_config.dequantize = True return @@ -468,12 +482,44 @@ def get_weight_conversions(self): return conversions def update_weight_conversions(self, weight_conversions, model, *, for_save=False): - from ..core_model_loading import WeightConverter - from ..integrations.finegrained.conversions import FineGrainedDeserialize + from ..core_model_loading import MergeModulelist, WeightConverter + from ..integrations.finegrained import FineGrainedExperts + from ..integrations.finegrained.conversions import FineGrainedDeserialize, FineGrainedScalarScales + dequantize = self.pre_quantized and self.quantization_config.dequantize + save_conversions = deepcopy(weight_conversions) if dequantize else [] conversions = super().update_weight_conversions(weight_conversions, model) + if dequantize: + self._dequantized_save_conversions = [] + for original in save_conversions: + loaded = next( + conversion + for conversion in conversions + if conversion._original_target_patterns == original._original_target_patterns + and all( + source in conversion._original_source_patterns for source in original._original_source_patterns + ) + ) + self._dequantized_save_conversions.append((original, loaded)) converters = [conversion for conversion in conversions if isinstance(conversion, WeightConverter)] + scalar_scales = [ + f"{name}.{proj}_scale_inv" + for name, module in model.named_modules() + if isinstance(module, FineGrainedExperts) + and module.quantization_config.checkpoint_format == "modelopt" + and module.quantization_config.method.value == "fp8" + and module.quantization_config.weight_block_size is None + for proj in (module.gate_up_name, "down_proj") + ] for converter in converters: + if ( + scalar_scales + and any(isinstance(op, MergeModulelist) for op in converter.operations) + and not any(isinstance(op, FineGrainedScalarScales) for op in converter.operations) + ): + reverse_converter = converter.reverse_transform() + if any(reverse_converter.rename_source_key(name)[1] is not None for name in scalar_scales): + converter.operations = [FineGrainedScalarScales(self), *converter.operations] if not any(isinstance(op, FineGrainedDeserialize) for op in converter.operations): converter.operations = [*converter.operations, FineGrainedDeserialize(self)] reverse = [converter.reverse_transform() for converter in converters] if for_save else [] diff --git a/src/transformers/quantizers/quantizer_blockfp8.py b/src/transformers/quantizers/quantizer_blockfp8.py index d13746c40caa..adf70bea37a4 100644 --- a/src/transformers/quantizers/quantizer_blockfp8.py +++ b/src/transformers/quantizers/quantizer_blockfp8.py @@ -35,7 +35,7 @@ def plan(self): def format_conversions(self, dequantize): conversions = [] - if self.quantization_config.checkpoint_format == "compressed-tensors": + if self.quantization_config.checkpoint_format in ("compressed-tensors", "modelopt"): # its block scale ships as `weight_scale`, and rides its weight's converter once renamed conversions.append(self._block_scale_rename()) if not dequantize and self.quantization_config.activation_scheme == "static": diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py index 4b7d3dfee206..adc6f66ae7be 100644 --- a/src/transformers/quantizers/quantizer_finegrained.py +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -14,11 +14,12 @@ """The fine-grained quantizer: a model whose modules are in different block-scaled formats, each part loaded by its own format's quantizer.""" +import re from typing import TYPE_CHECKING from ..utils.quantization_config import BlockFp8Config from .block_scaled_utils import BlockScaledQuantizer -from .quantizers_utils import embedding_pattern, get_module_from_name +from .quantizers_utils import embedding_pattern, get_module_from_name, module_kind, plan_entry_for if TYPE_CHECKING: @@ -40,6 +41,7 @@ def __init__(self, quantization_config: "FineGrainedConfig", **kwargs): # the embedding tables `modules_to_convert` names hold FP8 rows with one per-tensor scale self._embedding_config = BlockFp8Config(weight_block_size=None) + self._resolved_plan = None configs = {id(config): config for config in self.plan().values()}.values() self._members = [AUTO_QUANTIZER_MAPPING[config.method.value](config, **kwargs) for config in configs] # the plan's embedding tables are the host's to name @@ -49,9 +51,50 @@ def __init__(self, quantization_config: "FineGrainedConfig", **kwargs): def plan(self): return { **dict.fromkeys(map(embedding_pattern, self.embeddings), self._embedding_config), - **self.quantization_config.dtype_plan, + **(self.quantization_config.dtype_plan if self._resolved_plan is None else self._resolved_plan), } + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): + from ..conversion_mapping import get_model_conversion_mapping + from ..core_model_loading import WeightRenaming + + # Match checkpoint patterns against the source namespace without rewriting the saved config. + # A real leaf also exercises rename rules anchored to a parameter name. + reverse = [ + rule.reverse_transform() + for rule in get_model_conversion_mapping(model)[::-1] + if isinstance(rule, WeightRenaming) + ] + plan = self.quantization_config.dtype_plan + resolved = dict(plan) + matched = set() + for name, module in model.named_modules(): + if module_kind(name, module) not in ("linear", "grouped_linear", "experts"): + continue + checkpoint_key = name + ".weight" + for rule in reverse: + checkpoint_key, _ = rule.rename_source_key(checkpoint_key) + checkpoint_name = checkpoint_key.removesuffix(".weight") + native = plan_entry_for(plan, name, catch_all=False) + source = plan_entry_for(plan, checkpoint_name, catch_all=False) + if native is not None and source is not None and native is not source: + raise ValueError( + f"Conflicting quantization entries for {name!r} and checkpoint module {checkpoint_name!r}" + ) + if source is not None and native is None: + resolved[rf"^{re.escape(name)}$"] = source + matched.update( + pattern for pattern in plan if re.search(pattern, name) or re.search(pattern, checkpoint_name) + ) + for pattern, config in plan.items(): + if config.checkpoint_format == "modelopt" and pattern not in matched: + raise ValueError( + f"ModelOpt target {pattern!r} matches no supported linear or expert module. " + "Per-expert assignments cannot target an architecture that stores fused expert tensors." + ) + self._resolved_plan = resolved + return super()._process_model_before_weight_loading(model, **kwargs) + def members(self): return self._members diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 923c7cad6670..68ea88ab9387 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import FUSED_STACKS, get_module_from_name if TYPE_CHECKING: @@ -33,7 +34,24 @@ class Mxfp4HfQuantizer(BlockScaledQuantizer): quantization_config: "Mxfp4Config" def format_conversions(self, dequantize): - return self._dequantize_conversions() if dequantize else self._mxfp4_conversions() + conversions = self._dequantize_conversions() if dequantize else self._mxfp4_conversions() + if self.quantization_config.checkpoint_format == "modelopt": + conversions = [self._block_scale_rename(), *conversions] + return conversions + + def get_param_ops(self, model, param_name): + from ..integrations.finegrained.conversions import FineGrainedViewPackedInt8 + from ..integrations.finegrained.layers import _FineGrainedModule + + ops = super().get_param_ops(model, param_name) + module, name = get_module_from_name(model, param_name) + if ( + self.quantization_config.checkpoint_format == "modelopt" + and isinstance(module, _FineGrainedModule) + and name in ("weight", *FUSED_STACKS) + ): + ops = [FineGrainedViewPackedInt8(self), *ops] + return ops def _dequantize_conversions(self): """The blocks regroup into the packed rows and the exponent-byte scales come back out of diff --git a/src/transformers/quantizers/quantizer_mxfp8.py b/src/transformers/quantizers/quantizer_mxfp8.py index a57bf3ff7adc..06808f4ac6fb 100644 --- a/src/transformers/quantizers/quantizer_mxfp8.py +++ b/src/transformers/quantizers/quantizer_mxfp8.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """MXFP8: `float8_e4m3fn` weights with a UE8M0 scale per 32. It ships the plain `weight` / `weight_scale_inv` -pair, so it needs nothing beyond what every fine-grained format runs.""" +pair; ModelOpt exports name that scale `weight_scale` instead.""" from typing import TYPE_CHECKING @@ -27,3 +27,10 @@ class Mxfp8HfQuantizer(BlockScaledQuantizer): """MXFP8 weights, W8A8, W8A4 or weight-only.""" quantization_config: "Mxfp8Config" + + def format_conversions(self, dequantize): + return ( + [self._block_scale_rename()] + if self.quantization_config.checkpoint_format in ("modelopt", "compressed-tensors") + else [] + ) diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index d6f245a111d4..867231bf7dc3 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -61,10 +61,11 @@ def get_param_ops(self, model, param_name): non-float tensor in the checkpoint's dtype, modelopt's uint8, and the reverse saves it back as that.""" from ..integrations.finegrained import FineGrainedLinear from ..integrations.finegrained.conversions import FineGrainedViewPackedInt8 + from ..integrations.finegrained.layers import _FineGrainedModule ops = super().get_param_ops(model, param_name) module, name = get_module_from_name(model, param_name) - if name in ("weight", *FUSED_STACKS): + if isinstance(module, _FineGrainedModule) and name in ("weight", *FUSED_STACKS): ops = [FineGrainedViewPackedInt8(self), *ops] # an experts module's globals are inverted in their converters, before the per-expert reductions elif isinstance(module, FineGrainedLinear) and name in ("weight_global_scale", "activation_scale"): diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index ffff4be62680..0e5586769843 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -15,10 +15,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy +import fnmatch import importlib.metadata import json import os -import re import warnings from dataclasses import dataclass from enum import Enum @@ -1699,8 +1699,9 @@ def subtree_patterns(globs: list[str]) -> list[str]: """ patterns = [] for glob in globs: - subtree = glob[:-2] if glob.endswith(".*") else glob.rstrip("*").rstrip(".") - patterns.append(r"(?:^|.*\.)" + re.escape(subtree) + r"(\..*)?$") + subtree = glob.removesuffix(".*") + pattern = fnmatch.translate(subtree).removesuffix(r"\Z").removesuffix(r"\z") + patterns.append(r"(?:^|.*\.)" + pattern + r"(\..*)?$") return patterns @@ -1812,8 +1813,10 @@ def post_init(self): # one spelling for the weights' own activation format if self.activation_format == self.method.value: self.activation_format = None - if self.checkpoint_format not in (None, "compressed-tensors"): - raise ValueError(f"checkpoint_format must be None or 'compressed-tensors'; got {self.checkpoint_format!r}") + if self.checkpoint_format not in (None, "compressed-tensors", "modelopt"): + raise ValueError( + f"checkpoint_format must be None, 'compressed-tensors', or 'modelopt'; got {self.checkpoint_format!r}" + ) def get_loading_attributes(self): named = {"activation_format": self.activation_format} if self._names_activation_format else {} @@ -1845,7 +1848,7 @@ class BlockFp8Config(BlockScaledConfig): """ method = QuantizationMethod.FP8 - activation_formats = (None, "fp8") + activation_formats = (None, "fp8", "bf16") def __init__( self, @@ -1867,6 +1870,10 @@ def post_init(self): self.activation_scheme = self.activation_scheme.lower() if self.activation_scheme not in ("dynamic", "static"): raise ValueError(f"activation_scheme must be 'dynamic' or 'static'; got {self.activation_scheme!r}") + if self.activation_format == "bf16" and self.activation_scheme == "static": + raise ValueError( + "Weight-only block-FP8 has no activation quantization; activation_scheme must be 'dynamic'" + ) 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.weight_block_size is not None and ( @@ -1940,7 +1947,7 @@ def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs) def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: """One `config_groups` entry as `dtype_plan` entries, one per target, 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. + format we do not serve. ModelOpt rejects such groups; compressed-tensors keeps its own backend. 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 @@ -1948,6 +1955,10 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: i.e. every linear), and an `re:` prefix introduces a pattern. Activations are read the same way, and only where the weights' kernels quantize to them. """ + try: + _check_block_scaled_features(spec) + except ValueError: + return None formats = { (4, "float", 16): "nvfp4", (4, "float", 32): "mxfp4", @@ -1992,42 +2003,117 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: return dict.fromkeys(targets or [".*"], config) +def _check_block_scaled_features(spec: dict): + """Refuse runtime operations that a block-scaled GEMM alone cannot reproduce.""" + for field in ( + "pre_quant_scale", + "lora_rank", + "has_zero_point", + "transform_config", + "sparsity_config", + "weight_scale_encoding", + ): + if spec.get(field): + raise ValueError(f"Unsupported block-scaled quantization feature {field}: {spec[field]!r}") + for name in ("weights", "input_activations"): + if options := spec.get(name): + if options.get("symmetric") is False or options.get("zero_point") is not None: + raise ValueError(f"Unsupported asymmetric quantization of {name}") + + +def _modelopt_format(spec: dict) -> BlockScaledConfig: + """Translate an algorithm without inferring a different format from unsupported metadata.""" + if not isinstance(spec, dict): + raise ValueError(f"Expected a modelopt format object; got {spec!r}") + _check_block_scaled_features(spec) + algo = spec.get("quant_algo") + formats = { + "NVFP4": (NVFP4Config, 16, None), + "W4A16_NVFP4": (NVFP4Config, 16, "bf16"), + "MXFP8": (Mxfp8Config, 32, None), + "MXFP4": (Mxfp4Config, 32, None), + "W4A16_MXFP4": (Mxfp4Config, 32, "bf16"), + } + if algo == "FP8": + return BlockFp8Config(weight_block_size=None, activation_scheme="static", checkpoint_format="modelopt") + if algo == "FP8_PB_WO": + block = spec.get("group_size", 128) + if not isinstance(block, int) or isinstance(block, bool) or block <= 0: + raise ValueError(f"FP8_PB_WO requires a positive group_size; got {block!r}") + return BlockFp8Config(weight_block_size=(block, block), activation_format="bf16", checkpoint_format="modelopt") + if algo not in formats: + raise ValueError(f"Unsupported modelopt quant_algo: {algo!r}") + config_class, group_size, activation_format = formats[algo] + if spec.get("group_size", group_size) != group_size: + raise ValueError(f"modelopt {algo} requires group_size={group_size}; got {spec['group_size']!r}") + return config_class(activation_format=activation_format, checkpoint_format="modelopt") + + def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: - """The config an NVIDIA modelopt export describes: its one format's, or a [`FineGrainedConfig`] when its - `config_groups` cover different modules with different formats. + """Read a ModelOpt format or its explicit per-module `MIXED_PRECISION` assignment. - `quant_algo` names what it exported; modelopt also covers FP8, INT4 AWQ, W4A8 and MXFP4, which are refused - rather than silently mis-read. The calibrated `input_scale` TENSORS are the loader's business, not the config's. + ModelOpt's `quant_cfg.json` nests these fields under `quantization`; HF configs can store them directly. + Unsupported groups and algorithms raise rather than loading their tensors as blanket NVFP4. """ - config_dict = dict(config_dict) - quant_algo = config_dict.pop("quant_algo", None) - if quant_algo != "NVFP4": - raise ValueError(f"modelopt checkpoints are supported for quant_algo in ['NVFP4']; got {quant_algo!r}") - logger.info("modelopt checkpoint exported with quant_algo='NVFP4'.") - if config_dict.pop("kv_cache_scheme", None): + config_dict = {**config_dict, **config_dict.get("quantization", {})} + _check_block_scaled_features(config_dict) + quant_algo = config_dict.get("quant_algo") + if config_dict.get("kv_cache_scheme") or config_dict.get("kv_cache_quant_algo"): 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." + "This modelopt checkpoint calibrates a quantized KV cache; the KV cache stays in the model's dtype, " + "so its scales are not loaded." ) - ignore = config_dict.pop("ignore", None) or config_dict.pop("exclude_modules", None) + ignore = config_dict.get("ignore") or config_dict.get("exclude_modules") modules_to_not_convert = subtree_patterns(ignore) if ignore is not None else None - dtype_plan = {".*": NVFP4Config()} - # the same thing described parametrically, taken only when every group maps, so a multi-group export is not - # flattened to one format - if groups := config_dict.pop("config_groups", None): - mapped = [plan_from_config_groups(spec) for spec in groups.values()] - if all(entries is not None for entries in mapped): - dtype_plan = { - pattern: config for entries in mapped if entries is not None for pattern, config in entries.items() - } - return _config_for_plan(dtype_plan, modules_to_not_convert) + if quant_algo == "MIXED_PRECISION": + layers = config_dict.get("quantized_layers") + if not isinstance(layers, dict) or not layers: + raise ValueError("modelopt MIXED_PRECISION requires a non-empty quantized_layers mapping") + dtype_plan = {} + for name, spec in layers.items(): + if not isinstance(name, str) or not name or "*" in name: + raise ValueError(f"modelopt quantized_layers requires concrete module names; got {name!r}") + try: + dtype_plan[subtree_patterns([name])[0]] = _modelopt_format(spec) + except ValueError as error: + raise ValueError(f"modelopt module {name!r}: {error}") from error + config = FineGrainedConfig(dtype_plan=dtype_plan, modules_to_not_convert=modules_to_not_convert) + return config + if groups := config_dict.get("config_groups"): + dtype_plan = {} + # Validate the declared algorithm too; groups cannot disguise an unsupported export. + _modelopt_format(config_dict) + for name, spec in groups.items(): + _check_block_scaled_features(spec) + spec = dict(spec) + targets = [] + for target in spec.get("targets") or ["Linear"]: + if target == "Linear" or target.startswith("re:"): + targets.append(target) + elif "." in target or target[:1].islower(): + targets.append("re:" + subtree_patterns([target])[0]) + else: + raise ValueError(f"Unsupported modelopt target {target!r} in config group {name!r}") + spec["targets"] = targets + entries = plan_from_config_groups(spec) + if entries is None: + raise ValueError(f"Unsupported modelopt config_groups entry {name!r}: {spec}") + for config in entries.values(): + config.checkpoint_format = "modelopt" + dtype_plan.update(entries) + return _config_for_plan(dtype_plan, modules_to_not_convert) + config = _modelopt_format(config_dict) + config.modules_to_not_convert = modules_to_not_convert + return config def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None: """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, per-channel, sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's `checkpoint_format` names the layout its tensors ship in, which a save writes back.""" - if config_dict.get("sparsity_config") or config_dict.get("transform_config"): + try: + _check_block_scaled_features(config_dict) + except ValueError: return None if config_dict.get("quantization_status", "compressed") != "compressed": return None diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py index 6694457b4f59..16c6ef6a3f4f 100644 --- a/tests/integrations/finegrained/test_layers.py +++ b/tests/integrations/finegrained/test_layers.py @@ -331,6 +331,27 @@ def test_eager_gated_block_scales_pair_each_half_with_its_own_blocks(self): ] torch.testing.assert_close(out, torch.stack(halves, dim=-1).flatten(-2)) + def test_eager_gated_per_tensor_scale_is_shared_by_gate_and_up(self): + experts = FineGrainedExperts(_Cfg(), quantization_config=BlockFp8Config(weight_block_size=None)) + experts._deepgemm_disabled = True + torch.manual_seed(0) + codes = torch.randn_like(experts.gate_up_proj, dtype=torch.float32).to(torch.float8_e4m3fn) + experts.gate_up_proj = torch.nn.Parameter(codes) + # On-the-fly quantization of a fused projection produces one scale per expert, + # whereas separately calibrated gate/up checkpoints can hold two. + scales = torch.tensor([0.125, 0.5, 2.0, 8.0]).view(-1, 1, 1) + experts.gate_up_proj_scale_inv = torch.nn.Parameter(scales) + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = _block_scaled_matmul_2d + x = torch.randn(3, _Cfg.hidden_size) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + for expert in range(_Cfg.num_local_experts): + with self.subTest(expert=expert): + actual = experts.expert_linear(x, "gate_up_proj", expert) + expected = x @ (codes[expert].float() * scales[expert]).T + torch.testing.assert_close(actual, expected) + def test_stacked_gate_up_is_gated_by_halves_and_kept_from_the_triton_chains(self): """The DeepGEMM backends load gate|up stacked: the host GLU splits the halves, and the triton chains, which read the rows interleaved, refuse such a module rather than pair the wrong rows.""" @@ -362,6 +383,38 @@ def test_swiglu_alpha_without_a_limit_runs_unclamped(self): gate, up = gate_up[..., 0::2], gate_up[..., 1::2] torch.testing.assert_close(experts._apply_gate(gate_up), (up + 1) * gate * torch.sigmoid(gate * 1.702)) + def test_grouped_linear_broadcasts_a_shared_per_tensor_scale(self): + from transformers.integrations.finegrained import FineGrainedGroupedLinear + + linear = FineGrainedGroupedLinear( + 8, 12, 3, quantization_config=BlockFp8Config(weight_block_size=None), has_bias=True + ) + torch.manual_seed(0) + linear.weight = torch.nn.Parameter(torch.randn(12, 8).to(torch.float8_e4m3fn)) + linear.weight_scale_inv.data.fill_(0.25) + linear.bias.data.copy_(torch.randn(12)) + kernel, _ = fake_finegrained_kernel() + + def grouped_matmul(a, b, activation_scale, scales, *, expert_start, **kwargs): + self.assertEqual(expert_start.tolist(), [0, 10, 20, 30, 30]) + return torch.cat( + [ + _block_scaled_matmul_2d( + a[expert_start[group] : expert_start[group + 1]], b[group], None, scales[group] + ) + for group in range(3) + ] + ) + + kernel.matmul_grouped = grouped_matmul + x = torch.randn(2, 5, 3, 8) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + actual = linear(x) + weight = linear.weight.float().view(3, 4, 8) * 0.25 + expected = torch.einsum("...gk,gnk->...gn", x, weight) + linear.bias.view(3, 4) + torch.testing.assert_close(actual, expected) + def test_grouped_linear_hands_over_its_stored_weight_and_activation_scales(self): """The groups split the weight as stored (packed FP4 halves K), and the calibrated activation scale goes where `FineGrainedLinear` hands it: NVFP4's as the activation global, a static FP8 one as `As`.""" diff --git a/tests/integrations/finegrained/test_loading_regressions.py b/tests/integrations/finegrained/test_loading_regressions.py new file mode 100644 index 000000000000..2a1ef9ad797f --- /dev/null +++ b/tests/integrations/finegrained/test_loading_regressions.py @@ -0,0 +1,334 @@ +# 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. +"""CPU loading, inference, and saving against local quantized checkpoints with real model mappings.""" + +import os +import tempfile +import unittest +from unittest import mock + +import torch +from safetensors.torch import load_file, save_file + +import transformers.conversion_mapping as conversion_mapping +from transformers import AutoModelForCausalLM, LlamaConfig, NVFP4Config, Qwen2MoeConfig +from transformers.core_model_loading import WeightRenaming +from transformers.integrations.finegrained.conversions import FineGrainedDequantize +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig, Mxfp4Config, Mxfp8Config + +from .utils import _load_and_save, _per_expert + + +@require_torch +class FineGrainedLoadingRegressionTest(unittest.TestCase): + def _checkpoint(self, directory, moe=False, mixed=False, modelopt_mixed=False): + config_class = Qwen2MoeConfig if moe else LlamaConfig + config = config_class( + vocab_size=32, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + **({"num_experts": 2, "num_experts_per_tok": 1, "moe_intermediate_size": 32} if moe else {}), + ) + model = AutoModelForCausalLM.from_config(config).to(torch.bfloat16).eval() + # Saving the model first supplies the architecture's checkpoint key layout, including split experts. + model.save_pretrained(directory) + filename = os.path.join(directory, "model.safetensors") + state = load_file(filename) + quantized = {} + for name, tensor in state.items(): + if tensor.ndim != 2 or name.endswith(("embed_tokens.weight", "lm_head.weight", "mlp.gate.weight")): + quantized[name] = tensor + continue + if modelopt_mixed and ".q_proj." not in name and ".mlp.experts." not in name: + quantized[name] = tensor + continue + stem = name.removesuffix("weight") + if modelopt_mixed and ".q_proj." in name: + quantized[name] = torch.full_like(tensor, 2, dtype=torch.float8_e4m3fn) + quantized[stem + "weight_scale"] = torch.full((1, 1), 1 / 128) + continue + ct = mixed and ".q_proj." in name + quantized[stem + ("weight_packed" if ct else "weight")] = torch.full( + (tensor.shape[0], tensor.shape[1] // 2), 0x22, dtype=torch.uint8 + ) + quantized[stem + "weight_scale"] = torch.ones( + tensor.shape[0], tensor.shape[1] // 16, dtype=torch.float8_e4m3fn + ) + quantized[stem + ("weight_global_scale" if ct else "weight_scale_2")] = torch.tensor( + 64.0 if ct else 1 / 64 + ) + config.quantization_config = ( + FineGrainedConfig( + dtype_plan={r"\.q_proj$": NVFP4Config(checkpoint_format="compressed-tensors"), ".*": NVFP4Config()}, + modules_to_not_convert=["lm_head", "mlp.gate"], + ) + if mixed + else NVFP4Config(modules_to_not_convert=["lm_head", "mlp.gate"]) + ) + if modelopt_mixed: + config.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"}, + "model.layers.0.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16}, + }, + } + config.save_pretrained(directory) + save_file(quantized, filename, metadata={"format": "pt"}) + + def test_dequantized_dense_model_respects_requested_dtype(self): + with tempfile.TemporaryDirectory() as directory: + self._checkpoint(directory) + model = AutoModelForCausalLM.from_pretrained( + directory, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float16, device_map="cpu" + ).eval() + self.assertEqual({parameter.dtype for parameter in model.parameters()}, {torch.float16}) + with torch.no_grad(): + self.assertTrue(torch.isfinite(model(torch.tensor([[1, 2, 3]])).logits).all()) + + def test_dequantized_moe_can_save_and_reload(self): + with tempfile.TemporaryDirectory() as directory, tempfile.TemporaryDirectory() as saved: + self._checkpoint(directory, moe=True) + model = AutoModelForCausalLM.from_pretrained( + directory, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float32, device_map="cpu" + ).eval() + model.save_pretrained(saved) + reloaded, info = AutoModelForCausalLM.from_pretrained(saved, output_loading_info=True) + self.assertEqual(info["missing_keys"], set()) + self.assertEqual(info["unexpected_keys"], set()) + self.assertFalse(hasattr(reloaded.config, "quantization_config")) + for key, value in model.state_dict().items(): + torch.testing.assert_close(reloaded.state_dict()[key], value) + with torch.no_grad(): + inputs = torch.tensor([[1, 2, 3]]) + torch.testing.assert_close(model(inputs).logits, reloaded(inputs).logits) + + def test_mixed_checkpoint_global_conventions_are_scoped_per_module(self): + with tempfile.TemporaryDirectory() as modelopt, tempfile.TemporaryDirectory() as mixed: + self._checkpoint(modelopt, moe=True) + self._checkpoint(mixed, moe=True, mixed=True) + plain = AutoModelForCausalLM.from_pretrained( + modelopt, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float32, device_map="cpu" + ) + loaded = AutoModelForCausalLM.from_pretrained( + mixed, + quantization_config=FineGrainedConfig( + dtype_plan={ + r"\.q_proj$": NVFP4Config(checkpoint_format="compressed-tensors"), + ".*": NVFP4Config(), + }, + modules_to_not_convert=["lm_head", "mlp.gate"], + dequantize=True, + ), + dtype=torch.float32, + device_map="cpu", + ) + for name, tensor in plain.state_dict().items(): + if name.endswith(("q_proj.weight", "experts.gate_up_proj", "experts.down_proj")): + torch.testing.assert_close(loaded.state_dict()[name], tensor) + torch.testing.assert_close(tensor, torch.full_like(tensor, 1 / 64)) + + def test_mixed_block_sizes_use_the_current_modules_grid(self): + # Both block sizes produce a 2x2 scale grid for 192x192 weights, but different boundaries. + quantizer = FineGrainedHfQuantizer( + FineGrainedConfig( + dtype_plan={ + "a$": BlockFp8Config(weight_block_size=(96, 96)), + "b$": BlockFp8Config(weight_block_size=(128, 128)), + }, + dequantize=True, + ) + ) + model = torch.nn.Module() + model.a = torch.nn.Linear(192, 192, bias=False) + model.b = torch.nn.Linear(192, 192, bias=False) + scales = torch.tensor([[1.0, 2.0], [4.0, 8.0]]) + for name, block in (("a", 96), ("b", 128)): + with self.subTest(module=name): + result = FineGrainedDequantize(quantizer).convert( + {"weight": torch.ones(192, 192).to(torch.float8_e4m3fn), "weight_scale_inv": scales}, + full_layer_name=f"{name}.weight", + model=model, + )[f"{name}.weight"] + expected = scales.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1)[:192, :192] + torch.testing.assert_close(result, expected) + + def test_modelopt_mixed_precision_checkpoint_loads_and_saves(self): + for rename_namespace in (False, True): + with ( + self.subTest(rename_namespace=rename_namespace), + tempfile.TemporaryDirectory() as directory, + tempfile.TemporaryDirectory() as saved, + ): + self._checkpoint(directory, moe=True, modelopt_mixed=True) + mapping = conversion_mapping.get_checkpoint_conversion_mapping("qwen2_moe") + if rename_namespace: + config = Qwen2MoeConfig.from_pretrained(directory) + layers = config.quantization_config["quantized_layers"] + config.quantization_config["quantized_layers"] = { + name.replace("model.layers.", "decoder.blocks."): spec for name, spec in layers.items() + } + config.quantization_config = AutoQuantizationConfig.from_dict(config.quantization_config).to_dict() + config.save_pretrained(directory) + filename = os.path.join(directory, "model.safetensors") + save_file( + { + name.replace("model.layers.", "decoder.blocks."): value + for name, value in load_file(filename).items() + }, + filename, + metadata={"format": "pt"}, + ) + mapping = [WeightRenaming(r"^decoder\.blocks\.", "model.layers."), *mapping] + with mock.patch.dict(conversion_mapping._checkpoint_conversion_mapping_cache, {"qwen2_moe": mapping}): + model, info = AutoModelForCausalLM.from_pretrained( + directory, + quantization_config=FineGrainedConfig(dequantize=True), + dtype=torch.float16, + device_map="cpu", + output_loading_info=True, + ) + self.assertEqual(info["missing_keys"], set()) + self.assertEqual(info["unexpected_keys"], set()) + self.assertEqual({parameter.dtype for parameter in model.parameters()}, {torch.float16}) + state = model.state_dict() + for name, tensor in state.items(): + if name.endswith(("q_proj.weight", "experts.gate_up_proj", "experts.down_proj")): + torch.testing.assert_close(tensor, torch.full_like(tensor, 1 / 64)) + original = load_file(os.path.join(directory, "model.safetensors")) + unlisted = "model.layers.0.self_attn.k_proj.weight" + checkpoint_key = ( + unlisted.replace("model.layers.", "decoder.blocks.") if rename_namespace else unlisted + ) + torch.testing.assert_close(state[unlisted], original[checkpoint_key].to(torch.float16)) + with torch.no_grad(): + inputs = torch.tensor([[1, 2, 3]]) + expected = model(inputs).logits + self.assertTrue(torch.isfinite(expected).all()) + model.save_pretrained(saved) + reloaded = AutoModelForCausalLM.from_pretrained(saved).eval() + with torch.no_grad(): + torch.testing.assert_close(reloaded(inputs).logits, expected) + + def test_modelopt_mx_checkpoint_loads_and_saves_dense_and_expert_weights(self): + for format_name, config_class in (("mxfp8", Mxfp8Config), ("mxfp4", Mxfp4Config)): + with self.subTest(format=format_name): + packed = format_name == "mxfp4" + checkpoint = _per_expert( + lambda rows, cols: { + "weight": ( + torch.full((rows, cols // 2), 0xB2, dtype=torch.uint8) + if packed + else torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn) + ), + "weight_scale": torch.full((rows, cols // 32), 121, dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, config_class(checkpoint_format="modelopt")) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + q_proj = model.model.layers[0].self_attn.q_proj + experts = model.model.layers[0].mlp.experts + if packed: + self.assertEqual(q_proj.weight.dtype, torch.int8) + self.assertEqual(experts.gate_up_proj.dtype, torch.int8) + torch.testing.assert_close( + q_proj.weight.view(torch.uint8), checkpoint["model.layers.0.self_attn.q_proj.weight"] + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + actual = saved[name] + if tensor.element_size() == 1: + torch.testing.assert_close(actual.view(torch.uint8), tensor.view(torch.uint8)) + else: + torch.testing.assert_close(actual, tensor) + + def test_modelopt_static_fp8_checkpoint_loads_and_saves_dense_and_experts(self): + checkpoint = _per_expert( + lambda rows, cols: { + "weight": torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn), + "weight_scale": torch.tensor(1 / 128), + "input_scale": torch.tensor(1 / 256), + } + ) + for name in checkpoint: + if name.endswith(".up_proj.weight_scale"): + checkpoint[name] = torch.tensor(1 / 64) + config = AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "FP8"}) + model, info, saved = _load_and_save(checkpoint, config) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + q_proj = model.model.layers[0].self_attn.q_proj + self.assertEqual(q_proj.weight.dtype, torch.float8_e4m3fn) + torch.testing.assert_close(q_proj.weight_scale_inv, torch.tensor(1 / 128)) + torch.testing.assert_close(q_proj.activation_scale, torch.tensor(1 / 256)) + experts = model.model.layers[0].mlp.experts + torch.testing.assert_close( + experts.gate_up_proj_scale_inv, + torch.tensor([[[1 / 128], [1 / 64]], [[1 / 128], [1 / 64]]]), + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + if name.endswith(".weight"): + torch.testing.assert_close(saved[name].view(torch.uint8), tensor.view(torch.uint8)) + else: + torch.testing.assert_close(saved[name], tensor) + + def test_compressed_tensors_mxfp8_scales_reach_dense_and_expert_modules(self): + config = AutoQuantizationConfig.from_dict( + { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32}, + "input_activations": { + "num_bits": 8, + "type": "float", + "strategy": "group", + "group_size": 32, + "dynamic": True, + }, + } + }, + } + ) + self.assertIsInstance(config, Mxfp8Config) + checkpoint = _per_expert( + lambda rows, cols: { + "weight": torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn), + "weight_scale": torch.full((rows, cols // 32), 121, dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, config) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + torch.testing.assert_close( + model.model.layers[0].self_attn.q_proj.weight_scale_inv, + checkpoint["model.layers.0.self_attn.q_proj.weight_scale"], + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + torch.testing.assert_close(saved[name].view(torch.uint8), tensor.view(torch.uint8)) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index c2016691ed77..11b2dbd5c455 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -85,7 +85,7 @@ class FormatSpec: name="fp8", config_cls=BlockFp8Config, quantizer_cls=BlockFp8HfQuantizer, - activation_formats=(None, "fp8"), + activation_formats=(None, "fp8", "bf16"), checkpoint_config=BlockFp8Config, weight_bits=8, checkpoint_keys=( @@ -242,12 +242,12 @@ def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): self.assertTrue(quantizer.quantization_config.dequantize) @require_accelerate - def test_a_gpu_below_compute_capability_8_9_dequantizes_a_quantized_checkpoint(self): + def test_unsupported_gpu_dequantizes_a_quantized_checkpoint(self): quantizer = _quantizer_for(self.spec.checkpoint_config()) quantizer.pre_quantized = True with ( mock.patch("torch.cuda.is_available", return_value=True), - mock.patch("torch.cuda.get_device_capability", return_value=(8, 0)), + mock.patch("torch.cuda.get_device_capability", return_value=(6, 0)), ): quantizer.validate_environment() self.assertTrue(quantizer.quantization_config.dequantize) diff --git a/tests/quantization/finegrained/test_config_readers.py b/tests/quantization/finegrained/test_config_readers.py new file mode 100644 index 000000000000..f5609784d5f0 --- /dev/null +++ b/tests/quantization/finegrained/test_config_readers.py @@ -0,0 +1,239 @@ +# 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. +"""Normalize producer metadata without silently changing module targets or quantization formats.""" + +import unittest + +import torch + +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import NVFP4Config + +from ...integrations.finegrained.utils import _compressed_tensors, _modelopt, _quantizer_for + + +@require_torch +class BlockScaledCheckpointReaderTest(unittest.TestCase): + """modelopt and compressed-tensors NVFP4 exports read as `NVFP4Config`, skip-lists bounded to the modules they name.""" + + 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, so the export is one format.""" + config = _modelopt( + 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"], + } + }, + ) + self.assertIsInstance(config, NVFP4Config) + self.assertIsNone(config.activation_format) + # NVFP4's `dynamic: False` is its calibrated activation global, not a static activation scale + self.assertEqual(config.activation_scheme, "dynamic") + + def test_an_unsupported_modelopt_group_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Unsupported modelopt config_groups entry 'group_0'"): + _modelopt( + config_groups={ + "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} + }, + ) + + def test_mixed_precision_maps_explicit_modules_and_leaves_others_unquantized(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig + + spec = { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "layers.8.attention.wq": {"quant_algo": "FP8_PB_WO"}, + "layers.9.experts": {"quant_algo": "NVFP4", "group_size": 16}, + }, + } + for fields in (spec, {"quantization": spec}): + with self.subTest(nested="quantization" in fields): + config = AutoQuantizationConfig.from_dict({"quant_method": "modelopt", **fields}) + self.assertIsInstance(config, FineGrainedConfig) + dense = plan_entry_for(config.dtype_plan, "model.layers.8.attention.wq") + self.assertIsInstance(dense, BlockFp8Config) + self.assertEqual(dense.activation_format, "bf16") + self.assertEqual(dense.weight_block_size, (128, 128)) + self.assertIsInstance(plan_entry_for(config.dtype_plan, "layers.9.experts.0.gate_proj"), NVFP4Config) + self.assertIsNone(plan_entry_for(config.dtype_plan, "layers.90.experts")) + self.assertIsNone(plan_entry_for(config.dtype_plan, "layers.8.attention.wk")) + restored = AutoQuantizationConfig.from_dict(config.to_dict()) + self.assertEqual(restored.to_dict(), config.to_dict()) + + def test_mixed_precision_rejects_incomplete_or_unsupported_assignments(self): + for fields, message in ( + ({}, "non-empty quantized_layers"), + ({"quantized_layers": {}}, "non-empty quantized_layers"), + ({"quantized_layers": {"layers.0": {"quant_algo": "W4A16_AWQ"}}}, "layers.0"), + ({"quantized_layers": {"layers.0": {"quant_algo": "NVFP4", "group_size": 32}}}, "group_size=16"), + ): + with self.subTest(fields=fields), self.assertRaisesRegex(ValueError, message): + AutoQuantizationConfig.from_dict( + {"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", **fields} + ) + + def test_internal_globs_match_skipped_subtrees(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + skip = _modelopt(exclude_modules=["model.layers.*.self_attn", "*.shared_experts.*"]).modules_to_not_convert + for name in ( + "model.layers.0.self_attn.q_proj", + "model.layers.12.self_attn", + "model.layers.0.mlp.shared_experts.up_proj", + ): + self.assertFalse(should_convert_module(name, skip), name) + for name in ("model.layers.0.mlp.experts", "model.layers.0.my_self_attn", "modelXlayers.0.self_attn"): + self.assertTrue(should_convert_module(name, skip), name) + + def test_modelopt_runtime_transforms_are_not_dropped(self): + from transformers.utils.quantization_config import read_compressed_tensors + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + for field, value in ( + ("pre_quant_scale", True), + ("lora_rank", 16), + ("has_zero_point", True), + ("weight_scale_encoding", "csf"), + ): + for placement in ("top", "group", "layer"): + with self.subTest(field=field, placement=placement), self.assertRaisesRegex(ValueError, field): + if placement == "top": + _modelopt(**{field: value}) + elif placement == "group": + _modelopt(config_groups={"g": {"weights": nvfp4, field: value}}) + else: + AutoQuantizationConfig.from_dict( + { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {"layers.0": {"quant_algo": "NVFP4", field: value}}, + } + ) + # The compressed-tensors reader likewise must leave unsupported transforms to its own backend. + self.assertIsNone(read_compressed_tensors(_compressed_tensors(nvfp4, nvfp4, pre_quant_scale=True))) + + def test_modelopt_groups_accept_exported_module_names(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = _modelopt( + config_groups={ + "g": { + "weights": {"num_bits": 4, "type": "float", "group_size": 16}, + "targets": ["model.layers.0.mlp.experts"], + } + } + ) + self.assertIsInstance(plan_entry_for(config.dtype_plan, "model.layers.0.mlp.experts"), NVFP4Config) + self.assertIsNone(plan_entry_for(config.dtype_plan, "model.layers.1.mlp.experts")) + + def test_a_compressed_tensors_ignore_entry_is_that_module_not_a_prefix(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + skip = AutoQuantizationConfig.from_dict( + _compressed_tensors(nvfp4, ignore=["model.layers.0.mlp.gate", "re:.*router$"]) + ).modules_to_not_convert + self.assertFalse(should_convert_module("model.layers.0.mlp.gate", skip)) + self.assertFalse(should_convert_module("model.layers.3.mlp.router", skip)) + self.assertTrue(should_convert_module("model.layers.0.mlp.gate_proj", skip)) + + 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 + + # as a modelopt checkpoint ships them, through the config that normalizes them + skip = AutoQuantizationConfig.from_dict( + { + "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.01.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) + + def test_a_compressed_tensors_nvfp4_checkpoint_reads_as_nvfp4(self): + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + config = AutoQuantizationConfig.from_dict(_compressed_tensors(nvfp4, nvfp4, ignore=["lm_head"])) + self.assertIsInstance(config, NVFP4Config) + self.assertEqual(config.checkpoint_format, "compressed-tensors") + + def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): + """Integer, per-channel FP8 and sparse checkpoints are compressed-tensors' own.""" + from transformers.utils.quantization_config import read_compressed_tensors + + int4 = {"num_bits": 4, "type": "int", "group_size": 128, "strategy": "group"} + channel = {"num_bits": 8, "type": "float", "strategy": "channel"} + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + for checkpoint in ( + _compressed_tensors(int4, fmt="pack-quantized"), + _compressed_tensors(channel, fmt="float-quantized"), + _compressed_tensors(nvfp4, sparsity_config={"format": "sparse-24-bitmask"}), + ): + self.assertIsNone(read_compressed_tensors(checkpoint)) + + def test_activations_the_weights_kernels_do_not_serve_stay_with_compressed_tensors(self): + """An activation scheme is read only where the weights' kernels serve it: NVFP4 weights with FP8 per-token + activations and FP8 weights with int8 activations stay compressed-tensors' own, as do + group-32 FP8 weights whose scales are not E8M0 (MXFP8's).""" + from transformers.utils.quantization_config import Mxfp8Config, read_compressed_tensors + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} + fp8_token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} + int8_token = {"num_bits": 8, "type": "int", "strategy": "token", "dynamic": True} + group_32 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32} + for name, checkpoint in ( + ("nvfp4 x fp8", _compressed_tensors(nvfp4, fp8_token)), + ("block-fp8 x int8", _compressed_tensors(block, int8_token, fmt="float-quantized")), + ( + "group-32 fp32 scales", + _compressed_tensors({**group_32, "scale_dtype": "torch.float32"}, group_32, fmt="float-quantized"), + ), + ): + with self.subTest(name): + self.assertIsNone(read_compressed_tensors(checkpoint)) + e8m0 = {**group_32, "scale_dtype": "torch.float8_e8m0fnu"} + self.assertIsInstance( + read_compressed_tensors(_compressed_tensors(e8m0, group_32, fmt="float-quantized")), Mxfp8Config + ) + + def test_unrepresentable_modelopt_targets_raise_before_loading(self): + model = torch.nn.Module() + model.linear = torch.nn.Linear(16, 16) + for target in ("unknown", "experts.0.gate_proj"): + config = AutoQuantizationConfig.from_dict( + { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {target: {"quant_algo": "NVFP4"}}, + } + ) + with self.subTest(target=target), self.assertRaisesRegex(ValueError, "matches no supported"): + _quantizer_for(config)._process_model_before_weight_loading(model) diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index 556ca11ac2f0..dd2b7d4330b0 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -17,6 +17,7 @@ import tempfile import unittest +from unittest.mock import patch import torch @@ -49,6 +50,32 @@ ) +class Mxfp4EnvironmentTest(unittest.TestCase): + def test_weight_only_mxfp4_keeps_existing_cuda_architectures(self): + for capability in ((7, 5), (8, 0), (8, 6)): + with self.subTest(capability=capability): + config = Mxfp4Config() + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=capability), + patch("transformers.quantizers.block_scaled_utils.is_accelerate_available", return_value=True), + ): + _quantizer_for(config).validate_environment(device_map="cuda") + self.assertFalse(config.dequantize) + + def test_quantized_activations_keep_the_fp8_hardware_requirement(self): + for activation_format in ("mxfp4", "mxfp8"): + with self.subTest(activation_format=activation_format): + config = Mxfp4Config(activation_format=activation_format) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=(8, 0)), + patch("transformers.quantizers.block_scaled_utils.is_accelerate_available", return_value=True), + ): + _quantizer_for(config).validate_environment(device_map="cuda") + self.assertTrue(config.dequantize) + + @require_torch_accelerator @require_kernels class Mxfp4GenerateTest(unittest.TestCase): diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index 35f3b885c411..6c75f4e83e25 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -32,9 +32,7 @@ E, _Cfg, _ChainAssertions, - _compressed_tensors, _load_and_save, - _modelopt, _per_expert, _quantizer_for, _targets, @@ -358,113 +356,6 @@ def test_global_scale_converters_carry_an_expert_index(self): ) -@require_torch -class Nvfp4CheckpointReaderTest(unittest.TestCase): - """modelopt and compressed-tensors NVFP4 exports read as `NVFP4Config`, skip-lists bounded to the modules they name.""" - - 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, so the export is one format.""" - config = _modelopt( - 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"], - } - }, - ) - self.assertIsInstance(config, NVFP4Config) - self.assertIsNone(config.activation_format) - # NVFP4's `dynamic: False` is its calibrated activation global, not a static activation scale - self.assertEqual(config.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.""" - config = _modelopt( - config_groups={ - "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} - }, - ) - self.assertIsInstance(config, NVFP4Config) - - def test_a_compressed_tensors_ignore_entry_is_that_module_not_a_prefix(self): - from transformers.quantizers.quantizers_utils import should_convert_module - - nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} - skip = AutoQuantizationConfig.from_dict( - _compressed_tensors(nvfp4, ignore=["model.layers.0.mlp.gate", "re:.*router$"]) - ).modules_to_not_convert - self.assertFalse(should_convert_module("model.layers.0.mlp.gate", skip)) - self.assertFalse(should_convert_module("model.layers.3.mlp.router", skip)) - self.assertTrue(should_convert_module("model.layers.0.mlp.gate_proj", skip)) - - 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 - - # as a modelopt checkpoint ships them, through the config that normalizes them - skip = AutoQuantizationConfig.from_dict( - { - "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) - - def test_a_compressed_tensors_nvfp4_checkpoint_reads_as_nvfp4(self): - nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} - config = AutoQuantizationConfig.from_dict(_compressed_tensors(nvfp4, nvfp4, ignore=["lm_head"])) - self.assertIsInstance(config, NVFP4Config) - self.assertEqual(config.checkpoint_format, "compressed-tensors") - - def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): - """Integer, per-channel FP8 and sparse checkpoints are compressed-tensors' own.""" - from transformers.utils.quantization_config import read_compressed_tensors - - int4 = {"num_bits": 4, "type": "int", "group_size": 128, "strategy": "group"} - channel = {"num_bits": 8, "type": "float", "strategy": "channel"} - nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} - for checkpoint in ( - _compressed_tensors(int4, fmt="pack-quantized"), - _compressed_tensors(channel, fmt="float-quantized"), - _compressed_tensors(nvfp4, sparsity_config={"format": "sparse-24-bitmask"}), - ): - self.assertIsNone(read_compressed_tensors(checkpoint)) - - def test_activations_the_weights_kernels_do_not_serve_stay_with_compressed_tensors(self): - """An activation scheme is read only where the weights' kernels serve it: NVFP4 weights with FP8 per-token - activations, weight-only block-FP8 and FP8 weights with int8 activations stay compressed-tensors' own, as do - group-32 FP8 weights whose scales are not E8M0 (MXFP8's).""" - from transformers.utils.quantization_config import Mxfp8Config, read_compressed_tensors - - nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} - block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} - fp8_token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} - int8_token = {"num_bits": 8, "type": "int", "strategy": "token", "dynamic": True} - group_32 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32} - for name, checkpoint in ( - ("nvfp4 x fp8", _compressed_tensors(nvfp4, fp8_token)), - ("block-fp8 weight-only", _compressed_tensors(block, fmt="float-quantized")), - ("block-fp8 x int8", _compressed_tensors(block, int8_token, fmt="float-quantized")), - ( - "group-32 fp32 scales", - _compressed_tensors({**group_32, "scale_dtype": "torch.float32"}, group_32, fmt="float-quantized"), - ), - ): - with self.subTest(name): - self.assertIsNone(read_compressed_tensors(checkpoint)) - e8m0 = {**group_32, "scale_dtype": "torch.float8_e8m0fnu"} - self.assertIsInstance( - read_compressed_tensors(_compressed_tensors(e8m0, group_32, fmt="float-quantized")), Mxfp8Config - ) - - class Nvfp4Test(BlockScaledFormatTesterMixin, unittest.TestCase): spec = NVFP4 From 903b70c80886815d57a90775a3b2948725cf4424 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 07:03:53 -0700 Subject: [PATCH 82/91] Reject conflicting producer quantization group targets --- src/transformers/utils/quantization_config.py | 18 +++++++- .../finegrained/test_config_readers.py | 41 ++++++++++++++++++- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 0e5586769843..db2ecd2663b1 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -2049,6 +2049,14 @@ def _modelopt_format(spec: dict) -> BlockScaledConfig: return config_class(activation_format=activation_format, checkpoint_format="modelopt") +def _merge_config_group_plan(dtype_plan: dict, entries: dict): + """Combine groups without making their order decide a target's quantization format.""" + for pattern, config in entries.items(): + if pattern in dtype_plan and dtype_plan[pattern].to_dict() != config.to_dict(): + raise ValueError(f"Conflicting quantization configs for target {pattern!r}") + dtype_plan[pattern] = config + + def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: """Read a ModelOpt format or its explicit per-module `MIXED_PRECISION` assignment. @@ -2100,7 +2108,7 @@ def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: raise ValueError(f"Unsupported modelopt config_groups entry {name!r}: {spec}") for config in entries.values(): config.checkpoint_format = "modelopt" - dtype_plan.update(entries) + _merge_config_group_plan(dtype_plan, entries) return _config_for_plan(dtype_plan, modules_to_not_convert) config = _modelopt_format(config_dict) config.modules_to_not_convert = modules_to_not_convert @@ -2124,7 +2132,13 @@ def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None mapped = [plan_from_config_groups(spec) for spec in groups.values()] if any(entries is None for entries in mapped): return None - dtype_plan = {pattern: config for entries in mapped if entries is not None for pattern, config in entries.items()} + dtype_plan = {} + try: + for entries in mapped: + if entries is not None: + _merge_config_group_plan(dtype_plan, entries) + except ValueError: + return None for config in dtype_plan.values(): config.checkpoint_format = "compressed-tensors" if config_dict.get("kv_cache_scheme"): diff --git a/tests/quantization/finegrained/test_config_readers.py b/tests/quantization/finegrained/test_config_readers.py index f5609784d5f0..4b695c7b0c95 100644 --- a/tests/quantization/finegrained/test_config_readers.py +++ b/tests/quantization/finegrained/test_config_readers.py @@ -19,7 +19,7 @@ from transformers.quantizers.auto import AutoQuantizationConfig from transformers.testing_utils import require_torch -from transformers.utils.quantization_config import NVFP4Config +from transformers.utils.quantization_config import NVFP4Config, read_compressed_tensors, read_modelopt from ...integrations.finegrained.utils import _compressed_tensors, _modelopt, _quantizer_for @@ -237,3 +237,42 @@ def test_unrepresentable_modelopt_targets_raise_before_loading(self): ) with self.subTest(target=target), self.assertRaisesRegex(ValueError, "matches no supported"): _quantizer_for(config)._process_model_before_weight_loading(model) + + +class FineGrainedConfigGroupReaderTest(unittest.TestCase): + def _groups(self, reverse=False, identical=False): + nvfp4 = {"targets": ["Linear"], "weights": {"num_bits": 4, "type": "float", "group_size": 16}} + fp8 = { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]}, + } + groups = [("nvfp4", nvfp4), ("other", nvfp4 if identical else fp8)] + return dict(reversed(groups) if reverse else groups) + + def test_modelopt_rejects_conflicting_group_targets_in_either_order(self): + for reverse in (False, True): + with ( + self.subTest(reverse=reverse), + self.assertRaisesRegex(ValueError, "Conflicting quantization configs for target"), + ): + read_modelopt({"quant_algo": "NVFP4", "config_groups": self._groups(reverse=reverse)}) + + def test_compressed_tensors_keeps_its_backend_for_conflicting_group_targets(self): + for reverse in (False, True): + with self.subTest(reverse=reverse): + self.assertIsNone( + read_compressed_tensors( + {"format": "float-quantized", "config_groups": self._groups(reverse=reverse)} + ) + ) + + def test_identical_group_targets_can_share_a_config(self): + for reverse in (False, True): + for reader, fields, checkpoint_format in ( + (read_modelopt, {"quant_algo": "NVFP4"}, "modelopt"), + (read_compressed_tensors, {"format": "float-quantized"}, "compressed-tensors"), + ): + with self.subTest(reverse=reverse, producer=checkpoint_format): + config = reader({**fields, "config_groups": self._groups(reverse=reverse, identical=True)}) + self.assertIsInstance(config, NVFP4Config) + self.assertEqual(config.checkpoint_format, checkpoint_format) From 3af3c2b9343f465fcbf4030b77e894face7e9f84 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 07:12:43 -0700 Subject: [PATCH 83/91] Preserve ModelOpt native exclusions and trailing-dot module prefixes --- src/transformers/utils/quantization_config.py | 14 ++++- .../finegrained/test_config_readers.py | 58 +++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index db2ecd2663b1..353492d5f30c 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1699,7 +1699,7 @@ def subtree_patterns(globs: list[str]) -> list[str]: """ patterns = [] for glob in globs: - subtree = glob.removesuffix(".*") + subtree = glob.removesuffix(".*").rstrip(".") pattern = fnmatch.translate(subtree).removesuffix(r"\Z").removesuffix(r"\z") patterns.append(r"(?:^|.*\.)" + pattern + r"(\..*)?$") return patterns @@ -2071,8 +2071,16 @@ def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: "This modelopt checkpoint calibrates a quantized KV cache; the KV cache stays in the model's dtype, " "so its scales are not loaded." ) - ignore = config_dict.get("ignore") or config_dict.get("exclude_modules") - modules_to_not_convert = subtree_patterns(ignore) if ignore is not None else None + # Producer exclusions are globs; an explicitly supplied native skip-list already uses + # the library's regex semantics and must not be escaped or discarded. + skip_fields = ("modules_to_not_convert", "ignore", "exclude_modules") + native_skip = config_dict.get("modules_to_not_convert") or [] + producer_skip = [*(config_dict.get("ignore") or []), *(config_dict.get("exclude_modules") or [])] + modules_to_not_convert = ( + list(dict.fromkeys([*native_skip, *subtree_patterns(producer_skip)])) + if any(config_dict.get(field) is not None for field in skip_fields) + else None + ) if quant_algo == "MIXED_PRECISION": layers = config_dict.get("quantized_layers") if not isinstance(layers, dict) or not layers: diff --git a/tests/quantization/finegrained/test_config_readers.py b/tests/quantization/finegrained/test_config_readers.py index 4b695c7b0c95..25782b173ade 100644 --- a/tests/quantization/finegrained/test_config_readers.py +++ b/tests/quantization/finegrained/test_config_readers.py @@ -276,3 +276,61 @@ def test_identical_group_targets_can_share_a_config(self): config = reader({**fields, "config_groups": self._groups(reverse=reverse, identical=True)}) self.assertIsInstance(config, NVFP4Config) self.assertEqual(config.checkpoint_format, checkpoint_format) + + def test_modelopt_trailing_separator_exclusions_match_only_the_named_layers(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + config = read_modelopt( + { + "quant_algo": "NVFP4", + "exclude_modules": ["layers.0.", "layers.61.", "self_attn", "lm_head"], + } + ) + for name in ( + "model.layers.0.mlp.experts", + "model.layers.0", + "model.layers.61.mlp.up_proj", + "model.layers.61", + "model.layers.3.self_attn.q_proj", + "lm_head", + ): + with self.subTest(skipped=name): + self.assertFalse(should_convert_module(name, config.modules_to_not_convert)) + for name in ( + "model.layers.00.mlp.up_proj", + "model.layers.610.mlp.up_proj", + "model.layers.6.mlp.up_proj", + "model.layers.1.mlp.up_proj", + ): + with self.subTest(converted=name): + self.assertTrue(should_convert_module(name, config.modules_to_not_convert)) + + def test_modelopt_merges_native_regex_skip_list_with_both_producer_exclusions(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + native_pattern = r"^model\.audio_tower(?:\.|$)" + config = read_modelopt( + { + "quant_algo": "NVFP4", + "modules_to_not_convert": [native_pattern], + "ignore": ["model.visual_tower.*"], + "exclude_modules": ["layers.0."], + } + ) + self.assertIn(native_pattern, config.modules_to_not_convert) + for name in ( + "model.audio_tower.layers.1.proj", + "model.visual_tower.layers.1.proj", + "model.layers.0.mlp.up_proj", + ): + with self.subTest(skipped=name): + self.assertFalse(should_convert_module(name, config.modules_to_not_convert)) + for name in ( + "model.audio_tower_extra.proj", + "model.visual_tower_extra.proj", + "model.layers.00.mlp.up_proj", + ): + with self.subTest(converted=name): + self.assertTrue(should_convert_module(name, config.modules_to_not_convert)) + restored = AutoQuantizationConfig.from_dict(config.to_dict()) + self.assertEqual(restored.modules_to_not_convert, config.modules_to_not_convert) From 63e39561ce71032759f37ca2cbe607861b732af4 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 07:26:28 -0700 Subject: [PATCH 84/91] Refresh Muse NVFP4 fixture expectations and verify normalized FP8 scales --- docs/source/en/quantization/finegrained.md | 4 ++-- .../block_scaled_checkpoint_tester.py | 18 +++++++++++++----- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 0e76cd1bdefc..e0ac139a6909 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -88,8 +88,8 @@ that cannot be represented by an architecture's fused expert module. AWQ smoothing (`pre_quant_scale`), SVD residuals, asymmetric zero points, sparse or transformed layouts, and compressed scale encodings such as `weight_scale_encoding="csf"` are rejected. NVFP4 embedding tables are not -supported by the FP8 embedding integration. Weight-only block-FP8 currently runs on dense linears and eager -experts; fused expert backends require kernel support for that activation mode. +supported by the FP8 embedding integration. Weight-only block-FP8 requires an updated kernel bundle with weight-only support for the selected dense or +expert backend. To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass `FineGrainedConfig(dequantize=True)`. diff --git a/tests/quantization/block_scaled_checkpoint_tester.py b/tests/quantization/block_scaled_checkpoint_tester.py index 730b2193db37..c6626db88693 100644 --- a/tests/quantization/block_scaled_checkpoint_tester.py +++ b/tests/quantization/block_scaled_checkpoint_tester.py @@ -135,9 +135,9 @@ def _checkpoint_bytes(path): "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], + [0.0430, -0.0131, -0.0304, 0.0415, -0.0082, -0.0029, 0.0042, -0.0498], + [0.0105, 0.0723, 0.0615, -0.0175, 0.0522, 0.0015, 0.0023, 0.0304], + [0.0089, -0.0251, 0.0131, 0.0234, 0.0243, -0.0283, -0.0703, 0.0664], ], } ), # fmt: skip @@ -259,7 +259,8 @@ def test_save_writes_the_checkpoint_layout_back(self): self.assertEqual(shape, source_shape) if dtype != source_dtype: # Loading honors initialized dtypes, including FP32 norms and BF16 - # MiniMax routers/patch embeddings. Float32 UE8M0 containers save natively. + # MiniMax routers/patch embeddings. FP8 scales normalize to FP32; + # float32 UE8M0 containers save natively. self.assertTrue( ( dtype == torch.float32 @@ -273,11 +274,18 @@ def test_save_writes_the_checkpoint_layout_back(self): ("gate.weight", "e_score_correction_bias", "patch_embedding.weight") ) ) + or ( + dtype == torch.float32 + and source_dtype == torch.bfloat16 + and key.endswith("_scale_inv") + and written.get(key.removesuffix("_scale_inv"), (None,))[0] == torch.float8_e4m3fn + ) or ( dtype == torch.float8_e8m0fnu and source_dtype == torch.float32 and key.endswith((".scale", "_scale_inv")) - ) + ), + f"{repo}: {key} unexpectedly changed dtype from {source_dtype} to {dtype}", ) expected = torch.frombuffer(bytearray(source_raw), dtype=source_dtype).to(dtype) self.assertEqual(raw, expected.view(torch.uint8).numpy().tobytes()) From 0a193c0b6bab850df6ac4f6dec7a18f12b10b3eb Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 08:31:51 -0700 Subject: [PATCH 85/91] Keep format rules in their quantizers and load packed NVFP4 embeddings --- docs/source/en/quantization/finegrained.md | 14 +-- docs/source/en/quantization/nvfp4.md | 4 +- .../distributed/tensor_parallel.py | 8 ++ .../integrations/finegrained/layers.py | 72 ++++++++++++--- .../quantizers/block_scaled_utils.py | 67 ++++++++------ .../quantizers/quantizer_blockfp8.py | 39 +++++++- .../quantizers/quantizer_finegrained.py | 4 +- .../quantizers/quantizer_nvfp4.py | 21 +++++ tests/integrations/finegrained/test_layers.py | 58 ++++++++++++ .../finegrained/test_loading_regressions.py | 90 ++++++++++++++++++- 10 files changed, 324 insertions(+), 53 deletions(-) diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index e0ac139a6909..425e42a10e20 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -50,7 +50,8 @@ The saved model records its plan, so it reloads with no config. module raise an error. - A module no entry matches, or that `modules_to_not_convert` names, stays in full precision. - `modules_to_convert` names embedding tables to hold in FP8 with one per-tensor scale, whatever the plan's formats. -- `dequantize=True` loads every entry dequantized to full precision. +- `dequantize=True` expands linear and expert weights to full precision. FP8 embedding tables stay compressed; + plans containing NVFP4 embedding tables currently reject this option. ## Checkpoints that mix formats @@ -83,13 +84,16 @@ For example, ModelOpt's per-module assignment can combine weight-only block-FP8 The model must support the checkpoint's architecture and parameter names. Format parsing alone does not add support for a new architecture. Module names are resolved through the model's checkpoint key renames. An -assignment that matches no supported linear or expert module raises an error, including per-expert assignments +assignment that matches no supported linear, expert, or embedding module raises an error, including per-expert assignments that cannot be represented by an architecture's fused expert module. AWQ smoothing (`pre_quant_scale`), SVD residuals, asymmetric zero points, sparse or transformed layouts, and -compressed scale encodings such as `weight_scale_encoding="csf"` are rejected. NVFP4 embedding tables are not -supported by the FP8 embedding integration. Weight-only block-FP8 requires an updated kernel bundle with weight-only support for the selected dense or -expert backend. +compressed scale encodings such as `weight_scale_encoding="csf"` are rejected. Explicit NVFP4 embedding targets +load as packed tables with E4M3 scales per 16 values and a shared FP32 multiplier. Lookup decodes only selected +rows, and saving preserves the checkpoint's row shards. These tables cannot currently be loaded with +`dequantize=True`, including the automatic CPU dequantization fallback. Supporting the embedding layout alone +does not make a checkpoint using unsupported scale encodings loadable. Weight-only block-FP8 requires an updated +kernel bundle with weight-only support for the selected dense or expert backend. To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass `FineGrainedConfig(dequantize=True)`. diff --git a/docs/source/en/quantization/nvfp4.md b/docs/source/en/quantization/nvfp4.md index 88373397e25e..33ece19b7e0d 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -52,7 +52,9 @@ output = model.generate(**inputs, max_new_tokens=20) print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -Pass `NVFP4Config(dequantize=True)` to load the weights dequantized to bf16 instead. +Pass `NVFP4Config(dequantize=True)` to load linear and expert weights dequantized to bf16 instead. +Checkpoints with explicit NVFP4 embedding targets keep those tables packed and decode only selected rows during +lookup; they currently reject `dequantize=True`, including the automatic CPU fallback. ## Quantize on the fly diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index b11a82393f10..3d7e468a8b64 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -199,6 +199,14 @@ def should_use_local_tensors(self, module): def validate_param(self, module, param, mesh, parameter_name=None): meta = module._parameters.get(param) + group_size = getattr(module, "_hf_quantized_embedding_group_size", None) + if group_size is not None and ( + module.embedding_dim % mesh.size() or module.embedding_dim // mesh.size() % group_size + ): + raise ValueError( + f"Quantized embedding width {module.embedding_dim} must split into whole groups of {group_size} " + f"for tensor parallel size {mesh.size()}." + ) gathers_output = isinstance(self.output_layouts, Replicate) if meta is None or not gathers_output: return diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index a51d6b58e7f5..20297f54672d 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -467,13 +467,8 @@ 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.""" - - _hf_quantized_needs_local_tp = True +class FineGrainedEmbedding(_FineGrainedModule, nn.Embedding): + """FP8 or NVFP4 table decoded only for the rows selected by a lookup.""" def __init__( self, @@ -482,24 +477,73 @@ def __init__( padding_idx: int | None = None, *, quantization_config: BlockScaledConfig | None = None, + dtype: torch.dtype | None = None, ): - super().__init__(num_embeddings, embedding_dim, padding_idx=padding_idx) + nn.Module.__init__(self) + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + if padding_idx is not None: + if not -num_embeddings <= padding_idx < num_embeddings: + raise ValueError("padding_idx must be within num_embeddings.") + padding_idx %= num_embeddings + self.padding_idx = padding_idx + self.max_norm = None + self.norm_type = 2.0 + self.scale_grad_by_freq = False + self.sparse = False self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) - 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, dtype=torch.float32), requires_grad=False) + if self.weight_format not in ("fp8", "nvfp4") or ( + self.weight_format == "fp8" + and self.quantization_config.weight_block_size is not None + and not self.quantization_config.modules_to_convert + ): + raise ValueError("Fine-grained embeddings support per-tensor FP8 or NVFP4 weights.") + self.compute_dtype = dtype or torch.get_default_dtype() + if self.weight_format == "nvfp4": + self._hf_quantized_embedding_group_size = 16 + if embedding_dim % 16: + raise ValueError("NVFP4 embedding width must be divisible by its scale group size (16).") + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim // 2, dtype=torch.int8), requires_grad=False + ) + self.weight_scale_inv = nn.Parameter( + torch.empty(num_embeddings, embedding_dim // 16, dtype=_FP8_DTYPE), requires_grad=False + ) + self.weight_global_scale = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) + else: + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=_FP8_DTYPE), requires_grad=False + ) + self.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) @classmethod def from_module( cls, module: nn.Embedding, quantization_config: BlockScaledConfig, model: nn.Module ) -> FineGrainedEmbedding: return cls( - module.num_embeddings, module.embedding_dim, module.padding_idx, quantization_config=quantization_config + module.num_embeddings, + module.embedding_dim, + module.padding_idx, + quantization_config=quantization_config, + dtype=module.weight.dtype, ) def forward(self, input: torch.Tensor) -> torch.Tensor: - rows = super().forward(input) - return rows.to(self.weight_scale.dtype) * self.weight_scale + if self.weight_format == "fp8": + rows = super().forward(input) + return rows.to(self.weight_scale.dtype) * self.weight_scale + packed = nn.functional.embedding(input, self.weight).view(torch.uint8) + scales = nn.functional.embedding(input.to(self.weight_scale_inv.device), self.weight_scale_inv).to( + device=packed.device, dtype=torch.float32 + ) + low, high = packed & 15, packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2).long() + values = 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, + ) + rows = values[codes] * scales.repeat_interleave(16, dim=-1) * self.weight_global_scale.to(packed.device) + return rows.to(self.compute_dtype) class FineGrainedGroupedLinear(FineGrainedLinear): diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py index 4ee0578e2b74..6935ae317273 100644 --- a/src/transformers/quantizers/block_scaled_utils.py +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -30,8 +30,9 @@ from .quantizers_utils import ( FUSED_STACKS, STACK_COMPANIONS, - embedding_pattern, get_module_from_name, + module_kind, + plan_entry_for, replace_with_quantized_modules, should_convert_module, ) @@ -59,6 +60,7 @@ class BlockScaledQuantizer(HfQuantizer): # narrowed once the modules are built: whether this format holds experts, and the up projection they stack # (`gate_up_proj` gated, `up_proj` not) holds_experts = True + holds_embeddings = False up_projections: tuple[str, ...] = ("gate_up_proj",) def __init__(self, quantization_config, **kwargs): @@ -95,6 +97,12 @@ def format_conversions(self, dequantize: bool) -> list: what folding this format's keys needs beyond the shared `weight` + `weight_scale_inv` chain.""" return [] + def format_weight_conversions(self, converters, model): + """Adjust the model's companion converters after the base quantizer has cloned them.""" + + def format_parallel_plan(self, model, plan): + """Set this format's companion placements after the shared parallel plan is built.""" + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: """The ops a tensor loading into `param_name` runs, for the layout this format's modules hold it in: an expert's stacked gate|up rows interleaved (weight, scale grid and bias share the row axis), a scale @@ -146,8 +154,7 @@ def _per_expert_sources(proj: str, slot: str) -> list[str]: return [f"experts.*.{half}.{slot}" for half in halves] def _block_scale_rename(self): - """A producer's `weight_scale` block scale onto the `weight_scale_inv` slot every module holds it in; the - embedding tables the plan quantizes hold their per-tensor scale under that name, and keep it.""" + """Rename block scales while leaving the host's scalar embedding scales untouched.""" from ..core_model_loading import WeightRenaming kept = "".join(rf"(? Date: Fri, 16 Oct 2026 10:11:09 -0700 Subject: [PATCH 86/91] Dispatch quantized embedding lookups through the required kernel API --- src/transformers/integrations/finegrained/layers.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index 20297f54672d..b6f2578d993a 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -97,6 +97,7 @@ class FineGrained: get_supported_act_fns: Callable get_supported_norms: Callable rms_norm_rows: Callable + embedding_lookup: Callable # Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env @@ -529,6 +530,16 @@ def from_module( ) def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.weight.device.type in ("cuda", "xpu") and not ( + self.weight_format == "fp8" and torch.compiler.is_compiling() + ): + scale = self.weight_scale if self.weight_format == "fp8" else self.weight_scale_inv + global_scale = None if self.weight_format == "fp8" else self.weight_global_scale + if input.device == self.weight.device == scale.device and ( + global_scale is None or global_scale.device == self.weight.device + ): + dtype = scale.dtype if global_scale is None else self.compute_dtype + return load_finegrained_kernel().embedding_lookup(self.weight, input, scale, global_scale, dtype) if self.weight_format == "fp8": rows = super().forward(input) return rows.to(self.weight_scale.dtype) * self.weight_scale From abe8b7c2db884ae065063cf4e8fecabc3a74f8bd Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 11:20:25 -0700 Subject: [PATCH 87/91] Reject unsupported quantized embedding vocabulary sharding --- docs/source/en/quantization/finegrained.md | 4 +++ .../distributed/tensor_parallel.py | 11 ++++++++ tests/integrations/finegrained/test_layers.py | 25 +++++++++++++++++++ 3 files changed, 40 insertions(+) diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index 425e42a10e20..a04371df6934 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -53,6 +53,10 @@ The saved model records its plan, so it reloads with no config. - `dequantize=True` expands linear and expert weights to full precision. FP8 embedding tables stay compressed; plans containing NVFP4 embedding tables currently reject this option. +Quantized embedding tables support column-wise tensor parallelism (`colwise` or `colwise_gather_output`). +Vocabulary sharding (`rowwise`) is currently unsupported. For NVFP4, each rank's embedding width must be divisible +by the scale group size of 16. + ## Checkpoints that mix formats These load as a [`FineGrainedConfig`] without one being passed: diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 3d7e468a8b64..adab74502cd3 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -293,6 +293,17 @@ def should_use_local_tensors(self, module): uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled() return use_local_quantized_path or uses_local_inference_kernel + def validate_param(self, module, param, mesh, parameter_name=None): + if ( + isinstance(module, torch.nn.Embedding) + and getattr(module, "_hf_quantized_needs_local_tp", False) + and mesh.size() > 1 + ): + raise ValueError( + "Rowwise tensor parallelism is not supported for quantized embeddings. " + "Use colwise or colwise_gather_output to shard the embedding width." + ) + def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py index aa1e8fff3597..02aa2518aadb 100644 --- a/tests/integrations/finegrained/test_layers.py +++ b/tests/integrations/finegrained/test_layers.py @@ -63,6 +63,16 @@ def test_nvfp4_lookup_uses_packed_nibbles_block_scales_and_shared_global(self): with self.assertRaises((IndexError, RuntimeError)): m(torch.tensor([8])) + def test_nvfp4_lookup_decodes_every_e2m1_value_without_a_device_lut(self): + m = fg.FineGrainedEmbedding(2, 16, quantization_config=NVFP4Config(), dtype=torch.float32) + m.weight.data.copy_(torch.arange(8, dtype=torch.uint8).mul(17).view(torch.int8).expand(2, 8)) + m.weight_scale_inv.data.fill_(1) + m.weight_global_scale.data.fill_(1) + expected = torch.tensor([0.0, 0.0, 0.5, 0.5, 1.0, 1.0, 1.5, 1.5, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 6.0, 6.0]) + torch.testing.assert_close(m(torch.tensor([1, 0])), expected.expand(2, 16), rtol=0, atol=0) + m.weight.data.bitwise_or_(-120) # Set both sign bits without changing the magnitude. + torch.testing.assert_close(m(torch.tensor([0])), -expected.unsqueeze(0), rtol=0, atol=0) + def test_storage_can_be_built_on_meta_without_a_full_precision_table(self): with torch.device("meta"): m = fg.FineGrainedEmbedding(320001536, 160, quantization_config=NVFP4Config()) @@ -85,6 +95,21 @@ def test_colwise_tp_requires_complete_scale_groups(self): with self.assertRaisesRegex(ValueError, "whole groups of 16"): style.validate_param(invalid, "weight", mesh) + def test_rowwise_tp_rejects_quantized_vocabulary_shards(self): + from transformers.distributed.tensor_parallel import RowwiseParallel + + mesh = mock.Mock() + mesh.size.return_value = 2 + style = RowwiseParallel() + style.validate_param(torch.nn.Embedding(8, 32), "weight", mesh) + for config in (BlockFp8Config(weight_block_size=None), NVFP4Config()): + with self.subTest(format=config.method): + table = fg.FineGrainedEmbedding(8, 32, quantization_config=config) + with self.assertRaisesRegex(ValueError, "Rowwise.*quantized embeddings"): + style.validate_param(table, "weight", mesh) + mesh.size.return_value = 1 + style.validate_param(table, "weight", mesh) + def test_non_mx_formats_do_not_require_ue8m0(self): from transformers.integrations.finegrained import conversions From 4272b401bbfb671ffbd7fa8949fa13c1cfd312ee Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Fri, 16 Oct 2026 12:07:22 -0700 Subject: [PATCH 88/91] Disable extra autotuning phases explicitly in correctness tests --- tests/kernels/conftest.py | 14 +++++++++++--- tests/quantization/block_scaled_tester.py | 12 ++++++++++-- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py index 31861ce9245c..6fab2cbf664f 100644 --- a/tests/kernels/conftest.py +++ b/tests/kernels/conftest.py @@ -20,8 +20,16 @@ @pytest.fixture(scope="session", autouse=True) def cheap_finegrained_tuning(): - # Use the minimum Bayesian search budget for correctness tests. Set this before the - # kernel bundle loads: its autotuners read the trial budget at import time. + # Measure one configuration without extra search phases for correctness tests. + # Set this before kernel import: autotuners read these controls at import time. # Subprocesses inherit it too, including distributed workers. - with mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}): + with mock.patch.dict( + os.environ, + { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", + "FINEGRAINED_AUTOTUNE_REFINE": "0", + "FINEGRAINED_AUTOTUNE_RUNOFF": "0", + }, + ): yield diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index 11b2dbd5c455..b948e4ac8878 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -52,8 +52,16 @@ def _cheap_tuning(): - """Use the minimum Bayesian search budget for correctness tests.""" - return mock.patch.dict(os.environ, {"FINEGRAINED_AUTOTUNE_TRIALS": "1"}) + """Measure one configuration without extra search phases for correctness tests.""" + return mock.patch.dict( + os.environ, + { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", + "FINEGRAINED_AUTOTUNE_REFINE": "0", + "FINEGRAINED_AUTOTUNE_RUNOFF": "0", + }, + ) @dataclass(frozen=True) From 3606b29c09297332cefeeff842b28f75e1b92443 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 00:59:43 -0700 Subject: [PATCH 89/91] Use finegrained refinement iteration limit in tests --- tests/kernels/conftest.py | 2 +- tests/quantization/block_scaled_tester.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py index 6fab2cbf664f..571e4a374495 100644 --- a/tests/kernels/conftest.py +++ b/tests/kernels/conftest.py @@ -28,7 +28,7 @@ def cheap_finegrained_tuning(): { "FINEGRAINED_AUTOTUNE_TRIALS": "1", "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", - "FINEGRAINED_AUTOTUNE_REFINE": "0", + "FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS": "0", "FINEGRAINED_AUTOTUNE_RUNOFF": "0", }, ): diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py index b948e4ac8878..79c099b5140d 100644 --- a/tests/quantization/block_scaled_tester.py +++ b/tests/quantization/block_scaled_tester.py @@ -58,7 +58,7 @@ def _cheap_tuning(): { "FINEGRAINED_AUTOTUNE_TRIALS": "1", "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", - "FINEGRAINED_AUTOTUNE_REFINE": "0", + "FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS": "0", "FINEGRAINED_AUTOTUNE_RUNOFF": "0", }, ) From 319658d7e722974a6a67b98b135dcdf6cbc05c65 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 05:29:16 -0700 Subject: [PATCH 90/91] Route rowwise FP8 through finegrained and migrate Llama4 experts --- docs/source/en/quantization/blockfp8.md | 11 ++- docs/source/en/quantization/fbgemm_fp8.md | 15 ++-- docs/source/en/quantization/finegrained.md | 3 +- src/transformers/conversion_mapping.py | 12 +++ src/transformers/core_model_loading.py | 23 ++++- .../distributed/sharding_utils.py | 5 +- .../integrations/finegrained/conversions.py | 54 +++++++++-- .../integrations/finegrained/layers.py | 29 ++++-- .../models/llama4/configuration_llama4.py | 9 +- .../models/llama4/modeling_llama4.py | 63 ++++++------- src/transformers/quantizers/auto.py | 11 +++ .../quantizers/quantizer_blockfp8.py | 75 +++++++++++++++- src/transformers/utils/quantization_config.py | 76 ++++++++++++---- .../finegrained/test_conversions.py | 90 +++++++++++++++++++ tests/integrations/finegrained/test_layers.py | 15 +++- tests/integrations/finegrained/utils.py | 8 +- tests/models/llama4/test_modeling_llama4.py | 68 ++++++++++++++ .../finegrained/test_config_readers.py | 45 +++++++++- .../utils/test_distributed_sharding_utils.py | 35 ++++++-- 19 files changed, 541 insertions(+), 106 deletions(-) diff --git a/docs/source/en/quantization/blockfp8.md b/docs/source/en/quantization/blockfp8.md index d6b98b20e84d..a392911efe72 100644 --- a/docs/source/en/quantization/blockfp8.md +++ b/docs/source/en/quantization/blockfp8.md @@ -17,7 +17,8 @@ rendered properly in your Markdown viewer. # Block FP8 Block FP8 stores weights in 8-bit floating point (`float8_e4m3fn`) with one scale per 2D block of the weight matrix -(`weight_block_size=(128, 128)` by default), or one scale per tensor (`weight_block_size=None`). Activations are +(`weight_block_size=(128, 128)` by default), one scale per tensor (`weight_block_size=None`), or one scale per +output channel (`weight_block_size="row"`). Rowwise FP8 uses dynamic per-token activation scales. Activations are quantized to FP8 too, per token and block at run time (`activation_scheme="dynamic"`), or against a calibrated scale shipped with the checkpoint (`activation_scheme="static"`). @@ -58,6 +59,9 @@ print(tokenizer.decode(output[0], skip_special_tokens=True)) Checkpoints exported by [compressed-tensors](./compressed_tensors) in block FP8 (`format="float-quantized"`) load as [`BlockFp8Config`] as well. +Rowwise FP8 from FBGEMM, ModelOpt (`FP8_PER_CHANNEL_PER_TOKEN`), and compressed-tensors (channel weights with +token activations) uses the same finegrained runtime. Producer-specific key and scale layouts are converted on load. + ## Quantize on the fly Pass a [`BlockFp8Config`] to [`~PreTrainedModel.from_pretrained`] to quantize a full-precision checkpoint as it loads. @@ -91,7 +95,10 @@ model = AutoModelForCausalLM.from_pretrained("llama-3.1-8b-fp8", device_map="aut ## Options -- `weight_block_size` is the `(rows, columns)` block each weight scale covers; `None` gives one scale per tensor. +- `weight_block_size` is the `(rows, columns)` block each weight scale covers; `None` gives one scale per tensor, + and `"row"` gives one FP32 scale per output channel. +- `activation_scale_ub` caps activation amax for dynamic rowwise W8A8; values beyond it saturate. FBGEMM checkpoints + retain their configured bound (1200 by default). Other rowwise formats leave it unset. - `activation_scheme="static"` quantizes activations against the calibrated scale a checkpoint ships (its `input_scale`), instead of deriving one per token at run time. - `scale_fmt="ue8m0"` holds the block scales as powers of two (`float8_e8m0fnu`, with a `uint8` exponent fallback when the dtype is unavailable), as DeepSeek-V4 ships them. diff --git a/docs/source/en/quantization/fbgemm_fp8.md b/docs/source/en/quantization/fbgemm_fp8.md index 503eb87dd087..a7347fcefcde 100644 --- a/docs/source/en/quantization/fbgemm_fp8.md +++ b/docs/source/en/quantization/fbgemm_fp8.md @@ -16,20 +16,23 @@ rendered properly in your Markdown viewer. # FBGEMM -[FBGEMM (Facebook GEneral Matrix Multiplication)](https://github.com/pytorch/FBGEMM) is a low-precision matrix multiplication library for small batch sizes and support for accuracy-loss minimizing techniques such as row-wise quantization and outlier-aware quantization. With FBGEMM, quantize a models weights to 8-bits/channel and the activations to 8-bits/token (also known as fp8 or w8a8). +[FBGEMM](https://github.com/pytorch/FBGEMM) FP8 checkpoints store weights in FP8 with one scale per output channel +and quantize activations per token (W8A8). Transformers reads this recipe into +[`BlockFp8Config`] with `weight_block_size="row"` and executes it through the +[finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels), including Llama4 experts. +The producer's `activation_scale_ub` is preserved: it caps activation amax before FP8 quantization. > [!TIP] > You need a GPU with [compute capability 9+](https://developer.nvidia.com/cuda-gpus#collapseOne) like a H100. -Install the FBGEMM_GPU package with the command below to ensure you have the latest version. +Install the kernel loader and Accelerate: ```bash -pip install --upgrade accelerate fbgemm-gpu torch +pip install --upgrade accelerate kernels ``` -If you're having installation issues, try installing the [nightly release](https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries:~:text=found%20here.-,Install%20the%20FBGEMM_GPU%20Package,-Install%20through%20PyTorch). - -Create a [`FbgemmFp8Config`] and pass it to [`~PreTrainedModel.from_pretrained`] to quantize a model to fp8. +[`FbgemmFp8Config`] remains a checkpoint-compatible entry point and normalizes to this rowwise configuration. +Pass it to [`~PreTrainedModel.from_pretrained`] to quantize a model to FP8. ```py from transformers import FbgemmFp8Config, AutoModelForCausalLM diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md index a04371df6934..02c57411547a 100644 --- a/docs/source/en/quantization/finegrained.md +++ b/docs/source/en/quantization/finegrained.md @@ -69,7 +69,8 @@ These load as a [`FineGrainedConfig`] without one being passed: reads as a plan with one entry per group target. - ModelOpt `quant_algo="MIXED_PRECISION"` reads the `quantized_layers` mapping. Each concrete module name covers that module and its descendants; modules absent from the mapping stay in full precision. Supported algorithms - are `FP8` (calibrated per-tensor FP8), `FP8_PB_WO` (block-FP8 weights with full-precision activations), `NVFP4`, `MXFP8`, `MXFP4`, and the + are `FP8` (calibrated per-tensor FP8), `FP8_PER_CHANNEL_PER_TOKEN` (rowwise FP8), + `FP8_PB_WO` (block-FP8 weights with full-precision activations), `NVFP4`, `MXFP8`, `MXFP4`, and the `W4A16_NVFP4` / `W4A16_MXFP4` weight-only variants. Fields can appear directly in the quantization config or inside its `quantization` object. Unsupported algorithms or config groups raise an error. diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 199b5d17f2b6..1a3fbb8c9f2d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1963,6 +1963,18 @@ def _build_checkpoint_conversion_mapping(): "UltrabertForQuestionAnswering": [ WeightRenaming(r"^ultrabert\.", "model."), ], + "llama4_text": [ + WeightConverter( + source_patterns="feed_forward.experts.gate_up_proj", + target_patterns="feed_forward.experts.gate_up_proj", + operations=[Transpose(1, 2)], + ), + WeightConverter( + source_patterns="feed_forward.experts.down_proj", + target_patterns="feed_forward.experts.down_proj", + operations=[Transpose(1, 2)], + ), + ], } mapping["Tipsv2DptForNormalEstimation"] = [ diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index f0c2bf3ab6e6..56ec9e48c02f 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1742,7 +1742,28 @@ def convert_and_load_state_dict_in_model( param_device = get_device(device_map, renamed_key, valid_torch_device=True) sharding_op = None if is_dtensor(empty_param): - sharding_op = DtensorShardOperation(empty_param) + source_axes = list(range(empty_param.ndim)) + if ( + isinstance(mapping, WeightConverter) + and mapping.operations + and isinstance(mapping.operations[0], Transpose) + and not mapping.operations[0].check_dims + ): + source_shape = tensor.shape if isinstance(tensor, torch.Tensor) else tensor.get_shape() + transpose_count = 0 + for op in mapping.operations: + if not isinstance(op, Transpose) or op.check_dims: + break + transpose_count += 1 + preserves_axes = all( + getattr(op, "preserves_sharding_axes", lambda *_: False)(model, renamed_key) + for op in mapping.operations[transpose_count:] + ) + if len(source_shape) == empty_param.ndim and preserves_axes: + # Read the source axes corresponding to the target placements, then transpose the local shard. + for op in mapping.operations[:transpose_count]: + source_axes[op.dim0], source_axes[op.dim1] = source_axes[op.dim1], source_axes[op.dim0] + sharding_op = DtensorShardOperation(empty_param, source_axis_order=tuple(source_axes)) # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations # directly on the device, as it will completely blow up the memory during the ops memory spike. So defer to "cpu", then diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 666fdb82b10c..0faa21b28588 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -88,10 +88,11 @@ class DtensorShardOperation: we skip it. Later, `MergeModulelist` will stack the owned expert we kept to create the rank's local shard """ - def __init__(self, param: DTensor): + def __init__(self, param: DTensor, source_axis_order: tuple[int, ...] | None = None): self.device_mesh = param.device_mesh self.placements = tuple(param.placements) self.param_ndim = param.ndim + self.source_axis_order = source_axis_order or tuple(range(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) @@ -129,7 +130,7 @@ def shard_tensor( for mesh_dim, placement in dim_placements: sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() - dim_idx = self._normalize_param_dim(placement.dim) + dim_idx = self.source_axis_order[self._normalize_param_dim(placement.dim)] planned_ops_by_dim[dim_idx].append((placement, rank, world_size)) # prepare the slices to fetch on disk for each tensor dimension. diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py index b1fc0a52acd6..5b8675eee563 100644 --- a/src/transformers/integrations/finegrained/conversions.py +++ b/src/transformers/integrations/finegrained/conversions.py @@ -131,6 +131,12 @@ class FineGrainedDeserialize(_FineGrainedOp): quantizing the full-precision rows. Biases only need their row layout. Saving reverses the layouts before splitting or renaming the tensors. """ + def preserves_sharding_axes(self, model, full_layer_name): + # Row-FP8 only reorders gate/up rows within the existing packed-shard intervals. + # Other formats can reshape packed weights or swizzle scale grids across axes. + _, module = finegrained_module(model, full_layer_name) + return module is not None and module.quantization_config.weight_block_size == "row" + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): if not self.inverse: input_dict = keyed_by_target(input_dict, target_patterns) @@ -241,6 +247,16 @@ def reshape(tensor): } +class FineGrainedRowScales(_FineGrainedOp): + """ModelOpt exports channel scales as vectors; the GEMMs hold a column grid.""" + + def convert(self, input_dict, **kwargs): + return { + key: single_tensor(value).squeeze(-1) if self.inverse else single_tensor(value).unsqueeze(-1) + for key, value in input_dict.items() + } + + 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 @@ -453,7 +469,9 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc held = getattr(module, scale_name) if module is not None else None format_spec = get_weight_format(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: + if module is not None and module.quantization_config.weight_block_size == "row": + block = (1, value.shape[-1]) + elif 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 @@ -464,7 +482,13 @@ def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torc 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]) + block = ( + (1, value.shape[-1]) + if config_block == "row" + else tuple(config_block) + if config_block + else (value.shape[-2], value.shape[-1]) + ) ue8m0 = ( held.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) if held is not None @@ -581,6 +605,10 @@ 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=None, inverse=False, *, transpose_scales=False): + super().__init__(hf_quantizer, inverse) + self.transpose_scales = transpose_scales + def _scale_pattern_for(self, weight_pattern: str) -> str: anchor = "(?=$)" if weight_pattern.endswith("(?=$)") else "$" if weight_pattern.endswith("$") else "" base = weight_pattern.removesuffix(anchor) if anchor else weight_pattern @@ -646,7 +674,7 @@ def _dequantize_one( global_scale: torch.Tensor | None = None, *, reciprocal_globals: bool = False, - weight_block_size: tuple[int, int] | None = None, + weight_block_size: tuple[int, int] | str | 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) @@ -656,15 +684,25 @@ def _dequantize_one( else: quantized_fp32 = quantized.to(torch.float32) rows, cols = quantized_fp32.shape[-2:] + if weight_block_size is None: + weight_block_size = getattr(self._config_for(None), "weight_block_size", None) + if weight_block_size == "row": + if scales.ndim < quantized_fp32.ndim: + scales = scales.unsqueeze(-1) + if self.transpose_scales: + scales = scales.transpose(-1, -2) # A per-tensor scale has no grid: one block over the whole matrix. scale_rows, scale_cols = scales.shape[-2:] if scales.ndim >= 2 else (1, 1) # Ceil-rounded grids do not uniquely identify a block size. Use this module's config, # including in mixed plans; infer from the grid only when no configured block matches. - if weight_block_size is None: - weight_block_size = getattr(self._config_for(None), "weight_block_size", None) - if weight_block_size and tuple(_cdiv(size, block) for size, block in zip((rows, cols), weight_block_size)) == ( - scale_rows, - scale_cols, + if ( + weight_block_size + and weight_block_size != "row" + and tuple(_cdiv(size, block) for size, block in zip((rows, cols), weight_block_size)) + == ( + scale_rows, + scale_cols, + ) ): block_m, block_n = weight_block_size elif rows % scale_rows or cols % scale_cols: diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py index b6f2578d993a..74352440e4ed 100644 --- a/src/transformers/integrations/finegrained/layers.py +++ b/src/transformers/integrations/finegrained/layers.py @@ -198,15 +198,16 @@ def get_weight_format(weight_format: str) -> WeightFormat: 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]: + block_size: tuple[int, int] | str | None = None, +) -> tuple[WeightFormat, torch.dtype, tuple[int, int | None] | 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.""" format_spec = get_weight_format(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 format_spec, scale_dtype, format_spec.scale_group if format_spec.scale_group is not None else block_size + scale_group = (1, None) if isinstance(block_size, str) else block_size + return format_spec, scale_dtype, format_spec.scale_group if format_spec.scale_group is not None else scale_group def _set_optional_parameter(module: nn.Module, name: str, tensor: torch.Tensor | None) -> None: @@ -221,7 +222,7 @@ def _alloc_expert_proj( proj_in: int, format_spec: WeightFormat, scale_dtype: torch.dtype, - scale_group: tuple[int, int] | None, + scale_group: tuple[int, int | None] | None, min_scale_out: int = 1, swizzled: bool = False, ) -> tuple[nn.Parameter, nn.Parameter]: @@ -235,7 +236,8 @@ def _alloc_expert_proj( 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]) + rows = max(_cdiv(proj_out, scale_group[0]), min_scale_out) + cols = _cdiv(proj_in, scale_group[1] or proj_in) 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) @@ -275,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, + activation_scale_ub: float | None = None, ) -> torch.Tensor: """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. @@ -346,6 +349,7 @@ def finegrained_linear( output_dtype=input.dtype, a_global_scale=input_global_scale, b_global_scale=weight_global_scale, + activation_scale_ub=activation_scale_ub, ) output = output.reshape(*original_shape[:-1], output.shape[-1]) if bias is not None: @@ -386,7 +390,8 @@ def activation_scheme(self) -> str: @property def block_size(self) -> tuple[int, int] | None: - return self.quantization_config.weight_block_size + block = self.quantization_config.weight_block_size + return None if isinstance(block, str) else block class FineGrainedLinear(_FineGrainedModule, nn.Linear): @@ -407,7 +412,7 @@ def __init__( # 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 format_spec, sf_dtype, scale_group = resolve_weight_format( - self.weight_format, self.quantization_config.scale_fmt, self.block_size + self.weight_format, self.quantization_config.scale_fmt, self.quantization_config.weight_block_size ) # the format's own dtype both decides whether there is a second level and says what it # is held in @@ -428,7 +433,9 @@ def __init__( 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) + else torch.empty( + _cdiv(out_features, scale_group[0]), _cdiv(in_features, scale_group[1] or in_features), dtype=sf_dtype + ) ) self.weight_scale_inv = nn.Parameter(scale, requires_grad=scale.dtype.is_floating_point) @@ -464,6 +471,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: allow_deepgemm=not self._deepgemm_disabled, weight_global_scale=self.weight_global_scale, input_global_scale=self.activation_scale if two_level else None, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), activation_format=self.activation_format, ) @@ -628,6 +636,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: activation_format=self.activation_format, a_global_scale=self.activation_scale if two_level else None, b_global_scale=self.weight_global_scale, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), ) y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) if self.has_bias: @@ -698,7 +707,7 @@ def __init__( self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) activation_format = self.activation_format format_spec, scale_dtype, scale_group = resolve_weight_format( - self.weight_format, self.quantization_config.scale_fmt, self.block_size + self.weight_format, self.quantization_config.scale_fmt, self.quantization_config.weight_block_size ) self.global_scale_dtype = format_spec.global_scale_dtype @@ -855,6 +864,7 @@ def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torc weight_global_scale=weight_global, input_global_scale=input_global, activation_format=self.activation_format, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), ) # Separately calibrated gate/up grids interleave per block, rather than per weight row: # run each half with its own blocks, then interleave the output columns. A single scale row @@ -924,6 +934,7 @@ def _moe_operands(self, kernel) -> dict: "swiglu_alpha": self.swiglu_alpha, "swiglu_limit": self.swiglu_limit, "activation_format": self.activation_format, + "activation_scale_ub": getattr(self.quantization_config, "activation_scale_ub", None), # 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, diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 020bceaa7249..5c33c4dc6e60 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -117,8 +117,8 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.feed_forward.shared_expert.gate_proj": "colwise", "layers.*.feed_forward.shared_expert.up_proj": "colwise", "layers.*.feed_forward.shared_expert.down_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise", # row because not linear - "layers.*.feed_forward.experts.down_proj": "colwise", # col because not linear + "layers.*.feed_forward.experts.gate_up_proj": "packed_colwise", + "layers.*.feed_forward.experts.down_proj": "rowwise", "layers.*.feed_forward.experts": "moe_tp_experts", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", @@ -129,8 +129,9 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear - "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear + "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", + "layers.*.feed_forward.experts.down_proj": "grouped_gemm", + "layers.*.feed_forward.experts": "moe_tp_experts", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", "layers.*.feed_forward.down_proj": "rowwise", diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index 81ea67db7260..7d48966f2336 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -26,6 +26,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations import use_experts_implementation +from ...integrations.moe import batched_mm_experts_forward from ...masking_utils import create_causal_mask, create_chunked_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -51,36 +53,24 @@ logger = logging.get_logger(__name__) +@use_experts_implementation class Llama4TextExperts(nn.Module): def __init__(self, config: Llama4TextConfig): super().__init__() self.num_experts = config.num_local_experts - self.intermediate_size = config.intermediate_size - self.hidden_size = config.hidden_size - self.expert_dim = self.intermediate_size - self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim)) - self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size))) + self.intermediate_dim = config.intermediate_size + self.hidden_dim = config.hidden_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """ - This should really not be run on a single machine, as we are reaching compute bound: - - the inputs are expected to be "sorted" per expert already. - - the weights are viewed with another dim, to match num_expert, 1, shape * num_tokens, shape - - Args: - hidden_states (torch.Tensor): (batch_size * token_num, hidden_size) - selected_experts (torch.Tensor): (batch_size * token_num, top_k) - routing_weights (torch.Tensor): (batch_size * token_num, top_k) - Returns: - torch.Tensor - """ - hidden_states = hidden_states.view(self.gate_up_proj.shape[0], -1, self.hidden_size) - gate_up = torch.bmm(hidden_states, self.gate_up_proj) - gate, up = gate_up.chunk(2, dim=-1) # not supported for DTensors - next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj) - next_states = next_states.view(-1, self.hidden_size) - return next_states + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + return batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) # Phi3MLP @@ -145,13 +135,11 @@ def __init__(self, config): def forward(self, hidden_states): router_logits = super().forward(hidden_states) - router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=1) - router_scores = torch.full_like(router_logits, float("-inf")).scatter_(1, router_indices, router_top_value) - router_scores = torch.nn.functional.sigmoid(router_scores.float()).to(router_scores.dtype) - return router_scores, router_logits + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) + router_scores = torch.sigmoid(router_top_value.float()).to(router_logits.dtype) + return router_logits, router_scores, router_indices -# @use_kernel_forward_from_hub("Llama4TextMoe") class Llama4TextMoe(nn.Module): def __init__(self, config): super().__init__() @@ -164,13 +152,14 @@ def __init__(self, config): def forward(self, hidden_states): hidden_states = hidden_states.reshape(-1, self.hidden_dim) - router_scores, router_logits = self.router(hidden_states) - routed_in = hidden_states.repeat(router_scores.shape[1], 1) - routed_in = routed_in * router_scores.transpose(0, 1).reshape(-1, 1) - routed_out = self.experts(routed_in) - out = self.shared_expert(hidden_states) - out.add_(routed_out.reshape(router_scores.shape[1], -1, routed_out.shape[-1]).sum(dim=0)) - return out, router_logits + router_logits, router_scores, router_indices = self.router(hidden_states) + # Llama4 weights each selected expert's input, before the nonlinear GLU. + routed_in = (hidden_states[:, None, :] * router_scores[:, :, None]).reshape(-1, self.hidden_dim) + routed_out = self.experts( + routed_in, router_indices.reshape(-1, 1), torch.ones_like(router_scores).reshape(-1, 1) + ) + routed_out = routed_out.reshape(-1, self.top_k, self.hidden_dim).sum(dim=1) + return self.shared_expert(hidden_states) + routed_out, router_logits # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Llama4Text @@ -480,7 +469,7 @@ class Llama4TextModel(Llama4PreTrainedModel): _can_record_outputs = { "attentions": Llama4TextAttention, "hidden_states": Llama4TextDecoderLayer, - "router_logits": OutputRecorder(Llama4Router, index=1), + "router_logits": OutputRecorder(Llama4Router, index=0), } def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 79ca2fde05bf..be8605c0c512 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -191,6 +191,13 @@ def from_dict(cls, quantization_config_dict: dict): # a modelopt export reads as its one format's config, or a `FineGrainedConfig` when it mixes them if quant_method == QuantizationMethod.MODELOPT: return read_modelopt(quantization_config_dict) + if quant_method == QuantizationMethod.FBGEMM_FP8: + return BlockFp8Config( + weight_block_size="row", + activation_scale_ub=quantization_config_dict.get("activation_scale_ub", 1200.0), + modules_to_not_convert=quantization_config_dict.get("modules_to_not_convert"), + checkpoint_format="fbgemm", + ) # so does a compressed-tensors checkpoint in formats these kernels serve if quant_method == QuantizationMethod.COMPRESSED_TENSORS: if (config := read_compressed_tensors(quantization_config_dict)) is not None: @@ -230,6 +237,8 @@ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwar # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = AutoQuantizationConfig.from_dict(quantization_config) + elif isinstance(quantization_config, FbgemmFp8Config): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config.to_dict()) # the deprecated config keeps its own, frozen integration if isinstance(quantization_config, FineGrainedFP8Config): from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer @@ -472,6 +481,8 @@ def finegrained_quantization_config(quantization_config, model_config): """The config a checkpoint really describes, with its model config in hand: block-FP8 whose checkpoint declares FP4 experts (`expert_dtype`, DeepSeek-V4 and V4.1) as block-FP8 over W4A4 MXFP4 experts. Any other config unchanged.""" + if isinstance(quantization_config, FbgemmFp8Config): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config.to_dict()) if not isinstance(quantization_config, BlockFp8Config): return quantization_config if _declared_expert_dtype(model_config) != "fp4": diff --git a/src/transformers/quantizers/quantizer_blockfp8.py b/src/transformers/quantizers/quantizer_blockfp8.py index 70d328cd30e6..ab4dd4edc7cc 100644 --- a/src/transformers/quantizers/quantizer_blockfp8.py +++ b/src/transformers/quantizers/quantizer_blockfp8.py @@ -32,11 +32,38 @@ def plan(self): embeddings = self.quantization_config.modules_to_convert or [] return {**dict.fromkeys(map(embedding_pattern, embeddings), self.quantization_config), **super().plan()} + def get_param_ops(self, model, param_name): + from ..integrations.finegrained.conversions import FineGrainedRowScales + + ops = super().get_param_ops(model, param_name) + if ( + self.quantization_config.weight_block_size == "row" + and self.quantization_config.checkpoint_format == "modelopt" + and param_name.endswith("_scale_inv") + ): + ops = [FineGrainedRowScales(self), *ops] + return ops + def format_parallel_plan(self, model, plan): from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard - from ..integrations.finegrained import FineGrainedEmbedding + from ..integrations.finegrained import FineGrainedEmbedding, FineGrainedExperts, FineGrainedLinear for name, module in model.named_modules(): + if ( + self.quantization_config.weight_block_size == "row" + and isinstance(module, FineGrainedExperts) + and self.member_of(module) is not None + and _get_parameter_plan(f"{name}.down_proj", plan) in ("rowwise", "packed_rowwise") + ): + plan.pop(f"{replace_layer_number_by_wildcard(name)}.down_proj_scale_inv", None) + if ( + self.quantization_config.weight_block_size == "row" + and isinstance(module, FineGrainedLinear) + and self.member_of(module) is not None + and _get_parameter_plan(f"{name}.weight", plan) + in ("rowwise", "rowwise_rep", "rowwise_split_input", "packed_rowwise") + ): + plan[f"{replace_layer_number_by_wildcard(name)}.weight_scale_inv"] = "replicate" if ( isinstance(module, FineGrainedEmbedding) and self.member_of(module) is not None @@ -45,9 +72,44 @@ def format_parallel_plan(self, model, plan): plan.setdefault(f"{replace_layer_number_by_wildcard(name)}.weight_scale", "replicate") def format_weight_conversions(self, converters, model): - from ..core_model_loading import MergeModulelist + from ..core_model_loading import MergeModulelist, Transpose from ..integrations.finegrained import FineGrainedExperts - from ..integrations.finegrained.conversions import FineGrainedScalarScales + from ..integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedScalarScales + + producer = self.quantization_config.checkpoint_format + if self.quantization_config.weight_block_size == "row" and producer in ("fbgemm", "modelopt"): + for converter in converters: + dequantize = next((op for op in converter.operations if isinstance(op, FineGrainedDequantize)), None) + if dequantize is None or not any( + isinstance(op, Transpose) and not op.check_dims and {op.dim0, op.dim1} == {1, 2} + for op in converter.operations + ): + continue + reverse = converter.reverse_transform() + if any( + reverse.rename_source_key(name)[1] is not None + and dequantize._config_for(name) is self.quantization_config + and (producer == "modelopt" or name.endswith(".down_proj")) + for name, _ in model.named_parameters() + ): + # Dequantization precedes the BMM transpose; these producers' scales already name output rows. + converter.operations = [ + FineGrainedDequantize(op.hf_quantizer, transpose_scales=True) if op is dequantize else op + for op in converter.operations + ] + if producer == "fbgemm" or (producer == "modelopt" and self.quantization_config.weight_block_size == "row"): + # ModelOpt leaves BMM channel vectors untransposed; FBGEMM's down grid is already rowwise. + scales = [ + f"{name}.{param}" + for name, module in model.named_modules() + if self.member_of(module) is not None + for param, _ in module.named_parameters(recurse=False) + if param.endswith("_scale_inv") and (producer == "modelopt" or param == "down_proj_scale_inv") + ] + for converter in converters: + reverse = converter.reverse_transform() + if any(reverse.rename_source_key(name)[1] is not None for name in scales): + converter.operations = [op for op in converter.operations if not isinstance(op, Transpose)] if ( self.quantization_config.checkpoint_format != "modelopt" @@ -70,9 +132,14 @@ def format_weight_conversions(self, converters, model): def format_conversions(self, dequantize): conversions = [] - if self.quantization_config.checkpoint_format in ("compressed-tensors", "modelopt"): + if self.quantization_config.checkpoint_format in ("compressed-tensors", "modelopt", "fbgemm"): # its block scale ships as `weight_scale`, and rides its weight's converter once renamed conversions.append(self._block_scale_rename()) + if self.quantization_config.checkpoint_format in ("fbgemm", "modelopt"): + from ..core_model_loading import WeightRenaming + + suffix = "weight_scale" if self.quantization_config.checkpoint_format == "modelopt" else "scale" + conversions.append(WeightRenaming(rf"(gate_up_proj|up_proj|down_proj)_{suffix}(?=$)", r"\1_scale_inv")) if not dequantize and self.quantization_config.activation_scheme == "static": conversions += self._activation_scale_conversions() return conversions diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 353492d5f30c..2a6130a84bf1 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -18,6 +18,7 @@ import fnmatch import importlib.metadata import json +import math import os import warnings from dataclasses import dataclass @@ -1296,7 +1297,7 @@ def is_quantization_compressed(self): class FbgemmFp8Config(QuantizationConfigMixin): """ This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using fbgemm fp8 quantization. + loaded from FBGEMM rowwise FP8 metadata. Execution uses the finegrained kernels. Args: activation_scale_ub (`float`, *optional*, defaults to 1200.0): @@ -1750,7 +1751,7 @@ class BlockScaledConfig(QuantizationConfigMixin): activation_formats: tuple[str | None, ...] = (None,) default_activation_format: str | None = None # fixed by the MX and NVFP4 formats; block-FP8 sets them - weight_block_size: tuple[int, int] | None = None + weight_block_size: tuple[int, int] | str | None = None scale_fmt: str = "float" activation_scheme: str = "dynamic" # the values those fields take in this format, which a checkpoint's `config.json` may spell out @@ -1813,9 +1814,9 @@ def post_init(self): # one spelling for the weights' own activation format if self.activation_format == self.method.value: self.activation_format = None - if self.checkpoint_format not in (None, "compressed-tensors", "modelopt"): + if self.checkpoint_format not in (None, "compressed-tensors", "modelopt", "fbgemm"): raise ValueError( - f"checkpoint_format must be None, 'compressed-tensors', or 'modelopt'; got {self.checkpoint_format!r}" + f"checkpoint_format must be None, 'compressed-tensors', 'modelopt', or 'fbgemm'; got {self.checkpoint_format!r}" ) def get_loading_attributes(self): @@ -1834,8 +1835,11 @@ class BlockFp8Config(BlockScaledConfig): per-tensor scale without a block), as DeepSeek-V3 and Qwen3 ship them; W8A8 with dynamic or calibrated activations. Args: - weight_block_size (`tuple[int, int]`, *optional*, defaults to `(128, 128)`): - The block each weight scale covers; `None` for one scale per tensor. + weight_block_size (`tuple[int, int] | str`, *optional*, defaults to `(128, 128)`): + The block each weight scale covers; `None` for one scale per tensor, or `"row"` for one scale + per output channel. Dynamic FP8 activations then use one scale per token. + activation_scale_ub (`float`, *optional*): + Maximum activation amax for rowwise FP8 quantization. Values beyond it saturate, as in FBGEMM. scale_fmt (`str`, *optional*, defaults to `"float"`): Storage dtype of the scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`, DeepSeek-V4). activation_scheme (`str`, *optional*, defaults to `"dynamic"`): @@ -1853,18 +1857,32 @@ class BlockFp8Config(BlockScaledConfig): def __init__( self, *, - weight_block_size: tuple[int, int] | None = (128, 128), + weight_block_size: tuple[int, int] | str | None = (128, 128), scale_fmt: str = "float", activation_scheme: str = "dynamic", modules_to_convert: list | None = None, + activation_scale_ub: float | None = None, **kwargs, ): - self.weight_block_size = tuple(weight_block_size) if weight_block_size is not None else None + self.weight_block_size = ( + weight_block_size + if isinstance(weight_block_size, str) or weight_block_size is None + else tuple(weight_block_size) + ) self.scale_fmt = scale_fmt self.activation_scheme = activation_scheme self.modules_to_convert = modules_to_convert + self.activation_scale_ub = activation_scale_ub super().__init__(**kwargs) + def get_loading_attributes(self): + attributes = super().get_loading_attributes() + if self.activation_scale_ub is not None or ( + self._names_activation_format and self.activation_format == "bf16" + ): + attributes["activation_scale_ub"] = self.activation_scale_ub + return attributes + def post_init(self): super().post_init() self.activation_scheme = self.activation_scheme.lower() @@ -1876,10 +1894,24 @@ def post_init(self): ) 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.weight_block_size is not None and ( - len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0 + if self.weight_block_size == "row": + if self.scale_fmt != "float": + raise ValueError("Rowwise FP8 requires FP32 weight scales") + elif isinstance(self.weight_block_size, str) or ( + self.weight_block_size is not None + and (len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0) ): - raise ValueError(f"weight_block_size must be two positive integers; got {self.weight_block_size}") + raise ValueError( + f"weight_block_size must be 'row', None, or two positive integers; got {self.weight_block_size}" + ) + if self.activation_scale_ub is not None and ( + self.weight_block_size != "row" + or not math.isfinite(self.activation_scale_ub) + or self.activation_scale_ub <= 0 + or self.activation_scheme != "dynamic" + or self.activation_format is not None + ): + raise ValueError("activation_scale_ub must be positive and requires dynamic rowwise W8A8 FP8") class Mxfp8Config(BlockScaledConfig): @@ -1967,8 +1999,7 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: } weights = spec.get("weights") or {} quant_method = formats.get((weights.get("num_bits"), weights.get("type"), weights.get("group_size"))) - # one FP8 scale per output channel has no block the kernels tile - if quant_method is None or weights.get("strategy") == "channel": + if quant_method is None: return None # MX block scales are E8M0 (some producers store its exponent bytes as uint8): group-32 values over another # scale dtype are not an MX format @@ -1984,6 +2015,18 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: else: # a class name we cannot resolve to module paths here return None activations = spec.get("input_activations") or {} + rowwise = weights.get("strategy") == "channel" + if rowwise: + if quant_method != "fp8": + return None + if activations and not ( + (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) == (8, "float", None) + and ( + (activations.get("strategy") == "token" and activations.get("dynamic") is True) + or (activations.get("strategy") == "tensor" and activations.get("dynamic") is False) + ) + ): + return None act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) # block-FP8's own activations are FP8 per token, or per 128-wide group as its kernels quantize them; # no activation quantization is weight-only @@ -1997,7 +2040,8 @@ def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: # second-level global (`input_scale`), which the loader carries on its own static = not activations.get("dynamic", True) and act_key == (8, "float", None) options.update( - activation_scheme="static" if static else "dynamic", weight_block_size=weights.get("block_structure") + activation_scheme="static" if static else "dynamic", + weight_block_size="row" if rowwise else weights.get("block_structure"), ) config = config_class(**options) return dict.fromkeys(targets or [".*"], config) @@ -2036,6 +2080,8 @@ def _modelopt_format(spec: dict) -> BlockScaledConfig: } if algo == "FP8": return BlockFp8Config(weight_block_size=None, activation_scheme="static", checkpoint_format="modelopt") + if algo == "FP8_PER_CHANNEL_PER_TOKEN": + return BlockFp8Config(weight_block_size="row", checkpoint_format="modelopt") if algo == "FP8_PB_WO": block = spec.get("group_size", 128) if not isinstance(block, int) or isinstance(block, bool) or block <= 0: @@ -2125,7 +2171,7 @@ def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None: """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, - per-channel, sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's + sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's `checkpoint_format` names the layout its tensors ship in, which a save writes back.""" try: _check_block_scaled_features(config_dict) diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py index 93cda729bac0..a1a160658e49 100644 --- a/tests/integrations/finegrained/test_conversions.py +++ b/tests/integrations/finegrained/test_conversions.py @@ -28,6 +28,7 @@ Chunk, Concatenate, PermuteForRope, + Transpose, WeightConverter, convert_and_load_state_dict_in_model, revert_weight_conversion, @@ -51,6 +52,7 @@ _load_and_save, _model_mapping, _MoeModel, + _per_expert, _quantizer_for, _replace, fake_finegrained_kernel, @@ -430,6 +432,94 @@ def test_per_half_globals_are_not_folded_past_a_glu_nonlinear_in_up(self): @require_torch class FineGrainedLoaderChainTest(unittest.TestCase): + def test_rowwise_fp8_checkpoint_scales_load_and_save_in_their_producer_layout(self): + """Producer vectors/grids and transposed BMM weights share one runtime row-scale layout.""" + + def tensors(rows, cols, vector): + return { + "weight": torch.randn(rows, cols).to(torch.float8_e4m3fn), + "weight_scale": torch.rand((rows,) if vector else (rows, 1)) + 0.1, + } + + for producer, transposed in ( + ("modelopt", False), + ("compressed-tensors", False), + ("fbgemm", False), + ("modelopt", True), + ("fbgemm", True), + ): + with self.subTest(producer=producer, transposed=transposed): + source = _per_expert(lambda rows, cols: tensors(rows, cols, producer == "modelopt")) + checkpoint, mapping = source, _model_mapping() + if transposed: + prefix = "model.layers.0.mlp.experts" + checkpoint = {key: value for key, value in source.items() if not key.startswith(prefix)} + mapping = [] + for fused, projections in ( + ("gate_up_proj", ("gate_proj", "up_proj")), + ("down_proj", ("down_proj",)), + ): + weight = torch.stack( + [ + torch.cat([source[f"{prefix}.{expert}.{proj}.weight"] for proj in projections]) + for expert in range(E) + ] + ) + scale = torch.stack( + [ + torch.cat( + [ + source[f"{prefix}.{expert}.{proj}.weight_scale"].reshape(-1, 1) + for proj in projections + ] + ) + for expert in range(E) + ] + ) + key = f"{prefix}.{fused}" + checkpoint[key] = weight.transpose(1, 2).contiguous() + suffix = "weight_scale" if producer == "modelopt" else "scale" + checkpoint[f"{key}_{suffix}"] = ( + scale.squeeze(-1) + if producer == "modelopt" + else scale.transpose(1, 2) + if fused == "gate_up_proj" + else scale + ) + mapping.append(WeightConverter(key, key, [Transpose(1, 2)])) + config = BlockFp8Config(weight_block_size="row", checkpoint_format=producer) + model, info, saved = _load_and_save(checkpoint, config, mapping) + self.assertEqual( + (info.missing_keys, info.unexpected_keys, info.mismatched_keys, info.conversion_errors), + (set(), set(), set(), {}), + ) + experts = model.model.layers[0].mlp.experts + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (E, 2 * I, 1)) + self.assertEqual(experts.down_proj_scale_inv.shape, (E, H, 1)) + self.assertEqual(set(saved), set(checkpoint)) + for key, value in checkpoint.items(): + torch.testing.assert_close(saved[key].float(), value.float(), rtol=0, atol=0) + dequantized, info, _ = _load_and_save( + checkpoint, + BlockFp8Config(weight_block_size="row", checkpoint_format=producer, dequantize=True), + mapping, + ) + self.assertEqual((info.missing_keys, info.unexpected_keys, info.conversion_errors), (set(), set(), {})) + for fused, projections in (("gate_up_proj", ("gate_proj", "up_proj")), ("down_proj", ("down_proj",))): + expected = torch.stack( + [ + torch.cat( + [ + source[f"model.layers.0.mlp.experts.{expert}.{proj}.weight"].float() + * source[f"model.layers.0.mlp.experts.{expert}.{proj}.weight_scale"].reshape(-1, 1) + for proj in projections + ] + ) + for expert in range(E) + ] + ) + torch.testing.assert_close(getattr(dequantized.model.layers[0].mlp.experts, fused), expected) + """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.""" diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py index 02aa2518aadb..dc9ab8e7f4b7 100644 --- a/tests/integrations/finegrained/test_layers.py +++ b/tests/integrations/finegrained/test_layers.py @@ -485,6 +485,7 @@ def test_grouped_linear_hands_over_its_stored_weight_and_activation_scales(self) for config, packed in ( (NVFP4Config(), True), (BlockFp8Config(weight_block_size=(16, 16), activation_scheme="static"), False), + (BlockFp8Config(weight_block_size="row", activation_scale_ub=1200.0), False), ): with self.subTest(weight_format=config.method.value): linear = FineGrainedGroupedLinear(64, 32, 2, quantization_config=config) @@ -497,6 +498,9 @@ def test_grouped_linear_hands_over_its_stored_weight_and_activation_scales(self) two_level = linear.weight_global_scale is not None self.assertIs(call.args[2], None if two_level else linear.activation_scale) self.assertIs(call.kwargs["a_global_scale"], linear.activation_scale if two_level else None) + self.assertEqual(call.kwargs["activation_scale_ub"], getattr(config, "activation_scale_ub", None)) + if config.weight_block_size == "row": + self.assertEqual(call.args[3].shape, (2, 16, 1)) @require_torch @@ -538,10 +542,10 @@ class FineGrainedParallelPlanTest(unittest.TestCase): "layers.*.self_attn.q_proj": "colwise", } - def _planned(self, impl=None, raw=None): + def _planned(self, impl=None, raw=None, config=None): from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer - quantizer = BlockFp8HfQuantizer(BlockFp8Config()) + quantizer = BlockFp8HfQuantizer(config or BlockFp8Config()) model = torch.nn.Module() model.config = SimpleNamespace(_experts_implementation=impl) model.layers = torch.nn.ModuleList([torch.nn.Module() for _ in range(4)]) @@ -549,7 +553,7 @@ def _planned(self, impl=None, raw=None): for layer in model.layers: layer.mlp = torch.nn.Module() layer.mlp.experts = FineGrainedExperts( - _Cfg(), quantization_config=NVFP4Config(activation_format="nvfp4"), has_bias=True + _Cfg(), quantization_config=config or NVFP4Config(activation_format="nvfp4"), has_bias=True ) plan = dict(raw or self.BASE_EP) quantizer.update_parallel_plan(model, plan) @@ -631,6 +635,11 @@ def style(name): _get_parameter_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), "moe_experts_colwise", ) + rowwise = self._planned(raw=base, config=BlockFp8Config(weight_block_size="row")) + self.assertEqual( + _get_parameter_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", rowwise), "moe_experts_packed_colwise" + ) + self.assertEqual(_get_parameter_plan("layers.3.mlp.experts.down_proj_scale_inv", rowwise), "moe_tp_experts") 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 diff --git a/tests/integrations/finegrained/utils.py b/tests/integrations/finegrained/utils.py index 364afe4d6878..a38456903a94 100644 --- a/tests/integrations/finegrained/utils.py +++ b/tests/integrations/finegrained/utils.py @@ -15,6 +15,7 @@ checkpoint-config builders, the mocked kernel bundle, and the tiny MoE checkpoint the conversion chains load and save back.""" +from copy import deepcopy from dataclasses import dataclass from types import SimpleNamespace from unittest import mock @@ -203,7 +204,7 @@ def _conversion_model(quantizer): return model -def _load_and_save(checkpoint, quantization_config): +def _load_and_save(checkpoint, quantization_config, weight_mapping=None): """`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() @@ -214,12 +215,15 @@ 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(), model), + weight_mapping=quantizer.update_weight_conversions( + _model_mapping() if weight_mapping is None else deepcopy(weight_mapping), model + ), hf_quantizer=quantizer, ) info, _ = convert_and_load_state_dict_in_model( model, {k: v.clone() for k, v in checkpoint.items()}, load_config ) + quantizer._process_model_after_weight_loading(model) model.hf_quantizer = quantizer saved = revert_weight_conversion(model, model.state_dict()) return model, info, saved diff --git a/tests/models/llama4/test_modeling_llama4.py b/tests/models/llama4/test_modeling_llama4.py index cd07907fe892..1a7d4516bdf1 100644 --- a/tests/models/llama4/test_modeling_llama4.py +++ b/tests/models/llama4/test_modeling_llama4.py @@ -13,11 +13,15 @@ # limitations under the License. """Testing suite for the PyTorch Llama4 model.""" +import tempfile import unittest +from parameterized import parameterized + from transformers import is_torch_available from transformers.testing_utils import ( Expectations, + require_torch, require_torch_large_accelerator, slow, torch_device, @@ -29,13 +33,77 @@ if is_torch_available(): import torch + from safetensors.torch import load_file from transformers import ( + Llama4ForCausalLM, Llama4ForConditionalGeneration, Llama4Processor, + Llama4TextConfig, ) +@require_torch +class Llama4TextMoeTest(unittest.TestCase): + def get_config(self, top_k=1): + return Llama4TextConfig( + vocab_size=97, + hidden_size=32, + intermediate_size=48, + intermediate_size_mlp=48, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_local_experts=4, + num_experts_per_tok=top_k, + max_position_embeddings=64, + attention_chunk_size=16, + ) + + @parameterized.expand([1, 2]) + def test_input_router_weights(self, top_k): + config = self.get_config(top_k) + model = Llama4ForCausalLM(config).to(torch_device).eval() + moe = model.model.layers[0].feed_forward + hidden_states = torch.randn(6, config.hidden_size, device=torch_device) + with torch.no_grad(): + output, logits = moe(hidden_states) + # Independent reference: the original expert-major BMM computes every expert, + # with zero inputs for experts outside the top-k selection. + scores = torch.zeros_like(logits) + values, indices = logits.topk(top_k, dim=-1) + scores.scatter_(1, indices, values.sigmoid()) + routed = hidden_states.unsqueeze(0) * scores.t().unsqueeze(-1) + gate, up = torch.bmm(routed, moe.experts.gate_up_proj.transpose(1, 2)).chunk(2, dim=-1) + routed = torch.bmm(moe.experts.act_fn(gate) * up, moe.experts.down_proj.transpose(1, 2)) + expected = moe.shared_expert(hidden_states) + routed.sum(dim=0) + torch.testing.assert_close(output, expected) + + def test_legacy_expert_checkpoint_layout(self): + model = Llama4ForCausalLM(self.get_config(top_k=2)).eval() + input_ids = torch.tensor([[1, 4, 7]]) + with tempfile.TemporaryDirectory() as checkpoint: + model.save_pretrained(checkpoint) + state = load_file(f"{checkpoint}/model.safetensors") + prefix = "model.layers.0.feed_forward.experts." + torch.testing.assert_close( + state[prefix + "gate_up_proj"], model.state_dict()[prefix + "gate_up_proj"].transpose(1, 2) + ) + torch.testing.assert_close( + state[prefix + "down_proj"], model.state_dict()[prefix + "down_proj"].transpose(1, 2) + ) + loaded = Llama4ForCausalLM.from_pretrained(checkpoint).eval() + with torch.no_grad(): + expected = model(input_ids, output_router_logits=True) + actual = loaded(input_ids, output_router_logits=True) + torch.testing.assert_close(actual.logits, expected.logits) + self.assertEqual(len(actual.router_logits), model.config.num_hidden_layers) + for actual_router, expected_router in zip(actual.router_logits, expected.router_logits): + self.assertEqual(actual_router.shape, (input_ids.numel(), model.config.num_local_experts)) + torch.testing.assert_close(actual_router, expected_router) + + @slow @require_torch_large_accelerator class Llama4IntegrationTest(MemoryCleanupMixin, unittest.TestCase): diff --git a/tests/quantization/finegrained/test_config_readers.py b/tests/quantization/finegrained/test_config_readers.py index 25782b173ade..879cf1995207 100644 --- a/tests/quantization/finegrained/test_config_readers.py +++ b/tests/quantization/finegrained/test_config_readers.py @@ -19,7 +19,13 @@ from transformers.quantizers.auto import AutoQuantizationConfig from transformers.testing_utils import require_torch -from transformers.utils.quantization_config import NVFP4Config, read_compressed_tensors, read_modelopt +from transformers.utils.quantization_config import ( + BlockFp8Config, + FbgemmFp8Config, + NVFP4Config, + read_compressed_tensors, + read_modelopt, +) from ...integrations.finegrained.utils import _compressed_tensors, _modelopt, _quantizer_for @@ -28,6 +34,37 @@ class BlockScaledCheckpointReaderTest(unittest.TestCase): """modelopt and compressed-tensors NVFP4 exports read as `NVFP4Config`, skip-lists bounded to the modules they name.""" + def test_rowwise_fp8_producers_share_the_finegrained_runtime(self): + from transformers.quantizers.auto import AutoHfQuantizer + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + channel = {"num_bits": 8, "type": "float", "strategy": "channel"} + token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} + tensor = {"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": False} + for checkpoint, activation_format, scheme, upper_bound in ( + ({"quant_method": "fbgemm_fp8", "activation_scale_ub": 1200.0}, None, "dynamic", 1200.0), + ({"quant_method": "modelopt", "quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"}, None, "dynamic", None), + (_compressed_tensors(channel, token, fmt="float-quantized"), None, "dynamic", None), + (_compressed_tensors(channel, tensor, fmt="float-quantized"), None, "static", None), + (_compressed_tensors(channel, fmt="float-quantized"), "bf16", "dynamic", None), + ): + with self.subTest(checkpoint=checkpoint): + config = AutoQuantizationConfig.from_dict(checkpoint) + self.assertIsInstance(config, BlockFp8Config) + self.assertEqual(config.weight_block_size, "row") + self.assertEqual(config.activation_format, activation_format) + self.assertEqual(config.activation_scheme, scheme) + self.assertEqual(config.activation_scale_ub, upper_bound) + self.assertIsInstance(AutoHfQuantizer.from_config(config), BlockFp8HfQuantizer) + self.assertEqual(AutoQuantizationConfig.from_dict(config.to_dict()).to_dict(), config.to_dict()) + self.assertIsInstance(AutoHfQuantizer.from_config(FbgemmFp8Config()), BlockFp8HfQuantizer) + checkpoint = AutoQuantizationConfig.from_dict(FbgemmFp8Config().to_dict()) + loading = AutoQuantizationConfig.from_dict(FbgemmFp8Config(activation_scale_ub=600.0).to_dict()) + self.assertEqual(AutoHfQuantizer.merge_quantization_configs(checkpoint, loading).activation_scale_ub, 600.0) + weight_only = AutoHfQuantizer.merge_quantization_configs(checkpoint, BlockFp8Config(activation_format="bf16")) + self.assertEqual(weight_only.activation_format, "bf16") + self.assertIsNone(weight_only.activation_scale_ub) + 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, so the export is one format.""" @@ -185,7 +222,7 @@ def test_a_compressed_tensors_nvfp4_checkpoint_reads_as_nvfp4(self): self.assertEqual(config.checkpoint_format, "compressed-tensors") def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): - """Integer, per-channel FP8 and sparse checkpoints are compressed-tensors' own.""" + """Integer, unsupported activation recipes and sparse checkpoints stay compressed-tensors' own.""" from transformers.utils.quantization_config import read_compressed_tensors int4 = {"num_bits": 4, "type": "int", "group_size": 128, "strategy": "group"} @@ -193,7 +230,9 @@ def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} for checkpoint in ( _compressed_tensors(int4, fmt="pack-quantized"), - _compressed_tensors(channel, fmt="float-quantized"), + _compressed_tensors( + channel, {"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": True}, fmt="float-quantized" + ), _compressed_tensors(nvfp4, sparsity_config={"format": "sparse-24-bitmask"}), ): self.assertIsNone(read_compressed_tensors(checkpoint)) diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index 35e897e3e219..decbad6b90b5 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import unittest +from types import SimpleNamespace +from unittest import mock import torch @@ -65,25 +67,25 @@ def __getitem__(self, name): ) -def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): +def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape, source_axis_order=None): """Build a DtensorShardOperation without requiring a real DTensor / distributed init. The axis-0 ownership cache is computed by mimicking ``compute_local_shape_and_global_offset`` for the leading dim only: locate the mesh dim that shards param dim 0 (if any) and use its local rank. """ - op = object.__new__(DtensorShardOperation) - op.device_mesh = mesh - op.placements = tuple(placements) - op.param_ndim = len(param_shape) - op._axis0_offset = 0 - op._axis0_local_size = local_shape[0] + offsets = [0] * len(param_shape) for mesh_dim, p in enumerate(placements): if hasattr(p, "dim") and (p.dim % len(param_shape)) == 0: sub = mesh[mesh.mesh_dim_names[mesh_dim]] if mesh.ndim > 1 else mesh - op._axis0_offset = sub.get_local_rank() * local_shape[0] + offsets[0] = sub.get_local_rank() * local_shape[0] break - return op + param = SimpleNamespace(device_mesh=mesh, placements=tuple(placements), ndim=len(param_shape), shape=param_shape) + with mock.patch( + "transformers.distributed.sharding_utils.compute_local_shape_and_global_offset", + return_value=(local_shape, offsets), + ): + return DtensorShardOperation(param, source_axis_order=source_axis_order) class TestDtensorShardOperation(unittest.TestCase): @@ -131,6 +133,21 @@ def test_no_shard_placements_returns_full_copy(self): op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + def test_transposed_checkpoint_reads_the_source_sharding_axis(self): + source = torch.arange(4 * 8 * 12).reshape(4, 8, 12) + target = source.transpose(1, 2) + for placement in (Shard(0), Shard(2), _StridedShard(1, split_factor=2)): + for rank in range(2): + with self.subTest(placement=placement, rank=rank): + mesh = FakeMesh(shape=2, rank=rank) + local_shape = list(target.shape) + local_shape[placement.dim] //= 2 + expected = _make_dtensor_shard_op(mesh, [placement], target.shape, local_shape).shard_tensor( + target + ) + reader = _make_dtensor_shard_op(mesh, [placement], target.shape, local_shape, (0, 2, 1)) + torch.testing.assert_close(reader.shard_tensor(source).transpose(1, 2), expected) + def test_1D_shard(self): tensor = torch.arange(16).reshape(4, 4).float() expected = { From 1d7444285cacac2537f3f6bd77e0d0fb5fd86808 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Sat, 17 Oct 2026 05:47:04 -0700 Subject: [PATCH 91/91] Fix finegrained CI docstrings, TP fixtures, and MuseSpark expert plans --- .../models/muse_spark/configuration_muse_spark.py | 7 +++++++ .../models/muse_spark/modular_muse_spark.py | 7 +++++++ src/transformers/utils/quantization_config.py | 4 ++-- tests/tensor_parallel/test_tensor_parallel.py | 14 +++++++------- 4 files changed, 23 insertions(+), 9 deletions(-) diff --git a/src/transformers/models/muse_spark/configuration_muse_spark.py b/src/transformers/models/muse_spark/configuration_muse_spark.py index f5e607c2348a..46e8bd9f4442 100644 --- a/src/transformers/models/muse_spark/configuration_muse_spark.py +++ b/src/transformers/models/muse_spark/configuration_muse_spark.py @@ -105,6 +105,12 @@ class MuseSparkTextConfig(PreTrainedConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", + # Post-expert normalization requires complete expert outputs before ranks combine them. + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", + "layers.*.mlp.experts": "moe_tp_experts", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -140,6 +146,7 @@ class MuseSparkTextConfig(PreTrainedConfig): post_norm_eps: float = 1e-8 layer_rope_theta: list[float | int] | None = None base_model_ep_plan = { + "layers.*.mlp.gate": "replicate", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", diff --git a/src/transformers/models/muse_spark/modular_muse_spark.py b/src/transformers/models/muse_spark/modular_muse_spark.py index e8f5a997f869..81ecda602c14 100644 --- a/src/transformers/models/muse_spark/modular_muse_spark.py +++ b/src/transformers/models/muse_spark/modular_muse_spark.py @@ -110,8 +110,15 @@ class MuseSparkTextConfig(MuseGlimmerTextConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", + # Post-expert normalization requires complete expert outputs before ranks combine them. + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", + "layers.*.mlp.experts": "moe_tp_experts", } base_model_ep_plan = { + "layers.*.mlp.gate": "replicate", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index 2a6130a84bf1..85c6da01b3c7 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -1838,8 +1838,6 @@ class BlockFp8Config(BlockScaledConfig): weight_block_size (`tuple[int, int] | str`, *optional*, defaults to `(128, 128)`): The block each weight scale covers; `None` for one scale per tensor, or `"row"` for one scale per output channel. Dynamic FP8 activations then use one scale per token. - activation_scale_ub (`float`, *optional*): - Maximum activation amax for rowwise FP8 quantization. Values beyond it saturate, as in FBGEMM. scale_fmt (`str`, *optional*, defaults to `"float"`): Storage dtype of the scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`, DeepSeek-V4). activation_scheme (`str`, *optional*, defaults to `"dynamic"`): @@ -1847,6 +1845,8 @@ class BlockFp8Config(BlockScaledConfig): scale from the checkpoint). modules_to_convert (`list`, *optional*): Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. + activation_scale_ub (`float`, *optional*): + Maximum activation amax for rowwise FP8 quantization. Values beyond it saturate, as in FBGEMM. kwargs (`dict[str, Any]`, *optional*): The options every block-scaled format takes, see [`BlockScaledConfig`]. """ diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 906c1200c197..924973aeab92 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import warnings +from types import SimpleNamespace from unittest.mock import patch import torch @@ -225,13 +226,12 @@ def _get_local_shape(self, global_shape, placement, world_size, rank): return tuple(local_shape) def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape): - op = object.__new__(DtensorShardOperation) - op.device_mesh = mesh - op.placements = (placement,) - op.param_ndim = len(param_shape) - op._axis0_offset = 0 - op._axis0_local_size = local_shape[0] - return op + param = SimpleNamespace(device_mesh=mesh, placements=(placement,), ndim=len(param_shape), shape=param_shape) + with patch( + "transformers.distributed.sharding_utils.compute_local_shape_and_global_offset", + return_value=(local_shape, (0,) * len(param_shape)), + ): + return DtensorShardOperation(param) def test_colwise_gather_output_rejects_indivisible_out_features(self): model = torch.nn.Module()