diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f54ca539f62..28f97fc2190 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,6 +28,9 @@ Changelog **Backward Breaking Changes** +- ``examples/hf_ptq`` no longer detects MTP layers by name. Weights the loader could not place -- an MTP head, an auxiliary tower -- are identified from Transformers' own accounting: the model is loaded with ``from_pretrained(..., output_loading_info=True)`` and the reported ``unexpected_keys`` (present in the checkpoint, not in the model's architecture) are recorded on the model and carried into the export unchanged. Everything the loader *did* place goes through the normal export path. This removes ``load_mtp_weights``, ``mtp_layer_prefixes_from_checkpoint`` and their support matrix of MTP storage conventions, along with ``_add_mtp_exclusions`` and the pre-quantization ``enable: False`` entries ``hf_ptq`` appended to the recipe's ``quant_cfg``. Two consequences: MTP layers now follow the recipe like any other module instead of being force-excluded by the script -- matching ``examples/megatron_bridge``, which has no MTP-specific code at all -- and ``quantization_config.ignore`` can no longer claim a layer is unquantized that the export in fact quantized. Recipes importing ``configs/ptq/units/default_disabled_quantizers`` still disable ``mtp.*``, so their behaviour is unchanged; a recipe omitting that unit will now quantize an MTP the model actually built. + +- ``examples/hf_ptq --vllm_fakequant_export`` now raises ``NotImplementedError`` when the checkpoint holds weights the model has no parameter for and a shard actually provides them (an MTP head, an auxiliary tower). The fake-quant exporter writes only model-backed state, so it would otherwise drop those weights silently -- and a fake-quant checkpoint is evaluated, where a missing head changes the score rather than failing loudly. Use the unified HF export, which carries them through. Buffers Transformers recomputes are not weights to lose: ``*.inv_freq`` is skipped even when a shard provides it, since older Llama/Mistral-lineage conversions do list it in the index and refusing an export over it would reject checkpoints that export correctly today. The check runs immediately after the model loads, not at export time, so an incompatible run fails before calibration rather than after it. - The ``modelopt.onnx.quantization.graph_utils`` module has been removed with no compatibility shim; update direct imports using this migration map: diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index be8300d4d67..7c9a3f3e9b0 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -22,7 +22,7 @@ import logging import os import warnings -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Iterable, Iterator from contextlib import contextmanager from dataclasses import dataclass from datetime import timedelta @@ -47,7 +47,11 @@ from modelopt.recipe import load_recipe from modelopt.torch.export.model_utils import is_multimodal_model -from modelopt.torch.export.plugins.hf_checkpoint_utils import copy_non_safetensor_files_from_ckpt +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_non_safetensor_files_from_ckpt, + copy_off_index_safetensors, +) +from modelopt.torch.utils.plugins.model_load_utils import record_unplaced_source_keys try: from huggingface_hub import snapshot_download @@ -422,142 +426,6 @@ def get_processor( return None -def get_inlined_mtp_prefixes(config: Any) -> list[str]: - """Turn an HF config into the list of state-dict prefixes for inlined-MTP layers.""" - # ``or 0``: some configs set num_nextn_predict_layers=None rather than omit it. - num_nextn = int(getattr(config, "num_nextn_predict_layers", 0) or 0) - if not num_nextn: - return [] - num_hidden = config.num_hidden_layers - return [f"model.layers.{i}" for i in range(num_hidden, num_hidden + num_nextn)] - - -def _keys_to_prefixes(keys: Iterable[str]) -> set[str]: - """Invert separate-file MTP keys into the prefixes the exporter needs for exclude_modules. - ``"mtp.fc.weight"`` → ``{"mtp"}``; ``"mtp.layers.0.q_proj.weight"`` → - ``{"mtp", "mtp.layers.0"}``. ``"model"`` top-level is dropped to avoid the - ``"model*"`` wildcard covering the whole backbone. - """ - prefixes: set[str] = set() - for key in keys: - parts = key.split(".") - if parts and parts[0] != "model": - prefixes.add(parts[0]) - for i, part in enumerate(parts): - if part == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit(): - prefixes.add(".".join(parts[: i + 2])) - break - return prefixes - - -def _load_tensors_matching( - model_dir: Path, predicate: Callable[[str], bool] -) -> dict[str, torch.Tensor]: - """Stream tensors satisfying ``predicate(key)`` from every safetensors - source in ``model_dir`` (indexed shards + standalone files, each opened - at most once). - """ - tensors: dict[str, torch.Tensor] = {} - seen_shards: set[str] = set() - - index_file = model_dir / "model.safetensors.index.json" - if index_file.exists(): - with open(index_file) as f: - weight_map = json.load(f)["weight_map"] - per_shard: dict[str, list[str]] = {} - for key, shard_name in weight_map.items(): - if predicate(key): - per_shard.setdefault(shard_name, []).append(key) - for shard_name, keys in per_shard.items(): - seen_shards.add(shard_name) - with safe_open(str(model_dir / shard_name), framework="pt", device="cpu") as f: - for k in keys: - tensors[k] = f.get_tensor(k) - - for shard in sorted(model_dir.glob("*.safetensors")): - if shard.name in seen_shards: - continue - with safe_open(str(shard), framework="pt", device="cpu") as f: - for k in f.keys(): # noqa: SIM118 - safe_open is not iterable - if predicate(k): - tensors[k] = f.get_tensor(k) - return tensors - - -def _apply_to_model_state_dict( - model: torch.nn.Module, tensors: dict[str, torch.Tensor] -) -> dict[str, torch.Tensor]: - """Load tensors with a slot in ``model.state_dict()`` in-place; return the - rest as orphans for ``extra_state_dict``. - """ - model_state = model.state_dict() - in_state_dict = {k: v for k, v in tensors.items() if k in model_state} - out_state_dict = {k: v for k, v in tensors.items() if k not in model_state} - if in_state_dict: - model.load_state_dict(in_state_dict, strict=False) - return out_state_dict - - -def mtp_layer_prefixes_from_checkpoint(model_path: str) -> list[str]: - """MTP exclude-prefixes from a checkpoint's safetensors index (``[]`` if none); reads no tensors. - - Local-index-only, matching :func:`load_mtp_weights`, so detection and re-attach stay in sync. - """ - index_file = Path(model_path) / "model.safetensors.index.json" - if not index_file.exists(): - return [] - weight_map = json.load(open(index_file))["weight_map"] - mtp_keys = [k for k, v in weight_map.items() if "mtp" in k or "mtp" in v] - return list(_keys_to_prefixes(mtp_keys)) - - -def load_mtp_weights( - model: torch.nn.Module, model_path: str -) -> tuple[list[str], dict[str, torch.Tensor]]: - """Detect and load MTP weights. Support matrix: - - Convention Architectures On-disk shape - ------------- ------------------------ ------------------------------- - inlined GLM-5.1 (``GlmMoeDsa``), ``model.layers.{N}.*`` - DeepSeek-V3 - separate-file GLM-4.7 standalone ``mtp.safetensors`` - separate-file Qwen3-Next indexed ``mtp.*`` tail shard - - Inlined ``N`` in ``[num_hidden, num_hidden + num_nextn_predict_layers)``; - may be orphaned at ``from_pretrained`` time if the HF class only builds - ``num_hidden`` decoders. - - Returns ``(prefixes, not_in_state_dict)``: ``prefixes`` populates - ``quantization_config.exclude_modules``; ``not_in_state_dict`` is fed to - ``export_hf_checkpoint(extra_state_dict=...)``. - """ - model_dir = Path(model_path) - - inlined_prefixes = set(get_inlined_mtp_prefixes(model.config)) - inlined_tuple = tuple(p + "." for p in inlined_prefixes) - - # Combined predicate covering both conventions in one pass. - def predicate(key: str) -> bool: - return key.startswith(inlined_tuple) or "mtp" in key - - tensors = _load_tensors_matching(model_dir, predicate) - if not tensors: - return [], {} - - separate_keys = [k for k in tensors if not k.startswith(inlined_tuple)] - prefixes = inlined_prefixes | _keys_to_prefixes(separate_keys) - - not_in_state_dict = _apply_to_model_state_dict(model, tensors) - - print( - f"✓ Detected {len(tensors)} MTP tensors under {sorted(prefixes)} " - f"(loaded into model: {len(tensors) - len(not_in_state_dict)}, " - f"orphaned: {len(not_in_state_dict)})" - ) - - return sorted(prefixes), not_in_state_dict - - def get_dtype(dtype): if dtype == "bf16": dtype = torch.bfloat16 @@ -590,7 +458,6 @@ def _unpack_compressed_linear_weights(model, ckpt_path=None): return from huggingface_hub import hf_hub_download - from safetensors import safe_open is_local = os.path.isdir(ckpt_path) @@ -737,6 +604,46 @@ def _fmt_max_memory(max_memory: dict) -> str: return "\n".join(parts) +def _resolved_local_dir(ckpt_path: str) -> str: + """Return the local directory ``ckpt_path`` names, resolving a hub id to its snapshot. + + The export re-reads the source checkpoint by path to carry over the weights the loader could + not place. Recording the hub id instead would leave it reading ``org/model``, which is not a + directory -- so every carried weight would be dropped with a warning. ``from_pretrained`` has + already populated the cache by the time this runs, so the lookup is local and offline. + """ + if Path(ckpt_path).is_dir(): + return str(ckpt_path) + if snapshot_download is None: + return str(ckpt_path) + try: + return snapshot_download(ckpt_path, local_files_only=True) + except Exception: + # No snapshot to point at; the export falls back to its own provenance handling. + return str(ckpt_path) + + +def _from_pretrained_recording(auto_class, ckpt_path, **kwargs): + """``from_pretrained`` that records what the loader could not place. + + ``output_loading_info=True`` makes Transformers return its own accounting of the load; + ``unexpected_keys`` -- keys present in the checkpoint but not in the model's architecture -- + is exactly the set the export has to carry over (an MTP head, an auxiliary tower). Taking it + from the loader means no name patterns and no second pass over the index, and it already + accounts for on-the-fly key conversion, which a set re-derived afterwards would have to + replay to avoid mistaking a renamed key for an unplaced one. + """ + model, loading_info = auto_class.from_pretrained(ckpt_path, output_loading_info=True, **kwargs) + unexpected = loading_info.get("unexpected_keys") or [] + record_unplaced_source_keys(model, _resolved_local_dir(ckpt_path), unexpected) + if unexpected: + print( + f"✓ {len(unexpected)} checkpoint key(s) the model has no parameter for " + f"(e.g. {min(unexpected)}); the export will carry them over unchanged." + ) + return model + + def get_model( ckpt_path, device="cuda", @@ -846,7 +753,8 @@ def has_pack_quantized_config(config): ) if is_speculative(hf_config): - model = AutoModelForCausalLM.from_pretrained( + model = _from_pretrained_recording( + AutoModelForCausalLM, ckpt_path, device_map=device_map, **model_kwargs, @@ -855,7 +763,8 @@ def has_pack_quantized_config(config): from modelopt.torch.quantization.plugins.huggingface import patch_compressed_linear_loading with patch_compressed_linear_loading(): - model = AutoModelForCausalLM.from_pretrained( + model = _from_pretrained_recording( + AutoModelForCausalLM, ckpt_path, device_map="auto", trust_remote_code=trust_remote_code, @@ -879,7 +788,8 @@ def has_pack_quantized_config(config): # materialization. Sequential keeps each shard's dequant on a single device # (the whole model lands on one GPU when it fits there). model_kwargs["quantization_config"] = Mxfp4Config(dequantize=True) - model = AutoModelForCausalLM.from_pretrained( + model = _from_pretrained_recording( + AutoModelForCausalLM, ckpt_path, device_map="cpu" if device == "cpu" else "sequential", **model_kwargs, @@ -965,7 +875,8 @@ def has_pack_quantized_config(config): model_kwargs2 = _apply_dtype_to_config(model_kwargs, config_dtype, architecture) if _disk_offload: model_kwargs2["offload_folder"] = offload_folder - model = auto_model_module.from_pretrained( + model = _from_pretrained_recording( + auto_model_module, ckpt_path, device_map=device_map, **model_kwargs2, @@ -1079,6 +990,7 @@ def copy_custom_model_files( export_path: str, trust_remote_code: bool = False, exclude_files: Iterable[str] | None = None, + copy_off_index_weights: bool = True, ): """Copy source checkpoint sidecar files to an HF PTQ export. @@ -1098,6 +1010,11 @@ def copy_custom_model_files( export_path: Path to the exported model directory trust_remote_code: Passed to HuggingFace model-ID resolution; does not control copying. exclude_files: Additional source file names to skip. + copy_off_index_weights: Copy safetensors the loader never opens (GLM-4.7's + ``mtp.safetensors``). Only the unified-HF export gives them meaning -- it seeds their + tensor names into ``quantization_config.ignore`` -- so a TensorRT-LLM export, whose + checkpoint is ``rank.safetensors`` plus its own ``config.json``, should pass False + rather than carry gigabytes nothing there reads. """ # Resolve the source path (handles both local paths and HF model IDs) resolved_source_path = _resolve_model_path(source_path, trust_remote_code) @@ -1130,6 +1047,13 @@ def copy_custom_model_files( exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, ) + # Safetensors the loader never opens are sidecars too: untouched by quantization and absent + # from the export, so copy them rather than leave them behind. Skipped by the call above, + # which excludes every *.safetensors to avoid re-emitting the unquantized source weights. + copied_weights = ( + copy_off_index_safetensors(source_dir, export_dir) if copy_off_index_weights else [] + ) + copied_files = [*copied_files, *copied_weights] if copied_files: for file_name in copied_files: print(f"Copied checkpoint sidecar file: {file_name}") diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 82dac668f92..1266857d844 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -41,9 +41,7 @@ is_enc_dec, is_nemotron_vl, layerwise_export_block, - load_mtp_weights, mlflow_run, - mtp_layer_prefixes_from_checkpoint, needs_checkpoint_path_update, recipe_layerwise_blocks, resolve_checkpoint_dir, @@ -621,11 +619,6 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, hf_config=hf_config, ) - # The FSDP2 loader drops MTP weights (re-attached BF16 at export); flag their prefixes now - # so the pre-quant exclusion below skips any MTP module from_config did build. - mtp_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) - if mtp_prefixes: - full_model._mtp_layer_prefixes = mtp_prefixes elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.pyt_ckpt_path, @@ -985,7 +978,15 @@ def export_quantized( ) # Copy custom model files (Python files and JSON configs) for TensorRT-LLM export - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) + # TRT-LLM checkpoints are rank.safetensors plus their own config; nothing + # there reads an off-index sidecar, and the exclude_modules seeding that gives + # one meaning happens only inside export_hf_checkpoint. + copy_custom_model_files( + args.pyt_ckpt_path, + export_path, + args.trust_remote_code, + copy_off_index_weights=False, + ) else: # Check arguments for unified_hf export format and set to default if unsupported arguments are provided assert args.sparsity_fmt == "dense", ( @@ -1001,20 +1002,19 @@ def export_quantized( # Load any missing weights from non-standard safetensors (handled in get_model for non-low-memory mode) # Store the MTP layer prefixes on the model for later exclusion from quantization if args.vllm_fakequant_export: + # save_pretrained inside the exporter writes model-backed state only; weights + # the loader could not place (an MTP head, an auxiliary tower) are carried over + # as an extra shard afterward -- see _carry_over_unplaced_weights. export_hf_vllm_fq_checkpoint( full_model, export_dir=export_path, inplace_mem_efficient=True ) else: - mtp_layer_prefixes, mtp_state_dict = load_mtp_weights( - full_model, args.pyt_ckpt_path - ) - if mtp_layer_prefixes: - full_model._mtp_layer_prefixes = mtp_layer_prefixes - + # Weights the loader could not place (an MTP head, an auxiliary tower) are + # carried over by the exporter from the keys recorded at load time; nothing + # architecture-specific is needed here. export_hf_checkpoint( full_model, export_dir=export_path, - extra_state_dict=mtp_state_dict, ) if args.qformat == "w4a16_nvfp4": @@ -1042,6 +1042,7 @@ def export_quantized( export_path, args.trust_remote_code, exclude_files=exclude_files, + copy_off_index_weights=not is_tensorrt_llm_export, ) args.checkpoint_exported = True @@ -1423,22 +1424,6 @@ def quantize_main( KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"], ) - # Exclude MTP layers from quantization if detected (e.g., GLM-4.7's layer 92). - # These layers are typically speculative decoding layers that should be exported as-is. - # Complementary to recipe `*mtp*` wildcards (name-match); this catches MTP layers - # identified by index. - mtp_layer_prefixes = getattr(full_model, "_mtp_layer_prefixes", None) - if args.layerwise_export and not mtp_layer_prefixes: - # Only the FSDP2 loader flags these before quantization, and the exclusions must - # be in quant_cfg before mtq.quantize converts the first layer. - mtp_layer_prefixes = mtp_layer_prefixes_from_checkpoint(args.pyt_ckpt_path) - if mtp_layer_prefixes: - quant_cfg = copy.deepcopy(quant_cfg) - for prefix in mtp_layer_prefixes: - pattern = f"*{prefix}*" - quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) - print(f"Excluding MTP layer from quantization: {pattern}") - # Before resolve_checkpoint_dir, which hashes the config: with the placeholder # still in it, two --export_path values would share one checkpoint dir. if args.layerwise_export: diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py index 9f34e701122..b62dd041a18 100644 --- a/modelopt/torch/export/__init__.py +++ b/modelopt/torch/export/__init__.py @@ -15,6 +15,22 @@ """Export package for Hugging Face and Megatron-based models.""" +# hf_checkpoint_utils lives under modelopt.torch.utils.plugins (general HF-checkpoint-file logic, +# not export-specific -- see its module docstring), but re-exported here since existing callers +# outside this package (e.g. modelopt.torch.puzzletron, examples/megatron_bridge) import these +# names from modelopt.torch.export directly. +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_hf_ckpt_remote_code, + copy_non_safetensor_files_from_ckpt, + copy_off_index_safetensors, + indexed_weight_map, + load_multimodal_components, + locate_source_keys, + off_index_safetensors_files, + resolve_checkpoint_file, + sanitize_hf_config_for_deployment, +) + from .convert_hf_config import * from .model_utils import * from .moe_utils import * diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 9767d5635a4..bf4eba20c19 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -44,10 +44,10 @@ _postprocess_single_tensor, get_quant_config, get_quantization_format, + seed_carried_over_exclusions, ) from .registry import ExportContext, PrepareMoEInputsRegistry from .unified_export_hf import ( - _add_mtp_exclusions, _dispatch_export_handler, _fuse_shared_input_modules, _prepare_moe_inputs, @@ -371,7 +371,6 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d model = self._ctx.model quant_config = self._quant_config - _add_mtp_exclusions(model, quant_config) # No gate/up sync here: export_layer did every layer, and the tail has no experts. if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None @@ -379,6 +378,17 @@ def finalize(self, extra_state_dict: dict[str, torch.Tensor] | None = None) -> d if self._name_mapper is not None and quant_config: with contextlib.suppress(Exception): revert_quant_config_names(quant_config.get("quantization", {}), self._name_mapper) + # After the reversal, not before: carried names are source-checkpoint names already, so + # passing them through the mapper would rewrite names that are correct as they stand. + # bind() snapshotted this config during calibration, so the carried set -- which + # export_hf_checkpoint records immediately before calling us -- is only visible now. + if quant_config: + seeded = seed_carried_over_exclusions(model, quant_config) + if seeded: + print( + f"Excluding {len(seeded)} carried-over module(s) from the layerwise " + f"quantization config (e.g. {seeded[0]})" + ) names = module_name_maps(model) # Recomputed, not snapshotted in __init__: calibration adds modules inside the diff --git a/modelopt/torch/export/plugins/__init__.py b/modelopt/torch/export/plugins/__init__.py index 8736ece1869..d54d423f38b 100644 --- a/modelopt/torch/export/plugins/__init__.py +++ b/modelopt/torch/export/plugins/__init__.py @@ -15,25 +15,12 @@ """Export package plugin.""" -from typing import Any - from modelopt.torch.utils import import_plugin with import_plugin("megatron_importer"): from .megatron_importer import * from .hf_spec_export import * - -with import_plugin("hf_checkpoint_utils"): - from .hf_checkpoint_utils import * - -if "sanitize_hf_config_for_deployment" not in globals(): - - def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) -> None: - """No-op fallback when Hugging Face checkpoint utilities are unavailable.""" - return None - - from .vllm_fakequant_hf import * with import_plugin("vllm_fakequant_megatron"): diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py deleted file mode 100644 index 9d508e58904..00000000000 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ /dev/null @@ -1,352 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Hugging Face checkpoint utility.""" - -import fnmatch -import json -import os -import shutil -import warnings -from collections.abc import Iterable -from pathlib import Path -from typing import Any - -import torch -from huggingface_hub import snapshot_download -from huggingface_hub.errors import LocalEntryNotFoundError -from safetensors.torch import safe_open -from tqdm import tqdm - -_HF_HUB_OFFLINE_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} - - -def _as_nonnegative_int(value: Any) -> int | None: - """Return ``value`` as an int when it is a non-negative integer.""" - if isinstance(value, bool): - return None - if isinstance(value, int) and value >= 0: - return value - return None - - -def _count_mtp_layer_prefixes(prefixes: list[Any] | tuple[Any, ...]) -> int | None: - """Count actual MTP layer prefixes, excluding broad prefixes like ``mtp``.""" - layer_prefixes = { - prefix - for prefix in prefixes - if isinstance(prefix, str) - and (parts := prefix.split(".")) - and len(parts) >= 2 - and parts[-2] == "layers" - and parts[-1].isdigit() - } - return len(layer_prefixes) or None - - -def _get_num_nextn_predict_layers(config_data: dict[str, Any], model: Any) -> int | None: - """Get the number of next-token-prediction layers from config metadata.""" - num_nextn_predict_layers = _as_nonnegative_int(config_data.get("num_nextn_predict_layers")) - if num_nextn_predict_layers is not None: - return num_nextn_predict_layers - - model_config = getattr(model, "config", None) - if model_config is not None: - num_nextn_predict_layers = _as_nonnegative_int( - getattr(model_config, "num_nextn_predict_layers", None) - ) - if num_nextn_predict_layers is not None: - return num_nextn_predict_layers - - mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) - if isinstance(mtp_layer_prefixes, (list, tuple)): - return _count_mtp_layer_prefixes(mtp_layer_prefixes) - - return None - - -def _get_rope_theta(config_data: dict[str, Any], model: Any) -> Any: - """Return rope_theta from exported config data or the in-memory model config.""" - rope_theta = config_data.get("rope_theta") - if rope_theta is not None: - return rope_theta - - model_config = getattr(model, "config", None) - if model_config is None: - return None - - return getattr(model_config, "rope_theta", None) - - -def _sanitize_llama3_rope_config(config_data: dict[str, Any], model: Any) -> None: - """Fill missing llama3 rope_theta in rope config metadata when available.""" - rope_theta = _get_rope_theta(config_data, model) - if rope_theta is None: - return - - for key in ("rope_parameters", "rope_scaling"): - rope_config = config_data.get(key) - if not isinstance(rope_config, dict): - continue - - rope_type = rope_config.get("rope_type", rope_config.get("type")) - if rope_type == "llama3" and "rope_theta" not in rope_config: - rope_config["rope_theta"] = rope_theta - - -def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) -> None: - """Sanitize exported Hugging Face config metadata for deployment runtimes. - - Fix conservative deployment-only config incompatibilities: - - * add missing llama3 ``rope_theta`` metadata when available; - * trim trailing MTP/next-token-prediction ``layer_types`` entries only when - the mismatch is exactly explained by next-token-prediction metadata. - """ - _sanitize_llama3_rope_config(config_data, model) - - num_hidden_layers = _as_nonnegative_int(config_data.get("num_hidden_layers")) - layer_types = config_data.get("layer_types") - if num_hidden_layers is None or not isinstance(layer_types, list): - return - - num_layer_types = len(layer_types) - if num_layer_types == num_hidden_layers: - return - - num_nextn_predict_layers = _get_num_nextn_predict_layers(config_data, model) - if ( - num_layer_types > num_hidden_layers - and num_nextn_predict_layers == num_layer_types - num_hidden_layers - ): - warnings.warn( - "Trimming config.layer_types from " - f"{num_layer_types} to {num_hidden_layers} entries so it matches " - "num_hidden_layers; the removed entries correspond to " - "num_nextn_predict_layers.", - stacklevel=2, - ) - config_data["layer_types"] = layer_types[:num_hidden_layers] - - -def _is_hf_hub_offline() -> bool: - return os.environ.get("HF_HUB_OFFLINE", "").strip().upper() in _HF_HUB_OFFLINE_TRUE_VALUES - - -def _copy_python_files(source_dir: Path, save_dir: Path) -> None: - for py_file in source_dir.glob("*.py"): - shutil.copy2(py_file, save_dir / py_file.name) - - -def copy_hf_ckpt_remote_code( - pretrained_model_path: str | os.PathLike, save_directory: str | os.PathLike -): - """Copy remote code from pretrained model to save directory. - - For models that keep configuration and modeling files as part of the checkpoint, - we need to copy them to the export directory for seamless integration with inference - frameworks. - - If ``pretrained_model_path`` is a local directory, Python files are copied directly. - If it's a HF Hub model ID (e.g. ``nvidia/NVIDIA-Nemotron-Nano-12B-v2``), the Hub - snapshot is resolved first and Python files are copied from that snapshot. When - ``HF_HUB_OFFLINE`` is set, the snapshot must already be available in the local - Hugging Face cache. - - Args: - pretrained_model_path: Local path to the pretrained model or HuggingFace Hub model ID. - save_directory: Path to the save directory. - """ - hf_checkpoint_path = Path(pretrained_model_path) - save_dir = Path(save_directory) - save_dir.mkdir(parents=True, exist_ok=True) - - if hf_checkpoint_path.is_dir(): - _copy_python_files(hf_checkpoint_path, save_dir) - else: - local_files_only = _is_hf_hub_offline() - try: - source_dir = Path( - snapshot_download( - repo_id=str(pretrained_model_path), - allow_patterns=["*.py"], - local_files_only=local_files_only, - ) - ) - except LocalEntryNotFoundError as exc: - if local_files_only: - raise RuntimeError( - f"Could not copy Python sidecar files for {pretrained_model_path!r} because " - "HF_HUB_OFFLINE is enabled and the files are not available in the local " - "Hugging Face cache. Populate the cache with the model's *.py files or pass " - "a local pretrained model directory." - ) from exc - raise - - _copy_python_files(source_dir, save_dir) - - -def load_multimodal_components( - pretrained_model_path: str | os.PathLike, - prefixes: tuple[str, ...] = ("multi_modal_projector", "vision_model"), -) -> dict[str, torch.Tensor]: - """Load multimodal components from safetensors file. - - Args: - pretrained_model_path: Directory or HuggingFace repo id of the pretrained model. - prefixes: Tensor key prefixes to select. Defaults to the LLaVA-style - ``multi_modal_projector`` / ``vision_model`` prefixes. Pass - ``("model.visual.",)`` for Qwen3-VL checkpoints. - - Returns: - A dictionary of multimodal components. - """ - hf_checkpoint_path = Path(pretrained_model_path) - if not hf_checkpoint_path.is_dir(): - # Also accept a repo id, which is what the example scripts pass to quantize.py. - # Fetched in two stages: the vision tower is a small fraction of a VLM checkpoint, so - # pulling every shard to keep a few would waste tens of GB. - local_files_only = _is_hf_hub_offline() - repo_id = str(pretrained_model_path) - try: - index_dir = Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=["model.safetensors.index.json"], - local_files_only=local_files_only, - ) - ) - except (LocalEntryNotFoundError, OSError, ValueError) as exc: - raise ValueError( - f"Invalid pretrained model path: {pretrained_model_path}. It should be a " - "directory or an available HuggingFace repo id." - ) from exc - - index_file = index_dir / "model.safetensors.index.json" - if index_file.is_file(): - try: - weight_map = json.loads(index_file.read_text())["weight_map"] - except (json.JSONDecodeError, KeyError) as exc: - raise ValueError(f"Malformed safetensors index in {repo_id}.") from exc - wanted = sorted( - {shard for key, shard in weight_map.items() if key.startswith(prefixes)} - ) - else: - wanted = ["model.safetensors"] # unsharded checkpoint - # Kept separate from the resolution failure above: a hub outage or a full disk here is - # retryable, not a bad path. - hf_checkpoint_path = Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=["model.safetensors.index.json", *wanted], - local_files_only=local_files_only, - ) - ) - - safetensors_file = Path(hf_checkpoint_path) / "model.safetensors" - safetensors_index_file = Path(hf_checkpoint_path) / "model.safetensors.index.json" - - multimodal_state_dict = {} - - if safetensors_file.is_file(): - print(f"Loading multimodal components from single file: {safetensors_file}") - with safe_open(safetensors_file, framework="pt") as f: - multimodal_keys = [ - key - for key in f.keys() # noqa: SIM118 - if key.startswith(prefixes) - ] - for key in tqdm(multimodal_keys, desc="Loading multimodal tensors"): - multimodal_state_dict[key] = f.get_tensor(key) - - elif safetensors_index_file.is_file(): - print(f"Loading multimodal components from sharded model: {hf_checkpoint_path}") - with open(safetensors_index_file) as f: - safetensors_index = json.load(f) - - all_shard_files = sorted( - { - shard - for key, shard in safetensors_index["weight_map"].items() - if key.startswith(prefixes) - } - ) - for shard_file in all_shard_files: - safetensors_filepath = Path(hf_checkpoint_path) / shard_file - with safe_open(safetensors_filepath, framework="pt") as f: - for key in f.keys(): # noqa: SIM118 - if key.startswith(prefixes): - multimodal_state_dict[key] = f.get_tensor(key) - - else: - print(f"Warning: No safetensors files found in {hf_checkpoint_path}") - - if not multimodal_state_dict: - raise ValueError( - f"No tensors under {prefixes} in {pretrained_model_path}; the vision tower would be " - "missing from the export. The checkpoint's prefixes have likely changed." - ) - - print(f"Successfully loaded {len(multimodal_state_dict)} multimodal tensors") - return multimodal_state_dict - - -def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: - return any(fnmatch.fnmatchcase(file_name, pattern) for pattern in patterns) - - -def copy_non_safetensor_files_from_ckpt( - src: str | os.PathLike, - dst: str | os.PathLike, - *, - exclude_files: Iterable[str] | None = None, - exclude_patterns: Iterable[str] | None = None, -) -> list[str]: - """Copy every non-safetensors file from a local HF checkpoint dir verbatim. - - Use as a baseline so tokenizer files, remote_code ``*.py``, README, LICENSE, etc. - are preserved from the source. Callers can exclude additional files or patterns when - copying after export-owned metadata has already been written. - - Args: - src: Source HF checkpoint directory. Must be a local path. - dst: Destination directory; created if missing. - exclude_files: Exact file names to skip. - exclude_patterns: Glob patterns for additional files to skip. - - Returns: - File names copied into ``dst``. - """ - if not os.path.isdir(src): - raise ValueError(f"Invalid source path: {src}. It should be a directory.") - exclude_files = set(exclude_files or ()) - exclude_patterns = tuple(exclude_patterns or ()) - copied_files = [] - os.makedirs(dst, exist_ok=True) - for entry in sorted(os.listdir(src)): - if entry in exclude_files or _matches_any_pattern(entry, exclude_patterns): - continue - sp = os.path.join(src, entry) - if not os.path.isfile(sp): - continue - if entry.endswith(".safetensors") or entry == "model.safetensors.index.json": - continue - try: - shutil.copy2(sp, dst) - except OSError as error: - warnings.warn(f"Failed to copy checkpoint sidecar {entry}: {error}") - continue - copied_files.append(entry) - return copied_files diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index f8e73c4fa2c..94ea8003200 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -15,6 +15,7 @@ """Export HuggingFace model to vLLM fakequant checkpoint.""" import copy +import json import logging import re import warnings @@ -25,6 +26,8 @@ import torch import torch.nn as nn +from safetensors import safe_open +from safetensors.torch import save_file import modelopt.torch.opt as mto from modelopt.torch.models import hf_model_type, is_moe @@ -48,7 +51,7 @@ from ..layer_utils import get_experts_list from ..quant_utils import get_quantization_format -from ..unified_export_hf import collect_shared_input_modules +from ..unified_export_hf import collect_shared_input_modules, read_unplaced_weights __all__ = [ "export_hf_vllm_fq_checkpoint", @@ -526,6 +529,46 @@ def _dummy_forward() -> None: return out, requant_weights +def _carry_over_unplaced_weights(export_dir: Path, model: nn.Module) -> None: + """Write checkpoint weights the model never held as an extra safetensors shard. + + ``save_pretrained`` only ever writes ``model.state_dict()`` (or a copy of it), so a + checkpoint weight with no parameter in the built model -- an MTP head, an auxiliary + tower -- is never in what it saves, regardless of whether ``state_dict=`` was passed + explicitly. :func:`read_unplaced_weights` reads those tensors back from the source + checkpoint; this writes them as their own shard and rebuilds the index from every + shard on disk (mirroring ``LayerwiseExporter._write_index``), which sidesteps the + single-file-vs-sharded distinction ``save_pretrained`` may have already chosen. + + A no-op when there is nothing to carry (the common case: most models have no + unplaced weights at all). + """ + extra = read_unplaced_weights(model) + if not extra: + return + + shard_name = "model-carried-over.safetensors" + save_file( + {k: v.detach().contiguous().cpu() for k, v in extra.items()}, str(export_dir / shard_name) + ) + print( + f"Carrying {len(extra)} checkpoint weight(s) the model has no parameter for into {shard_name}" + ) + + weight_map: dict[str, str] = {} + total_size = 0 + for shard in sorted(export_dir.glob("*.safetensors")): + with safe_open(str(shard), framework="pt") as f: + for key in f.keys(): # noqa: SIM118 -- safe_open has no __iter__ + weight_map[key] = shard.name + with open(shard, "rb") as fh: + header_len = int.from_bytes(fh.read(8), "little") + total_size += shard.stat().st_size - 8 - header_len + + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + (export_dir / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + + def export_hf_vllm_fq_checkpoint( model: nn.Module, export_dir: Path | str, @@ -705,6 +748,14 @@ def export_hf_vllm_fq_checkpoint( else: model.save_pretrained(export_dir, state_dict=clean_sd, save_modelopt_state=False) + # Step 4: carry over checkpoint weights the model never held (an MTP head, an + # auxiliary tower). save_pretrained above wrote only model-backed state, same as + # export_hf_checkpoint's post_state_dict; this writes the rest as an extra shard, + # the one thing save_pretrained's state_dict= path cannot do for the + # inplace_mem_efficient branch (it deliberately omits state_dict= there -- see the + # comment above -- so there is no state_dict to merge extras into). + _carry_over_unplaced_weights(export_dir, model) + finally: if not inplace_mem_efficient: for wq, orig_rotate in wqs_to_restore: diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index d3f2b2f6b56..c36369bab19 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -15,6 +15,7 @@ """Utils for quantization including scaling factors adjustments.""" +import fnmatch import logging from collections import defaultdict from collections.abc import Generator @@ -1662,6 +1663,76 @@ def preprocess_linear_fusion(modules: list[torch.nn.Module], resmooth_only=False module.weight_quantizer.amax = weight_amax +def seed_carried_over_exclusions(model: nn.Module, quant_config: dict) -> list[str]: + """Add carried-weight module names to an already-built ``quant_config``'s exclusions. + + The single place carried weights reach ``exclude_modules``, for both exporters. + :func:`get_quant_config` calls it once the per-layer pass is done, and the layerwise exporter + calls it again from ``finalize()`` -- it snapshots its config during ``bind()``, while + calibration is still running and long before ``export_hf_checkpoint`` records what it carried, + so it has no chance to see them any earlier. Without it a layerwise export copies GLM-4.7's + ``mtp.safetensors`` into the checkpoint with nothing in ``exclude_modules``: the same + NVBug 5718750 failure this pass exists to prevent, reached through the other exporter. + + Exclusions are exact module names rather than prefix wildcards. That keeps both exporters + emitting the same thing, and a literal can never over-match a module the export did in fact + quantize -- the risk :func:`_prefix_wildcard_summarize_exclude_modules` has to guard against by + consulting ``quantized_layers``, which is unavailable by the time the layerwise path runs. + + Returns the names it added. No-op when the export is not uniformly quantized -- there is no + single ``quant_algo`` for a deployment framework to misapply, so there is nothing to exclude + a weight from. + """ + names = _get_carried_over_module_names(model) + if not names: + return [] + quantization = quant_config.get("quantization") + if not isinstance(quantization, dict): + return [] + if quantization.get("quant_algo") in (None, QUANTIZATION_NONE, "MIXED_PRECISION"): + return [] + exclude_modules = quantization.setdefault("exclude_modules", []) + added = [n for n in names if not any(fnmatch.fnmatch(n, p) for p in exclude_modules)] + if added: + exclude_modules.extend(added) + exclude_modules.sort() + return added + + +def _get_carried_over_module_names(model: nn.Module) -> list[str]: + """Return module names for checkpoint weights carried over without a module. + + Weights the loader could not place -- an MTP head, an auxiliary tower -- are copied into + the export verbatim from the source checkpoint (see + :func:`modelopt.torch.export.unified_export_hf.read_unplaced_weights`). They + have no module in the live model, so the quantizer walk in :func:`get_quant_config` cannot + see them and would leave them out of ``exclude_modules`` even though their original-precision + weight is written to the checkpoint. A deployment framework then reads the top-level + ``quant_algo`` and tries to load e.g. an MTP ``eh_proj`` as an FP8 weight. + + This is the same failure the MoE-router pass above exists to prevent -- only the reason the + module is invisible differs (no quantizer there, no module at all here). + + Prefers ``_modelopt_carried_over_names``, which the export records once it knows what it + actually wrote -- carried tensors plus the off-index sidecars copied verbatim. Those sidecars + are never ``unexpected_keys``, so the unplaced list alone would miss GLM-4.7's + ``mtp.safetensors`` and leave its tensors in the export with nothing in ``exclude_modules``. + + A state-dict key is ``.``, so the owning module is the key with + its last component removed. Keys without a dot are top-level tensors with no module and are + skipped. + """ + keys = getattr(model, "_modelopt_carried_over_names", None) + if keys is None: + # Export has not recorded yet (or this model never went through it). The recorded unplaced + # list is the best available answer; it is wider than what gets written, so it can name a + # module the export did not emit. That way round is harmless -- a deployment framework + # ignores an exclusion it finds no weight for, but fails loading one it was never told + # about. + keys = getattr(model, "_modelopt_unplaced_source_keys", None) or [] + return sorted({key.rsplit(".", 1)[0] for key in keys if "." in key}) + + def _get_unquantized_moe_router_names(model: nn.Module) -> list[str]: """Return the names of MoE router/gate submodules left in original precision. @@ -1834,6 +1905,15 @@ def get_quant_config( # Process per layer quantization config dict quant_config["quantization"].update(process_layer_quant_config(layer_config_dict)) + # Carried weights are seeded AFTER the per-layer pass, through the same helper the layerwise + # exporter calls. Seeding them into layer_config_dict instead would route them through + # _prefix_wildcard_summarize_exclude_modules and emit wildcards here, while the layerwise path + # -- which can only act once its config is already built -- emits literals: one model, two + # exporters, two different-looking quantization_config.ignore. The summarizer cannot serve both, + # because it needs `quantized_layers` to avoid a wildcard swallowing a quantized module and + # process_layer_quant_config pops that key before returning. + seed_carried_over_exclusions(model, quant_config) + weight_quant_algo = quant_config["quantization"].get("quant_algo") needs_layerwise_kv_metadata = bool(kv_cache_quantized_layers) and ( weight_quant_algo is None or len(kv_cache_formats) > 1 diff --git a/modelopt/torch/export/shard_cast_utils.py b/modelopt/torch/export/shard_cast_utils.py index 8da2761a36e..e38b577caa1 100644 --- a/modelopt/torch/export/shard_cast_utils.py +++ b/modelopt/torch/export/shard_cast_utils.py @@ -36,6 +36,7 @@ mxfp4_to_nvfp4_global_amax, mxfp4_to_nvfp4_per_block_amax, ) +from modelopt.torch.utils.plugins.hf_checkpoint_utils import resolve_checkpoint_file if TYPE_CHECKING: from collections.abc import Callable, Collection @@ -52,7 +53,6 @@ _MXFP4_BLOCK = 32 _MXFP4_BYTES_PER_BLOCK = 16 _NVFP4_BLOCK = 16 -_MAX_CHECKPOINT_METADATA_BYTES = 128 * 1024 * 1024 def dequantize_mxfp4_to_bf16( @@ -193,46 +193,6 @@ def link_or_copy(src: Path, dst: Path) -> None: shutil.copy2(src, dst) -def _is_relative_to(path: Path, root: Path) -> bool: - return path == root or root in path.parents - - -def _snapshot_blob_root(source_root: Path) -> Path | None: - if source_root.parent.name != "snapshots": - return None - blob_root = source_root.parent.parent / "blobs" - return blob_root.resolve(strict=True) if blob_root.is_dir() else None - - -def _allowed_source_roots(src_dir: Path) -> list[Path]: - source_root = src_dir.resolve(strict=True) - allowed_roots = [source_root] - if blob_root := _snapshot_blob_root(source_root): - allowed_roots.append(blob_root) - return allowed_roots - - -def resolve_checkpoint_file( - src_dir: Path, - relative_path: str | Path, - *, - max_bytes: int | None = _MAX_CHECKPOINT_METADATA_BYTES, -) -> Path: - """Resolve a contained regular checkpoint file and optionally bound its size.""" - src = src_dir / relative_path - try: - resolved_src = src.resolve(strict=True) - except OSError as exc: - raise ValueError(f"checkpoint source is not a readable regular file: {src}") from exc - if not resolved_src.is_file(): - raise ValueError(f"checkpoint source must resolve to a regular file: {src}") - if not any(_is_relative_to(resolved_src, root) for root in _allowed_source_roots(src_dir)): - raise ValueError(f"checkpoint source is outside the checkpoint directory: {src}") - if max_bytes is not None and resolved_src.stat().st_size > max_bytes: - raise ValueError(f"checkpoint source exceeds the {max_bytes}-byte size limit: {src}") - return resolved_src - - def _collect_aux_files( src_dir: Path, *, diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4aeddac0770..28edade5d7b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -36,6 +36,11 @@ from safetensors.torch import save_file from modelopt.torch.models import hf_model_type, is_moe +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + locate_source_keys, + off_index_safetensors_files, + sanitize_hf_config_for_deployment, +) from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales @@ -73,6 +78,7 @@ quantizer_attr_names, ) from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload +from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.distributed import is_fsdp2_model from modelopt.torch.utils.perf import maybe_clear_cuda_cache @@ -87,7 +93,7 @@ from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import get_experts_list, is_layernorm, is_quantlinear, sync_moe_gate_up_amax from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model -from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment +from .plugins import SpeculativeDecodingExporter, has_spec_opt from .quant_aware_conversion import ( build_reverse_name_mapper, revert_quant_config_names, @@ -915,22 +921,6 @@ def _prepare_moe_inputs( handler(name, sub_module, prepare_ctx) -def _add_mtp_exclusions(model: nn.Module, quant_config: dict) -> None: - """Add MTP layer prefixes to exclude_modules if they were excluded from quantization. - - This ensures they appear in ``quantization_config["ignore"]`` in ``config.json``. - """ - mtp_layer_prefixes = getattr(model, "_mtp_layer_prefixes", None) - if mtp_layer_prefixes: - exclude_modules = quant_config["quantization"].setdefault("exclude_modules", []) - for prefix in mtp_layer_prefixes: - # Add wildcard pattern to exclude all submodules under this MTP layer - pattern = f"{prefix}*" - if pattern not in exclude_modules: - exclude_modules.append(pattern) - print(f"Adding MTP layer to quantization_config ignore: {pattern}") - - def _warn_on_unsynced_moe_gate_up(model: nn.Module) -> None: """Safety net for gate/up weight quantizer amaxes that resmoothing did not reach. @@ -1014,8 +1004,6 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) - _add_mtp_exclusions(model, quant_config) - _warn_on_unsynced_moe_gate_up(model) # Merge per-side input_quantizer amaxes BEFORE export, so the retained tied weight's single @@ -1622,6 +1610,213 @@ def _revert_quant_config_names_best_effort( return hf_quant_config +def _source_checkpoint(model: nn.Module) -> str | None: + """Where the model's original checkpoint lives, or ``None`` if nothing is known. + + Prefers ``_modelopt_source_checkpoint`` (recorded at load time by + ``record_unplaced_source_keys``). A model that reached export without going through that path + still knows its own provenance via ``config._name_or_path``, and the several places this is + asked must agree on the answer -- otherwise, for instance, a weight is carried but its sidecar + tensors never reach ``exclude_modules`` because the two halves disagreed about where the + checkpoint was. + """ + return getattr(model, "_modelopt_source_checkpoint", None) or getattr( + getattr(model, "config", None), "_name_or_path", None + ) + + +def carryable_unplaced_keys(model: nn.Module) -> list[str]: + """Unplaced checkpoint keys a shard actually provides. + + ``_modelopt_unplaced_source_keys`` answers "does the model have a parameter for this key", + which is a wider question than "is there a tensor to carry". Checkpoints routinely list keys + no shard backs -- a stale ``rotary_emb.inv_freq`` buffer, leftovers after a + ``conversion_mapping`` rename -- and those carry nothing. Callers that want to know whether + real weights would be dropped must ask this, not the raw recorded list. + + Best-effort by design: it is used to decide how loudly to complain, so an unreadable index + answers "nothing to carry" rather than raising from inside a diagnostic. + """ + ckpt = _source_checkpoint(model) + if not ckpt or not Path(ckpt).is_dir(): + return [] + keys = unplaced_keys_for(model, ckpt) + if not keys: + return [] + try: + located = locate_source_keys(ckpt, keys) + except Exception: + return [] + return sorted(located) + + +def off_index_tensor_names(model: nn.Module) -> list[str]: + """Tensor names in the checkpoint's off-index safetensors sidecars. + + Those files (GLM-4.7's ``mtp.safetensors``) are copied into the export verbatim rather than + loaded, so they are never ``unexpected_keys`` and :func:`carryable_unplaced_keys` cannot see + them -- yet their tensors land in the export in original precision exactly like a carried + weight, and must reach ``exclude_modules`` the same way. Before this mechanism existed + ``_add_mtp_exclusions`` covered them by globbing for ``mtp*``. + + Reads safetensors headers only, never tensor data, and stays silent when the sidecars or the + library cannot be read: an absent exclusion is a deployment problem, but so is an export that + dies while computing one. + """ + ckpt = _source_checkpoint(model) + if not ckpt or not Path(ckpt).is_dir(): + return [] + try: + names: list[str] = [] + for file_name in off_index_safetensors_files(ckpt): + with safe_open(str(Path(ckpt) / file_name), framework="pt") as f: + names.extend(f.keys()) + return sorted(set(names)) + except Exception: + return [] + + +def unplaced_keys_for(model: nn.Module, ckpt: "str | Path") -> list[str]: + """Every source key the model has no parameter for, from both accountings. + + Neither source is sufficient alone, and each covers the other's blind spot: + + * The **structural** pass (``unplaced_source_keys``) walks the checkpoint index and asks, for + each source key, whether the built model has a parameter for it -- resolving the key through + Transformers' renames and converters first. Being structural it is unaffected by + ``_keys_to_ignore_on_load_unexpected``, and being source-keyed it names tensors that exist on + disk. It cannot see a key the index omits. + * The **recorded** set is the loader's own ``unexpected_keys``. It sees tensors the index omits, + because the loader enumerates the contents of every shard it opens. But Transformers filters + it through the architecture's ignore rules -- Qwen3-Next drops ``^mtp.*``, DeepSeek-V3 and GLM + drop their MTP prefixes -- so for exactly the MTP heads this mechanism exists to carry it can + come back EMPTY. It can also report post-conversion target names (a fused + ``...experts.gate_up_proj``) that exist in no shard, and one such name can stand for several + source tensors. + + Taking the union means an architecture's ignore rules cannot hide a weight, a fused name cannot + strand the tensors behind it, and a key missing from the index is still carried. Recorded names + that no shard backs are dropped by :func:`locate_source_keys`, which is the right outcome: the + structural pass has already named the real source keys for any converted target. + + The structural pass needs ``model_load_utils``, which imports transformers and accelerate at + module scope. Where those are absent the recorded set stands alone -- degraded, but no worse + than before this function existed. + """ + recorded = getattr(model, "_modelopt_unplaced_source_keys", None) or [] + structural: list[str] = [] + try: + from modelopt.torch.utils.plugins.model_load_utils import unplaced_source_keys + + structural = unplaced_source_keys(model, str(ckpt)) + except Exception as exc: + # Warn regardless of whether anything was recorded: a non-empty recorded set is not + # proof the union is complete, since the structural pass is what catches keys an + # architecture's ignore rules dropped from the recorded set in the first place. Staying + # quiet here just because recorded happens to be non-empty is the same silent + # degradation this function exists to avoid. + warnings.warn( + f"Could not derive unplaced source keys structurally ({exc}); relying on the " + "loader's report, which its architecture may have filtered." + ) + return sorted({*structural, *recorded}) + + +# Placeholder for keys_only resolution: the name is real, the tensor is never read. +_NO_TENSOR: Any = None + + +def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[str, torch.Tensor]: + """Read back checkpoint weights the model never loaded, so the export stays complete. + + A checkpoint can hold parameters the built model has no home for -- an MTP head, an auxiliary + tower -- which means quantization never sees them and they would be missing from the exported + checkpoint unless they are copied across verbatim. ``parallel_load_and_prepare_fsdp2`` records + which keys those were (:attr:`_modelopt_unplaced_source_keys`) and where they came from; this + reads them on demand rather than holding them in memory from load to export. + + Deliberately architecture-agnostic: the question asked at load time was "does the model have a + parameter for this checkpoint key", not "is this an MTP head", so anything the model did not + load is carried through. Returns an empty dict when the model was not loaded that way. + + Best-effort: a checkpoint that cannot be re-read warns rather than failing the export, since + the rest of the weights are already correct. + + With ``keys_only`` the shard lookup still runs -- so the answer matches what a real read would + carry -- but no tensor is loaded and the values are ``None``. Ranks that do not write the extra + state use this: the FSDP2 writer only emits ``extra_state_dict`` from rank 0, so having every + rank materialize an MTP head (10 GB+ in bf16 on a large MoE) is host memory read and dropped. + """ + # The recorded list is never treated as final, empty or not. An earlier revision returned + # early when the loader recorded [], reasoning that it "had already answered" -- but + # Transformers filters unexpected_keys through _keys_to_ignore_on_load_unexpected, and + # Qwen3-Next ignores ^mtp.*, so for exactly the MTP heads this exists to carry the report comes + # back empty and the export silently shipped without them. See unplaced_keys_for. + # + # These are pure filesystem checks and deliberately sit ABOVE the try below: deciding that + # there is nothing to carry must not depend on an optional import succeeding. + ckpt = _source_checkpoint(model) + if not ckpt or not Path(ckpt).is_dir(): + # A hub id rather than a local path, or no provenance at all -- nothing to read. Quiet + # when nothing was recorded, because there is no reason to think anything is missing. But + # if the loader DID report unplaced keys, they are about to be dropped, and dropping them + # silently is the failure this whole path exists to prevent. + recorded = getattr(model, "_modelopt_unplaced_source_keys", None) or [] + if recorded: + warnings.warn( + f"Could not copy {len(recorded)} unplaced source weight(s) into the export: " + f"the source checkpoint is not a readable directory ({ckpt}); " + "the checkpoint will be missing them." + ) + return {} + # A checkpoint with no safetensors at all (pytorch_model.bin) has nothing this path can read. + # Returning quietly is right: there is nothing to carry, so warning about weights "missing" + # from the export would be alarming and wrong. + if ( + not (Path(ckpt) / "model.safetensors.index.json").exists() + and not (Path(ckpt) / "model.safetensors").exists() + ): + return {} + + try: + keys = unplaced_keys_for(model, ckpt) + if not keys: + return {} + + by_file: dict[str, list[str]] = {} + for k, shard in locate_source_keys(ckpt, list(keys)).items(): + by_file.setdefault(shard, []).append(k) + + if keys_only: + # The caller wants the names, not the bytes: every rank needs an identical key list + # to build an identical quant config, but only the writing rank should pay the memory. + return dict.fromkeys( + (k for shard_keys in by_file.values() for k in shard_keys), _NO_TENSOR + ) + + out: dict[str, torch.Tensor] = {} + for shard, shard_keys in by_file.items(): + with safe_open(str(Path(ckpt) / shard), framework="pt") as f: + for k in shard_keys: + out[k] = f.get_tensor(k) + except Exception as exc: + # ``keys`` can still be None here -- the failure may land before it is resolved (a failed + # import, or unplaced_source_keys itself raising). Counting it unconditionally would throw + # from inside the handler whose whole job is to leave the export standing. + count = f"{len(keys)} " if keys is not None else "" + warnings.warn( + f"Could not copy {count}unplaced source weight(s) into the export ({exc}); " + "the checkpoint will be missing them." + ) + return {} + if out and not keys_only: + print_rank_0( + f"Carrying {len(out)} source weight(s) the model never loaded into the export " + f"(e.g. {min(out)})" + ) + return out + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1657,6 +1852,32 @@ def export_hf_checkpoint( :func:`_postprocess_safetensors` for diffusion model exports. See its docstring for supported keys. """ + # Weights the model never loaded (MTP head, auxiliary tower, ...) are copied straight from the + # source so the exported checkpoint is the complete model. Merged here, ahead of the path + # dispatch, so the gather and no-gather writers behave identically. An explicit extra_state_dict + # wins on conflict: the caller asked for that tensor by name. + # Only the rank that writes extra_state_dict should read it. _export_fsdp2_checkpoint_streaming + # emits it from rank 0 alone, so on the other ranks a full read is host memory spent and thrown + # away -- and the carried set is the large stuff. The KEY list is still resolved everywhere, + # because get_quant_config runs per rank and the configs have to agree. + _writes_extra = ( + not ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and is_fsdp2_model(model) + ) + or torch.distributed.get_rank() == 0 + ) + _carried = read_unplaced_weights(model, keys_only=not _writes_extra) + if _writes_extra and _carried: + extra_state_dict = {**_carried, **(extra_state_dict or {})} + # Everything the export writes in original precision straight from the source, by either + # mechanism: tensors carried above, and the off-index sidecars copied verbatim alongside. + # get_quant_config reads this to seed exclude_modules; recorded here because it runs before + # that, and because only this point knows what was actually written rather than what was + # merely unplaced. + model._modelopt_carried_over_names = sorted({*_carried, *off_index_tensor_names(model)}) + from .layerwise_export import LAYERWISE_EXPORTER_ATTR exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 32477173e51..7388486fe6c 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -52,7 +52,6 @@ ) from .registry import ExportContext from .unified_export_hf import ( - _add_mtp_exclusions, _dispatch_export_handler, _prepare_model_for_export, _prepare_moe_inputs, @@ -372,8 +371,6 @@ def _export_transformers_checkpoint_streaming( quant_config = get_quant_config(model, is_modelopt_qlora=is_modelopt_qlora) - _add_mtp_exclusions(model, quant_config) - _warn_on_unsynced_moe_gate_up(model) # --- Per-tensor constants --- diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 9e5facda2c3..215a62ff85e 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -38,13 +38,13 @@ from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 - -from .convert_hf_config import convert_hf_quant_config_format -from .plugins.hf_checkpoint_utils import ( +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_hf_ckpt_remote_code, copy_non_safetensor_files_from_ckpt, load_multimodal_components, ) + +from .convert_hf_config import convert_hf_quant_config_format from .plugins.mcore_common import ( all_mcore_hf_export_mapping, all_mcore_hf_vision_passthrough_mapping, diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 4a830bc265a..d2750ad2aef 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -637,17 +637,21 @@ def restore_draft_precision(self, checkpoint_dir=None): def _reload_draft_weights_at_stored_precision(self, checkpoint_dir): """Copy the draft's tensors back out of the checkpoint at the dtype they were saved in.""" - # Imported here rather than at module scope: that module pulls in accelerate, - # huggingface_hub and torch.distributed.tensor. - from modelopt.torch.utils.plugins.model_load_utils import ( + # Imported here rather than at module scope to keep this file's own import-time + # footprint minimal; hf_checkpoint_utils itself only needs huggingface_hub + safetensors. + from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + indexed_weight_map, read_safetensors_subset, - weight_map_for, ) + weight_map = indexed_weight_map(checkpoint_dir) + if not weight_map: + raise RuntimeError( + f"No safetensors checkpoint at {checkpoint_dir} " + "(expected model.safetensors or model.safetensors.index.json)." + ) prefix = "dflash_module." - stored = read_safetensors_subset( - checkpoint_dir, weight_map_for(checkpoint_dir), lambda k: k.startswith(prefix) - ) + stored = read_safetensors_subset(checkpoint_dir, weight_map, lambda k: k.startswith(prefix)) own = dict(self.dflash_module.named_parameters()) with torch.no_grad(): for key, saved in stored.items(): diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py new file mode 100644 index 00000000000..033716a9b6f --- /dev/null +++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py @@ -0,0 +1,662 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Hugging Face checkpoint utility. + +General-purpose logic about the on-disk shape of an HF checkpoint (index, shards, sidecars) -- +not export-specific, so it lives under ``modelopt.torch.utils.plugins`` alongside +``model_load_utils`` rather than under ``modelopt.torch.export``. Deliberately independent of +that module's ``transformers``/``accelerate`` module-scope imports; only needs +``huggingface_hub`` + ``safetensors``. +""" + +import contextlib +import fnmatch +import json +import os +import re +import shutil +import warnings +from collections.abc import Callable, Iterable +from pathlib import Path +from typing import Any + +import torch +from huggingface_hub import snapshot_download +from huggingface_hub.errors import LocalEntryNotFoundError +from safetensors.torch import safe_open +from tqdm import tqdm + +_HF_HUB_OFFLINE_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} + + +def _as_nonnegative_int(value: Any) -> int | None: + """Return ``value`` as an int when it is a non-negative integer.""" + if isinstance(value, bool): + return None + if isinstance(value, int) and value >= 0: + return value + return None + + +def _get_num_nextn_predict_layers(config_data: dict[str, Any], model: Any) -> int | None: + """Get the number of next-token-prediction layers from config metadata.""" + num_nextn_predict_layers = _as_nonnegative_int(config_data.get("num_nextn_predict_layers")) + if num_nextn_predict_layers is not None: + return num_nextn_predict_layers + + model_config = getattr(model, "config", None) + if model_config is not None: + num_nextn_predict_layers = _as_nonnegative_int( + getattr(model_config, "num_nextn_predict_layers", None) + ) + if num_nextn_predict_layers is not None: + return num_nextn_predict_layers + + return None + + +def _get_rope_theta(config_data: dict[str, Any], model: Any) -> Any: + """Return rope_theta from exported config data or the in-memory model config.""" + rope_theta = config_data.get("rope_theta") + if rope_theta is not None: + return rope_theta + + model_config = getattr(model, "config", None) + if model_config is None: + return None + + return getattr(model_config, "rope_theta", None) + + +def _sanitize_llama3_rope_config(config_data: dict[str, Any], model: Any) -> None: + """Fill missing llama3 rope_theta in rope config metadata when available.""" + rope_theta = _get_rope_theta(config_data, model) + if rope_theta is None: + return + + for key in ("rope_parameters", "rope_scaling"): + rope_config = config_data.get(key) + if not isinstance(rope_config, dict): + continue + + rope_type = rope_config.get("rope_type", rope_config.get("type")) + if rope_type == "llama3" and "rope_theta" not in rope_config: + rope_config["rope_theta"] = rope_theta + + +def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) -> None: + """Sanitize exported Hugging Face config metadata for deployment runtimes. + + Fix conservative deployment-only config incompatibilities: + + * add missing llama3 ``rope_theta`` metadata when available; + * trim trailing MTP/next-token-prediction ``layer_types`` entries only when + the mismatch is exactly explained by next-token-prediction metadata. + """ + _sanitize_llama3_rope_config(config_data, model) + + num_hidden_layers = _as_nonnegative_int(config_data.get("num_hidden_layers")) + layer_types = config_data.get("layer_types") + if num_hidden_layers is None or not isinstance(layer_types, list): + return + + num_layer_types = len(layer_types) + if num_layer_types == num_hidden_layers: + return + + num_nextn_predict_layers = _get_num_nextn_predict_layers(config_data, model) + if ( + num_layer_types > num_hidden_layers + and num_nextn_predict_layers == num_layer_types - num_hidden_layers + ): + warnings.warn( + "Trimming config.layer_types from " + f"{num_layer_types} to {num_hidden_layers} entries so it matches " + "num_hidden_layers; the removed entries correspond to " + "num_nextn_predict_layers.", + stacklevel=2, + ) + config_data["layer_types"] = layer_types[:num_hidden_layers] + + +def _is_hf_hub_offline() -> bool: + return os.environ.get("HF_HUB_OFFLINE", "").strip().upper() in _HF_HUB_OFFLINE_TRUE_VALUES + + +def _copy_python_files(source_dir: Path, save_dir: Path) -> None: + for py_file in source_dir.glob("*.py"): + shutil.copy2(py_file, save_dir / py_file.name) + + +def copy_hf_ckpt_remote_code( + pretrained_model_path: str | os.PathLike, save_directory: str | os.PathLike +): + """Copy remote code from pretrained model to save directory. + + For models that keep configuration and modeling files as part of the checkpoint, + we need to copy them to the export directory for seamless integration with inference + frameworks. + + If ``pretrained_model_path`` is a local directory, Python files are copied directly. + If it's a HF Hub model ID (e.g. ``nvidia/NVIDIA-Nemotron-Nano-12B-v2``), the Hub + snapshot is resolved first and Python files are copied from that snapshot. When + ``HF_HUB_OFFLINE`` is set, the snapshot must already be available in the local + Hugging Face cache. + + Args: + pretrained_model_path: Local path to the pretrained model or HuggingFace Hub model ID. + save_directory: Path to the save directory. + """ + hf_checkpoint_path = Path(pretrained_model_path) + save_dir = Path(save_directory) + save_dir.mkdir(parents=True, exist_ok=True) + + if hf_checkpoint_path.is_dir(): + _copy_python_files(hf_checkpoint_path, save_dir) + else: + local_files_only = _is_hf_hub_offline() + try: + source_dir = Path( + snapshot_download( + repo_id=str(pretrained_model_path), + allow_patterns=["*.py"], + local_files_only=local_files_only, + ) + ) + except LocalEntryNotFoundError as exc: + if local_files_only: + raise RuntimeError( + f"Could not copy Python sidecar files for {pretrained_model_path!r} because " + "HF_HUB_OFFLINE is enabled and the files are not available in the local " + "Hugging Face cache. Populate the cache with the model's *.py files or pass " + "a local pretrained model directory." + ) from exc + raise + + _copy_python_files(source_dir, save_dir) + + +def load_multimodal_components( + pretrained_model_path: str | os.PathLike, + prefixes: tuple[str, ...] = ("multi_modal_projector", "vision_model"), +) -> dict[str, torch.Tensor]: + """Load multimodal components from safetensors file. + + Args: + pretrained_model_path: Directory or HuggingFace repo id of the pretrained model. + prefixes: Tensor key prefixes to select. Defaults to the LLaVA-style + ``multi_modal_projector`` / ``vision_model`` prefixes. Pass + ``("model.visual.",)`` for Qwen3-VL checkpoints. + + Returns: + A dictionary of multimodal components. + """ + hf_checkpoint_path = Path(pretrained_model_path) + if not hf_checkpoint_path.is_dir(): + # Also accept a repo id, which is what the example scripts pass to quantize.py. + # Fetched in two stages: the vision tower is a small fraction of a VLM checkpoint, so + # pulling every shard to keep a few would waste tens of GB. + local_files_only = _is_hf_hub_offline() + repo_id = str(pretrained_model_path) + try: + index_dir = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=["model.safetensors.index.json"], + local_files_only=local_files_only, + ) + ) + except (LocalEntryNotFoundError, OSError, ValueError) as exc: + raise ValueError( + f"Invalid pretrained model path: {pretrained_model_path}. It should be a " + "directory or an available HuggingFace repo id." + ) from exc + + index_file = index_dir / "model.safetensors.index.json" + if index_file.is_file(): + try: + weight_map = json.loads(index_file.read_text())["weight_map"] + except (json.JSONDecodeError, KeyError) as exc: + raise ValueError(f"Malformed safetensors index in {repo_id}.") from exc + wanted = sorted( + {shard for key, shard in weight_map.items() if key.startswith(prefixes)} + ) + else: + wanted = ["model.safetensors"] # unsharded checkpoint + # Kept separate from the resolution failure above: a hub outage or a full disk here is + # retryable, not a bad path. + hf_checkpoint_path = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=["model.safetensors.index.json", *wanted], + local_files_only=local_files_only, + ) + ) + + safetensors_file = Path(hf_checkpoint_path) / "model.safetensors" + safetensors_index_file = Path(hf_checkpoint_path) / "model.safetensors.index.json" + + multimodal_state_dict = {} + + if safetensors_file.is_file(): + print(f"Loading multimodal components from single file: {safetensors_file}") + with safe_open(safetensors_file, framework="pt") as f: + multimodal_keys = [ + key + for key in f.keys() # noqa: SIM118 + if key.startswith(prefixes) + ] + for key in tqdm(multimodal_keys, desc="Loading multimodal tensors"): + multimodal_state_dict[key] = f.get_tensor(key) + + elif safetensors_index_file.is_file(): + print(f"Loading multimodal components from sharded model: {hf_checkpoint_path}") + with open(safetensors_index_file) as f: + safetensors_index = json.load(f) + + all_shard_files = sorted( + { + shard + for key, shard in safetensors_index["weight_map"].items() + if key.startswith(prefixes) + } + ) + for shard_file in all_shard_files: + safetensors_filepath = Path(hf_checkpoint_path) / shard_file + with safe_open(safetensors_filepath, framework="pt") as f: + for key in f.keys(): # noqa: SIM118 + if key.startswith(prefixes): + multimodal_state_dict[key] = f.get_tensor(key) + + else: + print(f"Warning: No safetensors files found in {hf_checkpoint_path}") + + if not multimodal_state_dict: + raise ValueError( + f"No tensors under {prefixes} in {pretrained_model_path}; the vision tower would be " + "missing from the export. The checkpoint's prefixes have likely changed." + ) + + print(f"Successfully loaded {len(multimodal_state_dict)} multimodal tensors") + return multimodal_state_dict + + +def _matches_any_pattern(file_name: str, patterns: tuple[str, ...]) -> bool: + return any(fnmatch.fnmatchcase(file_name, pattern) for pattern in patterns) + + +# Standard HF weight-file names: ``model.safetensors`` or ``model-00001-of-00005.safetensors``. +_IS_MAIN_WEIGHT_SHARD = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors") + +# Off-index files that re-ship weights rather than add new ones. ``consolidated*.safetensors`` +# is Mistral's second full copy of the model (vLLM's mistral load-format looks for it BY NAME, +# so copying it into an export is not inert -- it can be served in place of the quantized +# weights). ``adapter_model.safetensors`` is a PEFT adapter, whose tensor names do not overlap +# the index, so only a name rule catches it. +_IS_WEIGHT_DUPLICATE = re.compile(r"(consolidated[^/]*|adapter_model)\.safetensors") + + +# --- What reaches the export without passing through quantization -------------------------- +# Two disjoint sets, distinguished by what the LOADER did with the file. That difference decides +# both how we find them and how we move them, so it is worth keeping straight: +# +# 1. UNPLACED weights. The loader opened the file and read the tensor, but the model had no +# parameter for it, so transformers reports it in ``unexpected_keys`` -- an MTP head the +# recipe did not quantize is the common case. Moved as TENSORS: located in whichever shard +# holds them and merged into the exporter's ``extra_state_dict``. +# Found by: ``read_unplaced_weights`` / ``carryable_unplaced_keys`` / ``locate_source_keys``. +# +# 2. OFF-INDEX sidecars. The index never names the file, so the loader never opened it and never +# had the chance to call anything unexpected -- GLM-4.7 keeps its MTP head in a standalone +# ``mtp.safetensors`` exactly this way. Moved as FILES: copied byte for byte, so no host +# memory is spent re-serialising tensors the export does not otherwise touch. +# Found by: ``off_index_safetensors_files`` / ``off_index_tensor_names``. +# +# The index is what separates them, and it is NOT an inventory of the checkpoint: a tensor missing +# from ``weight_map`` but sitting in a shard the index names for other tensors is set 1, not set 2. +# See ``record_unplaced_source_keys`` for the loading behaviour this rests on. +# +# Both sets must reach ``quantization_config.ignore``, or a deployment framework reads the +# top-level ``quant_algo`` and tries to load an original-precision weight as a quantized one +# (NVBug 5718750). ``_modelopt_carried_over_names`` is their union, recorded by the export once it +# knows what it actually wrote. + +# --- HF checkpoint layout: what counts as a file inside a checkpoint ------------------- +# A hub snapshot stores every entry as a symlink into a sibling ``blobs/`` directory, so +# "inside the checkpoint" has to mean the snapshot dir OR that blob root. Getting this wrong +# in either direction is costly: reject links and no hub checkpoint works, follow them blindly +# and a checkpoint can name any file on the host. +_MAX_CHECKPOINT_METADATA_BYTES = 128 * 1024 * 1024 + + +def _is_relative_to(path: Path, root: Path) -> bool: + return path == root or root in path.parents + + +def _snapshot_blob_root(source_root: Path) -> Path | None: + if source_root.parent.name != "snapshots": + return None + blob_root = source_root.parent.parent / "blobs" + return blob_root.resolve(strict=True) if blob_root.is_dir() else None + + +def _allowed_source_roots(src_dir: Path) -> list[Path]: + source_root = src_dir.resolve(strict=True) + allowed_roots = [source_root] + if blob_root := _snapshot_blob_root(source_root): + allowed_roots.append(blob_root) + return allowed_roots + + +def resolve_checkpoint_file( + src_dir: Path, + relative_path: str | Path, + *, + max_bytes: int | None = _MAX_CHECKPOINT_METADATA_BYTES, +) -> Path: + """Resolve a contained regular checkpoint file and optionally bound its size.""" + src = src_dir / relative_path + try: + resolved_src = src.resolve(strict=True) + except OSError as exc: + raise ValueError(f"checkpoint source is not a readable regular file: {src}") from exc + if not resolved_src.is_file(): + raise ValueError(f"checkpoint source must resolve to a regular file: {src}") + if not any(_is_relative_to(resolved_src, root) for root in _allowed_source_roots(src_dir)): + raise ValueError(f"checkpoint source is outside the checkpoint directory: {src}") + if max_bytes is not None and resolved_src.stat().st_size > max_bytes: + raise ValueError(f"checkpoint source exceeds the {max_bytes}-byte size limit: {src}") + return resolved_src + + +def off_index_safetensors_files(src: "str | os.PathLike") -> list[str]: + """Safetensors files in a checkpoint that model loading never opens. + + Transformers reads the shards named in ``model.safetensors.index.json`` -- or the single + ``model.safetensors`` when there is no index -- and nothing else. A checkpoint may ship more: + GLM-4.7 keeps its MTP head in a standalone ``mtp.safetensors``. Those tensors are never loaded, + never quantized, and never reported as ``unexpected_keys`` (the loader did not see them to call + them unexpected), so they are not "the unquantized source weights" the export must avoid + re-emitting -- they are untouched sidecars that happen to be in safetensors format. See + :func:`~modelopt.torch.utils.plugins.model_load_utils.record_unplaced_source_keys` for why + "never opened" is a property of the FILE rather than of the individual tensor: a tensor the + index omits from a file it DOES open is reported, and is carried rather than copied. + + Files named like a main weight shard are excluded whatever the index says. An index that is + empty, partial or malformed would otherwise make the source weights look off-index, and + copying those into an export would leave unquantized weights beside the quantized ones. + + Files that re-ship the indexed weights are excluded too, by name for the conventions we know + (``consolidated.safetensors``, ``adapter_model.safetensors``) and by tensor-name overlap for + the ones we do not. See :func:`_without_reshipped_weights`. + """ + d = Path(src) + if not d.is_dir(): + return [] + index_file = d / "model.safetensors.index.json" + indexed_tensors: set[str] = set() + if index_file.exists(): + with open(index_file) as f: + weight_map = json.load(f).get("weight_map", {}) + read_by_loader = set(weight_map.values()) + indexed_tensors = set(weight_map) + else: + read_by_loader = {"model.safetensors"} + single = d / "model.safetensors" + if single.exists(): + # Same reason the indexed branch fills this in: without the loaded tensor names, + # _without_reshipped_weights has nothing to compare against and silently becomes a + # no-op, leaving a second full copy under an unrecognised name to be copied verbatim. + with contextlib.suppress(Exception), safe_open(str(single), framework="pt") as f: + indexed_tensors = set(f.keys()) + + candidates = [ + f.name + for f in d.glob("*.safetensors") + if f.name not in read_by_loader + and not _IS_MAIN_WEIGHT_SHARD.fullmatch(f.name) + and not _IS_WEIGHT_DUPLICATE.fullmatch(f.name) + ] + return sorted(_without_reshipped_weights(d, candidates, indexed_tensors)) + + +def _without_reshipped_weights( + d: Path, candidates: list[str], indexed_tensors: set[str] +) -> list[str]: + """Drop candidates that re-ship weights the index already covers. + + The name rules above only catch conventions we know. A checkpoint free to invent its own + filename can still carry a second copy of the indexed weights, and copying that into an + export puts unquantized tensors beside the quantized ones. Overlapping tensor names are the + general signal: a genuine sidecar (an MTP head) holds names the index does not have, which + is exactly why the loader never placed them. + + Best-effort. Reads safetensors headers, never tensor data, and keeps any candidate whose + header cannot be read: refusing to copy a real sidecar because of an unreadable header would + silently drop weights from the export, which is the failure this whole path exists to avoid. + """ + if not candidates or not indexed_tensors: + return candidates + + kept = [] + for name in candidates: + try: + with safe_open(str(d / name), framework="pt") as f: + names = set(f.keys()) + except Exception: + kept.append(name) + continue + if names and names <= indexed_tensors: + warnings.warn( + f"Skipping {name}: it re-ships {len(names)} weight(s) the checkpoint index " + "already covers, so copying it would duplicate unquantized weights." + ) + continue + kept.append(name) + return kept + + +def indexed_weight_map(ckpt: "str | Path") -> dict[str, str]: + """``param name -> shard file`` as the checkpoint INDEX declares it. + + Named for what it returns rather than for the question callers want answered. For a sharded + checkpoint this is ``weight_map`` verbatim: it describes what the loader will look for, not + what the shards physically contain, and the two differ -- a tensor present in a shard but + absent from the index is invisible here. :func:`locate_source_keys` exists to cover that gap + and should be preferred by anything asking "where does this key actually live". Only the + single-file case is exhaustive, because there is no index for it to disagree with. + + Independent of the loader's dependencies, deliberately -- only stdlib + safetensors, not + transformers/accelerate, so it stays answerable in the partial-install environments where + those are absent. + + Returns ``{}``, not an exception, when neither an index nor a single-file checkpoint exists: + right for callers that treat "nothing recorded" as legitimate (e.g. :func:`locate_source_keys` + below). Callers for whom that is a genuine error (a missing checkpoint, not merely nothing + recorded) should check for the empty result and raise themselves. + + The indexed case, which is every sharded checkpoint, is a stdlib JSON read and needs nothing. + Only a single-file ``model.safetensors`` needs safetensors, and only to list its keys. + """ + index = Path(ckpt) / "model.safetensors.index.json" + if index.exists(): + with open(index) as f: + return json.load(f).get("weight_map", {}) + single = Path(ckpt) / "model.safetensors" + if single.exists(): + with safe_open(str(single), framework="pt") as f: + return dict.fromkeys(f.keys(), "model.safetensors") + return {} + + +def read_safetensors_subset( + ckpt_path: "str | Path", + weight_map: dict, + select: Callable[[str], bool], +) -> dict: + """Read tensors whose name satisfies ``select`` from safetensors files. + + Groups param names by file to avoid re-opening. Returns CPU tensors. + Uses ``safe_open`` so only the requested tensors' bytes are read. + + ``get_tensor`` returns a zero-copy view into the mmap'd file; the bytes are + not actually read from disk until first touched. We ``clone()`` here to force + the read eagerly, while this function runs (each rank reading its own layers + in parallel, for FSDP2 callers). Without it the read is deferred to a later + per-source broadcast, which is serialized across ranks and silently destroys + the read parallelism such callers exist to provide. + """ + by_file: dict[str, list[str]] = {} + for name, file in weight_map.items(): + if select(name): + by_file.setdefault(file, []).append(name) + + state: dict[str, torch.Tensor] = {} + for file, names in by_file.items(): + with safe_open(os.path.join(str(ckpt_path), file), framework="pt", device="cpu") as f: + for name in names: + state[name] = f.get_tensor(name).clone() + return state + + +def locate_source_keys(ckpt: "str | Path", keys: list[str]) -> dict[str, str]: + """Map each key to the safetensors file holding it, for keys the index may not list. + + ``model.safetensors.index.json`` is not a complete inventory of the checkpoint. Transformers + enumerates the CONTENTS of each shard it opens, so it reports a tensor the index omits -- an + MTP head stored inside a main shard is exactly that case, and it reaches + ``_modelopt_unplaced_source_keys`` like any other unplaced key. Resolving purely through + ``weight_map`` then finds no shard for it and drops it silently, which is the failure the + carry-over exists to prevent. + + The index answers for everything it lists, at no cost. Only the leftovers trigger a header + scan -- names, never tensor data -- and only when there are any, which is the rare case. + """ + # indexed_weight_map, not anything from model_load_utils: that module imports transformers + # and accelerate at module scope, and reaching for it here would make this answer "nothing to + # carry" wherever they are absent -- the partial-install environments -- silencing the + # --vllm_fakequant_export guard in exactly the case it exists for. Pinned by + # test_carryable_unplaced_keys_works_without_the_loader_dependencies, which caught this. + if not Path(ckpt).is_dir(): + # A checkpoint that is not there is a different failure from a key that is not in it, and + # the caller's handler already says the right thing about the first ("could not copy ... + # the checkpoint will be missing them"). Collapsing both into "in no safetensors file" + # would report a missing directory as a missing tensor. + raise ValueError(f"source checkpoint is not a directory: {ckpt}") + + weight_map = indexed_weight_map(ckpt) + located = {k: weight_map[k] for k in keys if k in weight_map} + remaining = {k for k in keys if k not in located} + if not remaining: + return located + + for shard in sorted(Path(ckpt).glob("*.safetensors")): + if not remaining: + break + try: + with safe_open(str(shard), framework="pt") as f: + found = remaining.intersection(f.keys()) + except Exception: + continue + located.update(dict.fromkeys(found, shard.name)) + remaining -= found + + if remaining: + warnings.warn( + f"{len(remaining)} checkpoint key(s) the loader reported are in no safetensors file " + f"of {ckpt} (e.g. {min(remaining)}); they cannot be carried into the export." + ) + return located + + +def copy_off_index_safetensors(src: "str | os.PathLike", dst: "str | os.PathLike") -> list[str]: + """Copy the safetensors files model loading never reads, verbatim. + + Copying beats reading them into a state dict and re-serialising: no host memory is spent on + tensors the export does not touch, the bytes and the file layout are preserved exactly, and a + consumer that finds them by filename (vLLM looks for the MTP sidecar) sees what it saw in the + source. See :func:`off_index_safetensors_files` for how "never reads" is decided. + """ + names = off_index_safetensors_files(src) + copied = [] + for name in names: + target = Path(dst) / name + if target.exists(): + continue + # copy2 follows symlinks, so a checkpoint shipping ``x.safetensors -> /etc/passwd`` would + # copy whatever that names into the export under an approved-looking name. The question + # is where the link LANDS, not whether it is a link: a Hugging Face snapshot stores every + # file as a symlink into ``../../blobs/``, so refusing links outright drops the + # sidecar of every hub-downloaded checkpoint -- the GLM-4.7 ``mtp.safetensors`` this path + # exists to carry included. + # + # resolve_checkpoint_file already draws that line and is tested against both shapes: it + # resolves strictly, demands a regular file, and demands the target sit under the + # checkpoint dir or its sibling ``blobs/``. max_bytes=None because its default bounds + # metadata, and these are weight files. + try: + source = resolve_checkpoint_file(Path(src), name, max_bytes=None) + except ValueError as exc: + warnings.warn(f"Skipping {name}: {exc}") + continue + shutil.copy2(source, target) + copied.append(name) + return copied + + +def copy_non_safetensor_files_from_ckpt( + src: str | os.PathLike, + dst: str | os.PathLike, + *, + exclude_files: Iterable[str] | None = None, + exclude_patterns: Iterable[str] | None = None, +) -> list[str]: + """Copy every non-safetensors file from a local HF checkpoint dir verbatim. + + Use as a baseline so tokenizer files, remote_code ``*.py``, README, LICENSE, etc. + are preserved from the source. Callers can exclude additional files or patterns when + copying after export-owned metadata has already been written. + + Args: + src: Source HF checkpoint directory. Must be a local path. + dst: Destination directory; created if missing. + exclude_files: Exact file names to skip. + exclude_patterns: Glob patterns for additional files to skip. + + Returns: + File names copied into ``dst``. + """ + if not os.path.isdir(src): + raise ValueError(f"Invalid source path: {src}. It should be a directory.") + exclude_files = set(exclude_files or ()) + exclude_patterns = tuple(exclude_patterns or ()) + copied_files = [] + os.makedirs(dst, exist_ok=True) + for entry in sorted(os.listdir(src)): + if entry in exclude_files or _matches_any_pattern(entry, exclude_patterns): + continue + sp = os.path.join(src, entry) + if not os.path.isfile(sp): + continue + if entry.endswith(".safetensors") or entry == "model.safetensors.index.json": + continue + try: + shutil.copy2(sp, dst) + except OSError as error: + warnings.warn(f"Failed to copy checkpoint sidecar {entry}: {error}") + continue + copied_files.append(entry) + return copied_files diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index cd66567fa9a..1efe02e5748 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -15,19 +15,18 @@ """HuggingFace-coupled FSDP2 model loading helpers.""" -import json import logging import os import re -from collections.abc import Callable from itertools import chain -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn -from accelerate import init_empty_weights from huggingface_hub import snapshot_download -from safetensors import safe_open + +if TYPE_CHECKING: + from collections.abc import Iterable from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict from torch.distributed.tensor import DTensor from transformers import AutoConfig, AutoModelForCausalLM @@ -44,56 +43,14 @@ fsdp2_wrap, is_initialized, ) +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + indexed_weight_map, + read_safetensors_subset, +) logger = logging.getLogger(__name__) -def read_safetensors_subset( - ckpt_path: str, - weight_map: dict, - select: Callable[[str], bool], -) -> dict: - """Read tensors whose name satisfies ``select`` from safetensors files. - - Groups param names by file to avoid re-opening. Returns CPU tensors. - Uses ``safe_open`` so only the requested tensors' bytes are read. - - ``get_tensor`` returns a zero-copy view into the mmap'd file; the bytes are - not actually read from disk until first touched. We ``clone()`` here to force - the read eagerly, while this function runs (each rank reading its own layers - in parallel). Without it the read is deferred to the later per-source - broadcast (``.to(device)``), which is serialized across ranks and silently - destroys the read parallelism this loader exists to provide. - """ - by_file: dict[str, list[str]] = {} - for name, file in weight_map.items(): - if select(name): - by_file.setdefault(file, []).append(name) - - state: dict[str, torch.Tensor] = {} - for file, names in by_file.items(): - with safe_open(os.path.join(ckpt_path, file), framework="pt", device="cpu") as f: - for name in names: - state[name] = f.get_tensor(name).clone() - return state - - -def weight_map_for(ckpt_path: str) -> dict[str, str]: - """Return the ``param_name → safetensors_file`` map for a local checkpoint directory.""" - index_path = os.path.join(ckpt_path, "model.safetensors.index.json") - single_file = os.path.join(ckpt_path, "model.safetensors") - if os.path.exists(index_path): - with open(index_path) as f: - return json.load(f)["weight_map"] - if os.path.exists(single_file): - with safe_open(single_file, framework="pt", device="cpu") as f: - return dict.fromkeys(f.keys(), "model.safetensors") - raise RuntimeError( - f"No safetensors checkpoint at {ckpt_path} " - "(expected model.safetensors or model.safetensors.index.json)." - ) - - def _resolve_checkpoint_dir(ckpt_path: str, rank: int) -> str: """Local dir for ``ckpt_path``; resolves an HF Hub ID (rank 0 downloads, others wait).""" if os.path.isdir(ckpt_path): @@ -150,7 +107,7 @@ def _conversion_plan(model: nn.Module) -> dict | None: "legacy_renames": legacy_renames, "renamings": renamings, "converters": converters, - "prefix": model.base_model_prefix, + "prefix": getattr(model, "base_model_prefix", ""), "meta_state_dict": model.state_dict(), } @@ -216,6 +173,8 @@ def build_meta_causal_lm( # Honor the override even when the caller passed in a pre-fetched config. hf_config._attn_implementation = attn_implementation dtype = getattr(hf_config, "torch_dtype", None) or torch.bfloat16 + from accelerate import init_empty_weights # only real callers of this function need it + with init_empty_weights(include_buffers=False): model = AutoModelForCausalLM.from_config( hf_config, torch_dtype=dtype, trust_remote_code=trust_remote_code @@ -297,20 +256,23 @@ def _broadcast_load_group( def _group_sources_by_layer( weight_map: dict, plan: dict | None, model_param_names: set[str], layer_prefixes: list[str] -) -> tuple[dict[int, list[str]], list[str], int]: +) -> tuple[dict[int, list[str]], list[str], list[str]]: """Bucket checkpoint keys by the decoder layer their converted target lives in. - Returns ``(layer_sources, non_layer_sources, skipped)``: ``layer_sources[i]`` holds the keys + Returns ``(layer_sources, non_layer_sources, unplaced)``: ``layer_sources[i]`` holds the keys targeting decoder layer ``i``, ``non_layer_sources`` holds root (embed/lm_head/norm) keys, and - ``skipped`` counts keys whose target isn't in the model (aux weights, e.g. an MTP head). + ``unplaced`` NAMES the keys whose target isn't in the model -- weights the built model has no + home for (an MTP head, an auxiliary tower). The names are kept, not just counted, so the export + can copy them through: PTQ never touches them, but the exported checkpoint is still expected to + contain them. """ layer_sources: dict[int, list[str]] = {i: [] for i in range(len(layer_prefixes))} non_layer_sources: list[str] = [] - skipped = 0 + unplaced: list[str] = [] for ckpt_key in weight_map: target = _resolve_target(plan, ckpt_key)[0] if plan else ckpt_key if target not in model_param_names: - skipped += 1 + unplaced.append(ckpt_key) continue for i, prefix in enumerate(layer_prefixes): if target.startswith(prefix): @@ -318,7 +280,76 @@ def _group_sources_by_layer( break else: non_layer_sources.append(ckpt_key) - return layer_sources, non_layer_sources, skipped + return layer_sources, non_layer_sources, unplaced + + +def record_unplaced_source_keys( + model: nn.Module, ckpt_path: str, unexpected_keys: "Iterable[str] | None" +) -> list[str]: + """Record the checkpoint keys the loader could not place, for the export to carry over. + + ``unexpected_keys`` is what ``from_pretrained(..., output_loading_info=True)`` reports: keys + found in the checkpoint but not expected by the model's architecture. That is the loader's own + accounting, produced while loading, so it already reflects any on-the-fly name conversion -- + unlike re-deriving the set afterwards, which has to replay the conversion plan to avoid + mistaking a renamed key for an unplaced one. + + How Transformers decides what it has seen + ----------------------------------------- + The index (``model.safetensors.index.json``) selects which FILES the loader opens, not which + TENSORS it sees. Within a file it opens, it enumerates every tensor present and reports the + ones the architecture does not expect. Two consequences, both load-bearing here: + + * A tensor missing from ``weight_map`` but physically present in a shard the index names for + OTHER tensors is still reported. An MTP head stored inside a main shard is exactly this + shape -- when MTP is not quantized the model never declares it, so it arrives here like any + other unplaced key. The corollary is the one that is easy to get wrong: the index is not an + inventory of the checkpoint, so looking such a key up in ``weight_map`` to find its file + returns nothing. :func:`~modelopt.torch.export.unified_export_hf._locate_source_keys` falls + back to scanning shard headers for precisely this reason; resolving through ``weight_map`` + alone used to drop these tensors from the export silently. + * A tensor in a file the index never names is NOT reported -- the loader never opened it, so + it had no opportunity to call anything unexpected. Those are handled by + :func:`~modelopt.torch.utils.plugins.hf_checkpoint_utils.copy_off_index_safetensors`, which + copies the file whole rather than paying host memory to re-serialise it. + + This is observed behaviour, established by experiment against transformers 5.3.0 (a shard + holding one tensor the index omitted reported it; a file the index never named reported + nothing), not a published contract. Nothing here depends on it holding: a key that the loader does report is + located by header scan whether or not the index lists it, and a file the loader ignores is + copied verbatim regardless. + + Prefer this over :func:`unplaced_source_keys` whenever the loading info is available. + """ + keys = sorted(unexpected_keys or []) + model._modelopt_unplaced_source_keys = keys + model._modelopt_source_checkpoint = str(ckpt_path) + return keys + + +def unplaced_source_keys(model: nn.Module, ckpt_path: str) -> list[str]: + """Checkpoint keys the built model has no parameter for. + + Fallback for loaders that do not surface their own accounting. Prefer + :func:`record_unplaced_source_keys` with ``from_pretrained(..., output_loading_info=True)``, + whose ``unexpected_keys`` is the same set computed by the loader itself. + + Architecture-agnostic by construction -- it asks whether a target parameter exists, not whether + the key looks like an MTP head or an auxiliary tower. + """ + weight_map = indexed_weight_map(ckpt_path) + if not weight_map: + raise RuntimeError( + f"No safetensors checkpoint at {ckpt_path} " + "(expected model.safetensors or model.safetensors.index.json)." + ) + plan = _conversion_plan(model) + model_param_names = {n for n, _ in chain(model.named_parameters(), model.named_buffers())} + return [ + ckpt_key + for ckpt_key in weight_map + if (_resolve_target(plan, ckpt_key)[0] if plan else ckpt_key) not in model_param_names + ] def parallel_load_and_prepare_fsdp2( @@ -350,7 +381,12 @@ def parallel_load_and_prepare_fsdp2( pass ``None`` to broadcast all of a source's layers at once). """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) - weight_map = weight_map_for(resolved_path) + weight_map = indexed_weight_map(resolved_path) + if not weight_map: + raise RuntimeError( + f"No safetensors checkpoint at {resolved_path} " + "(expected model.safetensors or model.safetensors.index.json)." + ) model = build_meta_causal_lm(resolved_path, trust_remote_code, attn_implementation, hf_config) @@ -367,13 +403,20 @@ def parallel_load_and_prepare_fsdp2( model_param_names = {n for n, _ in chain(model.named_parameters(), model.named_buffers())} # Bucket each checkpoint key by its target's decoder layer (root params go to non_layer_sources). - layer_sources, non_layer_sources, skipped = _group_sources_by_layer( + layer_sources, non_layer_sources, unplaced = _group_sources_by_layer( weight_map, plan, model_param_names, layer_prefixes ) - if skipped: + if unplaced: logger.debug( - "skipping %d checkpoint keys not present in the model (e.g. MTP head)", skipped + "%d checkpoint keys have no parameter in the built model (e.g. an MTP head); " + "recorded for the export to copy through verbatim", + len(unplaced), ) + # Recorded on the model so export can read them back from the source without being told where + # it came from. Not a state dict: holding these tensors from load to export would waste the + # memory this loader exists to save. + model._modelopt_unplaced_source_keys = unplaced + model._modelopt_source_checkpoint = resolved_path _materialize_meta_model(model, torch.device("cpu") if cpu_offload else device) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index 23e67821371..eae7c2ad38c 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -30,6 +30,16 @@ def get_free_port(): return port +def _init_backend(backend): + """Mirror ``modelopt.torch.utils.distributed.setup``, which pairs NCCL with a CPU backend. + + Bare ``"nccl"`` registers no backend for CPU tensors, so a collective on a CPU-resident + shard -- gathering an FSDP2 ``cpu_offload`` parameter, say -- fails with "No backend type + associated with device type cpu" in tests but not in production. + """ + return "cpu:gloo,cuda:nccl" if backend == "nccl" else backend + + def init_process(rank, size, job=None, backend="gloo", port=None): """Initialize the distributed environment.""" @@ -45,7 +55,7 @@ def init_process(rank, size, job=None, backend="gloo", port=None): # We need to use a different port for each tests to avoid conflicts os.environ["MASTER_PORT"] = port - dist.init_process_group(backend, rank=rank, world_size=size) + dist.init_process_group(_init_backend(backend), rank=rank, world_size=size) if backend == "nccl" and torch.cuda.is_available(): torch.cuda.set_device(rank) torch.manual_seed(1234) @@ -140,7 +150,7 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard os.environ["LOCAL_RANK"] = str(rank) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) - dist.init_process_group(backend, rank=rank, world_size=world_size) + dist.init_process_group(_init_backend(backend), rank=rank, world_size=world_size) if backend == "nccl" and torch.cuda.is_available(): torch.cuda.set_device(rank) torch.manual_seed(1234) diff --git a/tests/examples/hf_ptq/test_carry_over_layouts.py b/tests/examples/hf_ptq/test_carry_over_layouts.py new file mode 100644 index 00000000000..d482deeabc1 --- /dev/null +++ b/tests/examples/hf_ptq/test_carry_over_layouts.py @@ -0,0 +1,174 @@ +# 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. + +"""Carry-over must catch every on-disk layout an unplaced weight can arrive in. + +These drive a real ``from_pretrained``, not a stub, because the thing under test is what the +Transformers loader does and does not report. The three MTP storage conventions come from the +support matrix the previous name-matching implementation carried: + + inlined GLM-5.1, DeepSeek-V3 ``model.layers.{N}.*`` past ``num_hidden_layers`` + standalone GLM-4.7 a ``mtp.safetensors`` the index does not reference + indexed shard Qwen3-Next an ``mtp.*`` tail shard listed in the index + +The standalone case is handled differently from the other two, and deliberately so. The loader +only opens shards named in the index, so tensors in an extra file produce no ``unexpected_keys`` +at all -- but they are also untouched by quantization and absent from the export, which makes them +sidecars rather than carry-over. They are copied verbatim, which costs no host memory and keeps +the bytes and filename a consumer looks for. The other two conventions land inside shards the +loader does read, so those keys ride the state dict. + +CPU-only and small -- the mechanism is about bookkeeping during load, so a GPU adds nothing. +""" + +import json + +import pytest +import torch +from _test_utils.examples.hf_ptq_example_utils import example_utils +from safetensors.torch import load_file, save_file +from transformers import AutoModelForCausalLM, LlamaConfig + +from modelopt.torch.utils.plugins.hf_checkpoint_utils import copy_off_index_safetensors + +NUM_HIDDEN_LAYERS = 2 + + +def _build_checkpoint(d): + """A tiny real model saved to disk, so the loader has something genuine to place.""" + cfg = LlamaConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=NUM_HIDDEN_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + AutoModelForCausalLM.from_config(cfg).save_pretrained(d, safe_serialization=True) + + +def _add_to_main_shard(d, extra): + f = d / "model.safetensors" + tensors = load_file(str(f)) + tensors.update(extra) + save_file(tensors, str(f), metadata={"format": "pt"}) + + +def _recorded(d): + model = example_utils._from_pretrained_recording(AutoModelForCausalLM, str(d)) + return model._modelopt_unplaced_source_keys + + +@pytest.fixture +def ckpt(tmp_path): + _build_checkpoint(tmp_path) + return tmp_path + + +def test_inlined_layer_past_num_hidden_layers(ckpt): + """GLM-5.1 / DeepSeek-V3: the head is decoder layer N, which the model never builds.""" + keys = [ + f"model.layers.{NUM_HIDDEN_LAYERS}.input_layernorm.weight", + f"model.layers.{NUM_HIDDEN_LAYERS}.mlp.gate_proj.weight", + ] + _add_to_main_shard(ckpt, {keys[0]: torch.zeros(32), keys[1]: torch.zeros(64, 32)}) + + assert set(_recorded(ckpt)) == set(keys) + + +def test_standalone_off_index_file_is_copied_not_carried(ckpt, tmp_path): + """GLM-4.7: a separate mtp.safetensors. The loader never opens it, so it contributes no + unexpected_keys -- it is preserved by copying the file, not by routing its tensors through + the state dict.""" + payload = {"mtp.fc.weight": torch.zeros(32, 32), "mtp.layers.0.enorm.weight": torch.zeros(32)} + save_file(payload, str(ckpt / "mtp.safetensors"), metadata={"format": "pt"}) + + assert _recorded(ckpt) == [], "the loader never saw it, so it cannot report it" + + export = tmp_path / "export" + export.mkdir() + assert copy_off_index_safetensors(ckpt, export) == ["mtp.safetensors"] + assert load_file(str(export / "mtp.safetensors")).keys() == payload.keys() + + +def test_indexed_shards_are_not_copied(ckpt, tmp_path): + """The guard on the above: shards the loader does read are the source weights, and copying + them into the export would sit alongside the quantized ones.""" + export = tmp_path / "export" + export.mkdir() + assert copy_off_index_safetensors(ckpt, export) == [] + + +def test_indexed_tail_shard(ckpt): + """Qwen3-Next: an mtp.* shard the index does reference, so the loader reads and rejects it.""" + main = ckpt / "model.safetensors" + base = load_file(str(main)) + first, tail = "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors" + mtp = {"mtp.fc.weight": torch.zeros(32, 32), "mtp.layers.0.enorm.weight": torch.zeros(32)} + save_file(base, str(ckpt / first), metadata={"format": "pt"}) + save_file(mtp, str(ckpt / tail), metadata={"format": "pt"}) + main.unlink() + weight_map = {**dict.fromkeys(base, first), **dict.fromkeys(mtp, tail)} + (ckpt / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": weight_map}) + ) + + assert set(_recorded(ckpt)) == set(mtp) + + +def test_auxiliary_tower_is_carried_too(ckpt): + """Nothing is keyed to MTP: any tensor the model has no home for is carried.""" + _add_to_main_shard(ckpt, {"aux_tower.blocks.0.weight": torch.zeros(8, 8)}) + + assert set(_recorded(ckpt)) == {"aux_tower.blocks.0.weight"} + + +def test_two_layouts_at_once(ckpt, tmp_path): + """An inlined head and an off-index sidecar together: each must be picked up by its own + mechanism, and neither by both -- a tensor carried *and* copied would be exported twice.""" + inlined = f"model.layers.{NUM_HIDDEN_LAYERS}.input_layernorm.weight" + _add_to_main_shard(ckpt, {inlined: torch.zeros(32)}) + save_file( + {"mtp.fc.weight": torch.zeros(32, 32)}, + str(ckpt / "mtp.safetensors"), + metadata={"format": "pt"}, + ) + + export = tmp_path / "export" + export.mkdir() + recorded = set(_recorded(ckpt)) + copied = copy_off_index_safetensors(ckpt, export) + + assert recorded == {inlined} + assert copied == ["mtp.safetensors"] + assert "mtp.fc.weight" not in recorded, "copied files must not also ride the state dict" + + +def test_clean_checkpoint_records_nothing(ckpt): + """No stray tensors -> an empty answer, distinct from never having asked.""" + model = example_utils._from_pretrained_recording(AutoModelForCausalLM, str(ckpt)) + + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == str(ckpt) + + +def test_recorded_keys_are_sorted_and_deduplicated(ckpt): + """The two sources can overlap; the export wants a stable, duplicate-free list.""" + _add_to_main_shard( + ckpt, {"zzz_orphan.weight": torch.zeros(4), "aaa_orphan.weight": torch.zeros(4)} + ) + + keys = _recorded(ckpt) + assert keys == sorted(keys) and len(keys) == len(set(keys)) diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index e532af09fed..8651183cb83 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -12,13 +12,14 @@ # 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. -"""End-to-end unit tests for ``examples/hf_ptq/example_utils.load_mtp_weights``. +"""Unit tests for ``examples/hf_ptq/example_utils`` helpers. -One test per supported on-disk MTP convention (inlined-orphaned, inlined-in-state-dict, -separate-file-standalone, separate-file-indexed) plus a negative case. +The per-MTP-convention tests are gone with ``load_mtp_weights``: weights the loader could not +place are now identified from Transformers' own ``unexpected_keys`` rather than by recognising +storage layouts, so there is no convention matrix left to enumerate. What remains covers the +recording path, checkpoint-path resolution, and the sidecar copy. """ -import json from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import patch @@ -29,22 +30,6 @@ from safetensors.torch import save_file -class _FakeModel: - """Stub exposing only the surface ``load_mtp_weights`` touches.""" - - def __init__(self, config, state_dict_keys): - self.config = config - self._sd = {k: torch.zeros(1) for k in state_dict_keys} - self.loaded = {} - - def state_dict(self): - return dict(self._sd) - - def load_state_dict(self, state_dict, strict=True): - self.loaded.update(state_dict) - self._sd.update(state_dict) - - def _write_safetensors(path, tensors): save_file(tensors, str(path), metadata={"format": "pt"}) @@ -138,121 +123,6 @@ def fake_from_pretrained(*args, **kwargs): ) -def test_load_mtp_weights_inlined_orphaned(tmp_path): - # GLM-5.1: HF builds only num_hidden decoders → MTP keys orphaned. - main_keys = ["model.embed_tokens.weight", "model.layers.0.x.weight"] - mtp_keys = ["model.layers.4.eh_proj.weight", "model.layers.4.enorm.weight"] - _write_safetensors( - tmp_path / "model.safetensors", - {k: torch.zeros(2, 2) for k in main_keys + mtp_keys}, - ) - - cfg = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=1) - model = _FakeModel(cfg, state_dict_keys=main_keys) - prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path)) - - assert prefixes == ["model.layers.4"] - assert set(orphans) == set(mtp_keys) - assert model.loaded == {} # nothing matched the (MTP-less) state_dict - - -def test_load_mtp_weights_inlined_in_state_dict(tmp_path): - # DeepSeek-V3 via trust_remote_code: MTP slots exist → keys loaded, no orphans. - main_keys = ["model.embed_tokens.weight"] - mtp_keys = ["model.layers.4.eh_proj.weight", "model.layers.4.enorm.weight"] - _write_safetensors( - tmp_path / "model.safetensors", - {k: torch.ones(2, 2) for k in main_keys + mtp_keys}, - ) - - cfg = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=1) - model = _FakeModel(cfg, state_dict_keys=main_keys + mtp_keys) - prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path)) - - assert prefixes == ["model.layers.4"] - assert orphans == {} - assert set(model.loaded) == set(mtp_keys) - - -def test_load_mtp_weights_separate_standalone_file(tmp_path): - # GLM-4.7: standalone mtp.safetensors with no shard index. - _write_safetensors( - tmp_path / "model.safetensors", {"model.embed_tokens.weight": torch.zeros(2, 2)} - ) - _write_safetensors( - tmp_path / "mtp.safetensors", - { - "mtp.fc.weight": torch.zeros(2, 2), - "mtp.layers.0.q_proj.weight": torch.zeros(2, 2), - }, - ) - - cfg = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=0) - model = _FakeModel(cfg, state_dict_keys=["model.embed_tokens.weight"]) - prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path)) - - assert set(prefixes) == {"mtp", "mtp.layers.0"} - assert set(orphans) == {"mtp.fc.weight", "mtp.layers.0.q_proj.weight"} - - -def test_load_mtp_weights_separate_indexed_shard(tmp_path): - # Qwen3-Next: mtp.* keys in a dedicated indexed tail shard (filename has no "mtp"). - main_shard = "model-00001-of-00002.safetensors" - mtp_shard = "model-00002-of-00002.safetensors" - _write_safetensors(tmp_path / main_shard, {"model.embed_tokens.weight": torch.zeros(2, 2)}) - mtp_tensors = { - "mtp.fc.weight": torch.zeros(2, 2), - "mtp.norm.weight": torch.zeros(2), - "mtp.layers.0.input_layernorm.weight": torch.zeros(2), - "mtp.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), - } - _write_safetensors(tmp_path / mtp_shard, mtp_tensors) - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps( - { - "weight_map": { - "model.embed_tokens.weight": main_shard, - **dict.fromkeys(mtp_tensors, mtp_shard), - } - } - ) - ) - - cfg = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=0) - model = _FakeModel(cfg, state_dict_keys=["model.embed_tokens.weight"]) - prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path)) - - assert set(prefixes) == {"mtp", "mtp.layers.0"} - assert set(orphans) == set(mtp_tensors) - - -def test_keys_to_prefixes_drops_model_top_level(): - # nvbug 6108133: inlined keys like "model.layers.92.X" must NOT emit "model" - # as a top-level prefix (would become "model*" excluding the whole backbone). - out = example_utils._keys_to_prefixes( - ["model.layers.92.eh_proj.weight", "mtp.fc.weight", "mtp.layers.0.q_proj.weight"] - ) - assert "model" not in out - assert out == {"mtp", "mtp.layers.0", "model.layers.92"} - - -def test_load_mtp_weights_no_mtp_returns_empty(tmp_path): - # Also pins the ``num_nextn_predict_layers=None`` regression: some configs - # set the field explicitly to None, which must not crash ``int(None)``. - _write_safetensors( - tmp_path / "model.safetensors", - { - "model.embed_tokens.weight": torch.zeros(2, 2), - "model.layers.0.x.weight": torch.zeros(2, 2), - }, - ) - cfg = SimpleNamespace(num_hidden_layers=4, num_nextn_predict_layers=None) - model = _FakeModel(cfg, state_dict_keys=["model.embed_tokens.weight"]) - prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path)) - assert prefixes == [] - assert orphans == {} - - # ---------- get_original_hf_quant_method ------------------------------------- # get_model uses this to detect native MXFP4 checkpoints (e.g. openai/gpt-oss-*) and load # them dequantized to BF16 GptOssExperts (so ModelOpt can quantize/export the experts). @@ -364,21 +234,23 @@ def from_config(config, **kwargs): return FakeModel() @staticmethod - def from_pretrained(*args, **kwargs): + def from_pretrained(*args, output_loading_info=False, **kwargs): calls["from_pretrained"] = kwargs assert "dtype" not in kwargs assert kwargs["torch_dtype"] is torch.float16 - return FakeModel() + m = FakeModel() + return (m, {"unexpected_keys": []}) if output_loading_info else m class FakeLlamaForCausalLM(FakeAutoModelForCausalLM): _from_config = FakeAutoModelForCausalLM.from_config @staticmethod - def from_pretrained(*args, **kwargs): + def from_pretrained(*args, output_loading_info=False, **kwargs): calls["from_pretrained"] = kwargs assert kwargs["dtype"] == "auto" assert "torch_dtype" not in kwargs - return FakeModel() + m = FakeModel() + return (m, {"unexpected_keys": []}) if output_loading_info else m monkeypatch.setattr( example_utils.AutoConfig, @@ -445,9 +317,10 @@ def _from_config(config, **kwargs): return FakeModel() @staticmethod - def from_pretrained(*args, **kwargs): + def from_pretrained(*args, output_loading_info=False, **kwargs): calls["from_pretrained"] = kwargs - return FakeModel() + m = FakeModel() + return (m, {"unexpected_keys": []}) if output_loading_info else m monkeypatch.setattr( example_utils.AutoConfig, "from_pretrained", lambda *args, **kwargs: hf_config @@ -540,9 +413,10 @@ def from_config(config, **kwargs): _from_config = from_config @staticmethod - def from_pretrained(*args, **kwargs): + def from_pretrained(*args, output_loading_info=False, **kwargs): used["path"] = tag - return FakeModel() + m = FakeModel() + return (m, {"unexpected_keys": []}) if output_loading_info else m return Fake @@ -691,3 +565,60 @@ def __init__(self, algorithm): def test_recipe_layerwise_blocks(recipe, expected): """Both recipe shapes normalize to dicts, so callers need no shape-aware access.""" assert example_utils.recipe_layerwise_blocks(recipe) == expected + + +# --- carry-over of weights the loader could not place ------------------------------------------- + + +class _StubAuto: + """Minimal stand-in for an ``AutoModelFor*`` class.""" + + def __init__(self, unexpected, model=None): + self._unexpected = unexpected + self._model = model if model is not None else SimpleNamespace() + self.saw_output_loading_info = None + self.saw_kwargs = None + + def from_pretrained(self, ckpt_path, output_loading_info=False, **kwargs): + self.saw_output_loading_info = output_loading_info + self.saw_kwargs = kwargs + return self._model, { + "missing_keys": [], + "unexpected_keys": list(self._unexpected), + "mismatched_keys": [], + "error_msgs": [], + } + + +def test_from_pretrained_recording_asks_the_loader_for_its_accounting(tmp_path): + """The point of this path: take unexpected_keys from the loader rather than re-deriving it.""" + auto = _StubAuto(["mtp.layers.0.eh_proj.weight", "mtp.fc.weight"]) + model = example_utils._from_pretrained_recording(auto, str(tmp_path), device_map="cpu") + + assert auto.saw_output_loading_info is True, "must request the loading info" + assert auto.saw_kwargs == {"device_map": "cpu"}, "caller kwargs must pass through untouched" + assert model._modelopt_unplaced_source_keys == [ + "mtp.fc.weight", + "mtp.layers.0.eh_proj.weight", + ], "recorded sorted, so export order is stable" + assert model._modelopt_source_checkpoint == str(tmp_path) + + +def test_from_pretrained_recording_is_quiet_when_everything_was_placed(tmp_path): + """Record an empty answer too, so the exporter can tell 'nothing to carry' from 'never asked'.""" + auto = _StubAuto([]) + model = example_utils._from_pretrained_recording(auto, str(tmp_path)) + + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == str(tmp_path) + + +def test_recording_is_architecture_agnostic(tmp_path): + """Nothing keys off the string 'mtp': an auxiliary tower is carried by the same rule.""" + auto = _StubAuto(["aux_tower.blocks.0.weight", "something_else.weight"]) + model = example_utils._from_pretrained_recording(auto, str(tmp_path)) + + assert model._modelopt_unplaced_source_keys == [ + "aux_tower.blocks.0.weight", + "something_else.weight", + ] diff --git a/tests/gpu/torch/export/test_export_carry_over.py b/tests/gpu/torch/export/test_export_carry_over.py new file mode 100644 index 00000000000..fa528f4cdf9 --- /dev/null +++ b/tests/gpu/torch/export/test_export_carry_over.py @@ -0,0 +1,266 @@ +# 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. +"""Carry-over of source weights the model could not place, end to end under FSDP2. + +A checkpoint can hold parameters the built model has no home for -- an MTP head, or a vision +tower that a language-model-only PTQ never touches. Quantization never sees those, so without an +explicit carry-over they are silently absent from the export: no error, plausible file sizes, a +checkpoint simply missing a component. + +Asserted through the public ``export_hf_checkpoint`` API, so these hold for whichever FSDP2 write +path it selects: + +1. every source key the model did not place is present in the export, with its source values; +2. a VLM whose language model alone was quantized keeps its vision tower -- unquantized, at full + shape, and carrying no stray scales. +""" + +import json +from functools import partial +from pathlib import Path + +import pytest +import torch +from _test_utils.torch.transformers_models import create_tiny_qwen3_dir, create_tiny_qwen3vl_dir +from safetensors.torch import load_file +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint +from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes +from modelopt.torch.utils.distributed import fsdp2_wrap, is_fsdp2_model + +# Small enough to force the tiny model across several shards, so the multi-file layout -- and the +# weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. +# Not smaller: each extra shard is another write + consolidation round trip on a toy model. +MAX_SHARD_SIZE = "512KB" + +# The default tiny Qwen3 is too degenerate to shard: head_dim = hidden_size / num_heads = 32/16 = 2, +# so q_norm/k_norm are 2-element tensors. Split over a 4-rank world most ranks get an empty chunk and +# the write planner rejects the coverage outright ("invalid fill tensor-volume"), taking the whole +# worker pool down with it. Size every FSDP2-sharded axis to stay comfortably divisible instead -- +# still a toy model, but one that survives a multi-rank split. Last dims are multiples of 16 so NVFP4 +# block quantization is also well defined. +TINY_KWARGS = { + "hidden_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 32, + "intermediate_size": 128, + "max_position_embeddings": 64, + "num_hidden_layers": 2, +} +TINY_MOE_KWARGS = {"moe_intermediate_size": 128, "num_experts": 8, "num_experts_per_tok": 2} + +# Suffixes the exporter ADDS to a checkpoint; every other exported key must correspond to a +# parameter of the source model. +_SCALE_SUFFIXES = ( + "weight_scale", + "weight_scale_2", + "weight_scale_inv", + "input_scale", + "pre_quant_scale", + "k_scale", + "v_scale", +) + +# Exported dtypes that mean "this weight was quantized" (NVFP4 packs two fp4 per uint8). +_QUANTIZED_DTYPES = {"F8_E4M3", "U8", "I8"} + + +def _safetensors_meta(directory: Path) -> dict[str, tuple[str, tuple[int, ...]]]: + """``{tensor_name: (dtype, shape)}`` across every ``*.safetensors`` in ``directory``. + + Reads the safetensors header directly (8-byte little-endian length, then JSON) so the check + depends only on what is on disk -- no loader, no dequantization, no GPU. + """ + out: dict[str, tuple[str, tuple[int, ...]]] = {} + for path in sorted(directory.glob("*.safetensors")): + with open(path, "rb") as fh: + header = json.loads(fh.read(int.from_bytes(fh.read(8), "little"))) + for name, spec in header.items(): + if name != "__metadata__": + out[name] = (spec["dtype"], tuple(spec["shape"])) + return out + + +def _is_scale(key: str) -> bool: + return key.endswith(_SCALE_SUFFIXES) + + +def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwargs): + """Load the tiny model on every rank, FSDP2-shard it, PTQ it, and export.""" + + with patch_fsdp_mp_dtypes(): + model = AutoModelForCausalLM.from_pretrained(src_dir, dtype=torch.bfloat16).to("cuda") + model.eval() + + fsdp2_wrap(model) + # The FSDP2 export path is selected by exactly this predicate, so assert it here: without it + # the export would silently fall back to the rank-0 gather and the test would pass while + # never touching the code under test. + assert is_fsdp2_model(model), "fsdp2_wrap did not shard the model" + assert torch.distributed.is_initialized() + torch.distributed.barrier() + + input_ids = torch.randint(0, model.config.vocab_size, (2, 8), device="cuda") + mtq.quantize(model, quant_cfg, lambda m: m(input_ids)) + torch.distributed.barrier() + + export_hf_checkpoint( + model, export_dir=export_dir, max_shard_size=MAX_SHARD_SIZE, **export_kwargs + ) + torch.distributed.barrier() + + +@pytest.mark.timeout(600) +def test_fsdp2_distributed_export_carries_unplaced_weights_from_provenance(dist_workers, tmp_path): + """Unplaced source weights are carried over even when nothing was recorded at load time. + + ``parallel_load_and_prepare_fsdp2`` records the checkpoint keys it could not place, but a plain + ``from_pretrained`` -- which is how most callers build the model, and what this suite uses -- + drops unexpected keys silently and records nothing. The export then has to ask the question + itself, from the model's own provenance (``config._name_or_path``), or the checkpoint comes out + missing weights that PTQ never had the chance to touch. + + Same fixture shape as the loader-side test: a checkpoint holding one more layer than the config + admits, so the last layer is present on disk with nowhere to load -- what an inlined MTP tail + does, without needing an MTP-capable architecture here. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + src_dir = Path( + create_tiny_qwen3_dir( + tmp_path, with_tokenizer=True, **{**TINY_KWARGS, "num_hidden_layers": 3} + ) + ) + cfg_path = src_dir / "config.json" + cfg = json.loads(cfg_path.read_text()) + orphan_idx = cfg["num_hidden_layers"] - 1 + cfg["num_hidden_layers"] = orphan_idx + # Qwen3 validates that layer_types matches num_hidden_layers, so trim it alongside; leaving it + # at the original length fails config construction before the model is ever built. + if isinstance(cfg.get("layer_types"), list): + cfg["layer_types"] = cfg["layer_types"][:orphan_idx] + cfg_path.write_text(json.dumps(cfg)) + orphan_prefix = f"model.layers.{orphan_idx}." + + src_meta = _safetensors_meta(src_dir) + orphans = sorted(k for k in src_meta if k.startswith(orphan_prefix)) + assert orphans, "fixture produced no orphaned weights, so this test would prove nothing" + + export_dir = tmp_path / "export_carry_provenance" + dist_workers.run( + partial( + _ptq_and_export, + src_dir=src_dir, + export_dir=export_dir, + quant_cfg=mtq.FP8_DEFAULT_CFG, + ) + ) + + exported = _safetensors_meta(export_dir) + assert exported, "nothing exported" + missing = [k for k in orphans if k not in exported] + assert not missing, ( + f"{len(missing)} source weight(s) the model never loaded are absent from the export " + f"(e.g. {missing[0]}); PTQ could not touch them, so they must be copied through" + ) + + merged: dict = {} + for shard in sorted(export_dir.glob("*.safetensors")): + merged.update(load_file(str(shard))) + source: dict = {} + for shard in sorted(src_dir.glob("*.safetensors")): + source.update(load_file(str(shard))) + for k in orphans: + assert torch.equal(merged[k].cpu(), source[k].cpu()), ( + f"'{k}' was carried over but its value changed; it should be a verbatim copy" + ) + + +def _ptq_language_model_and_export(rank, size, *, src_dir, export_dir): + """Quantize ONLY the VLM's language model, then export the whole model.""" + + with patch_fsdp_mp_dtypes(): + model = AutoModelForImageTextToText.from_pretrained(src_dir, dtype=torch.bfloat16).to( + "cuda" + ) + model.eval() + fsdp2_wrap(model) + assert is_fsdp2_model(model), "fsdp2_wrap did not shard the VLM" + torch.distributed.barrier() + + language_model = getattr(model.model, "language_model", None) + assert language_model is not None, "fixture has no model.language_model to target" + + # A hardcoded token bound overruns a tiny fixture's embedding and surfaces as an opaque + # device-side assert, so take the vocab from the config. + text_config = getattr(model.config, "text_config", None) + vocab = getattr(text_config, "vocab_size", None) or model.config.vocab_size + input_ids = torch.randint(0, int(vocab), (1, 8), device="cuda") + mtq.quantize(language_model, mtq.FP8_DEFAULT_CFG, lambda m: model(input_ids=input_ids)) + torch.distributed.barrier() + + export_hf_checkpoint(model, export_dir=export_dir, max_shard_size=MAX_SHARD_SIZE) + torch.distributed.barrier() + + +@pytest.mark.timeout(600) +def test_fsdp2_distributed_export_of_vlm_keeps_vision_tower(dist_workers, tmp_path): + """Quantizing only a VLM's language model must still export the COMPLETE model. + + This is the common VLM recipe: the vision tower is left alone and only the language model is + quantized. The exported checkpoint is nonetheless expected to be the whole model -- a checkpoint + holding just the language half loads as a broken VLM rather than failing outright, so nothing + downstream would flag it. + + The writer builds its state dict from the model it is handed, so the risk is not that the vision + tower is dropped but that the paths keyed on quantizer state -- the expert split, the scale + postprocess, the reverse conversion -- mishandle a subtree that has none. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + src_dir = Path(create_tiny_qwen3vl_dir(tmp_path)) + export_dir = tmp_path / "export_vlm" + dist_workers.run( + partial(_ptq_language_model_and_export, src_dir=src_dir, export_dir=export_dir) + ) + + exported = _safetensors_meta(export_dir) + assert exported, "nothing exported" + + src_meta = _safetensors_meta(src_dir) + vision_src = sorted(k for k in src_meta if ".visual." in k or ".vision" in k) + assert vision_src, "fixture has no vision tower, so this test would prove nothing" + + missing = [k for k in vision_src if k not in exported] + assert not missing, ( + f"{len(missing)} vision-tower tensor(s) absent from the export of a VLM whose language " + f"model was quantized (e.g. {missing[0]}) -- the checkpoint is not the complete model" + ) + # The untouched tower must come through at its original precision, with no scales invented. + for k in vision_src: + dtype, shape = exported[k] + assert dtype not in _QUANTIZED_DTYPES, f"unquantized '{k}' exported as {dtype}" + assert shape == src_meta[k][1], f"'{k}': shape {shape}, source {src_meta[k][1]}" + stray = sorted(k for k in exported if (".visual." in k or ".vision" in k) and _is_scale(k)) + assert not stray, f"vision tower was not quantized but carries scales: {stray}" + + # ...and the language model must actually be quantized, or the test passes vacuously. + quantized = [k for k, (d, _) in exported.items() if d in _QUANTIZED_DTYPES] + assert quantized, "no tensor was quantized -- the language-model PTQ did not take effect" diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 7c972319d97..3296ccbfc86 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -22,6 +22,7 @@ import pytest import torch from _test_utils.torch.transformers_models import create_tiny_llama_dir +from safetensors.torch import load_file from torch.distributed.tensor import DTensor from modelopt.torch.export.unified_export_hf import export_hf_checkpoint @@ -99,3 +100,74 @@ def test_parallel_load_and_export(dist_workers, tmp_path, cpu_offload): cpu_offload=cpu_offload, ) ) + + +def _test_carry_over_unplaced_weights(rank, size, ckpt_dir, export_dir, orphan_prefix): + """Load, export, and require the weights the model could not place to survive.""" + device = torch.device(f"cuda:{rank}") + model = parallel_load_and_prepare_fsdp2(ckpt_dir, device, rank, size) + + unplaced = getattr(model, "_modelopt_unplaced_source_keys", None) + assert unplaced, ( + "the loader placed every checkpoint key; the fixture was supposed to leave an orphaned " + "layer behind, so this test would pass vacuously" + ) + assert any(k.startswith(orphan_prefix) for k in unplaced), ( + f"orphaned '{orphan_prefix}*' weights were not recorded as unplaced; got {sorted(unplaced)[:5]}" + ) + + export_hf_checkpoint(model, export_dir=export_dir, dtype=torch.bfloat16) + + if rank == 0: + exported: dict = {} + for shard in sorted(os.listdir(export_dir)): + if shard.endswith(".safetensors"): + exported.update(load_file(os.path.join(export_dir, shard))) + source: dict = {} + for shard in sorted(os.listdir(ckpt_dir)): + if shard.endswith(".safetensors"): + source.update(load_file(os.path.join(ckpt_dir, shard))) + + orphans = sorted(k for k in source if k.startswith(orphan_prefix)) + assert orphans, "fixture produced no orphaned weights" + for k in orphans: + assert k in exported, ( + f"'{k}' is in the source checkpoint and was never loaded into the model, so PTQ " + f"could not touch it -- it must be copied into the export, but it is missing" + ) + assert torch.equal(exported[k].cpu(), source[k].cpu()), ( + f"'{k}' was carried over but its value changed; it should be a verbatim copy" + ) + + +def test_carry_over_unplaced_weights(dist_workers, tmp_path): + """Weights the built model has no home for must still reach the exported checkpoint. + + A checkpoint can carry parameters the model class does not build -- an MTP head is the common + case (HF builds only ``num_hidden_layers`` decoders, leaving an inlined MTP tail orphaned), but + an auxiliary tower or draft head behaves the same way. Quantization never sees them, so nothing + downstream would notice their absence: the export just comes out quietly incomplete. + + The fixture reproduces that shape directly. Build a checkpoint with one more layer than the + config admits, so the final layer's weights are present on disk with nowhere to load to. + """ + ckpt_dir = create_tiny_llama_dir(tmp_path, vocab_size=VOCAB_SIZE, num_hidden_layers=3) + + # Truncate the config so the last layer becomes unplaceable -- the same situation an inlined + # MTP tail creates, without needing an MTP-capable architecture in the test suite. + cfg_path = os.path.join(str(ckpt_dir), "config.json") + with open(cfg_path) as f: + cfg = json.load(f) + orphan_idx = cfg["num_hidden_layers"] - 1 + cfg["num_hidden_layers"] = orphan_idx + with open(cfg_path, "w") as f: + json.dump(cfg, f) + + dist_workers.run( + partial( + _test_carry_over_unplaced_weights, + ckpt_dir=str(ckpt_dir), + export_dir=str(tmp_path / "export_carry"), + orphan_prefix=f"model.layers.{orphan_idx}.", + ) + ) diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 911567f75e5..8a8a7ed93f2 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -39,12 +39,14 @@ QUANTIZATION_W4A8_AWQ, ) from modelopt.torch.export.quant_utils import ( + _get_carried_over_module_names, _has_large_fp8_scale, get_kv_cache_scaling_factor, get_quant_config, get_quantization_format, postprocess_state_dict, process_layer_quant_config, + seed_carried_over_exclusions, uses_iq_quantization, ) from modelopt.torch.quantization.nn import ( @@ -773,3 +775,165 @@ def test_moe_router_names_handle_root_module(): block = _FakeMoEBlock(hidden=16) # name == "" for the root module; the router must be "gate", not ".gate". assert _get_unquantized_moe_router_names(block) == ["gate"] + + +def test_carried_over_weights_are_excluded_from_quantization(): + """A carried MTP head has no module, so nothing in the quantizer walk can see it. + + Its original-precision weight is copied into the export verbatim, so it must still reach + ``exclude_modules`` -- otherwise a deployment framework reads the top-level ``quant_algo`` + and tries to load ``eh_proj`` as an NVFP4 weight. Same failure as the MoE router above, + different cause: no quantizer there, no module at all here. + """ + hidden = 16 + model = _FakeMoEModel(hidden=hidden) + mtq.quantize(model, _nvfp4_all_linears_config, lambda m: m(torch.randn(2, hidden))) + + # The loader could not place these; the exporter copies them straight from the checkpoint. + model._modelopt_unplaced_source_keys = [ + "model.mtp.eh_proj.weight", + "model.mtp.eh_proj.bias", + "model.mtp.embed_tokens.weight", + ] + + quant_config = get_quant_config(model) + assert quant_config["quantization"]["quant_algo"] == "NVFP4" + + exclude_modules = quant_config["quantization"]["exclude_modules"] + for carried in ("model.mtp.eh_proj", "model.mtp.embed_tokens"): + assert any(fnmatch.fnmatch(carried, pattern) for pattern in exclude_modules), ( + f"carried weight {carried!r} missing from exclude_modules: {exclude_modules}" + ) + # The quantized experts must NOT be excluded. + assert not any(fnmatch.fnmatch("block.experts.0", pattern) for pattern in exclude_modules), ( + f"Quantized expert wrongly excluded: {exclude_modules}" + ) + + +def test_carried_over_module_names_strip_parameter_and_dedup(): + """Keys are ``.``; two params of one module yield one module name.""" + model = torch.nn.Module() + # No attribute at all -- the common case, and every non-carry-over caller. + assert _get_carried_over_module_names(model) == [] + + model._modelopt_unplaced_source_keys = [ + "a.b.weight", + "a.b.bias", # same module as above + "c.weight", + "toplevel", # no dot: no owning module, skipped + ] + assert _get_carried_over_module_names(model) == ["a.b", "c"] + + +def test_carried_over_names_prefer_what_the_export_actually_wrote(): + """Off-index sidecars never appear in unexpected_keys, so the unplaced list alone misses them. + + GLM-4.7 ships its MTP head in a standalone ``mtp.safetensors`` that the loader never opens. + The export copies it verbatim, so its tensors are in the checkpoint in original precision and + must reach ``exclude_modules`` -- the same requirement as a carried weight, via the other + mechanism. The export records both under ``_modelopt_carried_over_names``. + """ + model = torch.nn.Module() + model._modelopt_unplaced_source_keys = ["model.mtp.eh_proj.weight"] + # What the export actually wrote: the carried weight plus the copied sidecar's tensors. + model._modelopt_carried_over_names = [ + "model.mtp.eh_proj.weight", + "model.mtp.embed_tokens.weight", # only in the sidecar + ] + assert _get_carried_over_module_names(model) == [ + "model.mtp.eh_proj", + "model.mtp.embed_tokens", + ] + + +def test_carried_over_names_fall_back_before_the_export_records(): + """Callers that never ran the export still get the wider unplaced answer.""" + model = torch.nn.Module() + model._modelopt_unplaced_source_keys = ["model.mtp.eh_proj.weight"] + assert _get_carried_over_module_names(model) == ["model.mtp.eh_proj"] + + # An export that wrote nothing is an answer, not a missing one -- do not fall back to the + # wider list and claim exclusions for weights the checkpoint does not contain. + model._modelopt_carried_over_names = [] + assert _get_carried_over_module_names(model) == [] + + +def test_seed_carried_over_exclusions_adds_missing_names(): + """The layerwise exporter snapshots its config before the carried set exists, so it re-seeds.""" + model = torch.nn.Module() + model._modelopt_carried_over_names = [ + "model.mtp.eh_proj.weight", + "model.mtp.embed_tokens.weight", + ] + cfg = {"quantization": {"quant_algo": "NVFP4", "exclude_modules": ["lm_head"]}} + added = seed_carried_over_exclusions(model, cfg) + + assert added == ["model.mtp.eh_proj", "model.mtp.embed_tokens"] + assert cfg["quantization"]["exclude_modules"] == [ + "lm_head", + "model.mtp.eh_proj", + "model.mtp.embed_tokens", + ] + + +def test_seed_carried_over_exclusions_respects_existing_wildcards(): + """A recipe that already excluded mtp* must not gain redundant per-module entries.""" + model = torch.nn.Module() + model._modelopt_carried_over_names = ["model.mtp.eh_proj.weight"] + cfg = {"quantization": {"quant_algo": "NVFP4", "exclude_modules": ["model.mtp*"]}} + + assert seed_carried_over_exclusions(model, cfg) == [] + assert cfg["quantization"]["exclude_modules"] == ["model.mtp*"] + + +def test_seed_carried_over_exclusions_noop_without_a_uniform_format(): + """No single quant_algo means nothing for a deployment framework to misapply.""" + model = torch.nn.Module() + model._modelopt_carried_over_names = ["model.mtp.eh_proj.weight"] + for algo in (None, "MIXED_PRECISION"): + cfg = {"quantization": {"quant_algo": algo}} + assert seed_carried_over_exclusions(model, cfg) == [] + assert "exclude_modules" not in cfg["quantization"] + + +def test_both_export_paths_exclude_carried_weights_identically(): + """The unified and layerwise exporters must emit the same exclusions for the same model. + + They reach exclude_modules at different times -- get_quant_config after its per-layer pass, + the layerwise exporter from finalize() because bind() snapshotted its config during + calibration -- so it is easy for them to drift into different formats (wildcards one side, + literals the other) for identical inputs. Both go through seed_carried_over_exclusions now; + this pins that they agree. + """ + hidden = 16 + model = _FakeMoEModel(hidden=hidden) + mtq.quantize(model, _nvfp4_all_linears_config, lambda m: m(torch.randn(2, hidden))) + model._modelopt_carried_over_names = [ + "model.mtp.eh_proj.weight", + "model.mtp.embed_tokens.weight", + ] + + # Unified: seeded inside get_quant_config. + unified = get_quant_config(model)["quantization"]["exclude_modules"] + + # Layerwise: the same config minus the carried names, re-seeded at finalize() time. + layerwise_cfg = { + "quantization": { + "quant_algo": "NVFP4", + "exclude_modules": [e for e in unified if not e.startswith("model.mtp")], + } + } + seed_carried_over_exclusions(model, layerwise_cfg) + + assert sorted(layerwise_cfg["quantization"]["exclude_modules"]) == sorted(unified) + + +def test_seeded_exclusions_are_literal_module_names(): + """Exact names, not prefix wildcards: a literal cannot over-match a quantized module.""" + model = torch.nn.Module() + model._modelopt_carried_over_names = ["model.mtp.eh_proj.weight"] + cfg = {"quantization": {"quant_algo": "NVFP4", "exclude_modules": []}} + + assert seed_carried_over_exclusions(model, cfg) == ["model.mtp.eh_proj"] + assert cfg["quantization"]["exclude_modules"] == ["model.mtp.eh_proj"] + assert not any("*" in e for e in cfg["quantization"]["exclude_modules"]) diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py deleted file mode 100644 index 08292c65f9f..00000000000 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ /dev/null @@ -1,331 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Tests for modelopt/torch/export/plugins/hf_checkpoint_utils.py""" - -from types import SimpleNamespace -from unittest.mock import patch - -import pytest - -pytest.importorskip("huggingface_hub") -hf_hub_errors = pytest.importorskip("huggingface_hub.errors") -LocalEntryNotFoundError = hf_hub_errors.LocalEntryNotFoundError - -from modelopt.torch.export import ( - copy_hf_ckpt_remote_code, - copy_non_safetensor_files_from_ckpt, - sanitize_hf_config_for_deployment, -) -from modelopt.torch.export.plugins import hf_checkpoint_utils - - -def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_path): - src_dir = tmp_path / "src" - src_dir.mkdir() - (src_dir / "model.safetensors").write_text("weights") - (src_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') - (src_dir / "pytorch_model.bin").write_text("weights") - (src_dir / "stats.npy").write_text("stats") - (src_dir / "reasoning_parser.py").write_text("parser") - - default_dst = tmp_path / "default" - copy_non_safetensor_files_from_ckpt(src_dir, default_dst) - assert not (default_dst / "model.safetensors").exists() - assert not (default_dst / "model.safetensors.index.json").exists() - assert (default_dst / "pytorch_model.bin").exists() - assert (default_dst / "stats.npy").exists() - - filtered_dst = tmp_path / "filtered" - copy_non_safetensor_files_from_ckpt( - src_dir, - filtered_dst, - exclude_patterns=("*.bin", "*.npy"), - ) - assert (filtered_dst / "reasoning_parser.py").exists() - assert not (filtered_dst / "pytorch_model.bin").exists() - assert not (filtered_dst / "stats.npy").exists() - - -def test_copy_non_safetensor_files_from_ckpt_continues_after_copy_failure(tmp_path, monkeypatch): - src_dir = tmp_path / "src" - src_dir.mkdir() - (src_dir / "bad.py").write_text("bad") - (src_dir / "good.py").write_text("good") - - original_copy2 = hf_checkpoint_utils.shutil.copy2 - - def copy2(source, *args, **kwargs): - if source.endswith("bad.py"): - raise PermissionError("unreadable") - return original_copy2(source, *args, **kwargs) - - monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2) - with pytest.warns(UserWarning, match="bad.py"): - copied_files = copy_non_safetensor_files_from_ckpt(src_dir, tmp_path / "dst") - - assert copied_files == ["good.py"] - - -def test_copy_hf_ckpt_remote_code_local_dir(tmp_path): - """copy_hf_ckpt_remote_code copies top-level .py files from a local directory.""" - src_dir = tmp_path / "src" - src_dir.mkdir() - (src_dir / "modeling_custom.py").write_text("# custom model") - (src_dir / "configuration_custom.py").write_text("# custom config") - (src_dir / "not_python.txt").write_text("not python") - (src_dir / "subdir").mkdir() - (src_dir / "subdir" / "nested.py").write_text("# nested — should not be copied") - - dst_dir = tmp_path / "dst" - dst_dir.mkdir() - - copy_hf_ckpt_remote_code(src_dir, dst_dir) - - assert (dst_dir / "modeling_custom.py").read_text() == "# custom model" - assert (dst_dir / "configuration_custom.py").read_text() == "# custom config" - assert not (dst_dir / "not_python.txt").exists(), "non-.py files should not be copied" - assert not (dst_dir / "nested.py").exists(), "nested .py files should not be copied" - - -def test_copy_hf_ckpt_remote_code_local_dir_no_py_files(tmp_path): - """copy_hf_ckpt_remote_code is a no-op when the local directory has no .py files.""" - src_dir = tmp_path / "src" - src_dir.mkdir() - (src_dir / "config.json").write_text("{}") - - dst_dir = tmp_path / "dst" - dst_dir.mkdir() - - copy_hf_ckpt_remote_code(src_dir, dst_dir) # should not raise - - assert list(dst_dir.iterdir()) == [], "no files should be copied" - - -def test_copy_hf_ckpt_remote_code_hub_id(tmp_path, monkeypatch): - """copy_hf_ckpt_remote_code copies .py files from the resolved Hub snapshot.""" - dst_dir = tmp_path / "dst" - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "modeling_custom.py").write_text("# custom model") - (snapshot_dir / "not_python.txt").write_text("not python") - - monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) - with patch( - "modelopt.torch.export.plugins.hf_checkpoint_utils.snapshot_download", - return_value=str(snapshot_dir), - ) as mock_sd: - copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-Nano-12B-v2", dst_dir) - - mock_sd.assert_called_once_with( - repo_id="nvidia/NVIDIA-Nemotron-Nano-12B-v2", - allow_patterns=["*.py"], - local_files_only=False, - ) - assert (dst_dir / "modeling_custom.py").read_text() == "# custom model" - assert not (dst_dir / "not_python.txt").exists(), "non-.py files should not be copied" - - -def test_copy_hf_ckpt_remote_code_hub_id_offline_uses_cache(tmp_path, monkeypatch): - """copy_hf_ckpt_remote_code resolves cached Hub snapshots when HF_HUB_OFFLINE is set.""" - dst_dir = tmp_path / "dst" - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "nemotron_reasoning_parser.py").write_text("# parser") - - monkeypatch.setenv("HF_HUB_OFFLINE", "1") - with patch( - "modelopt.torch.export.plugins.hf_checkpoint_utils.snapshot_download", - return_value=str(snapshot_dir), - ) as mock_sd: - copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", dst_dir) - - mock_sd.assert_called_once_with( - repo_id="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", - allow_patterns=["*.py"], - local_files_only=True, - ) - assert (dst_dir / "nemotron_reasoning_parser.py").read_text() == "# parser" - - -def test_copy_hf_ckpt_remote_code_hub_id_offline_missing_cache_raises(tmp_path, monkeypatch): - """copy_hf_ckpt_remote_code raises a clear error when offline cache is missing.""" - monkeypatch.setenv("HF_HUB_OFFLINE", "1") - with ( - patch( - "modelopt.torch.export.plugins.hf_checkpoint_utils.snapshot_download", - side_effect=LocalEntryNotFoundError("missing"), - ), - pytest.raises(RuntimeError, match="HF_HUB_OFFLINE"), - ): - copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", tmp_path / "dst") - - -def test_sanitize_hf_config_for_deployment_trims_nextn_layer_types(): - """Drop MTP/next-token-prediction layer types from exported config.json.""" - hidden_layer_types = ["full_attention"] * 45 - nextn_layer_types = ["nextn_predict"] * 3 - config_data = { - "num_hidden_layers": 45, - "num_nextn_predict_layers": 3, - "layer_types": hidden_layer_types + nextn_layer_types, - } - - with pytest.warns(UserWarning, match="Trimming config.layer_types"): - sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace()) - - assert config_data["layer_types"] == hidden_layer_types - - -def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_parameters(): - """Transformers 5.x requires rope_theta inside llama3 rope_parameters.""" - config_data = { - "rope_theta": 500000, - "rope_parameters": { - "rope_type": "llama3", - "factor": 8.0, - "original_max_position_embeddings": 4096, - "low_freq_factor": 1.0, - "high_freq_factor": 4.0, - }, - } - - sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) - - assert config_data["rope_parameters"]["rope_theta"] == 500000 - - -def test_sanitize_hf_config_for_deployment_uses_model_rope_theta_for_rope_parameters(): - """Use model.config.rope_theta when save_pretrained omits the top-level field.""" - config_data = { - "rope_parameters": { - "rope_type": "llama3", - "factor": 8.0, - "original_max_position_embeddings": 4096, - "low_freq_factor": 1.0, - "high_freq_factor": 4.0, - }, - } - model = SimpleNamespace(config=SimpleNamespace(rope_theta=500000)) - - sanitize_hf_config_for_deployment(config_data, model) - - assert config_data["rope_parameters"]["rope_theta"] == 500000 - - -def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_scaling(): - """Legacy rope_scaling metadata is normalized for llama3 configs as well.""" - config_data = { - "rope_scaling": { - "type": "llama3", - "factor": 8.0, - "original_max_position_embeddings": 4096, - "low_freq_factor": 1.0, - "high_freq_factor": 4.0, - }, - } - model = SimpleNamespace(config=SimpleNamespace(rope_theta=500000)) - - sanitize_hf_config_for_deployment(config_data, model) - - assert config_data["rope_scaling"]["rope_theta"] == 500000 - - -def test_sanitize_hf_config_for_deployment_keeps_existing_rope_theta(): - """Existing rope_theta in rope metadata is not overwritten.""" - config_data = { - "rope_theta": 500000, - "rope_parameters": { - "rope_type": "llama3", - "rope_theta": 1000000, - }, - } - - sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) - - assert config_data["rope_parameters"]["rope_theta"] == 1000000 - - -def test_sanitize_hf_config_for_deployment_ignores_non_llama3_rope_parameters(): - """Only llama3 RoPE parameters need the Transformers 5.x compatibility fix.""" - config_data = { - "rope_theta": 500000, - "rope_parameters": { - "rope_type": "default", - }, - } - - sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) - - assert "rope_theta" not in config_data["rope_parameters"] - - -def test_sanitize_hf_config_for_deployment_uses_model_config_nextn_count(): - """Handle exports where save_pretrained omits num_nextn_predict_layers.""" - config_data = { - "num_hidden_layers": 2, - "layer_types": ["full_attention", "linear_attention", "nextn_predict"], - } - model = SimpleNamespace(config=SimpleNamespace(num_nextn_predict_layers=1)) - - with pytest.warns(UserWarning, match="Trimming config.layer_types"): - sanitize_hf_config_for_deployment(config_data, model=model) - - assert config_data["layer_types"] == ["full_attention", "linear_attention"] - - -def test_sanitize_hf_config_for_deployment_counts_mtp_layer_prefixes(): - """Do not count broad MTP exclude prefixes as prediction layers.""" - config_data = { - "num_hidden_layers": 2, - "layer_types": ["full_attention", "linear_attention", "nextn_predict"], - } - model = SimpleNamespace(_mtp_layer_prefixes=["mtp", "mtp.layers.0"]) - - with pytest.warns(UserWarning, match="Trimming config.layer_types"): - sanitize_hf_config_for_deployment(config_data, model=model) - - assert config_data["layer_types"] == ["full_attention", "linear_attention"] - - -def test_sanitize_hf_config_for_deployment_ignores_broad_mtp_prefix_only(): - """Do not infer prediction-layer count from a broad exclude prefix alone.""" - config_data = { - "num_hidden_layers": 2, - "layer_types": ["full_attention", "linear_attention", "nextn_predict"], - } - model = SimpleNamespace(_mtp_layer_prefixes=["mtp"]) - - sanitize_hf_config_for_deployment(config_data, model=model) - - assert config_data["layer_types"] == ["full_attention", "linear_attention", "nextn_predict"] - - -def test_sanitize_hf_config_for_deployment_keeps_unexplained_layer_type_mismatch(): - """Do not rewrite config when extra layer types are not explained by nextn metadata.""" - config_data = { - "num_hidden_layers": 2, - "num_nextn_predict_layers": 1, - "layer_types": ["full_attention", "linear_attention", "extra_a", "extra_b"], - } - - sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace()) - - assert config_data["layer_types"] == [ - "full_attention", - "linear_attention", - "extra_a", - "extra_b", - ] diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 94f6dee5053..4827762a8d2 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -15,6 +15,8 @@ """Tests for tied-weight helpers in unified_export_hf.""" +import importlib.util +import json from types import SimpleNamespace import pytest @@ -23,6 +25,7 @@ make_tied_linear_pair, wrap_in_parent_with_tied_keys, ) +from safetensors.torch import save_file import modelopt.torch.quantization as mtq from modelopt.torch.export.model_utils import ( @@ -35,7 +38,7 @@ postprocess_state_dict, sync_tied_input_amax, ) -from modelopt.torch.export.unified_export_hf import _resolve_export_dtype +from modelopt.torch.export.unified_export_hf import _resolve_export_dtype, read_unplaced_weights from modelopt.torch.quantization.nn import TensorQuantizer @@ -571,3 +574,414 @@ def test_fuse_prequant_layernorm_fuses_and_removes_pre_quant_scale(): for module in modules: assert not hasattr(module.input_quantizer, "_pre_quant_scale") assert module.fused_with_prequant + + +# --- carrying over weights the loader could not place -------------------------------------------- + + +# ``read_unplaced_weights`` imports model_load_utils inside its try block, before it +# looks at the recorded keys, and model_load_utils imports accelerate at module scope. With +# accelerate absent every call raises ImportError, gets caught, warns and returns {} -- so these +# tests would still "pass" while exercising none of the logic they name. +requires_accelerate = pytest.mark.skipif( + importlib.util.find_spec("accelerate") is None, + reason="carry-over goes through model_load_utils, which requires accelerate", +) + + +class _ProvenanceModel(torch.nn.Module): + """A model carrying only what the carry-over reads: recorded keys and a source path.""" + + def __init__(self, keys=None, ckpt=None, name_or_path=None): + super().__init__() + self.lin = torch.nn.Linear(2, 2) + # Sentinel already-placed weights this file's checkpoints use to stand in for "the + # model loaded this one fine": real parameters, so the union's structural pass does + # not also flag them as unplaced (it only knows a checkpoint key by whether the model + # has a matching parameter, not by which test wrote it). + self.a = torch.nn.Linear(1, 1) + self.other = torch.nn.Linear(1, 1) + if keys is not None: + self._modelopt_unplaced_source_keys = keys + if ckpt is not None: + self._modelopt_source_checkpoint = str(ckpt) + if name_or_path is not None: + self.config = SimpleNamespace(_name_or_path=str(name_or_path)) + + +@requires_accelerate +def test_carry_over_returns_nothing_when_the_loader_placed_everything(): + """A recorded empty list means the question was asked and answered -- do not re-derive.""" + model = _ProvenanceModel(keys=[], ckpt="/anywhere") + assert read_unplaced_weights(model) == {} + + +@requires_accelerate +def test_carry_over_returns_nothing_without_provenance(): + assert read_unplaced_weights(_ProvenanceModel()) == {} + + +@requires_accelerate +def test_carry_over_returns_nothing_for_a_hub_id_rather_than_a_local_path(): + """``_name_or_path`` is often a hub id; there is nothing on disk to re-read.""" + assert read_unplaced_weights(_ProvenanceModel(name_or_path="org/Some-Model")) == {} + + +@requires_accelerate +def test_carry_over_reads_the_recorded_keys_off_disk(tmp_path): + save_file( + {"kept.weight": torch.arange(4, dtype=torch.float32), "other.weight": torch.zeros(2)}, + str(tmp_path / "model.safetensors"), + ) + model = _ProvenanceModel(keys=["kept.weight"], ckpt=tmp_path) + + carried = read_unplaced_weights(model) + + assert list(carried) == ["kept.weight"] + assert torch.equal(carried["kept.weight"], torch.arange(4, dtype=torch.float32)) + + +@requires_accelerate +def test_carry_over_warns_and_keeps_the_export_alive_when_the_checkpoint_is_unreadable(tmp_path): + """Best-effort: the rest of the weights are already correct, so this must not abort the export.""" + model = _ProvenanceModel(keys=["kept.weight"], ckpt=tmp_path / "does-not-exist") + with pytest.warns(UserWarning, match="Could not copy"): + assert read_unplaced_weights(model) == {} + + +@requires_accelerate +def test_carry_over_handler_survives_failing_before_the_keys_are_known(tmp_path, monkeypatch): + """The failure can land while ``keys`` is still None, and the handler must not throw itself. + + Nothing was recorded here, so the keys are derived from provenance -- and that derivation is + what fails. Counting the keys unconditionally in the warning would raise TypeError from inside + the very handler that exists to keep the export standing. + """ + from modelopt.torch.utils.plugins import model_load_utils + + def _boom(*args, **kwargs): + raise RuntimeError("cannot read the index") + + monkeypatch.setattr(model_load_utils, "unplaced_source_keys", _boom) + # A real directory holding safetensors, so the derivation is actually reached: a checkpoint + # with no safetensors short-circuits earlier, before anything can fail. + (tmp_path / "model.safetensors.index.json").write_text('{"weight_map": {}}') + model = _ProvenanceModel(name_or_path=tmp_path) + + # The message changed with the union derivation, deliberately: the structural pass failing + # is not the same as "could not copy N weights" -- with nothing recorded we do not know that + # anything is missing, only that we could not check. What must still hold is that the handler + # does not throw from inside itself. + with pytest.warns(UserWarning, match="Could not derive unplaced source keys"): + assert read_unplaced_weights(model) == {} + + +def test_carry_over_is_quiet_for_a_checkpoint_with_no_safetensors(tmp_path, recwarn): + """A pytorch_model.bin checkpoint has nothing this path can read, and nothing to carry. + + Warning that weights are "missing from the export" would be alarming and wrong -- there are + no unplaced weights, only a format this reader does not handle. + """ + (tmp_path / "pytorch_model.bin").write_bytes(b"not safetensors") + model = _ProvenanceModel(name_or_path=tmp_path) + + assert read_unplaced_weights(model) == {} + assert [w for w in recwarn if "Could not copy" in str(w.message)] == [] + + +def test_carry_over_is_quiet_without_the_loader_dependencies(tmp_path, monkeypatch, recwarn): + """The quiet path must not depend on the loader's optional imports. + + model_load_utils imports transformers and accelerate at module scope, and the partial-install + environments have neither. If the short-circuit sat below that import, the ImportError would + land in the handler and emit the very "will be missing them" warning it exists to avoid. + + Blocking safetensors here would prove nothing: this module binds ``safe_open`` at import time, + so patching sys.modules afterwards cannot affect it. + """ + import sys as _sys + + monkeypatch.setitem(_sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None) + monkeypatch.setitem(_sys.modules, "transformers", None) + monkeypatch.setitem(_sys.modules, "accelerate", None) + (tmp_path / "pytorch_model.bin").write_bytes(b"not safetensors") + model = _ProvenanceModel(name_or_path=tmp_path) + + assert read_unplaced_weights(model) == {} + assert [w for w in recwarn if "Could not copy" in str(w.message)] == [] + + +def test_carryable_unplaced_keys_skips_keys_no_shard_backs(tmp_path): + """Unplaced != carryable. A stale buffer listed by the model is not a weight to lose. + + ``--vllm_fakequant_export`` refuses to run when real weights would be dropped, so this + distinction decides whether working exports keep working. + """ + from modelopt.torch.export.unified_export_hf import carryable_unplaced_keys + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"model.mtp.eh_proj.weight": "mtp-0001.safetensors"}}' + ) + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = [ + "model.mtp.eh_proj.weight", # a shard has it -- losing it matters + "model.layers.0.self_attn.rotary_emb.inv_freq", # nothing backs it + ] + assert carryable_unplaced_keys(model) == ["model.mtp.eh_proj.weight"] + + +def test_carryable_unplaced_keys_is_quiet_when_nothing_was_recorded(): + """No provenance, no answer -- and no exception from a diagnostic helper.""" + from modelopt.torch.export.unified_export_hf import carryable_unplaced_keys + + assert carryable_unplaced_keys(torch.nn.Module()) == [] + + +def test_carryable_unplaced_keys_works_without_the_loader_dependencies(tmp_path, monkeypatch): + """The shard-backed question must be answerable where transformers/accelerate are absent. + + model_load_utils imports them at module scope, so routing through it would make this return + "nothing to carry" in the partial-install environments -- and the --vllm_fakequant_export + guard would then stay silent on a checkpoint whose weights it really would drop. + """ + import sys as _sys + + from modelopt.torch.export.unified_export_hf import carryable_unplaced_keys + + for mod in ("transformers", "accelerate", "huggingface_hub", "safetensors"): + monkeypatch.setitem(_sys.modules, mod, None) + monkeypatch.setitem(_sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None) + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"model.mtp.eh_proj.weight": "mtp-0001.safetensors"}}' + ) + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = [ + "model.mtp.eh_proj.weight", + "model.layers.0.self_attn.rotary_emb.inv_freq", + ] + assert carryable_unplaced_keys(model) == ["model.mtp.eh_proj.weight"] + + +def test_layerwise_finalize_sees_the_carried_keys(tmp_path, monkeypatch): + """The layerwise fix depends on an ordering: record the carried set, THEN call finalize(). + + LayerwiseExporter.bind() snapshots its quant config during calibration, so finalize() is the + only point where it can learn what the export carried. If export_hf_checkpoint ever records + _modelopt_carried_over_names after dispatching to the exporter -- or stops recording it on + that path -- the sidecar exclusions silently go missing again, with nothing else to catch it. + """ + from modelopt.torch.export import unified_export_hf as uehf + from modelopt.torch.export.layerwise_export import LAYERWISE_EXPORTER_ATTR + + seen = {} + + class _Exporter: + def finalize(self, extra_state_dict=None): + seen["keys"] = getattr(model, "_modelopt_carried_over_names", "") + return {} + + model = torch.nn.Module() + setattr(model, LAYERWISE_EXPORTER_ATTR, _Exporter()) + # Accepts **kwargs because the call site passes keys_only: non-writing ranks resolve + # names without reading tensors, and a stub that ignores that would hide a signature drift. + monkeypatch.setattr(uehf, "read_unplaced_weights", lambda m, **kw: {}) + monkeypatch.setattr(uehf, "off_index_tensor_names", lambda m: ["model.mtp.eh_proj.weight"]) + + uehf.export_hf_checkpoint(model, export_dir=tmp_path) + + assert seen["keys"] == ["model.mtp.eh_proj.weight"], ( + f"finalize() saw {seen['keys']!r}; the carried set must be recorded before dispatch" + ) + + +def test_carries_a_key_the_index_does_not_list(tmp_path): + """A tensor inside a main shard but absent from weight_map must still be carried. + + model.safetensors.index.json is not a complete inventory. Transformers enumerates the contents + of each shard it opens, so it reports such a tensor in unexpected_keys and it reaches + _modelopt_unplaced_source_keys -- but resolving the file purely through weight_map finds + nothing and used to drop it silently, with the --vllm_fakequant_export guard staying quiet + too because it shared that lookup. An MTP head stored in a main shard is exactly this shape: + when MTP is not quantized the loader never places it, so it is the case the carry-over exists + for. + """ + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import ( + carryable_unplaced_keys, + read_unplaced_weights, + ) + + shard, extra = "model-00001-of-00001.safetensors", "model.mtp.eh_proj.weight" + save_file({"a.weight": torch.zeros(2), extra: torch.full((2,), 7.0)}, str(tmp_path / shard)) + # The index deliberately omits `extra`. + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = [extra] + + carried = read_unplaced_weights(model) + assert extra in carried, f"un-indexed key dropped from the export: {sorted(carried)}" + assert torch.equal(carried[extra], torch.full((2,), 7.0)) + + # The guard must see it too, or it stays silent on the very weights that would be lost. + assert carryable_unplaced_keys(model) == [extra] + + +def test_warns_when_a_recorded_key_is_in_no_shard(tmp_path): + """A key in neither the index nor any file is reported, not silently ignored.""" + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import read_unplaced_weights + + shard = "model-00001-of-00001.safetensors" + save_file({"a.weight": torch.zeros(2)}, str(tmp_path / shard)) + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = ["ghost.weight"] + + with pytest.warns(UserWarning, match="in no safetensors file"): + assert read_unplaced_weights(model) == {} + + +def test_carry_over_works_without_the_loader_dependencies(tmp_path, monkeypatch): + """Recorded keys must carry where transformers/accelerate are absent. + + Only the `keys is None` fallback needs model_load_utils. Importing it unconditionally made the + recorded-keys path -- which needs nothing from it -- fail in the partial-install environments, + where the handler reported "could not copy" and dropped every carried weight. The sibling test + pins this for carryable_unplaced_keys; this pins the carry itself, which is what actually writes. + """ + import sys as _sys + + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import read_unplaced_weights + + for mod in ("transformers", "accelerate", "huggingface_hub"): + monkeypatch.setitem(_sys.modules, mod, None) + monkeypatch.setitem(_sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None) + + shard = "model-00001-of-00001.safetensors" + save_file({"model.mtp.eh_proj.weight": torch.full((2,), 3.0)}, str(tmp_path / shard)) + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"model.mtp.eh_proj.weight": "model-00001-of-00001.safetensors"}}' + ) + + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = ["model.mtp.eh_proj.weight"] + + carried = read_unplaced_weights(model) + assert "model.mtp.eh_proj.weight" in carried, "recorded keys must not need the loader imports" + + +def test_union_survives_an_architecture_that_ignores_its_mtp_keys(tmp_path, monkeypatch): + """An empty loader report must not mean "nothing to carry". + + Transformers filters unexpected_keys through _keys_to_ignore_on_load_unexpected, and + Qwen3-Next ignores ^mtp.*, DeepSeek-V3 and GLM their own MTP prefixes. So for exactly the + heads this mechanism exists to carry, the report comes back EMPTY -- and an earlier revision + treated [] as authoritative and shipped the export without them. The structural pass has no + such blind spot, because it asks whether the model has a parameter rather than what the + loader chose to mention. + """ + # monkeypatch.setattr on a dotted path imports the module for real to resolve it, and + # model_load_utils genuinely needs transformers (this test pins behaviour of ITS loader, + # ignore rules included) -- so there is nothing meaningful to assert without it. + pytest.importorskip("transformers") + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import read_unplaced_weights + + shard, mtp = "model-00001-of-00001.safetensors", "mtp.fc.weight" + save_file({"a.weight": torch.zeros(2), mtp: torch.full((2,), 5.0)}, str(tmp_path / shard)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"a.weight": shard, mtp: shard}}) + ) + + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + # What an architecture with an ^mtp.* ignore rule leaves behind. + model._modelopt_unplaced_source_keys = [] + + monkeypatch.setattr( + "modelopt.torch.utils.plugins.model_load_utils.unplaced_source_keys", + lambda m, c: [mtp], + ) + carried = read_unplaced_weights(model) + assert mtp in carried, f"an ignored MTP key was dropped from the export: {sorted(carried)}" + assert torch.equal(carried[mtp], torch.full((2,), 5.0)) + + +def test_union_prefers_source_keys_over_converted_names(tmp_path, monkeypatch): + """A fused target name is not a checkpoint key and cannot be located. + + The loader can report a post-conversion name (``...experts.gate_up_proj``) that exists in no + shard, standing for several source tensors. The structural pass resolves in the other + direction -- source key through the converters to a target -- so it names the tensors that are + actually on disk, and the union carries them. + """ + # Same reasoning as test_union_survives_an_architecture_that_ignores_its_mtp_keys above: + # the converter resolution this test pins is transformers', not ours. + pytest.importorskip("transformers") + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import read_unplaced_weights + + shard = "model-00001-of-00001.safetensors" + g = "model.layers.1.mlp.experts.0.gate_proj.weight" + u = "model.layers.1.mlp.experts.0.up_proj.weight" + save_file({g: torch.ones(2), u: torch.full((2,), 2.0)}, str(tmp_path / shard)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {g: shard, u: shard}}) + ) + + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + # The loader reports the FUSED name, which is in neither the index nor any shard. + model._modelopt_unplaced_source_keys = ["model.layers.1.mlp.experts.gate_up_proj"] + + monkeypatch.setattr( + "modelopt.torch.utils.plugins.model_load_utils.unplaced_source_keys", + lambda m, c: [g, u], + ) + with pytest.warns(UserWarning, match="in no safetensors file"): + carried = read_unplaced_weights(model) + assert {g, u} <= set(carried), f"fused name stranded its sources: {sorted(carried)}" + assert torch.equal(carried[g], torch.ones(2)) + assert torch.equal(carried[u], torch.full((2,), 2.0)) + + +def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeypatch): + """Where transformers/accelerate are absent the structural pass cannot run. + + The recorded set then stands alone -- worse than the union, but the export must still carry + what it can rather than refuse outright. + """ + import sys as _sys + + from safetensors.torch import save_file + + from modelopt.torch.export.unified_export_hf import read_unplaced_weights + + shard, key = "model-00001-of-00001.safetensors", "model.mtp.eh_proj.weight" + save_file({key: torch.full((2,), 3.0)}, str(tmp_path / shard)) + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {key: shard}})) + + monkeypatch.setitem(_sys.modules, "modelopt.torch.utils.plugins.model_load_utils", None) + model = _ProvenanceModel(name_or_path=tmp_path) + model._modelopt_source_checkpoint = str(tmp_path) + model._modelopt_unplaced_source_keys = [key] + + carried = read_unplaced_weights(model) + assert key in carried, "recorded keys must still carry when the structural pass is unavailable" diff --git a/tests/unit/torch/export/test_vllm_fakequant_hf.py b/tests/unit/torch/export/test_vllm_fakequant_hf.py new file mode 100644 index 00000000000..795fc8a2961 --- /dev/null +++ b/tests/unit/torch/export/test_vllm_fakequant_hf.py @@ -0,0 +1,106 @@ +# 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. + +"""Tests for the vLLM fake-quant exporter's checkpoint-weight carry-over.""" + +import json + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.export.plugins import vllm_fakequant_hf as vfq + + +def test_carry_over_unplaced_weights_writes_extra_shard_and_reindexes(tmp_path, monkeypatch): + """A single-file base checkpoint (the common, unindexed case) gains a shard and an index. + + export_hf_vllm_fq_checkpoint's inplace_mem_efficient path deliberately never passes an + explicit state_dict= to save_pretrained (it would crash on meta tensors for offloaded + params -- see the comment at its call site), so there is no state_dict to merge unplaced + weights into the way export_hf_checkpoint does. _carry_over_unplaced_weights instead writes + them as their own shard after save_pretrained has already run, and rebuilds the index from + every shard on disk so a checkpoint that started as a single unindexed file still ends up + correctly indexed once a second file exists. + """ + save_file({"model.embed.weight": torch.zeros(4, 4)}, str(tmp_path / "model.safetensors")) + assert not (tmp_path / "model.safetensors.index.json").exists() + + extra = { + "mtp.eh_proj.weight": torch.full((2, 2), 3.0), + "mtp.norm.weight": torch.full((2,), 5.0), + } + monkeypatch.setattr(vfq, "read_unplaced_weights", lambda model: extra) + + vfq._carry_over_unplaced_weights(tmp_path, model=torch.nn.Module()) + + shard_path = tmp_path / "model-carried-over.safetensors" + assert shard_path.exists() + + index = json.loads((tmp_path / "model.safetensors.index.json").read_text()) + weight_map = index["weight_map"] + assert weight_map == { + "model.embed.weight": "model.safetensors", + "mtp.eh_proj.weight": "model-carried-over.safetensors", + "mtp.norm.weight": "model-carried-over.safetensors", + } + assert index["metadata"]["total_size"] > 0 + + with safe_open(str(shard_path), framework="pt") as f: + assert set(f.keys()) == set(extra) + assert torch.equal(f.get_tensor("mtp.eh_proj.weight"), extra["mtp.eh_proj.weight"]) + + +def test_carry_over_unplaced_weights_extends_an_existing_sharded_index(tmp_path, monkeypatch): + """A checkpoint save_pretrained already sharded keeps its own entries after the extra shard lands.""" + save_file({"a.weight": torch.zeros(2)}, str(tmp_path / "model-00001-of-00002.safetensors")) + save_file({"b.weight": torch.zeros(2)}, str(tmp_path / "model-00002-of-00002.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {"total_size": 16}, + "weight_map": { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + }, + } + ) + ) + + extra = {"mtp.eh_proj.weight": torch.full((2, 2), 3.0)} + monkeypatch.setattr(vfq, "read_unplaced_weights", lambda model: extra) + + vfq._carry_over_unplaced_weights(tmp_path, model=torch.nn.Module()) + + weight_map = json.loads((tmp_path / "model.safetensors.index.json").read_text())["weight_map"] + assert weight_map["a.weight"] == "model-00001-of-00002.safetensors" + assert weight_map["b.weight"] == "model-00002-of-00002.safetensors" + assert weight_map["mtp.eh_proj.weight"] == "model-carried-over.safetensors" + + +def test_carry_over_unplaced_weights_is_a_true_noop_when_nothing_to_carry(tmp_path, monkeypatch): + """No unplaced weights means the checkpoint on disk is untouched, not merely unchanged in content. + + Writing a same-content index anyway would still be a behavior change for the (common) case of + a model with nothing to carry: every plain export would gain an index.json it did not have + before. + """ + save_file({"model.embed.weight": torch.zeros(4, 4)}, str(tmp_path / "model.safetensors")) + monkeypatch.setattr(vfq, "read_unplaced_weights", lambda model: {}) + + vfq._carry_over_unplaced_weights(tmp_path, model=torch.nn.Module()) + + assert not (tmp_path / "model-carried-over.safetensors").exists() + assert not (tmp_path / "model.safetensors.index.json").exists() diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py new file mode 100644 index 00000000000..44bc3fa3544 --- /dev/null +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -0,0 +1,609 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for modelopt/torch/utils/plugins/hf_checkpoint_utils.py""" + +import json +import sys +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from safetensors.torch import save_file + +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_off_index_safetensors, + indexed_weight_map, + off_index_safetensors_files, + read_safetensors_subset, +) + +pytest.importorskip("huggingface_hub") +hf_hub_errors = pytest.importorskip("huggingface_hub.errors") +LocalEntryNotFoundError = hf_hub_errors.LocalEntryNotFoundError + +from modelopt.torch.utils.plugins import hf_checkpoint_utils +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_hf_ckpt_remote_code, + copy_non_safetensor_files_from_ckpt, + sanitize_hf_config_for_deployment, +) + + +def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_path): + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "model.safetensors").write_text("weights") + (src_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + (src_dir / "pytorch_model.bin").write_text("weights") + (src_dir / "stats.npy").write_text("stats") + (src_dir / "reasoning_parser.py").write_text("parser") + + default_dst = tmp_path / "default" + copy_non_safetensor_files_from_ckpt(src_dir, default_dst) + assert not (default_dst / "model.safetensors").exists() + assert not (default_dst / "model.safetensors.index.json").exists() + assert (default_dst / "pytorch_model.bin").exists() + assert (default_dst / "stats.npy").exists() + + filtered_dst = tmp_path / "filtered" + copy_non_safetensor_files_from_ckpt( + src_dir, + filtered_dst, + exclude_patterns=("*.bin", "*.npy"), + ) + assert (filtered_dst / "reasoning_parser.py").exists() + assert not (filtered_dst / "pytorch_model.bin").exists() + assert not (filtered_dst / "stats.npy").exists() + + +def test_copy_non_safetensor_files_from_ckpt_continues_after_copy_failure(tmp_path, monkeypatch): + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "bad.py").write_text("bad") + (src_dir / "good.py").write_text("good") + + original_copy2 = hf_checkpoint_utils.shutil.copy2 + + def copy2(source, *args, **kwargs): + if source.endswith("bad.py"): + raise PermissionError("unreadable") + return original_copy2(source, *args, **kwargs) + + monkeypatch.setattr(hf_checkpoint_utils.shutil, "copy2", copy2) + with pytest.warns(UserWarning, match="bad.py"): + copied_files = copy_non_safetensor_files_from_ckpt(src_dir, tmp_path / "dst") + + assert copied_files == ["good.py"] + + +def test_copy_hf_ckpt_remote_code_local_dir(tmp_path): + """copy_hf_ckpt_remote_code copies top-level .py files from a local directory.""" + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "modeling_custom.py").write_text("# custom model") + (src_dir / "configuration_custom.py").write_text("# custom config") + (src_dir / "not_python.txt").write_text("not python") + (src_dir / "subdir").mkdir() + (src_dir / "subdir" / "nested.py").write_text("# nested — should not be copied") + + dst_dir = tmp_path / "dst" + dst_dir.mkdir() + + copy_hf_ckpt_remote_code(src_dir, dst_dir) + + assert (dst_dir / "modeling_custom.py").read_text() == "# custom model" + assert (dst_dir / "configuration_custom.py").read_text() == "# custom config" + assert not (dst_dir / "not_python.txt").exists(), "non-.py files should not be copied" + assert not (dst_dir / "nested.py").exists(), "nested .py files should not be copied" + + +def test_copy_hf_ckpt_remote_code_local_dir_no_py_files(tmp_path): + """copy_hf_ckpt_remote_code is a no-op when the local directory has no .py files.""" + src_dir = tmp_path / "src" + src_dir.mkdir() + (src_dir / "config.json").write_text("{}") + + dst_dir = tmp_path / "dst" + dst_dir.mkdir() + + copy_hf_ckpt_remote_code(src_dir, dst_dir) # should not raise + + assert list(dst_dir.iterdir()) == [], "no files should be copied" + + +def test_copy_hf_ckpt_remote_code_hub_id(tmp_path, monkeypatch): + """copy_hf_ckpt_remote_code copies .py files from the resolved Hub snapshot.""" + dst_dir = tmp_path / "dst" + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "modeling_custom.py").write_text("# custom model") + (snapshot_dir / "not_python.txt").write_text("not python") + + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + with patch( + "modelopt.torch.utils.plugins.hf_checkpoint_utils.snapshot_download", + return_value=str(snapshot_dir), + ) as mock_sd: + copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-Nano-12B-v2", dst_dir) + + mock_sd.assert_called_once_with( + repo_id="nvidia/NVIDIA-Nemotron-Nano-12B-v2", + allow_patterns=["*.py"], + local_files_only=False, + ) + assert (dst_dir / "modeling_custom.py").read_text() == "# custom model" + assert not (dst_dir / "not_python.txt").exists(), "non-.py files should not be copied" + + +def test_copy_hf_ckpt_remote_code_hub_id_offline_uses_cache(tmp_path, monkeypatch): + """copy_hf_ckpt_remote_code resolves cached Hub snapshots when HF_HUB_OFFLINE is set.""" + dst_dir = tmp_path / "dst" + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "nemotron_reasoning_parser.py").write_text("# parser") + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with patch( + "modelopt.torch.utils.plugins.hf_checkpoint_utils.snapshot_download", + return_value=str(snapshot_dir), + ) as mock_sd: + copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", dst_dir) + + mock_sd.assert_called_once_with( + repo_id="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + allow_patterns=["*.py"], + local_files_only=True, + ) + assert (dst_dir / "nemotron_reasoning_parser.py").read_text() == "# parser" + + +def test_copy_hf_ckpt_remote_code_hub_id_offline_missing_cache_raises(tmp_path, monkeypatch): + """copy_hf_ckpt_remote_code raises a clear error when offline cache is missing.""" + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with ( + patch( + "modelopt.torch.utils.plugins.hf_checkpoint_utils.snapshot_download", + side_effect=LocalEntryNotFoundError("missing"), + ), + pytest.raises(RuntimeError, match="HF_HUB_OFFLINE"), + ): + copy_hf_ckpt_remote_code("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", tmp_path / "dst") + + +def test_sanitize_hf_config_for_deployment_trims_nextn_layer_types(): + """Drop MTP/next-token-prediction layer types from exported config.json.""" + hidden_layer_types = ["full_attention"] * 45 + nextn_layer_types = ["nextn_predict"] * 3 + config_data = { + "num_hidden_layers": 45, + "num_nextn_predict_layers": 3, + "layer_types": hidden_layer_types + nextn_layer_types, + } + + with pytest.warns(UserWarning, match="Trimming config.layer_types"): + sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace()) + + assert config_data["layer_types"] == hidden_layer_types + + +def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_parameters(): + """Transformers 5.x requires rope_theta inside llama3 rope_parameters.""" + config_data = { + "rope_theta": 500000, + "rope_parameters": { + "rope_type": "llama3", + "factor": 8.0, + "original_max_position_embeddings": 4096, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + }, + } + + sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) + + assert config_data["rope_parameters"]["rope_theta"] == 500000 + + +def test_sanitize_hf_config_for_deployment_uses_model_rope_theta_for_rope_parameters(): + """Use model.config.rope_theta when save_pretrained omits the top-level field.""" + config_data = { + "rope_parameters": { + "rope_type": "llama3", + "factor": 8.0, + "original_max_position_embeddings": 4096, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + }, + } + model = SimpleNamespace(config=SimpleNamespace(rope_theta=500000)) + + sanitize_hf_config_for_deployment(config_data, model) + + assert config_data["rope_parameters"]["rope_theta"] == 500000 + + +def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_scaling(): + """Legacy rope_scaling metadata is normalized for llama3 configs as well.""" + config_data = { + "rope_scaling": { + "type": "llama3", + "factor": 8.0, + "original_max_position_embeddings": 4096, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + }, + } + model = SimpleNamespace(config=SimpleNamespace(rope_theta=500000)) + + sanitize_hf_config_for_deployment(config_data, model) + + assert config_data["rope_scaling"]["rope_theta"] == 500000 + + +def test_sanitize_hf_config_for_deployment_keeps_existing_rope_theta(): + """Existing rope_theta in rope metadata is not overwritten.""" + config_data = { + "rope_theta": 500000, + "rope_parameters": { + "rope_type": "llama3", + "rope_theta": 1000000, + }, + } + + sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) + + assert config_data["rope_parameters"]["rope_theta"] == 1000000 + + +def test_sanitize_hf_config_for_deployment_ignores_non_llama3_rope_parameters(): + """Only llama3 RoPE parameters need the Transformers 5.x compatibility fix.""" + config_data = { + "rope_theta": 500000, + "rope_parameters": { + "rope_type": "default", + }, + } + + sanitize_hf_config_for_deployment(config_data, SimpleNamespace(config=SimpleNamespace())) + + assert "rope_theta" not in config_data["rope_parameters"] + + +def test_sanitize_hf_config_for_deployment_uses_model_config_nextn_count(): + """Handle exports where save_pretrained omits num_nextn_predict_layers.""" + config_data = { + "num_hidden_layers": 2, + "layer_types": ["full_attention", "linear_attention", "nextn_predict"], + } + model = SimpleNamespace(config=SimpleNamespace(num_nextn_predict_layers=1)) + + with pytest.warns(UserWarning, match="Trimming config.layer_types"): + sanitize_hf_config_for_deployment(config_data, model=model) + + assert config_data["layer_types"] == ["full_attention", "linear_attention"] + + +def test_sanitize_hf_config_for_deployment_keeps_unexplained_layer_type_mismatch(): + """Do not rewrite config when extra layer types are not explained by nextn metadata.""" + config_data = { + "num_hidden_layers": 2, + "num_nextn_predict_layers": 1, + "layer_types": ["full_attention", "linear_attention", "extra_a", "extra_b"], + } + + sanitize_hf_config_for_deployment(config_data, model=SimpleNamespace()) + + assert config_data["layer_types"] == [ + "full_attention", + "linear_attention", + "extra_a", + "extra_b", + ] + + +# --- off-index safetensors: files model loading never opens -------------------------------------- + + +def _shard(path, name): + (path / name).write_text("tensors") + + +def test_no_index_means_only_model_safetensors_is_read(tmp_path): + _shard(tmp_path, "model.safetensors") + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == [] + + +def test_a_standalone_sidecar_is_off_index(tmp_path): + """GLM-4.7 ships its MTP head as mtp.safetensors, which the loader never opens.""" + _shard(tmp_path, "model.safetensors") + _shard(tmp_path, "mtp.safetensors") + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == ["mtp.safetensors"] + + +def test_indexed_shards_are_read_and_sidecars_are_not(tmp_path): + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a": "model-00001-of-00002.safetensors",' + ' "b": "model-00002-of-00002.safetensors"}}' + ) + for name in ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"): + _shard(tmp_path, name) + _shard(tmp_path, "mtp.safetensors") + + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == ["mtp.safetensors"] + + +@pytest.mark.parametrize( + "index", + [ + '{"weight_map": {}}', + "{}", + '{"weight_map": {"a": "model-00001-of-00002.safetensors"}}', + ], + ids=["empty-map", "no-map-key", "partial-map"], +) +def test_main_weight_shards_are_never_off_index_whatever_the_index_says(tmp_path, index): + """An empty, partial or malformed index must not make the real weights look like sidecars. + + Copying those into an export would leave the unquantized source weights sitting beside the + quantized ones -- a checkpoint that loads and is silently wrong. + """ + (tmp_path / "model.safetensors.index.json").write_text(index) + for name in ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"): + _shard(tmp_path, name) + + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == [] + + +def test_unsharded_main_weights_are_never_off_index(tmp_path): + (tmp_path / "model.safetensors.index.json").write_text('{"weight_map": {}}') + _shard(tmp_path, "model.safetensors") + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == [] + + +def test_results_are_sorted(tmp_path): + _shard(tmp_path, "model.safetensors") + for name in ("zeta.safetensors", "alpha.safetensors", "mtp.safetensors"): + _shard(tmp_path, name) + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path) == [ + "alpha.safetensors", + "mtp.safetensors", + "zeta.safetensors", + ] + + +def test_a_missing_directory_is_not_an_error(tmp_path): + assert hf_checkpoint_utils.off_index_safetensors_files(tmp_path / "nope") == [] + + +def test_copy_moves_only_the_sidecars_and_preserves_bytes(tmp_path): + src, dst = tmp_path / "src", tmp_path / "dst" + src.mkdir() + dst.mkdir() + (src / "model.safetensors.index.json").write_text( + '{"weight_map": {"a": "model-00001-of-00001.safetensors"}}' + ) + _shard(src, "model-00001-of-00001.safetensors") + (src / "mtp.safetensors").write_text("mtp-bytes") + + assert hf_checkpoint_utils.copy_off_index_safetensors(src, dst) == ["mtp.safetensors"] + assert (dst / "mtp.safetensors").read_text() == "mtp-bytes" + # the real weights are the export's job, not a verbatim copy + assert not (dst / "model-00001-of-00001.safetensors").exists() + + +def test_copy_does_not_overwrite_what_the_export_already_wrote(tmp_path): + src, dst = tmp_path / "src", tmp_path / "dst" + src.mkdir() + dst.mkdir() + _shard(src, "model.safetensors") + (src / "mtp.safetensors").write_text("source") + (dst / "mtp.safetensors").write_text("already exported") + + assert hf_checkpoint_utils.copy_off_index_safetensors(src, dst) == [] + assert (dst / "mtp.safetensors").read_text() == "already exported" + + +def _write_st(path, tensors): + """Minimal real safetensors file so header reads work.""" + save_file({k: torch.zeros(1) for k in tensors}, str(path)) + + +def test_off_index_skips_mistral_consolidated_copy(tmp_path): + """Mistral ships consolidated.safetensors: a SECOND full copy of the indexed weights. + + Copying it into an export puts unquantized weights beside the quantized ones, and vLLM's + mistral load-format looks for that filename specifically -- so it can be served instead of + what we quantized. + """ + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + _write_st(tmp_path / "model-00001-of-00001.safetensors", ["a.weight"]) + _write_st(tmp_path / "consolidated.safetensors", ["a.weight"]) + + assert off_index_safetensors_files(tmp_path) == [] + + +def test_off_index_skips_peft_adapter(tmp_path): + """A PEFT adapter's tensor names do NOT overlap the index, so only the name rule catches it.""" + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + _write_st(tmp_path / "model-00001-of-00001.safetensors", ["a.weight"]) + _write_st(tmp_path / "adapter_model.safetensors", ["base_model.a.lora_A.weight"]) + + assert off_index_safetensors_files(tmp_path) == [] + + +def test_off_index_skips_unknown_name_that_reships_indexed_weights(tmp_path): + """The name rules only know the conventions we have seen; overlap catches the rest.""" + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors",' + ' "b.weight": "model-00001-of-00001.safetensors"}}' + ) + _write_st(tmp_path / "model-00001-of-00001.safetensors", ["a.weight", "b.weight"]) + _write_st(tmp_path / "backup-copy.safetensors", ["a.weight", "b.weight"]) + + assert off_index_safetensors_files(tmp_path) == [] + + +def test_off_index_still_keeps_a_genuine_mtp_sidecar(tmp_path): + """The whole point: a real sidecar holds names the index does NOT have, and must be kept.""" + + (tmp_path / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + _write_st(tmp_path / "model-00001-of-00001.safetensors", ["a.weight"]) + _write_st(tmp_path / "mtp.safetensors", ["model.mtp.eh_proj.weight"]) + + assert off_index_safetensors_files(tmp_path) == ["mtp.safetensors"] + + +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX symlink support") +def test_copies_a_symlinked_sidecar_from_a_hub_cache_layout(tmp_path): + """A hub-downloaded checkpoint stores EVERY file as a symlink into ``../../blobs/``. + + Rejecting symlinks outright therefore skips the sidecar of every checkpoint loaded by hub id + -- including the GLM-4.7 ``mtp.safetensors`` this path exists to carry -- which is the + silent-missing-MTP failure the carry-over was written to prevent. The guard must look at what + the link resolves to, not at whether it is a link. + """ + blobs = tmp_path / "blobs" + snapshot = tmp_path / "snapshots" / "deadbeef" + blobs.mkdir(parents=True) + snapshot.mkdir(parents=True) + + save_file({"model.mtp.eh_proj.weight": torch.zeros(1)}, str(blobs / "sha123")) + save_file({"a.weight": torch.zeros(1)}, str(blobs / "sha456")) + (snapshot / "mtp.safetensors").symlink_to("../../blobs/sha123") + (snapshot / "model-00001-of-00001.safetensors").symlink_to("../../blobs/sha456") + (snapshot / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + + assert off_index_safetensors_files(snapshot) == ["mtp.safetensors"] + + dst = tmp_path / "export" + dst.mkdir() + assert copy_off_index_safetensors(snapshot, dst) == ["mtp.safetensors"] + assert (dst / "mtp.safetensors").is_file() + + +def _indexed_ckpt(src): + src.mkdir(parents=True, exist_ok=True) + save_file({"a.weight": torch.zeros(1)}, str(src / "model-00001-of-00001.safetensors")) + (src / "model.safetensors.index.json").write_text( + '{"weight_map": {"a.weight": "model-00001-of-00001.safetensors"}}' + ) + return src + + +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX symlink support") +def test_skips_a_sidecar_whose_link_dangles(tmp_path): + """A link to nothing copies nothing rather than raising out of the export.""" + src = _indexed_ckpt(tmp_path / "ckpt") + (src / "mtp.safetensors").symlink_to(tmp_path / "does-not-exist") + + dst = tmp_path / "export" + dst.mkdir() + with pytest.warns(UserWarning, match="not a readable regular file"): + assert copy_off_index_safetensors(src, dst) == [] + assert not (dst / "mtp.safetensors").exists() + + +@pytest.mark.skipif(sys.platform == "win32", reason="requires POSIX symlink support") +def test_skips_a_sidecar_pointing_outside_the_checkpoint(tmp_path): + """The original hardening, restored: a link out of the tree is refused, not followed. + + Accepting hub blobs must not mean accepting any target at all -- a checkpoint shipping + ``mtp.safetensors -> /etc/passwd`` would otherwise land that file in the export under a name + that looks like model weights. + """ + outside = tmp_path / "secret.txt" + outside.write_text("not model weights") + src = _indexed_ckpt(tmp_path / "ckpt") + (src / "mtp.safetensors").symlink_to(outside) + + dst = tmp_path / "export" + dst.mkdir() + with pytest.warns(UserWarning, match="outside the checkpoint directory"): + assert copy_off_index_safetensors(src, dst) == [] + assert not (dst / "mtp.safetensors").exists() + + +# --- indexed_weight_map / read_safetensors_subset (moved from model_load_utils.py, which used +# to duplicate indexed_weight_map's own index/single-file logic) -------------------------------- + + +def test_indexed_weight_map_sharded(tmp_path): + save_file({"a.weight": torch.zeros(2)}, str(tmp_path / "shard1.safetensors")) + save_file({"b.weight": torch.zeros(2)}, str(tmp_path / "shard2.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": {"a.weight": "shard1.safetensors", "b.weight": "shard2.safetensors"}} + ) + ) + + assert indexed_weight_map(str(tmp_path)) == { + "a.weight": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + +def test_indexed_weight_map_single_file(tmp_path): + save_file( + {"a.weight": torch.zeros(2), "b.weight": torch.zeros(2)}, + str(tmp_path / "model.safetensors"), + ) + + assert indexed_weight_map(str(tmp_path)) == { + "a.weight": "model.safetensors", + "b.weight": "model.safetensors", + } + + +def test_indexed_weight_map_missing_returns_empty(tmp_path): + """Neither an index nor a single-file checkpoint: {}, not an exception. + + Right for indexed_weight_map's own callers (e.g. locate_source_keys), which treat "nothing + recorded" as legitimate. Callers for whom a missing checkpoint is a genuine error (FSDP2 + parallel loading, the structural unplaced-keys fallback, DFlash's precision reload) check for + the empty result and raise themselves. + """ + assert indexed_weight_map(str(tmp_path)) == {} + + +def test_read_safetensors_subset(tmp_path): + save_file( + {"a.weight": torch.tensor([1.0, 2.0]), "a.bias": torch.tensor([3.0])}, + str(tmp_path / "shard1.safetensors"), + ) + save_file({"b.weight": torch.tensor([4.0])}, str(tmp_path / "shard2.safetensors")) + weight_map = { + "a.weight": "shard1.safetensors", + "a.bias": "shard1.safetensors", + "b.weight": "shard2.safetensors", + } + + result = read_safetensors_subset(str(tmp_path), weight_map, lambda n: n.startswith("a.")) + + assert set(result.keys()) == {"a.weight", "a.bias"} + assert torch.equal(result["a.weight"], torch.tensor([1.0, 2.0])) + assert torch.equal(result["a.bias"], torch.tensor([3.0])) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index 323fb92a568..428caba0617 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -15,12 +15,9 @@ """Pure-function tests for ``modelopt.torch.utils.plugins.model_load_utils``.""" -import json - import pytest import torch from packaging.version import Version -from safetensors.torch import save_file pytest.importorskip("accelerate") @@ -28,62 +25,10 @@ _conversion_plan, _convert_keys, _resolve_target, - read_safetensors_subset, - weight_map_for, + record_unplaced_source_keys, ) -def test_weight_map_for_sharded(tmp_path): - save_file({"a.weight": torch.zeros(2)}, str(tmp_path / "shard1.safetensors")) - save_file({"b.weight": torch.zeros(2)}, str(tmp_path / "shard2.safetensors")) - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps( - {"weight_map": {"a.weight": "shard1.safetensors", "b.weight": "shard2.safetensors"}} - ) - ) - - assert weight_map_for(str(tmp_path)) == { - "a.weight": "shard1.safetensors", - "b.weight": "shard2.safetensors", - } - - -def test_weight_map_for_single_file(tmp_path): - save_file( - {"a.weight": torch.zeros(2), "b.weight": torch.zeros(2)}, - str(tmp_path / "model.safetensors"), - ) - - assert weight_map_for(str(tmp_path)) == { - "a.weight": "model.safetensors", - "b.weight": "model.safetensors", - } - - -def test_weight_map_for_missing(tmp_path): - with pytest.raises(RuntimeError, match="No safetensors checkpoint"): - weight_map_for(str(tmp_path)) - - -def test_read_safetensors_subset(tmp_path): - save_file( - {"a.weight": torch.tensor([1.0, 2.0]), "a.bias": torch.tensor([3.0])}, - str(tmp_path / "shard1.safetensors"), - ) - save_file({"b.weight": torch.tensor([4.0])}, str(tmp_path / "shard2.safetensors")) - weight_map = { - "a.weight": "shard1.safetensors", - "a.bias": "shard1.safetensors", - "b.weight": "shard2.safetensors", - } - - result = read_safetensors_subset(str(tmp_path), weight_map, lambda n: n.startswith("a.")) - - assert set(result.keys()) == {"a.weight", "a.bias"} - assert torch.equal(result["a.weight"], torch.tensor([1.0, 2.0])) - assert torch.equal(result["a.bias"], torch.tensor([3.0])) - - def _build_tiny_qwen3_moe(): """A tiny meta-init Qwen3-MoE (fused ``gate_up_proj`` + ``down_proj`` experts) for converter tests.""" transformers = pytest.importorskip("transformers") @@ -154,3 +99,32 @@ def test_checkpoint_key_converter_multisource_expert_fusion(): for e in range(n_exp): assert _resolve_target(plan, f"{prefix}mlp.experts.{e}.gate_proj.weight")[0] == gname assert _resolve_target(plan, f"{prefix}mlp.experts.{e}.up_proj.weight")[0] == gname + + +# --- the loader's own accounting of what it could not place -------------------------------------- + + +def test_record_unplaced_source_keys_stores_sorted_keys_and_provenance(): + model = torch.nn.Linear(2, 2) + keys = record_unplaced_source_keys(model, "/ckpt/path", ["z.weight", "a.weight"]) + + assert keys == ["a.weight", "z.weight"] + assert model._modelopt_unplaced_source_keys == ["a.weight", "z.weight"] + assert model._modelopt_source_checkpoint == "/ckpt/path" + + +def test_record_unplaced_source_keys_distinguishes_none_from_empty(): + """An empty list is an ANSWER -- the loader placed everything -- not "nobody asked". + + The export re-derives the set only when the attribute is absent, so recording [] has to stick; + treating it as falsy-and-therefore-unknown would make every clean load pay a re-derivation. + """ + model = torch.nn.Linear(2, 2) + assert record_unplaced_source_keys(model, "/ckpt/path", None) == [] + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == "/ckpt/path" + + +def test_record_unplaced_source_keys_accepts_any_iterable(): + model = torch.nn.Linear(2, 2) + assert record_unplaced_source_keys(model, "/c", iter(["b", "a"])) == ["a", "b"]