Preserve HF PTQ checkpoint sidecar files [NV BUG 6491822]

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:

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

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"]