From da68df5c5324214a69d62ca5b0ae4c1b61a72f65 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 4 Aug 2026 19:08:15 -0700 Subject: [PATCH] Preserve HF PTQ checkpoint sidecars Signed-off-by: Jennifer Chen --- examples/hf_ptq/example_utils.py | 107 ++++++++---------- .../export/plugins/hf_checkpoint_utils.py | 61 ++++++++-- tests/examples/hf_ptq/test_example_utils.py | 77 +++++++++++++ 3 files changed, 177 insertions(+), 68 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d74ffb34efb..5e1976310a8 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -20,7 +20,6 @@ import json import logging import os -import shutil import warnings from collections.abc import Callable, Iterable from dataclasses import dataclass @@ -44,6 +43,7 @@ ) 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 try: from huggingface_hub import snapshot_download @@ -56,6 +56,27 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS = [ + "*.jinja", + "*.json", + "*.md", + "*.model", + "*.py", + "*.tiktoken", + "*.txt", + "LICENSE*", + "NOTICE*", +] +_HF_PTQ_EXPORT_OWNED_FILES = { + "config.json", + "hf_quant_config.json", + "quant_config.json", + "quantization_config.json", + "quantize_config.json", + "recipe.yaml", + "recipe.yml", +} + @dataclass class DistributedState: @@ -895,11 +916,13 @@ def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False try: local_path = snapshot_download( repo_id=model_name_or_path, - allow_patterns=["*.py", "*.json"], # Only download Python files and config + allow_patterns=_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, ) return local_path except Exception as e: - print(f"Warning: Could not download model files using snapshot_download: {e}") + print( + f"Warning: Could not download checkpoint sidecars using snapshot_download: {e}" + ) # Fallback: try to find in HuggingFace cache from transformers.utils import TRANSFORMERS_CACHE @@ -935,48 +958,23 @@ def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): - """Copy processor/tokenizer artifacts (and, with trust_remote_code, custom code) to export. - - Processor and tokenizer *data* artifacts -- e.g. a VLM's ``preprocessor_config.json``, - ``merges.txt``/``vocab.json``, and the processor helper modules -- are needed by the - deployment stack (vLLM/SGLang) even when the model itself runs on native (non-remote) - transformers code. transformers 5.x restructured many VLM configs and no longer - re-saves these on ``save_pretrained`` for models loaded natively, so without copying - them a native-path export is missing e.g. ``preprocessor_config.json`` and fails to - load (``Can't load image processor``). These are copied regardless of - ``trust_remote_code``. Executable model/config code (``modeling*.py``, - ``configuration_*.py``, ``tokenization_*.py``, and other custom JSON) is only meaningful - with ``trust_remote_code`` and is copied only then. ``config.json`` and - ``model.safetensors.index.json`` are always skipped (handled by the export itself). + """Copy source checkpoint sidecar files to an HF PTQ export. + + The HF PTQ script writes ModelOpt-owned metadata and quantized weights first, then + copies source checkpoint sidecars so tokenizer/processor files, remote-code modules, + README assets, parser plugins, and similar deployment files are preserved for both + native and ``trust_remote_code`` loads. Weight and weight-index files are skipped + to avoid copying the unquantized source weights. Export-owned metadata (``config.json``, + ``hf_quant_config.json``) and stale source quantization metadata are also skipped. + Source tokenizer, processor, and generation files intentionally still win because + Transformers may not regenerate all metadata in the source format. Args: source_path: Path to the original model directory or HuggingFace model ID export_path: Path to the exported model directory - trust_remote_code: Whether trust_remote_code was used (gates the executable code files) + trust_remote_code: Whether trust_remote_code was used when resolving HuggingFace model + IDs """ - # Deployment-critical processor/tokenizer artifacts: safe to copy regardless of - # trust_remote_code (data + processor helpers, not model code). - always_copy_patterns = [ - "preprocessor_config.json", - "processor_config.json", - "image_processing*.py", - "processing_*.py", - "video_processing*.py", - "feature_extraction_*.py", - "added_tokens.json", - "special_tokens_map.json", - "vocab.json", - "merges.txt", - "tokenizer.model", - ] - # Executable custom model/config code + other custom JSON: only used with trust_remote_code. - code_patterns = [ - "configuration_*.py", - "modeling*.py", - "tokenization_*.py", - "*.json", - ] - # Resolve the source path (handles both local paths and HF model IDs) resolved_source_path = _resolve_model_path(source_path, trust_remote_code) @@ -997,29 +995,18 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod print(f"Warning: Export directory {export_path} does not exist") return - patterns = [*always_copy_patterns, *(code_patterns if trust_remote_code else [])] - - copied_files: list[str] = [] - for pattern in patterns: - for file_path in source_dir.glob(pattern): - if file_path.is_file(): - # Skip config.json and model.safetensors.index.json as they're handled separately - if file_path.name in ["config.json", "model.safetensors.index.json"]: - continue - if file_path.name in copied_files: # e.g. matched by both pattern lists - continue - dest_path = export_dir / file_path.name - try: - shutil.copy2(file_path, dest_path) - copied_files.append(file_path.name) - print(f"Copied custom model file: {file_path.name}") - except Exception as e: - print(f"Warning: Failed to copy {file_path.name}: {e}") + copied_files = copy_non_safetensor_files_from_ckpt( + source_dir, + export_dir, + exclude_files=_HF_PTQ_EXPORT_OWNED_FILES, + ) if copied_files: - print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") + for file_name in copied_files: + print(f"Copied checkpoint sidecar file: {file_name}") + print(f"Successfully copied {len(copied_files)} checkpoint sidecar files to {export_path}") else: - print("No custom model files found to copy") + print("No checkpoint sidecar files found to copy") def _layerwise_checkpoint_dir_location(algorithm) -> tuple[str, str] | None: diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index ddae8e2b409..95250384943 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -15,10 +15,12 @@ """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 @@ -29,6 +31,31 @@ from tqdm import tqdm _HF_HUB_OFFLINE_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +_HF_CHECKPOINT_WEIGHT_FILE_PATTERNS = ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + "*.ckpt", + "*.gguf", + "*.h5", + "*.msgpack", + "*.npy", + "*.npz", + "*.onnx", + "*.pb", + "*.pickle", + "*.pkl", + "*.pt", + "*.pth", + "*.tar", + "*.tar.bz2", + "*.tar.gz", + "*.tar.xz", + "*.tflite", + "*.tgz", + "*.zip", +) def _as_nonnegative_int(value: Any) -> int | None: @@ -253,25 +280,43 @@ def load_multimodal_components( return multimodal_state_dict -def copy_non_safetensor_files_from_ckpt(src: str | os.PathLike, dst: str | os.PathLike): - """Copy every non-safetensors file from a local HF checkpoint dir verbatim. +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, +) -> list[str]: + """Copy every non-weight sidecar 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. The caller is expected to overwrite the files - modelopt owns (``config.json``, ``generation_config.json``, ``hf_quant_config.json``, - ``preprocessor_config.json``) after this step. + are preserved from the source. Callers can pass files through ``exclude_files`` 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 in addition to weights and weight indexes. + + 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 ()) + copied_files = [] os.makedirs(dst, exist_ok=True) - for entry in os.listdir(src): + for entry in sorted(os.listdir(src)): + if entry in exclude_files or _matches_any_pattern( + entry, _HF_CHECKPOINT_WEIGHT_FILE_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 shutil.copy2(sp, dst) + copied_files.append(entry) + return copied_files diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 00621ec6125..8c1a88fc690 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -49,6 +49,83 @@ def _write_safetensors(path, tensors): save_file(tensors, str(path), metadata={"format": "pt"}) +def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): + source_dir = tmp_path / "source" + export_dir = tmp_path / "export" + source_dir.mkdir() + export_dir.mkdir() + + source_files = { + "super_v3_reasoning_parser.py": "class Parser: pass\n", + "modeling_custom.py": "class Model: pass\n", + "README.md": "# Source model\n", + "LICENSE": "license text\n", + "chat_template.jinja": "{{ messages }}\n", + "generation_config.json": '{"source": "generation"}\n', + "config.json": '{"source": "config"}\n', + "hf_quant_config.json": '{"source": "quant"}\n', + "quant_config.json": '{"source": "stale quant"}\n', + "quantize_config.json": '{"source": "stale quant"}\n', + "recipe.yaml": "quantize: {}\n", + "model.safetensors.index.json": '{"weight_map": {}}\n', + "model-00001-of-00001.safetensors": "source weights\n", + "model.gguf": "source weights\n", + } + for file_name, contents in source_files.items(): + (source_dir / file_name).write_text(contents) + + (export_dir / "config.json").write_text('{"export": "config"}\n') + (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') + (export_dir / "hf_quant_config.json").write_text('{"export": "quant"}\n') + + example_utils.copy_custom_model_files(str(source_dir), str(export_dir), trust_remote_code=False) + + for file_name in [ + "super_v3_reasoning_parser.py", + "modeling_custom.py", + "README.md", + "LICENSE", + "chat_template.jinja", + "generation_config.json", + ]: + assert (export_dir / file_name).read_text() == source_files[file_name] + + assert (export_dir / "config.json").read_text() == '{"export": "config"}\n' + assert (export_dir / "hf_quant_config.json").read_text() == '{"export": "quant"}\n' + assert not (export_dir / "quant_config.json").exists() + assert not (export_dir / "quantize_config.json").exists() + assert not (export_dir / "recipe.yaml").exists() + assert not (export_dir / "model.safetensors.index.json").exists() + assert not (export_dir / "model-00001-of-00001.safetensors").exists() + assert not (export_dir / "model.gguf").exists() + + +def test_resolve_model_path_snapshot_download_stays_allowlisted(monkeypatch, tmp_path): + snapshot_dir = tmp_path / "snapshot" + + def fake_snapshot_download(**kwargs): + assert kwargs == { + "repo_id": "org/model", + "allow_patterns": example_utils._HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS, + } + return str(snapshot_dir) + + def fake_from_pretrained(*args, **kwargs): + assert (args, kwargs) == (("org/model",), {"trust_remote_code": False}) + return SimpleNamespace(_name_or_path="org/model") + + monkeypatch.setattr( + example_utils.AutoConfig, + "from_pretrained", + fake_from_pretrained, + ) + monkeypatch.setattr(example_utils, "snapshot_download", fake_snapshot_download) + + assert example_utils._resolve_model_path("org/model", trust_remote_code=False) == str( + snapshot_dir + ) + + 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"]