hf export: download trust_remote_code .py files via snapshot_download

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

Author: j-rausch

examples/megatron_bridge/distill.py

@@ -338,10 +338,23 @@ def _build_model_provider(hf_path):
                 show_progress=True,
                 strict=True,
             )
-            # Copy config.json from student_hf_path (handles both local paths and HF model IDs)
-            AutoConfig.from_pretrained(
+            # Save config.json from student_hf_path (handles both local paths and HF model IDs)
+            config = AutoConfig.from_pretrained(
                 args.student_hf_path, trust_remote_code=args.trust_remote_code
-            ).save_pretrained(args.hf_export_path)
+            )
+            config.save_pretrained(args.hf_export_path)
+
+            # Download trust_remote_code .py files from HF hub so the exported
+            # checkpoint can be loaded with trust_remote_code=True from a local path.
+            # save_pretrained only writes config.json, not the modeling code.
+            if hasattr(config, "auto_map") and isinstance(config.auto_map, dict):
+                from huggingface_hub import snapshot_download
+
+                snapshot_download(
+                    repo_id=args.student_hf_model,
+                    local_dir=args.hf_export_path,
+                    allow_patterns=["*.py"],
+                )
 
 
 if __name__ == "__main__":