From 6afc76c29e3cb6fd1e42eb48f021d6f325ce0815 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:25:52 +0000 Subject: [PATCH 1/5] Dequantize block-FP8 checkpoints as they load expert-parallel A native block-FP8 checkpoint, such as DeepSeek-V3's, is dequantized to bf16 while it loads expert-parallel, through transformers' FineGrainedFP8Config(dequantize=True); on transformers older than 5.10.1 the loader patches that dequantization for a partial last block, the destination dtype and the routed experts' fused converter. A block-FP8 checkpoint keeps its dense part replicated, since DCP reads it as stored. A per-tensor FP8 checkpoint, which names no weight_block_size, is refused, and so is an FP8 config object whose checkpoint has no config.json to say which it is. This is the PR's 25 commits in one, re-parented with #2742 onto the expert-parallel loader's [2/3]: the FP8 refusals, with their read of the checkpoint's config.json, run inside the loader's checks guard, so a rank refusing alone fails every rank. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 341 +++++++++++++++- .../torch/utils/test_fp8_dequant_patch.py | 369 ++++++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 139 +++++-- 3 files changed, 800 insertions(+), 49 deletions(-) create mode 100644 tests/unit/torch/utils/test_fp8_dequant_patch.py diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c2438c34e46..c7749ecbcbb 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -19,6 +19,7 @@ import copy import dataclasses import inspect +import json import logging import os import re @@ -1582,18 +1583,36 @@ def _decide_distribution( f"({num_experts})." ) # Why the expert-parallel load cannot shard the dense part, if it cannot: the export cannot write - # it sharded within several expert-data-parallel replicas. + # it sharded within several expert-data-parallel replicas, and Distributed Checkpoint, which + # loads it sharded, reads a block-fp8 checkpoint's tensors as stored, with no dequantization. unshardable, remedy = None, "" + block_fp8 = ep_size > 1 and _ckpt_quant_method(hf_config, ckpt_path) == "fp8" if ep_size > 1 and world_size // ep_size > 1: unshardable = ( f"within each of its {world_size // ep_size} expert-data-parallel replicas, a layout " "the export cannot write yet" ) + # No EP degree lets a block-fp8 dense part shard (below), so none is offered for one. remedy = ( - " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, " + " Set MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU, or start from an " + "unquantized copy of the model." + if block_fp8 + else " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, " "MODELOPT_EP_DEGREE=1 to FSDP2-shard the whole model without expert parallelism, or " "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) + elif block_fp8: + # Expert-parallel only, as its remedy holds there alone: the expert-parallel load is the one + # that dequantizes block FP8 as it loads. Without expert parallelism the dense part is the + # whole model, and the FSDP2 load refuses a quantized checkpoint itself (_load_sharded). + unshardable = ( + "from a block-fp8 checkpoint: Distributed Checkpoint reads it as stored, with no " + "dequantization" + ) + remedy = ( + " Set MODELOPT_FSDP_SIZE=1 to replicate it, which dequantizes it as it loads, or start " + "from an unquantized copy of the model." + ) if dense == "auto": # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size # share of the routed experts it does. @@ -1863,6 +1882,227 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st ) +_FINEGRAINED_FP8_PATCHED = False +_FP8_DEQUANT_DTYPE = None # model compute dtype the dequant must emit (set by the call site) + + +def _patch_finegrained_fp8_dequant(target_dtype=None): + """Correct transformers' block-FP8 dequant (``Fp8Dequantize._dequantize_one``) for DeepSeek-V3. + + transformers 5.9's ``_dequantize_one`` infers the block geometry by *dividing* the weight shape + by the scale-grid shape and requires exact divisibility:: + + block_m = rows // scale_rows # assumes rows == scale_rows * block_m exactly + + But DeepSeek-V3 block-FP8 uses a fixed 128x128 block with a *ceil* scale grid, so a weight whose + dimension isn't a multiple of 128 has a partial last block: e.g. ``kv_a_proj_with_mqa`` is + (576, 7168) with a (5, 56) grid -> ``576 % 5 != 0`` raises ``ValueError``, and degenerate/empty + scale grids from the merged/sharded expert path hit ``rows % 0`` -> ``ZeroDivisionError``. + + Replace it with a divisibility-free tile-and-crop: broadcast each per-block scale over its + ``block`` rows/cols (block size from the real ``weight_block_size`` when the grid is inexact, + else derived), then crop to the weight shape -- correct for both full and partial last blocks, + 2-D dense linears and (leading-dim-broadcast) per-expert tensors. Idempotent. + + transformers >= 5.10.1 keeps its own dequant. Upstream fixed two of the show-stoppers there -- the + loader's expert ``tensor_idx`` check became ``any(op is MergeModulelist ...)`` (bug 1) and + ``_dequantize_one`` gained an ``output_dtype`` the loader passes as the destination dtype (bug 3) -- + and it added native FP4-unpack + ue8m0 (E8M0) scale decoding that our simpler override does NOT do + (installing it on >= 5.10.1 would BREAK ue8m0 checkpoints like DeepSeek-V4). It still derives the + block from the scale grid (bug 2), which a partial last block makes wrong: 576 rows under a 5-row + grid do not divide at all, and 448 rows under a 4-row grid divide into blocks of 112 instead of 128. + So there a weight whose grid is the configured block's ceil grid and that has a partial last block + is padded to whole blocks for it and the result cropped back. + """ + global _FINEGRAINED_FP8_PATCHED, _FP8_DEQUANT_DTYPE + if isinstance(target_dtype, str): + # A config can name its dtype, or carry "auto" for the checkpoint's own, and Tensor.to reads a + # string as a device: take the torch.dtype it names, and for "auto" the bfloat16 default. + name = target_dtype.rsplit(".", 1)[-1] + target_dtype = torch.bfloat16 if name == "auto" else getattr(torch, name) + if target_dtype is not None: + _FP8_DEQUANT_DTYPE = target_dtype # keep current even if already patched + if _FINEGRAINED_FP8_PATCHED: + return + import transformers + from packaging.version import parse + from transformers.integrations import finegrained_fp8 as _fp8 + + if parse(transformers.__version__) >= parse("5.10.1"): + # Re-importing this module loses its flag, and the wrapper calls the method it replaces: + # the class keeps a flag of its own, so the wrapper never wraps itself. + if getattr(_fp8.Fp8Dequantize, "_modelopt_dequant_patched", False): + _FINEGRAINED_FP8_PATCHED = True + return + _stock_dequantize_one = _fp8.Fp8Dequantize._dequantize_one + fp8_dtypes = { + getattr(torch, name) + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") + if hasattr(torch, name) + } + + def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): + if quantized.dim() < 2 or scales.dim() < 2 or quantized.dtype not in fp8_dtypes: + return _stock_dequantize_one(self, quantized, scales, *args, **kwargs) + rows, cols = quantized.shape[-2:] + block = getattr( + getattr(self.hf_quantizer, "quantization_config", None), "weight_block_size", None + ) or (128, 128) + ceil_grid = (-(-rows // block[0]), -(-cols // block[1])) + if tuple(scales.shape[-2:]) != ceil_grid or ( + rows % block[0] == 0 and cols % block[1] == 0 + ): + return _stock_dequantize_one(self, quantized, scales, *args, **kwargs) + pad_rows, pad_cols = ceil_grid[0] * block[0] - rows, ceil_grid[1] * block[1] - cols + # Padded in its own fp8 dtype, through a same-size uint8 view (pad has no fp8 kernel), so + # transformers' dequant gets the input it is written for, without an fp32 copy. + padded = torch.nn.functional.pad( + quantized.view(torch.uint8), (0, pad_cols, 0, pad_rows) + ).view(quantized.dtype) + out = _stock_dequantize_one(self, padded, scales, *args, **kwargs) + return out[..., :rows, :cols] + + _fp8.Fp8Dequantize._dequantize_one = _dequantize_partial_blocks + _fp8.Fp8Dequantize._modelopt_dequant_patched = True + _FINEGRAINED_FP8_PATCHED = True + return + + def _dequantize_one(self, quantized, scales, output_dtype=None): + # Block FP8 only: the releases this replaces have no FP4 unpack (5.10.1 added it). + if quantized.dim() < 2 or scales.dim() < 2: + raise ValueError( + f"block-FP8 dequant: weight {tuple(quantized.shape)} with scales " + f"{tuple(scales.shape)} is not a block-scaled 2-D weight." + ) + q = quantized.to(torch.float32, copy=True) + rows, cols = q.shape[-2:] + scale_rows, scale_cols = scales.shape[-2:] + # Emit the MODEL compute dtype (bf16), NOT the scale dtype. DeepSeek's weight_scale_inv is F32; + # transformers' original out_dtype logic would return F32 weights, which then run against bf16 + # activations -> "mat1 and mat2 have the same dtype" RuntimeError in the calibration matmuls. + # A loader that passes the destination dtype, as transformers' does from 5.10.1, knows it best. + out_dtype = output_dtype or _FP8_DEQUANT_DTYPE or torch.bfloat16 + if scale_rows == 0 or scale_cols == 0 or scales.numel() == 0: + # Only an empty weight has no scales. A non-empty one with an empty grid was sliced + # apart from its scales (as 5.9's expert tensor_idx bug did): emitting it unscaled would + # load wrong weights without a word. + if q.numel() == 0: + return q.to(out_dtype) + raise ValueError( + f"block-FP8 dequant: weight {tuple(q.shape)} has an empty scale grid " + f"{tuple(scales.shape)}; the weight and its scales were sliced apart." + ) + wbs = getattr( + getattr(self.hf_quantizer, "quantization_config", None), "weight_block_size", None + ) or (128, 128) + + def _blk(dim, sgrid, cfg_blk): + # A ceil scale grid at the config block size (DeepSeek's uniform 128x128, incl. partial last + # block e.g. 576 -> 5 blocks) is the common case -- trust the config block. Only a grid that + # is NOT the ceil-at-config-block (a coarser/per-tensor scale) derives the block by division; + # never infer from exact divisibility alone (576 % 5 and 9 % 3 both mislead). + if sgrid == -(-dim // cfg_blk): + return cfg_blk + if dim % sgrid == 0: + return dim // sgrid + # Neither: no block is known to place it, and guessing would scale rows wrongly. + raise ValueError( + f"block-FP8 dequant: a scale grid of {sgrid} fits neither the configured block " + f"{cfg_blk} nor an even split of {dim}." + ) + + block_m = _blk(rows, scale_rows, wbs[0]) + block_n = _blk(cols, scale_cols, wbs[1]) + # Scale q, its own copy, in place: its whole column blocks through a view, then a partial + # last one. The scales are tiled over the rows only, so no full-size fp32 tile is built. + s = scales.to(torch.float32).repeat_interleave(block_m, dim=-2)[..., :rows, :] + whole = cols // block_n + q[..., : whole * block_n].unflatten(-1, (whole, block_n)).mul_(s[..., :whole, None]) + if whole < scale_cols: + q[..., whole * block_n :].mul_(s[..., whole:]) + return q.to(out_dtype) + + # --- Fix the routed-expert block-FP8 load corruption (transformers 5.9) --------------------- + # The FP8 quantizer's update_weight_conversions PREPENDS Fp8Dequantize to each converter's + # operations, so for the fused-MoE experts converter operations[0] becomes Fp8Dequantize instead + # of the original MergeModulelist. The state-dict loader only assigns a per-expert tensor_idx when + # ``isinstance(mapping.operations[0], MergeModulelist)`` (core_model_loading.py) -- with the prepend + # that check fails, tensor_idx becomes None, and GroupedGemmParallel.shard_tensor then row-slices + # each per-expert (2048,7168) weight to param[rank*local:(rank+1)*local] = (16,7168) and slices the + # scales full/empty per rank. Result: gate_up_proj loads as (256,32,7168), a grouped-GEMM contraction + # mismatch at calibration. Fix: FUSE the prepended Fp8Dequantize into the expert MergeModulelist (a + # MergeModulelist subclass that dequantizes each per-expert (weight,scale) pair, dropping scales, then + # stacks) so operations[0] IS a MergeModulelist again -- the loader assigns tensor_idx (whole-expert + # selection + EP shard) and Concatenate still runs after scales are dropped. Idempotent (flag-guarded). + try: + from transformers.core_model_loading import MergeModulelist as _MergeModulelist + from transformers.core_model_loading import WeightConverter as _WeightConverter + from transformers.quantizers.quantizer_finegrained_fp8 import ( + FineGrainedFP8HfQuantizer as _FineGrainedFP8HfQuantizer, + ) + + # Fuse the prepended Fp8Dequantize INTO the expert MergeModulelist: dequantize each per-expert + # (weight, scale) pair -- dropping scales -- then stack, exactly as the intended + # [Fp8Dequantize, MergeModulelist, Concatenate] chain does, but packaged as a single op that + # *is a* MergeModulelist. That keeps operations[0] a MergeModulelist so the loader assigns each + # per-expert its tensor_idx (whole-expert selection + EP shard) AND Concatenate still runs after + # dequant has dropped the scales (Concatenate cats every remaining source pattern, so it must not + # see scale keys). Net execution is identical to upstream's intent; only the op layout changes. + class _DequantThenMerge(_MergeModulelist): + def __init__(self, deq, dim=0): + super().__init__(dim) + self._deq = deq + + def convert(self, input_dict, source_patterns=None, target_patterns=None, **kwargs): + deq_out = self._deq.convert( + input_dict, + source_patterns=source_patterns, + target_patterns=target_patterns, + **kwargs, + ) + return super().convert( + deq_out, + source_patterns=source_patterns, + target_patterns=target_patterns, + **kwargs, + ) + + # A transformers without the hook prepends no dequant to the expert converters: nothing to fuse. + if hasattr(_FineGrainedFP8HfQuantizer, "update_weight_conversions") and not getattr( + _FineGrainedFP8HfQuantizer, "_modelopt_uwc_patched", False + ): + _orig_uwc = _FineGrainedFP8HfQuantizer.update_weight_conversions + + def _uwc_fixed(self, weight_conversions): + convs = _orig_uwc(self, weight_conversions) + for conv in convs: + ops = getattr(conv, "operations", None) + # Match the expert converter update_weight_conversions produced: + # [Fp8Dequantize, MergeModulelist, ]. Fuse ops[0]+ops[1]. + if ( + isinstance(conv, _WeightConverter) + and ops + and len(ops) >= 2 + and isinstance(ops[0], _fp8.Fp8Dequantize) + and isinstance(ops[1], _MergeModulelist) + ): + deq, mml = ops[0], ops[1] + ops[:2] = [_DequantThenMerge(deq, dim=mml.dim)] + return convs + + _FineGrainedFP8HfQuantizer.update_weight_conversions = _uwc_fixed + _FineGrainedFP8HfQuantizer._modelopt_uwc_patched = True + except Exception as e: + # Unpatched, every routed expert would load sliced apart from its scales (see above). + raise RuntimeError( + f"block-FP8 load: could not fuse transformers' expert dequant into its merge ({e})." + ) from e + + # Rebound once the fuse is in, so a refused patch leaves transformers as it was. + _fp8.Fp8Dequantize._dequantize_one = _dequantize_one + _FINEGRAINED_FP8_PATCHED = True + + def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: """FSDP2-shard the dense part over ``fsdp_mesh``, block by block; return the blocks wrapped. @@ -1935,22 +2175,58 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None). # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's. -_EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None}) +# fp8 is block FP8; _ckpt_quant_method calls per-tensor FP8 fp8-per-tensor, which the load refuses. +_EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None, "fp8"}) -def _ckpt_quant_method(cfg): +def _ckpt_quant_method(cfg, ckpt_path: str | None = None): """Return the checkpoint's native quantization method, or ``None``. Return the checkpoint's native quantization method (e.g. ``"fp8"`` for block-fp8 DeepSeek-V3), or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's - ``text_config`` or ``llm_config``. + ``text_config`` or ``llm_config``. Per-tensor FP8 (AutoFP8's), which names the same method as + block FP8 but no ``weight_block_size``, is ``"fp8-per-tensor"``. transformers' FP8 config object + fills in a 128x128 block whatever the checkpoint names, so for one the ``config.json`` in + ``ckpt_path`` says whether the checkpoint names a block; without a readable one (a Hub ID, say) + it is ``"fp8-unverified"``, for the load to refuse rather than guess. """ for c in (cfg, getattr(cfg, "text_config", None), getattr(cfg, "llm_config", None)): qc = getattr(c, "quantization_config", None) m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) if m: # transformers' config objects carry a QuantizationMethod enum, whose str() is not its value. - return str(getattr(m, "value", m)) + m = str(getattr(m, "value", m)) + if not isinstance(qc, dict) and m == "fp8": + qc = _config_json_quantization(ckpt_path) + if qc is None: + return "fp8-unverified" + block = ( + qc.get("weight_block_size") + if isinstance(qc, dict) + else getattr(qc, "weight_block_size", None) + ) + # transformers reads per-tensor FP8 as 128x128 blocks and loads its weights unscaled. + return "fp8-per-tensor" if m == "fp8" and not block else m + return None + + +def _config_json_quantization(ckpt_path: str | None) -> dict | None: + """The ``quantization_config`` the checkpoint's own ``config.json`` names, or ``None``. + + Read as stored: the top level, else a VLM's text or LLM config, where :func:`_ckpt_quant_method` + looks too. + """ + if not ckpt_path: + return None + try: + with open(os.path.join(ckpt_path, "config.json")) as f: + raw = json.load(f) + except (OSError, ValueError): + return None + for c in (raw, raw.get("text_config"), raw.get("llm_config")): + qc = c.get("quantization_config") if isinstance(c, dict) else None + if isinstance(qc, dict): + return qc return None @@ -2082,13 +2358,16 @@ def load_and_prepare_ep( ``num_experts / ep_size`` routed experts and a replica of the dense part, and tokens reach their experts' ranks by all-to-all, so every rank must calibrate on data of its own. - It loads unquantized checkpoints, on transformers 5.7 or later. A refusal on any rank -- of its - transformers, the config, the checkpoint, its GPU, or in the load -- raises on every rank. In - the load that rests on ``from_pretrained`` running no collective, as transformers' - expert-parallel load does not on 5.7 through 5.18. + It loads unquantized checkpoints on transformers 5.7 or later, and block-FP8 ones, which it + dequantizes as they load, on 5.8 or later. A refusal on any rank -- of its transformers, the + config, the checkpoint, its GPU, or in the load -- raises on every rank. In the load that rests + on ``from_pretrained`` running no collective, as transformers' expert-parallel load does not on + 5.7 through 5.18. With ``plan.dense == "fsdp"`` the dense part is FSDP2-sharded over the ``ep`` group instead, - each rank loading only its shard: for a model whose dense part does not fit one GPU. + each rank loading only its shard: for a model whose dense part does not fit one GPU. A block-FP8 + checkpoint cannot be loaded that way: Distributed Checkpoint reads it as stored, with no + dequantization. The model builds with its LM head where an auto class builds one; ``require_lm_head=False`` builds a VLM without it. Without ``hf_config`` every rank reads the checkpoint's config itself; @@ -2137,7 +2416,27 @@ def load_and_prepare_ep( if attn_implementation is not None: config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) - quant_method = _ckpt_quant_method(hf_config) + quant_method = _ckpt_quant_method(hf_config, ckpt_path) + if plan.is_fsdp_dense and quant_method == "fp8": + raise NotImplementedError( + f"{plan!r}: the dense part of a block-fp8 checkpoint cannot be sharded: Distributed " + "Checkpoint reads it as stored, with no dequantization. Plan dense='replicate'." + ) + if quant_method == "fp8-per-tensor": + # Not "cannot dequantize fp8", which it does: the block size it needs is what is missing. + raise NotImplementedError( + "The expert-parallel load dequantizes block-fp8 checkpoints, which name their block in " + "quantization_config.weight_block_size, and this fp8 checkpoint names none: per-tensor " + "scales (AutoFP8's) would load as 128x128 blocks, leaving the weights unscaled. Start " + "from an unquantized copy of the model." + ) + if quant_method == "fp8-unverified": + raise NotImplementedError( + "This fp8 config came in as transformers' config object, whose 128x128 block is a " + f"default, and {ckpt_path!r} holds no config.json to say whether the checkpoint names " + "one: per-tensor scales would load as blocks, leaving the weights unscaled. Hand the " + "config as AutoConfig.from_pretrained reads it, or a local checkpoint." + ) if quant_method not in _EP_SOURCE_FORMATS: # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) @@ -2274,6 +2573,24 @@ def gpu() -> tuple[bool, bool, str | None]: from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation + # Native block-FP8 checkpoints (e.g. DeepSeek-V3: quantization_config.quant_method="fp8", e4m3 with + # weight_block_size) must be DEQUANTIZED to the model dtype on load: ModelOpt PTQ inserts/exports quantizers on + # plain bf16 Linears, and the EP grouped_gemm experts expect bf16 fused params -- not fp8 blocks with + # weight_scale_inv. Mirror get_model's mxfp4 dequant-on-load, using FineGrainedFP8Config(dequantize=True). + # (compressed-tensors pack-quantized experts, e.g. Kimi, are handled separately by the on-the-fly + # decompressor plugin; here we only cover the native fp8 case.) + if quant_method == "fp8": + from transformers import FineGrainedFP8Config + + # transformers keeps the checkpoint's quantization_config, its weight_block_size included, + # and takes only the loading attributes (dequantize) from this one. + from_pretrained_kwargs["quantization_config"] = FineGrainedFP8Config(dequantize=True) + _patch_finegrained_fp8_dequant( + dtype + ) # correct transformers' block-fp8 dequant; emit model dtype + print_rank_0( + f"[ep] native block-fp8 checkpoint -> FineGrainedFP8Config(dequantize=True) -> {dtype}" + ) # Swap the plan to the torch_a2a styles and select the torch_a2a experts forward: tokens reach # their experts' ranks by all-to-all around the quantizable grouped GEMM. register_torch_a2a() # idempotent; register the experts forward + the two styles, right before use diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py new file mode 100644 index 00000000000..d9e378f66c1 --- /dev/null +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""``_patch_finegrained_fp8_dequant``: transformers' block-FP8 dequant, fixed for DeepSeek-V3. + +transformers 5.9's ``Fp8Dequantize._dequantize_one`` finds the block size by dividing the weight +by the scale grid, which breaks on DeepSeek-V3's partial last blocks (a 576-row weight under a +5-row grid of 128-row blocks). The patch installs a tile-and-crop dequant instead, and only below +transformers 5.10.1, so the tests present an older version to it. The dequant is checked against a +plain blockwise reference: ``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``. +""" + +import sys +from types import SimpleNamespace + +import pytest +import torch +from packaging.version import Version + +transformers = pytest.importorskip("transformers") +finegrained_fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8") +quantizer_fp8 = pytest.importorskip("transformers.quantizers.quantizer_finegrained_fp8") + +from transformers import FineGrainedFP8Config +from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter + +import modelopt.torch.utils.plugins.model_load_utils as model_load_utils + +if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("torch has no float8_e4m3fn", allow_module_level=True) + +# The installed version: the fixture below makes the patch see an older one. +_INSTALLED_TRANSFORMERS = Version(transformers.__version__) +Fp8Dequantize = finegrained_fp8.Fp8Dequantize +FineGrainedFP8HfQuantizer = quantizer_fp8.FineGrainedFP8HfQuantizer + +_MISSING = object() +# What the patch may rebind below transformers 5.10.1: the dequant method, and the FP8 quantizer's +# converter hook plus the flag guarding it. From 5.10.1 it wraps the dequant method, guarded by a +# flag of its own. +_PATCH_TARGETS = ( + (Fp8Dequantize, "_dequantize_one"), + (FineGrainedFP8HfQuantizer, "update_weight_conversions"), + (FineGrainedFP8HfQuantizer, "_modelopt_uwc_patched"), +) +_RESTORED = (*_PATCH_TARGETS, (Fp8Dequantize, "_modelopt_dequant_patched")) + + +@pytest.fixture +def pre_5_10_transformers(monkeypatch): + """Unpatched transformers, reporting a version the patch targets; restored afterwards.""" + saved = [(cls, name, vars(cls).get(name, _MISSING)) for cls, name in _RESTORED] + monkeypatch.setattr(transformers, "__version__", "5.9.0") + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + monkeypatch.setattr(model_load_utils, "_FP8_DEQUANT_DTYPE", None) + yield + for cls, name, value in saved: + if value is not _MISSING: + setattr(cls, name, value) + elif name in vars(cls): + delattr(cls, name) + + +def _patched_state(): + return [vars(cls).get(name, _MISSING) for cls, name in _PATCH_TARGETS] + + +def _dequant_op(block_size): + """The dequant op of an FP8 quantizer configured with ``weight_block_size=block_size``.""" + config = SimpleNamespace(weight_block_size=block_size) + return Fp8Dequantize(SimpleNamespace(quantization_config=config)) + + +def _block_fp8(shape, grid, seed=0): + """A random fp8 weight and its fp32 ``scale_inv`` grid.""" + gen = torch.Generator().manual_seed(seed) + weight = (4 * torch.randn(shape, generator=gen)).to(torch.float8_e4m3fn) + return weight, torch.rand(grid, generator=gen) + 0.5 + + +def _ceil_grid(shape, block): + *lead, rows, cols = shape + return (*lead, -(-rows // block[0]), -(-cols // block[1])) + + +def _reference(weight, scale_inv, block): + """``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``, in fp32.""" + rows, cols = weight.shape[-2:] + scale = scale_inv[..., torch.arange(rows) // block[0], :][..., torch.arange(cols) // block[1]] + return weight.float() * scale + + +@pytest.mark.parametrize( + ("shape", "block", "configured", "grid"), + [ + ((8, 12), (4, 4), (4, 4), None), + ((10, 13), (4, 4), (4, 4), None), + ((8, 3), (4, 4), (4, 4), None), + ((9, 8), (4, 4), (4, 4), None), + ((6, 10), (2, 4), (2, 4), None), + ((3, 10, 13), (4, 4), (4, 4), None), + ((576, 7168), (128, 128), (128, 128), None), + ((8, 8), (4, 2), (128, 128), (2, 4)), + ((130, 260), (128, 128), None, (2, 3)), + ((0, 13), (4, 4), (4, 4), None), + ], + ids=[ + "whole-blocks", + "partial-last-blocks", + "narrower-than-a-block", + "grid-divides-the-rows", + "non-square-blocks", + "per-expert", + "deepseek-kv-a-proj", + "grid-finer-than-configured", + "nothing-configured", + "empty-weight", + ], +) +def test_dequant_matches_the_blockwise_reference( + pre_5_10_transformers, shape, block, configured, grid +): + """``configured`` is the quantization config's block size, ``grid`` the scale grid when it is not + that block's ceil grid. ``grid-divides-the-rows``: 9 rows of 4-row blocks make a 3-row grid, + which also divides 9, yet the blocks are 4 rows, not 3. A grid that is not the configured + block's ceil grid sets the block by division; with nothing configured, DeepSeek's 128 x 128 + applies.""" + weight, scale_inv = _block_fp8(shape, grid or _ceil_grid(shape, block)) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + + out = _dequant_op(configured)._dequantize_one(weight, scale_inv) + + assert torch.equal(out, _reference(weight, scale_inv, block)) + + +def test_the_dequantize_config_keeps_the_checkpoints_block_size(): + """The expert-parallel load passes ``FineGrainedFP8Config(dequantize=True)``, and every block + decision of the dequant reads ``weight_block_size`` off the quantizer's config: transformers + keeps the checkpoint's, and takes only the loading attributes from the one passed.""" + from transformers.quantizers.auto import AutoHfQuantizer + + checkpoint = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [64, 64], + } + merged = AutoHfQuantizer.merge_quantization_configs( + checkpoint, FineGrainedFP8Config(dequantize=True) + ) + + assert list(merged.weight_block_size) == [64, 64] + assert merged.dequantize + + +@pytest.mark.parametrize( + ("target_dtype", "output_dtype", "expected"), + [ + (None, None, torch.bfloat16), + (torch.float16, None, torch.float16), + ("float16", None, torch.float16), + ("auto", None, torch.bfloat16), + (torch.bfloat16, torch.float16, torch.float16), + ], + ids=["unrequested", "requested", "named", "auto", "passed-by-the-loader"], +) +def test_dequant_emits_the_requested_model_dtype( + pre_5_10_transformers, target_dtype, output_dtype, expected +): + """Not the scales' dtype: DeepSeek's fp32 ``weight_scale_inv`` would give fp32 weights that + then meet bf16 activations in the calibration matmuls. Unrequested, it is bf16. A config can + name its dtype, or say "auto", which is bf16 too. A loader that passes the destination dtype, + as transformers' does from 5.10.1, gets that one.""" + weight, scale_inv = _block_fp8((10, 13), (3, 4)) + model_load_utils._patch_finegrained_fp8_dequant(target_dtype) + + out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv, output_dtype=output_dtype) + + assert out.dtype == expected + assert torch.equal(out, _reference(weight, scale_inv, (4, 4)).to(expected)) + + +@pytest.mark.parametrize( + ("grid", "match"), + [ + ((0, 4), "empty scale grid"), + ((3, 0), "empty scale grid"), + ((4, 4), "fits neither"), + ((4,), "not a block-scaled 2-D weight"), + ], + ids=["no-row-scales", "no-col-scales", "grid-no-block-places", "one-dim-scales"], +) +def test_dequant_refuses_a_scale_grid_it_cannot_place(pre_5_10_transformers, grid, match): + """Rather than load wrong weights: an empty grid for a non-empty weight (the weight and its + scales were sliced apart, as 5.9's expert ``tensor_idx`` bug did), a 4-row grid over 10 rows, + which no 4-row block (the configured one) or even split gives, or scales that are no grid.""" + weight, scale_inv = _block_fp8((10, 13), grid) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + + with pytest.raises(ValueError, match=match): + _dequant_op((4, 4))._dequantize_one(weight, scale_inv) + + +def test_the_expert_converter_leads_with_a_merge_that_dequantizes_each_expert( + pre_5_10_transformers, +): + """transformers prepends ``Fp8Dequantize`` to the routed-expert converter, so its first op is + no longer the ``MergeModulelist`` the loader keys each expert's ``tensor_idx`` on, and the EP + load then slices every expert by rows. The patch fuses the two into one ``MergeModulelist`` + that dequantizes each expert with its own scales and drops them before ``Concatenate``.""" + if not hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions"): + pytest.skip("this transformers' FP8 quantizer prepends no dequant to the expert converter") + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + experts = 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)], + ) + quantizer = FineGrainedFP8HfQuantizer( + FineGrainedFP8Config(weight_block_size=(4, 4), dequantize=True), pre_quantized=True + ) + + [converter] = [ + conv + for conv in quantizer.update_weight_conversions([experts]) + if conv.target_patterns == ["mlp.experts.gate_up_proj"] + ] + + assert isinstance(converter.operations[0], MergeModulelist) + blocks = { # every expert's weight and scales differ, so a mixed-up pair cannot pass + name: [_block_fp8((8, 12), (2, 3), seed=3 * p + e) for e in range(3)] + for p, name in enumerate(("gate_proj", "up_proj")) + } + tensors = {} # what the loader collects for each of the converter's source patterns + for name, pairs in blocks.items(): + tensors[f"mlp.experts.*.{name}.weight$"] = [weight for weight, _ in pairs] + tensors[f"mlp.experts.*.{name}.weight_scale_inv$"] = [scale_inv for _, scale_inv in pairs] + for op in converter.operations: + tensors = op.convert( + tensors, + source_patterns=converter.source_patterns, + target_patterns=converter.target_patterns, + ) + [fused] = tensors.values() + expected = torch.cat( + [ + torch.stack([_reference(weight, scale_inv, (4, 4)) for weight, scale_inv in pairs]) + for pairs in blocks.values() + ], + dim=1, + ) + assert torch.equal(fused, expected.to(fused.dtype)) + + +def test_patching_again_keeps_the_install_and_takes_the_new_dtype(pre_5_10_transformers): + """Every block-fp8 load patches again, with its own model dtype.""" + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + installed = _patched_state() + + model_load_utils._patch_finegrained_fp8_dequant(torch.float16) + + assert _patched_state() == installed + weight, scale_inv = _block_fp8((10, 13), (3, 4)) + assert _dequant_op((4, 4))._dequantize_one(weight, scale_inv).dtype == torch.float16 + + +@pytest.mark.parametrize( + ("version", "cls", "name"), + [ + ("5.9.0", FineGrainedFP8HfQuantizer, "update_weight_conversions"), + ("5.10.1", Fp8Dequantize, "_dequantize_one"), + ], + ids=["quantizer-hook", "dequant-wrapper"], +) +def test_a_reimported_module_does_not_wrap_a_patch_twice( + pre_5_10_transformers, monkeypatch, version, cls, name +): + """Re-importing ``model_load_utils`` loses its module flag; the patched class carries its own: + the FP8 quantizer for its converter hook, and from 5.10.1 the dequant op for the wrapper that + calls the method it replaced.""" + if not hasattr(cls, name) or Version(version) > _INSTALLED_TRANSFORMERS: + pytest.skip(f"this transformers has no {cls.__name__}.{name} to patch as {version} would") + monkeypatch.setattr(transformers, "__version__", version) + model_load_utils._patch_finegrained_fp8_dequant() + installed = vars(cls)[name] + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + + model_load_utils._patch_finegrained_fp8_dequant() + + assert vars(cls)[name] is installed + + +def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, monkeypatch): + """The patch rebinds only the dequant method and, where the quantizer has one, its converter + hook and guard flag. Putting those back restores the stock classes exactly, after which the + patch installs again once the module flag is cleared.""" + has_hook = hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions") + classes = (Fp8Dequantize, FineGrainedFP8HfQuantizer) + stock = {cls: dict(vars(cls)) for cls in classes} + + model_load_utils._patch_finegrained_fp8_dequant() + + changed = { + (cls, name) + for cls in classes + for name in stock[cls].keys() | vars(cls).keys() + if vars(cls).get(name, _MISSING) is not stock[cls].get(name, _MISSING) + } + assert changed == set(_PATCH_TARGETS if has_hook else _PATCH_TARGETS[:1]) + for cls, name in changed: + if name in stock[cls]: + setattr(cls, name, stock[cls][name]) + else: + delattr(cls, name) + assert all(dict(vars(cls)) == stock[cls] for cls in classes) + + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + model_load_utils._patch_finegrained_fp8_dequant() + assert vars(Fp8Dequantize).get("_dequantize_one") is not stock[Fp8Dequantize].get( + "_dequantize_one" + ) + + +@pytest.mark.parametrize( + ("shape", "grid"), + [((576, 7168), (5, 56)), ((192, 256), (2, 2)), ((512, 256), (4, 2))], + ids=["partial", "partial-dividing-the-grid", "whole"], +) +def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( + pre_5_10_transformers, monkeypatch, shape, grid +): + """Upstream fixed the expert converter there and decodes ue8m0 and FP4 scales, which the override + does not, so its dequant is kept and only wrapped. A partial last block (DeepSeek-V3's 576-row + kv_a_proj_with_mqa under a 5-row grid of 128-row blocks, or 192 rows that its own dequant would + split into two blocks of 96) is padded to whole blocks for it and cropped back; a weight of whole + blocks goes to it as is.""" + if Version("5.10.1") > _INSTALLED_TRANSFORMERS: + pytest.skip("needs transformers' own dequant from 5.10.1") + monkeypatch.setattr(transformers, "__version__", "5.10.1") + hook = vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) + weight, scale_inv = _block_fp8(shape, grid) + + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + out = _dequant_op((128, 128))._dequantize_one(weight, scale_inv, torch.float32) + + assert torch.equal(out, _reference(weight, scale_inv, (128, 128))) + assert vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) is hook + + +def test_an_expert_converter_it_cannot_fuse_stops_the_load(pre_5_10_transformers, monkeypatch): + """Unpatched, every routed expert would load sliced apart from its scales. The refusal leaves + transformers as it was.""" + monkeypatch.setitem(sys.modules, "transformers.core_model_loading", None) + stock = _patched_state() + + with pytest.raises(RuntimeError, match="could not fuse transformers' expert dequant"): + model_load_utils._patch_finegrained_fp8_dequant() + assert _patched_state() == stock diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index cd2ce19eb91..5aaa5c99ab8 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -25,6 +25,7 @@ import copy import dataclasses import inspect +import json import math import os from contextlib import ExitStack, contextmanager @@ -622,21 +623,41 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( @needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + ("num_experts", "quantization_config", "ep_size", "match"), + [ + (8, {"quant_method": "fp8", "weight_block_size": [128, 128]}, 4, "block-fp8"), + (2, None, 2, "2 expert-data-parallel replicas"), + ( + 2, + {"quant_method": "fp8", "weight_block_size": [128, 128]}, + 2, + "2 expert-data-parallel replicas", + ), + ], + ids=["block-fp8", "several-ep-groups", "block-fp8-several-ep-groups"], +) def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( - run_as, serve, gpu_memory, tmp_path + run_as, serve, gpu_memory, tmp_path, num_experts, quantization_config, ep_size, match ): - """The expert-parallel load cannot shard the dense part within several EP groups, which the - export cannot write: 2 experts over 4 ranks. Left to decide, the plan replicates it however - big, warning with the way out; asked to shard it, it refuses. Without expert parallelism the - rule does not apply.""" + """The expert-parallel load cannot shard the dense part of a block-fp8 checkpoint, which DCP + cannot dequantize, nor within several EP groups, which the export cannot write. Left to decide, + the plan replicates it however big, warning with the way out; asked to shard it, it refuses, + naming the knob that replicates it, and for block FP8 no EP degree, which cannot help. Without + expert parallelism neither rule applies.""" run_as(4) - serve(_qwen3_moe(num_experts=2)) + config = _qwen3_moe(num_experts=num_experts) + if quantization_config: + config.quantization_config = quantization_config + serve(config) gpu_memory(1) - with pytest.warns(UserWarning, match="may run out of memory. Set MODELOPT_EP_DEGREE"): - assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2) - with pytest.raises(NotImplementedError, match="2 expert-data-parallel replicas"): + with pytest.warns(UserWarning, match="may run out of memory"): + assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, ep_size) + with pytest.raises(NotImplementedError, match=f"{match}.*MODELOPT_FSDP_SIZE=1") as refused: resolve_distribution_plan(str(tmp_path), dense="fsdp") + assert ("MODELOPT_EP_DEGREE" in str(refused.value)) is (quantization_config is None) + assert ("unquantized copy" in str(refused.value)) is (quantization_config is not None) assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( 4, dense="fsdp" ) @@ -1902,6 +1923,7 @@ def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts": "moe_tp_experts", } +_BLOCK_FP8 = {"quant_method": "fp8", "weight_block_size": [128, 128]} # DeepSeek-V3's format class _Mesh: @@ -1994,10 +2016,18 @@ def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a( load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) -def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, gloo_world_of_one): +@pytest.mark.parametrize( + "quantization_config", + [None, _BLOCK_FP8], + ids=["unquantized", "block-fp8"], +) +def test_a_hub_id_is_made_local_before_the_weights_load( + monkeypatch, gloo_world_of_one, quantization_config +): """The export reads the tensors the load could not place back from the checkpoint it records, which must be a local directory: a Hub ID there leaves the export nothing to read. The load - resolves it once the plan and the config pass, before anything reads a weight.""" + resolves it once the plan and the config pass, before anything reads a weight: a checkpoint + unquantized or in a format the load dequantizes (block FP8) passes the format check.""" class _ResolvedError(Exception): """Stops the load where it resolves the checkpoint.""" @@ -2014,12 +2044,12 @@ def resolve(path, **hub_options): monkeypatch.setattr(moe, "require_transformers", lambda *args: None) monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", resolve) + config = SimpleNamespace(base_model_ep_plan=_EP_PLAN) + if quantization_config is not None: + config.quantization_config = quantization_config + with pytest.raises(_ResolvedError): - load_and_prepare_ep( - "org/model", - DistributionPlan(2, ep_size=2), - hf_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN), - ) + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) assert resolved == ["org/model"] @@ -2060,6 +2090,18 @@ def resolve(path, **hub_options): ) for method in ("gptq", "modelopt", "compressed-tensors", "mxfp4") ), + ( + SimpleNamespace( + quantization_config={"quant_method": "fp8", "activation_scheme": "static"} + ), + NotImplementedError, + "this fp8 checkpoint names none", + ), + ( + SimpleNamespace(quantization_config=FineGrainedFP8Config()), + NotImplementedError, + "whose 128x128 block is a default", + ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( SimpleNamespace(base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}), @@ -2089,6 +2131,8 @@ def resolve(path, **hub_options): "modelopt", "compressed-tensors", "mxfp4", + "per-tensor-fp8", + "fp8-config-object-of-a-hub-id", "no-ep-plan", "no-routed-experts", "experts-container-only", @@ -2316,28 +2360,24 @@ def test_expert_containers_and_their_children_get_the_expert_state(): @pytest.mark.parametrize( ("config", "expected"), [ - (SimpleNamespace(quantization_config={"quant_method": "fp8"}), "fp8"), + (SimpleNamespace(quantization_config=_BLOCK_FP8), "fp8"), (SimpleNamespace(quantization_config=SimpleNamespace(quant_method="mxfp4")), "mxfp4"), + (SimpleNamespace(text_config=SimpleNamespace(quantization_config=_BLOCK_FP8)), "fp8"), + (SimpleNamespace(llm_config=SimpleNamespace(quantization_config=_BLOCK_FP8)), "fp8"), ( SimpleNamespace( - text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) - ), - "fp8", - ), - ( - SimpleNamespace( - llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + quantization_config={"quant_method": "mxfp4"}, + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), ), - "fp8", + "mxfp4", ), + (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8-unverified"), ( SimpleNamespace( - quantization_config={"quant_method": "mxfp4"}, - text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), + quantization_config={"quant_method": "fp8", "activation_scheme": "static"} ), - "mxfp4", + "fp8-per-tensor", ), - (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), (SimpleNamespace(text_config=None), None), ], @@ -2348,6 +2388,7 @@ def test_expert_containers_and_their_children_get_the_expert_state(): "vlm-llm-config", "top-level-wins", "transformers-config-object", + "per-tensor-fp8", "no-method", "unquantized", ], @@ -2355,10 +2396,34 @@ def test_expert_containers_and_their_children_get_the_expert_state(): def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expected): """The top-level config first, else a VLM's text or LLM config. A loaded model's config holds transformers' config object, not the JSON dict, and callers compare against ``"fp8"`` / - ``"mxfp4"``: it reports the plain method name, not its enum.""" + ``"mxfp4"``: it reports the plain method name, not its enum (an FP8 object, with no checkpoint + to check its block against, is ``"fp8-unverified"``). ``"fp8"`` is block FP8 only: per-tensor + FP8 names the same method without a block size.""" assert _ckpt_quant_method(config) == expected +@pytest.mark.parametrize( + ("on_disk", "expected"), + [ + ({"quant_method": "fp8", "activation_scheme": "static"}, "fp8-per-tensor"), + ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8"), + (None, "fp8-unverified"), + ], + ids=["per-tensor", "block", "no-config-json"], +) +def test_an_fp8_config_object_is_read_against_the_checkpoints_own_config( + tmp_path, on_disk, expected +): + """transformers' FP8 config object names a 128x128 block whatever the checkpoint does, so the + checkpoint's own ``config.json`` says whether it names one: per-tensor FP8 handed in as an + object is still refused. Without that file there is no telling, which the load refuses.""" + if on_disk is not None: + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": on_disk})) + config = SimpleNamespace(quantization_config=FineGrainedFP8Config()) + + assert _ckpt_quant_method(config, str(tmp_path)) == expected + + def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, which shards the experts over every rank when EP is below the world size. During the load it gets @@ -2669,17 +2734,17 @@ def test_expert_param_matcher(plan, name, expected): # ------------------------------------------------------------------ a block-fp8 checkpoint -def test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout( - monkeypatch, gloo_world_of_one -): - """This load does not dequantize block-fp8 checkpoints at all, so that is the refusal a plan - sharding their dense part meets, not one advising to replicate it.""" +def test_a_block_fp8_checkpoints_dense_part_is_not_sharded(monkeypatch, gloo_world_of_one): + """DCP cannot dequantize block-fp8 dense weights, so a plan sharding them is refused before the + load starts, where from_pretrained would dequantize them as it replicated them.""" monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True) monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2) monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - config = SimpleNamespace(quantization_config={"quant_method": "fp8"}) + config = SimpleNamespace( + quantization_config={"quant_method": "fp8", "weight_block_size": [128, 128]} + ) - with pytest.raises(NotImplementedError, match="cannot dequantize a 'fp8' checkpoint"): + with pytest.raises(NotImplementedError, match="block-fp8 checkpoint cannot be sharded"): load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) From 823823d5d991f09c712cbd7cb41098776b9da53e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 06:07:21 +0000 Subject: [PATCH 2/5] Leave weights that are not block FP8 to transformers' own dequant Below transformers 5.10.1 the block-FP8 override replaced the dequant for every weight. From 5.9, transformers' own dequant unpacks packed FP4 (float4_e2m1fn_x2, or int8) with [1, 32] scales, as DeepSeek-V4 stores its routed experts under a block-FP8 config; the override read each packed byte as one value and returned the weight at half its width, without an error. It now hands any weight that is not fp8 to transformers' dequant, in the model dtype, as the 5.10.1 wrapper does, and the docstring no longer says these releases cannot unpack FP4. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 47 ++++++++++++------- .../torch/utils/test_fp8_dequant_patch.py | 32 +++++++++++-- 2 files changed, 56 insertions(+), 23 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c7749ecbcbb..bc4b5b27d46 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1904,15 +1904,17 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): else derived), then crop to the weight shape -- correct for both full and partial last blocks, 2-D dense linears and (leading-dim-broadcast) per-expert tensors. Idempotent. + The override takes block-FP8 weights only. Any other weight goes to transformers' own dequant, + which from 5.9 unpacks packed FP4 (``float4_e2m1fn_x2``, or ``int8``) itself: DeepSeek-V4 declares + block FP8 but stores its routed experts that way, with ``[1, 32]`` scales. + transformers >= 5.10.1 keeps its own dequant. Upstream fixed two of the show-stoppers there -- the loader's expert ``tensor_idx`` check became ``any(op is MergeModulelist ...)`` (bug 1) and - ``_dequantize_one`` gained an ``output_dtype`` the loader passes as the destination dtype (bug 3) -- - and it added native FP4-unpack + ue8m0 (E8M0) scale decoding that our simpler override does NOT do - (installing it on >= 5.10.1 would BREAK ue8m0 checkpoints like DeepSeek-V4). It still derives the - block from the scale grid (bug 2), which a partial last block makes wrong: 576 rows under a 5-row - grid do not divide at all, and 448 rows under a 4-row grid divide into blocks of 112 instead of 128. - So there a weight whose grid is the configured block's ceil grid and that has a partial last block - is padded to whole blocks for it and the result cropped back. + ``_dequantize_one`` gained an ``output_dtype`` the loader passes as the destination dtype (bug 3). + It still derives the block from the scale grid (bug 2), which a partial last block makes wrong: 576 + rows under a 5-row grid do not divide at all, and 448 rows under a 4-row grid divide into blocks of + 112 instead of 128. So there a block-FP8 weight whose grid is the configured block's ceil grid and + that has a partial last block is padded to whole blocks for it and the result cropped back. """ global _FINEGRAINED_FP8_PATCHED, _FP8_DEQUANT_DTYPE if isinstance(target_dtype, str): @@ -1928,6 +1930,11 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): from packaging.version import parse from transformers.integrations import finegrained_fp8 as _fp8 + fp8_dtypes = { + getattr(torch, name) + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") + if hasattr(torch, name) + } if parse(transformers.__version__) >= parse("5.10.1"): # Re-importing this module loses its flag, and the wrapper calls the method it replaces: # the class keeps a flag of its own, so the wrapper never wraps itself. @@ -1935,11 +1942,6 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): _FINEGRAINED_FP8_PATCHED = True return _stock_dequantize_one = _fp8.Fp8Dequantize._dequantize_one - fp8_dtypes = { - getattr(torch, name) - for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") - if hasattr(torch, name) - } def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): if quantized.dim() < 2 or scales.dim() < 2 or quantized.dtype not in fp8_dtypes: @@ -1967,8 +1969,22 @@ def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): _FINEGRAINED_FP8_PATCHED = True return + # The release's own dequant, for the weights the override leaves to it. Releases without one + # (5.5) dequantize inline and never call the override. + stock = vars(_fp8.Fp8Dequantize).get("_dequantize_one") + def _dequantize_one(self, quantized, scales, output_dtype=None): - # Block FP8 only: the releases this replaces have no FP4 unpack (5.10.1 added it). + # Emit the MODEL compute dtype (bf16), NOT the scale dtype. DeepSeek's weight_scale_inv is F32; + # transformers' original out_dtype logic would return F32 weights, which then run against bf16 + # activations -> "mat1 and mat2 have the same dtype" RuntimeError in the calibration matmuls. + # A loader that passes the destination dtype, as transformers' does from 5.10.1, knows it best. + out_dtype = output_dtype or _FP8_DEQUANT_DTYPE or torch.bfloat16 + if quantized.dtype not in fp8_dtypes: + # Block FP8 only: read as fp8, packed FP4 (two values a byte) would come out at half its + # width and wrong. transformers' own dequant unpacks it. + if stock is None: + raise ValueError(f"block-FP8 dequant: a {quantized.dtype} weight is not block FP8.") + return stock(self, quantized, scales).to(out_dtype) if quantized.dim() < 2 or scales.dim() < 2: raise ValueError( f"block-FP8 dequant: weight {tuple(quantized.shape)} with scales " @@ -1977,11 +1993,6 @@ def _dequantize_one(self, quantized, scales, output_dtype=None): q = quantized.to(torch.float32, copy=True) rows, cols = q.shape[-2:] scale_rows, scale_cols = scales.shape[-2:] - # Emit the MODEL compute dtype (bf16), NOT the scale dtype. DeepSeek's weight_scale_inv is F32; - # transformers' original out_dtype logic would return F32 weights, which then run against bf16 - # activations -> "mat1 and mat2 have the same dtype" RuntimeError in the calibration matmuls. - # A loader that passes the destination dtype, as transformers' does from 5.10.1, knows it best. - out_dtype = output_dtype or _FP8_DEQUANT_DTYPE or torch.bfloat16 if scale_rows == 0 or scale_cols == 0 or scales.numel() == 0: # Only an empty weight has no scales. A non-empty one with an empty grid was sliced # apart from its scales (as 5.9's expert tensor_idx bug did): emitting it unscaled would diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index d9e378f66c1..87f177628ef 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -212,6 +212,28 @@ def test_dequant_refuses_a_scale_grid_it_cannot_place(pre_5_10_transformers, gri _dequant_op((4, 4))._dequantize_one(weight, scale_inv) +@pytest.mark.parametrize("packed_dtype", ["float4_e2m1fn_x2", "int8"]) +def test_a_weight_that_is_not_fp8_keeps_transformers_dequant(pre_5_10_transformers, packed_dtype): + """DeepSeek-V4 declares block FP8 but stores its routed experts as packed FP4, two values a byte, + with ``[1, 32]`` scales, which transformers' own dequant unpacks from 5.9. Read as fp8, a + (256, 512) weight would come out (256, 256) and wrong: it goes to transformers' dequant, and + comes out in the model dtype.""" + stock = vars(Fp8Dequantize).get("_dequantize_one") + dtype = getattr(torch, packed_dtype, None) + if stock is None or dtype is None: + pytest.skip("needs transformers' own FP4 unpack (5.9 on) and torch's packed dtype") + gen = torch.Generator().manual_seed(0) + packed = torch.randint(0, 256, (256, 256), dtype=torch.uint8, generator=gen).view(dtype) + scales = torch.rand(256, 16, generator=gen) + 0.5 # one per 32 unpacked values + expected = stock(_dequant_op((128, 128)), packed, scales) + + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + out = _dequant_op((128, 128))._dequantize_one(packed, scales) + + assert out.shape == (256, 512) + assert torch.equal(out, expected.to(torch.bfloat16)) + + def test_the_expert_converter_leads_with_a_merge_that_dequantizes_each_expert( pre_5_10_transformers, ): @@ -340,11 +362,11 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( pre_5_10_transformers, monkeypatch, shape, grid ): - """Upstream fixed the expert converter there and decodes ue8m0 and FP4 scales, which the override - does not, so its dequant is kept and only wrapped. A partial last block (DeepSeek-V3's 576-row - kv_a_proj_with_mqa under a 5-row grid of 128-row blocks, or 192 rows that its own dequant would - split into two blocks of 96) is padded to whole blocks for it and cropped back; a weight of whole - blocks goes to it as is.""" + """Upstream fixed the expert converter there and emits the destination dtype, so its dequant is + kept and only wrapped. A partial last block (DeepSeek-V3's 576-row kv_a_proj_with_mqa under a + 5-row grid of 128-row blocks, or 192 rows that its own dequant would split into two blocks of + 96) is padded to whole blocks for it and cropped back; a weight of whole blocks goes to it as + is.""" if Version("5.10.1") > _INSTALLED_TRANSFORMERS: pytest.skip("needs transformers' own dequant from 5.10.1") monkeypatch.setattr(transformers, "__version__", "5.10.1") From e78c5fd2930f8961914a7c5fd917eabfb79b351a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 08:08:14 +0000 Subject: [PATCH 3/5] Move the block-FP8 dequant tests into test_model_load_utils.py test_fp8_dequant_patch.py covers model_load_utils' block-FP8 dequant patch alone, so its tests join that module's test file, as CONTRIBUTING.md asks, as one section with their fixtures, marks and parametrization. Its module-level skip on a torch without float8_e4m3fn becomes a mark on its tests. Signed-off-by: Shengliang Xu --- .../torch/utils/test_fp8_dequant_patch.py | 391 ------------------ .../unit/torch/utils/test_model_load_utils.py | 376 +++++++++++++++++ 2 files changed, 376 insertions(+), 391 deletions(-) delete mode 100644 tests/unit/torch/utils/test_fp8_dequant_patch.py diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py deleted file mode 100644 index 87f177628ef..00000000000 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ /dev/null @@ -1,391 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -"""``_patch_finegrained_fp8_dequant``: transformers' block-FP8 dequant, fixed for DeepSeek-V3. - -transformers 5.9's ``Fp8Dequantize._dequantize_one`` finds the block size by dividing the weight -by the scale grid, which breaks on DeepSeek-V3's partial last blocks (a 576-row weight under a -5-row grid of 128-row blocks). The patch installs a tile-and-crop dequant instead, and only below -transformers 5.10.1, so the tests present an older version to it. The dequant is checked against a -plain blockwise reference: ``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``. -""" - -import sys -from types import SimpleNamespace - -import pytest -import torch -from packaging.version import Version - -transformers = pytest.importorskip("transformers") -finegrained_fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8") -quantizer_fp8 = pytest.importorskip("transformers.quantizers.quantizer_finegrained_fp8") - -from transformers import FineGrainedFP8Config -from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter - -import modelopt.torch.utils.plugins.model_load_utils as model_load_utils - -if not hasattr(torch, "float8_e4m3fn"): - pytest.skip("torch has no float8_e4m3fn", allow_module_level=True) - -# The installed version: the fixture below makes the patch see an older one. -_INSTALLED_TRANSFORMERS = Version(transformers.__version__) -Fp8Dequantize = finegrained_fp8.Fp8Dequantize -FineGrainedFP8HfQuantizer = quantizer_fp8.FineGrainedFP8HfQuantizer - -_MISSING = object() -# What the patch may rebind below transformers 5.10.1: the dequant method, and the FP8 quantizer's -# converter hook plus the flag guarding it. From 5.10.1 it wraps the dequant method, guarded by a -# flag of its own. -_PATCH_TARGETS = ( - (Fp8Dequantize, "_dequantize_one"), - (FineGrainedFP8HfQuantizer, "update_weight_conversions"), - (FineGrainedFP8HfQuantizer, "_modelopt_uwc_patched"), -) -_RESTORED = (*_PATCH_TARGETS, (Fp8Dequantize, "_modelopt_dequant_patched")) - - -@pytest.fixture -def pre_5_10_transformers(monkeypatch): - """Unpatched transformers, reporting a version the patch targets; restored afterwards.""" - saved = [(cls, name, vars(cls).get(name, _MISSING)) for cls, name in _RESTORED] - monkeypatch.setattr(transformers, "__version__", "5.9.0") - monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) - monkeypatch.setattr(model_load_utils, "_FP8_DEQUANT_DTYPE", None) - yield - for cls, name, value in saved: - if value is not _MISSING: - setattr(cls, name, value) - elif name in vars(cls): - delattr(cls, name) - - -def _patched_state(): - return [vars(cls).get(name, _MISSING) for cls, name in _PATCH_TARGETS] - - -def _dequant_op(block_size): - """The dequant op of an FP8 quantizer configured with ``weight_block_size=block_size``.""" - config = SimpleNamespace(weight_block_size=block_size) - return Fp8Dequantize(SimpleNamespace(quantization_config=config)) - - -def _block_fp8(shape, grid, seed=0): - """A random fp8 weight and its fp32 ``scale_inv`` grid.""" - gen = torch.Generator().manual_seed(seed) - weight = (4 * torch.randn(shape, generator=gen)).to(torch.float8_e4m3fn) - return weight, torch.rand(grid, generator=gen) + 0.5 - - -def _ceil_grid(shape, block): - *lead, rows, cols = shape - return (*lead, -(-rows // block[0]), -(-cols // block[1])) - - -def _reference(weight, scale_inv, block): - """``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``, in fp32.""" - rows, cols = weight.shape[-2:] - scale = scale_inv[..., torch.arange(rows) // block[0], :][..., torch.arange(cols) // block[1]] - return weight.float() * scale - - -@pytest.mark.parametrize( - ("shape", "block", "configured", "grid"), - [ - ((8, 12), (4, 4), (4, 4), None), - ((10, 13), (4, 4), (4, 4), None), - ((8, 3), (4, 4), (4, 4), None), - ((9, 8), (4, 4), (4, 4), None), - ((6, 10), (2, 4), (2, 4), None), - ((3, 10, 13), (4, 4), (4, 4), None), - ((576, 7168), (128, 128), (128, 128), None), - ((8, 8), (4, 2), (128, 128), (2, 4)), - ((130, 260), (128, 128), None, (2, 3)), - ((0, 13), (4, 4), (4, 4), None), - ], - ids=[ - "whole-blocks", - "partial-last-blocks", - "narrower-than-a-block", - "grid-divides-the-rows", - "non-square-blocks", - "per-expert", - "deepseek-kv-a-proj", - "grid-finer-than-configured", - "nothing-configured", - "empty-weight", - ], -) -def test_dequant_matches_the_blockwise_reference( - pre_5_10_transformers, shape, block, configured, grid -): - """``configured`` is the quantization config's block size, ``grid`` the scale grid when it is not - that block's ceil grid. ``grid-divides-the-rows``: 9 rows of 4-row blocks make a 3-row grid, - which also divides 9, yet the blocks are 4 rows, not 3. A grid that is not the configured - block's ceil grid sets the block by division; with nothing configured, DeepSeek's 128 x 128 - applies.""" - weight, scale_inv = _block_fp8(shape, grid or _ceil_grid(shape, block)) - model_load_utils._patch_finegrained_fp8_dequant(torch.float32) - - out = _dequant_op(configured)._dequantize_one(weight, scale_inv) - - assert torch.equal(out, _reference(weight, scale_inv, block)) - - -def test_the_dequantize_config_keeps_the_checkpoints_block_size(): - """The expert-parallel load passes ``FineGrainedFP8Config(dequantize=True)``, and every block - decision of the dequant reads ``weight_block_size`` off the quantizer's config: transformers - keeps the checkpoint's, and takes only the loading attributes from the one passed.""" - from transformers.quantizers.auto import AutoHfQuantizer - - checkpoint = { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "weight_block_size": [64, 64], - } - merged = AutoHfQuantizer.merge_quantization_configs( - checkpoint, FineGrainedFP8Config(dequantize=True) - ) - - assert list(merged.weight_block_size) == [64, 64] - assert merged.dequantize - - -@pytest.mark.parametrize( - ("target_dtype", "output_dtype", "expected"), - [ - (None, None, torch.bfloat16), - (torch.float16, None, torch.float16), - ("float16", None, torch.float16), - ("auto", None, torch.bfloat16), - (torch.bfloat16, torch.float16, torch.float16), - ], - ids=["unrequested", "requested", "named", "auto", "passed-by-the-loader"], -) -def test_dequant_emits_the_requested_model_dtype( - pre_5_10_transformers, target_dtype, output_dtype, expected -): - """Not the scales' dtype: DeepSeek's fp32 ``weight_scale_inv`` would give fp32 weights that - then meet bf16 activations in the calibration matmuls. Unrequested, it is bf16. A config can - name its dtype, or say "auto", which is bf16 too. A loader that passes the destination dtype, - as transformers' does from 5.10.1, gets that one.""" - weight, scale_inv = _block_fp8((10, 13), (3, 4)) - model_load_utils._patch_finegrained_fp8_dequant(target_dtype) - - out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv, output_dtype=output_dtype) - - assert out.dtype == expected - assert torch.equal(out, _reference(weight, scale_inv, (4, 4)).to(expected)) - - -@pytest.mark.parametrize( - ("grid", "match"), - [ - ((0, 4), "empty scale grid"), - ((3, 0), "empty scale grid"), - ((4, 4), "fits neither"), - ((4,), "not a block-scaled 2-D weight"), - ], - ids=["no-row-scales", "no-col-scales", "grid-no-block-places", "one-dim-scales"], -) -def test_dequant_refuses_a_scale_grid_it_cannot_place(pre_5_10_transformers, grid, match): - """Rather than load wrong weights: an empty grid for a non-empty weight (the weight and its - scales were sliced apart, as 5.9's expert ``tensor_idx`` bug did), a 4-row grid over 10 rows, - which no 4-row block (the configured one) or even split gives, or scales that are no grid.""" - weight, scale_inv = _block_fp8((10, 13), grid) - model_load_utils._patch_finegrained_fp8_dequant(torch.float32) - - with pytest.raises(ValueError, match=match): - _dequant_op((4, 4))._dequantize_one(weight, scale_inv) - - -@pytest.mark.parametrize("packed_dtype", ["float4_e2m1fn_x2", "int8"]) -def test_a_weight_that_is_not_fp8_keeps_transformers_dequant(pre_5_10_transformers, packed_dtype): - """DeepSeek-V4 declares block FP8 but stores its routed experts as packed FP4, two values a byte, - with ``[1, 32]`` scales, which transformers' own dequant unpacks from 5.9. Read as fp8, a - (256, 512) weight would come out (256, 256) and wrong: it goes to transformers' dequant, and - comes out in the model dtype.""" - stock = vars(Fp8Dequantize).get("_dequantize_one") - dtype = getattr(torch, packed_dtype, None) - if stock is None or dtype is None: - pytest.skip("needs transformers' own FP4 unpack (5.9 on) and torch's packed dtype") - gen = torch.Generator().manual_seed(0) - packed = torch.randint(0, 256, (256, 256), dtype=torch.uint8, generator=gen).view(dtype) - scales = torch.rand(256, 16, generator=gen) + 0.5 # one per 32 unpacked values - expected = stock(_dequant_op((128, 128)), packed, scales) - - model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) - out = _dequant_op((128, 128))._dequantize_one(packed, scales) - - assert out.shape == (256, 512) - assert torch.equal(out, expected.to(torch.bfloat16)) - - -def test_the_expert_converter_leads_with_a_merge_that_dequantizes_each_expert( - pre_5_10_transformers, -): - """transformers prepends ``Fp8Dequantize`` to the routed-expert converter, so its first op is - no longer the ``MergeModulelist`` the loader keys each expert's ``tensor_idx`` on, and the EP - load then slices every expert by rows. The patch fuses the two into one ``MergeModulelist`` - that dequantizes each expert with its own scales and drops them before ``Concatenate``.""" - if not hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions"): - pytest.skip("this transformers' FP8 quantizer prepends no dequant to the expert converter") - model_load_utils._patch_finegrained_fp8_dequant(torch.float32) - experts = 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)], - ) - quantizer = FineGrainedFP8HfQuantizer( - FineGrainedFP8Config(weight_block_size=(4, 4), dequantize=True), pre_quantized=True - ) - - [converter] = [ - conv - for conv in quantizer.update_weight_conversions([experts]) - if conv.target_patterns == ["mlp.experts.gate_up_proj"] - ] - - assert isinstance(converter.operations[0], MergeModulelist) - blocks = { # every expert's weight and scales differ, so a mixed-up pair cannot pass - name: [_block_fp8((8, 12), (2, 3), seed=3 * p + e) for e in range(3)] - for p, name in enumerate(("gate_proj", "up_proj")) - } - tensors = {} # what the loader collects for each of the converter's source patterns - for name, pairs in blocks.items(): - tensors[f"mlp.experts.*.{name}.weight$"] = [weight for weight, _ in pairs] - tensors[f"mlp.experts.*.{name}.weight_scale_inv$"] = [scale_inv for _, scale_inv in pairs] - for op in converter.operations: - tensors = op.convert( - tensors, - source_patterns=converter.source_patterns, - target_patterns=converter.target_patterns, - ) - [fused] = tensors.values() - expected = torch.cat( - [ - torch.stack([_reference(weight, scale_inv, (4, 4)) for weight, scale_inv in pairs]) - for pairs in blocks.values() - ], - dim=1, - ) - assert torch.equal(fused, expected.to(fused.dtype)) - - -def test_patching_again_keeps_the_install_and_takes_the_new_dtype(pre_5_10_transformers): - """Every block-fp8 load patches again, with its own model dtype.""" - model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) - installed = _patched_state() - - model_load_utils._patch_finegrained_fp8_dequant(torch.float16) - - assert _patched_state() == installed - weight, scale_inv = _block_fp8((10, 13), (3, 4)) - assert _dequant_op((4, 4))._dequantize_one(weight, scale_inv).dtype == torch.float16 - - -@pytest.mark.parametrize( - ("version", "cls", "name"), - [ - ("5.9.0", FineGrainedFP8HfQuantizer, "update_weight_conversions"), - ("5.10.1", Fp8Dequantize, "_dequantize_one"), - ], - ids=["quantizer-hook", "dequant-wrapper"], -) -def test_a_reimported_module_does_not_wrap_a_patch_twice( - pre_5_10_transformers, monkeypatch, version, cls, name -): - """Re-importing ``model_load_utils`` loses its module flag; the patched class carries its own: - the FP8 quantizer for its converter hook, and from 5.10.1 the dequant op for the wrapper that - calls the method it replaced.""" - if not hasattr(cls, name) or Version(version) > _INSTALLED_TRANSFORMERS: - pytest.skip(f"this transformers has no {cls.__name__}.{name} to patch as {version} would") - monkeypatch.setattr(transformers, "__version__", version) - model_load_utils._patch_finegrained_fp8_dequant() - installed = vars(cls)[name] - monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) - - model_load_utils._patch_finegrained_fp8_dequant() - - assert vars(cls)[name] is installed - - -def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, monkeypatch): - """The patch rebinds only the dequant method and, where the quantizer has one, its converter - hook and guard flag. Putting those back restores the stock classes exactly, after which the - patch installs again once the module flag is cleared.""" - has_hook = hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions") - classes = (Fp8Dequantize, FineGrainedFP8HfQuantizer) - stock = {cls: dict(vars(cls)) for cls in classes} - - model_load_utils._patch_finegrained_fp8_dequant() - - changed = { - (cls, name) - for cls in classes - for name in stock[cls].keys() | vars(cls).keys() - if vars(cls).get(name, _MISSING) is not stock[cls].get(name, _MISSING) - } - assert changed == set(_PATCH_TARGETS if has_hook else _PATCH_TARGETS[:1]) - for cls, name in changed: - if name in stock[cls]: - setattr(cls, name, stock[cls][name]) - else: - delattr(cls, name) - assert all(dict(vars(cls)) == stock[cls] for cls in classes) - - monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) - model_load_utils._patch_finegrained_fp8_dequant() - assert vars(Fp8Dequantize).get("_dequantize_one") is not stock[Fp8Dequantize].get( - "_dequantize_one" - ) - - -@pytest.mark.parametrize( - ("shape", "grid"), - [((576, 7168), (5, 56)), ((192, 256), (2, 2)), ((512, 256), (4, 2))], - ids=["partial", "partial-dividing-the-grid", "whole"], -) -def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( - pre_5_10_transformers, monkeypatch, shape, grid -): - """Upstream fixed the expert converter there and emits the destination dtype, so its dequant is - kept and only wrapped. A partial last block (DeepSeek-V3's 576-row kv_a_proj_with_mqa under a - 5-row grid of 128-row blocks, or 192 rows that its own dequant would split into two blocks of - 96) is padded to whole blocks for it and cropped back; a weight of whole blocks goes to it as - is.""" - if Version("5.10.1") > _INSTALLED_TRANSFORMERS: - pytest.skip("needs transformers' own dequant from 5.10.1") - monkeypatch.setattr(transformers, "__version__", "5.10.1") - hook = vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) - weight, scale_inv = _block_fp8(shape, grid) - - model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) - out = _dequant_op((128, 128))._dequantize_one(weight, scale_inv, torch.float32) - - assert torch.equal(out, _reference(weight, scale_inv, (128, 128))) - assert vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) is hook - - -def test_an_expert_converter_it_cannot_fuse_stops_the_load(pre_5_10_transformers, monkeypatch): - """Unpatched, every routed expert would load sliced apart from its scales. The refusal leaves - transformers as it was.""" - monkeypatch.setitem(sys.modules, "transformers.core_model_loading", None) - stock = _patched_state() - - with pytest.raises(RuntimeError, match="could not fuse transformers' expert dequant"): - model_load_utils._patch_finegrained_fp8_dequant() - assert _patched_state() == stock diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 5aaa5c99ab8..c918dbb6b39 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -28,6 +28,7 @@ import json import math import os +import sys from contextlib import ExitStack, contextmanager from types import SimpleNamespace from unittest import mock @@ -46,7 +47,9 @@ transformers = pytest.importorskip("transformers") import transformers.core_model_loading as core_model_loading import transformers.distributed as hf_distributed +import transformers.integrations.finegrained_fp8 as finegrained_fp8 import transformers.modeling_utils as modeling_utils +import transformers.quantizers.quantizer_finegrained_fp8 as quantizer_fp8 from _test_utils.torch.transformers_models import ( _TinyQwen35Tokenizer, get_tiny_llama, @@ -73,6 +76,7 @@ T5Config, WhisperConfig, ) +from transformers.core_model_loading import Concatenate, MergeModulelist, WeightConverter from transformers.distributed import DistributedConfig import modelopt.torch.models.hf as models_hf @@ -3374,3 +3378,375 @@ def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch): _ep_aware_caching_allocator_warmup(4, _A2A_PLAN), ): assert not hasattr(modeling_utils, "get_total_byte_count") + + +# ----------------------------------------------------- the block-FP8 dequant, fixed for DeepSeek-V3 +# transformers 5.9's ``Fp8Dequantize._dequantize_one`` finds the block size by dividing the weight +# by the scale grid, which breaks on DeepSeek-V3's partial last blocks (a 576-row weight under a +# 5-row grid of 128-row blocks). ``_patch_finegrained_fp8_dequant`` installs a tile-and-crop dequant +# instead, and only below transformers 5.10.1, so the tests present an older version to it. The +# dequant is checked against a plain blockwise reference: +# ``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``. + +needs_float8 = pytest.mark.skipif( + not hasattr(torch, "float8_e4m3fn"), reason="torch has no float8_e4m3fn" +) +# The installed version: the fixture below makes the patch see an older one. +_INSTALLED_TRANSFORMERS = Version(transformers.__version__) +Fp8Dequantize = finegrained_fp8.Fp8Dequantize +FineGrainedFP8HfQuantizer = quantizer_fp8.FineGrainedFP8HfQuantizer + +_MISSING = object() +# What the patch may rebind below transformers 5.10.1: the dequant method, and the FP8 quantizer's +# converter hook plus the flag guarding it. From 5.10.1 it wraps the dequant method, guarded by a +# flag of its own. +_PATCH_TARGETS = ( + (Fp8Dequantize, "_dequantize_one"), + (FineGrainedFP8HfQuantizer, "update_weight_conversions"), + (FineGrainedFP8HfQuantizer, "_modelopt_uwc_patched"), +) +_RESTORED = (*_PATCH_TARGETS, (Fp8Dequantize, "_modelopt_dequant_patched")) + + +@pytest.fixture +def pre_5_10_transformers(monkeypatch): + """Unpatched transformers, reporting a version the patch targets; restored afterwards.""" + saved = [(cls, name, vars(cls).get(name, _MISSING)) for cls, name in _RESTORED] + monkeypatch.setattr(transformers, "__version__", "5.9.0") + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + monkeypatch.setattr(model_load_utils, "_FP8_DEQUANT_DTYPE", None) + yield + for cls, name, value in saved: + if value is not _MISSING: + setattr(cls, name, value) + elif name in vars(cls): + delattr(cls, name) + + +def _patched_state(): + return [vars(cls).get(name, _MISSING) for cls, name in _PATCH_TARGETS] + + +def _dequant_op(block_size): + """The dequant op of an FP8 quantizer configured with ``weight_block_size=block_size``.""" + config = SimpleNamespace(weight_block_size=block_size) + return Fp8Dequantize(SimpleNamespace(quantization_config=config)) + + +def _block_fp8(shape, grid, seed=0): + """A random fp8 weight and its fp32 ``scale_inv`` grid.""" + gen = torch.Generator().manual_seed(seed) + weight = (4 * torch.randn(shape, generator=gen)).to(torch.float8_e4m3fn) + return weight, torch.rand(grid, generator=gen) + 0.5 + + +def _ceil_grid(shape, block): + *lead, rows, cols = shape + return (*lead, -(-rows // block[0]), -(-cols // block[1])) + + +def _reference(weight, scale_inv, block): + """``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``, in fp32.""" + rows, cols = weight.shape[-2:] + scale = scale_inv[..., torch.arange(rows) // block[0], :][..., torch.arange(cols) // block[1]] + return weight.float() * scale + + +@needs_float8 +@pytest.mark.parametrize( + ("shape", "block", "configured", "grid"), + [ + ((8, 12), (4, 4), (4, 4), None), + ((10, 13), (4, 4), (4, 4), None), + ((8, 3), (4, 4), (4, 4), None), + ((9, 8), (4, 4), (4, 4), None), + ((6, 10), (2, 4), (2, 4), None), + ((3, 10, 13), (4, 4), (4, 4), None), + ((576, 7168), (128, 128), (128, 128), None), + ((8, 8), (4, 2), (128, 128), (2, 4)), + ((130, 260), (128, 128), None, (2, 3)), + ((0, 13), (4, 4), (4, 4), None), + ], + ids=[ + "whole-blocks", + "partial-last-blocks", + "narrower-than-a-block", + "grid-divides-the-rows", + "non-square-blocks", + "per-expert", + "deepseek-kv-a-proj", + "grid-finer-than-configured", + "nothing-configured", + "empty-weight", + ], +) +def test_dequant_matches_the_blockwise_reference( + pre_5_10_transformers, shape, block, configured, grid +): + """``configured`` is the quantization config's block size, ``grid`` the scale grid when it is not + that block's ceil grid. ``grid-divides-the-rows``: 9 rows of 4-row blocks make a 3-row grid, + which also divides 9, yet the blocks are 4 rows, not 3. A grid that is not the configured + block's ceil grid sets the block by division; with nothing configured, DeepSeek's 128 x 128 + applies.""" + weight, scale_inv = _block_fp8(shape, grid or _ceil_grid(shape, block)) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + + out = _dequant_op(configured)._dequantize_one(weight, scale_inv) + + assert torch.equal(out, _reference(weight, scale_inv, block)) + + +@needs_float8 +def test_the_dequantize_config_keeps_the_checkpoints_block_size(): + """The expert-parallel load passes ``FineGrainedFP8Config(dequantize=True)``, and every block + decision of the dequant reads ``weight_block_size`` off the quantizer's config: transformers + keeps the checkpoint's, and takes only the loading attributes from the one passed.""" + from transformers.quantizers.auto import AutoHfQuantizer + + checkpoint = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [64, 64], + } + merged = AutoHfQuantizer.merge_quantization_configs( + checkpoint, FineGrainedFP8Config(dequantize=True) + ) + + assert list(merged.weight_block_size) == [64, 64] + assert merged.dequantize + + +@needs_float8 +@pytest.mark.parametrize( + ("target_dtype", "output_dtype", "expected"), + [ + (None, None, torch.bfloat16), + (torch.float16, None, torch.float16), + ("float16", None, torch.float16), + ("auto", None, torch.bfloat16), + (torch.bfloat16, torch.float16, torch.float16), + ], + ids=["unrequested", "requested", "named", "auto", "passed-by-the-loader"], +) +def test_dequant_emits_the_requested_model_dtype( + pre_5_10_transformers, target_dtype, output_dtype, expected +): + """Not the scales' dtype: DeepSeek's fp32 ``weight_scale_inv`` would give fp32 weights that + then meet bf16 activations in the calibration matmuls. Unrequested, it is bf16. A config can + name its dtype, or say "auto", which is bf16 too. A loader that passes the destination dtype, + as transformers' does from 5.10.1, gets that one.""" + weight, scale_inv = _block_fp8((10, 13), (3, 4)) + model_load_utils._patch_finegrained_fp8_dequant(target_dtype) + + out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv, output_dtype=output_dtype) + + assert out.dtype == expected + assert torch.equal(out, _reference(weight, scale_inv, (4, 4)).to(expected)) + + +@needs_float8 +@pytest.mark.parametrize( + ("grid", "match"), + [ + ((0, 4), "empty scale grid"), + ((3, 0), "empty scale grid"), + ((4, 4), "fits neither"), + ((4,), "not a block-scaled 2-D weight"), + ], + ids=["no-row-scales", "no-col-scales", "grid-no-block-places", "one-dim-scales"], +) +def test_dequant_refuses_a_scale_grid_it_cannot_place(pre_5_10_transformers, grid, match): + """Rather than load wrong weights: an empty grid for a non-empty weight (the weight and its + scales were sliced apart, as 5.9's expert ``tensor_idx`` bug did), a 4-row grid over 10 rows, + which no 4-row block (the configured one) or even split gives, or scales that are no grid.""" + weight, scale_inv = _block_fp8((10, 13), grid) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + + with pytest.raises(ValueError, match=match): + _dequant_op((4, 4))._dequantize_one(weight, scale_inv) + + +@needs_float8 +@pytest.mark.parametrize("packed_dtype", ["float4_e2m1fn_x2", "int8"]) +def test_a_weight_that_is_not_fp8_keeps_transformers_dequant(pre_5_10_transformers, packed_dtype): + """DeepSeek-V4 declares block FP8 but stores its routed experts as packed FP4, two values a byte, + with ``[1, 32]`` scales, which transformers' own dequant unpacks from 5.9. Read as fp8, a + (256, 512) weight would come out (256, 256) and wrong: it goes to transformers' dequant, and + comes out in the model dtype.""" + stock = vars(Fp8Dequantize).get("_dequantize_one") + dtype = getattr(torch, packed_dtype, None) + if stock is None or dtype is None: + pytest.skip("needs transformers' own FP4 unpack (5.9 on) and torch's packed dtype") + gen = torch.Generator().manual_seed(0) + packed = torch.randint(0, 256, (256, 256), dtype=torch.uint8, generator=gen).view(dtype) + scales = torch.rand(256, 16, generator=gen) + 0.5 # one per 32 unpacked values + expected = stock(_dequant_op((128, 128)), packed, scales) + + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + out = _dequant_op((128, 128))._dequantize_one(packed, scales) + + assert out.shape == (256, 512) + assert torch.equal(out, expected.to(torch.bfloat16)) + + +@needs_float8 +def test_the_expert_converter_leads_with_a_merge_that_dequantizes_each_expert( + pre_5_10_transformers, +): + """transformers prepends ``Fp8Dequantize`` to the routed-expert converter, so its first op is + no longer the ``MergeModulelist`` the loader keys each expert's ``tensor_idx`` on, and the EP + load then slices every expert by rows. The patch fuses the two into one ``MergeModulelist`` + that dequantizes each expert with its own scales and drops them before ``Concatenate``.""" + if not hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions"): + pytest.skip("this transformers' FP8 quantizer prepends no dequant to the expert converter") + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + experts = 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)], + ) + quantizer = FineGrainedFP8HfQuantizer( + FineGrainedFP8Config(weight_block_size=(4, 4), dequantize=True), pre_quantized=True + ) + + [converter] = [ + conv + for conv in quantizer.update_weight_conversions([experts]) + if conv.target_patterns == ["mlp.experts.gate_up_proj"] + ] + + assert isinstance(converter.operations[0], MergeModulelist) + blocks = { # every expert's weight and scales differ, so a mixed-up pair cannot pass + name: [_block_fp8((8, 12), (2, 3), seed=3 * p + e) for e in range(3)] + for p, name in enumerate(("gate_proj", "up_proj")) + } + tensors = {} # what the loader collects for each of the converter's source patterns + for name, pairs in blocks.items(): + tensors[f"mlp.experts.*.{name}.weight$"] = [weight for weight, _ in pairs] + tensors[f"mlp.experts.*.{name}.weight_scale_inv$"] = [scale_inv for _, scale_inv in pairs] + for op in converter.operations: + tensors = op.convert( + tensors, + source_patterns=converter.source_patterns, + target_patterns=converter.target_patterns, + ) + [fused] = tensors.values() + expected = torch.cat( + [ + torch.stack([_reference(weight, scale_inv, (4, 4)) for weight, scale_inv in pairs]) + for pairs in blocks.values() + ], + dim=1, + ) + assert torch.equal(fused, expected.to(fused.dtype)) + + +@needs_float8 +def test_patching_again_keeps_the_install_and_takes_the_new_dtype(pre_5_10_transformers): + """Every block-fp8 load patches again, with its own model dtype.""" + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + installed = _patched_state() + + model_load_utils._patch_finegrained_fp8_dequant(torch.float16) + + assert _patched_state() == installed + weight, scale_inv = _block_fp8((10, 13), (3, 4)) + assert _dequant_op((4, 4))._dequantize_one(weight, scale_inv).dtype == torch.float16 + + +@needs_float8 +@pytest.mark.parametrize( + ("version", "cls", "name"), + [ + ("5.9.0", FineGrainedFP8HfQuantizer, "update_weight_conversions"), + ("5.10.1", Fp8Dequantize, "_dequantize_one"), + ], + ids=["quantizer-hook", "dequant-wrapper"], +) +def test_a_reimported_module_does_not_wrap_a_patch_twice( + pre_5_10_transformers, monkeypatch, version, cls, name +): + """Re-importing ``model_load_utils`` loses its module flag; the patched class carries its own: + the FP8 quantizer for its converter hook, and from 5.10.1 the dequant op for the wrapper that + calls the method it replaced.""" + if not hasattr(cls, name) or Version(version) > _INSTALLED_TRANSFORMERS: + pytest.skip(f"this transformers has no {cls.__name__}.{name} to patch as {version} would") + monkeypatch.setattr(transformers, "__version__", version) + model_load_utils._patch_finegrained_fp8_dequant() + installed = vars(cls)[name] + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + + model_load_utils._patch_finegrained_fp8_dequant() + + assert vars(cls)[name] is installed + + +@needs_float8 +def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, monkeypatch): + """The patch rebinds only the dequant method and, where the quantizer has one, its converter + hook and guard flag. Putting those back restores the stock classes exactly, after which the + patch installs again once the module flag is cleared.""" + has_hook = hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions") + classes = (Fp8Dequantize, FineGrainedFP8HfQuantizer) + stock = {cls: dict(vars(cls)) for cls in classes} + + model_load_utils._patch_finegrained_fp8_dequant() + + changed = { + (cls, name) + for cls in classes + for name in stock[cls].keys() | vars(cls).keys() + if vars(cls).get(name, _MISSING) is not stock[cls].get(name, _MISSING) + } + assert changed == set(_PATCH_TARGETS if has_hook else _PATCH_TARGETS[:1]) + for cls, name in changed: + if name in stock[cls]: + setattr(cls, name, stock[cls][name]) + else: + delattr(cls, name) + assert all(dict(vars(cls)) == stock[cls] for cls in classes) + + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + model_load_utils._patch_finegrained_fp8_dequant() + assert vars(Fp8Dequantize).get("_dequantize_one") is not stock[Fp8Dequantize].get( + "_dequantize_one" + ) + + +@needs_float8 +@pytest.mark.parametrize( + ("shape", "grid"), + [((576, 7168), (5, 56)), ((192, 256), (2, 2)), ((512, 256), (4, 2))], + ids=["partial", "partial-dividing-the-grid", "whole"], +) +def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( + pre_5_10_transformers, monkeypatch, shape, grid +): + """Upstream fixed the expert converter there and emits the destination dtype, so its dequant is + kept and only wrapped. A partial last block (DeepSeek-V3's 576-row kv_a_proj_with_mqa under a + 5-row grid of 128-row blocks, or 192 rows that its own dequant would split into two blocks of + 96) is padded to whole blocks for it and cropped back; a weight of whole blocks goes to it as + is.""" + if Version("5.10.1") > _INSTALLED_TRANSFORMERS: + pytest.skip("needs transformers' own dequant from 5.10.1") + monkeypatch.setattr(transformers, "__version__", "5.10.1") + hook = vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) + weight, scale_inv = _block_fp8(shape, grid) + + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) + out = _dequant_op((128, 128))._dequantize_one(weight, scale_inv, torch.float32) + + assert torch.equal(out, _reference(weight, scale_inv, (128, 128))) + assert vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) is hook + + +@needs_float8 +def test_an_expert_converter_it_cannot_fuse_stops_the_load(pre_5_10_transformers, monkeypatch): + """Unpatched, every routed expert would load sliced apart from its scales. The refusal leaves + transformers as it was.""" + monkeypatch.setitem(sys.modules, "transformers.core_model_loading", None) + stock = _patched_state() + + with pytest.raises(RuntimeError, match="could not fuse transformers' expert dequant"): + model_load_utils._patch_finegrained_fp8_dequant() + assert _patched_state() == stock From ce7dd38ed19912a2c3024911128434da9aa617da Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:55:20 +0000 Subject: [PATCH 4/5] Require transformers 5.8 to load a block-FP8 checkpoint expert-parallel The expert-parallel load takes transformers 5.7, but the hooks block FP8 is dequantized through, Fp8Dequantize._dequantize_one and the FP8 quantizer's update_weight_conversions, first appear in 5.8. On 5.7 the experts' converter claims their fp8 weights before the dequant can, so a block-FP8 MoE whose dimensions divide into whole blocks (Qwen3-MoE's FP8, say) would load its experts unscaled. The load now refuses block FP8 below 5.8, before anything loads. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 5 +++++ tests/unit/torch/utils/test_model_load_utils.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index bc4b5b27d46..097900a0d2a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2428,6 +2428,11 @@ def load_and_prepare_ep( config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) quant_method = _ckpt_quant_method(hf_config, ckpt_path) + if quant_method == "fp8": + # 5.8 added the hooks block-fp8 is dequantized through (Fp8Dequantize._dequantize_one, + # update_weight_conversions). On 5.7 the experts' converter takes their fp8 weights before + # the dequant can, and they would load unscaled. + moe.require_transformers("5.8", "Dequantizing a block-fp8 checkpoint expert-parallel") if plan.is_fsdp_dense and quant_method == "fp8": raise NotImplementedError( f"{plan!r}: the dense part of a block-fp8 checkpoint cannot be sharded: Distributed " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index c918dbb6b39..eec3b8249c6 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2020,6 +2020,21 @@ def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a( load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) +def test_a_block_fp8_checkpoint_needs_the_transformers_that_dequantizes_its_experts( + monkeypatch, gloo_world_of_one +): + """transformers 5.8 added the hooks block-fp8 is dequantized through (``_dequantize_one``, + ``update_weight_conversions``). On 5.7 the experts' converter takes their fp8 weights before the + dequant can, so they would load unscaled: refused before anything loads.""" + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr("transformers.__version__", "5.7.0") + config = SimpleNamespace(base_model_ep_plan=_EP_PLAN, quantization_config=_BLOCK_FP8) + + with pytest.raises(RuntimeError, match=r"requires transformers >= 5\.8, but 5\.7\.0"): + load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=config) + + @pytest.mark.parametrize( "quantization_config", [None, _BLOCK_FP8], From 939bf6e99aaa50961aa4740073491bc9881894b7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:14:55 +0000 Subject: [PATCH 5/5] Read a config.json that holds no JSON object as one that names no FP8 block _config_json_quantization returns None for a config.json it cannot read, so an FP8 config object meets the load's fp8-unverified refusal. A config.json holding valid JSON other than an object reached .get first and raised AttributeError instead; it now takes the same route. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 2 ++ tests/unit/torch/utils/test_model_load_utils.py | 15 ++++++++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 097900a0d2a..6927c2352c8 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2234,6 +2234,8 @@ def _config_json_quantization(ckpt_path: str | None) -> dict | None: raw = json.load(f) except (OSError, ValueError): return None + if not isinstance(raw, dict): # valid JSON, but no config + return None for c in (raw, raw.get("text_config"), raw.get("llm_config")): qc = c.get("quantization_config") if isinstance(c, dict) else None if isinstance(qc, dict): diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index eec3b8249c6..3067353fd04 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2424,20 +2424,25 @@ def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expect @pytest.mark.parametrize( ("on_disk", "expected"), [ - ({"quant_method": "fp8", "activation_scheme": "static"}, "fp8-per-tensor"), - ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8"), + ( + {"quantization_config": {"quant_method": "fp8", "activation_scheme": "static"}}, + "fp8-per-tensor", + ), + ({"quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]}}, "fp8"), (None, "fp8-unverified"), + ([], "fp8-unverified"), ], - ids=["per-tensor", "block", "no-config-json"], + ids=["per-tensor", "block", "no-config-json", "config-json-not-an-object"], ) def test_an_fp8_config_object_is_read_against_the_checkpoints_own_config( tmp_path, on_disk, expected ): """transformers' FP8 config object names a 128x128 block whatever the checkpoint does, so the checkpoint's own ``config.json`` says whether it names one: per-tensor FP8 handed in as an - object is still refused. Without that file there is no telling, which the load refuses.""" + object is still refused. Without that file, or with one that holds no JSON object, there is no + telling, which the load refuses.""" if on_disk is not None: - (tmp_path / "config.json").write_text(json.dumps({"quantization_config": on_disk})) + (tmp_path / "config.json").write_text(json.dumps(on_disk)) config = SimpleNamespace(quantization_config=FineGrainedFP8Config()) assert _ckpt_quant_method(config, str(tmp_path)) == expected