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/models/_remote_code.py b/modelopt/torch/models/_remote_code.py new file mode 100644 index 00000000000..dbd63838f70 --- /dev/null +++ b/modelopt/torch/models/_remote_code.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Loading the modeling code a checkpoint ships with itself.""" + +import glob +import os +import shutil +import tempfile +import warnings + + +def stage_full_remote_closure() -> None: + """Make transformers stage every top-level module of a local checkpoint's remote code. + + transformers' ``get_cached_module_file`` copies a remote module and its *direct* relative imports + into ``transformers_modules///``, but does not recurse. Remote code more than one import + deep -- Kimi-K2.5's ``modeling_kimi_k25 -> modeling_deepseek -> configuration_deepseek`` -- is + then staged incomplete, and the later recursive walk in ``get_relative_import_files`` fails with + ``FileNotFoundError``, from ``get_class_from_dynamic_module`` and ``from_pretrained`` alike. + + Wrap ``get_cached_module_file`` so that, loading from a local checkpoint dir, it copies every + sibling ``*.py`` into the stage dir it just populated; modules in a subdirectory of the + checkpoint are not copied. Idempotent. + """ + import transformers.dynamic_module_utils as dmu + from transformers.utils import HF_MODULES_CACHE + + if getattr(dmu.get_cached_module_file, "_modelopt_full_closure", False): + return + orig = dmu.get_cached_module_file + + def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs): + resolved = orig(pretrained_model_name_or_path, module_file, *args, **kwargs) + src = str(pretrained_model_name_or_path) + if os.path.isdir(src): + staged = str(resolved) + if not os.path.isabs(staged): + staged = os.path.join(HF_MODULES_CACHE, staged) + stage_dir = os.path.dirname(staged) + try: + for py in glob.glob(os.path.join(src, "*.py")): + dst = os.path.join(stage_dir, os.path.basename(py)) + if not os.path.exists(dst): + # Every rank loading the checkpoint stages into this dir and imports from it: + # renamed into place (atomic), a file appears there whole, never half written. + fd, tmp = tempfile.mkstemp(dir=stage_dir) + os.close(fd) + shutil.copy(py, tmp) + os.replace(tmp, dst) + except OSError as e: # a failed copy only matters if a deeper import needs the file + warnings.warn(f"Could not stage the remote code of {src}: {e}") + return resolved + + setattr(_staged, "_modelopt_full_closure", True) + dmu.get_cached_module_file = _staged diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 3b5826bd39c..3254372090f 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -2291,6 +2291,13 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int: _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"}) +def _decompressed_as_loaded(cfg) -> bool: + """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint.""" + from modelopt.torch.models import _onfly_decompress + + return _onfly_decompress.is_marked(cfg) + + def _ckpt_quant_method(cfg, ckpt_path: str | None = None): """Return the checkpoint's native quantization method, or ``None``. @@ -2605,7 +2612,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A "one: per-tensor scales would load as blocks, leaving the weights unscaled. Hand the " "config as AutoConfig.from_pretrained reads it, or a local checkpoint." ) - if quant_method not in _EP_SOURCE_FORMATS: + if quant_method not in _EP_SOURCE_FORMATS and not _decompressed_as_loaded(hf_config): # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights. dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f)) raise NotImplementedError( @@ -2615,6 +2622,18 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A ) # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares. hf_config = copy.deepcopy(hf_config) + if _decompressed_as_loaded(hf_config): + # The decompressor hands the experts over dequantized. Left on the config (Kimi-K2.5's config + # class copies its text config's to the top level), the checkpoint's quantization_config + # would build transformers' compressed-tensors quantizer, which from 5.13 claims the same + # expert converters first. Deleted rather than set to None: before 5.11 transformers takes + # any quantization_config attribute, None too, for a quantized checkpoint. Every level + # _ckpt_quant_method reads is cleared. + from modelopt.torch.utils.moe import _SUB_CONFIG_ATTRS + + for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): + if cfg is not None and hasattr(cfg, "quantization_config"): + del cfg.quantization_config ep_cfg = _ep_config(hf_config) # for VLMs the plan lives on .text_config ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {}) if not ep_plan: @@ -2638,6 +2657,12 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read. # Resolved once the plan and the config pass, so a refused load downloads nothing. ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank(), **hub_options) + if _decompressed_as_loaded(hf_config): + # Only the experts decompress as they load, and the config cannot say what else the + # checkpoint compressed, or in which dtypes: what it stores is checked before anything loads. + from modelopt.torch.models import _onfly_decompress + + _onfly_decompress.check_compressed_keys(hf_config, ckpt_path) # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two # factorizations of the same world: diff --git a/noxfile.py b/noxfile.py index 96232a2d381..3b333aa0112 100644 --- a/noxfile.py +++ b/noxfile.py @@ -68,7 +68,12 @@ def _cov_args(): @nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS]) def unit(session, torch_ver, tf_ver): """Unit tests — parametrized over torch and transformers versions.""" - session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]") + # compressed-tensors runs the real-library tests of on-the-fly expert decompression. It needs + # torch>=2.10, so the older legs skip those tests; it goes in with the torchvision pin so the + # resolver keeps the leg's torch. Not in dev-test: the transformers-free partial_unit(torch) + # lane installs that extra, and compressed-tensors pulls in transformers. + extra = [] if torch_ver in ("torch_28", "torch_29") else ["compressed-tensors"] + session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", *extra) tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) 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() + ) diff --git a/tests/unit/torch/models/test_remote_code.py b/tests/unit/torch/models/test_remote_code.py new file mode 100644 index 00000000000..d324773e47b --- /dev/null +++ b/tests/unit/torch/models/test_remote_code.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Staging every module of a local checkpoint's remote code.""" + +import shutil +from types import SimpleNamespace + +import pytest + +dmu = pytest.importorskip("transformers.dynamic_module_utils") + +from modelopt.torch.models import _remote_code +from modelopt.torch.models._remote_code import stage_full_remote_closure + + +def test_every_module_of_a_local_checkpoint_is_staged(monkeypatch, tmp_path): + """transformers stages a remote module and its direct imports only, so the third module of + a -> b -> c is missing until the wrap copies the checkpoint's other modules. Installing again + wraps the loader once.""" + ckpt, stage = tmp_path / "ckpt", tmp_path / "stage" + ckpt.mkdir() + stage.mkdir() + for name, body in ( + ("a.py", "from .b import B\n"), + ("b.py", "from .c import C\n"), + ("c.py", ""), + ): + (ckpt / name).write_text(body) + + def stage_one_hop(path, module_file, *args, **kwargs): + for name in (module_file, "b.py"): + shutil.copy(ckpt / name, stage / name) + return str(stage / module_file) + + monkeypatch.setattr(dmu, "get_cached_module_file", stage_one_hop) + stage_full_remote_closure() + wrapped = dmu.get_cached_module_file + stage_full_remote_closure() + + assert dmu.get_cached_module_file is wrapped + assert dmu.get_cached_module_file(str(ckpt), "a.py") == str(stage / "a.py") + assert sorted(p.name for p in stage.iterdir()) == ["a.py", "b.py", "c.py"] + + +def test_a_module_is_staged_whole_or_not_at_all(monkeypatch, tmp_path): + """Every rank stages into the same dir and imports what it finds there, so a module must not + appear before its copy completes: a copy cut short leaves nothing behind for a rank to take as + staged, and the next load stages the whole module.""" + ckpt, stage = tmp_path / "ckpt", tmp_path / "stage" + ckpt.mkdir() + stage.mkdir() + (ckpt / "a.py").write_text("from .b import B\n") + (ckpt / "b.py").write_text("B = 1\n" * 100) + + def stage_only_the_module(path, module_file, *args, **kwargs): + shutil.copy(ckpt / module_file, stage / module_file) + return str(stage / module_file) + + def copy_cut_short(src, dst): + with open(src) as f, open(dst, "w") as out: + out.write(f.read()[:10]) + raise OSError(28, "No space left on device") + + monkeypatch.setattr(dmu, "get_cached_module_file", stage_only_the_module) + stage_full_remote_closure() + with monkeypatch.context() as m: + m.setattr(_remote_code, "shutil", SimpleNamespace(copy=copy_cut_short)) + with pytest.warns(UserWarning, match="No space left"): + dmu.get_cached_module_file(str(ckpt), "a.py") + assert not (stage / "b.py").exists() + + dmu.get_cached_module_file(str(ckpt), "a.py") + assert (stage / "b.py").read_text() == (ckpt / "b.py").read_text() diff --git a/tests/unit/torch/utils/test_ep_loader.py b/tests/unit/torch/utils/test_ep_loader.py index 93c1d969c12..2d7db3d50d5 100644 --- a/tests/unit/torch/utils/test_ep_loader.py +++ b/tests/unit/torch/utils/test_ep_loader.py @@ -564,6 +564,64 @@ def auto_class(hf_config, require_lm_head=True, trust_remote_code=False): assert classes == [True] +def test_a_compressed_checkpoint_a_spec_decompresses_as_it_loads_passes( + ep_world_of_two, monkeypatch, tmp_path +): + """A model spec that installs the on-the-fly decompressor (Kimi-K2.5's) marks the config; that + checkpoint passes the format check and loads without its quantization_config, at every level, + which would build transformers' own compressed-tensors quantizer to compete with the + decompressor. The same checkpoint unmarked is refused, and the caller's config keeps its own. + So is one that compresses more than its experts, by the names it stores, before it loads.""" + from safetensors.torch import save_file + + from modelopt.torch.models import _onfly_decompress + + class _LoadedError(Exception): + """Stops the load at from_pretrained.""" + + loaded = {} + + class _Loader: + @staticmethod + def from_pretrained(path, **kwargs): + loaded.update(kwargs) + raise _LoadedError + + monkeypatch.setattr( + "torch.distributed.device_mesh.init_device_mesh", lambda *args, **kwargs: _Mesh() + ) + monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: _Loader) + stored_as = {"weight_packed": torch.int32, "weight_scale": torch.bfloat16} + expert, attention = "model.layers.1.mlp.experts.0.up_proj", "model.layers.1.self_attn.o_proj" + for name, compressed in (("experts", [expert]), ("more", [expert, attention])): + stored = { + f"{m}.{n}": torch.ones(1, dtype=d) for m in compressed for n, d in stored_as.items() + } + (tmp_path / name).mkdir() + save_file(stored, str(tmp_path / name / "model.safetensors")) + quantization_config = {"quant_method": "compressed-tensors"} + text = SimpleNamespace(quantization_config=quantization_config, base_model_ep_plan=_EP_PLAN) + llm = SimpleNamespace(quantization_config=quantization_config) + config = SimpleNamespace( + quantization_config=quantization_config, text_config=text, llm_config=llm + ) + plan = DistributionPlan(2, ep_size=2) + + with pytest.raises(NotImplementedError, match="cannot dequantize"): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) + _onfly_decompress._mark(config, tuple(stored_as)) + with pytest.raises(NotImplementedError, match="1 of them outside the experts"): + load_and_prepare_ep(str(tmp_path / "more"), plan, hf_config=config) + assert not loaded + with pytest.raises(_LoadedError): + load_and_prepare_ep(str(tmp_path / "experts"), plan, hf_config=config) + + for level in (loaded["config"], loaded["config"].text_config, loaded["config"].llm_config): + assert not hasattr(level, "quantization_config") + assert config.quantization_config is text.quantization_config is quantization_config + assert llm.quantization_config is quantization_config + + # ------------------------------------------------------------------------ quant wrappers