diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c2438c34e46..6927c2352c8 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,238 @@ 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. + + 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). + 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): + # 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 + + 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. + if getattr(_fp8.Fp8Dequantize, "_modelopt_dequant_patched", False): + _FINEGRAINED_FP8_PATCHED = True + return + _stock_dequantize_one = _fp8.Fp8Dequantize._dequantize_one + + 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 + + # 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): + # 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 " + 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:] + 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 +2186,60 @@ 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 + 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): + return qc return None @@ -2082,13 +2371,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 +2429,32 @@ 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 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 " + "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 +2591,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_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index cd2ce19eb91..3067353fd04 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -25,8 +25,10 @@ import copy import dataclasses import inspect +import json import math import os +import sys from contextlib import ExitStack, contextmanager from types import SimpleNamespace from unittest import mock @@ -45,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, @@ -72,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 @@ -622,21 +627,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 +1927,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 +2020,33 @@ 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): +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], + 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 +2063,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 +2109,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 +2150,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 +2379,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"}) + quantization_config={"quant_method": "mxfp4"}, + text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), ), - "fp8", + "mxfp4", ), + (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8-unverified"), ( SimpleNamespace( - llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) + quantization_config={"quant_method": "fp8", "activation_scheme": "static"} ), - "fp8", + "fp8-per-tensor", ), - ( - SimpleNamespace( - quantization_config={"quant_method": "mxfp4"}, - text_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}), - ), - "mxfp4", - ), - (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), (SimpleNamespace(text_config=None), None), ], @@ -2348,6 +2407,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 +2415,39 @@ 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"), + [ + ( + {"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", "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, 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(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 +2758,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) @@ -3309,3 +3398,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