Skip to content

[Bugfix] Enable essential python files export in llm_ptq examples by adding a customized copy function #351

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

Merged
merged 5 commits into from
Sep 26, 2025
Merged

[Bugfix] Enable essential python files export in llm_ptq examples by adding a customized copy function #351

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
39ccfad
add a copy_custom_model_files function in hf_ptq to copy essential py…
Edwardf0t1 Sep 21, 2025
3b0f98a
move customized copy function to example utils
Edwardf0t1 Sep 24, 2025
4aa4285
add a util function to resolve model path
Edwardf0t1 Sep 25, 2025
9d305fa
add a util function to resolve model path
Edwardf0t1 Sep 25, 2025
21fa47f
copy over json files in export path as well
Edwardf0t1 Sep 25, 2025
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
58 changes: 58 additions & 0 deletions examples/llm_ptq/hf_ptq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,10 @@
import argparse
import copy
import random
import shutil
import time
import warnings
from pathlib import Path
from typing import Any

import numpy as np
Expand Down Expand Up @@ -83,6 +85,56 @@
mto.enable_huggingface_checkpointing()


def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False):
"""Copy custom model files (configuration_*.py, modeling_*.py, etc.) from source to export directory.

Args:
source_path: Path to the original model directory
export_path: Path to the exported model directory
trust_remote_code: Whether trust_remote_code was used (only copy files if True)
"""
if not trust_remote_code:
return

source_dir = Path(source_path)
export_dir = Path(export_path)

if not source_dir.exists():
print(f"Warning: Source directory {source_path} does not exist")
return

if not export_dir.exists():
print(f"Warning: Export directory {export_path} does not exist")
return

# Common patterns for custom model files that need to be copied
custom_file_patterns = [
"configuration_*.py",
"modeling_*.py",
"tokenization_*.py",
"processing_*.py",
"image_processing_*.py",
"feature_extraction_*.py",
]

copied_files = []
for pattern in custom_file_patterns:
for file_path in source_dir.glob(pattern):
if file_path.is_file():
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}")

if copied_files:
print(f"Successfully copied {len(copied_files)} custom model files to {export_path}")
else:
print("No custom model files found to copy")


Copy link

Choose a reason for hiding this comment

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

⚠️ Potential issue

Make file copy resilient: handle nested packages via auto_map, preserve structure, and guard same-dir.

Current logic only copies top-level files and may miss packages/modules referenced by HF auto_map (common in custom repos like Phi). It can also double-copy names and doesn’t handle source==export paths.

Apply this refactor:

