From 17884f632379dfa406d142506f3ef557cbea190e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:00:51 +0000 Subject: [PATCH 01/20] 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 +++++++++++ modelopt/torch/utils/moe.py | 5 + noxfile.py | 7 +- .../torch/models/test_onfly_decompress.py | 764 ++++++++++++++++++ 4 files changed, 1250 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 00000000000..42cc1075c9d --- /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/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 29282c5f4cf..0bc4d9d313c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -55,6 +55,11 @@ def require_transformers(min_version: str, feature: str) -> None: ) +# The nested configs a VLM's expert-parallel plan and expert count are read from. A spec attaching +# its plan anywhere else would be read as having none, and the MoE planned as a dense model. +_SUB_CONFIG_ATTRS = ("text_config", "llm_config") + + def _ep_config(hf_config): """Return the config that carries the EP plan / expert count. diff --git a/noxfile.py b/noxfile.py index 8528f7136e0..193293aa38b 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 00000000000..33e62eae12f --- /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 d7c4d42c384f7366b29ccc937dcad24cb88bae61 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:01:06 +0000 Subject: [PATCH 02/20] 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 42cc1075c9d..464d1699c5a 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 8fc97d8fef90ae8d8790f30f26b5a142567770ac Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:01:10 +0000 Subject: [PATCH 03/20] 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 193293aa38b..427a4d6dbad 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 79cb8e61ba0033420d403ba2a4d963065b97d4f9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:42:52 +0000 Subject: [PATCH 04/20] 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 464d1699c5a..81c9283a3b4 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 33e62eae12f..a25802b0ca2 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 835af6c0658e1a11a12c66b729c580830c8d3f3e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:43:37 +0000 Subject: [PATCH 05/20] 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 81c9283a3b4..81c58c7db32 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 a25802b0ca2..e966b3504d1 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 0d47a96d7640bd057f1890703b8ae39bcf854df0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:44:39 +0000 Subject: [PATCH 06/20] 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 81c58c7db32..d1fecbd56ce 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 e966b3504d1..d37684156bf 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 4ef070ad884de49f9182e8cc5082a4ddb55ff997 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:47:43 +0000 Subject: [PATCH 07/20] 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 d1fecbd56ce..dadee897a32 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 41acc0bfd881be7bc65e054b69897599fb6006a8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:38 +0000 Subject: [PATCH 08/20] 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 dadee897a32..5f36f330309 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 d37684156bf..24730d76b5f 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 2c90e043362cdaebe0a611ccbc5434313fb35a06 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:42 +0000 Subject: [PATCH 09/20] 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 5f36f330309..ec5a8ef9a58 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 24730d76b5f..95b6cd01bc6 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 c2756cc5b943721a96cba917c408e0c84b4b40f8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:46 +0000 Subject: [PATCH 10/20] 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 ec5a8ef9a58..679ca1f0ef0 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 95b6cd01bc6..d0cba935981 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 b647b1dcd15e8e2f19bf83c0f836a7c8bba63b99 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:50 +0000 Subject: [PATCH 11/20] 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 679ca1f0ef0..a0f601c718e 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 d0cba935981..b1cc3d4a41b 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 cdc2e2444a52892e8ebfe297a14ca37c95a6af96 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:07:55 +0000 Subject: [PATCH 12/20] 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 a0f601c718e..e3b7c15056d 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 5f9f379b9553b6b8002a73e710b5405310b09155 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 19:35:48 +0000 Subject: [PATCH 13/20] 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 e3b7c15056d..38fed405eff 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 b1cc3d4a41b..931f02b32c6 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 0a711b95f5046a6522a2b8b66179db538684ca46 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:53:35 +0000 Subject: [PATCH 14/20] 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 38fed405eff..3c2af77b3a2 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 931f02b32c6..3b6e1e64b2c 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 6cd12d7353a16fff2b1aa5d32c82ece4c24dd6b8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 20:53:41 +0000 Subject: [PATCH 15/20] 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 3c2af77b3a2..2ccd7c96087 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 3b6e1e64b2c..c2712bd371d 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 485f6b63fe6421ad99b9be143d084f26be8a72f5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:22:40 +0000 Subject: [PATCH 16/20] 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 c2712bd371d..cf299a91991 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 5d0928df4941cebc70aded9edd1a7805ad198bff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:30:56 +0000 Subject: [PATCH 17/20] 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 2ccd7c96087..e2ab2bce52e 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 cf299a91991..2c680fbe1f9 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 32ca3d3ae47b4def4c2dafb723fea13686887708 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:40:26 +0000 Subject: [PATCH 18/20] 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 2c680fbe1f9..2f17987b4a8 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 7e678d2ed70003d25c185fcdab8a943265bed16c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 21:42:55 +0000 Subject: [PATCH 19/20] 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 e2ab2bce52e..27bc81d24d5 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 9deea3be035873acb653299819400e99b31aafd8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:34:59 +0000 Subject: [PATCH 20/20] 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 27bc81d24d5..1cbcbb9ea2e 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