From da68df5c5324214a69d62ca5b0ae4c1b61a72f65 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Tue, 4 Aug 2026 19:08:15 -0700 Subject: [PATCH 1/5] 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"] From cec2df98296636d0bbe1cb5b31f69e4403f05be4 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 5 Aug 2026 10:33:29 -0700 Subject: [PATCH 2/5] Preserve export generation config Signed-off-by: Jennifer Chen --- examples/hf_ptq/example_utils.py | 41 +++++++++++++++++-- examples/hf_ptq/hf_ptq.py | 13 ++++-- .../export/plugins/hf_checkpoint_utils.py | 38 ++++------------- tests/examples/hf_ptq/test_example_utils.py | 8 ++++ .../torch/export/test_hf_checkpoint_utils.py | 31 +++++++++++++- 5 files changed, 92 insertions(+), 39 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 5e1976310a8..d05339f7e9c 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -67,6 +67,31 @@ "LICENSE*", "NOTICE*", ] +_HF_PTQ_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", +) _HF_PTQ_EXPORT_OWNED_FILES = { "config.json", "hf_quant_config.json", @@ -957,7 +982,12 @@ def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False return model_name_or_path -def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False): +def copy_custom_model_files( + source_path: str, + export_path: str, + trust_remote_code: bool = False, + exclude_files: Iterable[str] | None = None, +): """Copy source checkpoint sidecar files to an HF PTQ export. The HF PTQ script writes ModelOpt-owned metadata and quantized weights first, then @@ -966,14 +996,16 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod 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. + Source tokenizer and processor files intentionally still win because Transformers may + not regenerate all metadata in the source format. Callers that write a generation config + can exclude it; the TensorRT-LLM export retains the source generation config. 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 when resolving HuggingFace model IDs + exclude_files: Additional source file names to skip. """ # Resolve the source path (handles both local paths and HF model IDs) resolved_source_path = _resolve_model_path(source_path, trust_remote_code) @@ -998,7 +1030,8 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod copied_files = copy_non_safetensor_files_from_ckpt( source_dir, export_dir, - exclude_files=_HF_PTQ_EXPORT_OWNED_FILES, + exclude_files=_HF_PTQ_EXPORT_OWNED_FILES | set(exclude_files or ()), + exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, ) if copied_files: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6a4bd476984..839589bf410 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -856,11 +856,12 @@ def export_quantized( print("This is normal for some VLM architectures that don't use AutoProcessor") start_time = time.time() - if ( + is_tensorrt_llm_export = ( model_type in ["t5", "bart", "whisper"] or args.sparsity_fmt != "dense" or "int8_sq" in args.qformat - ): + ) + if is_tensorrt_llm_export: if ( args.inference_tensor_parallel != 1 or args.inference_pipeline_parallel != 1 ) and args.qformat == "nvfp4_svdquant": @@ -932,7 +933,13 @@ def export_quantized( # from the source checkpoint take precedence over regenerated ones (which may # differ in format due to newer transformers versions). if args.dist_state.is_main: - copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code) + exclude_files = None if is_tensorrt_llm_export else {"generation_config.json"} + copy_custom_model_files( + args.pyt_ckpt_path, + export_path, + args.trust_remote_code, + exclude_files=exclude_files, + ) end_time = time.time() print_rank_0( diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index 95250384943..e72188beead 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -31,31 +31,6 @@ 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: @@ -289,17 +264,19 @@ def copy_non_safetensor_files_from_ckpt( dst: str | os.PathLike, *, exclude_files: Iterable[str] | None = None, + exclude_patterns: Iterable[str] | None = None, ) -> list[str]: - """Copy every non-weight sidecar file from a local HF checkpoint dir verbatim. + """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 pass files through ``exclude_files`` when + 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 in addition to weights and weight indexes. + exclude_files: Exact file names to skip. + exclude_patterns: Glob patterns for additional files to skip. Returns: File names copied into ``dst``. @@ -307,12 +284,11 @@ def copy_non_safetensor_files_from_ckpt( 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, _HF_CHECKPOINT_WEIGHT_FILE_PATTERNS - ): + 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): diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 8c1a88fc690..4d7570a389c 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -99,6 +99,14 @@ def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): assert not (export_dir / "model-00001-of-00001.safetensors").exists() assert not (export_dir / "model.gguf").exists() + (export_dir / "generation_config.json").write_text('{"export": "generation"}\n') + example_utils.copy_custom_model_files( + str(source_dir), + str(export_dir), + exclude_files={"generation_config.json"}, + ) + assert (export_dir / "generation_config.json").read_text() == '{"export": "generation"}\n' + def test_resolve_model_path_snapshot_download_stays_allowlisted(monkeypatch, tmp_path): snapshot_dir = tmp_path / "snapshot" diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index f3be2564312..df4e5e8714f 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -24,7 +24,36 @@ hf_hub_errors = pytest.importorskip("huggingface_hub.errors") LocalEntryNotFoundError = hf_hub_errors.LocalEntryNotFoundError -from modelopt.torch.export import copy_hf_ckpt_remote_code, sanitize_hf_config_for_deployment +from modelopt.torch.export 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 / "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 (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_hf_ckpt_remote_code_local_dir(tmp_path): From e18d5e4c1185649fedfe791deb9d3c80bdaeecd7 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 5 Aug 2026 11:00:19 -0700 Subject: [PATCH 3/5] Restore safetensor sidecar exclusion Signed-off-by: Jennifer Chen --- modelopt/torch/export/plugins/hf_checkpoint_utils.py | 2 ++ tests/unit/torch/export/test_hf_checkpoint_utils.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index e72188beead..9373673d83a 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -293,6 +293,8 @@ def copy_non_safetensor_files_from_ckpt( 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/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index df4e5e8714f..01fe33c64f8 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -35,6 +35,7 @@ def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_ 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") @@ -42,6 +43,7 @@ def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_ 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() From 874a57c82ee64ade0e304dc427bf4d10b5e9cb36 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 5 Aug 2026 13:11:45 -0700 Subject: [PATCH 4/5] Continue sidecar copy after errors Signed-off-by: Jennifer Chen --- .../export/plugins/hf_checkpoint_utils.py | 6 +++++- .../torch/export/test_hf_checkpoint_utils.py | 21 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/plugins/hf_checkpoint_utils.py b/modelopt/torch/export/plugins/hf_checkpoint_utils.py index 9373673d83a..54058c742db 100644 --- a/modelopt/torch/export/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/export/plugins/hf_checkpoint_utils.py @@ -295,6 +295,10 @@ def copy_non_safetensor_files_from_ckpt( continue if entry.endswith(".safetensors") or entry == "model.safetensors.index.json": continue - shutil.copy2(sp, dst) + 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/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index 01fe33c64f8..08292c65f9f 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -29,6 +29,7 @@ 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): @@ -58,6 +59,26 @@ def test_copy_non_safetensor_files_from_ckpt_supports_additional_exclusions(tmp_ 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" From d669f9c1944f8429d2b1c3321646d0fe1eda28aa Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 5 Aug 2026 13:13:10 -0700 Subject: [PATCH 5/5] Preserve exported tokenizer config Signed-off-by: Jennifer Chen --- examples/hf_ptq/example_utils.py | 14 +++++++++----- tests/examples/hf_ptq/test_example_utils.py | 4 ++++ 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index d05339f7e9c..71336993779 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -997,14 +997,14 @@ def copy_custom_model_files( 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 and processor files intentionally still win because Transformers may - not regenerate all metadata in the source format. Callers that write a generation config - can exclude it; the TensorRT-LLM export retains the source generation config. + not regenerate all metadata in the source format. The exported ``tokenizer_config.json`` + wins when it has a separate chat template. Callers that write a generation config can + exclude it; the TensorRT-LLM export retains the source generation config. 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 when resolving HuggingFace model - IDs + trust_remote_code: Passed to HuggingFace model-ID resolution; does not control copying. exclude_files: Additional source file names to skip. """ # Resolve the source path (handles both local paths and HF model IDs) @@ -1027,10 +1027,14 @@ def copy_custom_model_files( print(f"Warning: Export directory {export_path} does not exist") return + exclude_files = _HF_PTQ_EXPORT_OWNED_FILES | set(exclude_files or ()) + if (export_dir / "chat_template.jinja").is_file(): + exclude_files.add("tokenizer_config.json") + copied_files = copy_non_safetensor_files_from_ckpt( source_dir, export_dir, - exclude_files=_HF_PTQ_EXPORT_OWNED_FILES | set(exclude_files or ()), + exclude_files=exclude_files, exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, ) diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 4d7570a389c..4e39f150f0a 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -61,6 +61,7 @@ def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): "README.md": "# Source model\n", "LICENSE": "license text\n", "chat_template.jinja": "{{ messages }}\n", + "tokenizer_config.json": '{"chat_template": "source"}\n', "generation_config.json": '{"source": "generation"}\n', "config.json": '{"source": "config"}\n', "hf_quant_config.json": '{"source": "quant"}\n', @@ -77,6 +78,8 @@ def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): (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') + (export_dir / "chat_template.jinja").write_text("{{ exported_messages }}\n") + (export_dir / "tokenizer_config.json").write_text('{"chat_template": "export"}\n') example_utils.copy_custom_model_files(str(source_dir), str(export_dir), trust_remote_code=False) @@ -92,6 +95,7 @@ def test_copy_custom_model_files_preserves_non_weight_sidecars(tmp_path): assert (export_dir / "config.json").read_text() == '{"export": "config"}\n' assert (export_dir / "hf_quant_config.json").read_text() == '{"export": "quant"}\n' + assert (export_dir / "tokenizer_config.json").read_text() == '{"chat_template": "export"}\n' assert not (export_dir / "quant_config.json").exists() assert not (export_dir / "quantize_config.json").exists() assert not (export_dir / "recipe.yaml").exists()