diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 6f6823dfc9a..fb9134c8cc4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1493,6 +1493,213 @@ 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 requires the + weight to divide evenly into the scale grid (bug 2), which DeepSeek-V3's partial last blocks do not, + so there a weight whose grid is the configured block's ceil grid is padded to whole blocks for it + and the result cropped back. + """ + global _FINEGRAINED_FP8_PATCHED, _FP8_DEQUANT_DTYPE + if target_dtype is not None: + _FP8_DEQUANT_DTYPE = target_dtype # keep current even if already patched + if _FINEGRAINED_FP8_PATCHED: + return + import torch + import transformers + from packaging.version import parse + from transformers.integrations import finegrained_fp8 as _fp8 + + if parse(transformers.__version__) >= parse("5.10.1"): + _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): + 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 ( + quantized.dtype not in fp8_dtypes + or scales.dim() < 2 + or tuple(scales.shape[-2:]) != ceil_grid + or (rows % ceil_grid[0] == 0 and cols % ceil_grid[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 = torch.nn.functional.pad(quantized.float(), (0, pad_cols, 0, pad_rows)) + out = _stock_dequantize_one(self, padded, scales, *args, **kwargs) + return out[..., :rows, :cols] + + _fp8.Fp8Dequantize._dequantize_one = _dequantize_partial_blocks + _FINEGRAINED_FP8_PATCHED = True + return + + def _dequantize_one(self, quantized, scales): + # FP4 path (int8 / float4_e2m1fn_x2 nibble-packed) unchanged: unpack to fp32 first. + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + if quantized.dtype == torch.int8 or ( + fp4_dtype is not None and quantized.dtype == fp4_dtype + ): + q = self._unpack_fp4(quantized) + else: + q = quantized.to(torch.float32) + 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. + out_dtype = _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]) + s = ( + scales.to(torch.float32) + .repeat_interleave(block_m, dim=-2) + .repeat_interleave(block_n, dim=-1) + ) + s = s[..., :rows, :cols] + return (q * s).to(out_dtype) + + _fp8.Fp8Dequantize._dequantize_one = _dequantize_one + + # --- 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) + and not isinstance(ops[1], _DequantThenMerge) + ): + 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 + + _FINEGRAINED_FP8_PATCHED = True + + def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: """FSDP2-shard the model's dense modules, leaving the routed experts to EP. @@ -1579,7 +1786,7 @@ 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). -_EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) +_EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) def _ckpt_quant_method(cfg): @@ -1883,13 +2090,32 @@ def load_and_prepare_ep( from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation - if _ckpt_quant_method(hf_config) == "mxfp4": + # Native block-FP8 checkpoints (e.g. DeepSeek-V3: quantization_config.quant_method="fp8", e4m3 with + # weight_block_size) must be DEQUANTIZED to bf16 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.) + _ckpt_qm = _ckpt_quant_method(hf_config) + if _ckpt_qm == "fp8": + from transformers import FineGrainedFP8Config + + from_pretrained_kwargs["quantization_config"] = FineGrainedFP8Config(dequantize=True) + _patch_finegrained_fp8_dequant( + dtype + ) # correct transformers' block-fp8 dequant; emit model dtype + if is_master(): + print( + "[ep] native block-fp8 checkpoint -> FineGrainedFP8Config(dequantize=True) -> bf16", + flush=True, + ) + elif _ckpt_qm == "mxfp4": # Native mxfp4 checkpoints (e.g. gpt-oss: fused `gate_up_proj_blocks`/`_scales` + biases) must be # DEQUANTIZED to bf16 on load so ModelOpt can insert nvfp4/fp8 quantizers on the fused bf16 experts # (the a2a grouped_gemm expects bf16). transformers dequantizes blocks->bf16 during from_pretrained. # NB: the single-process loader needed device_map="sequential" to dodge a CUDA illegal-access in the # mxfp4 dequant under multi-GPU device_map="auto"; the EP path uses device_mesh (no device_map), so - # watch for an EP-shard-awareness issue in the dequant convert-op. + # watch for an EP-shard-awareness issue in the dequant convert-op (analogous to the fp8 patch above). from transformers import Mxfp4Config from_pretrained_kwargs["quantization_config"] = Mxfp4Config(dequantize=True) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 18a3b195bb0..659330ae44f 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -126,10 +126,12 @@ def test_the_loader_refuses_a_plan_for_another_world(gloo_world_of_one): load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2)) -def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch): +@pytest.mark.parametrize("quant_method", [None, "fp8", "mxfp4"]) +def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, quant_method): """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, MXFP4) passes the format check.""" class _ResolvedError(Exception): """Stops the load where it resolves the checkpoint.""" @@ -146,10 +148,12 @@ def resolve(path, rank): monkeypatch.setattr(moe, "require_transformers", lambda *args: None) monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) + config = SimpleNamespace() + if quant_method is not None: + config.quantization_config = {"quant_method": quant_method} + with pytest.raises(_ResolvedError): - load_and_prepare_ep( - "org/model", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace() - ) + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) assert resolved == ["org/model"] 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..906076f2672 --- /dev/null +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -0,0 +1,320 @@ +# 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(finegrained_fp8, "Fp8Dequantize"): + pytest.skip("transformers has no Fp8Dequantize", allow_module_level=True) +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: the dequant method, and the FP8 quantizer's converter hook plus the +# flag guarding it. +_PATCH_TARGETS = ( + (Fp8Dequantize, "_dequantize_one"), + (FineGrainedFP8HfQuantizer, "update_weight_conversions"), + (FineGrainedFP8HfQuantizer, "_modelopt_uwc_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 _PATCH_TARGETS] + 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), + ((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", + "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)) + + +@pytest.mark.parametrize( + ("target_dtype", "expected"), + [(None, torch.bfloat16), (torch.float16, torch.float16)], + ids=["unrequested", "requested"], +) +def test_dequant_emits_the_requested_model_dtype(pre_5_10_transformers, target_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.""" + 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) + + 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")], + ids=["no-row-scales", "no-col-scales", "grid-no-block-places"], +) +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), or a 4-row grid over 10 rows, + which no 4-row block (the configured one) or even split gives.""" + 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") + stock_dequant = Fp8Dequantize._dequantize_one + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) + # This transformers' chain passes output_dtype, which the patched dequant (written for the + # versions it installs on) does not take; its arithmetic is tested above. + Fp8Dequantize._dequantize_one = stock_dequant + 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 + + +def test_a_reimported_module_does_not_wrap_the_quantizer_hook_twice( + pre_5_10_transformers, monkeypatch +): + """Re-importing ``model_load_utils`` loses its module flag; the quantizer class carries its own.""" + if not hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions"): + pytest.skip("this transformers' FP8 quantizer has no update_weight_conversions") + model_load_utils._patch_finegrained_fp8_dequant() + hook = vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) + + model_load_utils._patch_finegrained_fp8_dequant() + + assert vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] is hook + + +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)), ((512, 256), (4, 2))], ids=["partial", "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) 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.""" + monkeypatch.setitem(sys.modules, "transformers.core_model_loading", None) + + with pytest.raises(RuntimeError, match="could not fuse transformers' expert dequant"): + model_load_utils._patch_finegrained_fp8_dequant()