diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 93bba2af51d..57e01a13939 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -21,7 +21,6 @@ import json import logging import os -import shutil import warnings from collections.abc import Callable, Iterable from contextlib import AbstractContextManager, nullcontext @@ -48,6 +47,7 @@ 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 try: from huggingface_hub import snapshot_download @@ -65,6 +65,52 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS = [ + "*.jinja", + "*.json", + "*.md", + "*.model", + "*.py", + "*.tiktoken", + "*.txt", + "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", + "quant_config.json", + "quantization_config.json", + "quantize_config.json", + "recipe.yaml", + "recipe.yml", +} + @dataclass class DistributedState: @@ -931,11 +977,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 @@ -970,49 +1018,31 @@ 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): - """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). +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 + 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 and processor files intentionally still win because Transformers may + 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 (gates the executable code files) + trust_remote_code: Passed to HuggingFace model-ID resolution; does not control copying. + exclude_files: Additional source file names to skip. """ - # 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) @@ -1033,29 +1063,23 @@ 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}") + 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=exclude_files, + exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS, + ) 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/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 6ddfba6614c..471805c8b22 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -861,11 +861,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": @@ -937,7 +938,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 ddae8e2b409..54058c742db 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 @@ -253,25 +255,50 @@ def load_multimodal_components( return multimodal_state_dict -def copy_non_safetensor_files_from_ckpt(src: str | os.PathLike, dst: str | os.PathLike): +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. 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 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 os.listdir(src): + 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 - 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/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 57b9676ee92..836b965f870 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -49,6 +49,95 @@ 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", + "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', + "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') + (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) + + 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 (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() + 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() + + (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" + + 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"] diff --git a/tests/unit/torch/export/test_hf_checkpoint_utils.py b/tests/unit/torch/export/test_hf_checkpoint_utils.py index f3be2564312..08292c65f9f 100644 --- a/tests/unit/torch/export/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/export/test_hf_checkpoint_utils.py @@ -24,7 +24,59 @@ 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, +) +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):