From ab100e0a4ebf76315075174ac59193fdacccd56f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 09:40:08 +0000 Subject: [PATCH 01/25] 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 | 183 ++++++++++++- .../torch/utils/test_fp8_dequant_patch.py | 246 ++++++++++++++++++ 2 files changed, 427 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 bed362fd106..6f162da1006 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1709,6 +1709,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. @@ -2144,13 +2305,31 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A from_pretrained_kwargs["tp_plan"] = "auto" if attn_implementation is not None: from_pretrained_kwargs["attn_implementation"] = attn_implementation - if quant_method == "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.) + if quant_method == "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 quant_method == "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 838cb2945aaec0fa923c5b6b49f7227925ba3f1b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 17:31:42 +0000 Subject: [PATCH 02/25] 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 6f162da1006..63a1a4c3c9c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2318,7 +2318,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A _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 65c7e7197159bf359772facd0ea686523a129c1c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:36:17 +0000 Subject: [PATCH 03/25] 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 f29a380851c583c04b4d72cf460545a2de326526 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:02:06 +0000 Subject: [PATCH 04/25] 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 097a656b678c0517d12a08e25bba4c265edf9415 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:50:00 +0000 Subject: [PATCH 05/25] 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 f74f307b773eed46d05d99e6fe964762c6eeb8c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:49:42 +0000 Subject: [PATCH 06/25] 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 63a1a4c3c9c..4ed8c0b75f3 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1768,9 +1768,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) @@ -1782,7 +1789,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( From 93f3db2cc81272e8f7b7cb2d1634528d14670e8d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 01:05:01 +0000 Subject: [PATCH 07/25] Dequantize block-FP8 partial last blocks on transformers 5.10.1 too From 5.10.1 the patch kept transformers' own dequant untouched, on the assumption that no shipping model has a partial last block there. DeepSeek-V3, V3.1 and Kimi-K2-Instruct do: kv_a_proj_with_mqa has 576 rows under a 5-row grid of 128-row blocks, and Fp8Dequantize still raises "not divisible by scale grid" for it, so their block-FP8 loads failed. The stock dequant is kept, for its ue8m0 and FP4 decoding and output_dtype, and wrapped: an FP8 weight whose grid is the configured block's ceil grid is padded to whole blocks for it and cropped back. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 41 +++++++++++++++---- .../torch/utils/test_fp8_dequant_patch.py | 24 +++++++++-- 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 4ed8c0b75f3..cf9b6875c09 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1731,27 +1731,54 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): else derived), then crop to the weight shape -- correct for both full and partial last blocks, 2-D dense linears and (leading-dim-broadcast) per-expert tensors. Idempotent. - transformers >= 5.10.1 ONLY: this is skipped. Upstream fixed the two show-stoppers there -- the + 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). 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. + (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"): - _FINEGRAINED_FP8_PATCHED = True # upstream handles it (and natively supports ue8m0/fp4) + _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 - 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. diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 21f4253639f..3dadc2441b8 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -26,6 +26,7 @@ import pytest import torch +from packaging.version import Version transformers = pytest.importorskip("transformers") finegrained_fp8 = pytest.importorskip("transformers.integrations.finegrained_fp8") @@ -41,6 +42,8 @@ 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 @@ -285,11 +288,24 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo ) -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.""" +@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") - stock = _patched_state() + 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 _patched_state() == stock + assert torch.equal(out, _reference(weight, scale_inv, (128, 128))) + assert vars(FineGrainedFP8HfQuantizer).get("update_weight_conversions", _MISSING) is hook From d89acfc502719932790cb17f7581a8ce406095e2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 01:05:07 +0000 Subject: [PATCH 08/25] Let a block-FP8 expert converter the patch cannot fuse stop the load When fusing transformers' prepended expert dequant into its merge failed, the patch printed a warning and went on. Every routed expert then loaded sliced apart from its scales. The failure now raises. A transformers without the update_weight_conversions hook prepends no dequant, so there is nothing to fuse and the patch skips it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 12 +++++++++--- tests/unit/torch/utils/test_fp8_dequant_patch.py | 9 +++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index cf9b6875c09..ae05bbf952c 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1881,7 +1881,10 @@ def convert(self, input_dict, source_patterns=None, target_patterns=None, **kwar **kwargs, ) - if not getattr(_FineGrainedFP8HfQuantizer, "_modelopt_uwc_patched", False): + # 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): @@ -1904,8 +1907,11 @@ def _uwc_fixed(self, weight_conversions): _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) + 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 diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 3dadc2441b8..906076f2672 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -22,6 +22,7 @@ plain blockwise reference: ``weight[..., i, j] * scale_inv[..., i // block_n, j // block_k]``. """ +import sys from types import SimpleNamespace import pytest @@ -309,3 +310,11 @@ def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( 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() From 44060c9d88d0950cd71636899eba339fb8c11507 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 01:09:02 +0000 Subject: [PATCH 09/25] Let the expert-parallel load take the block-FP8 checkpoints it dequantizes The load dequantizes native block-FP8 checkpoints as it loads, so they join the formats it accepts. 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 ae05bbf952c..ca976e64489 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2014,7 +2014,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). # mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. -_EP_SOURCE_FORMATS = frozenset({None, "mxfp4"}) +_EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) def _ckpt_quant_method(cfg): From faf195c18a271df258a8243e900550f8de40f952 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 05:26:36 +0000 Subject: [PATCH 10/25] Test that the formats the expert-parallel load dequantizes pass its check The hub-ID test now loads an unquantized, a block-FP8 and an MXFP4 config: each must reach the checkpoint's resolution, past the format check, so dropping either format from the accepted set fails it. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_ep_loader.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index df4656abcf0..847875ee6a0 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -169,10 +169,12 @@ def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a(m 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): +@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.""" @@ -189,12 +191,12 @@ def resolve(path, rank): monkeypatch.setattr(moe, "require_transformers", lambda *args: None) monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) + config = SimpleNamespace(base_model_ep_plan=_EP_PLAN) + 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(base_model_ep_plan=_EP_PLAN), - ) + load_and_prepare_ep("org/model", DistributionPlan(2, ep_size=2), hf_config=config) assert resolved == ["org/model"] From 9a1ec03a0949cfa5a7c9c2ffc769889786869140 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 19:28:41 +0000 Subject: [PATCH 11/25] Pad every partial-block FP8 weight for transformers' own dequant From transformers 5.10.1 the expert-parallel load wraps transformers' block-FP8 dequant, which derives the block size from the scale grid, and padded a weight to whole blocks only when it did not divide into that grid. A partial last block can still divide evenly: 448 rows under a 4-row grid of 128-row blocks go to the dequant as blocks of 112, and every row group after the first gets the wrong scale, silently. The wrapper now pads whenever the weight is not block-aligned, which is a no-op for whole blocks. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 11 ++++++----- tests/unit/torch/utils/test_fp8_dequant_patch.py | 9 ++++++--- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ca976e64489..1c6bcd54af1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1735,10 +1735,11 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): 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. + (installing it on >= 5.10.1 would BREAK ue8m0 checkpoints like DeepSeek-V4). It still derives the + block from the scale grid (bug 2), which a partial last block makes wrong: 576 rows under a 5-row + grid do not divide at all, and 448 rows under a 4-row grid divide into blocks of 112 instead of 128. + So there a weight whose grid is the configured block's ceil grid and that has a partial last block + is padded to whole blocks for it and the result cropped back. """ global _FINEGRAINED_FP8_PATCHED, _FP8_DEQUANT_DTYPE if target_dtype is not None: @@ -1768,7 +1769,7 @@ def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): 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) + 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 diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 906076f2672..0ec55b44173 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -290,15 +290,18 @@ def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, mo @pytest.mark.parametrize( - ("shape", "grid"), [((576, 7168), (5, 56)), ((512, 256), (4, 2))], ids=["partial", "whole"] + ("shape", "grid"), + [((576, 7168), (5, 56)), ((192, 256), (2, 2)), ((512, 256), (4, 2))], + ids=["partial", "partial-dividing-the-grid", "whole"], ) def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( pre_5_10_transformers, monkeypatch, shape, grid ): """Upstream fixed the expert converter there and decodes ue8m0 and FP4 scales, which the override does not, so its dequant is kept and only wrapped. A partial last block (DeepSeek-V3's 576-row - kv_a_proj_with_mqa under a 5-row grid of 128-row blocks) is padded to whole blocks for it and - cropped back; a weight of whole blocks goes to it as is.""" + kv_a_proj_with_mqa under a 5-row grid of 128-row blocks, or 192 rows that its own dequant would + split into two blocks of 96) is padded to whole blocks for it and cropped back; a weight of whole + blocks goes to it as is.""" if Version("5.10.1") > _INSTALLED_TRANSFORMERS: pytest.skip("needs transformers' own dequant from 5.10.1") monkeypatch.setattr(transformers, "__version__", "5.10.1") From 20a0ec3944cc5b46644240b38f84e73617a6562d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:32:05 +0000 Subject: [PATCH 12/25] Scale the FP8 dequant in place, and rebind it only once the fuse is in The pre-5.10.1 dequant multiplies its own fp32 copy of the weight by the tiled scales in place, so a weight holds two full-size fp32 temporaries instead of three while it loads. The dequant is rebound after the expert converter's fuse, so a refused patch leaves transformers as it was. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 14 +++++++------- tests/unit/torch/utils/test_fp8_dequant_patch.py | 5 ++++- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1c6bcd54af1..2efd2272c83 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1789,7 +1789,7 @@ def _dequantize_one(self, quantized, scales): ): q = self._unpack_fp4(quantized) else: - q = quantized.to(torch.float32) + q = quantized.to(torch.float32, copy=True) rows, cols = q.shape[-2:] scale_rows, scale_cols = scales.shape[-2:] # Emit the MODEL compute dtype (bf16), NOT the scale dtype. DeepSeek's weight_scale_inv is F32; @@ -1827,15 +1827,13 @@ def _blk(dim, sgrid, cfg_blk): block_m = _blk(rows, scale_rows, wbs[0]) block_n = _blk(cols, scale_cols, wbs[1]) - s = ( + # In place on q, its own copy: the product would be a third full-size fp32 temporary. + q *= ( scales.to(torch.float32) .repeat_interleave(block_m, dim=-2) - .repeat_interleave(block_n, dim=-1) + .repeat_interleave(block_n, dim=-1)[..., :rows, :cols] ) - s = s[..., :rows, :cols] - return (q * s).to(out_dtype) - - _fp8.Fp8Dequantize._dequantize_one = _dequantize_one + 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 @@ -1914,6 +1912,8 @@ def _uwc_fixed(self, weight_conversions): 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 diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 0ec55b44173..f4d3a77de7c 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -316,8 +316,11 @@ def test_from_5_10_1_its_own_dequant_also_takes_partial_last_blocks( 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.""" + """Unpatched, every routed expert would load sliced apart from its scales. The refusal leaves + transformers as it was.""" monkeypatch.setitem(sys.modules, "transformers.core_model_loading", None) + stock = _patched_state() with pytest.raises(RuntimeError, match="could not fuse transformers' expert dequant"): model_load_utils._patch_finegrained_fp8_dequant() + assert _patched_state() == stock From 3ca078824d5531ed77b874fbe8065a712c47f945 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:27:37 +0000 Subject: [PATCH 13/25] Replicate block-FP8's dense part, and dequantize with row-tiled scales The rule that a block-FP8 checkpoint's dense part is not sharded lands here, with the load that makes such a checkpoint loadable, rather than in [10/21], where it preempted the format refusal. Left to decide, the plan replicates that dense part and says why; asked to shard it, the decision refuses, naming MODELOPT_FSDP_SIZE=1, and the loader refuses a plan that shards it. The pre-5.10.1 dequant scales its fp32 copy of a weight in place, through a view of its whole column blocks and then the partial last one, against scales tiled over the rows only. It builds no full-size fp32 scale tile: an (18432, 7168) weight peaks at 6 bytes per element instead of 8. A weight narrower than one block joins the reference table. The expert converter's fuse no longer skips a merge it already fused, a layout transformers' update_weight_conversions never produces (it builds fresh converters on every call); fusing it again would dequantize once all the same. The load's log names the dtype the dequant emits. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 35 +++++++++++++------ .../torch/utils/test_ep_dense_sharding.py | 8 ++--- .../torch/utils/test_fp8_dequant_patch.py | 2 ++ .../utils/test_resolve_distribution_plan.py | 26 +++++++++----- 4 files changed, 47 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2efd2272c83..9e60497a674 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1364,7 +1364,8 @@ 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, cannot dequantize a block-fp8 checkpoint's F8_E8M0 scales. unshardable, remedy = None, "" if ep_size > 1 and world_size // ep_size > 1: unshardable = ( @@ -1375,6 +1376,12 @@ def _decide_distribution( " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) + elif ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8": + unshardable = "from a block-fp8 checkpoint: Distributed Checkpoint cannot dequantize its F8_E8M0 scales" + 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. @@ -1827,12 +1834,13 @@ def _blk(dim, sgrid, cfg_blk): block_m = _blk(rows, scale_rows, wbs[0]) block_n = _blk(cols, scale_cols, wbs[1]) - # In place on q, its own copy: the product would be a third full-size fp32 temporary. - q *= ( - scales.to(torch.float32) - .repeat_interleave(block_m, dim=-2) - .repeat_interleave(block_n, dim=-1)[..., :rows, :cols] - ) + # 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) --------------------- @@ -1898,7 +1906,6 @@ def _uwc_fixed(self, weight_conversions): 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)] @@ -2199,7 +2206,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A With ``plan.dense == "fsdp"`` the dense part is instead FSDP2-sharded over the ``ep`` group and each rank loads only its shard from the checkpoint, never the whole of it -- for a model whose dense part - does not fit one GPU. + does not fit one GPU. A block-fp8 checkpoint cannot be loaded that way: Distributed Checkpoint + cannot dequantize its F8_E8M0 scales. EP CANNOT be applied to a meta-init model (transformers' ``GroupedGemmParallel.shard_tensor`` slices materialized tensors), so this loads real weights via ``from_pretrained`` over the ``ep`` @@ -2270,6 +2278,11 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) quant_method = _ckpt_quant_method(hf_config) + 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, since Distributed " + "Checkpoint cannot dequantize its F8_E8M0 scales. Plan dense='replicate'." + ) 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)) @@ -2353,7 +2366,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 bf16 on load: ModelOpt PTQ inserts/exports quantizers on + # 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 @@ -2367,7 +2380,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ) # correct transformers' block-fp8 dequant; emit model dtype if is_master(): print( - "[ep] native block-fp8 checkpoint -> FineGrainedFP8Config(dequantize=True) -> bf16", + f"[ep] native block-fp8 checkpoint -> FineGrainedFP8Config(dequantize=True) -> {dtype}", flush=True, ) elif quant_method == "mxfp4": diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 78c7c21a068..047155b3203 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -107,15 +107,15 @@ 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): - """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): + """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(torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 2) monkeypatch.setattr(moe, "require_transformers", lambda *args: None) config = SimpleNamespace(quantization_config={"quant_method": "fp8"}) - 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) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index f4d3a77de7c..c5117576adf 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -107,6 +107,7 @@ def _reference(weight, scale_inv, block): [ ((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), @@ -118,6 +119,7 @@ def _reference(weight, scale_inv, block): ids=[ "whole-blocks", "partial-last-blocks", + "narrower-than-a-block", "grid-divides-the-rows", "non-square-blocks", "per-expert", diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 069c53c6a34..2a79d004bc6 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -473,20 +473,28 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( @needs_qwen3_moe_ep_plan +@pytest.mark.parametrize( + ("num_experts", "quant_method", "ep_size", "match"), + [(8, "fp8", 4, "block-fp8"), (2, None, 2, "2 expert-data-parallel replicas")], + ids=["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, quant_method, 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. Without expert parallelism neither rule applies.""" run_as(4) - serve(_qwen3_moe(num_experts=2)) + config = _qwen3_moe(num_experts=num_experts) + if quant_method: + config.quantization_config = {"quant_method": quant_method} + 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"): resolve_distribution_plan(str(tmp_path), dense="fsdp") assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( 4, dense="fsdp" From a05fa8640dee37245abcf88c7e260a9c819da57a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:11:01 +0000 Subject: [PATCH 14/25] Refuse per-tensor FP8 checkpoints the block-FP8 dequant cannot load quant_method "fp8" is not block FP8 alone: per-tensor FP8 checkpoints (AutoFP8's: a scalar weight_scale per weight, no weight_block_size) name the same method. The expert-parallel load admitted them, and transformers' dequant reads them as 128x128 blocks with no weight_scale_inv, so their weights load unscaled without an error (on transformers 5.18 a tiny Llama loaded that way has max |w| 448 where the original has 0.079). _ckpt_quant_method now calls an fp8 checkpoint without a weight_block_size "fp8-per-tensor", which the load refuses, before anything loads, as a format it cannot dequantize; the decision's block-FP8 rule no longer applies to it either. The decision, the loader's refusal and its docstring give the reason a block-FP8 dense part cannot be sharded as Distributed Checkpoint reading it as stored, with no dequantization: its scales' format was never the blocker (DeepSeek-V3's are fp32). A comment says where the dequant's block size comes from: transformers keeps the checkpoint's quantization_config and takes only the loading attributes from the one the load passes. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 28 +++++++--- .../torch/utils/test_ep_dense_sharding.py | 4 +- tests/unit/torch/utils/test_ep_loader.py | 55 ++++++++++++------- .../utils/test_resolve_distribution_plan.py | 13 +++-- 4 files changed, 68 insertions(+), 32 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 9e60497a674..62402be7a6e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1365,7 +1365,7 @@ def _decide_distribution( ) # 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, and Distributed Checkpoint, which - # loads it sharded, cannot dequantize a block-fp8 checkpoint's F8_E8M0 scales. + # loads it sharded, reads a block-fp8 checkpoint's tensors as stored, with no dequantization. unshardable, remedy = None, "" if ep_size > 1 and world_size // ep_size > 1: unshardable = ( @@ -1377,7 +1377,10 @@ def _decide_distribution( "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) elif ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8": - unshardable = "from a block-fp8 checkpoint: Distributed Checkpoint cannot dequantize its F8_E8M0 scales" + 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." @@ -2022,6 +2025,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). # mxfp4 is for experts stored in the standard layout: GPT-OSS's are transposed, which the load refuses. +# fp8 is block FP8; _ckpt_quant_method calls per-tensor FP8 fp8-per-tensor, which the load refuses. _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) @@ -2030,14 +2034,22 @@ def _ckpt_quant_method(cfg): 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"``. """ 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)) + 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 @@ -2207,7 +2219,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A With ``plan.dense == "fsdp"`` the dense part is instead FSDP2-sharded over the ``ep`` group and each rank loads only its shard from the checkpoint, never the whole of it -- for a model whose dense part does not fit one GPU. A block-fp8 checkpoint cannot be loaded that way: Distributed Checkpoint - cannot dequantize its F8_E8M0 scales. + reads it as stored, with no dequantization. EP CANNOT be applied to a meta-init model (transformers' ``GroupedGemmParallel.shard_tensor`` slices materialized tensors), so this loads real weights via ``from_pretrained`` over the ``ep`` @@ -2280,8 +2292,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A quant_method = _ckpt_quant_method(hf_config) 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, since Distributed " - "Checkpoint cannot dequantize its F8_E8M0 scales. Plan dense='replicate'." + 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 not in _EP_SOURCE_FORMATS: # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. @@ -2374,6 +2386,8 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 diff --git a/tests/unit/torch/utils/test_ep_dense_sharding.py b/tests/unit/torch/utils/test_ep_dense_sharding.py index 047155b3203..09b2e02a776 100644 --- a/tests/unit/torch/utils/test_ep_dense_sharding.py +++ b/tests/unit/torch/utils/test_ep_dense_sharding.py @@ -113,7 +113,9 @@ def test_a_block_fp8_checkpoints_dense_part_is_not_sharded(monkeypatch): monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) monkeypatch.setattr(torch.distributed, "get_world_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="block-fp8 checkpoint cannot be sharded"): load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 847875ee6a0..5128a1ee111 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -105,6 +105,7 @@ def _model_with_decoder_at(prefix): # ------------------------------------------------------------------------------- the plan _EP_PLAN = {"layers.*.mlp.gate": "ep_router", "layers.*.mlp.experts": "moe_tp_experts"} +_BLOCK_FP8 = {"quant_method": "fp8", "weight_block_size": [128, 128]} # DeepSeek-V3's format class _Mesh: @@ -169,8 +170,12 @@ def test_the_loader_refuses_a_transformers_that_builds_no_model_with_torch_a2a(m load_and_prepare_ep("unused", DistributionPlan(2, ep_size=2), hf_config=SimpleNamespace()) -@pytest.mark.parametrize("quant_method", [None, "fp8", "mxfp4"]) -def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, quant_method): +@pytest.mark.parametrize( + "quantization_config", + [None, _BLOCK_FP8, {"quant_method": "mxfp4"}], + ids=["unquantized", "block-fp8", "mxfp4"], +) +def test_a_hub_id_is_made_local_before_the_weights_load(monkeypatch, 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: a checkpoint @@ -192,8 +197,8 @@ def resolve(path, rank): monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) config = SimpleNamespace(base_model_ep_plan=_EP_PLAN) - if quant_method is not None: - config.quantization_config = {"quant_method": quant_method} + 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=config) @@ -212,6 +217,13 @@ def resolve(path, rank): ) for method in ("gptq", "modelopt", "compressed-tensors") ), + ( + SimpleNamespace( + quantization_config={"quant_method": "fp8", "activation_scheme": "static"} + ), + NotImplementedError, + "cannot dequantize a 'fp8-per-tensor'", + ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( SimpleNamespace(base_model_ep_plan={"layers.*.self_attn.q_proj": "colwise"}), @@ -219,7 +231,14 @@ def resolve(path, rank): "styles no routed experts", ), ], - ids=["gptq", "modelopt", "compressed-tensors", "no-ep-plan", "no-routed-experts"], + ids=[ + "gptq", + "modelopt", + "compressed-tensors", + "per-tensor-fp8", + "no-ep-plan", + "no-routed-experts", + ], ) def test_a_config_the_load_cannot_take_is_refused_before_anything_loads( ep_world_of_two, config, error, match @@ -660,20 +679,10 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): @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={"quant_method": "fp8"}) - ), - "fp8", - ), - ( - SimpleNamespace( - llm_config=SimpleNamespace(quantization_config={"quant_method": "fp8"}) - ), - "fp8", - ), + (SimpleNamespace(text_config=SimpleNamespace(quantization_config=_BLOCK_FP8)), "fp8"), + (SimpleNamespace(llm_config=SimpleNamespace(quantization_config=_BLOCK_FP8)), "fp8"), ( SimpleNamespace( quantization_config={"quant_method": "mxfp4"}, @@ -682,6 +691,12 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): "mxfp4", ), (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), + ( + SimpleNamespace( + quantization_config={"quant_method": "fp8", "activation_scheme": "static"} + ), + "fp8-per-tensor", + ), (SimpleNamespace(quantization_config={"weight_block_size": [128, 128]}), None), (SimpleNamespace(text_config=None), None), ], @@ -692,6 +707,7 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): "vlm-llm-config", "top-level-wins", "transformers-config-object", + "per-tensor-fp8", "no-method", "unquantized", ], @@ -699,7 +715,8 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): 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. ``"fp8"`` is block FP8 only: + per-tensor FP8 names the same method without a block size.""" assert _ckpt_quant_method(config) == expected diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 2a79d004bc6..06936986755 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -474,12 +474,15 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( @needs_qwen3_moe_ep_plan @pytest.mark.parametrize( - ("num_experts", "quant_method", "ep_size", "match"), - [(8, "fp8", 4, "block-fp8"), (2, None, 2, "2 expert-data-parallel replicas")], + ("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"), + ], ids=["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, num_experts, quant_method, ep_size, match + run_as, serve, gpu_memory, tmp_path, num_experts, quantization_config, ep_size, match ): """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, @@ -487,8 +490,8 @@ def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( naming the knob that replicates it. Without expert parallelism neither rule applies.""" run_as(4) config = _qwen3_moe(num_experts=num_experts) - if quant_method: - config.quantization_config = {"quant_method": quant_method} + if quantization_config: + config.quantization_config = quantization_config serve(config) gpu_memory(1) From e85dbbec755cd5fa01ffe488b8888d96a06acee1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:31:32 +0000 Subject: [PATCH 15/25] Take a block-FP8 dequant dtype the config names as a string A config the load is handed can name its dtype, or carry "auto" for the checkpoint's own, and the dequant's Tensor.to reads a string as a device. The dtype a string names is now taken as the torch dtype, and "auto" as the bfloat16 default the dequant emits unrequested. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 +++++- tests/unit/torch/utils/test_fp8_dequant_patch.py | 12 +++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 62402be7a6e..dd21cddde49 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1752,11 +1752,15 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): 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 torch import transformers from packaging.version import parse from transformers.integrations import finegrained_fp8 as _fp8 diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index c5117576adf..887dcd22fec 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -147,12 +147,18 @@ def test_dequant_matches_the_blockwise_reference( @pytest.mark.parametrize( ("target_dtype", "expected"), - [(None, torch.bfloat16), (torch.float16, torch.float16)], - ids=["unrequested", "requested"], + [ + (None, torch.bfloat16), + (torch.float16, torch.float16), + ("float16", torch.float16), + ("auto", torch.bfloat16), + ], + ids=["unrequested", "requested", "named", "auto"], ) 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.""" + 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.""" weight, scale_inv = _block_fp8((10, 13), (3, 4)) model_load_utils._patch_finegrained_fp8_dequant(target_dtype) From ac1ba085a4b27c93dd5e259738f290b1ed1542a2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:32:24 +0000 Subject: [PATCH 16/25] Drop the FP8 dequant tests' skip for a transformers without Fp8Dequantize Every transformers the repo supports ships Fp8Dequantize (5.5 does), so the skip never fires, and the load dereferences the class on the same releases. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_fp8_dequant_patch.py | 2 -- 1 file changed, 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 887dcd22fec..923bc2ea781 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -38,8 +38,6 @@ 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) From a1e7f7e587a3850a6ac3207c3df910ca800564c6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:56:30 +0000 Subject: [PATCH 17/25] Take the destination dtype the loader passes to the block-FP8 dequant The override installs below transformers 5.10.1 and took exactly (quantized, scales): 5.8.0 to 5.10.0 call it that way, and 5.10.1, where the gate switches to transformers' own dequant, is the first to pass output_dtype. A caller passing it would have stopped the load with a bare TypeError. It now takes output_dtype and, when given, emits it: the loader passes the destination dtype, which the module-level model dtype only reconstructs. The expert-converter test runs the patched dequant through this transformers' chain, which passes output_dtype, instead of restoring the stock one around it. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 5 ++-- .../torch/utils/test_fp8_dequant_patch.py | 26 +++++++++---------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index dd21cddde49..f663628bcc4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1795,7 +1795,7 @@ def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): _FINEGRAINED_FP8_PATCHED = True return - def _dequantize_one(self, quantized, scales): + def _dequantize_one(self, quantized, scales, output_dtype=None): # 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 ( @@ -1809,7 +1809,8 @@ def _dequantize_one(self, quantized, scales): # 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 + # A loader that passes the destination dtype, as transformers' does from 5.10.1, knows it best. + out_dtype = output_dtype or _FP8_DEQUANT_DTYPE or torch.bfloat16 if scale_rows == 0 or scale_cols == 0 or scales.numel() == 0: # Only an empty weight has no scales. A non-empty one with an empty grid was sliced # apart from its scales (as 5.9's expert tensor_idx bug did): emitting it unscaled would diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 923bc2ea781..f3a55078831 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -144,23 +144,27 @@ def test_dequant_matches_the_blockwise_reference( @pytest.mark.parametrize( - ("target_dtype", "expected"), + ("target_dtype", "output_dtype", "expected"), [ - (None, torch.bfloat16), - (torch.float16, torch.float16), - ("float16", torch.float16), - ("auto", torch.bfloat16), + (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"], + ids=["unrequested", "requested", "named", "auto", "passed-by-the-loader"], ) -def test_dequant_emits_the_requested_model_dtype(pre_5_10_transformers, target_dtype, expected): +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.""" + 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) + 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)) @@ -191,11 +195,7 @@ def test_the_expert_converter_leads_with_a_merge_that_dequantizes_each_expert( 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", From b5faa81623fbe55d440fe4e9c1d2cb8e6848e34a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:57:15 +0000 Subject: [PATCH 18/25] Pad a partial-block FP8 weight in its own dtype for transformers' dequant From transformers 5.10.1 a partial last block is padded to whole blocks for transformers' own dequant. The padding cast the weight to fp32 first, so the dequant got an input dtype it never sees from the loader, and the load an fp32 copy it does not need. The weight is now padded through a same-size uint8 view, in its own fp8 dtype; the zero byte is zero in every fp8 format. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f663628bcc4..ee07b3919be 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1787,7 +1787,11 @@ def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): ): 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)) + # 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] From e0d0e88a412ffa6c6ff78f7d2e323ce5ecf74e56 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:58:20 +0000 Subject: [PATCH 19/25] Keep the 5.10.1 dequant wrapper from wrapping itself on a re-import From transformers 5.10.1 the patch wraps transformers' dequant, calling the method it replaces, and only this module's flag kept it from wrapping again. Re-importing the module loses that flag, so the wrapper wrapped itself, one frame per re-import. The dequant op now carries a flag of its own, as the FP8 quantizer does for its converter hook. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 6 ++++ .../torch/utils/test_fp8_dequant_patch.py | 33 +++++++++++++------ 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ee07b3919be..ac6a7bc9eb2 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1766,6 +1766,11 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): from transformers.integrations import finegrained_fp8 as _fp8 if parse(transformers.__version__) >= parse("5.10.1"): + # Re-importing this module loses its flag, and the wrapper calls the method it replaces: + # the class keeps a flag of its own, so the wrapper never wraps itself. + if getattr(_fp8.Fp8Dequantize, "_modelopt_dequant_patched", False): + _FINEGRAINED_FP8_PATCHED = True + return _stock_dequantize_one = _fp8.Fp8Dequantize._dequantize_one fp8_dtypes = { getattr(torch, name) @@ -1796,6 +1801,7 @@ def _dequantize_partial_blocks(self, quantized, 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 diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index f3a55078831..42296c5ba52 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -47,19 +47,21 @@ 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. +# 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 _PATCH_TARGETS] + 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) @@ -249,19 +251,30 @@ def test_patching_again_keeps_the_install_and_takes_the_new_dtype(pre_5_10_trans 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 +@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 quantizer class carries its own.""" - if not hasattr(FineGrainedFP8HfQuantizer, "update_weight_conversions"): - pytest.skip("this transformers' FP8 quantizer has no update_weight_conversions") + """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() - hook = vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] + installed = vars(cls)[name] monkeypatch.setattr(model_load_utils, "_FINEGRAINED_FP8_PATCHED", False) model_load_utils._patch_finegrained_fp8_dequant() - assert vars(FineGrainedFP8HfQuantizer)["update_weight_conversions"] is hook + assert vars(cls)[name] is installed def test_undoing_the_patch_restores_stock_transformers(pre_5_10_transformers, monkeypatch): From 5787881c94802ee0dcf0c1cbb053ffd0601223e4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:59:36 +0000 Subject: [PATCH 20/25] Test that the dequantize config keeps the checkpoint's block size The expert-parallel load passes FineGrainedFP8Config(dequantize=True), and both dequant branches read weight_block_size off the quantizer's config, relying on transformers to keep the checkpoint's quantization_config and take only the loading attributes from the one passed. A test now pins that merge on whichever transformers runs the suite. Signed-off-by: Shengliang Xu --- .../torch/utils/test_fp8_dequant_patch.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 42296c5ba52..429ff454ae3 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -145,6 +145,25 @@ def test_dequant_matches_the_blockwise_reference( assert torch.equal(out, _reference(weight, scale_inv, block)) +def test_the_dequantize_config_keeps_the_checkpoints_block_size(): + """The expert-parallel load passes ``FineGrainedFP8Config(dequantize=True)``, and every block + decision of the dequant reads ``weight_block_size`` off the quantizer's config: transformers + keeps the checkpoint's, and takes only the loading attributes from the one passed.""" + from transformers.quantizers.auto import AutoHfQuantizer + + checkpoint = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [64, 64], + } + merged = AutoHfQuantizer.merge_quantization_configs( + checkpoint, FineGrainedFP8Config(dequantize=True) + ) + + assert list(merged.weight_block_size) == [64, 64] + assert merged.dequantize + + @pytest.mark.parametrize( ("target_dtype", "output_dtype", "expected"), [ From b1f4a678661faa75c9554eaca686a2c12004e623 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:59:55 +0000 Subject: [PATCH 21/25] Say why the block-fp8 dense rule applies under expert parallelism only Its reason, that Distributed Checkpoint reads block FP8 as stored, holds at any EP degree, but its remedy, replicating the dense part to dequantize it as it loads, holds only where the expert-parallel load runs. Without expert parallelism the FSDP2 load refuses a quantized checkpoint itself. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ac6a7bc9eb2..7016a91c0ce 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1377,6 +1377,9 @@ def _decide_distribution( "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) elif ep_size > 1 and _ckpt_quant_method(hf_config) == "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" From 7ee59e8466492d3d6c179dbd78f37afd5c5d127c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 08:00:37 +0000 Subject: [PATCH 22/25] Name the missing weight_block_size when refusing a per-tensor FP8 checkpoint An fp8 checkpoint without weight_block_size was told the load "cannot dequantize a 'fp8-per-tensor' checkpoint" while listing fp8 among what it dequantizes, and the refusal never named the block size that tells the two apart. It now says the load dequantizes block-fp8 checkpoints, which name their block in quantization_config.weight_block_size, and that this one names none. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 8 ++++++++ tests/unit/torch/utils/test_ep_loader.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7016a91c0ce..e5e4af1a5d1 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2313,6 +2313,14 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A 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 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)) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 5128a1ee111..df1b0dd83c4 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -222,7 +222,7 @@ def resolve(path, rank): quantization_config={"quant_method": "fp8", "activation_scheme": "static"} ), NotImplementedError, - "cannot dequantize a 'fp8-per-tensor'", + "this fp8 checkpoint names none", ), (SimpleNamespace(), ValueError, "declare base_model_ep_plan"), ( From 25144cf862fb6b0c99e2b890b371d9b0357eca5b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 08:27:24 +0000 Subject: [PATCH 23/25] Read an FP8 config object's block from the checkpoint's own config A per-tensor FP8 checkpoint names no weight_block_size, which is how the load tells it from block FP8 and refuses it. transformers' FP8 config object fills in a 128x128 block whatever the checkpoint names, so a per-tensor checkpoint whose config came in as that object, as a loaded model's does, read as block FP8 and would load unscaled. For an FP8 config object, the block now comes from the checkpoint's own config.json, where both the decision and the load can read it. The plan decision also stops offering an EP degree for a block-fp8 dense part sharded within several expert-data-parallel replicas: no EP degree lets that dense part shard. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 39 ++++++++++++++++--- tests/unit/torch/utils/test_ep_loader.py | 21 ++++++++++ .../utils/test_resolve_distribution_plan.py | 14 +++++-- 3 files changed, 66 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e5e4af1a5d1..89b7599c7c0 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -18,6 +18,7 @@ import contextlib import copy import inspect +import json import logging import os import re @@ -1367,16 +1368,20 @@ def _decide_distribution( # 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, or " + " Set MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." + if block_fp8 + else " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, or " "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." ) - elif ep_size > 1 and _ckpt_quant_method(hf_config) == "fp8": + 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). @@ -2047,13 +2052,15 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) -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``. Per-tensor FP8 (AutoFP8's), which names the same method as - block FP8 but no ``weight_block_size``, is ``"fp8-per-tensor"``. + 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 it, the object's block stands. """ for c in (cfg, getattr(cfg, "text_config", None), getattr(cfg, "llm_config", None)): qc = getattr(c, "quantization_config", None) @@ -2061,6 +2068,8 @@ def _ckpt_quant_method(cfg): if m: # transformers' config objects carry a QuantizationMethod enum, whose str() is not its value. m = str(getattr(m, "value", m)) + if not isinstance(qc, dict) and m == "fp8": + qc = _config_json_quantization(ckpt_path) or qc block = ( qc.get("weight_block_size") if isinstance(qc, dict) @@ -2071,6 +2080,26 @@ def _ckpt_quant_method(cfg): return None +def _config_json_quantization(ckpt_path: str | None) -> dict | None: + """The ``quantization_config`` the checkpoint's own ``config.json`` names, or ``None``. + + Read as stored: the top level, else a VLM's text or LLM config, where :func:`_ckpt_quant_method` + looks too. + """ + if not ckpt_path: + return None + try: + with open(os.path.join(ckpt_path, "config.json")) as f: + raw = json.load(f) + except (OSError, ValueError): + return None + for c in (raw, raw.get("text_config"), raw.get("llm_config")): + qc = c.get("quantization_config") if isinstance(c, dict) else None + if isinstance(qc, dict): + return qc + return None + + def _transformers_mesh_builder(): """The module whose ``initialize_distributed_mesh`` ``from_pretrained`` calls, or ``None``. @@ -2307,7 +2336,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A if attn_implementation is not None: config_kwargs["attn_implementation"] = attn_implementation hf_config = AutoConfig.from_pretrained(ckpt_path, **config_kwargs) - quant_method = _ckpt_quant_method(hf_config) + quant_method = _ckpt_quant_method(hf_config, ckpt_path) if plan.is_fsdp_dense and quant_method == "fp8": raise NotImplementedError( f"{plan!r}: the dense part of a block-fp8 checkpoint cannot be sharded: Distributed " diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index df1b0dd83c4..91e87dbfcf0 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -24,6 +24,7 @@ import copy import dataclasses +import json from contextlib import nullcontext from types import SimpleNamespace @@ -720,6 +721,26 @@ def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expect assert _ckpt_quant_method(config) == expected +@pytest.mark.parametrize( + ("on_disk", "expected"), + [ + ({"quant_method": "fp8", "activation_scheme": "static"}, "fp8-per-tensor"), + ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8"), + ], + ids=["per-tensor", "block"], +) +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.""" + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": on_disk})) + config = SimpleNamespace(quantization_config=FineGrainedFP8Config()) + + assert _ckpt_quant_method(config, str(tmp_path)) == expected + + def test_from_pretrained_builds_its_mesh_from_the_ep_mesh_it_is_given(): """Newer transformers ignores ``from_pretrained``'s ``device_mesh`` and builds a world-sized one, which shards the experts over every rank when EP is below the world size. During the load it gets diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 06936986755..922f395e345 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -478,8 +478,14 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit( [ (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"], + 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, num_experts, quantization_config, ep_size, match @@ -487,7 +493,8 @@ def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( """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. Without expert parallelism neither rule applies.""" + 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) config = _qwen3_moe(num_experts=num_experts) if quantization_config: @@ -497,8 +504,9 @@ def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( 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"): + 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 resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan( 4, dense="fsdp" ) From a9bc6519db53548ab9e300937d57e740915e2b52 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 08:27:54 +0000 Subject: [PATCH 24/25] Check a weight's dims before the FP8 dequant patches unpack its shape Both patches unpacked the last two dims first, so a weight or scale grid of fewer than two dims failed with a bare unpack error from inside the patched method. From 5.10.1 such a call now goes to transformers' own dequant, as every case the patch does not place does; below it, the override refuses it by name. The override also drops its FP4 branch: it called an unpack transformers only added in 5.10.1, so it could not run on the releases the override replaces. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 24 +++++++++---------- .../torch/utils/test_fp8_dequant_patch.py | 13 ++++++---- 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 89b7599c7c0..369c8f2bf64 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1787,16 +1787,15 @@ def _patch_finegrained_fp8_dequant(target_dtype=None): } 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 ( - quantized.dtype not in fp8_dtypes - or scales.dim() < 2 - or tuple(scales.shape[-2:]) != ceil_grid - or (rows % block[0] == 0 and cols % block[1] == 0) + 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 @@ -1814,14 +1813,13 @@ def _dequantize_partial_blocks(self, quantized, scales, *args, **kwargs): return def _dequantize_one(self, quantized, scales, output_dtype=None): - # 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, copy=True) + # Block FP8 only: the releases this replaces have no FP4 unpack (5.10.1 added it). + if quantized.dim() < 2 or scales.dim() < 2: + raise ValueError( + f"block-FP8 dequant: weight {tuple(quantized.shape)} with scales " + f"{tuple(scales.shape)} is not a block-scaled 2-D weight." + ) + q = quantized.to(torch.float32, copy=True) rows, cols = q.shape[-2:] scale_rows, scale_cols = scales.shape[-2:] # Emit the MODEL compute dtype (bf16), NOT the scale dtype. DeepSeek's weight_scale_inv is F32; diff --git a/tests/unit/torch/utils/test_fp8_dequant_patch.py b/tests/unit/torch/utils/test_fp8_dequant_patch.py index 429ff454ae3..d9e378f66c1 100644 --- a/tests/unit/torch/utils/test_fp8_dequant_patch.py +++ b/tests/unit/torch/utils/test_fp8_dequant_patch.py @@ -193,13 +193,18 @@ def test_dequant_emits_the_requested_model_dtype( @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"], + [ + ((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), or a 4-row grid over 10 rows, - which no 4-row block (the configured one) or even split gives.""" + 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) From 94ea7f37d2fe86762eea9c6ad37179373de01559 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 08:44:51 +0000 Subject: [PATCH 25/25] Refuse an FP8 config object whose checkpoint config cannot be read The per-tensor check reads an FP8 config object's block from the checkpoint's own config.json, but a path without one, a Hub ID the load has not made local yet, left the object's default 128x128 block standing, and a per-tensor checkpoint would still load unscaled. Such a config is now "fp8-unverified", which the load refuses, naming the two ways in: the config as AutoConfig.from_pretrained reads it, or a local checkpoint. Reading the file decides it on every rank alike, without fetching anything. For a block-fp8 dense part within several expert-data-parallel replicas, the remedy also names starting from an unquantized copy, which on the warning path is the one that changes the outcome. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 17 ++++++++++++--- tests/unit/torch/utils/test_ep_loader.py | 21 +++++++++++++------ .../utils/test_resolve_distribution_plan.py | 1 + 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 369c8f2bf64..abdab66a157 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1376,7 +1376,8 @@ def _decide_distribution( ) # No EP degree lets a block-fp8 dense part shard (below), so none is offered for one. remedy = ( - " Set MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." + " 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, or " "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU." @@ -2058,7 +2059,8 @@ def _ckpt_quant_method(cfg, ckpt_path: str | None = None): ``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 it, the object's block stands. + ``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) @@ -2067,7 +2069,9 @@ def _ckpt_quant_method(cfg, ckpt_path: str | None = None): # transformers' config objects carry a QuantizationMethod enum, whose str() is not its value. m = str(getattr(m, "value", m)) if not isinstance(qc, dict) and m == "fp8": - qc = _config_json_quantization(ckpt_path) or qc + qc = _config_json_quantization(ckpt_path) + if qc is None: + return "fp8-unverified" block = ( qc.get("weight_block_size") if isinstance(qc, dict) @@ -2348,6 +2352,13 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "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)) diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 91e87dbfcf0..7971645d6d2 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -225,6 +225,11 @@ def resolve(path, rank): 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"}), @@ -237,6 +242,7 @@ def resolve(path, rank): "modelopt", "compressed-tensors", "per-tensor-fp8", + "fp8-config-object-of-a-hub-id", "no-ep-plan", "no-routed-experts", ], @@ -691,7 +697,7 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): ), "mxfp4", ), - (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8"), + (SimpleNamespace(quantization_config=FineGrainedFP8Config()), "fp8-unverified"), ( SimpleNamespace( quantization_config={"quant_method": "fp8", "activation_scheme": "static"} @@ -716,8 +722,9 @@ def test_an_expert_parallel_model_without_its_mesh_is_refused(): 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. ``"fp8"`` is block FP8 only: - per-tensor FP8 names the same method without a block size.""" + ``"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 @@ -726,16 +733,18 @@ def test_quant_method_comes_from_the_config_or_a_vlms_text_config(config, expect [ ({"quant_method": "fp8", "activation_scheme": "static"}, "fp8-per-tensor"), ({"quant_method": "fp8", "weight_block_size": [128, 128]}, "fp8"), + (None, "fp8-unverified"), ], - ids=["per-tensor", "block"], + ids=["per-tensor", "block", "no-config-json"], ) def test_an_fp8_config_object_is_read_against_the_checkpoints_own_config( tmp_path, on_disk, expected ): """transformers' FP8 config object names a 128x128 block whatever the checkpoint does, so the checkpoint's own ``config.json`` says whether it names one: per-tensor FP8 handed in as an - object is still refused.""" - (tmp_path / "config.json").write_text(json.dumps({"quantization_config": on_disk})) + object is still refused. Without that file there is no telling, which the load refuses.""" + if on_disk is not None: + (tmp_path / "config.json").write_text(json.dumps({"quantization_config": on_disk})) config = SimpleNamespace(quantization_config=FineGrainedFP8Config()) assert _ckpt_quant_method(config, str(tmp_path)) == expected diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 922f395e345..1681c8b5d54 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -507,6 +507,7 @@ def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated( 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" )