Copy remote code files to output hf ckpt

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

modelopt/torch/export/plugins/hf_checkpoint_utils.py

@@ -21,42 +21,38 @@
 from pathlib import Path
 
 import torch
-from huggingface_hub import hf_hub_download, list_repo_files
+from huggingface_hub import snapshot_download
 from safetensors.torch import safe_open
 from tqdm import tqdm
 
 
-def copy_remote_code(
-    pretrained_model_path: str | os.PathLike,
-    save_directory: str | os.PathLike,
-):
+def copy_remote_code(pretrained_model_path: str | os.PathLike, save_directory: str | os.PathLike):
     """Copy remote code from pretrained model to save directory.
 
     For models that keep configuration and modeling files as part of the checkpoint,
     we need to copy them to the export directory for seamless integration with inference
     frameworks.
 
     If ``pretrained_model_path`` is a local directory, Python files are copied directly.
-    If it is a HuggingFace Hub model ID, Python files are downloaded from the Hub first.
+    If it's a HF Hub model ID (e.g. ``nvidia/NVIDIA-Nemotron-Nano-12B-v2``), files are downloaded from the Hub.
 
     Args:
         pretrained_model_path: Local path to the pretrained model or HuggingFace Hub model ID.
         save_directory: Path to the save directory.
     """
     hf_checkpoint_path = Path(pretrained_model_path)
     save_dir = Path(save_directory)
+    save_dir.mkdir(parents=True, exist_ok=True)
 
     if hf_checkpoint_path.is_dir():
         for py_file in hf_checkpoint_path.glob("*.py"):
-            if py_file.is_file():
-                shutil.copy(py_file, save_dir / py_file.name)
+            shutil.copy2(py_file, save_dir / py_file.name)
     else:
