From 3b2ea90d9d122c8c4f8fc27a8491b0b82fa6de90 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:00:51 +0000 Subject: [PATCH 01/45] Decompress compressed-tensors MoE experts 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. install_onfly_decompress rewrites the per-expert converters of the model built from the config it marks: each .weight source becomes the stored keys (weight_packed, weight_scale, ...), and a DecompressExperts op ahead of the merge dequantizes each owned expert to bf16 with the compressed_tensors library, driven by the checkpoint's own quantization_config. Non-local experts arrive as None and are skipped. A checkpoint must compress its experts alone, in one config group, in floats bf16 holds exactly; check_compressed_keys checks the stored names and dtypes. The hooks are process-wide but rewrite only a marked config's converters. to_native_backbone_config rebuilds a VLM's bundled text config as a native transformers backbone config, carrying the mark. utils/moe.py gains _SUB_CONFIG_ATTRS, the nested configs the expert- parallel load reads, word for word as the model specs' change adds it. CI's unit sessions install compressed-tensors, for the tests against the real library. Moved out of the on-the-fly decompression change, whose history is archived. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 475 +++++++++++ noxfile.py | 7 +- .../torch/models/test_onfly_decompress.py | 764 ++++++++++++++++++ 3 files changed, 1245 insertions(+), 1 deletion(-) create mode 100644 modelopt/torch/models/_onfly_decompress.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 000000000000..42cc1075c9d4 --- /dev/null +++ b/modelopt/torch/models/_onfly_decompress.py @@ -0,0 +1,475 @@ +# 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. + +"""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 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`` +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 os + +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) + native = CONFIG_MAPPING[model_type].from_dict(d) + 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, stored_as) + 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, 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"): + 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 checkpoint's one config + 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) + if quantization_config is None: + return None + try: + from compressed_tensors.compressors.base import BaseCompressor + from compressed_tensors.quantization import QuantizationConfig + except ImportError as e: + # No pip extra of ModelOpt installs it: only a compressed checkpoint needs it. + raise ImportError( + "On-the-fly decompression of a compressed-tensors checkpoint needs the " + "compressed-tensors package. Install it with: pip install compressed-tensors" + ) from e + + qconfig = QuantizationConfig.model_validate(quantization_config) + groups = qconfig.config_groups + 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)}); on-the-fly decompression handles a single group, " + "compressing only the experts. Decompress the checkpoint first." + ) + [scheme] = groups.values() + # 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 " + 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) + 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, 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] + 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] + 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] = [ + 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 + ] + 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 + + +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) # 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( + source_patterns=new_sources, + target_patterns=tgt, + operations=[DecompressExperts(), *conv.operations], + ) + # 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(): + """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, 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 + + if getattr(cml._materialize_copy, "_onfly_patched", False): + return + orig = cml._materialize_copy + + # Any other argument passes through, positionally as the loader passes the ones it has, or by + # keyword: the patch sits on every later load's copy path. + def _materialize_copy(tensor, device=None, dtype=None, *args, **kwargs): + t = tensor[...] + 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 + return orig(t, device, dtype, *args, **kwargs) + + setattr(_materialize_copy, "_onfly_patched", True) + cml._materialize_copy = _materialize_copy + + +# 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, stored_as): + """Mark ``hf_config``'s load for the installed hooks; ``from_pretrained``'s copy of it keeps the mark.""" + setattr(hf_config, _MARK, stored_as) + + +def is_marked(hf_config) -> bool: + """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): + """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. + 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 marks that config too, and installing for a checkpoint compressed differently raises. Only a + 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, 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. + 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: + # A config its prepare dropped quantization_config from after installing (Kimi-K3's), + # 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( + 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" + ) + # The same checkpoint through a config of its own, which its load must find marked too. + _mark(hf_config, installed[2]) + return + 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 + _install_keep_int_native() + orig = mu.get_model_conversion_mapping + + def patched(*args, **kwargs): + convs = orig(*args, **kwargs) + config = getattr(args[0] if args else kwargs["model"], "config", 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 + out, rewritten = [], 0 + for c in convs: + # 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 + ): + 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) + mu.get_model_conversion_mapping = patched + _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, 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. 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. + """ + 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_files: + 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." + ) + 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/noxfile.py b/noxfile.py index 4fa9fc107afa..bcf70c6ac32a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -66,7 +66,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) 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 000000000000..33e62eae12fe --- /dev/null +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -0,0 +1,764 @@ +# 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 +from safetensors.torch import save_file + +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +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 transformers import DeepseekV3Config, PretrainedConfig +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import ( + Chunk, + Concatenate, + 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 ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + monkeypatch.setattr(owner, name, getattr(owner, name)) + + +class _ScaledCompressor: + """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): + stored = state_dict.get("weight_packed", state_dict.get("weight")) + return {"weight": stored.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_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)], + ) + 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) + + 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(): + """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 + + +@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) + 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." + 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, 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") + 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, but + one stored for some experts and not others is refused.""" + 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 + # 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") +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 + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_copy_parameter_a_release_adds_passes_through(monkeypatch): + """The patch sits on every later load's copy path, so a parameter a release adds reaches the + loader's copy, passed positionally, as the loader passes the ones it has, or by keyword.""" + seen = [] + + def release_copy(tensor, device=None, dtype=None, non_blocking=False): + seen.append(non_blocking) + return tensor + + monkeypatch.setattr(cml, "_materialize_copy", release_copy) + onfly._install_keep_int_native() + + cml._materialize_copy(torch.ones(2), "cpu", torch.bfloat16, True) + cml._materialize_copy(torch.ones(2), "cpu", torch.bfloat16, non_blocking=True) + + assert seen == [True, True] + + +# ------------------------------------------------------------------------ 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, 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", + ["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)] + ) + 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, shared_expert, gate_up, down] + calls = [] + + def stock_mapping(model, key_mapping=None, hf_quantizer=None, add_legacy=True): + calls.append((model, key_mapping, hf_quantizer, add_legacy)) + 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) + config = SimpleNamespace() + onfly.install_onfly_decompress(config) + model = SimpleNamespace(config=config) + + # 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", False)] + assert len(out) == len(stock) + 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$", + ] + + +@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 + 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 = [] + + 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, cml._materialize_copy + + onfly.install_onfly_decompress(SimpleNamespace()) + installed = hooks() + again = SimpleNamespace() + onfly.install_onfly_decompress(again) + + assert hooks() == installed + assert len(reads) == 2 + [conv] = mu.get_model_conversion_mapping(SimpleNamespace(config=again)) + 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; 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 = { + "a": (_ScaledCompressor, None, PACKED, "stub-format"), + "b": (_ScaledCompressor, None, (*PACKED, "weight_zero_point"), "stub-format"), + None: None, + } + + def read(cfg): + 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( + mu, "get_model_conversion_mapping", lambda model, **kw: [copy.deepcopy(down)] + ) + monkeypatch.setattr(onfly, "_read_compression", read) + + first = SimpleNamespace(checkpoint="a") + onfly.install_onfly_decompress(first) + installed = mu.get_model_conversion_mapping + 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"): + onfly.install_onfly_decompress(SimpleNamespace(checkpoint="malformed")) + + 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): + 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()) + ) + 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") +@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 loaded == "another-model": + config = DeepseekV3Config(num_hidden_layers=1) + 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] == [ + f"mlp.experts.*.gate_proj.{stored}", + f"mlp.experts.*.down_proj.{stored}", + ] + 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) + + +@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 + + +@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"}, + ), + ( + 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", "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 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 + + +@pytest.mark.usefixtures("restore_loader_hooks") +@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. 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, + WeightConverter.add_tensor, + WeightConverter.materialize_tensors, + cml._materialize_copy, + ) + + with pytest.raises(error, match=match): + onfly.install_onfly_decompress(config) + + assert ( + mu.get_model_conversion_mapping, + WeightConverter.add_tensor, + WeightConverter.materialize_tensors, + cml._materialize_copy, + ) == 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 +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( + ("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"), + ( + (*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", + "experts-alone-float-format", + "every-linear", + "all-but-the-experts", + "an-expert-left-uncompressed", + "nothing-compressed", + "an-fp32-scale", + ], +) +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. 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, tmp_path) + else: + with pytest.raises(NotImplementedError, match=match): + onfly.check_compressed_keys(config, tmp_path) + + +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", "error", "match"), + [ + ( + {"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", "one-group-not-the-experts", "none"], +) +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, 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(error, match=match): + onfly._read_compression(config) + + +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", 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 + 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 7cccba27ba98bced42f72eeec39d843ab0dfc583 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:01:06 +0000 Subject: [PATCH 02/45] Import _SUB_CONFIG_ATTRS at module level in the decompressor modelopt.torch.utils.moe imports nothing from the model registry at module level, so the decompressor can import the constant at the top of the file: importing the registry, moe, the decompressor or a Kimi adapter on its own shows no cycle. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 42cc1075c9d4..464d1699c5a8 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -34,6 +34,8 @@ import torch +from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS + def to_native_backbone_config( vlm_config, @@ -75,8 +77,6 @@ def to_native_backbone_config( 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] From 9ac1abfb5e9bbc4592d20ac89dacdb0c4aa7f4f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:01:10 +0000 Subject: [PATCH 03/45] Install compressed-tensors in every unit session The torch 2.8 and 2.9 legs it was kept out of are gone from the unit sessions, so the condition and the comment about older legs are dead. Signed-off-by: Shengliang Xu --- noxfile.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/noxfile.py b/noxfile.py index bcf70c6ac32a..f4a8ecb726a8 100644 --- a/noxfile.py +++ b/noxfile.py @@ -66,12 +66,11 @@ 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.""" - # 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) + # compressed-tensors runs the real-library tests of on-the-fly expert decompression; 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. + session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", "compressed-tensors") tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) From 6972dc598575e14f50351f6f1fd484718600e020 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:42:52 +0000 Subject: [PATCH 04/45] Refuse an unmarked config by name when checking the stored tensors check_compressed_keys read the names the mark holds through a bare next(), so a config the decompressor never marked, or one whose mark the hook has already cleared, raised an unexplained StopIteration. stored_names() reads them, and refuses an unmarked config with a ValueError that says what is missing. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 12 +++++++++--- tests/unit/torch/models/test_onfly_decompress.py | 9 +++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 464d1699c5a8..81c9283a3b43 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -317,6 +317,14 @@ def is_marked(hf_config) -> bool: return any(getattr(c, _MARK, None) is not None for c in _config_levels(hf_config)) +def stored_names(hf_config) -> tuple[str, ...]: + """The names each compressed weight of ``hf_config``'s checkpoint is stored under, from its mark.""" + for c in _config_levels(hf_config): + if (stored_as := getattr(c, _MARK, None)) is not None: + return stored_as + raise ValueError("The config is not marked for on-the-fly decompression: install it first.") + + 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. @@ -433,9 +441,7 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: 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 - ) + stored_as = stored_names(hf_config) tensor_files = checkpoint_tensor_files(ckpt_dir) compressed: set[str] = set() expert_weights: set[str] = set() diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 33e62eae12fe..a25802b0ca28 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -638,6 +638,15 @@ def test_a_checkpoint_must_compress_its_experts_and_nothing_else(tmp_path, store onfly.check_compressed_keys(config, tmp_path) +def test_an_unmarked_config_has_no_stored_names_to_check(tmp_path): + """The check reads what the checkpoint stores by the names the mark holds, so a config the + decompressor never marked is refused, by name, rather than with a bare StopIteration.""" + save_file({f"{EXPERT}.weight_packed": torch.zeros(1)}, str(tmp_path / "model.safetensors")) + + with pytest.raises(ValueError, match="not marked for on-the-fly decompression"): + onfly.check_compressed_keys(SimpleNamespace(text_config=SimpleNamespace()), tmp_path) + + class _RemoteTextConfig(PretrainedConfig): """Stands in for the remote-code DeepSeek-V3 config a VLM bundles as its ``text_config``.""" From 69d0186269c05f10ecf0f75a8278fff41d9f6579 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:43:37 +0000 Subject: [PATCH 05/45] Decompress with the one config group's own format, and refuse one with no compressor compressed-tensors writes "mixed-precision" at a config's top level once a group names its own format, and the registry lookup of that name failed inside the library. The one group's own format is read first, and a format the library registers no compressor for is refused here, with this module's "decompress the checkpoint first". Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 8 ++++++- .../torch/models/test_onfly_decompress.py | 24 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 81c9283a3b43..81c58c7db32a 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -151,8 +151,14 @@ def _read_compression(hf_config): 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 + # A group's own format wins: the config's top level then says "mixed-precision". + fmt = getattr(scheme, "format", None) or qconfig.format fmt = fmt.value if hasattr(fmt, "value") else str(fmt) + if fmt not in BaseCompressor.registered_names(): + raise NotImplementedError( + f"install_onfly_decompress: compressed-tensors registers no compressor for the format " + f"{fmt!r}, so on-the-fly decompression cannot read it. Decompress the checkpoint first." + ) compressor = BaseCompressor.get_value_from_registry(fmt) param_names = tuple(compressor.compression_param_names(scheme)) return compressor, scheme, param_names, fmt diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index a25802b0ca28..e966b3504d14 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -745,6 +745,30 @@ def test_a_checkpoint_whose_experts_are_not_its_one_compressed_group_is_refused( onfly._read_compression(config) +@pytest.mark.parametrize( + ("top", "group", "match"), + [("mixed-precision", "pack-quantized", None), ("fakequant", None, "no compressor")], + ids=["the-group-names-it", "none-registered"], +) +def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group, match): + """compressed-tensors writes "mixed-precision" at the top level once a group names its own + format; the one group's format is the one to decompress with. A format nothing registers is + refused by this module, rather than failing inside the library's registry.""" + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) + experts = _int4(True, EXPERTS) + if group is not None: + experts["format"] = group + config = SimpleNamespace(quantization_config={**_pack_quantized(moe=experts), "format": top}) + + if match is None: + assert onfly._read_compression(config)[3] == group + else: + with pytest.raises(NotImplementedError, match=match): + onfly._read_compression(config) + + 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).""" From 694282c700308dbaab835df8847e9f1b7ccc8595 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:44:39 +0000 Subject: [PATCH 06/45] Refuse on-the-fly decompression for a load in a dtype other than bf16 The experts decompress to bf16, and the loader sets each parameter from what the conversion hands it, without a cast: a checkpoint loading in fp16 or fp32 would end up with bf16 experts beside the rest of the model, failing far from here. install_onfly_decompress refuses such a config before anything is patched. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 8 ++++++++ tests/unit/torch/models/test_onfly_decompress.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 81c58c7db32a..d1fecbd56ce7 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -351,6 +351,14 @@ def install_onfly_decompress(hf_config): 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") + dtype = getattr(hf_config, "dtype", None) + if dtype is not None and str(dtype).removeprefix("torch.") not in ("bfloat16", "auto"): + # The experts would load in bf16 beside the rest in this dtype: the loader sets each + # parameter from what the conversion hands it, without a cast. + raise NotImplementedError( + f"install_onfly_decompress: the experts decompress to bfloat16, but this config " + f"loads in {dtype}. Load the checkpoint in bfloat16." + ) 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. diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index e966b3504d14..d37684156bf9 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -543,15 +543,21 @@ def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_conf [ (SimpleNamespace(text_config=SimpleNamespace()), ValueError, "quantization_config"), ({"quantization_config": {"format": "pack-quantized"}}, TypeError, "not a dict"), + ( + SimpleNamespace(dtype=torch.float16, quantization_config={"format": "pack-quantized"}), + NotImplementedError, + "loads in torch.float16", + ), ], - ids=["nothing-to-decompress", "raw-dict"], + ids=["nothing-to-decompress", "raw-dict", "not-bf16"], ) 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. Nor for a raw - dict, which cannot carry the mark its load would be found by.""" + dict, which cannot carry the mark its load would be found by, or for a load in a dtype other + than the bf16 the experts decompress to: the loader would set them beside the rest uncast.""" monkeypatch.setitem(sys.modules, "compressed_tensors", None) hooks = ( mu.get_model_conversion_mapping, From 0d5a7dcc4b787112310f2178723ae4394c930f6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:47:43 +0000 Subject: [PATCH 07/45] Compress the decompressor's docstrings and comments Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 132 ++++++--------------- 1 file changed, 37 insertions(+), 95 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index d1fecbd56ce7..dadee897a328 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -15,19 +15,10 @@ """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 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`` -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. +Each rank decompresses only its own experts into the fused-experts layout, so expert-parallel PTQ +runs on the original compressed checkpoint. The ``compressed_tensors`` library does the math, +driven by the checkpoint's ``quantization_config``. Needs transformers 5.8: earlier releases' +``Concatenate`` looks up the sibling keys the decompression consumed. """ import os @@ -42,18 +33,11 @@ def to_native_backbone_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. + """The transformers-native backbone config for a VLM config's bundled ``text_config``. + + Drops the remote ``auto_map`` and the stale ``quantization_config``, and names the native model + type and class (hf_ptq's VLM detection reads ``architectures``). The decompression mark carries + over. Defaults target the Kimi DeepSeek-V3 backbone. """ from transformers.models.auto.configuration_auto import CONFIG_MAPPING @@ -102,14 +86,10 @@ def _find_quantization_config(hf_config): 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 checkpoint's one config - 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.""" + """``(compressor, scheme, param_names, format)`` for the checkpoint's compressed-tensors config. + + ``param_names`` are the keys one weight is stored under. The one config group must target the + experts; several groups are refused. ``None`` for a config with no ``quantization_config``.""" # 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) @@ -167,10 +147,8 @@ def _read_compression(hf_config): def _make_op_class(compressor, scheme, param_names): """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. + It subclasses ``MergeModulelist`` only so the loader shards experts at load, handing it ``None`` + for each expert another rank owns; the merge later in the chain still does the stacking. """ from transformers.core_model_loading import MergeModulelist @@ -178,23 +156,10 @@ def _make_op_class(compressor, scheme, param_names): 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. - """ + """Decompress each owned expert's stored params to bf16, under the first stored key.""" def __init__(self): - super().__init__( - dim=0 - ) # dim unused (convert overridden); needed so isinstance -> shard_index + super().__init__(dim=0) # dim unused: convert is overridden @torch.no_grad() def convert(self, input_dict, *args, **kwargs): @@ -206,11 +171,8 @@ def convert(self, input_dict, *args, **kwargs): 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). + # A checkpoint may store a subset of the scheme's params (Kimi-K3's mxfp4 stores no + # weight_global_scale): only those present go to the compressor. for pn, suf in others: sib = input_dict.pop(stem + suf + marker, None) if sib is not None: @@ -224,8 +186,7 @@ def convert(self, input_dict, *args, **kwargs): 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. + # Experts other ranks own arrive as None: only this rank's are decompressed. out[pat] = [ compressor.decompress({pn: lists[pn][i] for pn in present}, scheme)[ "weight" @@ -245,10 +206,10 @@ def convert(self, input_dict, *args, **kwargs): def _compressedify(conv, compressor, scheme, param_names): - """Rewrite a per-expert WeightConverter (.weight sources) into a compressed-aware one. + """Rewrite a per-expert converter to read the stored keys, with ``DecompressExperts`` first. - 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. + Each ``.weight`` source becomes its stored keys, grouped by key with the anchor first, which + keeps gate before up for the merge. """ from transformers.core_model_loading import WeightConverter @@ -278,16 +239,11 @@ def _compressedify(conv, compressor, scheme, param_names): def _install_keep_int_native(): - """Patch _materialize_copy so INTEGER source tensors are never cast to a FLOATING dtype during load. + """Patch ``_materialize_copy`` so an integer source tensor is never cast to a float dtype. - 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, 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. + Without a quantizer the loader casts every source to the parameter's dtype, which would destroy + packed weights, shapes and zero-points. Process-wide: only a compressed checkpoint loaded without + its quantizer asks for such a cast. """ import transformers.core_model_loading as cml @@ -332,19 +288,12 @@ def stored_names(hf_config) -> tuple[str, ...]: 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. - 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 marks that config too, and installing for a checkpoint compressed differently raises. Only a - 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, and the expert-parallel load checks the checkpoint against it - (:func:`check_compressed_keys`).""" + """Make ``from_pretrained`` decompress the experts of the checkpoint ``hf_config`` describes. + + Call before ``from_pretrained``. The hooks are process-wide but rewrite only the converters of a + model built from a config they marked: one compressed checkpoint per process, installing again + for it marks that config too, and one compressed differently is refused. + """ from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import require_transformers @@ -439,17 +388,10 @@ def patched(*args, **kwargs): 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. 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. + Only the experts decompress, and the load drops the config any other compressed weight would need: + the experts alone must be compressed, every one of them. The load may cast their float tensors to + bf16 before decompressing, so one bf16 cannot hold exactly (an fp32 scale) is refused too. Reads + only the tensors' names and dtypes. """ from safetensors import safe_open From 5017b8bffe7ec0352d39ce78a8ac730daddb81e3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:38 +0000 Subject: [PATCH 08/45] Reverse a decompressing converter as the fusion it rewrote from_pretrained keeps the converters it ran as model._weight_conversions, and a save reverses them: transformers' save_pretrained, and ModelOpt's export, which names the expanded experts back after the checkpoint. The rewritten converter reversed into its stored names (weight_packed, weight_scale, ...), so the export's expert-leaf arity check refused it and the export fell back to in-memory names: gate_proj/up_proj/down_proj for a checkpoint whose experts are w1/w3/w2, Kimi-K3's among them. The rewritten converter is now a WeightConverter subclass that reverses as the converter it rewrote: the decompression is load-only. Real from_pretrained loads of a tiny compressed Mixtral, INT4 and MXFP4, replace the converter-only round trip: the experts load exactly, save back as the checkpoint's own w1/w3/w2 weights, and export under those names. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 30 ++++- .../torch/models/test_onfly_decompress.py | 123 ++++++++++++++---- 2 files changed, 124 insertions(+), 29 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index dadee897a328..5f36f3303092 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -21,6 +21,7 @@ ``Concatenate`` looks up the sibling keys the decompression consumed. """ +import functools import os import torch @@ -205,14 +206,34 @@ def convert(self, input_dict, *args, **kwargs): return DecompressExperts +@functools.cache +def _decompressing_converter_class(): + """The ``WeightConverter`` subclass a rewritten converter is, reversing as the one it rewrote.""" + from transformers.core_model_loading import WeightConverter + + class DecompressingConverter(WeightConverter): + """Reads the stored keys, and reverses as the converter it rewrote. + + A save reverses the converters ``from_pretrained`` kept as ``model._weight_conversions`` + (transformers' ``save_pretrained``, and ModelOpt's export naming the experts back after + the checkpoint): the decompression is load-only, so the experts go back under their + checkpoint's ``.weight`` names, as the original fusion has them. + """ + + __slots__ = ("_reversible",) + + def reverse_transform(self): + return self._reversible.reverse_transform() + + return DecompressingConverter + + def _compressedify(conv, compressor, scheme, param_names): """Rewrite a per-expert converter to read the stored keys, with ``DecompressExperts`` first. Each ``.weight`` source becomes its stored keys, grouped by key with the anchor first, which - keeps gate before up for the merge. + keeps gate before up for the merge. The rewritten converter reverses as ``conv`` does. """ - 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: @@ -225,11 +246,12 @@ def _compressedify(conv, compressor, scheme, param_names): 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( + new = _decompressing_converter_class()( source_patterns=new_sources, target_patterns=tgt, operations=[DecompressExperts(), *conv.operations], ) + new._reversible = conv # 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"): diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index d37684156bf9..24730d76b5fe 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -17,18 +17,20 @@ 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. +with a stand-in compressor, so they need neither a checkpoint nor ``compressed_tensors``. The tests +that use the real ``compressed_tensors`` registry, among them real ``from_pretrained`` loads of tiny +compressed checkpoints, skip when it cannot be imported. """ import copy +import json import sys from concurrent.futures import Future from types import SimpleNamespace import pytest import torch -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file transformers = pytest.importorskip("transformers") @@ -39,7 +41,13 @@ import transformers.core_model_loading as cml import transformers.modeling_utils as mu -from transformers import DeepseekV3Config, PretrainedConfig +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + DeepseekV3Config, + MixtralConfig, + PretrainedConfig, +) from transformers.conversion_mapping import get_checkpoint_conversion_mapping from transformers.core_model_loading import ( Chunk, @@ -775,29 +783,94 @@ def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group onfly._read_compression(config) -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", reason="needs compressed-tensors (CI's unit session installs it)" +# The compressions the published Kimi checkpoints use: INT4 (Kimi-K2.5) and MXFP4 (Kimi-K3), each +# with the values its format holds exactly, and a scale it holds exactly. +_INT4 = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} +_FP4_E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6]) +_FORMATS = { + "pack-quantized": (_INT4, torch.arange(-8, 8) * 0.5, 0.5), + "mxfp4-pack-quantized": ( + {**_INT4, "type": "float", "scale_dtype": "torch.uint8"}, + torch.cat([_FP4_E2M1, -_FP4_E2M1[1:]]), + 1.0, + ), +} + + +def _compressed_mixtral(path, fmt): + """Save a tiny Mixtral whose experts ``fmt`` compresses; return its config and dense experts. + + Mixtral's checkpoint names its experts ``w1``/``w3``/``w2``, which the model fuses and the export + names back: names a save gets wrong would show. + """ + weights, values, scale = _FORMATS[fmt] + config = MixtralConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=2, + num_experts_per_tok=1, + tie_word_embeddings=False, + ) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(path) + quantization_config = { + "quant_method": "compressed-tensors", + "format": fmt, + "config_groups": {"experts": {"targets": ["Linear"], "weights": weights, "format": fmt}}, + } + compressor, scheme, _, _ = onfly._read_compression( + SimpleNamespace(quantization_config=quantization_config) ) - config = SimpleNamespace(quantization_config=_pack_quantized(experts=_int4(True, ["Linear"]))) - compressor, scheme, names, _ = onfly._read_compression(config) - n_experts, inter, hidden = 3, 4, 64 + state = load_file(str(path / "model.safetensors")) 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() + dense = {} + for key in sorted(k for k in list(state) if ".experts." in k): # ...experts..w1.weight + rows, cols = state.pop(key).shape + dense[key] = values[torch.randint(0, len(values), (rows, cols), generator=gen)] + scales = torch.full((rows, cols // 32), scale) + stored = compressor.compress({"weight": dense[key], "weight_scale": scales}, scheme) + for name, tensor in stored.items(): # floats as bf16, as the Kimi checkpoints store them + state[key.removesuffix("weight") + name] = ( + tensor.bfloat16() if tensor.is_floating_point() else tensor + ) + save_file(state, str(path / "model.safetensors"), metadata={"format": "pt"}) + config_json = json.loads((path / "config.json").read_text()) + config_json["quantization_config"] = quantization_config + (path / "config.json").write_text(json.dumps(config_json)) + return AutoConfig.from_pretrained(path), dense - 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() +@pytest.mark.usefixtures("restore_loader_hooks") +@pytest.mark.parametrize("fmt", sorted(_FORMATS)) +def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp_path, fmt): + """A real ``from_pretrained`` through the installed decompressor: the fused experts hold the + checkpoint's weights exactly. The converters the load kept reverse as the original fusion, so a + save writes the experts back as the checkpoint's own ``w1``/``w3``/``w2`` weights, decompressed, + and the export names its expanded experts after them.""" + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" ) + from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper + + config, dense = _compressed_mixtral(tmp_path, fmt) + onfly.install_onfly_decompress(config) + del config.quantization_config # as the expert-parallel load does: no quantizer competes + + model = AutoModelForCausalLM.from_pretrained(tmp_path, config=config, dtype=torch.bfloat16) + + experts, stored = model.model.layers[0].mlp.experts, "model.layers.0.block_sparse_moe.experts" + for e in range(2): + gate_up = torch.cat([dense[f"{stored}.{e}.w1.weight"], dense[f"{stored}.{e}.w3.weight"]]) + assert torch.equal(experts.gate_up_proj[e], gate_up.bfloat16()) + assert torch.equal(experts.down_proj[e], dense[f"{stored}.{e}.w2.weight"].bfloat16()) + saved = cml.revert_weight_conversion(model, model.state_dict()) + for key, weight in dense.items(): + assert torch.equal(saved[key], weight.bfloat16()) + export_name = build_reverse_name_mapper(model, tensor_keys=True) + assert [ + export_name(f"model.layers.0.mlp.experts.1.{part}.weight") + for part in ("gate_proj", "up_proj", "down_proj") + ] == [f"{stored}.1.{hub}.weight" for hub in ("w1", "w3", "w2")] From fb9685c7b80f0d4d945a3d32c33d4ac5c78db1d8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:42 +0000 Subject: [PATCH 09/45] Refuse a checkpoint whose compressed weights store different names The load hands the decompression each expert's stored tensors by the order it collects them, not by expert id, so equal counts do not make a weight and a scale the same expert's: packed weights for experts 0 and 1 beside scales for experts 1 and 2 would pair expert 0's weight with expert 1's scale. check_compressed_keys, which still sees the names, now refuses a checkpoint unless every compressed weight stores the same names; a name no weight stores (Kimi-K3's weight_global_scale) is not read. A scale stored without its packed weight is refused the same way. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 19 ++++++++++++++++++ .../torch/models/test_onfly_decompress.py | 20 +++++++++++++++++-- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 5f36f3303092..ec5a8ef9a580 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -423,8 +423,11 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: tensor_files = checkpoint_tensor_files(ckpt_dir) compressed: set[str] = set() expert_weights: set[str] = set() + stems_storing: dict[str, set[str]] = {} for name in tensor_files: stem, _, stored = name.rpartition(".") + if stored in stored_as: + stems_storing.setdefault(stored, set()).add(stem) if stored in stored_as and stored != "weight": compressed.add(stem) elif stored == "weight" and ".experts." in name: @@ -438,6 +441,22 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: "on-the-fly decompression loads the experts alone, every one compressed, so the rest " "would load newly initialized. Decompress the checkpoint first." ) + # The load hands the decompression each expert's stored tensors by the order it collects them, + # not by expert id, so every compressed weight must store the same names: one missing a scale + # would pair a later expert's weight with an earlier one's scale. A name no weight stores + # (Kimi-K3's weight_global_scale) is simply not read. + unpaired = sorted( + (stored, sorted(compressed - stems)) + for stored, stems in stems_storing.items() + if stems & compressed and compressed - stems + ) + if unpaired: + stored, missing = unpaired[0] + raise NotImplementedError( + f"The checkpoint stores {stored} for some compressed weights but not for {len(missing)} " + f"others {missing[:1]}: the decompression pairs each expert's stored tensors by position, " + "so every compressed weight must store the same names. Decompress the checkpoint first." + ) by_file: dict[str, list[str]] = {} for name, file in tensor_files.items(): stem, _, stored = name.rpartition(".") diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 24730d76b5fe..95b6cd01bc65 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -585,7 +585,8 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( ) == hooks -EXPERT = "model.layers.1.mlp.experts.0.gate_proj" +EXPERTS = [f"model.layers.1.mlp.experts.{e}.gate_proj" for e in range(3)] +EXPERT = EXPERTS[0] 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 @@ -621,6 +622,17 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( [f"{EXPERT}.{n}" for n in (*PACKED, "weight_global_scale")], r"store weight_global_scale \(F32\), which bf16 cannot hold exactly", ), + ( + PACKED, + [f"{EXPERT}.{n}" for n in PACKED] + [f"{EXPERTS[1]}.weight_scale"], + "stores weight_packed for some compressed weights but not for 1 others", + ), + ( + PACKED, + [f"{EXPERTS[e]}.{n}" for e in (0, 1) for n in ("weight_packed", "weight_shape")] + + [f"{EXPERTS[e]}.weight_scale" for e in (1, 2)], + "stores weight_packed for some compressed weights but not for 1 others", + ), ], ids=[ "experts-alone", @@ -630,13 +642,17 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( "an-expert-left-uncompressed", "nothing-compressed", "an-fp32-scale", + "a-scale-alone", + "equal-counts-other-experts", ], ) 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. A float - stored for the experts must be one bf16 holds exactly: the load may cast it to bf16 first.""" + stored for the experts must be one bf16 holds exactly: the load may cast it to bf16 first. And + every compressed weight stores the same names: the load pairs them by position, not by expert, + so a scale stored for other experts in equal number would pair with the wrong weights.""" tensors = { name: torch.zeros(1, dtype=STORED_DTYPES.get(name.rpartition(".")[2], torch.bfloat16)) for name in names From bc76e5874851365091a8d40934061e84d96f32c3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:46 +0000 Subject: [PATCH 10/45] Keep integer sources native only in a decompressing load The copy patch kept every integer source native for every later load in the process, so an unmarked model whose checkpoint stores an integer for a float parameter would load it uncast. The patched conversion mapping, which from_pretrained asks for before it reads a tensor, now records whether the load under way is one whose converters it rewrote, and only that load keeps integers native. A real load of an unmarked model after installing casts as before. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 17 +++++--- .../torch/models/test_onfly_decompress.py | 40 +++++++++++++++++-- 2 files changed, 49 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index ec5a8ef9a580..679ca1f0ef01 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -260,12 +260,17 @@ def _compressedify(conv, compressor, scheme, param_names): return new +# Whether the load under way is one whose converters the hooks rewrote: from_pretrained asks for the +# conversion mapping before it reads a tensor. Module state, not a context variable: the loader +# materializes tensors in its worker threads. +_LOAD = {"decompressing": False} + + def _install_keep_int_native(): - """Patch ``_materialize_copy`` so an integer source tensor is never cast to a float dtype. + """Patch ``_materialize_copy`` so a decompressing load never casts an integer source to a float. Without a quantizer the loader casts every source to the parameter's dtype, which would destroy - packed weights, shapes and zero-points. Process-wide: only a compressed checkpoint loaded without - its quantizer asks for such a cast. + packed weights, shapes and zero-points. Every other load casts as before. """ import transformers.core_model_loading as cml @@ -277,7 +282,8 @@ def _install_keep_int_native(): # keyword: the patch sits on every later load's copy path. def _materialize_copy(tensor, device=None, dtype=None, *args, **kwargs): t = tensor[...] - if dtype is not None and dtype.is_floating_point and not t.is_floating_point(): + floating = dtype is not None and dtype.is_floating_point + if _LOAD["decompressing"] and floating and not t.is_floating_point(): dtype = None # keep packed/shape/zero-point/g_idx integers native return orig(t, device, dtype, *args, **kwargs) @@ -369,7 +375,8 @@ 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 not is_marked(config): + _LOAD["decompressing"] = is_marked(config) + if not _LOAD["decompressing"]: # Not the compressed checkpoint: rewritten, an uncompressed MoE's experts would match no # stored key and load newly initialized without an error. return convs diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 95b6cd01bc65..d0cba9359811 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -45,6 +45,7 @@ AutoConfig, AutoModelForCausalLM, DeepseekV3Config, + LlamaConfig, MixtralConfig, PretrainedConfig, ) @@ -64,9 +65,10 @@ @pytest.fixture def restore_loader_hooks(monkeypatch): - """Undo the process-wide transformers hooks the decompressor installs.""" + """Undo the process-wide transformers hooks the decompressor installs, and its load state.""" for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) + monkeypatch.setattr(onfly, "_LOAD", dict(onfly._LOAD)) class _ScaledCompressor: @@ -235,10 +237,12 @@ def decompress(cls, state_dict, scheme): @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.""" + packed weights, shapes and zero-points. In a decompressing load integers stay integers; floats + are still cast, and integer-to-integer casts still happen. Any other load casts as before.""" onfly._install_keep_int_native() packed = torch.tensor([[2**31 - 15, -(2**31) + 3]], dtype=torch.int32) # not exact in bf16 + assert cml._materialize_copy(packed, dtype=torch.bfloat16).dtype == torch.bfloat16 + onfly._LOAD["decompressing"] = True kept = cml._materialize_copy(packed, dtype=torch.bfloat16) @@ -890,3 +894,33 @@ def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp export_name(f"model.layers.0.mlp.experts.1.{part}.weight") for part in ("gate_proj", "up_proj", "down_proj") ] == [f"{stored}.1.{hub}.weight" for hub in ("w1", "w3", "w2")] + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_load_the_decompressor_did_not_mark_casts_as_before(tmp_path): + """The hooks are process-wide, but only a decompressing load keeps an integer source integer: + a later load of a model no config of which was marked casts it to its float parameter.""" + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) + onfly.install_onfly_decompress( + SimpleNamespace(quantization_config=_pack_quantized(experts=_int4(True, ["Linear"]))) + ) + config = LlamaConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + tie_word_embeddings=False, + ) + AutoModelForCausalLM.from_config(config).save_pretrained(tmp_path) + state = load_file(str(tmp_path / "model.safetensors")) + state["model.norm.weight"] = torch.full_like(state["model.norm.weight"], 3, dtype=torch.int8) + save_file(state, str(tmp_path / "model.safetensors"), metadata={"format": "pt"}) + + model = AutoModelForCausalLM.from_pretrained(tmp_path, dtype=torch.bfloat16) + + assert model.model.norm.weight.dtype == torch.bfloat16 + assert torch.equal(model.model.norm.weight, torch.full_like(model.model.norm.weight, 3)) From 0ef52f53d62083d7be9b57dcae02ec1e543408fb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:50 +0000 Subject: [PATCH 11/45] Refuse a dtype other than bf16 at any config level The install checked only the top-level config's dtype. A VLM config with none there and float16 on its text config passed, and the backbone config to_native_backbone_config derives from that text config keeps float16, with the experts decompressed to bf16 beside it. Every level must now be bf16 or unset; "auto" is refused too, since it reads the dtype from the checkpoint. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 18 ++++++++++-------- .../unit/torch/models/test_onfly_decompress.py | 13 +++++++++++-- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 679ca1f0ef01..a0f601c718e4 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -328,14 +328,16 @@ def install_onfly_decompress(hf_config): 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") - dtype = getattr(hf_config, "dtype", None) - if dtype is not None and str(dtype).removeprefix("torch.") not in ("bfloat16", "auto"): - # The experts would load in bf16 beside the rest in this dtype: the loader sets each - # parameter from what the conversion hands it, without a cast. - raise NotImplementedError( - f"install_onfly_decompress: the experts decompress to bfloat16, but this config " - f"loads in {dtype}. Load the checkpoint in bfloat16." - ) + # The experts would load in bf16 beside the rest in another dtype: the loader sets each + # parameter from what the conversion hands it, without a cast. Every level counts: a backbone + # config derived from the text config (to_native_backbone_config) keeps its dtype. + for level in _config_levels(hf_config): + dtype = getattr(level, "dtype", None) + if dtype is not None and str(dtype).removeprefix("torch.") != "bfloat16": + raise NotImplementedError( + f"install_onfly_decompress: the experts decompress to bfloat16, but this config " + f"loads in {dtype}. Load the checkpoint in bfloat16." + ) 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. diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index d0cba9359811..b1cc3d4a41b4 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -560,8 +560,16 @@ def test_the_quantization_config_is_found_at_the_top_level_then_on_the_text_conf NotImplementedError, "loads in torch.float16", ), + ( + SimpleNamespace( + text_config=SimpleNamespace(dtype=torch.float16), + quantization_config={"format": "pack-quantized"}, + ), + NotImplementedError, + "loads in torch.float16", + ), ], - ids=["nothing-to-decompress", "raw-dict", "not-bf16"], + ids=["nothing-to-decompress", "raw-dict", "not-bf16", "not-bf16-text-config"], ) def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( monkeypatch, config, error, match @@ -569,7 +577,8 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( """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. Nor for a raw dict, which cannot carry the mark its load would be found by, or for a load in a dtype other - than the bf16 the experts decompress to: the loader would set them beside the rest uncast.""" + than the bf16 the experts decompress to, at any config level (a backbone config derived from the + text config keeps its dtype): the loader would set them beside the rest uncast.""" monkeypatch.setitem(sys.modules, "compressed_tensors", None) hooks = ( mu.get_model_conversion_mapping, From ba83e7011ad411bd347ece1ebb7e7dab45263573 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:55 +0000 Subject: [PATCH 12/45] Import the decompressor's dependency-light helpers at module level print_rank_0, require_transformers and safetensors' safe_open add no optional dependency at import. compressed-tensors, the version-dependent transformers modules and the checkpoint plugin stay lazy. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index a0f601c718e4..e3b7c15056db 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -25,8 +25,10 @@ import os import torch +from safetensors import safe_open -from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS +from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS, require_transformers def to_native_backbone_config( @@ -322,9 +324,6 @@ def install_onfly_decompress(hf_config): model built from a config they marked: one compressed checkpoint per process, installing again for it marks that config too, and one compressed differently is refused. """ - 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") @@ -424,8 +423,7 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: bf16 before decompressing, so one bf16 cannot hold exactly (an fp32 scale) is refused too. Reads only the tensors' names and dtypes. """ - from safetensors import safe_open - + # A plugin: it imports huggingface_hub, which this module needs nowhere else. from modelopt.torch.utils.plugins.hf_checkpoint_utils import checkpoint_tensor_files stored_as = stored_names(hf_config) From 41b67941a5a704cf1a08fe4694e023d2c6426b32 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:35:48 +0000 Subject: [PATCH 13/45] Refuse unknown stored names and non-weight sources before a tensor is read The stored names come from the installed compressed-tensors, not the checkpoint: a name a newer release added under a compressed weight would match no rewritten source, load as an unexpected key, and leave the experts decompressed without it. check_compressed_keys now refuses it. A per-expert converter source that is no weight was handed to the decompress op, which refused it from a loader thread as an unconsumed key; _compressedify now refuses it where the converter is rewritten. stored_names says what it returns. Comments tightened to keep the change within its size. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 64 ++++++++++--------- .../torch/models/test_onfly_decompress.py | 32 ++++++---- 2 files changed, 55 insertions(+), 41 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index e3b7c15056db..38fed405eff5 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -124,10 +124,8 @@ def _read_compression(hf_config): "compressing only the experts. Decompress the checkpoint first." ) [scheme] = groups.values() - # 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). + # Regex targets alone must name the experts. A class target such as Linear does not say which + # modules it compressed: check_compressed_keys checks the names the checkpoint stores instead. 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 " @@ -198,8 +196,8 @@ def convert(self, input_dict, *args, **kwargs): if anchor_list[i] is not None ] 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. + # A stored param without its 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)}" ) @@ -214,13 +212,9 @@ def _decompressing_converter_class(): from transformers.core_model_loading import WeightConverter class DecompressingConverter(WeightConverter): - """Reads the stored keys, and reverses as the converter it rewrote. - - A save reverses the converters ``from_pretrained`` kept as ``model._weight_conversions`` - (transformers' ``save_pretrained``, and ModelOpt's export naming the experts back after - the checkpoint): the decompression is load-only, so the experts go back under their - checkpoint's ``.weight`` names, as the original fusion has them. - """ + """Reads the stored keys, and reverses as the converter it rewrote: the decompression is + load-only, so a save (``save_pretrained``, ModelOpt's export) names the experts back after + the checkpoint's ``.weight`` keys.""" __slots__ = ("_reversible",) @@ -240,12 +234,15 @@ def _compressedify(conv, compressor, scheme, param_names): 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) # DecompressExperts refuses it (none expected) + if not base.endswith(".weight"): + # Only a weight is stored compressed: another source would reach the merge as stored. + raise NotImplementedError( + f"install_onfly_decompress: the per-expert converter for {conv.target_patterns} " + f"reads {p!r}, which is no weight; on-the-fly decompression rewrites weights only." + ) + stem = base[: -len(".weight")] + for pn in param_names: + by_pname[pn].append(stem + "." + pn + "$") new_sources = [s for pn in param_names for s in by_pname[pn]] tgt = getattr(conv, "_original_target_patterns", conv.target_patterns) new = _decompressing_converter_class()( @@ -262,9 +259,8 @@ def _compressedify(conv, compressor, scheme, param_names): return new -# Whether the load under way is one whose converters the hooks rewrote: from_pretrained asks for the -# conversion mapping before it reads a tensor. Module state, not a context variable: the loader -# materializes tensors in its worker threads. +# Whether the load under way is one whose converters the hooks rewrote (from_pretrained asks for the +# mapping before it reads a tensor). Module state: the loader materializes in worker threads. _LOAD = {"decompressing": False} @@ -310,7 +306,7 @@ def is_marked(hf_config) -> bool: def stored_names(hf_config) -> tuple[str, ...]: - """The names each compressed weight of ``hf_config``'s checkpoint is stored under, from its mark.""" + """The names the installed compressed-tensors stores each weight under for ``hf_config``'s scheme.""" for c in _config_levels(hf_config): if (stored_as := getattr(c, _MARK, None)) is not None: return stored_as @@ -327,9 +323,8 @@ def install_onfly_decompress(hf_config): 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") - # The experts would load in bf16 beside the rest in another dtype: the loader sets each - # parameter from what the conversion hands it, without a cast. Every level counts: a backbone - # config derived from the text config (to_native_backbone_config) keeps its dtype. + # The loader sets each parameter from what the conversion hands it, uncast. Every level counts: + # a backbone config derived from the text config (to_native_backbone_config) keeps its dtype. for level in _config_levels(hf_config): dtype = getattr(level, "dtype", None) if dtype is not None and str(dtype).removeprefix("torch.") != "bfloat16": @@ -448,10 +443,19 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: "on-the-fly decompression loads the experts alone, every one compressed, so the rest " "would load newly initialized. Decompress the checkpoint first." ) - # The load hands the decompression each expert's stored tensors by the order it collects them, - # not by expert id, so every compressed weight must store the same names: one missing a scale - # would pair a later expert's weight with an earlier one's scale. A name no weight stores - # (Kimi-K3's weight_global_scale) is simply not read. + # The names are what the installed compressed-tensors gives the scheme: one a newer release added + # would match no rewritten source, and the experts would decompress without it. + unknown = [n for n in sorted(tensor_files) if n.rpartition(".")[0] in compressed] + unknown = [n for n in unknown if n.rpartition(".")[2] not in stored_as] + if unknown: + raise NotImplementedError( + f"The checkpoint stores {len(unknown)} tensors under its compressed weights that this " + f"compressed-tensors release does not name for the scheme {unknown[:1]}: the decompression " + f"reads only {list(stored_as)}. Upgrade compressed-tensors, or decompress it first." + ) + # The load pairs each expert's stored tensors by the order it collects them, not by expert id: one + # missing a scale would pair a later expert's weight with an earlier one's scale. A name no weight + # stores (Kimi-K3's weight_global_scale) is simply not read. unpaired = sorted( (stored, sorted(compressed - stems)) for stored, stems in stems_storing.items() diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index b1cc3d4a41b4..931f02b32c6d 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -117,21 +117,23 @@ def _load(converters, checkpoint): # ------------------------------------------------------------- rewriting the per-expert converters -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.""" +def test_only_weight_sources_are_rewritten_and_stored_params_need_their_weight(): + """Only ``.weight`` sources name a compressed tensor: a converter reading another is refused + where it is rewritten, before any tensor is read. A stored param without its weight 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)], ) - 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) + with pytest.raises(NotImplementedError, match=r"reads .*gate_proj\.bias.*weights only"): + onfly._compressedify(conv, _ScaledCompressor, None, PACKED) - with pytest.raises(ValueError, match=r"unconsumed checkpoint keys .*gate_proj\.bias"): - _load([onfly._compressedify(conv, _ScaledCompressor, None, PACKED)], checkpoint) + gate_up, _ = _stock_expert_converters() + scales_alone = {"model.layers.0.mlp.experts.0.gate_proj.weight_scale": torch.ones(2, 2)} + with pytest.raises(ValueError, match=r"unconsumed checkpoint keys .*gate_proj\.weight_scale"): + _load([onfly._compressedify(gate_up, _ScaledCompressor, None, PACKED)], scales_alone) def test_a_rewritten_converter_keeps_the_scope_transformers_set_on_it(): @@ -640,6 +642,11 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( [f"{EXPERT}.{n}" for n in PACKED] + [f"{EXPERTS[1]}.weight_scale"], "stores weight_packed for some compressed weights but not for 1 others", ), + ( + PACKED, + [f"{EXPERT}.{n}" for n in (*PACKED, "weight_zero_point")], + "stores 1 tensors under its compressed weights that this compressed-tensors release", + ), ( PACKED, [f"{EXPERTS[e]}.{n}" for e in (0, 1) for n in ("weight_packed", "weight_shape")] @@ -656,6 +663,7 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( "nothing-compressed", "an-fp32-scale", "a-scale-alone", + "a-name-this-release-does-not-know", "equal-counts-other-experts", ], ) @@ -664,8 +672,10 @@ def test_a_checkpoint_must_compress_its_experts_and_nothing_else(tmp_path, store 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. A float stored for the experts must be one bf16 holds exactly: the load may cast it to bf16 first. And - every compressed weight stores the same names: the load pairs them by position, not by expert, - so a scale stored for other experts in equal number would pair with the wrong weights.""" + every compressed weight stores the same names, all of them names the installed compressed-tensors + gives the scheme: the load pairs them by position, not by expert, so a scale stored for other + experts in equal number would pair with the wrong weights, and a name it does not know would be + left out of the decompression.""" tensors = { name: torch.zeros(1, dtype=STORED_DTYPES.get(name.rpartition(".")[2], torch.bfloat16)) for name in names From a9df1accf7b1403aa899eb375d30f3bc69a654b3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:53:35 +0000 Subject: [PATCH 14/45] Keep packed integers uncast where an expert-parallel load shards them From transformers 5.14 an expert-parallel load reads each expert's tensors through DtensorShardOperation.shard_tensor, which casts them to the parameter's dtype as it shards them onto the mesh, and never reaches the _materialize_copy the keep-integers patch sat on. A compressed checkpoint's packed weights (INT4's int32, MXFP4's uint8) then reached the decompression as bf16 and failed there, as a 2-GPU Kimi-K3 load on transformers 5.18 did. The patch now sits on spawn_materialize, where the loader starts both the copy and the shard, in every release from 5.8: a decompressing load drops the float dtype for an integer source there, judging a safetensors slice by its stored dtype without reading it. A test shards an int32 tensor onto a one-rank CPU mesh through transformers' own DtensorShardOperation; it fails without the change. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 33 ++++--- .../torch/models/test_onfly_decompress.py | 90 +++++++++++++------ 2 files changed, 85 insertions(+), 38 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 38fed405eff5..3c2af77b3a2f 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -265,28 +265,35 @@ def _compressedify(conv, compressor, scheme, param_names): def _install_keep_int_native(): - """Patch ``_materialize_copy`` so a decompressing load never casts an integer source to a float. + """Patch ``spawn_materialize`` so a decompressing load never casts an integer source to a float. - Without a quantizer the loader casts every source to the parameter's dtype, which would destroy - packed weights, shapes and zero-points. Every other load casts as before. + Without a quantizer the loader casts every source to the parameter's dtype as it reads it: in its + copy, or, from transformers 5.14, in ``DtensorShardOperation.shard_tensor`` as it shards it onto + the expert-parallel mesh. Either would destroy packed weights, shapes and zero-points, so the + dtype is dropped where both start. Every other load casts as before. """ import transformers.core_model_loading as cml - if getattr(cml._materialize_copy, "_onfly_patched", False): + if getattr(cml.spawn_materialize, "_onfly_patched", False): return - orig = cml._materialize_copy + orig = cml.spawn_materialize - # Any other argument passes through, positionally as the loader passes the ones it has, or by - # keyword: the patch sits on every later load's copy path. - def _materialize_copy(tensor, device=None, dtype=None, *args, **kwargs): - t = tensor[...] + # Any other argument passes through, positionally or by keyword: a release may add some. + def spawn_materialize(thread_pool, tensor, device=None, dtype=None, *args, **kwargs): floating = dtype is not None and dtype.is_floating_point - if _LOAD["decompressing"] and floating and not t.is_floating_point(): + if _LOAD["decompressing"] and floating and not _stores_floats(tensor): dtype = None # keep packed/shape/zero-point/g_idx integers native - return orig(t, device, dtype, *args, **kwargs) + return orig(thread_pool, tensor, device, dtype, *args, **kwargs) - setattr(_materialize_copy, "_onfly_patched", True) - cml._materialize_copy = _materialize_copy + setattr(spawn_materialize, "_onfly_patched", True) + cml.spawn_materialize = spawn_materialize + + +def _stores_floats(tensor) -> bool: + """Whether a checkpoint tensor holds floats; a safetensors slice is judged without reading it.""" + if hasattr(tensor, "get_dtype"): + return tensor.get_dtype().startswith(("F", "BF")) + return tensor.is_floating_point() # Marks a config whose model's converters the installed hooks rewrite (holds the names each diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 931f02b32c6d..3b6e1e64b2c6 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -30,6 +30,7 @@ import pytest import torch +from safetensors import safe_open from safetensors.torch import load_file, save_file transformers = pytest.importorskip("transformers") @@ -66,7 +67,7 @@ @pytest.fixture def restore_loader_hooks(monkeypatch): """Undo the process-wide transformers hooks the decompressor installs, and its load state.""" - for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) monkeypatch.setattr(onfly, "_LOAD", dict(onfly._LOAD)) @@ -236,43 +237,82 @@ def decompress(cls, state_dict, scheme): decompress.convert(short) +PACKED_INT32 = torch.tensor([[2**31 - 15, -(2**31) + 3]], dtype=torch.int32) # not exact in bf16 + + +def _read(path, sharding_op=None, dtype=torch.bfloat16): + """Read ``path``'s one tensor the way the loader does: a safetensors slice, through spawn_materialize.""" + with safe_open(str(path), framework="pt") as stored: + [key] = stored.keys() + kwargs = {"sharding_op": sharding_op, "tensor_idx": 0} if sharding_op else {} + return cml.spawn_materialize(None, stored.get_slice(key), "cpu", dtype, **kwargs)() + + +@pytest.mark.usefixtures("restore_loader_hooks") +def test_integer_checkpoint_tensors_keep_their_dtype_on_load(tmp_path): + """Without a quantizer the loader casts every tensor to the model's float dtype as it reads it, + which destroys packed weights, shapes and zero-points. In a decompressing load integers stay + integers; floats are still cast, and integer-to-integer casts still happen. Any other load casts + as before.""" + save_file({"packed": PACKED_INT32}, str(tmp_path / "packed.safetensors")) + onfly._install_keep_int_native() + assert _read(tmp_path / "packed.safetensors").dtype == torch.bfloat16 + onfly._LOAD["decompressing"] = True + + kept = _read(tmp_path / "packed.safetensors") + + assert kept.dtype == torch.int32 + assert torch.equal(kept, PACKED_INT32) + assert _read(tmp_path / "packed.safetensors", dtype=torch.int64).dtype == torch.int64 + assert ( + cml.spawn_materialize(None, torch.ones(2), "cpu", torch.bfloat16)().dtype == torch.bfloat16 + ) + assert cml.spawn_materialize(None, PACKED_INT32, "cpu", torch.bfloat16)().dtype == torch.int32 + + @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. In a decompressing load integers stay integers; floats - are still cast, and integer-to-integer casts still happen. Any other load casts as before.""" +def test_an_integer_sharded_onto_the_expert_mesh_keeps_its_dtype(tmp_path, gloo_world_of_one): + """From transformers 5.14 an expert-parallel load shards each expert's tensor onto the mesh as it + reads it, and the shard op casts it to the parameter's dtype, not the copy: a decompressing load + keeps the packed integers there too.""" + sharding_utils = pytest.importorskip("transformers.distributed.sharding_utils") + from torch.distributed.tensor import Shard, distribute_tensor + + experts = torch.zeros(2, 4, 8, dtype=torch.bfloat16) # a fused expert param, sharded by expert + shard = sharding_utils.DtensorShardOperation( + distribute_tensor(experts, gloo_world_of_one, [Shard(0)]) + ) + save_file({"packed": PACKED_INT32}, str(tmp_path / "packed.safetensors")) onfly._install_keep_int_native() - packed = torch.tensor([[2**31 - 15, -(2**31) + 3]], dtype=torch.int32) # not exact in bf16 - assert cml._materialize_copy(packed, dtype=torch.bfloat16).dtype == torch.bfloat16 onfly._LOAD["decompressing"] = True - kept = cml._materialize_copy(packed, dtype=torch.bfloat16) + kept = _read(tmp_path / "packed.safetensors", sharding_op=shard) 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 + assert torch.equal(kept, PACKED_INT32) @pytest.mark.usefixtures("restore_loader_hooks") -def test_a_copy_parameter_a_release_adds_passes_through(monkeypatch): - """The patch sits on every later load's copy path, so a parameter a release adds reaches the - loader's copy, passed positionally, as the loader passes the ones it has, or by keyword.""" +def test_a_read_parameter_a_release_adds_passes_through(monkeypatch): + """The patch sits on every later load's reads, so a parameter past the dtype (5.14's sharding_op + and tensor_idx, or one a release adds) reaches the loader's own, positionally or by keyword.""" seen = [] - def release_copy(tensor, device=None, dtype=None, non_blocking=False): - seen.append(non_blocking) + def release_spawn( + thread_pool, tensor, device=None, dtype=None, sharding_op=None, tensor_idx=None + ): + seen.append((sharding_op, tensor_idx)) return tensor - monkeypatch.setattr(cml, "_materialize_copy", release_copy) + monkeypatch.setattr(cml, "spawn_materialize", release_spawn) onfly._install_keep_int_native() - cml._materialize_copy(torch.ones(2), "cpu", torch.bfloat16, True) - cml._materialize_copy(torch.ones(2), "cpu", torch.bfloat16, non_blocking=True) + cml.spawn_materialize(None, torch.ones(2), "cpu", torch.bfloat16, "op", 3) + cml.spawn_materialize( + None, torch.ones(2), "cpu", torch.bfloat16, sharding_op="op", tensor_idx=3 + ) - assert seen == [True, True] + assert seen == [("op", 3), ("op", 3)] # ------------------------------------------------------------------------ installing the rewrite @@ -361,7 +401,7 @@ def stock_mapping(model, **kwargs): monkeypatch.setattr(onfly, "_read_compression", read) def hooks(): - return mu.get_model_conversion_mapping, cml._materialize_copy + return mu.get_model_conversion_mapping, cml.spawn_materialize onfly.install_onfly_decompress(SimpleNamespace()) installed = hooks() @@ -586,7 +626,7 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( mu.get_model_conversion_mapping, WeightConverter.add_tensor, WeightConverter.materialize_tensors, - cml._materialize_copy, + cml.spawn_materialize, ) with pytest.raises(error, match=match): @@ -596,7 +636,7 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( mu.get_model_conversion_mapping, WeightConverter.add_tensor, WeightConverter.materialize_tensors, - cml._materialize_copy, + cml.spawn_materialize, ) == hooks From 9998913e95792350678508f3f699f24d1b71f76a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:53:41 +0000 Subject: [PATCH 15/45] Rewrite experts' weights only, let an expert's bias load, check the floor first From Claude's review: _compressedify refuses a source that is no expert weight, so a dense weight fused with the experts' cannot be rewritten to stored names it does not have and load newly initialized. The stored-name check refuses only a weight's names the installed compressed-tensors does not give the scheme, so an expert's uncompressed bias, loaded by its own converter, passes. The install checks the transformers floor before it needs compressed-tensors, so the first error names what to fix first. Comments tightened. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 39 +++++++++--------- .../torch/models/test_onfly_decompress.py | 40 +++++++++++++------ 2 files changed, 46 insertions(+), 33 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 3c2af77b3a2f..2ccd7c960879 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -93,8 +93,7 @@ def _read_compression(hf_config): ``param_names`` are the keys one weight is stored under. The one config group must target the experts; several groups are refused. ``None`` for a config with no ``quantization_config``.""" - # 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. + # With no config, compressed_tensors need not even be installed. quantization_config = _find_quantization_config(hf_config) if quantization_config is None: return None @@ -116,8 +115,7 @@ def _read_compression(hf_config): "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. + # Only the expert converters are rewritten: 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)}); on-the-fly decompression handles a single group, " @@ -181,8 +179,6 @@ def convert(self, input_dict, *args, **kwargs): 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" @@ -196,8 +192,7 @@ def convert(self, input_dict, *args, **kwargs): if anchor_list[i] is not None ] if input_dict: - # A stored param without its weight: passed on, the merge would fuse it into the - # wrong slot or fail far from here. + # A stored param without its weight: the merge would fuse it into the wrong slot. raise ValueError( f"DecompressExperts: unconsumed checkpoint keys {sorted(input_dict)}" ) @@ -234,11 +229,13 @@ def _compressedify(conv, compressor, scheme, param_names): by_pname = {pn: [] for pn in param_names} for p in conv.source_patterns: base = p.removesuffix("$") - if not base.endswith(".weight"): - # Only a weight is stored compressed: another source would reach the merge as stored. + if not base.endswith(".weight") or ".experts." not in base: + # Only the experts' weights are stored compressed: rewritten, a dense weight would match + # no stored key and load newly initialized; a bias would reach the merge as stored. raise NotImplementedError( f"install_onfly_decompress: the per-expert converter for {conv.target_patterns} " - f"reads {p!r}, which is no weight; on-the-fly decompression rewrites weights only." + f"reads {p!r}, which is no expert weight; on-the-fly decompression rewrites the " + "experts' weights only." ) stem = base[: -len(".weight")] for pn in param_names: @@ -251,8 +248,7 @@ def _compressedify(conv, compressor, scheme, param_names): operations=[DecompressExperts(), *conv.operations], ) new._reversible = conv - # 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. + # What else transformers sets on a converter it maps, 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)) @@ -339,8 +335,9 @@ def install_onfly_decompress(hf_config): f"install_onfly_decompress: the experts decompress to bfloat16, but this config " f"loads in {dtype}. Load the checkpoint in bfloat16." ) - compression = _read_compression(hf_config) + # Before compressed-tensors is needed: a release below the floor is what to fix first. require_transformers("5.8", "on-the-fly compressed-tensors decompression") + compression = _read_compression(hf_config) # 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 @@ -348,9 +345,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: - # A config its prepare dropped quantization_config from after installing (Kimi-K3's), - # already marked, may be prepared again: nothing new to decompress. An unmarked one would - # load its experts through the stock converters, newly initialized. + # A marked config its prepare dropped quantization_config from (Kimi-K3's) may be prepared + # again; an unmarked one would load its experts newly initialized. if not is_marked(hf_config): raise ValueError( "install_onfly_decompress: config has no compressed-tensors " @@ -395,8 +391,7 @@ def patched(*args, **kwargs): else: out.append(c) if not rewritten: - # The stock converters would read the experts' plain .weight keys, which the checkpoint - # does not store. + # The stock converters would read .weight keys 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." @@ -450,10 +445,12 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: "on-the-fly decompression loads the experts alone, every one compressed, so the rest " "would load newly initialized. Decompress the checkpoint first." ) - # The names are what the installed compressed-tensors gives the scheme: one a newer release added - # would match no rewritten source, and the experts would decompress without it. + # The names are what the installed compressed-tensors gives the scheme: a weight's name a newer + # release added would match no rewritten source, and the experts would decompress without it. A + # bias loads through its own converter. unknown = [n for n in sorted(tensor_files) if n.rpartition(".")[0] in compressed] unknown = [n for n in unknown if n.rpartition(".")[2] not in stored_as] + unknown = [n for n in unknown if n.rpartition(".")[2].startswith("weight")] if unknown: raise NotImplementedError( f"The checkpoint stores {len(unknown)} tensors under its compressed weights that this " diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 3b6e1e64b2c6..c2712bd371d2 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -118,18 +118,19 @@ def _load(converters, checkpoint): # ------------------------------------------------------------- rewriting the per-expert converters -def test_only_weight_sources_are_rewritten_and_stored_params_need_their_weight(): - """Only ``.weight`` sources name a compressed tensor: a converter reading another is refused - where it is rewritten, before any tensor is read. A stored param without its weight 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)], - ) - with pytest.raises(NotImplementedError, match=r"reads .*gate_proj\.bias.*weights only"): - onfly._compressedify(conv, _ScaledCompressor, None, PACKED) +def test_only_expert_weight_sources_are_rewritten_and_stored_params_need_their_weight(): + """Only the experts' ``.weight`` sources name a compressed tensor: a converter reading another -- + an expert's bias, or a dense weight fused with the experts' -- is refused where it is rewritten, + before any tensor is read. A stored param without its weight still reaches the decompress op, + which refuses it rather than hand the merge a tensor it would fuse into the wrong slot.""" + for other in ("mlp.experts.*.gate_proj.bias", "mlp.gate_proj.weight"): + conv = WeightConverter( + source_patterns=["mlp.experts.*.gate_proj.weight", other], + target_patterns="mlp.experts.gate_proj", + operations=[MergeModulelist(dim=0)], + ) + with pytest.raises(NotImplementedError, match=r"which is no expert weight"): + onfly._compressedify(conv, _ScaledCompressor, None, PACKED) gate_up, _ = _stock_expert_converters() scales_alone = {"model.layers.0.mlp.experts.0.gate_proj.weight_scale": torch.ones(2, 2)} @@ -640,6 +641,19 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( ) == hooks +@pytest.mark.usefixtures("restore_loader_hooks") +def test_a_transformers_below_the_floor_is_named_before_compressed_tensors(monkeypatch): + """On a transformers the decompression cannot run on, the first error names the release, not + the compressed-tensors package a user would install only to hit the release next.""" + monkeypatch.setitem(sys.modules, "compressed_tensors", None) + monkeypatch.setattr(transformers, "__version__", "5.5.4") + + with pytest.raises(RuntimeError, match=r"transformers >= 5\.8"): + onfly.install_onfly_decompress( + SimpleNamespace(quantization_config={"format": "pack-quantized"}) + ) + + EXPERTS = [f"model.layers.1.mlp.experts.{e}.gate_proj" for e in range(3)] EXPERT = EXPERTS[0] ATTENTION = "model.layers.1.self_attn.q_proj" @@ -687,6 +701,7 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( [f"{EXPERT}.{n}" for n in (*PACKED, "weight_zero_point")], "stores 1 tensors under its compressed weights that this compressed-tensors release", ), + (PACKED, [f"{EXPERT}.{n}" for n in (*PACKED, "bias")], None), ( PACKED, [f"{EXPERTS[e]}.{n}" for e in (0, 1) for n in ("weight_packed", "weight_shape")] @@ -704,6 +719,7 @@ def test_a_config_with_nothing_to_decompress_leaves_the_loader_alone( "an-fp32-scale", "a-scale-alone", "a-name-this-release-does-not-know", + "an-expert-bias-loads-apart", "equal-counts-other-experts", ], ) From 35b43ed0723f321502213f2679035e1e57949ebf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:22:40 +0000 Subject: [PATCH 16/45] Start each decompressor test with no decompressing load under way A real decompressing load leaves the load flag set until the next load, and the fixture copied whatever an earlier test left, so the integer test failed when it ran after one in the same session. The fixture now starts each test from a clean flag and restores the module's own after. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_onfly_decompress.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index c2712bd371d2..cf299a91991f 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -66,10 +66,11 @@ @pytest.fixture def restore_loader_hooks(monkeypatch): - """Undo the process-wide transformers hooks the decompressor installs, and its load state.""" + """Undo the process-wide transformers hooks the decompressor installs, and start the test with + no decompressing load under way, whatever an earlier test's load left.""" for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) - monkeypatch.setattr(onfly, "_LOAD", dict(onfly._LOAD)) + monkeypatch.setattr(onfly, "_LOAD", {"decompressing": False}) class _ScaledCompressor: From a2229ca9af901160b5d9007f52e6c846e44937d9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:30:56 +0000 Subject: [PATCH 17/45] Say what the stored-name check covers, and why a grouped bias is refused From Claude's review: the unknown-name check covers weight* names only, which every compressor uses today, and says so. The rewrite's refusal of a source that is no expert weight explains the bias case: a bias belongs in a converter of its own, and grouped with the experts' weights it has no stored key either. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 13 ++++++++----- tests/unit/torch/models/test_onfly_decompress.py | 2 +- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 2ccd7c960879..e2ab2bce52e7 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -231,11 +231,13 @@ def _compressedify(conv, compressor, scheme, param_names): base = p.removesuffix("$") if not base.endswith(".weight") or ".experts." not in base: # Only the experts' weights are stored compressed: rewritten, a dense weight would match - # no stored key and load newly initialized; a bias would reach the merge as stored. + # no stored key and load newly initialized. A bias is expected in a converter of its own + # (check_compressed_keys lets one through); grouped with the weight, the merge takes it. raise NotImplementedError( f"install_onfly_decompress: the per-expert converter for {conv.target_patterns} " - f"reads {p!r}, which is no expert weight; on-the-fly decompression rewrites the " - "experts' weights only." + f"reads {p!r} alongside the experts' weights, which is no compressed tensor; " + "on-the-fly decompression rewrites a converter that reads the experts' weights only. " + "Decompress the checkpoint first." ) stem = base[: -len(".weight")] for pn in param_names: @@ -446,8 +448,9 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: "would load newly initialized. Decompress the checkpoint first." ) # The names are what the installed compressed-tensors gives the scheme: a weight's name a newer - # release added would match no rewritten source, and the experts would decompress without it. A - # bias loads through its own converter. + # release added would match no rewritten source, and the experts would decompress without it. + # Only weight* names are checked (every compressor's today), so an expert's bias, which loads + # through its own converter, is not refused. unknown = [n for n in sorted(tensor_files) if n.rpartition(".")[0] in compressed] unknown = [n for n in unknown if n.rpartition(".")[2] not in stored_as] unknown = [n for n in unknown if n.rpartition(".")[2].startswith("weight")] diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index cf299a91991f..2c680fbe1f94 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -130,7 +130,7 @@ def test_only_expert_weight_sources_are_rewritten_and_stored_params_need_their_w target_patterns="mlp.experts.gate_proj", operations=[MergeModulelist(dim=0)], ) - with pytest.raises(NotImplementedError, match=r"which is no expert weight"): + with pytest.raises(NotImplementedError, match=r"alongside the experts. weights"): onfly._compressedify(conv, _ScaledCompressor, None, PACKED) gate_up, _ = _stock_expert_converters() From da1c6e85feb3a7eb4206262faa9cbd5f2b84262c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:40:26 +0000 Subject: [PATCH 18/45] Import the DTensor helpers at module scope in the decompressor's tests torch.distributed.tensor is a core import with no optional dependency, so the shard-path test imports it with the module's other imports. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_onfly_decompress.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 2c680fbe1f94..2f17987b4a8e 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -32,6 +32,7 @@ import torch from safetensors import safe_open from safetensors.torch import load_file, save_file +from torch.distributed.tensor import Shard, distribute_tensor transformers = pytest.importorskip("transformers") @@ -278,8 +279,6 @@ def test_an_integer_sharded_onto_the_expert_mesh_keeps_its_dtype(tmp_path, gloo_ reads it, and the shard op casts it to the parameter's dtype, not the copy: a decompressing load keeps the packed integers there too.""" sharding_utils = pytest.importorskip("transformers.distributed.sharding_utils") - from torch.distributed.tensor import Shard, distribute_tensor - experts = torch.zeros(2, 4, 8, dtype=torch.bfloat16) # a fused expert param, sharded by expert shard = sharding_utils.DtensorShardOperation( distribute_tensor(experts, gloo_world_of_one, [Shard(0)]) From e9266837465d0de04ac215321c6d934b6943b139 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:42:55 +0000 Subject: [PATCH 19/45] Tighten the decompressor's docstrings to keep the change within its size No behaviour change: shorter module, converter and hook docstrings, and the source-dtype check folded into the hook it serves. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 29 +++++++--------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index e2ab2bce52e7..27bc81d24d59 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -15,10 +15,8 @@ """Decompress a compressed-tensors MoE checkpoint's experts as ``from_pretrained`` loads them. -Each rank decompresses only its own experts into the fused-experts layout, so expert-parallel PTQ -runs on the original compressed checkpoint. The ``compressed_tensors`` library does the math, -driven by the checkpoint's ``quantization_config``. Needs transformers 5.8: earlier releases' -``Concatenate`` looks up the sibling keys the decompression consumed. +Each rank decompresses only its own experts, with the ``compressed_tensors`` library, so expert-parallel +PTQ runs on the original checkpoint. Needs transformers 5.8, whose ``Concatenate`` reads only its keys. """ import functools @@ -207,9 +205,7 @@ def _decompressing_converter_class(): from transformers.core_model_loading import WeightConverter class DecompressingConverter(WeightConverter): - """Reads the stored keys, and reverses as the converter it rewrote: the decompression is - load-only, so a save (``save_pretrained``, ModelOpt's export) names the experts back after - the checkpoint's ``.weight`` keys.""" + """Reverses as the converter it rewrote: a save names the experts back as ``.weight`` keys.""" __slots__ = ("_reversible",) @@ -265,10 +261,8 @@ def _compressedify(conv, compressor, scheme, param_names): def _install_keep_int_native(): """Patch ``spawn_materialize`` so a decompressing load never casts an integer source to a float. - Without a quantizer the loader casts every source to the parameter's dtype as it reads it: in its - copy, or, from transformers 5.14, in ``DtensorShardOperation.shard_tensor`` as it shards it onto - the expert-parallel mesh. Either would destroy packed weights, shapes and zero-points, so the - dtype is dropped where both start. Every other load casts as before. + Without a quantizer the loader casts each source to its parameter's dtype in the copy, or from + transformers 5.14 in ``DtensorShardOperation.shard_tensor``: both start here. Other loads cast. """ import transformers.core_model_loading as cml @@ -278,8 +272,10 @@ def _install_keep_int_native(): # Any other argument passes through, positionally or by keyword: a release may add some. def spawn_materialize(thread_pool, tensor, device=None, dtype=None, *args, **kwargs): - floating = dtype is not None and dtype.is_floating_point - if _LOAD["decompressing"] and floating and not _stores_floats(tensor): + # A safetensors slice is judged by its stored dtype, without reading it. + stored = tensor.get_dtype() if hasattr(tensor, "get_dtype") else None + floats = stored.startswith(("F", "BF")) if stored else tensor.is_floating_point() + if _LOAD["decompressing"] and dtype is not None and dtype.is_floating_point and not floats: dtype = None # keep packed/shape/zero-point/g_idx integers native return orig(thread_pool, tensor, device, dtype, *args, **kwargs) @@ -287,13 +283,6 @@ def spawn_materialize(thread_pool, tensor, device=None, dtype=None, *args, **kwa cml.spawn_materialize = spawn_materialize -def _stores_floats(tensor) -> bool: - """Whether a checkpoint tensor holds floats; a safetensors slice is judged without reading it.""" - if hasattr(tensor, "get_dtype"): - return tensor.get_dtype().startswith(("F", "BF")) - return tensor.is_floating_point() - - # Marks a config whose model's converters the installed hooks rewrite (holds the names each # compressed weight is stored under). _MARK = "_modelopt_onfly_decompress" From 00c6e498caecc658382135b7373662ec3cfc414a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:34:59 +0000 Subject: [PATCH 20/45] Return the tensor names check_compressed_keys read, for its caller to reuse From Claude's review of #2747: the expert-parallel loader read the checkpoint's tensor names and shard headers again, for a check of its own, right after this one read them. The map is now returned instead of thrown away. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 27bc81d24d59..1cbcbb9ea2e2 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -403,13 +403,13 @@ def patched(*args, **kwargs): ) -def check_compressed_keys(hf_config, ckpt_dir) -> None: +def check_compressed_keys(hf_config, ckpt_dir) -> dict[str, str]: """Refuse a checkpoint the decompression would load wrong, by the tensors ``ckpt_dir`` stores. Only the experts decompress, and the load drops the config any other compressed weight would need: the experts alone must be compressed, every one of them. The load may cast their float tensors to bf16 before decompressing, so one bf16 cannot hold exactly (an fp32 scale) is refused too. Reads - only the tensors' names and dtypes. + only the tensors' names and dtypes, and returns the ``tensor name -> file`` map it read. """ # A plugin: it imports huggingface_hub, which this module needs nowhere else. from modelopt.torch.utils.plugins.hf_checkpoint_utils import checkpoint_tensor_files @@ -483,3 +483,4 @@ def check_compressed_keys(hf_config, ckpt_dir) -> None: "to, before decompressing, shifting every weight they scale. Decompress the checkpoint " "first." ) + return tensor_files From 9cf8a9a0156c3aa3cf9b39714900985be1081883 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:43:16 +0000 Subject: [PATCH 21/45] Refuse a compressed-tensors older than 0.16, whose compressors take no scheme The decompressor asks the registered compressor for its parameter names and decompresses one weight with it, both classmethods by the weight's scheme from compressed-tensors 0.16 on. Before that a compressor is built from the whole config and states its names as a property (to 0.14) or not at all (0.15), so on 0.13 the setup failed with "'property' object is not callable". Such a release is now refused before anything is read, naming 0.16.0 and the upgrade; the tests that need the real library skip below it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 25 +++++++++-- .../torch/models/test_onfly_decompress.py | 41 +++++++++++++++++-- 2 files changed, 58 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 1cbcbb9ea2e2..51e84c4d6522 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -16,18 +16,26 @@ """Decompress a compressed-tensors MoE checkpoint's experts as ``from_pretrained`` loads them. Each rank decompresses only its own experts, with the ``compressed_tensors`` library, so expert-parallel -PTQ runs on the original checkpoint. Needs transformers 5.8, whose ``Concatenate`` reads only its keys. +PTQ runs on the original checkpoint. Needs transformers 5.8, whose ``Concatenate`` reads only its keys, +and compressed-tensors 0.16, whose compressors decompress one weight by its scheme. """ import functools import os import torch +from packaging.version import parse from safetensors import safe_open from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS, require_transformers +# The first compressed-tensors whose compressors decompress one weight by its scheme, through the +# classmethods compression_param_names(scheme) and decompress(state_dict, scheme). Before it, a +# compressor is built from the whole config, and states its parameter names as a property (to 0.14) +# or not at all (0.15). +_MIN_COMPRESSED_TENSORS = "0.16.0" + def to_native_backbone_config( vlm_config, @@ -96,14 +104,23 @@ def _read_compression(hf_config): if quantization_config is None: return None try: - from compressed_tensors.compressors.base import BaseCompressor - from compressed_tensors.quantization import QuantizationConfig + import compressed_tensors except ImportError as e: # No pip extra of ModelOpt installs it: only a compressed checkpoint needs it. raise ImportError( "On-the-fly decompression of a compressed-tensors checkpoint needs the " - "compressed-tensors package. Install it with: pip install compressed-tensors" + "compressed-tensors package. Install it with: " + f"pip install 'compressed-tensors>={_MIN_COMPRESSED_TENSORS}'" ) from e + have = compressed_tensors.__version__ + if parse(have) < parse(_MIN_COMPRESSED_TENSORS): + raise RuntimeError( + "On-the-fly decompression of a compressed-tensors checkpoint requires " + f"compressed-tensors >= {_MIN_COMPRESSED_TENSORS}, but {have} is installed. Upgrade " + f"with: pip install -U 'compressed-tensors>={_MIN_COMPRESSED_TENSORS}'." + ) + from compressed_tensors.compressors.base import BaseCompressor + from compressed_tensors.quantization import QuantizationConfig qconfig = QuantizationConfig.model_validate(quantization_config) groups = qconfig.config_groups diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 2f17987b4a8e..ee0019283621 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -846,7 +846,9 @@ def test_a_checkpoint_whose_experts_are_not_its_one_compressed_group_is_refused( 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)" + "compressed_tensors", + minversion=onfly._MIN_COMPRESSED_TENSORS, + reason="needs compressed-tensors (CI's unit session installs it)", ) config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) @@ -864,7 +866,9 @@ def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group format; the one group's format is the one to decompress with. A format nothing registers is refused by this module, rather than failing inside the library's registry.""" pytest.importorskip( - "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + "compressed_tensors", + minversion=onfly._MIN_COMPRESSED_TENSORS, + reason="needs compressed-tensors (CI's unit session installs it)", ) experts = _int4(True, EXPERTS) if group is not None: @@ -878,6 +882,31 @@ def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group onfly._read_compression(config) +def test_a_compressed_tensors_before_per_scheme_compressors_is_refused_by_its_release(monkeypatch): + """compressed-tensors 0.13 states a compressor's parameter names as a property, which a call + by scheme fails on deep in the setup; such a release is refused first, naming the one needed.""" + ct = pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) + from compressed_tensors.compressors.base import BaseCompressor + + class _PropertyForm: + """A compressor as compressed-tensors 0.13 registers one.""" + + @property + def compression_param_names(self): + return ("weight_packed", "weight_scale") + + monkeypatch.setattr(ct, "__version__", "0.13.0") + monkeypatch.setattr( + BaseCompressor, "get_value_from_registry", staticmethod(lambda name: _PropertyForm) + ) + config = SimpleNamespace(quantization_config=_pack_quantized(moe=_int4(True, EXPERTS))) + + with pytest.raises(RuntimeError, match=r"compressed-tensors >= 0\.16\.0, but 0\.13\.0"): + onfly._read_compression(config) + + # The compressions the published Kimi checkpoints use: INT4 (Kimi-K2.5) and MXFP4 (Kimi-K3), each # with the values its format holds exactly, and a scale it holds exactly. _INT4 = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} @@ -946,7 +975,9 @@ def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp save writes the experts back as the checkpoint's own ``w1``/``w3``/``w2`` weights, decompressed, and the export names its expanded experts after them.""" pytest.importorskip( - "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + "compressed_tensors", + minversion=onfly._MIN_COMPRESSED_TENSORS, + reason="needs compressed-tensors (CI's unit session installs it)", ) from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper @@ -976,7 +1007,9 @@ def test_a_load_the_decompressor_did_not_mark_casts_as_before(tmp_path): """The hooks are process-wide, but only a decompressing load keeps an integer source integer: a later load of a model no config of which was marked casts it to its float parameter.""" pytest.importorskip( - "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + "compressed_tensors", + minversion=onfly._MIN_COMPRESSED_TENSORS, + reason="needs compressed-tensors (CI's unit session installs it)", ) onfly.install_onfly_decompress( SimpleNamespace(quantization_config=_pack_quantized(experts=_int4(True, ["Linear"]))) From f70b6b71e8fd0abf06f4d8877866312749934714 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:48:38 +0000 Subject: [PATCH 22/45] Check the compressed-tensors floor with transformers' require_version The same refusal in four lines instead of seventeen, keeping this PR within its size budget: require_version reads the installed distribution and names the release needed beside the upgrade. The test stands the old release in through the distribution version it reads, and still fails on the property-form compressor without the check. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 29 +++++-------------- .../torch/models/test_onfly_decompress.py | 20 ++++++++----- 2 files changed, 21 insertions(+), 28 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 51e84c4d6522..b88e3208ec6e 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -16,26 +16,18 @@ """Decompress a compressed-tensors MoE checkpoint's experts as ``from_pretrained`` loads them. Each rank decompresses only its own experts, with the ``compressed_tensors`` library, so expert-parallel -PTQ runs on the original checkpoint. Needs transformers 5.8, whose ``Concatenate`` reads only its keys, -and compressed-tensors 0.16, whose compressors decompress one weight by its scheme. +PTQ runs on the original checkpoint. Needs transformers 5.8, whose ``Concatenate`` reads only its keys. """ import functools import os import torch -from packaging.version import parse from safetensors import safe_open from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS, require_transformers -# The first compressed-tensors whose compressors decompress one weight by its scheme, through the -# classmethods compression_param_names(scheme) and decompress(state_dict, scheme). Before it, a -# compressor is built from the whole config, and states its parameter names as a property (to 0.14) -# or not at all (0.15). -_MIN_COMPRESSED_TENSORS = "0.16.0" - def to_native_backbone_config( vlm_config, @@ -104,23 +96,18 @@ def _read_compression(hf_config): if quantization_config is None: return None try: - import compressed_tensors + from compressed_tensors.compressors.base import BaseCompressor + from compressed_tensors.quantization import QuantizationConfig except ImportError as e: # No pip extra of ModelOpt installs it: only a compressed checkpoint needs it. raise ImportError( "On-the-fly decompression of a compressed-tensors checkpoint needs the " - "compressed-tensors package. Install it with: " - f"pip install 'compressed-tensors>={_MIN_COMPRESSED_TENSORS}'" + "compressed-tensors package. Install it with: pip install compressed-tensors" ) from e - have = compressed_tensors.__version__ - if parse(have) < parse(_MIN_COMPRESSED_TENSORS): - raise RuntimeError( - "On-the-fly decompression of a compressed-tensors checkpoint requires " - f"compressed-tensors >= {_MIN_COMPRESSED_TENSORS}, but {have} is installed. Upgrade " - f"with: pip install -U 'compressed-tensors>={_MIN_COMPRESSED_TENSORS}'." - ) - from compressed_tensors.compressors.base import BaseCompressor - from compressed_tensors.quantization import QuantizationConfig + from transformers.utils.versions import require_version + + # From 0.16, a compressor's parameter names and decompress are classmethods taking the scheme. + require_version("compressed-tensors>=0.16", "pip install -U 'compressed-tensors>=0.16'") qconfig = QuantizationConfig.model_validate(quantization_config) groups = qconfig.config_groups diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index ee0019283621..3da219a176d3 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -23,6 +23,7 @@ """ import copy +import importlib.metadata import json import sys from concurrent.futures import Future @@ -847,7 +848,7 @@ def test_a_checkpoint_whose_experts_are_not_its_one_compressed_group_is_refused( experts no group targets are stored as plain weights the rewritten converters would not find.""" pytest.importorskip( "compressed_tensors", - minversion=onfly._MIN_COMPRESSED_TENSORS, + minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) config = SimpleNamespace(quantization_config=_pack_quantized(**groups)) @@ -867,7 +868,7 @@ def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group refused by this module, rather than failing inside the library's registry.""" pytest.importorskip( "compressed_tensors", - minversion=onfly._MIN_COMPRESSED_TENSORS, + minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) experts = _int4(True, EXPERTS) @@ -885,7 +886,7 @@ def test_the_format_is_the_one_group_s_own_and_must_have_a_compressor(top, group def test_a_compressed_tensors_before_per_scheme_compressors_is_refused_by_its_release(monkeypatch): """compressed-tensors 0.13 states a compressor's parameter names as a property, which a call by scheme fails on deep in the setup; such a release is refused first, naming the one needed.""" - ct = pytest.importorskip( + pytest.importorskip( "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" ) from compressed_tensors.compressors.base import BaseCompressor @@ -897,13 +898,18 @@ class _PropertyForm: def compression_param_names(self): return ("weight_packed", "weight_scale") - monkeypatch.setattr(ct, "__version__", "0.13.0") + installed = importlib.metadata.version + monkeypatch.setattr( + importlib.metadata, + "version", + lambda name: "0.13.0" if name == "compressed-tensors" else installed(name), + ) monkeypatch.setattr( BaseCompressor, "get_value_from_registry", staticmethod(lambda name: _PropertyForm) ) config = SimpleNamespace(quantization_config=_pack_quantized(moe=_int4(True, EXPERTS))) - with pytest.raises(RuntimeError, match=r"compressed-tensors >= 0\.16\.0, but 0\.13\.0"): + with pytest.raises(ImportError, match=r"compressed-tensors>=0\.16 is required.*==0\.13\.0"): onfly._read_compression(config) @@ -976,7 +982,7 @@ def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp and the export names its expanded experts after them.""" pytest.importorskip( "compressed_tensors", - minversion=onfly._MIN_COMPRESSED_TENSORS, + minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper @@ -1008,7 +1014,7 @@ def test_a_load_the_decompressor_did_not_mark_casts_as_before(tmp_path): a later load of a model no config of which was marked casts it to its float parameter.""" pytest.importorskip( "compressed_tensors", - minversion=onfly._MIN_COMPRESSED_TENSORS, + minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) onfly.install_onfly_decompress( From 7894a5603d0f4dffa7b38790831baf681e367751 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 20:05:25 +0000 Subject: [PATCH 23/45] Look at a source only in a decompressing load, and import the name mapper at module scope The loader hook probed every source's dtype before reading the flag, so in a load the decompressor never marked it still read each tensor's dtype, and raised on a source with no dtype to read. It now reads the flag first and leaves any other load alone; the hook test hands it such a source. The test's reverse name mapper is imported with the module's other ModelOpt imports, after the transformers gate. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 10 +++++----- tests/unit/torch/models/test_onfly_decompress.py | 5 +++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index b88e3208ec6e..1200032a42e1 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -276,11 +276,11 @@ def _install_keep_int_native(): # Any other argument passes through, positionally or by keyword: a release may add some. def spawn_materialize(thread_pool, tensor, device=None, dtype=None, *args, **kwargs): - # A safetensors slice is judged by its stored dtype, without reading it. - stored = tensor.get_dtype() if hasattr(tensor, "get_dtype") else None - floats = stored.startswith(("F", "BF")) if stored else tensor.is_floating_point() - if _LOAD["decompressing"] and dtype is not None and dtype.is_floating_point and not floats: - dtype = None # keep packed/shape/zero-point/g_idx integers native + if _LOAD["decompressing"] and dtype is not None and dtype.is_floating_point: + # A safetensors slice is judged by its stored dtype, without reading it. + stored = tensor.get_dtype() if hasattr(tensor, "get_dtype") else None + if not (stored.startswith(("F", "BF")) if stored else tensor.is_floating_point()): + dtype = None # keep packed/shape/zero-point/g_idx integers native return orig(thread_pool, tensor, device, dtype, *args, **kwargs) setattr(spawn_materialize, "_onfly_patched", True) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 3da219a176d3..7271a21052c5 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -61,6 +61,7 @@ WeightRenaming, ) +from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper from modelopt.torch.models import _onfly_decompress as onfly PACKED = ("weight_packed", "weight_scale", "weight_shape") # what pack-quantized stores per weight @@ -261,6 +262,8 @@ def test_integer_checkpoint_tensors_keep_their_dtype_on_load(tmp_path): save_file({"packed": PACKED_INT32}, str(tmp_path / "packed.safetensors")) onfly._install_keep_int_native() assert _read(tmp_path / "packed.safetensors").dtype == torch.bfloat16 + # Outside a decompressing load the hook looks at nothing it is handed. + cml.spawn_materialize(None, object(), "cpu", torch.bfloat16) onfly._LOAD["decompressing"] = True kept = _read(tmp_path / "packed.safetensors") @@ -985,8 +988,6 @@ def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) - from modelopt.torch.export.quant_aware_conversion import build_reverse_name_mapper - config, dense = _compressed_mixtral(tmp_path, fmt) onfly.install_onfly_decompress(config) del config.quantization_config # as the expert-parallel load does: no quantizer competes From a652f45a3c33a6b31c18bc8e7f820680d3521c2e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Mon, 2 Nov 2026 23:47:57 +0000 Subject: [PATCH 24/45] Name the reason beside each in-function import of the decompression The repository accepts an import inside a function only for an optional dependency or a circular import, with a comment naming the reason. compressed-tensors and transformers are both optional here: the module imports without them, and core_model_loading is new in transformers 5. Each such import, in the module and in its test, now says so. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 8 +++++++- tests/unit/torch/models/test_onfly_decompress.py | 1 + 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 1200032a42e1..64161f143388 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -40,6 +40,7 @@ def to_native_backbone_config( type and class (hf_ptq's VLM detection reads ``architectures``). The decompression mark carries over. Defaults target the Kimi DeepSeek-V3 backbone. """ + # transformers is optional: this module imports without it. from transformers.models.auto.configuration_auto import CONFIG_MAPPING text_config = getattr(vlm_config, "text_config", None) @@ -95,15 +96,17 @@ def _read_compression(hf_config): quantization_config = _find_quantization_config(hf_config) if quantization_config is None: return None + # Optional: no pip extra of ModelOpt installs compressed-tensors, as only a compressed + # checkpoint needs it. try: from compressed_tensors.compressors.base import BaseCompressor from compressed_tensors.quantization import QuantizationConfig except ImportError as e: - # No pip extra of ModelOpt installs it: only a compressed checkpoint needs it. raise ImportError( "On-the-fly decompression of a compressed-tensors checkpoint needs the " "compressed-tensors package. Install it with: pip install compressed-tensors" ) from e + # transformers is optional: this module imports without it. from transformers.utils.versions import require_version # From 0.16, a compressor's parameter names and decompress are classmethods taking the scheme. @@ -151,6 +154,7 @@ def _make_op_class(compressor, scheme, param_names): It subclasses ``MergeModulelist`` only so the loader shards experts at load, handing it ``None`` for each expert another rank owns; the merge later in the chain still does the stacking. """ + # transformers is optional, and core_model_loading is new in transformers 5. from transformers.core_model_loading import MergeModulelist anchor = "." + param_names[0] @@ -206,6 +210,7 @@ def convert(self, input_dict, *args, **kwargs): @functools.cache def _decompressing_converter_class(): """The ``WeightConverter`` subclass a rewritten converter is, reversing as the one it rewrote.""" + # transformers is optional, and core_model_loading is new in transformers 5. from transformers.core_model_loading import WeightConverter class DecompressingConverter(WeightConverter): @@ -268,6 +273,7 @@ def _install_keep_int_native(): Without a quantizer the loader casts each source to its parameter's dtype in the copy, or from transformers 5.14 in ``DtensorShardOperation.shard_tensor``: both start here. Other loads cast. """ + # transformers is optional, and core_model_loading is new in transformers 5. import transformers.core_model_loading as cml if getattr(cml.spawn_materialize, "_onfly_patched", False): diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 7271a21052c5..87913a8c6422 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -892,6 +892,7 @@ def test_a_compressed_tensors_before_per_scheme_compressors_is_refused_by_its_re pytest.importorskip( "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" ) + # Optional dependency: imported after importorskip so the test skips without it. from compressed_tensors.compressors.base import BaseCompressor class _PropertyForm: From f73538e6598be0fbbd2f72960ddd47dcc7e65e5c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:22:14 +0000 Subject: [PATCH 25/45] Load a checkpoint its model spec decompresses through the expert-parallel loader load_and_prepare_ep refused a compressed-tensors checkpoint at its format check. A config the on-the-fly decompressor marked now passes it, and on that path the load deletes quantization_config from its copy of the config at every level _ckpt_quant_method reads, so transformers builds no compressed-tensors quantizer to compete with the decompressor (from 5.13 it claims the experts' converters itself). Before anything loads, check_compressed_keys checks the names and dtypes the checkpoint stores: the experts alone compressed, every one of them, in floats bf16 holds. Both run inside the loader's collective checks: the format check and the deletion with the config checks, the stored-tensor check right after the checkpoint is made local, so a refusal on one rank fails every rank. Moved here from the on-the-fly decompression change, whose history is archived; the decompression itself is its own change below this one. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 26 ++++++++- .../unit/torch/utils/test_model_load_utils.py | 58 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d33f87d5ba33..d64b95c88b87 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1998,6 +1998,13 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None: return ValueError(refusal) if refusal is not None else None +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): """Return the checkpoint's native quantization method, or ``None``. @@ -2153,7 +2160,7 @@ def load_and_prepare_ep( # checkpoint of the model, nor a plan of its own, would load. check_distributed_support(hf_config.model_type) quant_method = _ckpt_quant_method(hf_config) - 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)) loads = "unquantized checkpoints" + ( @@ -2165,6 +2172,16 @@ def load_and_prepare_ep( ) # 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. Every level _ckpt_quant_method reads is cleared. + for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): + 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: @@ -2224,6 +2241,13 @@ def gpu() -> tuple[bool, bool, str | None]: # refused load downloads nothing. with _refused_together("the local checkpoint", alone=False): _, ckpt_path = ensure_local_checkpoint(ckpt_path, **hub_options) + if _decompressed_as_loaded(hf_config): + # Only the experts decompress as they load, and the config cannot say what else the + # 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, ckpt_path) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 70da8ef18ea3..538d45ced2b1 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2964,6 +2964,64 @@ 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( + 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 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. + 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): + """Stops the load at from_pretrained.""" + + loaded = {} + + 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) + 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])): + 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(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) + 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(str(tmp_path / "experts"), plan, hf_config=config) + _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 + with pytest.raises(_LoadedError): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=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 + + # ----------------------------------------------------------------------- _seed_parallel_states From c44a269b4821b38b2f6f1f2476aa9614779bf712 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:52:05 +0000 Subject: [PATCH 26/45] Refuse a decompressed load whose model leaves compressed tensors unplaced The export copies the checkpoint tensors the model has no parameter for as stored, and the expert-parallel load of a decompressed checkpoint drops the quantization_config that would read a compressed one. An MTP layer's compressed experts, in no layer the model builds, would reach the export packed. The load now refuses once it knows the unplaced tensors, inside its collective checks. _ckpt_quant_method reads the same config levels as the deletion, from _SUB_CONFIG_ATTRS, and the deletion pops from the instance dict. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 18 +++++-- .../unit/torch/utils/test_model_load_utils.py | 53 +++++++++++++++++++ 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index d64b95c88b87..f26001b013bd 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2012,7 +2012,7 @@ def _ckpt_quant_method(cfg): or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's ``text_config`` or ``llm_config``. """ - for c in (cfg, getattr(cfg, "text_config", None), getattr(cfg, "llm_config", None)): + for c in (cfg, *(getattr(cfg, attr, None) for attr in _SUB_CONFIG_ATTRS)): qc = getattr(c, "quantization_config", None) m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) if m: @@ -2180,8 +2180,8 @@ def load_and_prepare_ep( # 5.11 transformers takes 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, attr, None) for attr in _SUB_CONFIG_ATTRS)): - if cfg is not None and hasattr(cfg, "quantization_config"): - del cfg.quantization_config + if cfg is not None: + vars(cfg).pop("quantization_config", None) 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: @@ -2239,6 +2239,7 @@ def gpu() -> tuple[bool, bool, str | None]: dtype = getattr(hf_config, "dtype", None) or torch.bfloat16 # The load and the export both read the local snapshot, made only once the checks pass, so a # refused load downloads nothing. + compressed_names: set[str] = set() with _refused_together("the local checkpoint", alone=False): _, ckpt_path = ensure_local_checkpoint(ckpt_path, **hub_options) if _decompressed_as_loaded(hf_config): @@ -2248,6 +2249,7 @@ def gpu() -> tuple[bool, bool, str | None]: from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). @@ -2324,7 +2326,15 @@ def gpu() -> tuple[bool, bool, str | None]: # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, # for the export to carry. - record_unplaced_source_keys(model, ckpt_path) + unplaced = record_unplaced_source_keys(model, ckpt_path) + # The export would copy them as stored, and this load dropped the quantization_config that + # reads a compressed one, such as an MTP layer's experts. + if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: + raise NotImplementedError( + f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " + f"for, such as {carried[0]}: the export would copy them as stored, with no " + "quantization_config to read them. Decompress the checkpoint first." + ) # The model carries the ep submesh, for calibration to re-install the all-to-all forwards # over (expert_parallel_mesh); transformers sets it differently by release. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 538d45ced2b1..295620dee76a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3022,6 +3022,59 @@ def from_pretrained(path, **kwargs): assert llm.quantization_config is quantization_config +@pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) +def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( + ep_world_of_two, monkeypatch, tmp_path, compressed +): + """The export copies the tensors the model has no parameter for as stored, and this load dropped + the quantization_config that reads a compressed one: an MTP layer's compressed experts are + refused, its uncompressed weights carried as before.""" + from safetensors.torch import save_file + + from modelopt.torch.models import _onfly_decompress + + class _PastError(Exception): + """Raised by the load's next check.""" + + def next_check(*args, **kwargs): + raise _PastError + + model = _model_with_decoder_at("model.") + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + model.config = copy.deepcopy(config) + return model + + mtp = "model.layers.2.mlp.experts.0.up_proj" # past the model's two layers + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} + names = list(stored_as) if compressed else ["weight"] + stored = {f"{m}.{n}": torch.ones(1, dtype=d) for m in (mtp,) for n, d in stored_as.items()} + save_file(stored, str(tmp_path / "model.safetensors")) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr( + model_load_utils, "unplaced_source_keys", lambda *args: [f"{mtp}.{n}" for n in names] + ) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + config = SimpleNamespace( + num_experts=8, + base_model_ep_plan=dict(_EP_PLAN), + quantization_config={"quant_method": "compressed-tensors"}, + ) + _onfly_decompress._mark(config, tuple(stored_as)) + + with pytest.raises(NotImplementedError if compressed else _PastError) as refused: + load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) + if compressed: + assert f"2 compressed tensors the model has no parameter for, such as {mtp}" in str( + refused.value + ) + + # ----------------------------------------------------------------------- _seed_parallel_states From 43689ac98d3a20fa92955a0ca324320783aa5792 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:54:10 +0000 Subject: [PATCH 27/45] Let the loader tests' stand-in unplaced-key record return a list, as the real one does Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 295620dee76a..e5a181cfa428 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2827,7 +2827,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) class _Config(transformers.PretrainedConfig): @@ -2889,7 +2889,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: None) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = Qwen3MoeConfig( hidden_size=16, From d9d4d718aac376cd620e232f513792822d5a98a3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:24:42 +0000 Subject: [PATCH 28/45] Pin that a real decompressed load places its experts' stored tensors The expert-parallel load refuses compressed tensors the model leaves unplaced, and relies on unplaced_source_keys placing the stored weight_packed / weight_scale / weight_shape of every expert the model builds. A CPU from_pretrained of a tiny compressed Qwen3-MoE, through the real decompressor, pins that: they resolve through the converters the load recorded (the stock mapping places none of them), and only the experts of a layer past num_hidden_layers stay unplaced. The two local modelopt.torch.models imports say why they are local, and the loader test's docstring says what its stand-in covers. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 3 +- .../unit/torch/utils/test_model_load_utils.py | 6 +- .../torch/utils/test_unplaced_source_keys.py | 77 +++++++++++++++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index f26001b013bd..ef73516d2940 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2000,6 +2000,7 @@ def _ep_load_refusal(hf_config, ep_plan: dict[str, str]) -> Exception | None: def _decompressed_as_loaded(cfg) -> bool: """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint.""" + # Importing modelopt.torch.models registers every model spec; only expert-parallel loads need them. from modelopt.torch.models import _onfly_decompress return _onfly_decompress.is_marked(cfg) @@ -2245,7 +2246,7 @@ def gpu() -> tuple[bool, bool, str | None]: if _decompressed_as_loaded(hf_config): # Only the experts decompress as they load, and the config cannot say what else the # checkpoint compressed, or in which dtypes: what it stores is checked before anything - # loads. + # loads. Imported here for the reason _decompressed_as_loaded gives. from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index e5a181cfa428..922d6d53e8ae 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3027,8 +3027,10 @@ def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place ep_world_of_two, monkeypatch, tmp_path, compressed ): """The export copies the tensors the model has no parameter for as stored, and this load dropped - the quantization_config that reads a compressed one: an MTP layer's compressed experts are - refused, its uncompressed weights carried as before.""" + the quantization_config that reads a compressed one: an unplaced tensor under a stored name the + mark holds (an MTP layer's compressed experts) is refused, one under the plain ``weight`` name is + not. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in + test_unplaced_source_keys.py.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 46a5e7ef6347..16a9368f23fe 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -391,3 +391,80 @@ def test_mxfp4_expert_blocks_folded_in_at_load_are_placed(tmp_path): model = _load(tmp_path, quantization_config=Mxfp4Config(dequantize=True), dtype=torch.bfloat16) assert model._modelopt_unplaced_source_keys == [STRAY] + + +def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatch): + """The on-the-fly decompressor feeds each expert's stored ``weight_packed``/``weight_scale``/ + ``weight_shape`` into the fused weights. A built layer's are placed, through the converters + ``from_pretrained`` recorded (the stock mapping places none of them), so the expert-parallel load + passes them; only a layer the model does not build (an MTP layer) keeps its own, which that load + refuses to carry.""" + pytest.importorskip( + "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + ) + import transformers.core_model_loading as cml + import transformers.modeling_utils as mu + from packaging.version import Version + from transformers import __version__ as transformers_version + + from modelopt.torch.models import _onfly_decompress + + if Version(transformers_version) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers >= 5.8") + # The decompressor's hooks are process-wide: undone after the test. + for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + monkeypatch.setattr(owner, name, getattr(owner, name)) + hidden, inter, experts = 64, 32, 2 + config = AutoConfig.for_model( + "qwen3_moe", + vocab_size=64, + hidden_size=hidden, + intermediate_size=64, + moe_intermediate_size=inter, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_experts=experts, + num_experts_per_tok=1, + tie_word_embeddings=False, + architectures=["Qwen3MoeForCausalLM"], + ) + group = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} + config.quantization_config = { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "config_groups": {"experts": {"targets": ["re:.*experts.*"], "weights": group}}, + } + compressor, scheme, stored_as, _ = _onfly_decompress._read_compression(config) + shapes = { + "gate_proj": (inter, hidden), + "up_proj": (inter, hidden), + "down_proj": (hidden, inter), + } + state, mtp = {}, [] + for layer in (0, 1): # layer 1 is past num_hidden_layers: an MTP layer, say + for e in range(experts): + for proj, (rows, cols) in shapes.items(): + weight = { + "weight": torch.ones(rows, cols), + "weight_scale": torch.ones(rows, cols // 32), + } + stem = f"model.layers.{layer}.mlp.experts.{e}.{proj}." + for name, tensor in compressor.compress(weight, scheme).items(): + state[stem + name] = tensor + if layer: + mtp.append(stem + name) + for name, tensor in AutoModelForCausalLM.from_config(config).state_dict().items(): + if ".experts." not in name: + state[name] = tensor.to(torch.bfloat16) + _save(tmp_path, config, state, config.quantization_config) + _onfly_decompress.install_onfly_decompress(config) + del config.quantization_config # as the expert-parallel load does once the hooks are in + + model = AutoModelForCausalLM.from_pretrained(str(tmp_path), config=config, dtype=torch.bfloat16) + if getattr(model, "_weight_conversions", None) is None: + pytest.skip("this transformers does not record the conversions from_pretrained ran") + + assert sorted(stored_as) == ["weight_packed", "weight_scale", "weight_shape"] + assert model_load_utils.unplaced_source_keys(model, str(tmp_path)) == sorted(mtp) From d6e7e00cc3ddb8ee9e851a0df4f16c45aeb2e098 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:42:15 +0000 Subject: [PATCH 29/45] Load compressed experts expert-parallel on 2 GPUs, each rank decompressing its own A tiny Qwen3-MoE whose experts are compressed-tensors INT4 (Kimi-K2.5's format) or MXFP4 (Kimi-K3's), installed for on-the-fly decompression and loaded through load_and_prepare_ep: each rank must hold exactly the dense model's weights for its own experts, every stored tensor must be placed, and the sharded model's logits must match the dense model loaded whole. The MXFP4 case goes through compressed-tensors, which the decompressor reads, not transformers' own MXFP4 dequantization. Each rank's checks are gathered before any rank fails, under bounded collectives. Signed-off-by: Shengliang Xu --- .../torch/utils/test_expert_parallel_load.py | 131 +++++++++++++++++- 1 file changed, 130 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 849305f4f6a8..1bdcbb4c9577 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -17,20 +17,30 @@ ``load_and_prepare_ep`` shards an MoE's routed experts over the ranks, and the ``torch_a2a`` experts forward sends each token to the rank that holds its expert. Every rank runs its own batch through the -sharded model, and the result must match the whole model run on that batch alone. +sharded model, and the result must match the whole model run on that batch alone. A checkpoint whose +experts are compressed-tensors compressed loads the same way, each rank decompressing its own. """ +import json +from types import SimpleNamespace + import pytest import torch +import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version +from safetensors.torch import load_file, save_file transformers = pytest.importorskip("transformers") if Version(transformers.__version__) < Version("5.7"): pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7", allow_module_level=True) +import transformers.core_model_loading as cml +import transformers.modeling_utils as mu from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +from modelopt.torch.models import _onfly_decompress from modelopt.torch.utils.distributed import ( DistributionPlan, distribution_plan, @@ -101,6 +111,125 @@ def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree): torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2) +# The compressions the published Kimi checkpoints use, INT4 (Kimi-K2.5) and MXFP4 (Kimi-K3), through +# compressed-tensors (quant_method "compressed-tensors"), which the on-the-fly decompressor reads; not +# transformers' own MXFP4 (quant_method "mxfp4"). The experts draw values each format holds exactly, +# at a power-of-two scale bf16 holds too, so they decompress to the dense model's weights exactly. +_INT4 = {"num_bits": 4, "type": "int", "symmetric": True, "strategy": "group", "group_size": 32} +_E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6]) +_COMPRESSIONS = { + "int4": ("pack-quantized", _INT4, torch.arange(-8, 8) / 64, 1 / 64), + "mxfp4": ( + "mxfp4-pack-quantized", + {**_INT4, "type": "float", "scale_dtype": "torch.uint8"}, + torch.cat([_E2M1, -_E2M1[1:]]) / 16, + 1 / 16, + ), +} + + +def _save_compressed_tiny_qwen3_moe(dense_dir, packed_dir, compression): + """The tiny Qwen3-MoE twice: whole in bf16, and with its experts compressed by ``compression``.""" + fmt, weights, values, scale = _COMPRESSIONS[compression] + _save_tiny_qwen3_moe(dense_dir) + state = load_file(str(dense_dir / "model.safetensors")) + gen = torch.Generator().manual_seed(0) + for key in [k for k in state if ".experts." in k]: # ...mlp.experts..gate_proj.weight + draw = torch.randint(0, len(values), state[key].shape, generator=gen) + state[key] = values[draw].to(torch.bfloat16) + save_file(state, str(dense_dir / "model.safetensors"), metadata={"format": "pt"}) + + quantization_config = { + "quant_method": "compressed-tensors", + "format": fmt, + "config_groups": {"experts": {"targets": ["Linear"], "weights": weights, "format": fmt}}, + } + compressor, scheme, _, _ = _onfly_decompress._read_compression( + SimpleNamespace(quantization_config=quantization_config) + ) + packed = {} + for key, tensor in state.items(): + if ".experts." not in key: + packed[key] = tensor + continue + rows, cols = tensor.shape + scales = torch.full((rows, cols // 32), scale) + stored = compressor.compress({"weight": tensor.float(), "weight_scale": scales}, scheme) + for name, value in stored.items(): # floats as bf16, as the Kimi checkpoints store them + packed[key.removesuffix("weight") + name] = ( + value.bfloat16() if value.is_floating_point() else value + ) + packed_dir.mkdir() + save_file(packed, str(packed_dir / "model.safetensors"), metadata={"format": "pt"}) + config = json.loads((dense_dir / "config.json").read_text()) + config["quantization_config"] = quantization_config + (packed_dir / "config.json").write_text(json.dumps(config)) + + +def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_dir): + device = torch.device(f"cuda:{rank}") + problems = [] + # The decompressor's hooks are process-wide, and the workers outlive the job. + hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + try: + with bounded_collectives(seconds=120): + config = AutoConfig.from_pretrained(packed_dir) + _onfly_decompress.install_onfly_decompress(config) # as a model spec's prepare does + model = load_and_prepare_ep( + packed_dir, DistributionPlan(size, ep_size=size), hf_config=config + ).eval() + whole = AutoModelForCausalLM.from_pretrained(dense_dir, dtype=torch.bfloat16) + whole = whole.to(device).eval() + + # This rank's experts, decompressed as they loaded, are the dense model's, exactly. + per_rank = NUM_EXPERTS // size + owned = slice(rank * per_rank, (rank + 1) * per_rank) + for i, (layer, whole_layer) in enumerate(zip(model.model.layers, whole.model.layers)): + for name in ("gate_up_proj", "down_proj"): + local = _local(getattr(layer.mlp.experts, name)) + if local.shape[0] != per_rank: + problems.append(f"layer {i} {name}: {local.shape[0]} experts held") + elif not torch.equal(local, getattr(whole_layer.mlp.experts, name)[owned]): + problems.append(f"layer {i} {name}: rank {rank}'s experts differ") + # Every stored tensor of the experts it built went into them. + if model._modelopt_unplaced_source_keys: + problems.append(f"unplaced: {model._modelopt_unplaced_source_keys[:3]}") + set_distribution_plan(model, DistributionPlan(size, ep_size=size)) + + generator = torch.Generator().manual_seed(rank) + input_ids = torch.randint(0, 128, (2, 16), generator=generator).to(device) + with torch.no_grad(): + sharded = model(input_ids).logits.float() + reference = whole(input_ids).logits.float() + if not torch.allclose(sharded, reference, atol=2e-2, rtol=2e-2): + problems.append(f"logits differ by {(sharded - reference).abs().max().item():.3g}") + + gathered = [None] * size + dist.all_gather_object(gathered, problems) + finally: + mu.get_model_conversion_mapping, cml._materialize_copy = hooks + found = [problem for rank_problems in gathered for problem in rank_problems] + assert not found, "\n".join(found[:20]) + + +@pytest.mark.parametrize("compression", sorted(_COMPRESSIONS)) +def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own( + need_2_gpus, dist_workers, tmp_path, compression +): + """A checkpoint whose experts are compressed (INT4, MXFP4) loads expert-parallel through the + on-the-fly decompressor: each rank holds exactly the dense model's weights for its own experts, + every stored tensor is placed, and the sharded model matches the dense one loaded whole. Each + rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" + pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + if Version(transformers.__version__) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers 5.8") + if NUM_EXPERTS % dist_workers.world_size: + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {dist_workers.world_size} GPUs") + dense_dir, packed_dir = tmp_path / "dense", tmp_path / "packed" + _save_compressed_tiny_qwen3_moe(dense_dir, packed_dir, compression) + dist_workers.run(_compressed_experts_load_on_their_own_ranks, str(dense_dir), str(packed_dir)) + + @pytest.mark.parametrize("ep_share", [1, 2], ids=["ep-world", "ep-half"]) def test_expert_parallel_model_matches_the_whole_model( need_2_gpus, dist_workers, tmp_path, ep_share From 6aa19c0d093a52cd77b9bf5199231bb75587e3f1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:54:03 +0000 Subject: [PATCH 30/45] Restore the decompressor's read hook where the tests undo its hooks The decompressor's keep-integers patch moved from _materialize_copy to spawn_materialize (#3137), so the 2-GPU job and the real-load unplaced test restore that hook. The GPU case's skip reason names the package it needs: compressed-tensors is not in CI's GPU image, where it skips. Signed-off-by: Shengliang Xu --- tests/gpu/torch/utils/test_expert_parallel_load.py | 9 ++++++--- tests/unit/torch/utils/test_unplaced_source_keys.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 1bdcbb4c9577..92ac059cb10a 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -170,7 +170,7 @@ def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_di device = torch.device(f"cuda:{rank}") problems = [] # The decompressor's hooks are process-wide, and the workers outlive the job. - hooks = (mu.get_model_conversion_mapping, cml._materialize_copy) + hooks = (mu.get_model_conversion_mapping, cml.spawn_materialize) try: with bounded_collectives(seconds=120): config = AutoConfig.from_pretrained(packed_dir) @@ -207,7 +207,7 @@ def _compressed_experts_load_on_their_own_ranks(rank, size, dense_dir, packed_di gathered = [None] * size dist.all_gather_object(gathered, problems) finally: - mu.get_model_conversion_mapping, cml._materialize_copy = hooks + mu.get_model_conversion_mapping, cml.spawn_materialize = hooks found = [problem for rank_problems in gathered for problem in rank_problems] assert not found, "\n".join(found[:20]) @@ -220,7 +220,10 @@ def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own on-the-fly decompressor: each rank holds exactly the dense model's weights for its own experts, every stored tensor is placed, and the sharded model matches the dense one loaded whole. Each rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" - pytest.importorskip("compressed_tensors", reason="the checkpoint's experts are compressed") + pytest.importorskip( + "compressed_tensors", + reason="compressed-tensors is needed: the checkpoint's experts are compressed", + ) if Version(transformers.__version__) < Version("5.8"): pytest.skip("the on-the-fly decompression needs transformers 5.8") if NUM_EXPERTS % dist_workers.world_size: diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 16a9368f23fe..4fd5da334061 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -412,7 +412,7 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc if Version(transformers_version) < Version("5.8"): pytest.skip("the on-the-fly decompression needs transformers >= 5.8") # The decompressor's hooks are process-wide: undone after the test. - for owner, name in ((cml, "_materialize_copy"), (mu, "get_model_conversion_mapping")): + for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) hidden, inter, experts = 64, 32, 2 config = AutoConfig.for_model( From c6e98e3e8183254c63dc282b68fe527356cc87b8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:22:45 +0000 Subject: [PATCH 31/45] Leave no decompressing load flag behind the real-load unplaced test The real decompressing load sets the decompressor's load flag, which outlives the load; the test now runs it on a flag of its own. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_unplaced_source_keys.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 4fd5da334061..f3b89452821a 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -414,6 +414,7 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc # The decompressor's hooks are process-wide: undone after the test. for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) + monkeypatch.setattr(_onfly_decompress, "_LOAD", {"decompressing": False}) hidden, inter, experts = 64, 32, 2 config = AutoConfig.for_model( "qwen3_moe", From 8e53be7e9e1b94aaacc9a540563883e5d1835151 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:32:40 +0000 Subject: [PATCH 32/45] Refuse a config whose class keeps quantization_config, and say why the matchers hold From Claude's review: the load clears quantization_config from each config level's instance dict, which leaves a class attribute or property in place, and transformers' own quantizer would then compete with the decompressor. The load now refuses such a config inside its collective checks. Comments say why "weight" is left out of the compressed names (a float format's anchor, and a plain weight the export has always carried) and why a suffix tells a carried compressed tensor (check_compressed_keys refused any outside the experts). Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 12 +++++++++++- tests/unit/torch/utils/test_model_load_utils.py | 12 +++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index ef73516d2940..b30cd6971a39 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2183,6 +2183,13 @@ def load_and_prepare_ep( for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): if cfg is not None: vars(cfg).pop("quantization_config", None) + if hasattr(cfg, "quantization_config"): + # A class attribute or property, which the instance dict cannot clear. + raise NotImplementedError( + f"{type(cfg).__name__} keeps quantization_config off its instance, so " + "this load cannot clear it, and transformers' own quantizer would " + "compete with the decompressor. Decompress the checkpoint first." + ) 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: @@ -2250,6 +2257,8 @@ def gpu() -> tuple[bool, bool, str | None]: from modelopt.torch.models import _onfly_decompress _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight + # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over @@ -2329,7 +2338,8 @@ def gpu() -> tuple[bool, bool, str | None]: # for the export to carry. unplaced = record_unplaced_source_keys(model, ckpt_path) # The export would copy them as stored, and this load dropped the quantization_config that - # reads a compressed one, such as an MTP layer's experts. + # reads a compressed one, such as an MTP layer's experts. A suffix is enough to tell them: + # check_compressed_keys refused any compressed tensor outside the experts. if carried := [key for key in unplaced if key.rpartition(".")[2] in compressed_names]: raise NotImplementedError( f"The checkpoint stores {len(carried)} compressed tensors the model has no parameter " diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 922d6d53e8ae..23763e2b273a 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2971,7 +2971,8 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( 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. - So is one that compresses more than its experts, by the names it stores, before it loads.""" + So is one that compresses more than its experts, by the names it stores, before it loads, and a + config whose class keeps quantization_config, which the load cannot clear.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress @@ -3021,6 +3022,15 @@ def from_pretrained(path, **kwargs): assert config.quantization_config is text.quantization_config is quantization_config assert llm.quantization_config is quantization_config + # A config class that keeps quantization_config as a class attribute cannot have it cleared. + typed_config = type("_Typed", (SimpleNamespace,), {"quantization_config": quantization_config}) + typed = typed_config(text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN)) + _onfly_decompress._mark(typed, tuple(stored_as)) + loaded.clear() + with pytest.raises(NotImplementedError, match="keeps quantization_config off its instance"): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=typed) + assert not loaded + @pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( From 1e9f214ab209c7d5d3da0b8a7eddeae4b9625c38 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:46:25 +0000 Subject: [PATCH 33/45] Walk a config's levels through one helper, and check them once all are cleared From Claude's review: the top config and its nested language-model configs were walked three times, in _ckpt_quant_method, in the load's quantization_config deletion and in the decompressor. moe.config_levels, beside _SUB_CONFIG_ATTRS, now gives those levels to all three. (#3135 adds the same _SUB_CONFIG_ATTRS; the helper beside it conflicts there once, at the Kimi base's rebuild.) The deletion now clears every level before it checks any: a config class that forwards its text config's quantization_config through a property holds none once the text config's is gone, and is no longer refused; a class attribute, or a property of its own, still is. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 18 +++++--------- modelopt/torch/utils/moe.py | 6 +++++ .../torch/utils/plugins/model_load_utils.py | 24 ++++++++++--------- .../unit/torch/utils/test_model_load_utils.py | 15 +++++++++++- 4 files changed, 39 insertions(+), 24 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index 64161f143388..d7f621d0e28f 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -26,7 +26,7 @@ from safetensors import safe_open from modelopt.torch.utils import print_rank_0 -from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS, require_transformers +from modelopt.torch.utils.moe import config_levels, require_transformers def to_native_backbone_config( @@ -61,19 +61,13 @@ 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.""" - 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, 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) + for c in config_levels(hf_config) if (found := getattr(c, "quantization_config", None)) is not None ), None, @@ -306,12 +300,12 @@ def _mark(hf_config, stored_as): def is_marked(hf_config) -> bool: """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)) + return any(getattr(c, _MARK, None) is not None for c in config_levels(hf_config)) def stored_names(hf_config) -> tuple[str, ...]: """The names the installed compressed-tensors stores each weight under for ``hf_config``'s scheme.""" - for c in _config_levels(hf_config): + for c in config_levels(hf_config): if (stored_as := getattr(c, _MARK, None)) is not None: return stored_as raise ValueError("The config is not marked for on-the-fly decompression: install it first.") @@ -329,7 +323,7 @@ def install_onfly_decompress(hf_config): raise TypeError("install_onfly_decompress: pass the model's config object, not a dict") # The loader sets each parameter from what the conversion hands it, uncast. Every level counts: # a backbone config derived from the text config (to_native_backbone_config) keeps its dtype. - for level in _config_levels(hf_config): + for level in config_levels(hf_config): dtype = getattr(level, "dtype", None) if dtype is not None and str(dtype).removeprefix("torch.") != "bfloat16": raise NotImplementedError( @@ -399,7 +393,7 @@ def patched(*args, **kwargs): ) # 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): + for c in config_levels(config): if hasattr(c, _MARK): delattr(c, _MARK) return out diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 15de325fd9f7..72e3579c4c1c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -64,6 +64,12 @@ def require_transformers(min_version: str, feature: str) -> None: _SUB_CONFIG_ATTRS = ("text_config", "llm_config") +def config_levels(hf_config) -> list: + """``hf_config`` and its nested language-model configs, those it has, which may carry its settings.""" + 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 _ep_config(hf_config): """Return the config that carries the EP plan / expert count. diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index b30cd6971a39..2c54d18b08ae 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -82,6 +82,7 @@ _spec_refuses_ep, apply_distributed_support, check_distributed_support, + config_levels, register_torch_a2a, ) from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -2013,7 +2014,7 @@ def _ckpt_quant_method(cfg): or ``None`` if the checkpoint isn't natively quantized. Checks the top config, then a VLM's ``text_config`` or ``llm_config``. """ - for c in (cfg, *(getattr(cfg, attr, None) for attr in _SUB_CONFIG_ATTRS)): + for c in config_levels(cfg): qc = getattr(c, "quantization_config", None) m = qc.get("quant_method") if isinstance(qc, dict) else getattr(qc, "quant_method", None) if m: @@ -2180,16 +2181,17 @@ def load_and_prepare_ep( # 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. Every level _ckpt_quant_method reads is cleared. - for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): - if cfg is not None: - vars(cfg).pop("quantization_config", None) - if hasattr(cfg, "quantization_config"): - # A class attribute or property, which the instance dict cannot clear. - raise NotImplementedError( - f"{type(cfg).__name__} keeps quantization_config off its instance, so " - "this load cannot clear it, and transformers' own quantizer would " - "compete with the decompressor. Decompress the checkpoint first." - ) + levels = config_levels(hf_config) + for cfg in levels: + vars(cfg).pop("quantization_config", None) + # Checked once every level is cleared: a property forwarding to a nested level's is gone + # then, while a class attribute, or a property of its own, is not. + if kept := [cfg for cfg in levels if hasattr(cfg, "quantization_config")]: + raise NotImplementedError( + f"{type(kept[0]).__name__} keeps quantization_config off its instance, so this " + "load cannot clear it, and transformers' own quantizer would compete with the " + "decompressor. Decompress the checkpoint first." + ) 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/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 23763e2b273a..94cdb0774122 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2972,7 +2972,8 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( 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. So is one that compresses more than its experts, by the names it stores, before it loads, and a - config whose class keeps quantization_config, which the load cannot clear.""" + config whose class keeps quantization_config, which the load cannot clear; one whose class only + forwards its text config's loads.""" from safetensors.torch import save_file from modelopt.torch.models import _onfly_decompress @@ -3031,6 +3032,18 @@ def from_pretrained(path, **kwargs): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=typed) assert not loaded + # One whose class forwards the text config's, as Kimi-K2.5's copies it, loads once both are gone. + class _Forwarding(SimpleNamespace): + @property + def quantization_config(self): + return self.text_config.quantization_config + + text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) + forwarding = _Forwarding(text_config=text) + _onfly_decompress._mark(forwarding, tuple(stored_as)) + with pytest.raises(_LoadedError): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) + @pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( From 4d4e412a9bf3c0da06e51ef4525dd489e3ae5340 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:10:13 +0000 Subject: [PATCH 34/45] Build the unplaced-refusal loader test's model from the shared test model The EP loader tests now build their stand-in MoE from the shared model_with_decoder_at, and the file's own copy is gone. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 94cdb0774122..61c49c18796f 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3064,7 +3064,7 @@ class _PastError(Exception): def next_check(*args, **kwargs): raise _PastError - model = _model_with_decoder_at("model.") + model = model_with_decoder_at("model.") class _Loader: @staticmethod From a674670e1153e6359ef69130416a210ba7ba01a3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:13:02 +0000 Subject: [PATCH 35/45] Refuse an MTP layer's compressed experts before the load, not after it From Claude's review: the load refused the compressed tensors its model leaves unplaced only once from_pretrained had read the whole checkpoint. An MTP layer, the usual case, is stored past the decoder layers the EP config says the model builds, and the checkpoint's tensor names show that before anything loads. Its compressed experts are now refused there; the check after the load stays for any other shape. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 19 +++++++++++++++ .../unit/torch/utils/test_model_load_utils.py | 23 +++++++++++++++---- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 2c54d18b08ae..3efdbdf13c36 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2262,6 +2262,25 @@ def gpu() -> tuple[bool, bool, str | None]: # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} + # An MTP layer is stored past the decoder layers the model builds, so its compressed + # experts would be refused only after the whole load, by the check that follows it. + # Refused here instead, before anything loads; that check stays for any other shape. + built = getattr(ep_cfg, "num_hidden_layers", None) + past = sorted( + name + for name in checkpoint_tensor_files(ckpt_path) + if name.rpartition(".")[2] in compressed_names + and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) + and built is not None + and int(layer[1]) >= built + ) + if past: + raise NotImplementedError( + f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " + f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " + "for them, and the export would copy them as stored, with no quantization_config " + "to read them. Decompress the checkpoint first." + ) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 61c49c18796f..b26b8221054b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3045,14 +3045,19 @@ def quantization_config(self): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) -@pytest.mark.parametrize("compressed", [False, True], ids=["uncompressed", "compressed"]) +@pytest.mark.parametrize( + ("compressed", "layers"), + [(False, 3), (True, 3), (True, 2)], + ids=["uncompressed", "compressed", "compressed-past-the-layers"], +) def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place( - ep_world_of_two, monkeypatch, tmp_path, compressed + ep_world_of_two, monkeypatch, tmp_path, compressed, layers ): """The export copies the tensors the model has no parameter for as stored, and this load dropped the quantization_config that reads a compressed one: an unplaced tensor under a stored name the mark holds (an MTP layer's compressed experts) is refused, one under the plain ``weight`` name is - not. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in + not. Compressed experts stored past the config's decoder layers are refused before anything + loads. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in test_unplaced_source_keys.py.""" from safetensors.torch import save_file @@ -3065,14 +3070,16 @@ def next_check(*args, **kwargs): raise _PastError model = model_with_decoder_at("model.") + loads = [] class _Loader: @staticmethod def from_pretrained(path, config, **kwargs): + loads.append(path) model.config = copy.deepcopy(config) return model - mtp = "model.layers.2.mlp.experts.0.up_proj" # past the model's two layers + mtp = "model.layers.2.mlp.experts.0.up_proj" # unplaced, as an MTP layer's are stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} names = list(stored_as) if compressed else ["weight"] stored = {f"{m}.{n}": torch.ones(1, dtype=d) for m in (mtp,) for n, d in stored_as.items()} @@ -3087,6 +3094,7 @@ def from_pretrained(path, config, **kwargs): monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( num_experts=8, + num_hidden_layers=layers, base_model_ep_plan=dict(_EP_PLAN), quantization_config={"quant_method": "compressed-tensors"}, ) @@ -3094,7 +3102,12 @@ def from_pretrained(path, config, **kwargs): with pytest.raises(NotImplementedError if compressed else _PastError) as refused: load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) - if compressed: + if layers == 2: + assert not loads + assert f"2 compressed tensors past the model's 2 decoder layers, such as {mtp}." in str( + refused.value + ) + elif compressed: assert f"2 compressed tensors the model has no parameter for, such as {mtp}" in str( refused.value ) From 5051122ad3f225a68291fbf809295e52cfea5d98 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:36:01 +0000 Subject: [PATCH 36/45] Refuse an MTP layer's experts from the names the compressed-key check read From Claude's review: the early refusal read the checkpoint's tensor names and every shard header a second time, right after check_compressed_keys had, and scanned them even for a config with no num_hidden_layers. It now takes the map check_compressed_keys returns, and skips the scan without a layer count. Its message says what it found, an MTP layer's typically, rather than assert the cause. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3efdbdf13c36..da459e3aebf6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2258,28 +2258,33 @@ def gpu() -> tuple[bool, bool, str | None]: # loads. Imported here for the reason _decompressed_as_loaded gives. from modelopt.torch.models import _onfly_decompress - _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) + tensor_files = _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) # Minus "weight": a float format's anchor is .weight itself, and an unplaced plain weight # is what the export has always carried as stored. compressed_names = set(_onfly_decompress.stored_names(hf_config)) - {"weight"} # An MTP layer is stored past the decoder layers the model builds, so its compressed # experts would be refused only after the whole load, by the check that follows it. - # Refused here instead, before anything loads; that check stays for any other shape. + # Refused here instead, before anything loads, from the names just read; that check + # stays for any other shape. transformers 5.18 builds no decoder layer past + # num_hidden_layers, so this refuses nothing that check would let through there. built = getattr(ep_cfg, "num_hidden_layers", None) - past = sorted( - name - for name in checkpoint_tensor_files(ckpt_path) - if name.rpartition(".")[2] in compressed_names - and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) - and built is not None - and int(layer[1]) >= built + past = ( + [] + if built is None + else sorted( + name + for name in tensor_files + if name.rpartition(".")[2] in compressed_names + and (layer := re.search(r"(?:^|\.)layers\.(\d+)\.", name)) + and int(layer[1]) >= built + ) ) if past: raise NotImplementedError( f"The checkpoint stores {len(past)} compressed tensors past the model's {built} " - f"decoder layers, such as {past[0]} (an MTP layer): the model has no parameter " - "for them, and the export would copy them as stored, with no quantization_config " - "to read them. Decompress the checkpoint first." + f"decoder layers, such as {past[0]} (an MTP layer's, typically): the model builds " + "no parameter for them, and the export would copy them as stored, with no " + "quantization_config to read them. Decompress the checkpoint first." ) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over From cb8ca2c5ff0f7301648c205e59bb0aeff975a103 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 04:24:59 +0000 Subject: [PATCH 37/45] Give the decompressing loader tests' stand-in configs a model type no spec names The expert-parallel load is to read the config's model type for its spec checks, so each config standing in for a checkpoint's names one, as the moved loader tests' do. save_file is already imported at the top of the file the tests moved into. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b26b8221054b..34070f847c0c 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2974,8 +2974,6 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( So is one that compresses more than its experts, by the names it stores, before it loads, and a config whose class keeps quantization_config, which the load cannot clear; one whose class only forwards its text config's loads.""" - from safetensors.torch import save_file - from modelopt.torch.models import _onfly_decompress class _LoadedError(Exception): @@ -3005,7 +3003,10 @@ def from_pretrained(path, **kwargs): text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) llm = SimpleNamespace(quantization_config=quantization_config) config = SimpleNamespace( - quantization_config=quantization_config, text_config=text, llm_config=llm + model_type=_NO_SPEC, + quantization_config=quantization_config, + text_config=text, + llm_config=llm, ) plan = DistributionPlan(2, ep_size=2) @@ -3025,7 +3026,9 @@ def from_pretrained(path, **kwargs): # A config class that keeps quantization_config as a class attribute cannot have it cleared. typed_config = type("_Typed", (SimpleNamespace,), {"quantization_config": quantization_config}) - typed = typed_config(text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN)) + typed = typed_config( + model_type=_NO_SPEC, text_config=SimpleNamespace(base_model_ep_plan=_EP_PLAN) + ) _onfly_decompress._mark(typed, tuple(stored_as)) loaded.clear() with pytest.raises(NotImplementedError, match="keeps quantization_config off its instance"): @@ -3039,7 +3042,7 @@ def quantization_config(self): return self.text_config.quantization_config text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) - forwarding = _Forwarding(text_config=text) + forwarding = _Forwarding(model_type=_NO_SPEC, text_config=text) _onfly_decompress._mark(forwarding, tuple(stored_as)) with pytest.raises(_LoadedError): load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=forwarding) @@ -3059,8 +3062,6 @@ def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place not. Compressed experts stored past the config's decoder layers are refused before anything loads. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in test_unplaced_source_keys.py.""" - from safetensors.torch import save_file - from modelopt.torch.models import _onfly_decompress class _PastError(Exception): @@ -3093,6 +3094,7 @@ def from_pretrained(path, config, **kwargs): ) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( + model_type=_NO_SPEC, num_experts=8, num_hidden_layers=layers, base_model_ep_plan=dict(_EP_PLAN), From 4101987a39d1afaf02c500b11f250beffc15f4ec Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 19:49:59 +0000 Subject: [PATCH 38/45] Skip the compressed-expert load tests below compressed-tensors 0.16 Signed-off-by: Shengliang Xu #3137 now refuses an older compressed-tensors, whose compressors take no scheme: the tests that load through the real library skip there, as the decompressor's own do, rather than fail on the refusal. --- tests/gpu/torch/utils/test_expert_parallel_load.py | 1 + tests/unit/torch/utils/test_unplaced_source_keys.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 92ac059cb10a..533147ac1b03 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -222,6 +222,7 @@ def test_compressed_experts_load_expert_parallel_each_rank_decompressing_its_own rank's checks are gathered before any rank fails, so a regression fails instead of hanging.""" pytest.importorskip( "compressed_tensors", + minversion="0.16", reason="compressed-tensors is needed: the checkpoint's experts are compressed", ) if Version(transformers.__version__) < Version("5.8"): diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index f3b89452821a..13fa3600e091 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -400,7 +400,9 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc passes them; only a layer the model does not build (an MTP layer) keeps its own, which that load refuses to carry.""" pytest.importorskip( - "compressed_tensors", reason="needs compressed-tensors (CI's unit session installs it)" + "compressed_tensors", + minversion="0.16", + reason="needs compressed-tensors (CI's unit session installs it)", ) import transformers.core_model_loading as cml import transformers.modeling_utils as mu From 1661bcfe5885c54cd984312af7b9570a8392be77 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 01:17:29 +0000 Subject: [PATCH 39/45] Import the decompressor at module scope in the loader tests Two loader tests imported the decompressor inside the test, and the unplaced-keys test also packaging's Version and transformers' version, though nothing makes them optional there: those files import transformers at module scope, and the decompressor imports nothing optional. They are module imports now, so an import error shows at collection. Only core_model_loading stays imported in its test, after the version check, with the reason beside it: transformers 4 has none. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 4 +--- tests/unit/torch/utils/test_unplaced_source_keys.py | 12 ++++++------ 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 34070f847c0c..b62a845b971b 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -75,7 +75,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist import modelopt.torch.utils.moe as moe -from modelopt.torch.models import get_spec +from modelopt.torch.models import _onfly_decompress, get_spec from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -2974,7 +2974,6 @@ def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( So is one that compresses more than its experts, by the names it stores, before it loads, and a config whose class keeps quantization_config, which the load cannot clear; one whose class only forwards its text config's loads.""" - from modelopt.torch.models import _onfly_decompress class _LoadedError(Exception): """Stops the load at from_pretrained.""" @@ -3062,7 +3061,6 @@ def test_a_decompressed_load_refuses_compressed_tensors_its_model_does_not_place not. Compressed experts stored past the config's decoder layers are refused before anything loads. Which keys go unplaced is stood in for here; a real decompressed load's are pinned in test_unplaced_source_keys.py.""" - from modelopt.torch.models import _onfly_decompress class _PastError(Exception): """Raised by the load's next check.""" diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 13fa3600e091..274c04cc9e66 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -25,6 +25,7 @@ import pytest import torch +from packaging.version import Version from safetensors.torch import load_file, save_file pytest.importorskip("accelerate") @@ -40,7 +41,9 @@ LlamaForCausalLM, Mxfp4Config, ) +from transformers import __version__ as transformers_version +from modelopt.torch.models import _onfly_decompress from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( @@ -404,15 +407,12 @@ def test_compressed_experts_decompressed_at_load_are_placed(tmp_path, monkeypatc minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) + if Version(transformers_version) < Version("5.8"): + pytest.skip("the on-the-fly decompression needs transformers >= 5.8") + # Imported once the release is known to have them: transformers 4 has no core_model_loading. import transformers.core_model_loading as cml import transformers.modeling_utils as mu - from packaging.version import Version - from transformers import __version__ as transformers_version - - from modelopt.torch.models import _onfly_decompress - if Version(transformers_version) < Version("5.8"): - pytest.skip("the on-the-fly decompression needs transformers >= 5.8") # The decompressor's hooks are process-wide: undone after the test. for owner, name in ((cml, "spawn_materialize"), (mu, "get_model_conversion_mapping")): monkeypatch.setattr(owner, name, getattr(owner, name)) From c92e23b951736a796c634d5fa27d251c3cd56cdd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 01:40:07 +0000 Subject: [PATCH 40/45] Move the native backbone config helper out to the Kimi adapters' PR to_native_backbone_config rebuilds a VLM's bundled text config as a native backbone config: the Kimi adapters' work, not the decompression's, and nothing in this change calls it. It moves, with its tests, to the change that adds the adapters calling it. The test that loads a VLM marked at either level keeps its VLM case here; the case that loads the native backbone on its own moves with the helper. Signed-off-by: Shengliang Xu --- modelopt/torch/models/_onfly_decompress.py | 34 +------------ .../torch/models/test_onfly_decompress.py | 48 ++----------------- 2 files changed, 4 insertions(+), 78 deletions(-) diff --git a/modelopt/torch/models/_onfly_decompress.py b/modelopt/torch/models/_onfly_decompress.py index d7f621d0e28f..24da3d1413f8 100644 --- a/modelopt/torch/models/_onfly_decompress.py +++ b/modelopt/torch/models/_onfly_decompress.py @@ -29,38 +29,6 @@ from modelopt.torch.utils.moe import config_levels, require_transformers -def to_native_backbone_config( - vlm_config, - model_type: str = "deepseek_v3", - architectures: "tuple[str, ...]" = ("DeepseekV3ForCausalLM",), -): - """The transformers-native backbone config for a VLM config's bundled ``text_config``. - - Drops the remote ``auto_map`` and the stale ``quantization_config``, and names the native model - type and class (hf_ptq's VLM detection reads ``architectures``). The decompression mark carries - over. Defaults target the Kimi DeepSeek-V3 backbone. - """ - # transformers is optional: this module imports without it. - 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) - native = CONFIG_MAPPING[model_type].from_dict(d) - 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, stored_as) - return native - - def _find_quantization_config(hf_config): """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.""" @@ -322,7 +290,7 @@ def install_onfly_decompress(hf_config): # 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") # The loader sets each parameter from what the conversion hands it, uncast. Every level counts: - # a backbone config derived from the text config (to_native_backbone_config) keeps its dtype. + # a backbone config derived from the text config keeps its dtype. for level in config_levels(hf_config): dtype = getattr(level, "dtype", None) if dtype is not None and str(dtype).removeprefix("torch.") != "bfloat16": diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 87913a8c6422..a8bae8f0da7a 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -510,13 +510,9 @@ def test_only_the_compressed_checkpoint_loads_through_the_rewritten_converters(m @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 -): +def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters(monkeypatch, marked): """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 + install marks the config it is handed. Either way the VLM's 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()) @@ -526,8 +522,7 @@ def test_a_vlm_marked_at_either_level_loads_through_the_rewritten_converters( 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)) + model = SimpleNamespace(config=copy.deepcopy(vlm)) out = mu.get_model_conversion_mapping(model) @@ -766,43 +761,6 @@ class _RemoteTextConfig(PretrainedConfig): 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 From 8390969a51ae6c641c1628bd261a202dee096168 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 02:28:02 +0000 Subject: [PATCH 41/45] Refuse compressed tensors in files the expert-parallel load never opens The export copies the safetensors files the load never opens as they are stored, such as a standalone mtp.safetensors. A decompressing load drops the quantization_config those files' compressed tensors need, so the export would carry tensors its own config cannot read. Such a file is now refused before anything loads, by its tensors' names, as the compressed tensors past the decoder layers already are; a plain one is carried as before. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 19 ++++++ .../unit/torch/utils/test_model_load_utils.py | 62 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index da459e3aebf6..29a7bdefc021 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -37,6 +37,7 @@ if TYPE_CHECKING: from collections.abc import Iterable +from safetensors.torch import safe_open from torch.distributed import device_mesh from torch.distributed.checkpoint import HuggingFaceStorageReader from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict @@ -89,6 +90,7 @@ checkpoint_tensor_files, ensure_local_checkpoint, indexed_weight_map, + off_index_safetensors_files, read_safetensors_subset, ) @@ -2286,6 +2288,23 @@ def gpu() -> tuple[bool, bool, str | None]: "no parameter for them, and the export would copy them as stored, with no " "quantization_config to read them. Decompress the checkpoint first." ) + # The files the load never opens (GLM-4.7's standalone mtp.safetensors, say) the export + # copies as stored, and this load drops the config a compressed tensor there needs: + # refused here too. A plain one is carried as before. + sidecar = [] + for file_name in off_index_safetensors_files(ckpt_path): + with safe_open(os.path.join(ckpt_path, file_name), framework="pt") as f: + sidecar += [ + f"{file_name}'s {name}" + for name in f.keys() # noqa: SIM118 - safe_open is not iterable + if name.rpartition(".")[2] in compressed_names + ] + if sidecar: + raise NotImplementedError( + f"The checkpoint stores {len(sidecar)} compressed tensors in files the load " + f"never opens, such as {sidecar[0]}: the export would copy them as stored, with " + "no quantization_config to read them. Decompress the checkpoint first." + ) # Experts are sharded over ``ep`` and replicated over ``edp``; the dense part is replicated over # the whole world (``dp``). diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index b62a845b971b..d3eabc1cba46 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3113,6 +3113,68 @@ def from_pretrained(path, config, **kwargs): ) +@pytest.mark.parametrize("compressed", [False, True], ids=["plain", "compressed"]) +def test_a_decompressed_load_refuses_compressed_tensors_in_a_file_it_never_opens( + ep_world_of_two, monkeypatch, tmp_path, compressed +): + """The export copies the safetensors files the load never opens as stored (GLM-4.7 keeps its MTP + head in a standalone ``mtp.safetensors``), and this load drops the quantization_config a + compressed tensor there needs: such a file is refused before anything loads. A plain one is + carried, and the load goes on.""" + + class _NextCheckError(Exception): + """Raised by the load's next check.""" + + def next_check(*args, **kwargs): + raise _NextCheckError + + model = model_with_decoder_at("model.") + loads = [] + + class _Loader: + @staticmethod + def from_pretrained(path, config, **kwargs): + loads.append(path) + model.config = copy.deepcopy(config) + return model + + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} + experts = "model.layers.0.mlp.experts.0.up_proj" + save_file( + {f"{experts}.{n}": torch.ones(1, dtype=d) for n, d in stored_as.items()}, + str(tmp_path / "model.safetensors"), + ) + mtp = "model.layers.1.mlp.experts.0.up_proj" + off_index = stored_as if compressed else {"weight": torch.bfloat16} + save_file( + {f"{mtp}.{n}": torch.ones(1, dtype=d) for n, d in off_index.items()}, + str(tmp_path / "mtp.safetensors"), + ) + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + monkeypatch.setattr(model_load_utils, "unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) + config = SimpleNamespace( + model_type=_NO_SPEC, + num_experts=8, + num_hidden_layers=2, + base_model_ep_plan=dict(_EP_PLAN), + quantization_config={"quant_method": "compressed-tensors"}, + ) + _onfly_decompress._mark(config, tuple(stored_as)) + + with pytest.raises(NotImplementedError if compressed else _NextCheckError) as refused: + load_and_prepare_ep(str(tmp_path), DistributionPlan(2, ep_size=2), hf_config=config) + assert bool(loads) != compressed + if compressed: + assert ( + f"2 compressed tensors in files the load never opens, such as mtp.safetensors's {mtp}." + in str(refused.value) + ) + + # ----------------------------------------------------------------------- _seed_parallel_states From 71bf5838eb7c7af74df46c29e7f1ce8558b0f60b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 02:31:01 +0000 Subject: [PATCH 42/45] Load a compressed checkpoint sharded, each expert with its own scales The real round trips stored one file with one scale shared by every expert, so experts paired with another's scales would have loaded unnoticed. Each expert now has scales of its own, and a sharded case stores the packed weights and the scales in two shards in opposite expert order. transformers 5.8 and 5.18 both sort the merged checkpoint's keys before the conversion, and the round trips pass on both; pairing the scales in reverse fails all four. Signed-off-by: Shengliang Xu --- .../torch/models/test_onfly_decompress.py | 42 ++++++++++++++----- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index a8bae8f0da7a..1880c0dfc742 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -889,11 +889,13 @@ def compression_param_names(self): } -def _compressed_mixtral(path, fmt): +def _compressed_mixtral(path, fmt, sharded=False): """Save a tiny Mixtral whose experts ``fmt`` compresses; return its config and dense experts. Mixtral's checkpoint names its experts ``w1``/``w3``/``w2``, which the model fuses and the export - names back: names a save gets wrong would show. + names back: names a save gets wrong would show. Each expert has a scale of its own, so experts + paired with another's scales would show too. ``sharded`` stores the experts' packed weights and + their scales in two shards in opposite expert order. """ weights, values, scale = _FORMATS[fmt] config = MixtralConfig( @@ -921,14 +923,29 @@ def _compressed_mixtral(path, fmt): dense = {} for key in sorted(k for k in list(state) if ".experts." in k): # ...experts..w1.weight rows, cols = state.pop(key).shape - dense[key] = values[torch.randint(0, len(values), (rows, cols), generator=gen)] - scales = torch.full((rows, cols // 32), scale) + factor = 2 ** int(key.split(".experts.")[1].split(".")[0]) # per expert, exact in both + dense[key] = values[torch.randint(0, len(values), (rows, cols), generator=gen)] * factor + scales = torch.full((rows, cols // 32), scale * factor) stored = compressor.compress({"weight": dense[key], "weight_scale": scales}, scheme) for name, tensor in stored.items(): # floats as bf16, as the Kimi checkpoints store them state[key.removesuffix("weight") + name] = ( tensor.bfloat16() if tensor.is_floating_point() else tensor ) - save_file(state, str(path / "model.safetensors"), metadata={"format": "pt"}) + if sharded: + # Expert 0's packed weights with expert 1's scales, and the other way around. + shards = {"model-00001-of-00002.safetensors": {}, "model-00002-of-00002.safetensors": {}} + for key, tensor in state.items(): + expert = key.split(".experts.")[1][0] if ".experts." in key else "0" + first = (expert == "0") == ("weight_scale" not in key) + shards[list(shards)[0 if first else 1]][key] = tensor + (path / "model.safetensors").unlink() + for name, shard in shards.items(): + save_file(shard, str(path / name), metadata={"format": "pt"}) + weight_map = {key: name for name, shard in shards.items() for key in shard} + index = {"metadata": {}, "weight_map": weight_map} + (path / "model.safetensors.index.json").write_text(json.dumps(index)) + else: + save_file(state, str(path / "model.safetensors"), metadata={"format": "pt"}) config_json = json.loads((path / "config.json").read_text()) config_json["quantization_config"] = quantization_config (path / "config.json").write_text(json.dumps(config_json)) @@ -936,18 +953,23 @@ def _compressed_mixtral(path, fmt): @pytest.mark.usefixtures("restore_loader_hooks") +@pytest.mark.parametrize("sharded", [False, True], ids=["one-file", "sharded"]) @pytest.mark.parametrize("fmt", sorted(_FORMATS)) -def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names(tmp_path, fmt): +def test_a_compressed_checkpoint_loads_exactly_and_saves_under_its_own_names( + tmp_path, fmt, sharded +): """A real ``from_pretrained`` through the installed decompressor: the fused experts hold the - checkpoint's weights exactly. The converters the load kept reverse as the original fusion, so a - save writes the experts back as the checkpoint's own ``w1``/``w3``/``w2`` weights, decompressed, - and the export names its expanded experts after them.""" + checkpoint's weights exactly, each expert's scaled by its own scales, also when the packed + weights and the scales sit in different shards in opposite expert order. The converters the load + kept reverse as the original fusion, so a save writes the experts back as the checkpoint's own + ``w1``/``w3``/``w2`` weights, decompressed, and the export names its expanded experts after + them.""" pytest.importorskip( "compressed_tensors", minversion="0.16", reason="needs compressed-tensors (CI's unit session installs it)", ) - config, dense = _compressed_mixtral(tmp_path, fmt) + config, dense = _compressed_mixtral(tmp_path, fmt, sharded) onfly.install_onfly_decompress(config) del config.quantization_config # as the expert-parallel load does: no quantizer competes From 02e72275a95a481d9a8e6ad8efc601d483b1daa9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 02:48:44 +0000 Subject: [PATCH 43/45] Read the checkpoint's tensor names once in a decompressing load check_compressed_keys returns the tensor name -> file map it read, but the unplaced-key check after the load read every shard header again: a second pass of up to a few seconds per rank for Kimi-K3's half a million names. unplaced_source_keys and record_unplaced_source_keys now take a map the load already read, and the expert-parallel load passes the one the decompression checks read. Other loads read it as before. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 23 +++++++++++++------ .../unit/torch/utils/test_model_load_utils.py | 6 ++--- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 29a7bdefc021..9707ca143dfb 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -401,7 +401,9 @@ def _group_sources_by_layer( return layer_sources, non_layer_sources, unplaced -def unplaced_source_keys(model: nn.Module, ckpt_path: str) -> list[str]: +def unplaced_source_keys( + model: nn.Module, ckpt_path: str, tensor_files: "dict[str, str] | None" = None +) -> list[str]: """Checkpoint tensors the loaded ``model`` has no parameter for, e.g. an MTP head. Every tensor in the weight files the loader opened @@ -424,8 +426,10 @@ def unplaced_source_keys(model: nn.Module, ckpt_path: str) -> list[str]: and a model our own loader built from its config has none, so the stock mapping answers there. A checkpoint without safetensors (``pytorch_model.bin``) has nothing the export can carry: ``[]``. + Pass ``tensor_files``, the checkpoint's ``tensor name -> file`` map, when the load already read it. """ - tensor_files = checkpoint_tensor_files(ckpt_path) + if tensor_files is None: + tensor_files = checkpoint_tensor_files(ckpt_path) if not tensor_files: return [] # Every name of a tied parameter: a checkpoint can store the one named_parameters() drops by @@ -448,7 +452,10 @@ def unplaced_source_keys(model: nn.Module, ckpt_path: str) -> list[str]: def record_unplaced_source_keys( - model: nn.Module, ckpt_path: str, keys: "Iterable[str] | None" = None + model: nn.Module, + ckpt_path: str, + keys: "Iterable[str] | None" = None, + tensor_files: "dict[str, str] | None" = None, ) -> list[str]: """Record on ``model`` the checkpoint tensors it has no parameter for, for the export to carry. @@ -456,10 +463,11 @@ def record_unplaced_source_keys( at load time, and the export reads it back (``_modelopt_unplaced_source_keys``, beside the checkpoint in ``_modelopt_source_checkpoint``) rather than deriving it again. Pass ``keys`` when the loader already found them while placing the weights; otherwise :func:`unplaced_source_keys` - derives them from the loaded model. Every loader of a model that may be exported calls this: - :func:`load_pretrained` and :func:`parallel_load_and_prepare_fsdp2` do. + derives them from the loaded model, from ``tensor_files`` when the load already read them. Every + loader of a model that may be exported calls this: :func:`load_pretrained` and + :func:`parallel_load_and_prepare_fsdp2` do. """ - keys = sorted(unplaced_source_keys(model, ckpt_path) if keys is None else keys) + keys = sorted(unplaced_source_keys(model, ckpt_path, tensor_files) if keys is None else keys) model._modelopt_unplaced_source_keys = keys model._modelopt_source_checkpoint = str(ckpt_path) return keys @@ -2252,6 +2260,7 @@ def gpu() -> tuple[bool, bool, str | None]: # The load and the export both read the local snapshot, made only once the checks pass, so a # refused load downloads nothing. compressed_names: set[str] = set() + tensor_files: dict[str, str] | None = None # read once, by the decompressing load's checks with _refused_together("the local checkpoint", alone=False): _, ckpt_path = ensure_local_checkpoint(ckpt_path, **hub_options) if _decompressed_as_loaded(hf_config): @@ -2381,7 +2390,7 @@ def gpu() -> tuple[bool, bool, str | None]: # The checkpoint tensors the model has no parameter for (an MTP head), worked out once, at load, # for the export to carry. - unplaced = record_unplaced_source_keys(model, ckpt_path) + unplaced = record_unplaced_source_keys(model, ckpt_path, tensor_files=tensor_files) # The export would copy them as stored, and this load dropped the quantization_config that # reads a compressed one, such as an MTP layer's experts. A suffix is enough to tell them: # check_compressed_keys refused any compressed tensor outside the experts. diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index d3eabc1cba46..0dc082518270 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -2778,7 +2778,7 @@ 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) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args, **kwargs: []) config = SimpleNamespace( model_type=_NO_SPEC, base_model_ep_plan={"layers.*.mlp.experts.*.gate_up_proj": "grouped_gemm"}, @@ -2827,7 +2827,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args, **kwargs: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) class _Config(transformers.PretrainedConfig): @@ -2889,7 +2889,7 @@ def next_check(*args, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: mesh ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args, **kwargs: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = Qwen3MoeConfig( hidden_size=16, From 3ae43d65874c7ac29967fca8c1848b253951d67f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 07:14:10 +0000 Subject: [PATCH 44/45] Pass the read tensor map by keyword, and only when there is one record_unplaced_source_keys handed unplaced_source_keys the map a load already read as a third positional argument, every time. A caller that replaces the function with its two-argument form then failed, and the export, which catches that and warns, carried no unplaced keys. The map now goes by keyword, and only when the load read one, so any other call is the two-argument one it always was. The export test's stand-in takes the keyword, and a new test pins both forms. The tests that load a safetensors file and write over it now copy the tensors out of it first, as a lazier safetensors would otherwise read a file being rewritten. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 10 ++++++- .../torch/utils/test_expert_parallel_load.py | 3 ++- .../torch/export/test_unified_export_hf.py | 4 ++- .../torch/models/test_onfly_decompress.py | 6 +++-- .../torch/utils/test_unplaced_source_keys.py | 26 +++++++++++++++++++ 5 files changed, 44 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 9707ca143dfb..ee23676b4c39 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -467,7 +467,15 @@ def record_unplaced_source_keys( loader of a model that may be exported calls this: :func:`load_pretrained` and :func:`parallel_load_and_prepare_fsdp2` do. """ - keys = sorted(unplaced_source_keys(model, ckpt_path, tensor_files) if keys is None else keys) + if keys is None: + # The map by keyword, and only when the load read it: a caller that replaced + # unplaced_source_keys with its two-argument form gets its keys as before. + keys = ( + unplaced_source_keys(model, ckpt_path) + if tensor_files is None + else unplaced_source_keys(model, ckpt_path, tensor_files=tensor_files) + ) + keys = sorted(keys) model._modelopt_unplaced_source_keys = keys model._modelopt_source_checkpoint = str(ckpt_path) return keys diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py index 533147ac1b03..f26bb6506a81 100644 --- a/tests/gpu/torch/utils/test_expert_parallel_load.py +++ b/tests/gpu/torch/utils/test_expert_parallel_load.py @@ -132,7 +132,8 @@ def _save_compressed_tiny_qwen3_moe(dense_dir, packed_dir, compression): """The tiny Qwen3-MoE twice: whole in bf16, and with its experts compressed by ``compression``.""" fmt, weights, values, scale = _COMPRESSIONS[compression] _save_tiny_qwen3_moe(dense_dir) - state = load_file(str(dense_dir / "model.safetensors")) + # Copied out of the file, which is written over below. + state = {k: v.clone() for k, v in load_file(str(dense_dir / "model.safetensors")).items()} gen = torch.Generator().manual_seed(0) for key in [k for k in state if ".experts." in k]: # ...mlp.experts..gate_proj.weight draw = torch.randint(0, len(values), state[key].shape, generator=gen) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 1b8d55373fae..12126c61a277 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -1080,7 +1080,9 @@ def test_an_unrecorded_model_derives_its_unplaced_keys_once(tmp_path, monkeypatc ) derived_from = [] monkeypatch.setattr( - model_load_utils, "unplaced_source_keys", lambda m, c: derived_from.append(c) or [mtp] + model_load_utils, + "unplaced_source_keys", + lambda m, c, tensor_files=None: derived_from.append(c) or [mtp], ) model = _ProvenanceModel(name_or_path=tmp_path) diff --git a/tests/unit/torch/models/test_onfly_decompress.py b/tests/unit/torch/models/test_onfly_decompress.py index 1880c0dfc742..3c68410d6fcd 100644 --- a/tests/unit/torch/models/test_onfly_decompress.py +++ b/tests/unit/torch/models/test_onfly_decompress.py @@ -918,7 +918,8 @@ def _compressed_mixtral(path, fmt, sharded=False): compressor, scheme, _, _ = onfly._read_compression( SimpleNamespace(quantization_config=quantization_config) ) - state = load_file(str(path / "model.safetensors")) + # Copied out of the file, which is written over below. + state = {k: v.clone() for k, v in load_file(str(path / "model.safetensors")).items()} gen = torch.Generator().manual_seed(0) dense = {} for key in sorted(k for k in list(state) if ".experts." in k): # ...experts..w1.weight @@ -1012,7 +1013,8 @@ def test_a_load_the_decompressor_did_not_mark_casts_as_before(tmp_path): tie_word_embeddings=False, ) AutoModelForCausalLM.from_config(config).save_pretrained(tmp_path) - state = load_file(str(tmp_path / "model.safetensors")) + # Copied out of the file, which is written over below. + state = {k: v.clone() for k, v in load_file(str(tmp_path / "model.safetensors")).items()} state["model.norm.weight"] = torch.full_like(state["model.norm.weight"], 3, dtype=torch.int8) save_file(state, str(tmp_path / "model.safetensors"), metadata={"format": "pt"}) diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 274c04cc9e66..1882b7c5cc88 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -51,6 +51,32 @@ parallel_load_and_prepare_fsdp2, ) + +def test_a_two_argument_unplaced_source_keys_still_records_the_keys(monkeypatch, tmp_path): + """The map a load already read reaches ``unplaced_source_keys`` by keyword, and only when there is + one, so a caller that replaced the function with its two-argument form still records keys.""" + seen = [] + + def two_arguments(model, ckpt_path): + seen.append(ckpt_path) + return ["mtp.fc.weight"] + + monkeypatch.setattr(model_load_utils, "unplaced_source_keys", two_arguments) + model = LlamaForCausalLM(_tiny_llama_config()) + + assert model_load_utils.record_unplaced_source_keys(model, str(tmp_path)) == ["mtp.fc.weight"] + assert seen == [str(tmp_path)] + read = {"mtp.fc.weight": "model.safetensors"} + monkeypatch.setattr( + model_load_utils, + "unplaced_source_keys", + lambda model, ckpt_path, tensor_files=None: sorted(tensor_files), + ) + assert model_load_utils.record_unplaced_source_keys( + model, str(tmp_path), tensor_files=read + ) == ["mtp.fc.weight"] + + # A layer past ``num_hidden_layers``: genuinely stray, so it must always be recorded. STRAY = "model.layers.1.input_layernorm.weight" From 8feea6bbe742dffef19ccfc3dad50329972e6480 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 3 Nov 2026 07:16:32 +0000 Subject: [PATCH 45/45] Let the decompressing-load tests' stand-ins take the read map A decompressing load now passes the tensor map it read by keyword, which the two loader tests' stand-ins for unplaced_source_keys did not accept. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_model_load_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 0dc082518270..3669a1ccf33d 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -3088,7 +3088,9 @@ def from_pretrained(path, config, **kwargs): ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) monkeypatch.setattr( - model_load_utils, "unplaced_source_keys", lambda *args: [f"{mtp}.{n}" for n in names] + model_load_utils, + "unplaced_source_keys", + lambda *args, **kwargs: [f"{mtp}.{n}" for n in names], ) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( @@ -3154,7 +3156,7 @@ def from_pretrained(path, config, **kwargs): "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() ) monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) - monkeypatch.setattr(model_load_utils, "unplaced_source_keys", lambda *args: []) + monkeypatch.setattr(model_load_utils, "unplaced_source_keys", lambda *args, **kwargs: []) monkeypatch.setattr(model_load_utils, "_seed_parallel_states", next_check) config = SimpleNamespace( model_type=_NO_SPEC,