-def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False):
-    """Copy custom model files (configuration_*.py, modeling_*.py, etc.) from source to export directory.
-
-    Args:
-        source_path: Path to the original model directory
-        export_path: Path to the exported model directory
-        trust_remote_code: Whether trust_remote_code was used (only copy files if True)
-    """
-    if not trust_remote_code:
-        return
-
-    source_dir = Path(source_path)
-    export_dir = Path(export_path)
-
-    if not source_dir.exists():
-        print(f"Warning: Source directory {source_path} does not exist")
-        return
-
-    if not export_dir.exists():
-        print(f"Warning: Export directory {export_path} does not exist")
-        return
-
-    # Common patterns for custom model files that need to be copied
-    custom_file_patterns = [
-        "configuration_*.py",
-        "modeling_*.py",
-        "tokenization_*.py",
-        "processing_*.py",
-        "image_processing_*.py",
-        "feature_extraction_*.py",
-    ]
-
-    copied_files = []
-    for pattern in custom_file_patterns:
-        for file_path in source_dir.glob(pattern):
-            if file_path.is_file():
-                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}")
-
-    if copied_files:
-        print(f"Successfully copied {len(copied_files)} custom model files to {export_path}")
-    else:
-        print("No custom model files found to copy")
+def copy_custom_model_files(source_path: str, export_path: str, trust_remote_code: bool = False):
+    """Copy custom model Python files from source to export directory.
+    - Supports nested packages referenced by config.json::auto_map
+    - Preserves package structure (subdirectories)
+    - No-op unless trust_remote_code is True
+    """
+    if not trust_remote_code:
+        return
+
+    source_dir = Path(source_path)
+    export_dir = Path(export_path)
+
+    if not source_dir.exists():
+        print(f"Warning: Source directory {source_path} does not exist")
+        return
+    if not export_dir.exists():
+        print(f"Warning: Export directory {export_path} does not exist")
+        return
+    if source_dir.resolve() == export_dir.resolve():
+        print("Info: Source and export directories are identical; skipping custom file copy")
+        return
+
+    # Common patterns for custom model files
+    custom_file_patterns = [
+        "configuration_*.py",
+        "modeling_*.py",
+        "tokenization_*.py",
+        "processing_*.py",
+        "image_processing_*.py",
+        "feature_extraction_*.py",
+        "__init__.py",
+    ]
+
+    candidate_paths: set[Path] = set()
+
+    # 1) Collect files by patterns recursively
+    for pattern in custom_file_patterns:
+        for file_path in source_dir.rglob(pattern):
+            if file_path.is_file():
+                candidate_paths.add(file_path)
+
+    # 2) Collect modules from config.json auto_map if present
+    cfg_path = source_dir / "config.json"
+    if cfg_path.exists():
+        try:
+            with cfg_path.open("r", encoding="utf-8") as f:
+                cfg = json.load(f)
+            auto_map = cfg.get("auto_map", {}) or {}
+
+            def iter_strings(x):
+                if isinstance(x, str):
+                    yield x
+                elif isinstance(x, list):
+                    for i in x:
+                        if isinstance(i, str):
+                            yield i
+                elif isinstance(x, dict):
+                    for i in x.values():
+                        if isinstance(i, str):
+                            yield i
+
+            modules = set()
+            for v in auto_map.values():
+                for s in iter_strings(v):
+                    s = s.split(":")[0]                     # optional "module:qualname"
+                    mod = s.rsplit(".", 1)[0]               # drop class
+                    modules.add(mod)
+
+            for mod in modules:
+                rel = Path(mod.replace(".", "/"))
+                mod_file = source_dir / (str(rel) + ".py")
+                pkg_dir = source_dir / rel
+                if mod_file.exists():
+                    candidate_paths.add(mod_file)
+                if pkg_dir.is_dir():
+                    for p in pkg_dir.rglob("*.py"):
+                        candidate_paths.add(p)
+        except Exception as e:
+            print(f"Warning: Failed to parse auto_map from config.json: {e}")
+
+    copied = 0
+    for file_path in sorted(candidate_paths):
+        try:
+            rel = file_path.relative_to(source_dir)
+            dest_path = export_dir / rel
+            dest_path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(file_path, dest_path)
+            copied += 1
+            print(f"Copied custom model file: {rel}")
+        except Exception as e:
+            print(f"Warning: Failed to copy {file_path}: {e}")
+
+    if copied:
+        print(f"Successfully copied {copied} custom model files to {export_path}")
+    else:
+        print("No custom model files found to copy")

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
In examples/llm_ptq/hf_ptq.py around lines 88 to 137, the current
copy_custom_model_files only copies top-level files and can double-copy names or
fail when source==export or when custom code lives in nested packages referenced
via HuggingFace auto_map; update it to detect and preserve directory structure
for both files and packages: skip work when source_path == export_path, traverse
matches for patterns plus any packages referenced in an optional auto_map (or
any subdirectory under source that contains matching files), compute each file's
relative path to source_dir and create the same subdirectory under export_dir
before copying, use shutil.copy2 for files and shutil.copytree for directories
(or a recursive copy that merges into export_dir) while avoiding overwriting
identical existing files (or only overwrite if different), and deduplicate by
tracking already-copied relative paths; ensure parent dirs are created and log
clear warnings on failures.

def auto_quantize(
model, qformat, auto_quantize_bits, calib_dataloader, calibrate_loop, batch_size=1
):
Expand Down Expand Up @@ -604,6 +656,9 @@ def output_decode(generated_ids, input_shape):
inference_tensor_parallel=args.inference_tensor_parallel,
inference_pipeline_parallel=args.inference_pipeline_parallel,
)

# Copy custom model files for TensorRT-LLM export as well
copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)
else:
# Check arguments for unified_hf export format and set to default if unsupported arguments are provided
assert args.sparsity_fmt == "dense", (
Expand All @@ -621,6 +676,9 @@ def output_decode(generated_ids, input_shape):
export_dir=export_path,
)

# Copy custom model files (configuration_*.py, modeling_*.py, etc.) if trust_remote_code is used
copy_custom_model_files(args.pyt_ckpt_path, export_path, args.trust_remote_code)

# Restore default padding and export the tokenizer as well.
if tokenizer is not None:
tokenizer.padding_side = default_padding_side
Expand Down
Loading