From f1469ed7741b87d75488167aff931037902c342c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 09:40:08 +0000 Subject: [PATCH 1/6] Dequantize block-FP8 checkpoints as they load expert-parallel A native block-FP8 checkpoint (e.g. DeepSeek-V3) is dequantized to bf16 while it loads, through transformers' FineGrainedFP8Config, since quantizers attach to bf16 weights. Its dequant is corrected on transformers < 5.10.1 for weights whose last 128x128 block is partial, and made to emit the model's dtype. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 184 ++++++++++++- .../torch/utils/test_fp8_dequant_patch.py | 246 ++++++++++++++++++ 2 files changed, 428 insertions(+), 2 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 de4ba15ee53..d519df22520 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1498,6 +1498,167 @@ 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 ONLY: this is skipped. Upstream fixed the two 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). The only thing still + unpatched upstream is the partial-block divisibility (bug 2), which no shipping model on >= 5.10.1 + hits here (DeepSeek-V4's dims are all 128-multiples); revisit if one does. + """ + 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 transformers + from packaging.version import parse + + if parse(transformers.__version__) >= parse("5.10.1"): + _FINEGRAINED_FP8_PATCHED = True # upstream handles it (and natively supports ue8m0/fp4) + return + import torch + from transformers.integrations import finegrained_fp8 as _fp8 + + 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 + # Degenerate/empty scale grid (sharded/merged expert artifact): emit unscaled rather than crash. + if scale_rows == 0 or scale_cols == 0 or scales.numel() == 0: + return q.to(out_dtype) + 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 + return dim // sgrid if sgrid and dim % sgrid == 0 else cfg_blk + + 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, + ) + + if 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: + print(f"[ep] WARNING: could not patch update_weight_conversions: {_e}", flush=True) + + _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. @@ -1875,13 +2036,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 mo_session.rank == 0: + 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_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py new file mode 100644 index 00000000000..d74af85d9ae --- /dev/null +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -0,0 +1,246 @@ +# 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]``. +""" + +from types import SimpleNamespace + +import pytest +import torch + +transformers = pytest.importorskip("transformers") +finegrained_fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8") +quantizer_fp8 = pytest.importorskip("transformers.quantizers.quantizer_finegrained_fp8") + +import modelopt.torch.utils.moe as moe + +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) + +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(moe, "_FINEGRAINED_FP8_PATCHED", False) + monkeypatch.setattr(moe, "_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): + """A random fp8 weight and its fp32 ``scale_inv`` grid.""" + gen = torch.Generator().manual_seed(0) + 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"), + [ + ((8, 12), (4, 4)), + ((10, 13), (4, 4)), + ((9, 8), (4, 4)), + ((6, 10), (2, 4)), + ((3, 10, 13), (4, 4)), + ((576, 7168), (128, 128)), + ], + ids=[ + "whole-blocks", + "partial-last-blocks", + "grid-divides-the-rows", + "non-square-blocks", + "per-expert", + "deepseek-kv-a-proj", + ], +) +def test_dequant_matches_the_blockwise_reference(pre_5_10_transformers, shape, block): + """``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.""" + weight, scale_inv = _block_fp8(shape, _ceil_grid(shape, block)) + moe._patch_finegrained_fp8_dequant(torch.float32) + + out = _dequant_op(block)._dequantize_one(weight, scale_inv) + + assert torch.equal(out, _reference(weight, scale_inv, block)) + + +@pytest.mark.parametrize( + ("configured", "shape", "grid", "block"), + [((128, 128), (8, 8), (2, 4), (4, 2)), (None, (130, 260), (2, 3), (128, 128))], + ids=["grid-finer-than-configured", "nothing-configured"], +) +def test_the_block_size_comes_from_the_grid_or_defaults_to_deepseeks( + pre_5_10_transformers, configured, shape, grid, block +): + """A grid that is not the configured block's ceil grid sets the block by division; with no + block size configured, DeepSeek's 128 x 128 applies.""" + weight, scale_inv = _block_fp8(shape, grid) + moe._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.bfloat16, torch.bfloat16), + (torch.float16, torch.float16), + (torch.float32, torch.float32), + ], +) +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)) + moe._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", [(0, 4), (3, 0)]) +def test_an_empty_scale_grid_passes_the_weight_through_unscaled(pre_5_10_transformers, grid): + """Sharded or merged expert loads can hand over an empty grid; it must not crash the load.""" + weight, scale_inv = _block_fp8((10, 13), grid) + moe._patch_finegrained_fp8_dequant(torch.float32) + + out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv) + + assert torch.equal(out, weight.float()) + + +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.""" + moe._patch_finegrained_fp8_dequant(torch.bfloat16) + installed = _patched_state() + + moe._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 ``moe`` 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") + moe._patch_finegrained_fp8_dequant() + hook = vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] + monkeypatch.setattr(moe, "_FINEGRAINED_FP8_PATCHED", False) + + moe._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} + + moe._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(moe, "_FINEGRAINED_FP8_PATCHED", False) + moe._patch_finegrained_fp8_dequant() + assert vars(Fp8Dequantize).get("_dequantize_one") is not stock[Fp8Dequantize].get( + "_dequantize_one" + ) + + +@pytest.mark.parametrize("version", ["5.10.1", "5.17.0"]) +def test_transformers_from_5_10_1_keeps_its_own_dequant( + pre_5_10_transformers, monkeypatch, version +): + """Upstream fixed the dequant there and decodes ue8m0 scales, which the patch does not.""" + monkeypatch.setattr(transformers, "__version__", version) + stock = _patched_state() + + moe._patch_finegrained_fp8_dequant(torch.bfloat16) + + assert _patched_state() == stock From 4ce96f9ae5134c41705d04d71a60e586dc321120 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:31:42 +0000 Subject: [PATCH 2/6] Log the block-fp8 dequant from rank 0 by the live process group, without the session Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d519df22520..f52f23a8daf 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2050,7 +2050,7 @@ def load_and_prepare_ep( _patch_finegrained_fp8_dequant( dtype ) # correct transformers' block-fp8 dequant; emit model dtype - if mo_session.rank == 0: + if is_master(): print( "[ep] native block-fp8 checkpoint -> FineGrainedFP8Config(dequantize=True) -> bf16", flush=True, From ca63311cac5127ac2ae5e62f83a732ffcb72f6fb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:36:17 +0000 Subject: [PATCH 3/6] Test the block-FP8 dequant patch where it lives, with the loader The patch and its module flags sit in model_load_utils with the expert-parallel loader that applies them, so the test drives and resets them there. Signed-off-by: Shengliang Xu --- .../torch/utils/test_fp8_dequant_patch.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index d74af85d9ae..f886d280ef0 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -31,7 +31,7 @@ finegrained_fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8") quantizer_fp8 = pytest.importorskip("transformers.quantizers.quantizer_finegrained_fp8") -import modelopt.torch.utils.moe as moe +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) @@ -56,8 +56,8 @@ 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(moe, "_FINEGRAINED_FP8_PATCHED", False) - monkeypatch.setattr(moe, "_FP8_DEQUANT_DTYPE", None) + 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: @@ -118,7 +118,7 @@ def test_dequant_matches_the_blockwise_reference(pre_5_10_transformers, shape, b """``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.""" weight, scale_inv = _block_fp8(shape, _ceil_grid(shape, block)) - moe._patch_finegrained_fp8_dequant(torch.float32) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) out = _dequant_op(block)._dequantize_one(weight, scale_inv) @@ -136,7 +136,7 @@ def test_the_block_size_comes_from_the_grid_or_defaults_to_deepseeks( """A grid that is not the configured block's ceil grid sets the block by division; with no block size configured, DeepSeek's 128 x 128 applies.""" weight, scale_inv = _block_fp8(shape, grid) - moe._patch_finegrained_fp8_dequant(torch.float32) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) out = _dequant_op(configured)._dequantize_one(weight, scale_inv) @@ -156,7 +156,7 @@ def test_dequant_emits_the_requested_model_dtype(pre_5_10_transformers, target_d """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)) - moe._patch_finegrained_fp8_dequant(target_dtype) + model_load_utils._patch_finegrained_fp8_dequant(target_dtype) out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv) @@ -168,7 +168,7 @@ def test_dequant_emits_the_requested_model_dtype(pre_5_10_transformers, target_d def test_an_empty_scale_grid_passes_the_weight_through_unscaled(pre_5_10_transformers, grid): """Sharded or merged expert loads can hand over an empty grid; it must not crash the load.""" weight, scale_inv = _block_fp8((10, 13), grid) - moe._patch_finegrained_fp8_dequant(torch.float32) + model_load_utils._patch_finegrained_fp8_dequant(torch.float32) out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv) @@ -177,10 +177,10 @@ def test_an_empty_scale_grid_passes_the_weight_through_unscaled(pre_5_10_transfo 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.""" - moe._patch_finegrained_fp8_dequant(torch.bfloat16) + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) installed = _patched_state() - moe._patch_finegrained_fp8_dequant(torch.float16) + model_load_utils._patch_finegrained_fp8_dequant(torch.float16) assert _patched_state() == installed weight, scale_inv = _block_fp8((10, 13), (3, 4)) @@ -190,14 +190,14 @@ def test_patching_again_keeps_the_install_and_takes_the_new_dtype(pre_5_10_trans def test_a_reimported_module_does_not_wrap_the_quantizer_hook_twice( pre_5_10_transformers, monkeypatch ): - """Re-importing ``moe`` loses its module flag; the quantizer class carries its own.""" + """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") - moe._patch_finegrained_fp8_dequant() + model_load_utils._patch_finegrained_fp8_dequant() hook = vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] - monkeypatch.setattr(moe, "_FINEGRAINED_FP8_PATCHED", False) + monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) - moe._patch_finegrained_fp8_dequant() + model_load_utils._patch_finegrained_fp8_dequant() assert vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] is hook @@ -210,7 +210,7 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo classes = (Fp8Dequantize, FineGrainedFP8HfQuantizer) stock = {cls: dict(vars(cls)) for cls in classes} - moe._patch_finegrained_fp8_dequant() + model_load_utils._patch_finegrained_fp8_dequant() changed = { (cls, name) @@ -226,8 +226,8 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo delattr(cls, name) assert all(dict(vars(cls)) == stock[cls] for cls in classes) - monkeypatch.setattr(moe, "_FINEGRAINED_FP8_PATCHED", False) - moe._patch_finegrained_fp8_dequant() + 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" ) @@ -241,6 +241,6 @@ def test_transformers_from_5_10_1_keeps_its_own_dequant( monkeypatch.setattr(transformers, "__version__", version) stock = _patched_state() - moe._patch_finegrained_fp8_dequant(torch.bfloat16) + model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) assert _patched_state() == stock From 14ed21830e85be788efe8882f1fd416904e5148d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:02:06 +0000 Subject: [PATCH 4/6] Run the fused block-FP8 expert converter in a test, not just its install The patch fuses the Fp8Dequantize that transformers prepends to the routed-expert converter into its MergeModulelist, so the loader still gives each expert its tensor_idx. Only the install was tested. The new test runs the patched update_weight_conversions on an expert converter and the resulting chain on per-expert block-FP8 weights. It fails if the converter no longer leads with a MergeModulelist, or if the fused op skips the dequant or the merge. transformers 5.18's chain passes an output_dtype the patched dequant does not take, so the chain runs that version's own dequant; the patched arithmetic has its own tests. Signed-off-by: Shengliang Xu --- .../torch/utils/test_fp8_dequant_patch.py | 62 ++++++++++++++++++- 1 file changed, 60 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index f886d280ef0..cf6386b227e 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -31,6 +31,9 @@ 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"): @@ -76,9 +79,9 @@ def _dequant_op(block_size): return Fp8Dequantize(SimpleNamespace(quantization_config=config)) -def _block_fp8(shape, grid): +def _block_fp8(shape, grid, seed=0): """A random fp8 weight and its fp32 ``scale_inv`` grid.""" - gen = torch.Generator().manual_seed(0) + 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 @@ -175,6 +178,61 @@ def test_an_empty_scale_grid_passes_the_weight_through_unscaled(pre_5_10_transfo assert torch.equal(out, weight.float()) +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) From 00b5939c40bfdb86a22ac7476b69a7d76f7c8adc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:50:00 +0000 Subject: [PATCH 5/6] Trim the block-FP8 dequant tests: block sizes join the reference table - The two block-size cases, a grid finer than the configured block and no block configured, become rows of the blockwise-reference table, which gains configured and grid columns. - The dtype test keeps its unrequested (bf16) and requested (fp16) cases; requesting bf16 or fp32 took the same branch. - The version test keeps the 5.10.1 boundary; 5.17.0 took the same branch. Signed-off-by: Shengliang Xu --- .../torch/utils/test_fp8_dequant_patch.py | 62 +++++++------------ 1 file changed, 23 insertions(+), 39 deletions(-) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index cf6386b227e..089ffa90e76 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -99,14 +99,16 @@ def _reference(weight, scale_inv, block): @pytest.mark.parametrize( - ("shape", "block"), + ("shape", "block", "configured", "grid"), [ - ((8, 12), (4, 4)), - ((10, 13), (4, 4)), - ((9, 8), (4, 4)), - ((6, 10), (2, 4)), - ((3, 10, 13), (4, 4)), - ((576, 7168), (128, 128)), + ((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)), ], ids=[ "whole-blocks", @@ -115,30 +117,19 @@ def _reference(weight, scale_inv, block): "non-square-blocks", "per-expert", "deepseek-kv-a-proj", + "grid-finer-than-configured", + "nothing-configured", ], ) -def test_dequant_matches_the_blockwise_reference(pre_5_10_transformers, shape, block): - """``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.""" - weight, scale_inv = _block_fp8(shape, _ceil_grid(shape, block)) - model_load_utils._patch_finegrained_fp8_dequant(torch.float32) - - out = _dequant_op(block)._dequantize_one(weight, scale_inv) - - assert torch.equal(out, _reference(weight, scale_inv, block)) - - -@pytest.mark.parametrize( - ("configured", "shape", "grid", "block"), - [((128, 128), (8, 8), (2, 4), (4, 2)), (None, (130, 260), (2, 3), (128, 128))], - ids=["grid-finer-than-configured", "nothing-configured"], -) -def test_the_block_size_comes_from_the_grid_or_defaults_to_deepseeks( - pre_5_10_transformers, configured, shape, grid, block +def test_dequant_matches_the_blockwise_reference( + pre_5_10_transformers, shape, block, configured, grid ): - """A grid that is not the configured block's ceil grid sets the block by division; with no - block size configured, DeepSeek's 128 x 128 applies.""" - weight, scale_inv = _block_fp8(shape, 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) @@ -148,12 +139,8 @@ def test_the_block_size_comes_from_the_grid_or_defaults_to_deepseeks( @pytest.mark.parametrize( ("target_dtype", "expected"), - [ - (None, torch.bfloat16), - (torch.bfloat16, torch.bfloat16), - (torch.float16, torch.float16), - (torch.float32, torch.float32), - ], + [(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 @@ -291,12 +278,9 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo ) -@pytest.mark.parametrize("version", ["5.10.1", "5.17.0"]) -def test_transformers_from_5_10_1_keeps_its_own_dequant( - pre_5_10_transformers, monkeypatch, version -): +def test_transformers_from_5_10_1_keeps_its_own_dequant(pre_5_10_transformers, monkeypatch): """Upstream fixed the dequant there and decodes ue8m0 scales, which the patch does not.""" - monkeypatch.setattr(transformers, "__version__", version) + monkeypatch.setattr(transformers, "__version__", "5.10.1") stock = _patched_state() model_load_utils._patch_finegrained_fp8_dequant(torch.bfloat16) From d3fc951743d38de746236ae963144ce3344fd09d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:49:42 +0000 Subject: [PATCH 6/6] Refuse block-FP8 scale grids the dequant patch cannot place The pre-5.10.1 dequant patch emitted a non-empty weight unscaled when its scale grid was empty. That only happens when a weight and its scales were sliced apart, as transformers 5.9's expert tensor_idx bug did on expert-parallel loads, so it loaded wrong weights without a word; upstream's own dequant fails on such a grid. It also fell back to the configured block for a grid neither that block nor an even split gives, scaling rows wrongly. Raise in both cases; an empty weight still passes through its empty grid. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 19 ++++++++++++++++--- .../torch/utils/test_fp8_dequant_patch.py | 19 +++++++++++++------ 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f52f23a8daf..6420e6f7fd6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1557,9 +1557,16 @@ def _dequantize_one(self, quantized, scales): # 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 - # Degenerate/empty scale grid (sharded/merged expert artifact): emit unscaled rather than crash. if scale_rows == 0 or scale_cols == 0 or scales.numel() == 0: - return q.to(out_dtype) + # 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) @@ -1571,7 +1578,13 @@ def _blk(dim, sgrid, cfg_blk): # never infer from exact divisibility alone (576 % 5 and 9 % 3 both mislead). if sgrid == -(-dim // cfg_blk): return cfg_blk - return dim // sgrid if sgrid and dim % sgrid == 0 else 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]) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 089ffa90e76..21f4253639f 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -109,6 +109,7 @@ def _reference(weight, scale_inv, block): ((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", @@ -119,6 +120,7 @@ def _reference(weight, scale_inv, block): "deepseek-kv-a-proj", "grid-finer-than-configured", "nothing-configured", + "empty-weight", ], ) def test_dequant_matches_the_blockwise_reference( @@ -154,15 +156,20 @@ def test_dequant_emits_the_requested_model_dtype(pre_5_10_transformers, target_d assert torch.equal(out, _reference(weight, scale_inv, (4, 4)).to(expected)) -@pytest.mark.parametrize("grid", [(0, 4), (3, 0)]) -def test_an_empty_scale_grid_passes_the_weight_through_unscaled(pre_5_10_transformers, grid): - """Sharded or merged expert loads can hand over an empty grid; it must not crash the load.""" +@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) - out = _dequant_op((4, 4))._dequantize_one(weight, scale_inv) - - assert torch.equal(out, weight.float()) + 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(