-        # Hub model ID: download any top-level .py files (custom modeling code)
-        repo_id = str(pretrained_model_path)
-        for filename in list_repo_files(repo_id):
-            if "/" not in filename and filename.endswith(".py"):
-                local_path = hf_hub_download(repo_id=repo_id, filename=filename)
-                shutil.copy(local_path, save_dir / filename)
+        snapshot_download(
+            repo_id=str(pretrained_model_path),
+            local_dir=str(save_dir),
+            allow_patterns=["*.py"],
+        )
 
 
 def load_multimodal_components(

modelopt/torch/puzzletron/tools/bypassed_training/init_child_from_parent.py

@@ -24,6 +24,8 @@
 import yaml
 from transformers import AutoModelForCausalLM
 
+from modelopt.torch.export.plugins.hf_checkpoint_utils import copy_remote_code
+
 from ...anymodel.model_descriptor import ModelDescriptor, ModelDescriptorFactory
 from ...anymodel.puzzformer import deci_x_patcher
 from ..checkpoint_utils import copy_tokenizer, load_state_dict
@@ -87,6 +89,9 @@ def init_child_from_parent(
         trust_remote_code=descriptor.requires_trust_remote_code(),
     )
 
+    if descriptor.requires_trust_remote_code():
+        copy_remote_code(parent_checkpoint_dir, output_checkpoint_dir)
+
     parent_model_config = load_model_config(
         parent_checkpoint_dir, trust_remote_code=descriptor.requires_trust_remote_code()
     )

tests/_test_utils/torch/puzzletron/utils.py

@@ -19,11 +19,11 @@
 import torch
 from _test_utils.torch.transformers_models import get_tiny_tokenizer
 from datasets import Dataset, DatasetDict
-from huggingface_hub import snapshot_download
 from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedTokenizerBase
 
 import modelopt.torch.puzzletron as mtpz
 import modelopt.torch.utils.distributed as dist
+from modelopt.torch.export.plugins.hf_checkpoint_utils import copy_remote_code
 
 
 def setup_test_model_and_data(
@@ -189,21 +189,11 @@ def create_and_save_small_hf_model(
                 submodule._tied_weights_keys = None
         model.save_pretrained(output_path, save_original_format=False)
 
-    # Save tokenizer
+    # Save tokenizer, config, and custom code files
     tokenizer.save_pretrained(output_path)
-
-    # Save config
     config.save_pretrained(output_path)
-
-    # Download trust_remote_code .py files from HF hub into the checkpoint directory so that
-    # force_cache_dynamic_modules can resolve classes from the local path.
-    # save_pretrained only saves weights + config, not these .py files.
     if hasattr(config, "auto_map") and isinstance(config.auto_map, dict):
-        snapshot_download(
-            repo_id=hf_model_name,
-            local_dir=output_path,
-            allow_patterns=["*.py"],
-        )
+        copy_remote_code(hf_model_name, output_path)
 
 
 def save_dummy_dataset(dataset_path: Path | str):

…torch/export/test_hf_checkpoint_utils.py …torch/export/test_hf_checkpoint_utils.pytests/gpu_megatron/torch/export/test_hf_checkpoint_utils.py renamed to tests/unit/torch/export/test_hf_checkpoint_utils.py tests/gpu_megatron/torch/export/test_hf_checkpoint_utils.py renamed to tests/unit/torch/export/test_hf_checkpoint_utils.py

@@ -56,54 +56,14 @@ def test_copy_remote_code_local_dir_no_py_files(tmp_path):
 
 
 def test_copy_remote_code_hub_id(tmp_path):
-    """copy_remote_code downloads and copies top-level .py files from a Hub model ID."""
+    """copy_remote_code delegates to snapshot_download for a Hub model ID."""
     dst_dir = tmp_path / "dst"
-    dst_dir.mkdir()
-
-    # Create a fake cached file that hf_hub_download would return
-    cached_py = tmp_path / "cached_modeling_custom.py"
-    cached_py.write_text("# custom hub model")
-
-    repo_files = [
-        "modeling_custom.py",  # top-level .py — should be downloaded
-        "config.json",  # non-.py — skip
-        "model.safetensors",  # non-.py — skip
-        "subdir/nested.py",  # subdirectory .py — skip (contains "/")
-    ]
-
-    with (
-        patch(
-            "modelopt.torch.export.plugins.hf_checkpoint_utils.list_repo_files",
-            return_value=repo_files,
-        ) as mock_list,
-        patch(
-            "modelopt.torch.export.plugins.hf_checkpoint_utils.hf_hub_download",
-            return_value=str(cached_py),
-        ) as mock_download,
-    ):
-        copy_remote_code("meta-llama/Llama-3.2-1B", dst_dir)
-
-    mock_list.assert_called_once_with("meta-llama/Llama-3.2-1B")
-    # Only the top-level .py should have been downloaded
-    mock_download.assert_called_once_with(
-        repo_id="meta-llama/Llama-3.2-1B", filename="modeling_custom.py"
-    )
-    assert (dst_dir / "modeling_custom.py").read_text() == "# custom hub model"
 
+    with patch("modelopt.torch.export.plugins.hf_checkpoint_utils.snapshot_download") as mock_sd:
+        copy_remote_code("nvidia/NVIDIA-Nemotron-Nano-12B-v2", dst_dir)
 
-def test_copy_remote_code_hub_id_no_py_files(tmp_path):
-    """copy_remote_code is a no-op when the Hub repo has no top-level .py files."""
-    dst_dir = tmp_path / "dst"
-    dst_dir.mkdir()
-
-    with (
-        patch(
-            "modelopt.torch.export.plugins.hf_checkpoint_utils.list_repo_files",
-            return_value=["config.json", "model.safetensors"],
-        ),
-        patch("modelopt.torch.export.plugins.hf_checkpoint_utils.hf_hub_download") as mock_download,
-    ):
-        copy_remote_code("meta-llama/Llama-3.2-1B", dst_dir)
-
-    mock_download.assert_not_called()
-    assert list(dst_dir.iterdir()) == []
+    mock_sd.assert_called_once_with(
+        repo_id="nvidia/NVIDIA-Nemotron-Nano-12B-v2",
+        local_dir=str(dst_dir),
+        allow_patterns=["*.py"],
+    )