Skip to content

Preserve HF PTQ checkpoint sidecar files [NV BUG 6491822] #2060

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jenchen13 wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Preserve HF PTQ checkpoint sidecar files [NV BUG 6491822] #2060

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
da68df5
Preserve HF PTQ checkpoint sidecars
jenchen13 Aug 5, 2026
cec2df9
Preserve export generation config
jenchen13 Aug 5, 2026
e18d5e4
Restore safetensor sidecar exclusion
jenchen13 Aug 5, 2026
874a57c
Continue sidecar copy after errors
jenchen13 Aug 5, 2026
d669f9c
Preserve exported tokenizer config
jenchen13 Aug 5, 2026
7e701d4
Merge branch 'main' into hf_ptq_sidecars
jenchen13 Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 81 additions & 61 deletions examples/hf_ptq/example_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import json
import logging
import os
import shutil
import warnings
from collections.abc import Callable, Iterable
from dataclasses import dataclass
Expand All @@ -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
Expand All @@ -56,6 +56,52 @@

SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]

_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS = [
Comment thread
jenchen13 marked this conversation as resolved.
"*.jinja",
"*.json",
"*.md",
"*.model",
"*.py",
"*.tiktoken",
"*.txt",
"LICENSE*",
"NOTICE*",
]
Comment thread
jenchen13 marked this conversation as resolved.
_HF_PTQ_WEIGHT_FILE_PATTERNS = (
"*.safetensors",
Comment thread
jenchen13 marked this conversation as resolved.
"*.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",
Comment thread
jenchen13 marked this conversation as resolved.
"hf_quant_config.json",
"quant_config.json",
"quantization_config.json",
"quantize_config.json",
"recipe.yaml",
"recipe.yml",
}
Comment thread
jenchen13 marked this conversation as resolved.
Comment on lines +104 to +112

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The name _HF_PTQ_EXPORT_OWNED_FILES conflates two different policies, which makes the set easy to misextend later.

config.json and hf_quant_config.json really are export-owned — ModelOpt writes them, so skipping the source copy is an overwrite-protection. But ModelOpt never writes quant_config.json, quantization_config.json, quantize_config.json, recipe.yaml, or recipe.yml; for those, skipping means the file is dropped entirely from the export. That's the intended behavior for an already-quantized source (stale compressed-tensors/llm-compressor metadata next to a fresh hf_quant_config.json confuses deployment stacks), but it's the opposite of "owned," and a future reader adding a name here can't tell which semantics they're getting.

Worth splitting so the intent is self-documenting, and so the recipe.yaml breadth is visible — recipe.yaml is a generic filename that a non-quantized source could plausibly use for something unrelated, and it will be silently discarded:

# Written by the export itself; the source copy must not clobber them.
_HF_PTQ_EXPORT_OWNED_FILES = {
    "config.json",
    "hf_quant_config.json",
}
# Quantization metadata from an already-quantized source (compressed-tensors,
# llm-compressor, MXFP4). Dropped, not overwritten: carrying it next to a freshly
# written hf_quant_config.json makes deployment stacks pick the wrong algorithm.
_HF_PTQ_STALE_QUANT_METADATA_FILES = {
    "quant_config.json",
    "quantization_config.json",
    "quantize_config.json",
    "recipe.yaml",
    "recipe.yml",
}

and at the call site: exclude_files=_HF_PTQ_EXPORT_OWNED_FILES | _HF_PTQ_STALE_QUANT_METADATA_FILES | set(exclude_files or ()).



@dataclass
class DistributedState:
Expand Down Expand Up @@ -895,11 +941,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,

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] For HF-Hub-ID sources, the widening this PR is about is capped by a second, independent whitelist — so sidecars are still silently dropped.

What happens. When --pyt_ckpt_path is a Hub ID, _resolve_model_path reaches this branch (config._name_or_path is the repo ID string, not a dir, so os.path.isdir() is False). The earlier load_model()/from_pretrained() only pulled config + weights + tokenizer files into the snapshot, so this snapshot_download call is what actually fetches the sidecars that copy_custom_model_files will later copy. Anything outside _HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS is never on disk, so the copy step can't preserve it — the file is dropped exactly as before, just via a different whitelist.

Concrete misses (all copied fine for a local source dir, dropped for a Hub ID):

  • *.yaml / *.yml — non-quantization config sidecars
  • *.jinja2 / *.j2 — repos that didn't adopt the .jinja extension
  • tokenizer.model.v3, tokenizer.model.v7 — Mistral-family tokenizers (*.model does not match tokenizer.model.v3 under fnmatch)
  • *.png / *.jpg — assets referenced by the README.md this PR now copies, so the copied README ends up with broken images

This isn't a regression (the old list was just ["*.py", "*.json"]), but it leaves the stated goal half-done for one of the two supported source forms, and it re-introduces the brittle-whitelist failure mode the PR set out to remove.

