copy custom modeling files to pruned checkpoint dirs; without them, trust_remote_code checkpoints are silently excluded from the replacement library

Signed-off-by: jrausch <jrausch@nvidia.com>

Author: j-rausch

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

@@ -30,6 +30,7 @@
 from ..checkpoint_utils_hf import (
     _get_auto_class_for_trust_remote_code,
     _save_checkpoint,
+    copy_remote_code_files,
     load_model_config,
 )
 from ..logger import mprint
@@ -87,6 +88,9 @@ def init_child_from_parent(
         trust_remote_code=descriptor.requires_trust_remote_code(),
     )
 
+    if descriptor.requires_trust_remote_code():
+        copy_remote_code_files(parent_checkpoint_dir, output_checkpoint_dir)
+
     parent_model_config = load_model_config(
         parent_checkpoint_dir, trust_remote_code=descriptor.requires_trust_remote_code()
     )

modelopt/torch/puzzletron/tools/checkpoint_utils_hf.py

@@ -22,6 +22,7 @@
 import dataclasses
 import fcntl
 import os
+import shutil
 import time
 from collections import defaultdict
 from collections.abc import Callable, Mapping
@@ -448,3 +449,26 @@ def save_model_config(model_config: PretrainedConfig, checkpoint_dir: Path | str
             for conf in model_config.block_configs
         ]
     model_config.save_pretrained(checkpoint_dir)
+
+
+def copy_remote_code_files(source_dir: Path | str, output_dir: Path | str) -> None:
+    """Copy custom code ``.py`` files from a trust-remote-code checkpoint into *output_dir*.
+
+    Models with dynamic modules (e.g. Nemotron-H) ship custom Python files
+    (``modeling_*.py``, ``configuration_*.py``, etc.) alongside ``config.json``.
+    When pruned checkpoints are saved via ``_save_checkpoint``, these files are not
+    automatically copied, so ``AutoConfig.from_pretrained(output_dir, trust_remote_code=True)``
+    fails and the checkpoint is silently excluded from the replacement library.
+
+    Only copies ``.py`` files from the root of *source_dir* (matching the layout of
+    HuggingFace Hub checkpoints with custom code).
+    """
+    source_dir = Path(source_dir)
+    output_dir = Path(output_dir)
+    if not source_dir.is_dir():
+        return
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for py_file in source_dir.glob("*.py"):
+        target = output_dir / py_file.name
+        if not target.exists():
+            shutil.copy2(py_file, target)

tests/unit/torch/puzzletron/test_copy_remote_code_files.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for copy_remote_code_files.
+
+Verifies that custom modeling .py files are copied from a source checkpoint
+directory to an output directory so that pruned checkpoints remain valid
+trust_remote_code checkpoints.
+"""
+
+import pytest
+
+pytest.importorskip("transformers")
+
+from modelopt.torch.puzzletron.tools.checkpoint_utils_hf import copy_remote_code_files
+
+
+def test_copies_only_py_files(tmp_path):
+    source = tmp_path / "source"
+    source.mkdir()
+    (source / "modeling_nemotron_h.py").write_text("class Model: pass")
+    (source / "configuration_nemotron_h.py").write_text("class Config: pass")
+    (source / "config.json").write_text("{}")
+    (source / "model.safetensors").write_bytes(b"\x00")
+
+    output = tmp_path / "child"
+
+    copy_remote_code_files(source, output)
+
+    assert (output / "modeling_nemotron_h.py").exists()
+    assert (output / "configuration_nemotron_h.py").exists()
+    assert not (output / "config.json").exists()
+    assert not (output / "model.safetensors").exists()
+
+
+def test_skips_existing_files(tmp_path):
+    source = tmp_path / "source"
+    source.mkdir()
+    (source / "modeling.py").write_text("new content")
+
+    output = tmp_path / "child"
+    output.mkdir()
+    (output / "modeling.py").write_text("existing content")
+
+    copy_remote_code_files(source, output)
+
+    assert (output / "modeling.py").read_text() == "existing content"
+
+
+def test_noop_for_missing_source(tmp_path):
+    copy_remote_code_files(tmp_path / "nonexistent", tmp_path / "output")
+
+    assert not (tmp_path / "output").exists()