From 4dab013f104b1098be5058054a0b22d7b105e77d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:53:21 +0000 Subject: [PATCH 01/16] Decompress compressed-tensors MoE experts on the fly as they load Expert-parallel PTQ loads a checkpoint with from_pretrained over the expert-parallel mesh, so each rank builds only its own experts. A checkpoint whose experts are compressed-tensors quantized -- Kimi-K2.5's INT4, Kimi-K3's MXFP4 -- cannot take that load as it is: transformers' fused-experts conversion expects each expert's plain .weight. _onfly_decompress rewrites those per-expert converters to read the checkpoint's stored keys (weight_packed, weight_scale, ...) and prepends a step that dequantizes each owned expert to bf16 before the merge, so PTQ runs on the original compressed checkpoint with no decompressed copy on disk. It is format-agnostic: the compressed_tensors library does the math, driven by the checkpoint's own quantization_config. It needs transformers 5.8. It also holds the helper that rebuilds a VLM's bundled text config as a native transformers backbone config. _remote_code.stage_full_remote_closure makes transformers stage every module of a local checkpoint's remote code, not only those one import deep, which loading the Kimi checkpoints' own modeling needs. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 356 ++++++++++++ modelopt/torch/models/_remote_code.py | 61 +++ .../torch/models/test_onfly_decompress.py | 517 ++++++++++++++++++ 3 files changed, 934 insertions(+) create mode 100644 modelopt/torch/models/_onfly_decompress.py create mode 100644 modelopt/torch/models/_remote_code.py create mode 100644 tests/unit/torch/models/test_onfly_decompress.py diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py new file mode 100644 index 00000000000..14d5ffac047 --- /dev/null +++ b/modelopt/torch/models/_onfly_decompress.py @@ -0,0 +1,356 @@ +# 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. + +"""On-the-fly decompression of compressed-tensors quantized MoE experts into the transformers-5.x +fused-experts layout, DURING from_pretrained — so a2a PTQ can run on the ORIGINAL compressed +checkpoint with no decompressed intermediate on disk, EP-sharded (each rank decompresses only its +own experts). + +FORMAT-AGNOSTIC. The dequant math is delegated to the ``compressed_tensors`` library and driven +entirely by the checkpoint's own ``quantization_config`` (the ``format`` string + the per-group +``QuantizationArgs``). Any format the library registers works with ZERO per-format code here: +``pack-quantized`` (INT, symmetric or asymmetric, any bit-width / group_size / strategy), +``nvfp4-pack-quantized`` (FP4 + FP8 block scale + FP32 global), ``naive-quantized`` / +``float-quantized`` / ``int-quantized`` (unpacked FP8/INT), etc. The two universal hooks the library +exposes per format are all we need: + - ``compressor.compression_param_names(scheme)`` -> the exact set of stored keys for one weight + (e.g. ``("weight_packed","weight_scale","weight_shape")`` for pack, ``("weight","weight_scale")`` + for float), i.e. which per-expert sibling tensors to pull. + - ``compressor.decompress({: tensor}, scheme)["weight"]`` -> the dequantized weight. +``compressed_tensors`` must be importable at load time, and transformers must be 5.8 or newer: earlier +releases' ``Concatenate`` looks up every source pattern, including the sibling keys the decompression +already consumed. + +The model-specific knobs remain parameters; defaults target moonshotai/Kimi-K2.7-Code (DeepSeek-V3 +backbone): + - install_onfly_decompress(hf_config) — auto-detects format + quant args from the config + - to_native_backbone_config(vlm_config, ...) — build the native backbone config + +Mechanism (mirrors transformers' FP8 quantizer ``update_weight_conversions``): the fused-MoE model +already ships a ``WeightConverter`` that fuses per-expert ``gate/up/down_proj.weight`` -> 3-D +``gate_up_proj``/``down_proj`` via ``MergeModulelist``+``Concatenate``. We rewrite those converters so +each ``.weight`` source becomes the checkpoint's actual stored keys for that weight +(``compression_param_names``), and PREPEND a ``DecompressExperts`` op that pairs those per expert and +dequantizes -> bf16 BEFORE the merge/concat collapse the per-expert structure. +``install_onfly_decompress()`` monkeypatches ``get_model_conversion_mapping``. +""" + +import torch + + +def to_native_backbone_config( + vlm_config, + model_type: str = "deepseek_v3", + architectures: "tuple[str, ...]" = ("DeepseekV3ForCausalLM",), +): + """Derive the transformers-native fused backbone config from a VLM config's ``text_config``, so + ``from_pretrained`` builds the native model directly from the ORIGINAL checkpoint's compressed + shards — NO preprocessed native-config "view" dir. + + A VLM's ``text_config`` is typically the bundled remote-code config (its ``auto_map`` points at a + ModuleList modeling file, and it carries the stale compressed-tensors ``quantization_config``). + Dropping ``auto_map`` + ``quantization_config`` and setting ``model_type`` to the native type makes + it resolve to the stock transformers native config (which, for the a2a path, supplies + ``base_model_ep_plan``). ``architectures`` must be set to the native class (NOT left None — hf_ptq's + VL-detection iterates ``config.architectures``). + + Defaults target the Kimi DeepSeek-V3 backbone; pass ``model_type`` / ``architectures`` for another. + """ + from transformers.models.auto.configuration_auto import CONFIG_MAPPING + + text_config = getattr(vlm_config, "text_config", None) + if text_config is None: + raise ValueError("to_native_backbone_config: config has no .text_config (not a VLM?)") + if model_type not in CONFIG_MAPPING: + raise ValueError(f"to_native_backbone_config: unknown native model_type {model_type!r}") + d = text_config.to_dict() if hasattr(text_config, "to_dict") else dict(text_config) + for key in ("auto_map", "quantization_config", "_name_or_path"): + d.pop(key, None) + d["model_type"] = model_type + d["architectures"] = list(architectures) + return CONFIG_MAPPING[model_type].from_dict(d) + + +def _find_quantization_config(hf_config): + """Locate the compressed-tensors ``quantization_config`` (top-level or under ``text_config``) and + return it as a plain dict.""" + qc = getattr(hf_config, "quantization_config", None) + if qc is None: + tc = getattr(hf_config, "text_config", None) + qc = getattr(tc, "quantization_config", None) if tc is not None else None + if qc is None and isinstance(hf_config, dict): + qc = hf_config.get("quantization_config") or (hf_config.get("text_config") or {}).get( + "quantization_config" + ) + if qc is None: + raise ValueError( + "install_onfly_decompress: config has no compressed-tensors quantization_config" + ) + if hasattr(qc, "to_dict"): + qc = qc.to_dict() + elif not isinstance(qc, dict): + qc = dict(qc) + return qc + + +def _read_compression(hf_config): + """Parse the checkpoint's compressed-tensors ``quantization_config`` and return + ``(compressor, scheme, param_names, format_str)`` via the ``compressed_tensors`` registry. + + ``compressor`` is the registered compressor for ``format`` (its ``decompress`` classmethod is the + dequant); ``scheme`` is the ``QuantizationScheme`` (supplies ``scheme.weights`` to ``decompress``); + ``param_names`` is the exact set of stored sibling keys per weight. The experts' group is chosen as + the sole config group, else the one whose ``targets`` reference "expert", else the first.""" + # Look the config up first: with none, there is nothing to decompress (ValueError), and + # compressed_tensors need not even be installed. + quantization_config = _find_quantization_config(hf_config) + from compressed_tensors.compressors.base import BaseCompressor + from compressed_tensors.quantization import QuantizationConfig + + qconfig = QuantizationConfig.model_validate(quantization_config) + groups = qconfig.config_groups + if len(groups) == 1: + scheme = next(iter(groups.values())) + else: + scheme = next( + ( + g + for g in groups.values() + if any("expert" in str(t).lower() for t in (getattr(g, "targets", None) or [])) + ), + next(iter(groups.values())), + ) + fmt = qconfig.format + fmt = fmt.value if hasattr(fmt, "value") else str(fmt) + compressor = BaseCompressor.get_value_from_registry(fmt) + param_names = tuple(compressor.compression_param_names(scheme)) + return compressor, scheme, param_names, fmt + + +def _make_op_class(compressor, scheme, param_names): + """Build the per-expert decompress op lazily (so this module imports without transformers present), + bound to a compressed_tensors compressor + scheme + the checkpoint's stored param names. The FIRST + param name is the "anchor" (the source key the decompressed list is emitted under, replacing the + original ``.weight`` so the downstream MergeModulelist/Concatenate fuse exactly as uncompressed). + + NOTE: ``DecompressExperts`` subclasses ``MergeModulelist`` so that it shards experts AT LOAD. The + transformers loader (``convert_and_load_state_dict_in_model``) only computes the per-expert + ``shard_index`` (``tensor_idx``) when ``isinstance(converter.operations[0], MergeModulelist)``; that + ``tensor_idx`` is what makes ``GroupedGemmParallel.shard_tensor`` keep ONLY this rank's local experts + (returning ``None`` for the rest). If our op were a plain ``ConversionOps`` prepended before the + original ``MergeModulelist``, ``operations[0]`` would not match, ``tensor_idx`` would be ``None``, and + every rank would materialize + decompress ALL experts (garbled, and a ~device-max transient -> OOM at + scale, e.g. Kimi-K2.7-Code EP32). Subclassing keeps ``operations[0]`` a ``MergeModulelist`` instance so + per-expert sharding fires; we still override ``convert`` to decompress (the ``dim``/merge behavior of + the base class is unused -- the ORIGINAL ``MergeModulelist`` later in the chain does the stacking).""" + from transformers.core_model_loading import MergeModulelist + + anchor = "." + param_names[0] + others = [(pn, "." + pn) for pn in param_names[1:]] + + class DecompressExperts(MergeModulelist): + """First op in a fused-experts WeightConverter chain. + + For each anchor source (e.g. ``*.weight_packed``), pair it with its sibling stored params (``*.weight_scale``, + ``*.weight_shape``, ``*.weight_zero_point``, ``*.weight_global_scale``, ... — whatever the format stores), + decompress each expert via ``compressed_tensors`` -> bf16, and emit the decompressed list UNDER the anchor key + (dropping the siblings). The downstream MergeModulelist/Concatenate then fuse as for an uncompressed checkpoint. + + Subclasses ``MergeModulelist`` purely so the loader recognizes ``operations[0]`` and shards experts + at load (see ``_make_op_class`` docstring); non-local experts arrive as ``None`` (from + ``GroupedGemmParallel.shard_tensor``) and are skipped, so only this rank's local experts are + decompressed. + """ + + def __init__(self): + super().__init__( + dim=0 + ) # dim unused (convert overridden); needed so isinstance -> shard_index + + @torch.no_grad() + def convert(self, input_dict, *args, **kwargs): + out = {} + for pat in list(input_dict.keys()): + marker = "$" if pat.endswith("$") else "" + base = pat[:-1] if marker else pat + if not base.endswith(anchor): + continue + stem = base[: -len(anchor)] + lists = {param_names[0]: input_dict.pop(pat)} + # Pop each sibling if PRESENT; a checkpoint may store only a SUBSET of the scheme's + # possible params (e.g. mxfp4-pack-quantized declares weight_global_scale via + # compression_param_names, but the public moonshotai/Kimi-K3 checkpoint stores only + # weight_packed + weight_scale). Missing optional siblings simply aren't passed to the + # compressor (whose decompress uses whatever the scheme requires). + for pn, suf in others: + sib = input_dict.pop(stem + suf + marker, None) + if sib is not None: + lists[pn] = sib + anchor_list = lists[param_names[0]] + present = [pn for pn in param_names if pn in lists] + # Non-local experts were sharded away at load (GroupedGemmParallel -> None); decompress + # only the local experts this rank owns. Keeps peak memory at the per-rank shard. + out[pat] = [ + compressor.decompress({pn: lists[pn][i] for pn in present}, scheme)[ + "weight" + ].to(torch.bfloat16) + for i in range(len(anchor_list)) + if anchor_list[i] is not None + ] + out.update(input_dict) # nothing else expected, but pass through defensively + return out + + return DecompressExperts + + +def _compressedify(conv, compressor, scheme, param_names): + """Rewrite a per-expert WeightConverter (.weight sources) into a compressed-aware one. + + each ``.weight`` -> the checkpoint's stored keys ``[.$ for pn in param_names]`` (anchor first, preserving + gate/up order so Concatenate fuses gate-then-up), and prepend DecompressExperts. + """ + from transformers.core_model_loading import WeightConverter + + DecompressExperts = _make_op_class(compressor, scheme, param_names) # noqa: N806 (a class, keeping its own name) + by_pname = {pn: [] for pn in param_names} + for p in conv.source_patterns: + base = p.removesuffix("$") + if base.endswith(".weight"): + stem = base[: -len(".weight")] + for pn in param_names: + by_pname[pn].append(stem + "." + pn + "$") + else: + by_pname[param_names[0]].append(p) # leave non-weight sources as-is (none expected) + new_sources = [s for pn in param_names for s in by_pname[pn]] + tgt = getattr(conv, "_original_target_patterns", conv.target_patterns) + return WeightConverter( + source_patterns=new_sources, + target_patterns=tgt, + operations=[DecompressExperts(), *conv.operations], + ) + + +# id(converter) -> {source_pattern: [expert indices in append order]}; used to re-sort experts numerically. +_EXPERT_ORDER: dict = {} + + +def _install_numeric_expert_ordering(): + """Patch WeightConverter.add_tensor / materialize_tensors so per-expert tensors are merged in NUMERIC index order. + + The loader appends in checkpoint-iteration order, which is LEXICOGRAPHIC for many indices (0,1,10,100,...,2,...) — + that scrambles fused-expert slots. A no-op for numerically-ordered checkpoints. + """ + import re as _re + + from transformers.core_model_loading import WeightConverter + + if getattr(WeightConverter.add_tensor, "_onfly_patched", False): + return + orig_add = WeightConverter.add_tensor + orig_mat = WeightConverter.materialize_tensors + + def add_tensor(self, target_key, source_key, source_pattern, future): + orig_add(self, target_key, source_key, source_pattern, future) + m = _re.search(r"experts\.(\d+)\.", source_key) + if m is not None: + _EXPERT_ORDER.setdefault(id(self), {}).setdefault(source_pattern, []).append( + int(m.group(1)) + ) + + def materialize_tensors(self): + order = _EXPERT_ORDER.pop(id(self), None) + if order: + for pat, idxs in order.items(): + lst = self.collected_tensors.get(pat) + if lst is not None and len(lst) == len(idxs): + self.collected_tensors[pat] = [ + f for _, f in sorted(zip(idxs, lst), key=lambda x: x[0]) + ] + return orig_mat(self) + + setattr(add_tensor, "_onfly_patched", True) + WeightConverter.add_tensor = add_tensor + WeightConverter.materialize_tensors = materialize_tensors + + +def _install_keep_int_native(): + """Patch _materialize_copy so INTEGER source tensors are never cast to a FLOATING dtype during load. + + Without an hf_quantizer the loader casts every source to the target param dtype (bf16); that destroys the integer + bit pattern of packed weights (`weight_packed`), shapes (`weight_shape`), zero-points and g_idx. Floats (the scale, + global scale, dense weights) still cast normally. + """ + import torch + import transformers.core_model_loading as cml + + if getattr(cml._materialize_copy, "_onfly_patched", False): + return + + def _materialize_copy(tensor, device=None, dtype=None): + t = tensor[...] + if ( + dtype is not None + and not t.is_floating_point() + and torch.empty((), dtype=dtype).is_floating_point() + ): + dtype = None # keep packed/shape/zero-point/g_idx integers native + if dtype is not None or device is not None: + t = t.to(device=device, dtype=dtype) + return t + + setattr(_materialize_copy, "_onfly_patched", True) + cml._materialize_copy = _materialize_copy + + +def install_onfly_decompress(hf_config): + """Monkeypatch get_model_conversion_mapping so the fused-MoE per-expert converters decompress the + checkpoint's compressed-tensors weights on the fly, in whatever ``format`` the checkpoint declares. + Idempotent. Call once before from_pretrained of a compressed checkpoint. ``hf_config`` is the model + config carrying the compressed-tensors ``quantization_config`` (top-level or under ``text_config``); + the format, bit-width, group size, strategy and stored param names are all read from it.""" + import transformers.modeling_utils as mu + from transformers.core_model_loading import WeightConverter + + from modelopt.torch.utils.moe import require_transformers + + if getattr(mu.get_model_conversion_mapping, "_onfly_patched", False): + return + # Raises ValueError for a config with nothing to decompress, before anything is patched. + compressor, scheme, param_names, fmt = _read_compression(hf_config) + require_transformers("5.8", "on-the-fly compressed-tensors decompression") + _install_numeric_expert_ordering() + _install_keep_int_native() + orig = mu.get_model_conversion_mapping + + def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): + convs = orig(model, key_mapping=key_mapping, hf_quantizer=hf_quantizer, *a, **k) + out = [] + for c in convs: + # only the per-expert MoE converters (have `.weight` sources under `experts`) + if isinstance(c, WeightConverter) and any( + "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns + ): + out.append(_compressedify(c, compressor, scheme, param_names)) + else: + out.append(c) + return out + + setattr(patched, "_onfly_patched", True) + mu.get_model_conversion_mapping = patched + print( + f"[onfly_decompress] installed: compressed_tensors '{fmt}' fused-MoE experts decompress " + f"(stored params={param_names}) -> bf16 during from_pretrained", + flush=True, + ) diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py new file mode 100644 index 00000000000..43780f39c65 --- /dev/null +++ b/modelopt/torch/models/_remote_code.py @@ -0,0 +1,61 @@ +# 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. + +"""Loading the modeling code a checkpoint ships with itself.""" + +import glob +import os +import shutil +import warnings + + +def stage_full_remote_closure() -> None: + """Make transformers stage every module of a local checkpoint's remote code, not just one hop. + + transformers' ``get_cached_module_file`` copies a remote module and its *direct* relative imports + into ``transformers_modules///``, but does not recurse. Remote code more than one import + deep -- Kimi-K2.5's ``modeling_kimi_k25 -> modeling_deepseek -> configuration_deepseek`` -- is + then staged incomplete, and the later recursive walk in ``get_relative_import_files`` fails with + ``FileNotFoundError``, from ``get_class_from_dynamic_module`` and ``from_pretrained`` alike. + + Wrap ``get_cached_module_file`` so that, loading from a local checkpoint dir, it copies every + sibling ``*.py`` into the stage dir it just populated. Idempotent. + """ + import transformers.dynamic_module_utils as dmu + from transformers.utils import HF_MODULES_CACHE + + if getattr(dmu.get_cached_module_file, "_modelopt_full_closure", False): + return + orig = dmu.get_cached_module_file + + def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): + resolved = orig(pretrained_model_name_or_path, module_file, *args, **kwargs) + src = str(pretrained_model_name_or_path) + if os.path.isdir(src): + staged = str(resolved) + if not os.path.isabs(staged): + staged = os.path.join(HF_MODULES_CACHE, staged) + stage_dir = os.path.dirname(staged) + try: + for py in glob.glob(os.path.join(src, "*.py")): + dst = os.path.join(stage_dir, os.path.basename(py)) + if not os.path.exists(dst): + shutil.copy(py, dst) + except OSError as e: # a failed copy only matters if a deeper import needs the file + warnings.warn(f"Could not stage the remote code of {src}: {e}") + return resolved + + setattr(_staged, "_modelopt_full_closure", True) + dmu.get_cached_module_file = _staged diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py new file mode 100644 index 00000000000..45a71568d58 --- /dev/null +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -0,0 +1,517 @@ +# 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. + +"""On-the-fly decompression of compressed-tensors MoE experts while ``from_pretrained`` loads. + +The tests drive transformers' own weight-loading objects the way its loader does -- +``rename_source_key`` routes a checkpoint key, ``add_tensor`` collects it, ``convert`` fuses it -- +with a stand-in compressor, so they need neither a checkpoint nor ``compressed_tensors``. The two +tests that use the real ``compressed_tensors`` registry skip when it cannot be imported. +""" + +import copy +import sys +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest +import torch + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.6"): + pytest.skip("needs the transformers >= 5.6 weight-loading API", allow_module_level=True) + +import transformers.core_model_loading as cml +import transformers.modeling_utils as mu +from transformers import DeepseekV3Config, PretrainedConfig +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import Chunk, MergeModulelist, WeightConverter, WeightRenaming + +from modelopt.torch.models import _onfly_decompress as onfly + +PACKED = ("weight_packed", "weight_scale", "weight_shape") # what pack-quantized stores per weight + + +@pytest.fixture +def restore_loader_hooks(monkeypatch): + """Undo the process-wide transformers hooks the decompressor installs.""" + for owner, name in ( + (WeightConverter, "add_tensor"), + (WeightConverter, "materialize_tensors"), + (cml, "_materialize_copy"), + (mu, "get_model_conversion_mapping"), + ): + monkeypatch.setattr(owner, name, getattr(owner, name)) + monkeypatch.setattr(onfly, "_EXPERT_ORDER", {}) + + +class _ScaledCompressor: + """Stands in for a compressed_tensors compressor: ``weight = weight_packed * weight_scale``.""" + + @classmethod + def decompress(cls, state_dict, scheme): + return {"weight": state_dict["weight_packed"].float() * state_dict["weight_scale"]} + + +def _stock_expert_converters(): + """transformers' own per-expert -> fused converters for DeepSeek-V3: ``(gate_up, down)``.""" + by_target = { + conv.target_patterns[0]: conv + for conv in get_checkpoint_conversion_mapping("deepseek_v3") + if isinstance(conv, WeightConverter) + } + return by_target["mlp.experts.gate_up_proj"], by_target["mlp.experts.down_proj"] + + +def _load(converters, checkpoint): + """Route, collect and fuse ``checkpoint`` the way transformers' loader does. + + The first converter whose sources match a key claims it; every fused target gets its own deep + copy of that converter, fed in checkpoint order with completed futures (async loading). + """ + per_target = {} + for key, tensor in checkpoint.items(): + for conv in converters: + target, pattern = conv.rename_source_key(key) + if pattern is not None: + break + else: + raise AssertionError(f"no converter claims {key}") + future = Future() + future.set_result(tensor) + per_target.setdefault(target, copy.deepcopy(conv)).add_tensor(target, key, pattern, future) + fused = {} + for target, conv in per_target.items(): + fused.update(conv.convert(target)) + return fused + + +# ------------------------------------------------------------- rewriting the per-expert converters + + +def test_weight_sources_become_the_stored_keys_anchor_first_and_gate_before_up(): + """Each ``.weight`` source becomes the checkpoint's stored keys, grouped anchor first with gate + ahead of up in every group (``Concatenate`` stacks in source order, so rows stay [gate; up]). + Targets are kept, and the decompress op runs first, as a ``MergeModulelist`` so that loaders + which only shard experts when the first op is one still do.""" + gate_up, _ = _stock_expert_converters() + rewritten = onfly._compressedify(gate_up, _ScaledCompressor, None, PACKED) + + assert rewritten.source_patterns == [ + "mlp.experts.*.gate_proj.weight_packed$", + "mlp.experts.*.up_proj.weight_packed$", + "mlp.experts.*.gate_proj.weight_scale$", + "mlp.experts.*.up_proj.weight_scale$", + "mlp.experts.*.gate_proj.weight_shape$", + "mlp.experts.*.up_proj.weight_shape$", + ] + assert rewritten.target_patterns == ["mlp.experts.gate_up_proj"] + decompress, *stock_ops = rewritten.operations + assert isinstance(decompress, MergeModulelist) + assert [repr(op) for op in stock_ops] == [repr(op) for op in gate_up.operations] + + +def test_stored_keys_are_anchored_so_a_weight_never_claims_its_scale(): + """Float formats store ``weight`` next to ``weight_scale``, and transformers regex-*searches* + source patterns: an unanchored ``gate_proj.weight`` would also claim ``gate_proj.weight_scale``.""" + gate_up, _ = _stock_expert_converters() + rewritten = onfly._compressedify(gate_up, _ScaledCompressor, None, ("weight", "weight_scale")) + key = "model.layers.0.mlp.experts.7." + fused = "model.layers.0.mlp.experts.gate_up_proj" + + assert rewritten.rename_source_key(key + "gate_proj.weight_scale") == ( + fused, + "mlp.experts.*.gate_proj.weight_scale$", + ) + assert rewritten.rename_source_key(key + "gate_proj.weight") == ( + fused, + "mlp.experts.*.gate_proj.weight$", + ) + assert rewritten.rename_source_key(key + "up_proj.weight_scale")[1] == ( + "mlp.experts.*.up_proj.weight_scale$" + ) + assert rewritten.rename_source_key(key + "down_proj.weight")[1] is None + + +def test_non_weight_sources_are_kept_verbatim(): + """Only ``.weight`` sources name a compressed tensor; any other source passes through as is.""" + conv = WeightConverter( + source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.gate_proj.bias"], + target_patterns="mlp.experts.gate_proj", + operations=[MergeModulelist(dim=0)], + ) + rewritten = onfly._compressedify(conv, _ScaledCompressor, None, PACKED) + + assert "mlp.experts.*.gate_proj.bias" in rewritten.source_patterns + assert sorted(p for p in rewritten.source_patterns if "weight" in p) == sorted( + f"mlp.experts.*.gate_proj.{name}$" for name in PACKED + ) + + +# ----------------------------------------------------------------- decompressing during the load + + +def test_packed_experts_decompress_then_fuse_like_an_uncompressed_checkpoint(): + """Through transformers' loader objects, per-expert packed weights and scales come out as the + fused bf16 ``gate_up_proj`` ([gate; up] per expert) and the stacked ``down_proj``.""" + n_experts, inter, hidden = 3, 2, 4 + shapes = {"gate": (inter, hidden), "up": (inter, hidden), "down": (hidden, inter)} + gen = torch.Generator().manual_seed(0) + dense, checkpoint = {}, {} + for e in range(n_experts): + for proj, shape in shapes.items(): + packed = torch.randint(-8, 8, shape, generator=gen, dtype=torch.int32) + scale = torch.rand(shape[0], 1, generator=gen) + 0.5 + dense[proj, e] = (packed.float() * scale).bfloat16() + key = f"model.layers.0.mlp.experts.{e}.{proj}_proj." + checkpoint |= { + key + "weight_packed": packed, + key + "weight_scale": scale, + key + "weight_shape": torch.tensor(shape), + } + converters = [ + onfly._compressedify(conv, _ScaledCompressor, None, PACKED) + for conv in _stock_expert_converters() + ] + + fused = _load(converters, checkpoint) + gate_up = fused.pop("model.layers.0.mlp.experts.gate_up_proj") + down = fused.pop("model.layers.0.mlp.experts.down_proj") + + assert not fused + assert gate_up.dtype == down.dtype == torch.bfloat16 + expected = [torch.cat([dense["gate", e], dense["up", e]]) for e in range(n_experts)] + assert torch.equal(gate_up, torch.stack(expected)) + assert torch.equal(down, torch.stack([dense["down", e] for e in range(n_experts)])) + + +def test_only_owned_experts_and_stored_siblings_reach_the_compressor(): + """Under EP the loader hands experts this rank does not own over as ``None``; they are skipped, + so a rank decompresses only its own. A checkpoint may store fewer params than the scheme + declares (no ``weight_global_scale`` in Kimi-K3's mxfp4); missing ones are not passed.""" + seen = [] + + class RecordingCompressor(_ScaledCompressor): + @classmethod + def decompress(cls, state_dict, scheme): + seen.append((sorted(state_dict), scheme)) + return super().decompress(state_dict, scheme) + + _, down = _stock_expert_converters() + names = ("weight_packed", "weight_scale", "weight_global_scale") + decompress = onfly._compressedify(down, RecordingCompressor, "scheme", names).operations[0] + packed = [torch.full((2, 2), e + 1, dtype=torch.int32) for e in range(3)] + scale = torch.full((2, 1), 0.5) + + out = decompress.convert( + { + "mlp.experts.*.down_proj.weight_packed$": [packed[0], None, packed[2]], + "mlp.experts.*.down_proj.weight_scale$": [scale, None, scale], + } + ) + + [(pattern, experts)] = out.items() + assert pattern == "mlp.experts.*.down_proj.weight_packed$" # the anchor, for the stock merge + assert [t.dtype for t in experts] == [torch.bfloat16] * 2 + assert [t[0, 0].item() for t in experts] == [0.5, 1.5] + assert seen == [(["weight_packed", "weight_scale"], "scheme")] * 2 + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_experts_merge_in_numeric_not_lexicographic_order(): + """Checkpoint keys iterated as strings yield experts 0, 1, 10, 11, 2, ...; the fused tensor must + still hold expert i in slot i (unfixed, expert 10 lands in slot 2).""" + onfly._install_numeric_expert_ordering() + _, down = _stock_expert_converters() + keys = {e: f"model.layers.0.mlp.experts.{e}.down_proj.weight" for e in range(12)} + checkpoint = {keys[e]: torch.full((2, 3), float(e)) for e in sorted(keys, key=str)} + assert list(checkpoint)[2] == keys[10] # the order the fix exists for + + fused = _load([down], checkpoint)["model.layers.0.mlp.experts.down_proj"] + + assert fused[:, 0, 0].tolist() == [float(e) for e in range(12)] + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_integer_checkpoint_tensors_keep_their_dtype_on_load(): + """Without a quantizer the loader casts every tensor to the model's float dtype, which destroys + packed weights, shapes and zero-points. Integers stay integers; floats are still cast, and + integer-to-integer casts still happen.""" + onfly._install_keep_int_native() + packed = torch.tensor([[2**31 - 15, -(2**31) + 3]], dtype=torch.int32) # not exact in bf16 + + kept = cml._materialize_copy(packed, dtype=torch.bfloat16) + + assert kept.dtype == torch.int32 + assert torch.equal(kept, packed) + assert cml._materialize_copy(torch.ones(2), dtype=torch.bfloat16).dtype == torch.bfloat16 + assert cml._materialize_copy(packed, dtype=torch.int64).dtype == torch.int64 + # the loader reaches the hook through spawn_materialize, which looks it up at call time + assert cml.spawn_materialize(None, packed, "cpu", torch.bfloat16)().dtype == torch.int32 + + +# ------------------------------------------------------------------------ installing the rewrite + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_install_rewrites_only_the_per_expert_weight_converters(monkeypatch): + """The loader's conversion list also holds renamings and other converters; only converters that + fuse per-expert ``.weight`` tensors are rewritten, and from_pretrained's arguments reach the + stock mapping unchanged.""" + renaming = WeightRenaming(r"experts\.(\d+)\.w1\.weight", r"experts.\1.gate_proj.weight") + qkv = WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + [Chunk(dim=0)], + ) + expert_bias = WeightConverter( + "mlp.experts.*.down_proj.bias", "mlp.experts.down_proj_bias", [MergeModulelist(dim=0)] + ) + gate_up, down = _stock_expert_converters() + stock = [renaming, qkv, expert_bias, gate_up, down] + calls = [] + + def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): + calls.append((model, key_mapping, hf_quantizer)) + return list(stock) + + compression = (_ScaledCompressor, None, PACKED, "stub-format") + monkeypatch.setattr(mu, "get_model_conversion_mapping", stock_mapping) + monkeypatch.setattr(onfly, "_read_compression", lambda cfg: compression) + onfly.install_onfly_decompress(SimpleNamespace()) + + # positionally, the way from_pretrained calls it + out = mu.get_model_conversion_mapping("model", {"old": "new"}, "quantizer") + + assert calls == [("model", {"old": "new"}, "quantizer")] + assert len(out) == len(stock) + assert all(kept is original for kept, original in zip(out[:3], stock[:3])) + assert [conv.source_patterns[0] for conv in out[3:]] == [ + "mlp.experts.*.gate_proj.weight_packed$", + "mlp.experts.*.down_proj.weight_packed$", + ] + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_installing_twice_wraps_the_loader_once(monkeypatch): + """``prepare`` may run more than once per process. A second install must not re-read the config + or wrap the hooks again: with a ``weight`` anchor the outer wrap would rewrite the rewritten + converter and decompress twice.""" + _, down = _stock_expert_converters() + reads = [] + + def read(cfg): + reads.append(cfg) + return _ScaledCompressor, None, ("weight", "weight_scale"), "stub-format" + + def stock_mapping(model, **kwargs): + return [copy.deepcopy(down)] + + monkeypatch.setattr(mu, "get_model_conversion_mapping", stock_mapping) + monkeypatch.setattr(onfly, "_read_compression", read) + + def hooks(): + return ( + mu.get_model_conversion_mapping, + WeightConverter.add_tensor, + WeightConverter.materialize_tensors, + cml._materialize_copy, + ) + + onfly.install_onfly_decompress(SimpleNamespace()) + installed = hooks() + onfly.install_onfly_decompress(SimpleNamespace()) + + assert hooks() == installed + assert len(reads) == 1 + [conv] = mu.get_model_conversion_mapping("model") + assert len(conv.operations) == len(down.operations) + 1 + + +# ------------------------------------------------------------------------ reading the checkpoint + + +@pytest.mark.parametrize( + ("config", "found"), + [ + ( + SimpleNamespace( + quantization_config={"at": "top"}, + text_config=SimpleNamespace(quantization_config={"at": "text"}), + ), + {"at": "top"}, + ), + ( + SimpleNamespace(text_config=SimpleNamespace(quantization_config={"at": "text"})), + {"at": "text"}, + ), + ({"text_config": {"quantization_config": {"at": "dict"}}}, {"at": "dict"}), + ( + SimpleNamespace(quantization_config=SimpleNamespace(to_dict=lambda: {"at": "obj"})), + {"at": "obj"}, + ), + ], + ids=["top-level-first", "text-config", "raw-dict", "config-object"], +) +def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_config(config, found): + """A VLM may carry its compressed-tensors config at the top level or on ``text_config``, as a + config object or a raw dict; it is always returned as a plain dict.""" + assert onfly._find_quantization_config(config) == found + + +def test_a_config_without_a_quantization_config_is_a_value_error(): + """Callers rely on exactly this error to mean "plain experts, nothing to decompress".""" + with pytest.raises(ValueError, match="quantization_config"): + onfly._find_quantization_config(SimpleNamespace(text_config=SimpleNamespace())) + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone(monkeypatch): + """By the time a bf16 checkpoint's caller sees that error, nothing may be patched yet -- and + compressed_tensors, which such a checkpoint never needs, need not be installed.""" + monkeypatch.setitem(sys.modules, "compressed_tensors", None) + hooks = ( + mu.get_model_conversion_mapping, + WeightConverter.add_tensor, + WeightConverter.materialize_tensors, + cml._materialize_copy, + ) + + with pytest.raises(ValueError, match="quantization_config"): + onfly.install_onfly_decompress(SimpleNamespace(text_config=SimpleNamespace())) + + assert ( + mu.get_model_conversion_mapping, + WeightConverter.add_tensor, + WeightConverter.materialize_tensors, + cml._materialize_copy, + ) == hooks + + +class _RemoteTextConfig(PretrainedConfig): + """Stands in for the remote-code DeepSeek-V3 config a VLM bundles as its ``text_config``.""" + + model_type = "remote_deepseek_v3" + + +@pytest.mark.parametrize("as_dict", [False, True], ids=["config-object", "raw-dict"]) +def test_the_native_backbone_config_is_a_clean_stock_deepseek_v3(as_dict): + """The bundled text config becomes transformers' own ``DeepseekV3Config`` with the same + architecture, minus ``auto_map`` (which routes back to the remote code) and the stale + ``quantization_config`` (which would build a competing quantizer). ``architectures`` names the + native class: hf_ptq's VLM detection iterates it.""" + fields = { + "hidden_size": 64, + "num_hidden_layers": 3, + "n_routed_experts": 8, + "moe_intermediate_size": 16, + "auto_map": {"AutoConfig": "configuration_deepseek.DeepseekV3Config"}, + "quantization_config": {"quant_method": "compressed-tensors"}, + } + if as_dict: + text = {**fields, "model_type": "remote_deepseek_v3"} + else: + text = _RemoteTextConfig(**fields) + + native = onfly.to_native_backbone_config(SimpleNamespace(text_config=text)) + + assert type(native) is DeepseekV3Config + assert native.model_type == "deepseek_v3" + assert native.architectures == ["DeepseekV3ForCausalLM"] + assert (native.hidden_size, native.num_hidden_layers) == (64, 3) + assert (native.n_routed_experts, native.moe_intermediate_size) == (8, 16) + assert getattr(native, "auto_map", None) is None + assert getattr(native, "quantization_config", None) is None + + +def test_the_native_backbone_config_needs_a_text_config_and_a_known_model_type(): + with pytest.raises(ValueError, match="text_config"): + onfly.to_native_backbone_config(SimpleNamespace()) + with pytest.raises(ValueError, match="unknown native model_type"): + onfly.to_native_backbone_config(SimpleNamespace(text_config={}), model_type="not_a_model") + + +# -------------------------------------------------------- with the real compressed_tensors library + + +def _int4(symmetric, targets): + """A compressed-tensors INT4, group-32 config group.""" + weights = {"num_bits": 4, "type": "int", "symmetric": symmetric, "strategy": "group"} + return {"targets": targets, "weights": {**weights, "group_size": 32}} + + +def _pack_quantized(**groups): + """A compressed-tensors ``quantization_config`` in the ``pack-quantized`` format.""" + return { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "config_groups": groups, + } + + +ATTN = ["re:.*self_attn.*"] +EXPERTS = ["re:.*experts.*"] + + +@pytest.mark.parametrize( + ("groups", "zero_point"), + [ + ({"attn": _int4(True, ATTN), "moe": _int4(False, EXPERTS)}, True), + ({"moe": _int4(True, EXPERTS), "attn": _int4(False, ATTN)}, False), + ({"attn": _int4(False, ATTN), "dense": _int4(True, ["Linear"])}, True), + ], + ids=["expert-group-second", "expert-group-first", "no-expert-group-takes-the-first"], +) +def test_the_expert_config_group_decides_the_stored_params(groups, zero_point): + """With several config groups the one targeting experts wins, else the first; an asymmetric + scheme stores a zero-point next to the packed weight, a symmetric one does not.""" + pytest.importorskip("compressed_tensors") + config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) + + _, scheme, names, fmt = onfly._read_compression(config) + + assert fmt == "pack-quantized" + assert names[: len(PACKED)] == PACKED + assert ("weight_zero_point" in names) is zero_point + assert scheme.weights.symmetric is not zero_point + + +def test_real_pack_quantized_experts_round_trip_exactly(): + """With the real compressed_tensors INT4 compressor, packed experts come back out of the fused + converter as the original weights (chosen to be exact in INT4 x 0.5 and in bf16).""" + pytest.importorskip("compressed_tensors") + config = SimpleNamespace(quantization_config=_pack_quantized(experts=_int4(True, ["Linear"]))) + compressor, scheme, names, _ = onfly._read_compression(config) + n_experts, inter, hidden = 3, 4, 64 + gen = torch.Generator().manual_seed(0) + dense, checkpoint = {}, {} + for e in range(n_experts): + for proj in ("gate", "up"): + dense[proj, e] = torch.randint(-8, 8, (inter, hidden), generator=gen).float() * 0.5 + scale = torch.full((inter, hidden // 32), 0.5) + stored = compressor.compress({"weight": dense[proj, e], "weight_scale": scale}, scheme) + key = f"model.layers.0.mlp.experts.{e}.{proj}_proj." + checkpoint |= {key + name: stored[name] for name in names} + gate_up, _ = _stock_expert_converters() + + fused = _load([onfly._compressedify(gate_up, compressor, scheme, names)], checkpoint) + + expected = [torch.cat([dense["gate", e], dense["up", e]]) for e in range(n_experts)] + assert torch.equal( + fused["model.layers.0.mlp.experts.gate_up_proj"], torch.stack(expected).bfloat16() + ) From 508052b5fe768f974ae3f4fc3a53ccff45d4e12e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:25:02 +0000 Subject: [PATCH 02/16] Skip the tests below the transformers release the code requires The module skipped below 5.6, but the code under test requires 5.8 and raises below it, so a 5.6 or 5.7 run failed instead of skipping. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_onfly_decompress.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 45a71568d58..523bff872e3 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -33,8 +33,8 @@ from packaging.version import Version -if Version(transformers.__version__) < Version("5.6"): - pytest.skip("needs the transformers >= 5.6 weight-loading API", allow_module_level=True) +if Version(transformers.__version__) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers >= 5.8", allow_module_level=True) import transformers.core_model_loading as cml import transformers.modeling_utils as mu From 3119a49c9d51c496cdacd6ce852643615867a619 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:46:32 +0000 Subject: [PATCH 03/16] Refuse a second, differently compressed checkpoint in one process The decompress hooks are process-wide and keep the first checkpoint's compressor, scheme and stored params, so a second checkpoint compressed differently was decompressed with the first one's settings without a word. Installing again now compares the compression: the same one is a no-op, as is a config its installing prepare already stripped of quantization_config, and a different one raises. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 29 ++++++++++++--- .../torch/models/test_onfly_decompress.py | 37 +++++++++++++++++-- 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 14d5ffac047..eba859086b0 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -317,18 +317,35 @@ def _materialize_copy(tensor, device=None, dtype=None): def install_onfly_decompress(hf_config): """Monkeypatch get_model_conversion_mapping so the fused-MoE per-expert converters decompress the checkpoint's compressed-tensors weights on the fly, in whatever ``format`` the checkpoint declares. - Idempotent. Call once before from_pretrained of a compressed checkpoint. ``hf_config`` is the model - config carrying the compressed-tensors ``quantization_config`` (top-level or under ``text_config``); - the format, bit-width, group size, strategy and stored param names are all read from it.""" + Call before from_pretrained of a compressed checkpoint. ``hf_config`` is the model config carrying + the compressed-tensors ``quantization_config`` (top-level or under ``text_config``); the format, + bit-width, group size, strategy and stored param names are all read from it. + + The hooks are process-wide, so a process decompresses one compressed checkpoint: installing again + for it is a no-op, and installing for a checkpoint compressed differently raises.""" import transformers.modeling_utils as mu from transformers.core_model_loading import WeightConverter from modelopt.torch.utils.moe import require_transformers - if getattr(mu.get_model_conversion_mapping, "_onfly_patched", False): + installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) + if installed is not None: + try: + compression = _read_compression(hf_config) + except ValueError: + # The prepare that installed the hooks drops quantization_config from its config and may + # run again on that same config: nothing new to decompress. + return + if compression != installed: + raise RuntimeError( + f"install_onfly_decompress: this process already decompresses a '{installed[3]}' " + f"checkpoint (stored params {installed[2]}); one compressed differently ('{compression[3]}', " + f"stored params {compression[2]}) needs a process of its own" + ) return # Raises ValueError for a config with nothing to decompress, before anything is patched. - compressor, scheme, param_names, fmt = _read_compression(hf_config) + compression = _read_compression(hf_config) + compressor, scheme, param_names, fmt = compression require_transformers("5.8", "on-the-fly compressed-tensors decompression") _install_numeric_expert_ordering() _install_keep_int_native() @@ -347,7 +364,7 @@ def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): out.append(c) return out - setattr(patched, "_onfly_patched", True) + setattr(patched, "_onfly_compression", compression) mu.get_model_conversion_mapping = patched print( f"[onfly_decompress] installed: compressed_tensors '{fmt}' fused-MoE experts decompress " diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 523bff872e3..195b5b1532d 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -309,8 +309,8 @@ def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): @pytest.mark.usefixtures("restore_loader_hooks") def test_installing_twice_wraps_the_loader_once(monkeypatch): - """``prepare`` may run more than once per process. A second install must not re-read the config - or wrap the hooks again: with a ``weight`` anchor the outer wrap would rewrite the rewritten + """``prepare`` may run more than once per process. Installing again for the same checkpoint must + not wrap the hooks again: with a ``weight`` anchor the outer wrap would rewrite the rewritten converter and decompress twice.""" _, down = _stock_expert_converters() reads = [] @@ -338,11 +338,42 @@ def hooks(): onfly.install_onfly_decompress(SimpleNamespace()) assert hooks() == installed - assert len(reads) == 1 + assert len(reads) == 2 [conv] = mu.get_model_conversion_mapping("model") assert len(conv.operations) == len(down.operations) + 1 +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_checkpoint_compressed_differently_needs_a_process_of_its_own(monkeypatch): + """The hooks are process-wide: decompressing a second checkpoint with the first one's scheme and + stored params would load it wrong without a word, so its install is refused and the hooks stay as + they are. A config the installing prepare already stripped of its quantization_config is the + first checkpoint again.""" + _, down = _stock_expert_converters() + compressions = { + "a": (_ScaledCompressor, None, PACKED, "stub-format"), + "b": (_ScaledCompressor, None, (*PACKED, "weight_zero_point"), "stub-format"), + } + + def read(cfg): + if cfg.checkpoint is None: + raise ValueError("install_onfly_decompress: config has no quantization_config") + return compressions[cfg.checkpoint] + + monkeypatch.setattr( + mu, "get_model_conversion_mapping", lambda model, **kw: [copy.deepcopy(down)] + ) + monkeypatch.setattr(onfly, "_read_compression", read) + + onfly.install_onfly_decompress(SimpleNamespace(checkpoint="a")) + installed = mu.get_model_conversion_mapping + onfly.install_onfly_decompress(SimpleNamespace(checkpoint=None)) + with pytest.raises(RuntimeError, match="needs a process of its own"): + onfly.install_onfly_decompress(SimpleNamespace(checkpoint="b")) + + assert mu.get_model_conversion_mapping is installed + + # ------------------------------------------------------------------------ reading the checkpoint From 92169def3eb750d4e1cae242b6c258425adc10e7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 23:27:38 +0000 Subject: [PATCH 04/16] Rewrite only the compressed checkpoint's own expert converters The installed decompressor rewrote the per-expert converters of every model from_pretrained loaded afterwards. An uncompressed MoE in the same process had its .weight sources renamed to the first checkpoint's packed names, matched no stored key, and loaded its experts newly initialized without an error. Installing now marks the config it is given, a repeat install for the same checkpoint marks its config too, and only a model built from a marked config gets the rewritten converters. from_pretrained loads from a deep copy of the config, which keeps the mark; the rewrite removes it from that copy so a later save does not write it into config.json. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 34 ++++++++- .../torch/models/test_onfly_decompress.py | 73 +++++++++++++++++-- 2 files changed, 98 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index eba859086b0..1cca6a308fd 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -44,7 +44,8 @@ each ``.weight`` source becomes the checkpoint's actual stored keys for that weight (``compression_param_names``), and PREPEND a ``DecompressExperts`` op that pairs those per expert and dequantizes -> bf16 BEFORE the merge/concat collapse the per-expert structure. -``install_onfly_decompress()`` monkeypatches ``get_model_conversion_mapping``. +``install_onfly_decompress()`` monkeypatches ``get_model_conversion_mapping``, for the model built from +the config it is given: every other model ``from_pretrained`` loads keeps its stock converters. """ import torch @@ -314,6 +315,18 @@ def _materialize_copy(tensor, device=None, dtype=None): cml._materialize_copy = _materialize_copy +# Marks a config whose model's converters the installed hooks rewrite (holds the checkpoint's format). +_MARK = "_modelopt_onfly_decompress" + + +def _mark(hf_config, fmt): + """Mark ``hf_config``'s load for the installed hooks; ``from_pretrained``'s copy of it keeps the mark.""" + if isinstance(hf_config, dict): + hf_config[_MARK] = fmt + else: + setattr(hf_config, _MARK, fmt) + + def install_onfly_decompress(hf_config): """Monkeypatch get_model_conversion_mapping so the fused-MoE per-expert converters decompress the checkpoint's compressed-tensors weights on the fly, in whatever ``format`` the checkpoint declares. @@ -322,7 +335,9 @@ def install_onfly_decompress(hf_config): bit-width, group size, strategy and stored param names are all read from it. The hooks are process-wide, so a process decompresses one compressed checkpoint: installing again - for it is a no-op, and installing for a checkpoint compressed differently raises.""" + for it marks that config too, and installing for a checkpoint compressed differently raises. Only a + model built from a marked config has its converters rewritten; any other model ``from_pretrained`` + loads in the process, an uncompressed MoE say, keeps its stock converters.""" import transformers.modeling_utils as mu from transformers.core_model_loading import WeightConverter @@ -333,8 +348,8 @@ def install_onfly_decompress(hf_config): try: compression = _read_compression(hf_config) except ValueError: - # The prepare that installed the hooks drops quantization_config from its config and may - # run again on that same config: nothing new to decompress. + # The prepare that installed the hooks drops quantization_config from its config, already + # marked, and may run again on that same config: nothing new to decompress. return if compression != installed: raise RuntimeError( @@ -342,6 +357,8 @@ def install_onfly_decompress(hf_config): f"checkpoint (stored params {installed[2]}); one compressed differently ('{compression[3]}', " f"stored params {compression[2]}) needs a process of its own" ) + # The same checkpoint through a config of its own, which its load must find marked too. + _mark(hf_config, installed[3]) return # Raises ValueError for a config with nothing to decompress, before anything is patched. compression = _read_compression(hf_config) @@ -353,6 +370,14 @@ def install_onfly_decompress(hf_config): def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): convs = orig(model, key_mapping=key_mapping, hf_quantizer=hf_quantizer, *a, **k) + config = getattr(model, "config", None) + if getattr(config, _MARK, None) is None: + # Not the compressed checkpoint: rewritten, an uncompressed MoE's experts would match no + # stored key and load newly initialized without an error. + return convs + # from_pretrained loads from a copy of the marked config; the mark must not reach the config + # a later save writes. + delattr(config, _MARK) out = [] for c in convs: # only the per-expert MoE converters (have `.weight` sources under `experts`) @@ -366,6 +391,7 @@ def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): setattr(patched, "_onfly_compression", compression) mu.get_model_conversion_mapping = patched + _mark(hf_config, fmt) print( f"[onfly_decompress] installed: compressed_tensors '{fmt}' fused-MoE experts decompress " f"(stored params={param_names}) -> bf16 during from_pretrained", diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 195b5b1532d..488028fbb88 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -293,12 +293,14 @@ def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): compression = (_ScaledCompressor, None, PACKED, "stub-format") monkeypatch.setattr(mu, "get_model_conversion_mapping", stock_mapping) monkeypatch.setattr(onfly, "_read_compression", lambda cfg: compression) - onfly.install_onfly_decompress(SimpleNamespace()) + config = SimpleNamespace() + onfly.install_onfly_decompress(config) + model = SimpleNamespace(config=config) # positionally, the way from_pretrained calls it - out = mu.get_model_conversion_mapping("model", {"old": "new"}, "quantizer") + out = mu.get_model_conversion_mapping(model, {"old": "new"}, "quantizer") - assert calls == [("model", {"old": "new"}, "quantizer")] + assert calls == [(model, {"old": "new"}, "quantizer")] assert len(out) == len(stock) assert all(kept is original for kept, original in zip(out[:3], stock[:3])) assert [conv.source_patterns[0] for conv in out[3:]] == [ @@ -335,11 +337,12 @@ def hooks(): onfly.install_onfly_decompress(SimpleNamespace()) installed = hooks() - onfly.install_onfly_decompress(SimpleNamespace()) + again = SimpleNamespace() + onfly.install_onfly_decompress(again) assert hooks() == installed assert len(reads) == 2 - [conv] = mu.get_model_conversion_mapping("model") + [conv] = mu.get_model_conversion_mapping(SimpleNamespace(config=again)) assert len(conv.operations) == len(down.operations) + 1 @@ -374,6 +377,66 @@ def read(cfg): assert mu.get_model_conversion_mapping is installed +def _install_for_deepseek_v3(monkeypatch): + """Install for a compressed DeepSeek-V3 config; a config without quantization_config reads as none.""" + compression = (_ScaledCompressor, None, PACKED, "stub-format") + + def read(cfg): + if getattr(cfg, "quantization_config", None) is None: + raise ValueError("install_onfly_decompress: config has no quantization_config") + return compression + + monkeypatch.setattr( + mu, "get_model_conversion_mapping", lambda model, **kw: list(_stock_expert_converters()) + ) + monkeypatch.setattr(onfly, "_read_compression", read) + config = DeepseekV3Config(num_hidden_layers=1) + config.quantization_config = {"format": "stub-format"} + onfly.install_onfly_decompress(config) + return config + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_another_model_loaded_in_the_process_keeps_its_stock_expert_converters(monkeypatch): + """The hooks outlive the compressed checkpoint's load. Rewritten, an uncompressed MoE's experts + would match no stored key, load newly initialized, and from_pretrained would not complain.""" + _install_for_deepseek_v3(monkeypatch) + other = SimpleNamespace(config=DeepseekV3Config(num_hidden_layers=1)) + + out = mu.get_model_conversion_mapping(other) + + stock = [conv.source_patterns for conv in _stock_expert_converters()] + assert [conv.source_patterns for conv in out] == stock + assert all(p.endswith(".weight") for patterns in stock for p in patterns) + + +@pytest.mark.usefixtures("restore_loader_hooks") +@pytest.mark.parametrize("repeat", ["stripped", "fresh"]) +def test_the_compressed_checkpoint_is_rewritten_through_the_copy_from_pretrained_loads( + monkeypatch, repeat +): + """from_pretrained builds the model from a deep copy of the config ``prepare`` handed over. That + copy must still select the rewrite -- whether ``prepare`` ran again on the config it stripped of + its quantization_config (Kimi-K3) or on a fresh config of the same checkpoint -- and the mark + must not reach the config a later save writes.""" + config = _install_for_deepseek_v3(monkeypatch) + if repeat == "stripped": + config.quantization_config = None + else: + config = DeepseekV3Config(num_hidden_layers=1) + config.quantization_config = {"format": "stub-format"} + onfly.install_onfly_decompress(config) + model = SimpleNamespace(config=copy.deepcopy(config)) + + out = mu.get_model_conversion_mapping(model) + + assert [conv.source_patterns[0] for conv in out] == [ + "mlp.experts.*.gate_proj.weight_packed$", + "mlp.experts.*.down_proj.weight_packed$", + ] + assert onfly._MARK not in model.config.to_dict() + + # ------------------------------------------------------------------------ reading the checkpoint From 25ba0d27a80d06d6b51b236da4985ff669537964 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:04:03 +0000 Subject: [PATCH 05/16] Say what the real compressed-tensors tests need to run They drive the real compressed-tensors registry and compressor, so they skip where the library is missing, as it is in a plain dev-test install. CI's unit session installs it; the skip reason now says so instead of leaving only an import error. With compressed-tensors 0.19.0 all four cases pass. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_onfly_decompress.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 488028fbb88..2104883cc06 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -574,7 +574,9 @@ def _pack_quantized(**groups): def test_the_expert_config_group_decides_the_stored_params(groups, zero_point): """With several config groups the one targeting experts wins, else the first; an asymmetric scheme stores a zero-point next to the packed weight, a symmetric one does not.""" - pytest.importorskip("compressed_tensors") + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) _, scheme, names, fmt = onfly._read_compression(config) @@ -588,7 +590,9 @@ def test_the_expert_config_group_decides_the_stored_params(groups, zero_point): def test_real_pack_quantized_experts_round_trip_exactly(): """With the real compressed_tensors INT4 compressor, packed experts come back out of the fused converter as the original weights (chosen to be exact in INT4 x 0.5 and in bf16).""" - pytest.importorskip("compressed_tensors") + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) config = SimpleNamespace(quantization_config=_pack_quantized(experts=_int4(True, ["Linear"]))) compressor, scheme, names, _ = onfly._read_compression(config) n_experts, inter, hidden = 3, 4, 64 From 3307733159b091a2b1e102793436ca47a231fafd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:14:34 +0000 Subject: [PATCH 06/16] Install compressed-tensors in the torch>=2.10 unit sessions The real-library tests of on-the-fly expert decompression skip without compressed-tensors, and no CI environment installs it, so CI never ran them. The unit sessions now install it with the leg's torchvision pin, so the resolver keeps the leg's torch. compressed-tensors 0.19 needs torch>=2.10, so the torch 2.8 and 2.9 legs leave it out: there it would upgrade torch or, held back by the pin, fall back to an old 0.15 release. It stays out of the dev-test extra: the transformers-free partial_unit(torch) lane installs that extra, and compressed-tensors pulls in transformers. Signed-off-by: Shengliang Xu --- noxfile.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index 96232a2d381..3b333aa0112 100644 --- a/noxfile.py +++ b/noxfile.py @@ -68,7 +68,12 @@ def _cov_args(): @nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS]) def unit(session, torch_ver, tf_ver): """Unit tests — parametrized over torch and transformers versions.""" - session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]") + # compressed-tensors runs the real-library tests of on-the-fly expert decompression. It needs + # torch>=2.10, so the older legs skip those tests; it goes in with the torchvision pin so the + # resolver keeps the leg's torch. Not in dev-test: the transformers-free partial_unit(torch) + # lane installs that extra, and compressed-tensors pulls in transformers. + extra = [] if torch_ver in ("torch_28", "torch_29") else ["compressed-tensors"] + session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", *extra) tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) From a2ae1dd9669e6109707a13fbf579769247fc6d6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:44:05 +0000 Subject: [PATCH 07/16] Trim the onfly decompress tests: fold pattern checks into the load The two tests that restated _compressedify's source patterns now ride the packed-load test, parametrized over the stored-param sets (pack-quantized; weight and weight_scale): a swapped gate/up, an unanchored pattern and a decompress op that is not a MergeModulelist each fail it. The other-model test and the copy-from_pretrained test become one test over the three loads. The helper-level ValueError test goes: the install-level test asserts the same error and that nothing was patched. The non-weight-sources and native-backbone guard tests stay, since nothing else catches a regression in their branches. Signed-off-by: Shengliang Xu --- .../torch/models/test_onfly_decompress.py | 130 ++++++------------ 1 file changed, 40 insertions(+), 90 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 2104883cc06..7614cd13774 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -61,11 +61,13 @@ def restore_loader_hooks(monkeypatch): class _ScaledCompressor: - """Stands in for a compressed_tensors compressor: ``weight = weight_packed * weight_scale``.""" + """Stands in for a compressed_tensors compressor: ``weight = stored weight * weight_scale``, + the stored weight being ``weight_packed`` (pack-quantized) or ``weight`` (float formats).""" @classmethod def decompress(cls, state_dict, scheme): - return {"weight": state_dict["weight_packed"].float() * state_dict["weight_scale"]} + stored = state_dict.get("weight_packed", state_dict.get("weight")) + return {"weight": stored.float() * state_dict["weight_scale"]} def _stock_expert_converters(): @@ -104,50 +106,6 @@ def _load(converters, checkpoint): # ------------------------------------------------------------- rewriting the per-expert converters -def test_weight_sources_become_the_stored_keys_anchor_first_and_gate_before_up(): - """Each ``.weight`` source becomes the checkpoint's stored keys, grouped anchor first with gate - ahead of up in every group (``Concatenate`` stacks in source order, so rows stay [gate; up]). - Targets are kept, and the decompress op runs first, as a ``MergeModulelist`` so that loaders - which only shard experts when the first op is one still do.""" - gate_up, _ = _stock_expert_converters() - rewritten = onfly._compressedify(gate_up, _ScaledCompressor, None, PACKED) - - assert rewritten.source_patterns == [ - "mlp.experts.*.gate_proj.weight_packed$", - "mlp.experts.*.up_proj.weight_packed$", - "mlp.experts.*.gate_proj.weight_scale$", - "mlp.experts.*.up_proj.weight_scale$", - "mlp.experts.*.gate_proj.weight_shape$", - "mlp.experts.*.up_proj.weight_shape$", - ] - assert rewritten.target_patterns == ["mlp.experts.gate_up_proj"] - decompress, *stock_ops = rewritten.operations - assert isinstance(decompress, MergeModulelist) - assert [repr(op) for op in stock_ops] == [repr(op) for op in gate_up.operations] - - -def test_stored_keys_are_anchored_so_a_weight_never_claims_its_scale(): - """Float formats store ``weight`` next to ``weight_scale``, and transformers regex-*searches* - source patterns: an unanchored ``gate_proj.weight`` would also claim ``gate_proj.weight_scale``.""" - gate_up, _ = _stock_expert_converters() - rewritten = onfly._compressedify(gate_up, _ScaledCompressor, None, ("weight", "weight_scale")) - key = "model.layers.0.mlp.experts.7." - fused = "model.layers.0.mlp.experts.gate_up_proj" - - assert rewritten.rename_source_key(key + "gate_proj.weight_scale") == ( - fused, - "mlp.experts.*.gate_proj.weight_scale$", - ) - assert rewritten.rename_source_key(key + "gate_proj.weight") == ( - fused, - "mlp.experts.*.gate_proj.weight$", - ) - assert rewritten.rename_source_key(key + "up_proj.weight_scale")[1] == ( - "mlp.experts.*.up_proj.weight_scale$" - ) - assert rewritten.rename_source_key(key + "down_proj.weight")[1] is None - - def test_non_weight_sources_are_kept_verbatim(): """Only ``.weight`` sources name a compressed tensor; any other source passes through as is.""" conv = WeightConverter( @@ -166,9 +124,18 @@ def test_non_weight_sources_are_kept_verbatim(): # ----------------------------------------------------------------- decompressing during the load -def test_packed_experts_decompress_then_fuse_like_an_uncompressed_checkpoint(): - """Through transformers' loader objects, per-expert packed weights and scales come out as the - fused bf16 ``gate_up_proj`` ([gate; up] per expert) and the stacked ``down_proj``.""" +@pytest.mark.parametrize( + "stored", [PACKED, ("weight", "weight_scale")], ids=["pack-quantized", "weight-and-scale"] +) +def test_packed_experts_decompress_then_fuse_like_an_uncompressed_checkpoint(stored): + """Through transformers' loader objects, per-expert compressed weights and scales come out as + the fused bf16 ``gate_up_proj`` ([gate; up] per expert) and the stacked ``down_proj``. + + The decompress op runs first, as a ``MergeModulelist``: loaders shard experts at load only when + the first op is one. A float format stores ``weight`` next to ``weight_scale``, and transformers + regex-*searches* source patterns, so only anchored patterns keep a weight from claiming its + scale. + """ n_experts, inter, hidden = 3, 2, 4 shapes = {"gate": (inter, hidden), "up": (inter, hidden), "down": (hidden, inter)} gen = torch.Generator().manual_seed(0) @@ -179,15 +146,14 @@ def test_packed_experts_decompress_then_fuse_like_an_uncompressed_checkpoint(): scale = torch.rand(shape[0], 1, generator=gen) + 0.5 dense[proj, e] = (packed.float() * scale).bfloat16() key = f"model.layers.0.mlp.experts.{e}.{proj}_proj." - checkpoint |= { - key + "weight_packed": packed, - key + "weight_scale": scale, - key + "weight_shape": torch.tensor(shape), - } + values = {"weight_packed": packed, "weight": packed, "weight_scale": scale} + values["weight_shape"] = torch.tensor(shape) + checkpoint |= {key + name: values[name] for name in stored} converters = [ - onfly._compressedify(conv, _ScaledCompressor, None, PACKED) + onfly._compressedify(conv, _ScaledCompressor, None, stored) for conv in _stock_expert_converters() ] + assert all(isinstance(conv.operations[0], MergeModulelist) for conv in converters) fused = _load(converters, checkpoint) gate_up = fused.pop("model.layers.0.mlp.experts.gate_up_proj") @@ -397,42 +363,32 @@ def read(cfg): @pytest.mark.usefixtures("restore_loader_hooks") -def test_another_model_loaded_in_the_process_keeps_its_stock_expert_converters(monkeypatch): - """The hooks outlive the compressed checkpoint's load. Rewritten, an uncompressed MoE's experts - would match no stored key, load newly initialized, and from_pretrained would not complain.""" - _install_for_deepseek_v3(monkeypatch) - other = SimpleNamespace(config=DeepseekV3Config(num_hidden_layers=1)) - - out = mu.get_model_conversion_mapping(other) - - stock = [conv.source_patterns for conv in _stock_expert_converters()] - assert [conv.source_patterns for conv in out] == stock - assert all(p.endswith(".weight") for patterns in stock for p in patterns) - - -@pytest.mark.usefixtures("restore_loader_hooks") -@pytest.mark.parametrize("repeat", ["stripped", "fresh"]) -def test_the_compressed_checkpoint_is_rewritten_through_the_copy_from_pretrained_loads( - monkeypatch, repeat -): - """from_pretrained builds the model from a deep copy of the config ``prepare`` handed over. That - copy must still select the rewrite -- whether ``prepare`` ran again on the config it stripped of - its quantization_config (Kimi-K3) or on a fresh config of the same checkpoint -- and the mark - must not reach the config a later save writes.""" +@pytest.mark.parametrize("loaded", ["another-model", "stripped-copy", "fresh-config"]) +def test_only_the_compressed_checkpoint_loads_through_the_rewritten_converters(monkeypatch, loaded): + """The hooks outlive the compressed checkpoint's load. Rewritten, another (uncompressed) MoE's + experts would match no stored key, load newly initialized, and from_pretrained would not + complain. from_pretrained builds the model from a deep copy of the config ``prepare`` handed + over; that copy must still select the rewrite, whether ``prepare`` ran again on the config it + stripped of its quantization_config (Kimi-K3) or on a fresh config of the same checkpoint, and + the mark must not reach the config a later save writes.""" config = _install_for_deepseek_v3(monkeypatch) - if repeat == "stripped": - config.quantization_config = None - else: + if loaded == "another-model": config = DeepseekV3Config(num_hidden_layers=1) - config.quantization_config = {"format": "stub-format"} - onfly.install_onfly_decompress(config) + else: + if loaded == "stripped-copy": + config.quantization_config = None + else: + config = DeepseekV3Config(num_hidden_layers=1) + config.quantization_config = {"format": "stub-format"} + onfly.install_onfly_decompress(config) model = SimpleNamespace(config=copy.deepcopy(config)) out = mu.get_model_conversion_mapping(model) + stored = "weight" if loaded == "another-model" else "weight_packed$" assert [conv.source_patterns[0] for conv in out] == [ - "mlp.experts.*.gate_proj.weight_packed$", - "mlp.experts.*.down_proj.weight_packed$", + f"mlp.experts.*.gate_proj.{stored}", + f"mlp.experts.*.down_proj.{stored}", ] assert onfly._MARK not in model.config.to_dict() @@ -468,12 +424,6 @@ def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_conf assert onfly._find_quantization_config(config) == found -def test_a_config_without_a_quantization_config_is_a_value_error(): - """Callers rely on exactly this error to mean "plain experts, nothing to decompress".""" - with pytest.raises(ValueError, match="quantization_config"): - onfly._find_quantization_config(SimpleNamespace(text_config=SimpleNamespace())) - - @pytest.mark.usefixtures("restore_loader_hooks") def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone(monkeypatch): """By the time a bf16 checkpoint's caller sees that error, nothing may be patched yet -- and From f9c1074858dcb3c6cab7ea1b6717f98fe7d2750b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 01:11:25 +0000 Subject: [PATCH 08/16] Let a compressed checkpoint a model spec decompresses pass the format check A model spec that installs the on-the-fly decompressor marks its config. Kimi-K2.5's keeps its compressed-tensors quantization_config at the top level, where the decompressor reads it, so the expert-parallel load's format check would refuse a checkpoint it does decompress. A marked config now passes; an unmarked compressed-tensors checkpoint is still refused. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 6 +++++ .../torch/utils/plugins/model_load_utils.py | 9 ++++++- tests/unit/torch/utils/test_ep_loader.py | 26 +++++++++++++++++++ 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 1cca6a308fd..76f03c9e6a9 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -327,6 +327,12 @@ def _mark(hf_config, fmt): setattr(hf_config, _MARK, fmt) +def is_marked(hf_config) -> bool: + """Whether :func:`install_onfly_decompress` marked ``hf_config`` (or its ``text_config``) for its hooks.""" + configs = (hf_config, getattr(hf_config, "text_config", None)) + return any(getattr(c, _MARK, None) is not None for c in configs if c is not None) + + def install_onfly_decompress(hf_config): """Monkeypatch get_model_conversion_mapping so the fused-MoE per-expert converters decompress the checkpoint's compressed-tensors weights on the fly, in whatever ``format`` the checkpoint declares. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 561a5dfb323..673f930f09f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2200,6 +2200,13 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) +def _decompressed_as_loaded(cfg) -> bool: + """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint.""" + from modelopt.torch.models import _onfly_decompress + + return _onfly_decompress.is_marked(cfg) + + def _ckpt_quant_method(cfg, ckpt_path: str | None = None): """Return the checkpoint's native quantization method, or ``None``. @@ -2513,7 +2520,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "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: + if quant_method not in _EP_SOURCE_FORMATS and not _decompressed_as_loaded(hf_config): # 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)) raise NotImplementedError( diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 1cb7a0dff2f..18809cf68ed 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -559,6 +559,32 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): assert classes == [True] +def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes(monkeypatch): + """A model spec that installs the on-the-fly decompressor (Kimi-K2.5's) marks the config; that + checkpoint passes the format check, and the same checkpoint unmarked does not.""" + from modelopt.torch.models import _onfly_decompress + + class _ResolvedPastFormatError(Exception): + """Stops the load where it resolves the checkpoint, past the format check.""" + + def resolve(path, rank): + raise _ResolvedPastFormatError + + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) + monkeypatch.setattr(moe, "require_transformers", lambda *args: None) + monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) + config = SimpleNamespace(quantization_config={"quant_method": "compressed-tensors"}) + plan = DistributionPlan(2, ep_size=2) + + with pytest.raises(NotImplementedError, match="cannot dequantize"): + load_and_prepare_ep("org/model", plan, hf_config=config) + _onfly_decompress._mark(config, "pack-quantized") + with pytest.raises(_ResolvedPastFormatError): + load_and_prepare_ep("org/model", plan, hf_config=config) + + # ------------------------------------------------------------------------ quant wrappers From 7bb2898e4b49fccf3663109280c74c0066cb1fde Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:10:13 +0000 Subject: [PATCH 09/16] Tell a compressed checkpoint's config that fails to parse from none The decompressor read "no quantization_config" off a ValueError, but compressed_tensors raises pydantic's ValidationError, also a ValueError, for a quantization_config it cannot parse (say, one written by a newer compressed-tensors). Taken for none, the second install of a process returned without marking the config, and the checkpoint's experts would load newly initialized. Reading a config with none now returns None, so a parse error goes through. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 27 ++++++++++--------- .../torch/models/test_onfly_decompress.py | 15 ++++++----- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 76f03c9e6a9..19afe608bb7 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -86,7 +86,7 @@ def to_native_backbone_config( def _find_quantization_config(hf_config): """Locate the compressed-tensors ``quantization_config`` (top-level or under ``text_config``) and - return it as a plain dict.""" + return it as a plain dict, or ``None`` for a config without one.""" qc = getattr(hf_config, "quantization_config", None) if qc is None: tc = getattr(hf_config, "text_config", None) @@ -96,9 +96,7 @@ def _find_quantization_config(hf_config): "quantization_config" ) if qc is None: - raise ValueError( - "install_onfly_decompress: config has no compressed-tensors quantization_config" - ) + return None if hasattr(qc, "to_dict"): qc = qc.to_dict() elif not isinstance(qc, dict): @@ -113,10 +111,13 @@ def _read_compression(hf_config): ``compressor`` is the registered compressor for ``format`` (its ``decompress`` classmethod is the dequant); ``scheme`` is the ``QuantizationScheme`` (supplies ``scheme.weights`` to ``decompress``); ``param_names`` is the exact set of stored sibling keys per weight. The experts' group is chosen as - the sole config group, else the one whose ``targets`` reference "expert", else the first.""" - # Look the config up first: with none, there is nothing to decompress (ValueError), and - # compressed_tensors need not even be installed. + the sole config group, else the one whose ``targets`` reference "expert", else the first. + ``None`` for a config with no ``quantization_config``: nothing to decompress.""" + # Look the config up first: with none, compressed_tensors need not even be installed. A config + # that fails to parse raises pydantic's ValidationError, a ValueError of its own. quantization_config = _find_quantization_config(hf_config) + if quantization_config is None: + return None from compressed_tensors.compressors.base import BaseCompressor from compressed_tensors.quantization import QuantizationConfig @@ -350,10 +351,9 @@ def install_onfly_decompress(hf_config): from modelopt.torch.utils.moe import require_transformers installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) + compression = _read_compression(hf_config) if installed is not None: - try: - compression = _read_compression(hf_config) - except ValueError: + if compression is None: # The prepare that installed the hooks drops quantization_config from its config, already # marked, and may run again on that same config: nothing new to decompress. return @@ -366,8 +366,11 @@ def install_onfly_decompress(hf_config): # The same checkpoint through a config of its own, which its load must find marked too. _mark(hf_config, installed[3]) return - # Raises ValueError for a config with nothing to decompress, before anything is patched. - compression = _read_compression(hf_config) + if compression is None: + # Nothing to decompress: refused before anything is patched. + raise ValueError( + "install_onfly_decompress: config has no compressed-tensors quantization_config" + ) compressor, scheme, param_names, fmt = compression require_transformers("5.8", "on-the-fly compressed-tensors decompression") _install_numeric_expert_ordering() diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 7614cd13774..8320ecf02e5 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -317,16 +317,19 @@ def test_a_checkpoint_compressed_differently_needs_a_process_of_its_own(monkeypa """The hooks are process-wide: decompressing a second checkpoint with the first one's scheme and stored params would load it wrong without a word, so its install is refused and the hooks stay as they are. A config the installing prepare already stripped of its quantization_config is the - first checkpoint again.""" + first checkpoint again; one whose quantization_config fails to parse is not, and its error + goes through.""" _, down = _stock_expert_converters() compressions = { "a": (_ScaledCompressor, None, PACKED, "stub-format"), "b": (_ScaledCompressor, None, (*PACKED, "weight_zero_point"), "stub-format"), + None: None, } def read(cfg): - if cfg.checkpoint is None: - raise ValueError("install_onfly_decompress: config has no quantization_config") + if cfg.checkpoint == "malformed": + # compressed_tensors' parse raises pydantic's ValidationError, a ValueError. + raise ValueError("1 validation error for QuantizationConfig") return compressions[cfg.checkpoint] monkeypatch.setattr( @@ -339,6 +342,8 @@ def read(cfg): onfly.install_onfly_decompress(SimpleNamespace(checkpoint=None)) with pytest.raises(RuntimeError, match="needs a process of its own"): onfly.install_onfly_decompress(SimpleNamespace(checkpoint="b")) + with pytest.raises(ValueError, match="validation error"): + onfly.install_onfly_decompress(SimpleNamespace(checkpoint="malformed")) assert mu.get_model_conversion_mapping is installed @@ -348,9 +353,7 @@ def _install_for_deepseek_v3(monkeypatch): compression = (_ScaledCompressor, None, PACKED, "stub-format") def read(cfg): - if getattr(cfg, "quantization_config", None) is None: - raise ValueError("install_onfly_decompress: config has no quantization_config") - return compression + return None if getattr(cfg, "quantization_config", None) is None else compression monkeypatch.setattr( mu, "get_model_conversion_mapping", lambda model, **kw: list(_stock_expert_converters()) From 31f8f6ed356daf20c8d2d6aec204a29b8762b1bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 21:04:08 +0000 Subject: [PATCH 10/16] Forward the conversion mapping's arguments as given, and drop an ordering patch transformers made moot The rewritten get_model_conversion_mapping forwards its arguments as it gets them: pinning two of them as keywords made a fourth positional collide with them. The rewritten expert converter keeps the scope transformers sets on a converter after building it (scope_prefix, base_model_prefix, force_cpu, whichever the release has), and the integer-keeping copy hook calls transformers' own copy instead of standing in for it. The patch that re-sorted experts numerically is gone: every transformers from 5.5 sorts the state dict by dot_natural_key before collecting, so experts already arrive in numeric order, and the patch only kept a process-wide id() map. stage_full_remote_closure gets a unit test, the install message prints on rank 0 only, and the module docstrings keep the contract, not the design notes. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 145 +++++------------- .../torch/models/test_onfly_decompress.py | 53 +++---- tests/unit/torch/models/test_remote_code.py | 53 +++++++ tests/unit/torch/utils/test_ep_loader.py | 4 +- 4 files changed, 114 insertions(+), 141 deletions(-) create mode 100644 tests/unit/torch/models/test_remote_code.py diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 19afe608bb7..59a32c8c2c2 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -13,39 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""On-the-fly decompression of compressed-tensors quantized MoE experts into the transformers-5.x -fused-experts layout, DURING from_pretrained — so a2a PTQ can run on the ORIGINAL compressed -checkpoint with no decompressed intermediate on disk, EP-sharded (each rank decompresses only its -own experts). - -FORMAT-AGNOSTIC. The dequant math is delegated to the ``compressed_tensors`` library and driven -entirely by the checkpoint's own ``quantization_config`` (the ``format`` string + the per-group -``QuantizationArgs``). Any format the library registers works with ZERO per-format code here: -``pack-quantized`` (INT, symmetric or asymmetric, any bit-width / group_size / strategy), -``nvfp4-pack-quantized`` (FP4 + FP8 block scale + FP32 global), ``naive-quantized`` / -``float-quantized`` / ``int-quantized`` (unpacked FP8/INT), etc. The two universal hooks the library -exposes per format are all we need: - - ``compressor.compression_param_names(scheme)`` -> the exact set of stored keys for one weight - (e.g. ``("weight_packed","weight_scale","weight_shape")`` for pack, ``("weight","weight_scale")`` - for float), i.e. which per-expert sibling tensors to pull. - - ``compressor.decompress({: tensor}, scheme)["weight"]`` -> the dequantized weight. -``compressed_tensors`` must be importable at load time, and transformers must be 5.8 or newer: earlier -releases' ``Concatenate`` looks up every source pattern, including the sibling keys the decompression -already consumed. - -The model-specific knobs remain parameters; defaults target moonshotai/Kimi-K2.7-Code (DeepSeek-V3 -backbone): - - install_onfly_decompress(hf_config) — auto-detects format + quant args from the config - - to_native_backbone_config(vlm_config, ...) — build the native backbone config - -Mechanism (mirrors transformers' FP8 quantizer ``update_weight_conversions``): the fused-MoE model -already ships a ``WeightConverter`` that fuses per-expert ``gate/up/down_proj.weight`` -> 3-D -``gate_up_proj``/``down_proj`` via ``MergeModulelist``+``Concatenate``. We rewrite those converters so -each ``.weight`` source becomes the checkpoint's actual stored keys for that weight -(``compression_param_names``), and PREPEND a ``DecompressExperts`` op that pairs those per expert and -dequantizes -> bf16 BEFORE the merge/concat collapse the per-expert structure. -``install_onfly_decompress()`` monkeypatches ``get_model_conversion_mapping``, for the model built from -the config it is given: every other model ``from_pretrained`` loads keeps its stock converters. +"""Decompress a compressed-tensors MoE checkpoint's experts as ``from_pretrained`` loads them. + +Each rank decompresses only its own experts into the transformers 5.x fused-experts layout, so +expert-parallel PTQ runs on the original compressed checkpoint, with no decompressed copy on disk. +The math is the ``compressed_tensors`` library's, driven by the checkpoint's ``quantization_config``: +``compressor.compression_param_names(scheme)`` names the stored keys of one weight, and +``compressor.decompress(...)["weight"]`` dequantizes them, so every format the library registers +works. It needs ``compressed_tensors`` importable and transformers 5.8 or newer (earlier releases' +``Concatenate`` looks up the sibling keys the decompression consumed). + +:func:`install_onfly_decompress` rewrites the per-expert converters ``get_model_conversion_mapping`` +returns, for the model built from the config it marks: each ``.weight`` source becomes the stored +keys, and a ``DecompressExperts`` op ahead of the merge dequantizes each expert. +:func:`to_native_backbone_config` builds a VLM's transformers-native backbone config. """ import torch @@ -142,21 +123,13 @@ def _read_compression(hf_config): def _make_op_class(compressor, scheme, param_names): - """Build the per-expert decompress op lazily (so this module imports without transformers present), - bound to a compressed_tensors compressor + scheme + the checkpoint's stored param names. The FIRST - param name is the "anchor" (the source key the decompressed list is emitted under, replacing the - original ``.weight`` so the downstream MergeModulelist/Concatenate fuse exactly as uncompressed). - - NOTE: ``DecompressExperts`` subclasses ``MergeModulelist`` so that it shards experts AT LOAD. The - transformers loader (``convert_and_load_state_dict_in_model``) only computes the per-expert - ``shard_index`` (``tensor_idx``) when ``isinstance(converter.operations[0], MergeModulelist)``; that - ``tensor_idx`` is what makes ``GroupedGemmParallel.shard_tensor`` keep ONLY this rank's local experts - (returning ``None`` for the rest). If our op were a plain ``ConversionOps`` prepended before the - original ``MergeModulelist``, ``operations[0]`` would not match, ``tensor_idx`` would be ``None``, and - every rank would materialize + decompress ALL experts (garbled, and a ~device-max transient -> OOM at - scale, e.g. Kimi-K2.7-Code EP32). Subclassing keeps ``operations[0]`` a ``MergeModulelist`` instance so - per-expert sharding fires; we still override ``convert`` to decompress (the ``dim``/merge behavior of - the base class is unused -- the ORIGINAL ``MergeModulelist`` later in the chain does the stacking).""" + """Build the per-expert decompress op, bound to a compressor, its scheme and the stored keys. + + The first stored key is the anchor the decompressed experts are emitted under, in place of the + ``.weight`` the downstream merge expects. The op subclasses ``MergeModulelist`` only because the + loader shards experts at load (``tensor_idx``, so non-local experts arrive as ``None``) only when + a converter's first op is one; the original merge later in the chain still does the stacking. + """ from transformers.core_model_loading import MergeModulelist anchor = "." + param_names[0] @@ -237,54 +210,17 @@ def _compressedify(conv, compressor, scheme, param_names): by_pname[param_names[0]].append(p) # leave non-weight sources as-is (none expected) new_sources = [s for pn in param_names for s in by_pname[pn]] tgt = getattr(conv, "_original_target_patterns", conv.target_patterns) - return WeightConverter( + new = WeightConverter( source_patterns=new_sources, target_patterns=tgt, operations=[DecompressExperts(), *conv.operations], ) - - -# id(converter) -> {source_pattern: [expert indices in append order]}; used to re-sort experts numerically. -_EXPERT_ORDER: dict = {} - - -def _install_numeric_expert_ordering(): - """Patch WeightConverter.add_tensor / materialize_tensors so per-expert tensors are merged in NUMERIC index order. - - The loader appends in checkpoint-iteration order, which is LEXICOGRAPHIC for many indices (0,1,10,100,...,2,...) — - that scrambles fused-expert slots. A no-op for numerically-ordered checkpoints. - """ - import re as _re - - from transformers.core_model_loading import WeightConverter - - if getattr(WeightConverter.add_tensor, "_onfly_patched", False): - return - orig_add = WeightConverter.add_tensor - orig_mat = WeightConverter.materialize_tensors - - def add_tensor(self, target_key, source_key, source_pattern, future): - orig_add(self, target_key, source_key, source_pattern, future) - m = _re.search(r"experts\.(\d+)\.", source_key) - if m is not None: - _EXPERT_ORDER.setdefault(id(self), {}).setdefault(source_pattern, []).append( - int(m.group(1)) - ) - - def materialize_tensors(self): - order = _EXPERT_ORDER.pop(id(self), None) - if order: - for pat, idxs in order.items(): - lst = self.collected_tensors.get(pat) - if lst is not None and len(lst) == len(idxs): - self.collected_tensors[pat] = [ - f for _, f in sorted(zip(idxs, lst), key=lambda x: x[0]) - ] - return orig_mat(self) - - setattr(add_tensor, "_onfly_patched", True) - WeightConverter.add_tensor = add_tensor - WeightConverter.materialize_tensors = materialize_tensors + # What else transformers sets on a converter when it maps a model -- the sub-model scope it + # matches keys under, whether it loads on the CPU -- whichever of them this release has. + for attr in ("scope_prefix", "base_model_prefix", "force_cpu"): + if hasattr(conv, attr): + setattr(new, attr, getattr(conv, attr)) + return new def _install_keep_int_native(): @@ -294,23 +230,17 @@ def _install_keep_int_native(): bit pattern of packed weights (`weight_packed`), shapes (`weight_shape`), zero-points and g_idx. Floats (the scale, global scale, dense weights) still cast normally. """ - import torch import transformers.core_model_loading as cml if getattr(cml._materialize_copy, "_onfly_patched", False): return + orig = cml._materialize_copy def _materialize_copy(tensor, device=None, dtype=None): t = tensor[...] - if ( - dtype is not None - and not t.is_floating_point() - and torch.empty((), dtype=dtype).is_floating_point() - ): + if dtype is not None and dtype.is_floating_point and not t.is_floating_point(): dtype = None # keep packed/shape/zero-point/g_idx integers native - if dtype is not None or device is not None: - t = t.to(device=device, dtype=dtype) - return t + return orig(t, device=device, dtype=dtype) setattr(_materialize_copy, "_onfly_patched", True) cml._materialize_copy = _materialize_copy @@ -348,6 +278,7 @@ def install_onfly_decompress(hf_config): import transformers.modeling_utils as mu from transformers.core_model_loading import WeightConverter + from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import require_transformers installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) @@ -373,13 +304,12 @@ def install_onfly_decompress(hf_config): ) compressor, scheme, param_names, fmt = compression require_transformers("5.8", "on-the-fly compressed-tensors decompression") - _install_numeric_expert_ordering() _install_keep_int_native() orig = mu.get_model_conversion_mapping - def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): - convs = orig(model, key_mapping=key_mapping, hf_quantizer=hf_quantizer, *a, **k) - config = getattr(model, "config", None) + def patched(*args, **kwargs): + convs = orig(*args, **kwargs) + config = getattr(args[0] if args else kwargs["model"], "config", None) if getattr(config, _MARK, None) is None: # Not the compressed checkpoint: rewritten, an uncompressed MoE's experts would match no # stored key and load newly initialized without an error. @@ -401,8 +331,7 @@ def patched(model, key_mapping=None, hf_quantizer=None, *a, **k): setattr(patched, "_onfly_compression", compression) mu.get_model_conversion_mapping = patched _mark(hf_config, fmt) - print( + print_rank_0( f"[onfly_decompress] installed: compressed_tensors '{fmt}' fused-MoE experts decompress " - f"(stored params={param_names}) -> bf16 during from_pretrained", - flush=True, + f"(stored params={param_names}) -> bf16 during from_pretrained" ) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 8320ecf02e5..85b87c907bd 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -50,14 +50,8 @@ @pytest.fixture def restore_loader_hooks(monkeypatch): """Undo the process-wide transformers hooks the decompressor installs.""" - for owner, name in ( - (WeightConverter, "add_tensor"), - (WeightConverter, "materialize_tensors"), - (cml, "_materialize_copy"), - (mu, "get_model_conversion_mapping"), - ): + for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) - monkeypatch.setattr(onfly, "_EXPERT_ORDER", {}) class _ScaledCompressor: @@ -121,6 +115,21 @@ def test_non_weight_sources_are_kept_verbatim(): ) +def test_a_rewritten_converter_keeps_the_scope_transformers_set_on_it(): + """transformers scopes a sub-model's converters (a VLM's language model) once it has built them; + the rewrite keeps whichever of those fields this release has.""" + _, down = _stock_expert_converters() + down = copy.deepcopy(down) + fields = {"scope_prefix": "language_model", "base_model_prefix": "model", "force_cpu": True} + fields = {name: value for name, value in fields.items() if hasattr(type(down), name)} + for name, value in fields.items(): + setattr(down, name, value) + + rewritten = onfly._compressedify(down, _ScaledCompressor, None, PACKED) + + assert {name: getattr(rewritten, name) for name in fields} == fields + + # ----------------------------------------------------------------- decompressing during the load @@ -198,21 +207,6 @@ def decompress(cls, state_dict, scheme): assert seen == [(["weight_packed", "weight_scale"], "scheme")] * 2 -@pytest.mark.usefixtures("restore_loader_hooks") -def test_experts_merge_in_numeric_not_lexicographic_order(): - """Checkpoint keys iterated as strings yield experts 0, 1, 10, 11, 2, ...; the fused tensor must - still hold expert i in slot i (unfixed, expert 10 lands in slot 2).""" - onfly._install_numeric_expert_ordering() - _, down = _stock_expert_converters() - keys = {e: f"model.layers.0.mlp.experts.{e}.down_proj.weight" for e in range(12)} - checkpoint = {keys[e]: torch.full((2, 3), float(e)) for e in sorted(keys, key=str)} - assert list(checkpoint)[2] == keys[10] # the order the fix exists for - - fused = _load([down], checkpoint)["model.layers.0.mlp.experts.down_proj"] - - assert fused[:, 0, 0].tolist() == [float(e) for e in range(12)] - - @pytest.mark.usefixtures("restore_loader_hooks") def test_integer_checkpoint_tensors_keep_their_dtype_on_load(): """Without a quantizer the loader casts every tensor to the model's float dtype, which destroys @@ -253,7 +247,7 @@ def test_install_rewrites_only_the_per_expert_weight_converters(monkeypatch): calls = [] def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): - calls.append((model, key_mapping, hf_quantizer)) + calls.append((model, key_mapping, hf_quantizer, add_legacy)) return list(stock) compression = (_ScaledCompressor, None, PACKED, "stub-format") @@ -263,10 +257,10 @@ def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): onfly.install_onfly_decompress(config) model = SimpleNamespace(config=config) - # positionally, the way from_pretrained calls it - out = mu.get_model_conversion_mapping(model, {"old": "new"}, "quantizer") + # positionally, the way from_pretrained calls it, and a later parameter too + out = mu.get_model_conversion_mapping(model, {"old": "new"}, "quantizer", False) - assert calls == [(model, {"old": "new"}, "quantizer")] + assert calls == [(model, {"old": "new"}, "quantizer", False)] assert len(out) == len(stock) assert all(kept is original for kept, original in zip(out[:3], stock[:3])) assert [conv.source_patterns[0] for conv in out[3:]] == [ @@ -294,12 +288,7 @@ def stock_mapping(model, **kwargs): monkeypatch.setattr(onfly, "_read_compression", read) def hooks(): - return ( - mu.get_model_conversion_mapping, - WeightConverter.add_tensor, - WeightConverter.materialize_tensors, - cml._materialize_copy, - ) + return mu.get_model_conversion_mapping, cml._materialize_copy onfly.install_onfly_decompress(SimpleNamespace()) installed = hooks() diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py new file mode 100644 index 00000000000..99e2c5dc169 --- /dev/null +++ b/tests/unit/torch/models/test_remote_code.py @@ -0,0 +1,53 @@ +# 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. + +"""Staging every module of a local checkpoint's remote code.""" + +import shutil + +import pytest + +dmu = pytest.importorskip("transformers.dynamic_module_utils") + +from modelopt.torch.models._remote_code import stage_full_remote_closure + + +def test_every_module_of_a_local_checkpoint_is_staged(monkeypatch, tmp_path): + """transformers stages a remote module and its direct imports only, so the third module of + a -> b -> c is missing until the wrap copies the checkpoint's other modules. Installing again + wraps the loader once.""" + ckpt, stage = tmp_path / "ckpt", tmp_path / "stage" + ckpt.mkdir() + stage.mkdir() + for name, body in ( + ("a.py", "from .b import B\n"), + ("b.py", "from .c import C\n"), + ("c.py", ""), + ): + (ckpt / name).write_text(body) + + def stage_one_hop(path, module_file, *args, **kwargs): + for name in (module_file, "b.py"): + shutil.copy(ckpt / name, stage / name) + return str(stage / module_file) + + monkeypatch.setattr(dmu, "get_cached_module_file", stage_one_hop) + stage_full_remote_closure() + wrapped = dmu.get_cached_module_file + stage_full_remote_closure() + + assert dmu.get_cached_module_file is wrapped + assert dmu.get_cached_module_file(str(ckpt), "a.py") == str(stage / "a.py") + assert sorted(p.name for p in stage.iterdir()) == ["a.py", "b.py", "c.py"] diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 18809cf68ed..c171861753a 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -575,7 +575,9 @@ def resolve(path, rank): monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(moe, "require_transformers", lambda *args: None) monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) - config = SimpleNamespace(quantization_config={"quant_method": "compressed-tensors"}) + config = SimpleNamespace( + quantization_config={"quant_method": "compressed-tensors"}, base_model_ep_plan=_EP_PLAN + ) plan = DistributionPlan(2, ep_size=2) with pytest.raises(NotImplementedError, match="cannot dequantize"): From 4c4d80d2ac2f0d07afd2509aabce52868180029f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:15:19 +0000 Subject: [PATCH 11/16] Refuse a compressed checkpoint with several groups or stray expert keys _read_compression took the config group whose targets mention "expert", else the first, as the experts' scheme. Only the expert converters are rewritten, so another group's weights loaded newly initialized, and a guessed group decompressed the experts with the wrong bits, group size or zero point. A checkpoint with more than one group is now refused before anything is patched. The Kimi checkpoints this path serves carry one group each (Kimi-K2.5, K2.6 and K2.7-Code INT4 pack-quantized, Kimi-K3 MXFP4), whose ignore list leaves everything but the routed experts in bf16. DecompressExperts handed a checkpoint key it did not consume on to the merge, which would fuse it into the wrong slot. It now raises, which transformers reports as a conversion error that fails the load. install_onfly_decompress checks the transformers 5.8 floor before it imports core_model_loading, so an older release gets that message rather than an ImportError, and _mark no longer writes a dict key no reader looks at. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 47 ++++++++++--------- .../torch/models/test_onfly_decompress.py | 39 +++++++-------- 2 files changed, 42 insertions(+), 44 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 59a32c8c2c2..50be5914f48 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -91,8 +91,8 @@ def _read_compression(hf_config): ``compressor`` is the registered compressor for ``format`` (its ``decompress`` classmethod is the dequant); ``scheme`` is the ``QuantizationScheme`` (supplies ``scheme.weights`` to ``decompress``); - ``param_names`` is the exact set of stored sibling keys per weight. The experts' group is chosen as - the sole config group, else the one whose ``targets`` reference "expert", else the first. + ``param_names`` is the exact set of stored sibling keys per weight. The checkpoint's one config + group is the experts' scheme; a checkpoint with several is refused. ``None`` for a config with no ``quantization_config``: nothing to decompress.""" # Look the config up first: with none, compressed_tensors need not even be installed. A config # that fails to parse raises pydantic's ValidationError, a ValueError of its own. @@ -104,17 +104,15 @@ def _read_compression(hf_config): qconfig = QuantizationConfig.model_validate(quantization_config) groups = qconfig.config_groups - if len(groups) == 1: - scheme = next(iter(groups.values())) - else: - scheme = next( - ( - g - for g in groups.values() - if any("expert" in str(t).lower() for t in (getattr(g, "targets", None) or [])) - ), - next(iter(groups.values())), + if len(groups) != 1: + # Only the expert converters are rewritten, so another group's weights would load newly + # initialized, and which group holds the experts would be a guess. + raise NotImplementedError( + f"install_onfly_decompress: the checkpoint has {len(groups)} compressed-tensors config " + f"groups ({', '.join(groups) or 'none'}); on-the-fly decompression handles a single " + "group, compressing only the experts. Decompress the checkpoint first." ) + [scheme] = groups.values() fmt = qconfig.format fmt = fmt.value if hasattr(fmt, "value") else str(fmt) compressor = BaseCompressor.get_value_from_registry(fmt) @@ -184,7 +182,12 @@ def convert(self, input_dict, *args, **kwargs): for i in range(len(anchor_list)) if anchor_list[i] is not None ] - out.update(input_dict) # nothing else expected, but pass through defensively + if input_dict: + # A stored param without its weight, or a source that is no weight: passed on, the + # merge would fuse it into the wrong slot or fail far from here. + raise ValueError( + f"DecompressExperts: unconsumed checkpoint keys {sorted(input_dict)}" + ) return out return DecompressExperts @@ -207,7 +210,7 @@ def _compressedify(conv, compressor, scheme, param_names): for pn in param_names: by_pname[pn].append(stem + "." + pn + "$") else: - by_pname[param_names[0]].append(p) # leave non-weight sources as-is (none expected) + by_pname[param_names[0]].append(p) # DecompressExperts refuses it (none expected) new_sources = [s for pn in param_names for s in by_pname[pn]] tgt = getattr(conv, "_original_target_patterns", conv.target_patterns) new = WeightConverter( @@ -252,10 +255,7 @@ def _materialize_copy(tensor, device=None, dtype=None): def _mark(hf_config, fmt): """Mark ``hf_config``'s load for the installed hooks; ``from_pretrained``'s copy of it keeps the mark.""" - if isinstance(hf_config, dict): - hf_config[_MARK] = fmt - else: - setattr(hf_config, _MARK, fmt) + setattr(hf_config, _MARK, fmt) def is_marked(hf_config) -> bool: @@ -275,14 +275,16 @@ def install_onfly_decompress(hf_config): for it marks that config too, and installing for a checkpoint compressed differently raises. Only a model built from a marked config has its converters rewritten; any other model ``from_pretrained`` loads in the process, an uncompressed MoE say, keeps its stock converters.""" - import transformers.modeling_utils as mu - from transformers.core_model_loading import WeightConverter - from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import require_transformers - installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) compression = _read_compression(hf_config) + require_transformers("5.8", "on-the-fly compressed-tensors decompression") + # Imported once the release is known to have them: transformers 4 has no core_model_loading. + import transformers.modeling_utils as mu + from transformers.core_model_loading import WeightConverter + + installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) if installed is not None: if compression is None: # The prepare that installed the hooks drops quantization_config from its config, already @@ -303,7 +305,6 @@ def install_onfly_decompress(hf_config): "install_onfly_decompress: config has no compressed-tensors quantization_config" ) compressor, scheme, param_names, fmt = compression - require_transformers("5.8", "on-the-fly compressed-tensors decompression") _install_keep_int_native() orig = mu.get_model_conversion_mapping diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 85b87c907bd..8f1bb084aa5 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -100,19 +100,21 @@ def _load(converters, checkpoint): # ------------------------------------------------------------- rewriting the per-expert converters -def test_non_weight_sources_are_kept_verbatim(): - """Only ``.weight`` sources name a compressed tensor; any other source passes through as is.""" +def test_a_source_the_decompression_does_not_consume_fails_the_load(): + """Only ``.weight`` sources name a compressed tensor. Another source still reaches the + decompress op, which refuses it rather than hand the merge a tensor it would fuse into the wrong + slot.""" conv = WeightConverter( source_patterns=["mlp.experts.*.gate_proj.weight", "mlp.experts.*.gate_proj.bias"], target_patterns="mlp.experts.gate_proj", operations=[MergeModulelist(dim=0)], ) - rewritten = onfly._compressedify(conv, _ScaledCompressor, None, PACKED) + key = "model.layers.0.mlp.experts.0.gate_proj." + checkpoint = {key + name: torch.ones(2, 2, dtype=torch.int32) for name in PACKED} + checkpoint[key + "bias"] = torch.ones(2) - assert "mlp.experts.*.gate_proj.bias" in rewritten.source_patterns - assert sorted(p for p in rewritten.source_patterns if "weight" in p) == sorted( - f"mlp.experts.*.gate_proj.{name}$" for name in PACKED - ) + with pytest.raises(ValueError, match=r"unconsumed checkpoint keys .*gate_proj\.bias"): + _load([onfly._compressedify(conv, _ScaledCompressor, None, PACKED)], checkpoint) def test_a_rewritten_converter_keeps_the_scope_transformers_set_on_it(): @@ -505,28 +507,23 @@ def _pack_quantized(**groups): @pytest.mark.parametrize( - ("groups", "zero_point"), + "groups", [ - ({"attn": _int4(True, ATTN), "moe": _int4(False, EXPERTS)}, True), - ({"moe": _int4(True, EXPERTS), "attn": _int4(False, ATTN)}, False), - ({"attn": _int4(False, ATTN), "dense": _int4(True, ["Linear"])}, True), + {"attn": _int4(True, ATTN), "moe": _int4(False, EXPERTS)}, + {"attn": _int4(False, ATTN), "dense": _int4(True, ["Linear"])}, ], - ids=["expert-group-second", "expert-group-first", "no-expert-group-takes-the-first"], + ids=["an-expert-group-among-others", "no-expert-group"], ) -def test_the_expert_config_group_decides_the_stored_params(groups, zero_point): - """With several config groups the one targeting experts wins, else the first; an asymmetric - scheme stores a zero-point next to the packed weight, a symmetric one does not.""" +def test_a_checkpoint_compressed_in_several_config_groups_is_refused(groups): + """Only the expert converters are rewritten: another group's weights would load newly + initialized, and a group guessed for the experts would decompress them with the wrong scheme.""" pytest.importorskip( "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" ) config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) - _, scheme, names, fmt = onfly._read_compression(config) - - assert fmt == "pack-quantized" - assert names[: len(PACKED)] == PACKED - assert ("weight_zero_point" in names) is zero_point - assert scheme.weights.symmetric is not zero_point + with pytest.raises(NotImplementedError, match="2 compressed-tensors config groups"): + onfly._read_compression(config) def test_real_pack_quantized_experts_round_trip_exactly(): From 462e599024cd7cae64627172dfb51a71ec9561fb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:10:57 +0000 Subject: [PATCH 12/16] Load a decompressed checkpoint without its compressed-tensors quantizer Kimi-K2.5's config class copies its text config's quantization_config to the top level, where the text_config swap leaves it, so the expert-parallel load handed from_pretrained a marked config that still carried it. That builds transformers' compressed-tensors quantizer, whose own expert decompression (from 5.13) claims the per-expert converters before the decompressor's hook sees them. On the marked path the load now deletes quantization_config from its copy of the config, at the top level and on text_config; deleted rather than set to None, which transformers before 5.11 takes for a quantized checkpoint. The decompressor's mark has one scope now: a config counts as marked when it or its text_config carries the mark, the hook reads it that way and clears it at both levels, and to_native_backbone_config carries it to the backbone config it derives. The hook read the top level only, so a VLM whose text_config was marked loaded stock converters while the format check passed it, and a backbone loaded on its own from the derived config was never rewritten. The checkpoint's one config group must target the experts (a class target or a regex naming them, as transformers' own compressed-tensors loader finds the experts' group), and a quantization_config with no groups is refused as nothing to decompress rather than with the several-groups message. stage_full_remote_closure copies each module to a temp file in the stage dir and renames it into place, so a rank never imports a module another rank is still writing. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 49 +++++++++++----- modelopt/torch/models/_remote_code.py | 8 ++- .../torch/utils/plugins/model_load_utils.py | 9 +++ .../torch/models/test_onfly_decompress.py | 57 ++++++++++++++++--- tests/unit/torch/models/test_remote_code.py | 33 +++++++++++ tests/unit/torch/utils/test_ep_loader.py | 40 ++++++++----- 6 files changed, 161 insertions(+), 35 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 50be5914f48..3811d4bbf57 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -62,7 +62,12 @@ def to_native_backbone_config( d.pop(key, None) d["model_type"] = model_type d["architectures"] = list(architectures) - return CONFIG_MAPPING[model_type].from_dict(d) + native = CONFIG_MAPPING[model_type].from_dict(d) + fmt = getattr(vlm_config, _MARK, None) or getattr(text_config, _MARK, None) + if fmt is not None: + # A backbone loaded from this config on its own decompresses as the VLM would. + _mark(native, fmt) + return native def _find_quantization_config(hf_config): @@ -92,8 +97,8 @@ def _read_compression(hf_config): ``compressor`` is the registered compressor for ``format`` (its ``decompress`` classmethod is the dequant); ``scheme`` is the ``QuantizationScheme`` (supplies ``scheme.weights`` to ``decompress``); ``param_names`` is the exact set of stored sibling keys per weight. The checkpoint's one config - group is the experts' scheme; a checkpoint with several is refused. - ``None`` for a config with no ``quantization_config``: nothing to decompress.""" + group is the experts' scheme; a checkpoint with several, or whose one group does not target the + experts, is refused. ``None`` for a config with no ``quantization_config``: nothing to decompress.""" # Look the config up first: with none, compressed_tensors need not even be installed. A config # that fails to parse raises pydantic's ValidationError, a ValueError of its own. quantization_config = _find_quantization_config(hf_config) @@ -104,15 +109,29 @@ def _read_compression(hf_config): qconfig = QuantizationConfig.model_validate(quantization_config) groups = qconfig.config_groups - if len(groups) != 1: + if not groups: + raise ValueError( + "install_onfly_decompress: the checkpoint's compressed-tensors quantization_config has " + "no config groups: nothing to decompress." + ) + if len(groups) > 1: # Only the expert converters are rewritten, so another group's weights would load newly # initialized, and which group holds the experts would be a guess. raise NotImplementedError( f"install_onfly_decompress: the checkpoint has {len(groups)} compressed-tensors config " - f"groups ({', '.join(groups) or 'none'}); on-the-fly decompression handles a single " - "group, compressing only the experts. Decompress the checkpoint first." + f"groups ({', '.join(groups)}); on-the-fly decompression handles a single group, " + "compressing only the experts. Decompress the checkpoint first." ) [scheme] = groups.values() + # The experts' group as transformers' own compressed-tensors loader finds it: a class target + # (the checkpoint's experts are Linear modules) or a regex naming the experts. Another group + # leaves the experts stored as plain weights, which the rewritten converters would not find. + if not any(not t.startswith("re:") or "experts" in t for t in scheme.targets): + raise NotImplementedError( + f"install_onfly_decompress: the checkpoint's compressed-tensors config group targets " + f"{scheme.targets}, not the experts; on-the-fly decompression handles a checkpoint " + "whose experts are what it compresses. Decompress the checkpoint first." + ) fmt = qconfig.format fmt = fmt.value if hasattr(fmt, "value") else str(fmt) compressor = BaseCompressor.get_value_from_registry(fmt) @@ -273,8 +292,10 @@ def install_onfly_decompress(hf_config): The hooks are process-wide, so a process decompresses one compressed checkpoint: installing again for it marks that config too, and installing for a checkpoint compressed differently raises. Only a - model built from a marked config has its converters rewritten; any other model ``from_pretrained`` - loads in the process, an uncompressed MoE say, keeps its stock converters.""" + model built from a marked config, or from one whose ``text_config`` is marked, has its converters + rewritten; any other model ``from_pretrained`` loads in the process, an uncompressed MoE say, + keeps its stock converters. :func:`to_native_backbone_config` carries the mark to the backbone + config it derives.""" from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import require_transformers @@ -287,8 +308,8 @@ def install_onfly_decompress(hf_config): installed = getattr(mu.get_model_conversion_mapping, "_onfly_compression", None) if installed is not None: if compression is None: - # The prepare that installed the hooks drops quantization_config from its config, already - # marked, and may run again on that same config: nothing new to decompress. + # A config its prepare dropped quantization_config from after installing (Kimi-K3's), + # already marked, may be prepared again: nothing new to decompress. return if compression != installed: raise RuntimeError( @@ -311,13 +332,15 @@ def install_onfly_decompress(hf_config): def patched(*args, **kwargs): convs = orig(*args, **kwargs) config = getattr(args[0] if args else kwargs["model"], "config", None) - if getattr(config, _MARK, None) is None: + if not is_marked(config): # Not the compressed checkpoint: rewritten, an uncompressed MoE's experts would match no # stored key and load newly initialized without an error. return convs # from_pretrained loads from a copy of the marked config; the mark must not reach the config - # a later save writes. - delattr(config, _MARK) + # a later save writes, at either level. + for c in (config, getattr(config, "text_config", None)): + if hasattr(c, _MARK): + delattr(c, _MARK) out = [] for c in convs: # only the per-expert MoE converters (have `.weight` sources under `experts`) diff --git a/modelopt/torch/models/_remote_code.py b/modelopt/torch/models/_remote_code.py index 43780f39c65..7db39497bab 100644 --- a/modelopt/torch/models/_remote_code.py +++ b/modelopt/torch/models/_remote_code.py @@ -18,6 +18,7 @@ import glob import os import shutil +import tempfile import warnings @@ -52,7 +53,12 @@ def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): for py in glob.glob(os.path.join(src, "*.py")): dst = os.path.join(stage_dir, os.path.basename(py)) if not os.path.exists(dst): - shutil.copy(py, dst) + # Every rank loading the checkpoint stages into this dir and imports from it: + # renamed into place (atomic), a file appears there whole, never half written. + fd, tmp = tempfile.mkstemp(dir=stage_dir) + os.close(fd) + shutil.copy(py, tmp) + os.replace(tmp, dst) except OSError as e: # a failed copy only matters if a deeper import needs the file warnings.warn(f"Could not stage the remote code of {src}: {e}") return resolved diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 673f930f09f..c18f130b80a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2530,6 +2530,15 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ) # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. hf_config = copy.deepcopy(hf_config) + if _decompressed_as_loaded(hf_config): + # The decompressor hands the experts over dequantized. Left on the config (Kimi-K2.5's config + # class copies its text config's to the top level), the checkpoint's quantization_config + # would build transformers' compressed-tensors quantizer, which from 5.13 claims the same + # expert converters first. Deleted rather than set to None: before 5.11 transformers takes + # any quantization_config attribute, None too, for a quantized checkpoint. + for cfg in (hf_config, getattr(hf_config, "text_config", None)): + if cfg is not None and hasattr(cfg, "quantization_config"): + del cfg.quantization_config ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 8f1bb084aa5..e46a71225fa 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -387,6 +387,36 @@ def test_only_the_compressed_checkpoint_loads_through_the_rewritten_converters(m assert onfly._MARK not in model.config.to_dict() +@pytest.mark.usefixtures("restore_loader_hooks") +@pytest.mark.parametrize("marked", ["vlm", "text-config"]) +@pytest.mark.parametrize("loaded", ["vlm", "backbone"]) +def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters( + monkeypatch, marked, loaded +): + """A VLM carries its compressed-tensors config at the top level or on its text config, and the + install marks the config it is handed. Whether the load builds the VLM, its text config replaced + by the native backbone's, or that backbone on its own, its expert converters are rewritten, and + neither level keeps the mark for a later save.""" + monkeypatch.setattr( + mu, "get_model_conversion_mapping", lambda model, **kw: list(_stock_expert_converters()) + ) + compression = (_ScaledCompressor, None, PACKED, "stub-format") + monkeypatch.setattr(onfly, "_read_compression", lambda cfg: compression) + text = _RemoteTextConfig(hidden_size=64, quantization_config={"format": "stub-format"}) + vlm = SimpleNamespace(text_config=text) + onfly.install_onfly_decompress(vlm if marked == "vlm" else text) + vlm.text_config = onfly.to_native_backbone_config(vlm) + model = SimpleNamespace(config=copy.deepcopy(vlm if loaded == "vlm" else vlm.text_config)) + + out = mu.get_model_conversion_mapping(model) + + assert [conv.source_patterns[0] for conv in out] == [ + "mlp.experts.*.gate_proj.weight_packed$", + "mlp.experts.*.down_proj.weight_packed$", + ] + assert not onfly.is_marked(model.config) + + # ------------------------------------------------------------------------ reading the checkpoint @@ -507,22 +537,35 @@ def _pack_quantized(**groups): @pytest.mark.parametrize( - "groups", + ("groups", "error", "match"), [ - {"attn": _int4(True, ATTN), "moe": _int4(False, EXPERTS)}, - {"attn": _int4(False, ATTN), "dense": _int4(True, ["Linear"])}, + ( + {"attn": _int4(True, ATTN), "moe": _int4(False, EXPERTS)}, + NotImplementedError, + "2 compressed-tensors config groups", + ), + ( + {"attn": _int4(False, ATTN), "dense": _int4(True, ["Linear"])}, + NotImplementedError, + "2 compressed-tensors config groups", + ), + ({"attn": _int4(True, ATTN)}, NotImplementedError, "not the experts"), + ({}, ValueError, "nothing to decompress"), ], - ids=["an-expert-group-among-others", "no-expert-group"], + ids=["an-expert-group-among-others", "no-expert-group", "one-group-not-the-experts", "none"], ) -def test_a_checkpoint_compressed_in_several_config_groups_is_refused(groups): +def test_a_checkpoint_whose_experts_are_not_its_one_compressed_group_is_refused( + groups, error, match +): """Only the expert converters are rewritten: another group's weights would load newly - initialized, and a group guessed for the experts would decompress them with the wrong scheme.""" + initialized, a group guessed for the experts would decompress them with the wrong scheme, and + experts no group targets are stored as plain weights the rewritten converters would not find.""" pytest.importorskip( "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" ) config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) - with pytest.raises(NotImplementedError, match="2 compressed-tensors config groups"): + with pytest.raises(error, match=match): onfly._read_compression(config) diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py index 99e2c5dc169..d324773e47b 100644 --- a/tests/unit/torch/models/test_remote_code.py +++ b/tests/unit/torch/models/test_remote_code.py @@ -16,11 +16,13 @@ """Staging every module of a local checkpoint's remote code.""" import shutil +from types import SimpleNamespace import pytest dmu = pytest.importorskip("transformers.dynamic_module_utils") +from modelopt.torch.models import _remote_code from modelopt.torch.models._remote_code import stage_full_remote_closure @@ -51,3 +53,34 @@ def stage_one_hop(path, module_file, *args, **kwargs): assert dmu.get_cached_module_file is wrapped assert dmu.get_cached_module_file(str(ckpt), "a.py") == str(stage / "a.py") assert sorted(p.name for p in stage.iterdir()) == ["a.py", "b.py", "c.py"] + + +def test_a_module_is_staged_whole_or_not_at_all(monkeypatch, tmp_path): + """Every rank stages into the same dir and imports what it finds there, so a module must not + appear before its copy completes: a copy cut short leaves nothing behind for a rank to take as + staged, and the next load stages the whole module.""" + ckpt, stage = tmp_path / "ckpt", tmp_path / "stage" + ckpt.mkdir() + stage.mkdir() + (ckpt / "a.py").write_text("from .b import B\n") + (ckpt / "b.py").write_text("B = 1\n" * 100) + + def stage_only_the_module(path, module_file, *args, **kwargs): + shutil.copy(ckpt / module_file, stage / module_file) + return str(stage / module_file) + + def copy_cut_short(src, dst): + with open(src) as f, open(dst, "w") as out: + out.write(f.read()[:10]) + raise OSError(28, "No space left on device") + + monkeypatch.setattr(dmu, "get_cached_module_file", stage_only_the_module) + stage_full_remote_closure() + with monkeypatch.context() as m: + m.setattr(_remote_code, "shutil", SimpleNamespace(copy=copy_cut_short)) + with pytest.warns(UserWarning, match="No space left"): + dmu.get_cached_module_file(str(ckpt), "a.py") + assert not (stage / "b.py").exists() + + dmu.get_cached_module_file(str(ckpt), "a.py") + assert (stage / "b.py").read_text() == (ckpt / "b.py").read_text() diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index c171861753a..f7beeefbd62 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -559,33 +559,45 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): assert classes == [True] -def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes(monkeypatch): +def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( + ep_world_of_two, monkeypatch +): """A model spec that installs the on-the-fly decompressor (Kimi-K2.5's) marks the config; that - checkpoint passes the format check, and the same checkpoint unmarked does not.""" + checkpoint passes the format check and loads without its quantization_config, at either level, + which would build transformers' own compressed-tensors quantizer to compete with the + decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own.""" from modelopt.torch.models import _onfly_decompress - class _ResolvedPastFormatError(Exception): - """Stops the load where it resolves the checkpoint, past the format check.""" + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" - def resolve(path, rank): - raise _ResolvedPastFormatError + loaded = {} - monkeypatch.setattr(dist, "is_initialized", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 2) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) - monkeypatch.setattr(moe, "require_transformers", lambda *args: None) - monkeypatch.setattr(model_load_utils, "_resolve_checkpoint_dir", resolve) - config = SimpleNamespace( - quantization_config={"quant_method": "compressed-tensors"}, base_model_ep_plan=_EP_PLAN + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + quantization_config = {"quant_method": "compressed-tensors"} + text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) + config = SimpleNamespace(quantization_config=quantization_config, text_config=text) plan = DistributionPlan(2, ep_size=2) with pytest.raises(NotImplementedError, match="cannot dequantize"): load_and_prepare_ep("org/model", plan, hf_config=config) _onfly_decompress._mark(config, "pack-quantized") - with pytest.raises(_ResolvedPastFormatError): + with pytest.raises(_LoadedError): load_and_prepare_ep("org/model", plan, hf_config=config) + assert not hasattr(loaded["config"], "quantization_config") + assert not hasattr(loaded["config"].text_config, "quantization_config") + assert config.quantization_config is text.quantization_config is quantization_config + # ------------------------------------------------------------------------ quant wrappers From 7a33c5def581b89ba00f42c1d239ecfa751991e3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 02:13:50 +0000 Subject: [PATCH 13/16] Refuse a compressed checkpoint unless it compresses exactly its experts The on-the-fly decompressor rewrites only the experts' converters, and the expert-parallel load drops the quantization_config that would load any other compressed weight. The checkpoint's one config group can target a class such as Linear, which does not say which modules it compressed: the common compressed-tensors layout compresses every Linear but lm_head, and an ignore list can leave the experts plain instead. Either way part of the model loaded newly initialized behind a missing-keys warning. The load now checks the names the checkpoint stores before anything loads: every weight stored compressed must be an expert's, and every expert's weight compressed. The mark now holds the names a compressed weight is stored under, which the check reads off the config. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints pass it. install_onfly_decompress refuses a raw dict before it patches anything, since the load finds its hooks by a mark a dict cannot carry, and the load deletes the quantization_config of llm_config too, every level _ckpt_quant_method reads. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 65 ++++++++++++---- .../torch/utils/plugins/model_load_utils.py | 15 +++- .../torch/models/test_onfly_decompress.py | 74 +++++++++++++++++-- tests/unit/torch/utils/test_ep_loader.py | 34 ++++++--- 4 files changed, 154 insertions(+), 34 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 3811d4bbf57..fd953a6dddf 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -63,10 +63,10 @@ def to_native_backbone_config( d["model_type"] = model_type d["architectures"] = list(architectures) native = CONFIG_MAPPING[model_type].from_dict(d) - fmt = getattr(vlm_config, _MARK, None) or getattr(text_config, _MARK, None) - if fmt is not None: + stored_as = getattr(vlm_config, _MARK, None) or getattr(text_config, _MARK, None) + if stored_as is not None: # A backbone loaded from this config on its own decompresses as the VLM would. - _mark(native, fmt) + _mark(native, stored_as) return native @@ -77,10 +77,6 @@ def _find_quantization_config(hf_config): if qc is None: tc = getattr(hf_config, "text_config", None) qc = getattr(tc, "quantization_config", None) if tc is not None else None - if qc is None and isinstance(hf_config, dict): - qc = hf_config.get("quantization_config") or (hf_config.get("text_config") or {}).get( - "quantization_config" - ) if qc is None: return None if hasattr(qc, "to_dict"): @@ -123,9 +119,10 @@ def _read_compression(hf_config): "compressing only the experts. Decompress the checkpoint first." ) [scheme] = groups.values() - # The experts' group as transformers' own compressed-tensors loader finds it: a class target - # (the checkpoint's experts are Linear modules) or a regex naming the experts. Another group - # leaves the experts stored as plain weights, which the rewritten converters would not find. + # A group of regex targets alone must name the experts, or it leaves them stored as plain weights + # the rewritten converters would not find. A class target such as Linear does not say which + # modules it compressed (the experts, and any Linear the ignore list leaves): the expert-parallel + # load checks the names the checkpoint stores instead (check_compressed_keys). if not any(not t.startswith("re:") or "experts" in t for t in scheme.targets): raise NotImplementedError( f"install_onfly_decompress: the checkpoint's compressed-tensors config group targets " @@ -268,13 +265,14 @@ def _materialize_copy(tensor, device=None, dtype=None): cml._materialize_copy = _materialize_copy -# Marks a config whose model's converters the installed hooks rewrite (holds the checkpoint's format). +# Marks a config whose model's converters the installed hooks rewrite (holds the names each +# compressed weight is stored under). _MARK = "_modelopt_onfly_decompress" -def _mark(hf_config, fmt): +def _mark(hf_config, stored_as): """Mark ``hf_config``'s load for the installed hooks; ``from_pretrained``'s copy of it keeps the mark.""" - setattr(hf_config, _MARK, fmt) + setattr(hf_config, _MARK, stored_as) def is_marked(hf_config) -> bool: @@ -295,10 +293,14 @@ def install_onfly_decompress(hf_config): model built from a marked config, or from one whose ``text_config`` is marked, has its converters rewritten; any other model ``from_pretrained`` loads in the process, an uncompressed MoE say, keeps its stock converters. :func:`to_native_backbone_config` carries the mark to the backbone - config it derives.""" + config it derives, and the expert-parallel load checks the checkpoint against it + (:func:`check_compressed_keys`).""" from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import require_transformers + if isinstance(hf_config, dict): + # The load finds the hooks by a mark on the config object, which a dict cannot carry. + raise TypeError("install_onfly_decompress: pass the model's config object, not a dict") compression = _read_compression(hf_config) require_transformers("5.8", "on-the-fly compressed-tensors decompression") # Imported once the release is known to have them: transformers 4 has no core_model_loading. @@ -318,7 +320,7 @@ def install_onfly_decompress(hf_config): f"stored params {compression[2]}) needs a process of its own" ) # The same checkpoint through a config of its own, which its load must find marked too. - _mark(hf_config, installed[3]) + _mark(hf_config, installed[2]) return if compression is None: # Nothing to decompress: refused before anything is patched. @@ -354,8 +356,39 @@ def patched(*args, **kwargs): setattr(patched, "_onfly_compression", compression) mu.get_model_conversion_mapping = patched - _mark(hf_config, fmt) + _mark(hf_config, param_names) print_rank_0( f"[onfly_decompress] installed: compressed_tensors '{fmt}' fused-MoE experts decompress " f"(stored params={param_names}) -> bf16 during from_pretrained" ) + + +def check_compressed_keys(hf_config, tensor_names) -> None: + """Refuse a checkpoint that compresses anything but its experts, or leaves an expert uncompressed. + + The hooks rewrite only the experts' converters, and the expert-parallel load drops the + quantization_config that would load any other compressed weight: a weight compressed outside the + experts would load newly initialized, and so would an expert stored uncompressed, which the + rewritten converters do not read. ``tensor_names`` are the checkpoint's. A weight counts as + compressed when it stores a name ``hf_config``'s mark holds other than ``weight``, which an + uncompressed weight is stored under too. + """ + configs = (hf_config, getattr(hf_config, "text_config", None)) + stored_as = next(getattr(c, _MARK) for c in configs if getattr(c, _MARK, None) is not None) + compressed: set[str] = set() + expert_weights: set[str] = set() + for name in tensor_names: + stem, _, stored = name.rpartition(".") + if stored in stored_as and stored != "weight": + compressed.add(stem) + elif stored == "weight" and ".experts." in name: + expert_weights.add(stem) + outside = sorted(stem for stem in compressed if ".experts." not in stem) + plain = sorted(expert_weights - compressed) + if outside or plain or not compressed: + raise NotImplementedError( + f"The checkpoint compresses {len(compressed)} weights, {len(outside)} of them outside the " + f"experts {outside[:1]}, and stores {len(plain)} expert weights uncompressed {plain[:1]}: " + "on-the-fly decompression loads the experts alone, every one compressed, so the rest " + "would load newly initialized. Decompress the checkpoint first." + ) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index c18f130b80a..e415999228a 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2535,8 +2535,13 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # class copies its text config's to the top level), the checkpoint's quantization_config # would build transformers' compressed-tensors quantizer, which from 5.13 claims the same # expert converters first. Deleted rather than set to None: before 5.11 transformers takes - # any quantization_config attribute, None too, for a quantized checkpoint. - for cfg in (hf_config, getattr(hf_config, "text_config", None)): + # any quantization_config attribute, None too, for a quantized checkpoint. Every level + # _ckpt_quant_method reads is cleared. + for cfg in ( + hf_config, + getattr(hf_config, "text_config", None), + getattr(hf_config, "llm_config", None), + ): if cfg is not None and hasattr(cfg, "quantization_config"): del cfg.quantization_config ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config @@ -2562,6 +2567,12 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. # Resolved once the plan and the config pass, so a refused load downloads nothing. ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank()) + if _decompressed_as_loaded(hf_config): + # Only the experts decompress as they load, and the config cannot say what else the + # checkpoint compressed: its stored names are checked before anything loads. + from modelopt.torch.models import _onfly_decompress + + _onfly_decompress.check_compressed_keys(hf_config, checkpoint_tensor_files(ckpt_path)) # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two # factorizations of the same world: diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index e46a71225fa..5f0bddeab13 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -434,24 +434,34 @@ def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters( SimpleNamespace(text_config=SimpleNamespace(quantization_config={"at": "text"})), {"at": "text"}, ), - ({"text_config": {"quantization_config": {"at": "dict"}}}, {"at": "dict"}), ( SimpleNamespace(quantization_config=SimpleNamespace(to_dict=lambda: {"at": "obj"})), {"at": "obj"}, ), ], - ids=["top-level-first", "text-config", "raw-dict", "config-object"], + ids=["top-level-first", "text-config", "config-object"], ) def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_config(config, found): """A VLM may carry its compressed-tensors config at the top level or on ``text_config``, as a - config object or a raw dict; it is always returned as a plain dict.""" + config object or a dict; it is always returned as a plain dict.""" assert onfly._find_quantization_config(config) == found @pytest.mark.usefixtures("restore_loader_hooks") -def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone(monkeypatch): +@pytest.mark.parametrize( + ("config", "error", "match"), + [ + (SimpleNamespace(text_config=SimpleNamespace()), ValueError, "quantization_config"), + ({"quantization_config": {"format": "pack-quantized"}}, TypeError, "not a dict"), + ], + ids=["nothing-to-decompress", "raw-dict"], +) +def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( + monkeypatch, config, error, match +): """By the time a bf16 checkpoint's caller sees that error, nothing may be patched yet -- and - compressed_tensors, which such a checkpoint never needs, need not be installed.""" + compressed_tensors, which such a checkpoint never needs, need not be installed. Nor for a raw + dict, which cannot carry the mark its load would be found by.""" monkeypatch.setitem(sys.modules, "compressed_tensors", None) hooks = ( mu.get_model_conversion_mapping, @@ -460,8 +470,8 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone(monkeypatch cml._materialize_copy, ) - with pytest.raises(ValueError, match="quantization_config"): - onfly.install_onfly_decompress(SimpleNamespace(text_config=SimpleNamespace())) + with pytest.raises(error, match=match): + onfly.install_onfly_decompress(config) assert ( mu.get_model_conversion_mapping, @@ -471,6 +481,56 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone(monkeypatch ) == hooks +EXPERT = "model.layers.1.mlp.experts.0.gate_proj" +ATTENTION = "model.layers.1.self_attn.q_proj" +SHARED_EXPERT = "model.layers.1.mlp.shared_experts.up_proj" # a dense MLP, not the routed experts + + +@pytest.mark.parametrize( + ("stored_as", "names", "match"), + [ + ( + PACKED, + [f"{EXPERT}.{n}" for n in PACKED] + [f"{ATTENTION}.weight", f"{SHARED_EXPERT}.weight"], + None, + ), + ( + ("weight", "weight_scale"), + [f"{EXPERT}.weight", f"{EXPERT}.weight_scale", f"{ATTENTION}.weight"], + None, + ), + (PACKED, [f"{p}.{n}" for p in (EXPERT, ATTENTION) for n in PACKED], "1 of them outside"), + (PACKED, [f"{EXPERT}.weight"] + [f"{ATTENTION}.{n}" for n in PACKED], "1 of them outside"), + ( + PACKED, + [f"{EXPERT}.{n}" for n in PACKED] + ["model.layers.1.mlp.experts.1.gate_proj.weight"], + "1 expert weights uncompressed", + ), + (PACKED, [f"{ATTENTION}.weight", f"{SHARED_EXPERT}.weight"], "compresses 0 weights"), + ], + ids=[ + "experts-alone", + "experts-alone-float-format", + "every-linear", + "all-but-the-experts", + "an-expert-left-uncompressed", + "nothing-compressed", + ], +) +def test_a_checkpoint_must_compress_its_experts_and_nothing_else(stored_as, names, match): + """The config cannot say which modules a class target such as Linear compressed, so the names + the checkpoint stores decide: anything else compressed, or an expert stored uncompressed, would + load newly initialized, without a quantizer to load it or a converter that reads it.""" + config = SimpleNamespace(text_config=SimpleNamespace()) + onfly._mark(config.text_config, stored_as) + + if match is None: + onfly.check_compressed_keys(config, names) + else: + with pytest.raises(NotImplementedError, match=match): + onfly.check_compressed_keys(config, names) + + class _RemoteTextConfig(PretrainedConfig): """Stands in for the remote-code DeepSeek-V3 config a VLM bundles as its ``text_config``.""" diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index f7beeefbd62..05b0e458d4e 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -560,12 +560,15 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( - ep_world_of_two, monkeypatch + ep_world_of_two, monkeypatch, tmp_path ): """A model spec that installs the on-the-fly decompressor (Kimi-K2.5's) marks the config; that - checkpoint passes the format check and loads without its quantization_config, at either level, + checkpoint passes the format check and loads without its quantization_config, at every level, which would build transformers' own compressed-tensors quantizer to compete with the - decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own.""" + decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own. + So is one that compresses more than its experts, by the names it stores, before it loads.""" + from safetensors.torch import save_file + from modelopt.torch.models import _onfly_decompress class _LoadedError(Exception): @@ -583,20 +586,33 @@ def from_pretrained(path, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + stored_as = ("weight_packed", "weight_scale") + expert, attention = "model.layers.1.mlp.experts.0.up_proj", "model.layers.1.self_attn.o_proj" + for name, compressed in (("experts", [expert]), ("more", [expert, attention])): + keys = [f"{module}.{n}" for module in compressed for n in stored_as] + (tmp_path / name).mkdir() + save_file({key: torch.ones(1) for key in keys}, str(tmp_path / name / "model.safetensors")) quantization_config = {"quant_method": "compressed-tensors"} text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) - config = SimpleNamespace(quantization_config=quantization_config, text_config=text) + llm = SimpleNamespace(quantization_config=quantization_config) + config = SimpleNamespace( + quantization_config=quantization_config, text_config=text, llm_config=llm + ) plan = DistributionPlan(2, ep_size=2) with pytest.raises(NotImplementedError, match="cannot dequantize"): - load_and_prepare_ep("org/model", plan, hf_config=config) - _onfly_decompress._mark(config, "pack-quantized") + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) + _onfly_decompress._mark(config, stored_as) + with pytest.raises(NotImplementedError, match="1 of them outside the experts"): + load_and_prepare_ep(str(tmp_path / "more"), plan, hf_config=config) + assert not loaded with pytest.raises(_LoadedError): - load_and_prepare_ep("org/model", plan, hf_config=config) + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) - assert not hasattr(loaded["config"], "quantization_config") - assert not hasattr(loaded["config"].text_config, "quantization_config") + for level in (loaded["config"], loaded["config"].text_config, loaded["config"].llm_config): + assert not hasattr(level, "quantization_config") assert config.quantization_config is text.quantization_config is quantization_config + assert llm.quantization_config is quantization_config # ------------------------------------------------------------------------ quant wrappers From b760f80694ce59ffb62ef644c201dc0ae121c6c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 03:53:43 +0000 Subject: [PATCH 14/16] Refuse decompression that rewrites no expert or narrows a stored float Three ways the on-the-fly decompression could still load a compressed checkpoint wrong without a word are now refused: - The conversion-mapping hook rewrote whatever per-expert converters it found, possibly none, and cleared the config's mark before rewriting. A marked model whose mapping has no per-expert converter to rewrite, whose experts would load newly initialized, is now refused, and the mark is cleared only once the converters are rewritten. - The load may cast a source tensor to the model's dtype before the decompression reads it; only integers are kept as stored. The expert-parallel load's check of what the checkpoint stores now also reads the dtypes of the compressed experts' tensors and refuses one bf16 cannot hold exactly (fp16, fp32, fp64), such as NVFP4's fp32 global scale. The Kimi-K2.5, K2.6, K2.7-Code and K3 checkpoints store int32, uint8 and bf16 there. - DecompressExperts pairs each expert's stored tensors by position, and now refuses stored params that disagree on how many experts they hold. The decompressor also finds its compressed-tensors config and mark on llm_config, as the rest of the expert-parallel load does, and the integer-keeping loader hook says why it stays process-wide. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 115 ++++++++++++++---- .../torch/utils/plugins/model_load_utils.py | 4 +- .../torch/models/test_onfly_decompress.py | 88 ++++++++++++-- tests/unit/torch/utils/test_ep_loader.py | 10 +- 4 files changed, 178 insertions(+), 39 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index fd953a6dddf..695fe6c609f 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -20,7 +20,8 @@ The math is the ``compressed_tensors`` library's, driven by the checkpoint's ``quantization_config``: ``compressor.compression_param_names(scheme)`` names the stored keys of one weight, and ``compressor.decompress(...)["weight"]`` dequantizes them, so every format the library registers -works. It needs ``compressed_tensors`` importable and transformers 5.8 or newer (earlier releases' +works whose stored floating-point tensors bf16 holds exactly (:func:`check_compressed_keys`). It +needs ``compressed_tensors`` importable and transformers 5.8 or newer (earlier releases' ``Concatenate`` looks up the sibling keys the decompression consumed). :func:`install_onfly_decompress` rewrites the per-expert converters ``get_model_conversion_mapping`` @@ -29,6 +30,8 @@ :func:`to_native_backbone_config` builds a VLM's transformers-native backbone config. """ +import os + import torch @@ -70,13 +73,25 @@ def to_native_backbone_config( return native +def _config_levels(hf_config) -> list: + """``hf_config`` and the nested language-model configs the expert-parallel load reads.""" + from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS + + nested = (getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS) + return [c for c in (hf_config, *nested) if c is not None] + + def _find_quantization_config(hf_config): - """Locate the compressed-tensors ``quantization_config`` (top-level or under ``text_config``) and - return it as a plain dict, or ``None`` for a config without one.""" - qc = getattr(hf_config, "quantization_config", None) - if qc is None: - tc = getattr(hf_config, "text_config", None) - qc = getattr(tc, "quantization_config", None) if tc is not None else None + """Locate the compressed-tensors ``quantization_config`` (top-level, else on the nested + language-model config) and return it as a plain dict, or ``None`` for a config without one.""" + qc = next( + ( + found + for c in _config_levels(hf_config) + if (found := getattr(c, "quantization_config", None)) is not None + ), + None, + ) if qc is None: return None if hasattr(qc, "to_dict"): @@ -189,6 +204,13 @@ def convert(self, input_dict, *args, **kwargs): lists[pn] = sib anchor_list = lists[param_names[0]] present = [pn for pn in param_names if pn in lists] + if any(len(lists[pn]) != len(anchor_list) for pn in present): + # Paired by position: an expert short of one would pair every later expert's + # weight with the next one's scale. + counts = ", ".join(f"{len(lists[pn])} {pn}" for pn in present) + raise ValueError( + f"DecompressExperts: {stem} stores {counts}, not one per expert" + ) # Non-local experts were sharded away at load (GroupedGemmParallel -> None); decompress # only the local experts this rank owns. Keeps peak memory at the per-rank shard. out[pat] = [ @@ -247,7 +269,12 @@ def _install_keep_int_native(): Without an hf_quantizer the loader casts every source to the target param dtype (bf16); that destroys the integer bit pattern of packed weights (`weight_packed`), shapes (`weight_shape`), zero-points and g_idx. Floats (the scale, - global scale, dense weights) still cast normally. + global scale, dense weights) still cast normally, so the expert-parallel load refuses a checkpoint whose compressed + experts store a float bf16 cannot hold exactly (:func:`check_compressed_keys`). + + Process-wide, unlike the converter rewrite: any later load keeps an integer tensor native where it would have cast + it to a floating-point dtype. Only a compressed checkpoint loaded without its quantizer asks for that cast: with a + quantizer the loader keeps integer tensors as stored, and an unquantized checkpoint stores none for a float param. """ import transformers.core_model_loading as cml @@ -276,9 +303,9 @@ def _mark(hf_config, stored_as): def is_marked(hf_config) -> bool: - """Whether :func:`install_onfly_decompress` marked ``hf_config`` (or its ``text_config``) for its hooks.""" - configs = (hf_config, getattr(hf_config, "text_config", None)) - return any(getattr(c, _MARK, None) is not None for c in configs if c is not None) + """Whether :func:`install_onfly_decompress` marked ``hf_config``, or its nested language-model + config, for its hooks.""" + return any(getattr(c, _MARK, None) is not None for c in _config_levels(hf_config)) def install_onfly_decompress(hf_config): @@ -338,20 +365,28 @@ def patched(*args, **kwargs): # Not the compressed checkpoint: rewritten, an uncompressed MoE's experts would match no # stored key and load newly initialized without an error. return convs - # from_pretrained loads from a copy of the marked config; the mark must not reach the config - # a later save writes, at either level. - for c in (config, getattr(config, "text_config", None)): - if hasattr(c, _MARK): - delattr(c, _MARK) - out = [] + out, rewritten = [], 0 for c in convs: # only the per-expert MoE converters (have `.weight` sources under `experts`) if isinstance(c, WeightConverter) and any( "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns ): out.append(_compressedify(c, compressor, scheme, param_names)) + rewritten += 1 else: out.append(c) + if not rewritten: + # The stock converters would read the experts' plain .weight keys, which the checkpoint + # does not store. + raise RuntimeError( + "install_onfly_decompress: the model's conversion mapping has no per-expert " + "converter to rewrite, so its compressed experts would load newly initialized." + ) + # from_pretrained loads from a copy of the marked config; the mark must not reach the config + # a later save writes, at any level. Cleared once the experts' converters are rewritten. + for c in _config_levels(config): + if hasattr(c, _MARK): + delattr(c, _MARK) return out setattr(patched, "_onfly_compression", compression) @@ -363,21 +398,32 @@ def patched(*args, **kwargs): ) -def check_compressed_keys(hf_config, tensor_names) -> None: - """Refuse a checkpoint that compresses anything but its experts, or leaves an expert uncompressed. +def check_compressed_keys(hf_config, ckpt_dir) -> None: + """Refuse a checkpoint the decompression would load wrong, by the tensors ``ckpt_dir`` stores. The hooks rewrite only the experts' converters, and the expert-parallel load drops the quantization_config that would load any other compressed weight: a weight compressed outside the experts would load newly initialized, and so would an expert stored uncompressed, which the - rewritten converters do not read. ``tensor_names`` are the checkpoint's. A weight counts as - compressed when it stores a name ``hf_config``'s mark holds other than ``weight``, which an - uncompressed weight is stored under too. + rewritten converters do not read. A weight counts as compressed when it stores a name + ``hf_config``'s mark holds other than ``weight``, which an uncompressed weight is stored under + too. + + The load may also cast a compressed weight's floating-point tensors to bf16, the dtype its + experts decompress to, before the decompression reads them: one that bf16 cannot hold exactly, + such as an fp32 global scale or an fp16 scale, would shift every weight it scales. Only the + tensors' names and dtypes are read. """ - configs = (hf_config, getattr(hf_config, "text_config", None)) - stored_as = next(getattr(c, _MARK) for c in configs if getattr(c, _MARK, None) is not None) + from safetensors import safe_open + + from modelopt.torch.utils.plugins.hf_checkpoint_utils import checkpoint_tensor_files + + stored_as = next( + getattr(c, _MARK) for c in _config_levels(hf_config) if getattr(c, _MARK, None) is not None + ) + tensor_files = checkpoint_tensor_files(ckpt_dir) compressed: set[str] = set() expert_weights: set[str] = set() - for name in tensor_names: + for name in tensor_files: stem, _, stored = name.rpartition(".") if stored in stored_as and stored != "weight": compressed.add(stem) @@ -392,3 +438,22 @@ def check_compressed_keys(hf_config, tensor_names) -> None: "on-the-fly decompression loads the experts alone, every one compressed, so the rest " "would load newly initialized. Decompress the checkpoint first." ) + by_file: dict[str, list[str]] = {} + for name, file in tensor_files.items(): + stem, _, stored = name.rpartition(".") + if stem in compressed and stored in stored_as: + by_file.setdefault(file, []).append(name) + inexact: set[str] = set() + for file, names in by_file.items(): + with safe_open(os.path.join(ckpt_dir, file), framework="pt") as f: + for name in names: + dtype = f.get_slice(name).get_dtype() + if dtype in ("F16", "F32", "F64"): + inexact.add(f"{name.rpartition('.')[2]} ({dtype})") + if inexact: + raise NotImplementedError( + f"The checkpoint's compressed experts store {', '.join(sorted(inexact))}, which bf16 " + "cannot hold exactly: the load may cast them to bf16, the dtype the experts decompress " + "to, before decompressing, shifting every weight they scale. Decompress the checkpoint " + "first." + ) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index e415999228a..b22a6e88d5e 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2569,10 +2569,10 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank()) if _decompressed_as_loaded(hf_config): # Only the experts decompress as they load, and the config cannot say what else the - # checkpoint compressed: its stored names are checked before anything loads. + # checkpoint compressed, or in which dtypes: what it stores is checked before anything loads. from modelopt.torch.models import _onfly_decompress - _onfly_decompress.check_compressed_keys(hf_config, checkpoint_tensor_files(ckpt_path)) + _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two # factorizations of the same world: diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 5f0bddeab13..a868afd5635 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -28,6 +28,7 @@ import pytest import torch +from safetensors.torch import save_file transformers = pytest.importorskip("transformers") @@ -180,7 +181,8 @@ def test_packed_experts_decompress_then_fuse_like_an_uncompressed_checkpoint(sto def test_only_owned_experts_and_stored_siblings_reach_the_compressor(): """Under EP the loader hands experts this rank does not own over as ``None``; they are skipped, so a rank decompresses only its own. A checkpoint may store fewer params than the scheme - declares (no ``weight_global_scale`` in Kimi-K3's mxfp4); missing ones are not passed.""" + declares (no ``weight_global_scale`` in Kimi-K3's mxfp4); missing ones are not passed, but + one stored for some experts and not others is refused.""" seen = [] class RecordingCompressor(_ScaledCompressor): @@ -207,6 +209,13 @@ def decompress(cls, state_dict, scheme): assert [t.dtype for t in experts] == [torch.bfloat16] * 2 assert [t[0, 0].item() for t in experts] == [0.5, 1.5] assert seen == [(["weight_packed", "weight_scale"], "scheme")] * 2 + # Paired by position, an expert missing one of them would shift every later expert's pairing. + short = { + "mlp.experts.*.down_proj.weight_packed$": packed, + "mlp.experts.*.down_proj.weight_scale$": [scale, scale], + } + with pytest.raises(ValueError, match="3 weight_packed, 2 weight_scale, not one per expert"): + decompress.convert(short) @pytest.mark.usefixtures("restore_loader_hooks") @@ -271,6 +280,24 @@ def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): ] +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_marked_model_with_no_expert_converter_to_rewrite_is_refused(monkeypatch): + """With nothing rewritten, the stock converters would look for plain ``.weight`` keys the + checkpoint does not store, and the experts would load newly initialized. The load is refused, + and the config, whose converters were not rewritten, keeps its mark.""" + renaming = WeightRenaming(r"experts\.(\d+)\.w1\.weight", r"experts.\1.gate_proj.weight") + monkeypatch.setattr(mu, "get_model_conversion_mapping", lambda model, **kw: [renaming]) + compression = (_ScaledCompressor, None, PACKED, "stub-format") + monkeypatch.setattr(onfly, "_read_compression", lambda cfg: compression) + config = SimpleNamespace() + onfly.install_onfly_decompress(config) + model = SimpleNamespace(config=copy.deepcopy(config)) + + with pytest.raises(RuntimeError, match="no per-expert converter to rewrite"): + mu.get_model_conversion_mapping(model) + assert onfly.is_marked(model.config) + + @pytest.mark.usefixtures("restore_loader_hooks") def test_installing_twice_wraps_the_loader_once(monkeypatch): """``prepare`` may run more than once per process. Installing again for the same checkpoint must @@ -417,6 +444,29 @@ def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters( assert not onfly.is_marked(model.config) +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_language_model_config_named_llm_config_is_read_as_text_config_is(monkeypatch): + """Some VLMs name their language-model config llm_config, which the expert-parallel load reads + as it reads text_config: marked there, the VLM loads through the rewritten converters, and the + mark is cleared.""" + monkeypatch.setattr( + mu, "get_model_conversion_mapping", lambda model, **kw: list(_stock_expert_converters()) + ) + compression = (_ScaledCompressor, None, PACKED, "stub-format") + monkeypatch.setattr(onfly, "_read_compression", lambda cfg: compression) + llm = SimpleNamespace() + onfly.install_onfly_decompress(llm) + model = SimpleNamespace(config=SimpleNamespace(llm_config=copy.deepcopy(llm))) + + out = mu.get_model_conversion_mapping(model) + + assert [conv.source_patterns[0] for conv in out] == [ + "mlp.experts.*.gate_proj.weight_packed$", + "mlp.experts.*.down_proj.weight_packed$", + ] + assert not onfly.is_marked(model.config) + + # ------------------------------------------------------------------------ reading the checkpoint @@ -434,16 +484,21 @@ def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters( SimpleNamespace(text_config=SimpleNamespace(quantization_config={"at": "text"})), {"at": "text"}, ), + ( + SimpleNamespace(llm_config=SimpleNamespace(quantization_config={"at": "llm"})), + {"at": "llm"}, + ), ( SimpleNamespace(quantization_config=SimpleNamespace(to_dict=lambda: {"at": "obj"})), {"at": "obj"}, ), ], - ids=["top-level-first", "text-config", "config-object"], + ids=["top-level-first", "text-config", "llm-config", "config-object"], ) def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_config(config, found): - """A VLM may carry its compressed-tensors config at the top level or on ``text_config``, as a - config object or a dict; it is always returned as a plain dict.""" + """A VLM may carry its compressed-tensors config at the top level or on its language-model + config (``text_config`` or ``llm_config``), as a config object or a dict; it is always returned + as a plain dict.""" assert onfly._find_quantization_config(config) == found @@ -484,6 +539,11 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( EXPERT = "model.layers.1.mlp.experts.0.gate_proj" ATTENTION = "model.layers.1.self_attn.q_proj" SHARED_EXPERT = "model.layers.1.mlp.shared_experts.up_proj" # a dense MLP, not the routed experts +STORED_DTYPES = { # as compressed-tensors stores them; every other tensor is bf16 here + "weight_packed": torch.int32, + "weight_shape": torch.int32, + "weight_global_scale": torch.float32, +} @pytest.mark.parametrize( @@ -507,6 +567,11 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( "1 expert weights uncompressed", ), (PACKED, [f"{ATTENTION}.weight", f"{SHARED_EXPERT}.weight"], "compresses 0 weights"), + ( + (*PACKED, "weight_global_scale"), + [f"{EXPERT}.{n}" for n in (*PACKED, "weight_global_scale")], + r"store weight_global_scale \(F32\), which bf16 cannot hold exactly", + ), ], ids=[ "experts-alone", @@ -515,20 +580,27 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( "all-but-the-experts", "an-expert-left-uncompressed", "nothing-compressed", + "an-fp32-scale", ], ) -def test_a_checkpoint_must_compress_its_experts_and_nothing_else(stored_as, names, match): +def test_a_checkpoint_must_compress_its_experts_and_nothing_else(tmp_path, stored_as, names, match): """The config cannot say which modules a class target such as Linear compressed, so the names the checkpoint stores decide: anything else compressed, or an expert stored uncompressed, would - load newly initialized, without a quantizer to load it or a converter that reads it.""" + load newly initialized, without a quantizer to load it or a converter that reads it. A float + stored for the experts must be one bf16 holds exactly: the load may cast it to bf16 first.""" + tensors = { + name: torch.zeros(1, dtype=STORED_DTYPES.get(name.rpartition(".")[2], torch.bfloat16)) + for name in names + } + save_file(tensors, str(tmp_path / "model.safetensors")) config = SimpleNamespace(text_config=SimpleNamespace()) onfly._mark(config.text_config, stored_as) if match is None: - onfly.check_compressed_keys(config, names) + onfly.check_compressed_keys(config, tmp_path) else: with pytest.raises(NotImplementedError, match=match): - onfly.check_compressed_keys(config, names) + onfly.check_compressed_keys(config, tmp_path) class _RemoteTextConfig(PretrainedConfig): diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 05b0e458d4e..729b86678a3 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -586,12 +586,14 @@ def from_pretrained(path, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - stored_as = ("weight_packed", "weight_scale") + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} expert, attention = "model.layers.1.mlp.experts.0.up_proj", "model.layers.1.self_attn.o_proj" for name, compressed in (("experts", [expert]), ("more", [expert, attention])): - keys = [f"{module}.{n}" for module in compressed for n in stored_as] + stored = { + f"{m}.{n}": torch.ones(1, dtype=d) for m in compressed for n, d in stored_as.items() + } (tmp_path / name).mkdir() - save_file({key: torch.ones(1) for key in keys}, str(tmp_path / name / "model.safetensors")) + save_file(stored, str(tmp_path / name / "model.safetensors")) quantization_config = {"quant_method": "compressed-tensors"} text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) llm = SimpleNamespace(quantization_config=quantization_config) @@ -602,7 +604,7 @@ def from_pretrained(path, **kwargs): with pytest.raises(NotImplementedError, match="cannot dequantize"): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) - _onfly_decompress._mark(config, stored_as) + _onfly_decompress._mark(config, tuple(stored_as)) with pytest.raises(NotImplementedError, match="1 of them outside the experts"): load_and_prepare_ep(str(tmp_path / "more"), plan, hf_config=config) assert not loaded From 3a15d133c93bedf9a25e544bb6be0ca9f135f08e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:25:04 +0000 Subject: [PATCH 15/16] Refuse an unmarked config with nothing to decompress once the hooks are in With the hooks already installed, a config without a compressed-tensors quantization_config returned as if its load would decompress. Only a config the installing prepare stripped, which carries the mark, is the same checkpoint again; an unmarked one would load its experts newly initialized, and is now refused as it is before the hooks go in. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 8 +++++++- tests/unit/torch/models/test_onfly_decompress.py | 11 ++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 695fe6c609f..aeb2c5bc214 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -338,7 +338,13 @@ def install_onfly_decompress(hf_config): if installed is not None: if compression is None: # A config its prepare dropped quantization_config from after installing (Kimi-K3's), - # already marked, may be prepared again: nothing new to decompress. + # already marked, may be prepared again: nothing new to decompress. An unmarked one would + # load its experts through the stock converters, newly initialized. + if not is_marked(hf_config): + raise ValueError( + "install_onfly_decompress: config has no compressed-tensors " + "quantization_config and is not marked: nothing would decompress it" + ) return if compression != installed: raise RuntimeError( diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index a868afd5635..7c20093c7be 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -335,7 +335,8 @@ def test_a_checkpoint_compressed_differently_needs_a_process_of_its_own(monkeypa """The hooks are process-wide: decompressing a second checkpoint with the first one's scheme and stored params would load it wrong without a word, so its install is refused and the hooks stay as they are. A config the installing prepare already stripped of its quantization_config is the - first checkpoint again; one whose quantization_config fails to parse is not, and its error + first checkpoint again; one with none that was never marked is refused, since nothing would + decompress its experts; one whose quantization_config fails to parse is not, and its error goes through.""" _, down = _stock_expert_converters() compressions = { @@ -355,9 +356,13 @@ def read(cfg): ) monkeypatch.setattr(onfly, "_read_compression", read) - onfly.install_onfly_decompress(SimpleNamespace(checkpoint="a")) + first = SimpleNamespace(checkpoint="a") + onfly.install_onfly_decompress(first) installed = mu.get_model_conversion_mapping - onfly.install_onfly_decompress(SimpleNamespace(checkpoint=None)) + first.checkpoint = None # its prepare dropped the quantization_config after installing + onfly.install_onfly_decompress(first) + with pytest.raises(ValueError, match="is not marked"): + onfly.install_onfly_decompress(SimpleNamespace(checkpoint=None)) with pytest.raises(RuntimeError, match="needs a process of its own"): onfly.install_onfly_decompress(SimpleNamespace(checkpoint="b")) with pytest.raises(ValueError, match="validation error"): From 49a43a2ea3c3576aaec9ddb043007be4cec55316 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:25:04 +0000 Subject: [PATCH 16/16] Rewrite only converters under .experts., not a shared expert's The filter matched "experts" anywhere in a source, so a converter fusing a shared expert's .weight tensors (mlp.shared_experts.*) would be rewritten to read compressed keys the checkpoint does not store, and the shared MLP would load newly initialized. It now matches .experts., as check_compressed_keys does. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 5 ++-- .../torch/models/test_onfly_decompress.py | 23 ++++++++++++++----- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index aeb2c5bc214..190cb96cfd8 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -373,9 +373,10 @@ def patched(*args, **kwargs): return convs out, rewritten = [], 0 for c in convs: - # only the per-expert MoE converters (have `.weight` sources under `experts`) + # only the per-expert MoE converters (have `.weight` sources under `.experts.`, which + # leaves a shared expert's `shared_experts.` out, as check_compressed_keys does) if isinstance(c, WeightConverter) and any( - "experts" in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns + ".experts." in p and p.endswith((".weight", ".weight$")) for p in c.source_patterns ): out.append(_compressedify(c, compressor, scheme, param_names)) rewritten += 1 diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 7c20093c7be..61119b5b44d 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -41,7 +41,13 @@ import transformers.modeling_utils as mu from transformers import DeepseekV3Config, PretrainedConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping -from transformers.core_model_loading import Chunk, MergeModulelist, WeightConverter, WeightRenaming +from transformers.core_model_loading import ( + Chunk, + Concatenate, + MergeModulelist, + WeightConverter, + WeightRenaming, +) from modelopt.torch.models import _onfly_decompress as onfly @@ -242,8 +248,8 @@ def test_integer_checkpoint_tensors_keep_their_dtype_on_load(): @pytest.mark.usefixtures("restore_loader_hooks") def test_install_rewrites_only_the_per_expert_weight_converters(monkeypatch): """The loader's conversion list also holds renamings and other converters; only converters that - fuse per-expert ``.weight`` tensors are rewritten, and from_pretrained's arguments reach the - stock mapping unchanged.""" + fuse per-expert ``.weight`` tensors are rewritten, not a shared expert's, and from_pretrained's + arguments reach the stock mapping unchanged.""" renaming = WeightRenaming(r"experts\.(\d+)\.w1\.weight", r"experts.\1.gate_proj.weight") qkv = WeightConverter( "self_attn.qkv_proj.weight", @@ -253,8 +259,13 @@ def test_install_rewrites_only_the_per_expert_weight_converters(monkeypatch): expert_bias = WeightConverter( "mlp.experts.*.down_proj.bias", "mlp.experts.down_proj_bias", [MergeModulelist(dim=0)] ) + shared_expert = WeightConverter( + ["mlp.shared_experts.gate_proj.weight", "mlp.shared_experts.up_proj.weight"], + "mlp.shared_experts.gate_up_proj", + [Concatenate(dim=0)], + ) gate_up, down = _stock_expert_converters() - stock = [renaming, qkv, expert_bias, gate_up, down] + stock = [renaming, qkv, expert_bias, shared_expert, gate_up, down] calls = [] def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): @@ -273,8 +284,8 @@ def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): assert calls == [(model, {"old": "new"}, "quantizer", False)] assert len(out) == len(stock) - assert all(kept is original for kept, original in zip(out[:3], stock[:3])) - assert [conv.source_patterns[0] for conv in out[3:]] == [ + assert all(kept is original for kept, original in zip(out[:4], stock[:4])) + assert [conv.source_patterns[0] for conv in out[4:]] == [ "mlp.experts.*.gate_proj.weight_packed$", "mlp.experts.*.down_proj.weight_packed$", ]