Suggested fix. _HF_PTQ_WEIGHT_FILE_PATTERNS is now comprehensive enough (it already covers *.gguf, *.npz, *.pkl, *.tar*, *.zip, *.h5, *.msgpack) to serve as the download's ignore list, which also collapses the two lists into one policy and removes the drift risk between them:

local_path = snapshot_download(
    repo_id=model_name_or_path,
    ignore_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS,
)

_HF_SIDECAR_DOWNLOAD_ALLOW_PATTERNS can then be deleted, and test_resolve_model_path_snapshot_download_stays_allowlisted updated to assert on ignore_patterns. If you'd rather keep an allow-list for the download path, please at least add *.yaml, *.yml, *.jinja2, tokenizer.model*, and a comment noting that the two lists must be kept in sync.

)
Comment thread
jenchen13 marked this conversation as resolved.
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
Expand Down Expand Up @@ -934,49 +982,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. 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: Whether trust_remote_code was used when resolving HuggingFace model
IDs
exclude_files: Additional source file names to skip.
Comment on lines 1040 to +1044

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The trust_remote_code docstring is now misleading about what the parameter does.

"Whether trust_remote_code was used when resolving HuggingFace model IDs" describes the plumbing but not the consequence, and the consequence is the part that changed in this PR: *.py files (modeling*.py, configuration_*.py, tokenization_*.py, custom reasoning parsers) are now copied into the export regardless of this flag. Previously they were gated behind it, and the removed docstring stated that gating explicitly. Since executable code in the export directory is a security-posture question, a reader checking "does this flag stop remote code from landing in my export?" should get a straight answer here rather than having to diff against main.

        trust_remote_code: Forwarded to ``AutoConfig.from_pretrained`` when ``source_path``
            is a HuggingFace model ID, so the source snapshot can be resolved. It does
            **not** gate which files are copied: remote-code modules present in the source
            checkpoint are copied either way, since deployment stacks need them to load a
            custom-architecture export.

Same for the sentence "Weight and weight-index files are skipped" — worth noting the skip is top-level only (the helper does not recurse), so a source with weights in a subdirectory is not covered.

"""
# 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)

@cjluo-nv cjluo-nv Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot comment.

The old docstring deliberately gated executable code (modeling*.py, configuration_*.py, tokenization_*.py) behind trust_remote_code; after this change those files are copied into the export unconditionally and trust_remote_code only affects HF-ID resolution. That's likely intentional given the goal, but please call it out in the PR body (it's a behavior change relevant to SECURITY.md guidance) and add a test pinning the new parity between trust_remote_code=True/False.


Expand All @@ -997,29 +1027,19 @@ 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 | set(exclude_files or ()),
exclude_patterns=_HF_PTQ_WEIGHT_FILE_PATTERNS,
)
Comment on lines +1070 to +1075

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] Delegating to the shared helper drops the per-file try/except, so one unreadable sidecar now aborts the entire export at the final step.

What changed. The old loop isolated failures per file:

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}")

copy_non_safetensor_files_from_ckpt has no equivalent guard — shutil.copy2 raises straight out of the loop.

Why it matters more now than before. Two effects compound:

  1. This PR deliberately widens the copy set from ~11 curated names to every non-weight top-level file, so the number of copy2 calls that can fail grows a lot, and now includes files nobody vetted (.gitattributes, root-owned files in a shared HF cache, README assets, arbitrary vendor sidecars).
  2. copy_custom_model_files is the last thing export_quantized does (hf_ptq.py:937), after calibration and after export_hf_checkpoint has written all the quantized weights. So the failure mode is: hours of PTQ complete, weights are on disk, then the run dies on a README.md copy and the user sees a traceback instead of Quantized model exported to: ....

Concrete triggers, all reachable in practice:

  • shutil.SameFileError when --export_path resolves to the same directory as --pyt_ckpt_path (in-place re-export). Previously warned per file and continued; now it's a hard crash.
  • PermissionError from copy2's copystat on a destination filesystem that rejects chmod/utime (bind-mounted volumes, NFS, some container overlays) — note the bytes are already written when copystat fails, so the file lands but the export still dies.
  • PermissionError/OSError reading a mode-600 file owned by another user in a shared HF_HOME cache.

Suggested fix. Keep the shared helper as-is (the Megatron caller wants strictness) and restore per-file tolerance in this example wrapper, since a missing sidecar is not worth discarding a completed quantization run. Either wrap the call:

try:
    copied_files = copy_non_safetensor_files_from_ckpt(...)
except Exception as e:
    print(f"Warning: Failed to copy checkpoint sidecar files: {e}")
    copied_files = []

or, better, add an opt-in continue_on_error: bool = False to copy_non_safetensor_files_from_ckpt that wraps just the shutil.copy2 in try/except and warns, and pass continue_on_error=True here. The second form preserves the "copy as much as possible" intent instead of giving up on the first bad file.


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:
Expand Down
13 changes: 10 additions & 3 deletions examples/hf_ptq/hf_ptq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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,
)
Comment on lines +941 to +947

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Export] The widened copy now lets the source tokenizer_config.json overwrite the exported one, which can leave the export with two chat templates in conflicting formats.

Why this is new. Before this PR, tokenizer_config.json was only copied when trust_remote_code was set (it matched code_patterns' "*.json", not always_copy_patterns). Now it is copied unconditionally on every export path. The docstring's "source tokenizer and processor files intentionally still win" is a reasonable policy in isolation, but it interacts badly with the transformers 4.x → 5.x chat-template relocation that the removed docstring in this same PR explicitly calls out.

Failure scenario (very common repo layout — Llama/Qwen/Mistral native, no --trust_remote_code):

  1. Source repo is transformers-4.x era: tokenizer_config.json carries an embedded "chat_template": "...", and there is no chat_template.jinja.
  2. tokenizer.save_pretrained(export_path) at hf_ptq.py:929 runs under transformers 5.x, which splits the template out: it writes chat_template.jinja and a tokenizer_config.json with the chat_template key removed.
  3. copy_custom_model_files then overwrites tokenizer_config.json with the source copy, re-introducing the embedded chat_template.
  4. The export directory now contains both chat_template.jinja and tokenizer_config.json["chat_template"].

Transformers treats a template defined in both places as a conflict rather than a benign duplicate — please verify against the transformers version this example targets, since the outcome is either a hard ValueError on AutoTokenizer.from_pretrained (breaking vLLM/SGLang serving of the exported checkpoint outright) or a version-dependent silent precedence choice. Either way the export has two sources of truth for the template where a plain save_pretrained would have produced one, and this isn't covered by the new tests.

Suggested fix. Make the exclusion decision reflect what the export actually wrote, rather than hardcoding a single filename. generation_config.json is already handled this way in spirit; extend it:

if args.dist_state.is_main:
    exclude_files = set()
    if not is_tensorrt_llm_export:
        # export_hf_checkpoint() writes a sanitized generation_config.json
        # (see _sanitize_generation_config_for_save); don't let the source clobber it.
        exclude_files.add("generation_config.json")
    if (Path(export_path) / "chat_template.jinja").exists():
        # tokenizer.save_pretrained() relocated the chat template out of
        # tokenizer_config.json; copying the source config back would restore a
        # duplicate template in the legacy embedded format.
        exclude_files.add("tokenizer_config.json")
    copy_custom_model_files(
        args.pyt_ckpt_path,
        export_path,
        args.trust_remote_code,
        exclude_files=exclude_files or None,
    )

If you'd rather keep source-wins for tokenizer_config.json, then also drop the exported chat_template.jinja in that case so exactly one template format survives — and please add a test fixture with an embedded chat_template in the source tokenizer_config.json asserting the export ends up with only one.


end_time = time.time()
print_rank_0(
Expand Down
33 changes: 28 additions & 5 deletions modelopt/torch/export/plugins/hf_checkpoint_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -253,25 +255,46 @@ 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]:
Comment thread
coderabbitai[bot] marked this conversation as resolved.
"""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
Comment on lines +265 to +272

@claude claude Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION] The removed sentence documented a contract the Megatron caller still depends on.

The old docstring said "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." That is exactly what unified_export_megatron.py:319-330 does — it calls this helper first, then self._hf_config.save_pretrained() and generation_config.save_pretrained() to overwrite. The replacement text ("Callers can exclude additional files or patterns when copying after export-owned metadata has already been written") describes only the new ordering used by hf_ptq, so the two callers now use opposite orderings and the docstring only documents one of them.

Both are valid, but a future caller reading this can't tell which one to follow. Suggest documenting both:

    Use as a baseline so tokenizer files, remote_code ``*.py``, README, LICENSE, etc.
    are preserved from the source. Two call orderings are supported:

    * copy first, then overwrite the files modelopt owns (``config.json``,
      ``generation_config.json``, ``hf_quant_config.json``) — see
      ``unified_export_megatron``;
    * write modelopt-owned metadata first, then copy with those names passed via
      ``exclude_files``/``exclude_patterns`` — see ``examples/hf_ptq``.

    Copies top-level regular files only; subdirectories are not traversed.

The "top-level only" note is worth adding regardless — it's the pre-existing behavior but it's not stated anywhere, and it's load-bearing now that the exclusion patterns are the only thing keeping source weights out.

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)
Comment thread
jenchen13 marked this conversation as resolved.
Outdated
copied_files.append(entry)
Comment thread
jenchen13 marked this conversation as resolved.
return copied_files
85 changes: 85 additions & 0 deletions tests/examples/hf_ptq/test_example_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,91 @@ 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()

(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"]
Expand Down
Loading
Loading