megatron_bridge: load the HF model from local disk; Megatron export copies its non-model files by shengliangxu · Pull Request #2610 · NVIDIA/Model-Optimizer

shengliangxu · 2026-09-30T21:55:15Z

What does this PR do?

Type of change: new feature (with a backward breaking change and a deprecation)

Summary of the series

Every pipeline that quantizes, prunes, sparsifies or distills a Hugging Face model now starts the same way: make the checkpoint local, then load from it. ensure_local_checkpoint returns a local directory as-is and downloads a Hub model ID in full if it is not cached yet -- once, on rank 0 -- and the model, tokenizer and processor are then loaded from that local directory.

Before, a Hub ID went straight to from_pretrained, which fetched only the files loading needed, so every later step that needed the source as files found its own way back to it: hf_ptq through a hand-rolled cache lookup (_resolve_model_path, whose TRANSFORMERS_CACHE glob picked the lexicographically largest commit hash rather than refs/main) and a hardcoded download allowlist, regenerating tokenizer files along the way; the Megatron export by fetching only *.py for a Hub ID; other callers of export_hf_checkpoint not at all.

With the source on local disk, exports follow one rule: an export writes the model files -- weights in any format and the metadata describing them -- and carries every other source file (non-model files: tokenizer, processor, remote code, chat templates, README, ...) over verbatim, from local disk only.

The rest of the series is merged: #2607, #2608, #2609 and #2611. This PR is its last part and is based on main.

This PR

Loading. The megatron_bridge quantize, export, prune, distillation and distillation-export scripts now make the HF model local first with ensure_local_checkpoint -- a local directory as-is, a Hub ID downloaded once -- and load the model, tokenizer and processor from that copy (args.local_checkpoint_path; student_local_checkpoint_path in distillation and its export). --hf_model_name_or_path / --student_hf_path stay as given, and the tokenizer path recorded in saved Megatron checkpoints stays the Hub ID, which remains valid on other hosts.

Exporting. The Megatron unified HF export copied every top-level non-safetensors file of a local source -- *.bin weights included -- but only *.py of a Hub ID, then regenerated the generation config, tokenizer and image processor. Now it writes its config and copies the source's non-model files with copy_non_model_files, from a local source only; a Hub ID gets a warning, as in the HF exporters. Breaking: callers passing a Hub ID no longer get those files (CHANGELOG entry with the fix: pass the path from ensure_local_checkpoint). prune_minitron and the distillation VLM export -- which every rank reaches -- copy from the main rank only, and the distillation LLM export now copies the files too.

copy_non_safetensor_files_from_ckpt lost its last internal user and is deprecated in favour of copy_non_model_files.

Usage

from modelopt.torch.export import ensure_local_checkpoint, export_mcore_gpt_to_hf

_, local_checkpoint_path = ensure_local_checkpoint("Qwen/Qwen3-8B")
export_mcore_gpt_to_hf(model, local_checkpoint_path, export_dir=export_dir, dtype=torch.bfloat16)

Testing

Pre-commit (ruff, mypy, bandit, rst, license) passes on the branch against main. CPU: tests/unit/torch/utils/test_hf_checkpoint_utils.py and tests/unit/torch/export, 483 passed (deprecation test included). GPU, left to CI: the existing tests/examples/megatron_bridge runs of quantize.py + export_quantized_megatron_to_hf.py, prune_minitron.py, and distill.py + export_distilled_megatron_to_hf.py (LLM and VLM) now pass the model as a Hub ID served from a local copy (fake_hub_checkpoint, single-rank runs). They check that loading and every export read the resolved copy and that the quantized Megatron checkpoint records the Hub ID as its tokenizer, adding no new test runs.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

Is this change backward compatible?: ❌ -- export_mcore_gpt_to_hf with a Hub-ID source no longer writes tokenizer, processor, generation config or remote-code files; pass a local directory (e.g. from ensure_local_checkpoint). copy_non_safetensor_files_from_ckpt is deprecated but still works.
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
Did you write any new necessary tests?: ✅ (deprecation); the exporter path is covered by the Megatron GPU tests in CI
Did you update Changelog?: ✅
Did you get Claude approval on this PR?: ❌

Additional Information

Based on main.

Summary by CodeRabbit

Exports
- Megatron-to-Hugging Face exports copy non-model files, such as tokenizer, processor, generation-configuration, and remote-code files, when the source checkpoint is available locally.
- Exports using a Hub model ID skip these files. Use a local checkpoint to include them.
Workflow Improvements
- Quantization, distillation, pruning, and export workflows resolve Hub checkpoints locally before loading them.
Deprecations
- The previous checkpoint file-copying utility now displays a deprecation warning. Use the replacement utility instead.

copy-pr-bot · 2026-09-30T21:55:18Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-09-30T21:55:23Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Last reviewed commit: 4a208a2
Latest commit, not reviewed: e1fc2c7

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cdd1ca4c-6857-4e9a-841c-788c911ac742

📥 Commits

Reviewing files that changed from the base of the PR and between e8e3d5a and 4a208a2.

📒 Files selected for processing (4)

tests/examples/megatron_bridge/conftest.py
tests/examples/megatron_bridge/test_distill.py
tests/examples/megatron_bridge/test_prune_minitron.py
tests/examples/megatron_bridge/test_quantize_export.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review.

📝 Walkthrough

Megatron Bridge scripts now resolve Hugging Face inputs to local checkpoints for model operations and export. Export paths copy non-model files from local sources. The unified exporter warns when the source is a Hub model ID. The older non-safetensor copy helper emits a deprecation warning.

Changes: Megatron HF checkpoint handling

Layer / File(s)	Summary
Resolve checkpoint inputs `examples/megatron_bridge/distill.py`, `examples/megatron_bridge/export_distilled_megatron_to_hf.py`, `examples/megatron_bridge/export_quantized_megatron_to_hf.py`, `examples/megatron_bridge/prune_minitron.py`, `examples/megatron_bridge/quantize.py`, `tests/examples/megatron_bridge/conftest.py`, `tests/examples/megatron_bridge/test_distill.py`, `tests/examples/megatron_bridge/test_prune_minitron.py`, `tests/examples/megatron_bridge/test_quantize_export.py`	The scripts resolve Hugging Face inputs to local checkpoint paths. They use those paths for model loading, processor loading, and export. Saved tokenizer paths use the Hub ID when available and otherwise use the local path. Integration tests cover Hub-ID inputs and files from the resolved checkpoint.
Copy non-model files during export `examples/megatron_bridge/export_distilled_megatron_to_hf.py`, `examples/megatron_bridge/prune_minitron.py`, `modelopt/torch/export/unified_export_megatron.py`, `CHANGELOG.rst`	Distilled and pruned exports copy non-model files from local sources. The unified exporter copies these files for local directories and warns when the source is a Hub model ID. The changelog records the behavior and local-source restriction.
Deprecate the non-safetensor copy helper `modelopt/torch/utils/plugins/hf_checkpoint_utils.py`, `tests/unit/torch/utils/test_hf_checkpoint_utils.py`, `CHANGELOG.rst`	The deprecated helper emits a `DeprecationWarning` and directs callers to `copy_non_model_files`. Tests expect the warning, and the changelog records the deprecation.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant MegatronBridgeScript
  participant ensure_local_checkpoint
  participant MegatronExporter
  participant copy_non_model_files
  MegatronBridgeScript->>ensure_local_checkpoint: Resolve Hugging Face input
  ensure_local_checkpoint-->>MegatronBridgeScript: Return local path and Hub ID
  MegatronBridgeScript->>MegatronExporter: Export using local checkpoint path
  MegatronExporter->>copy_non_model_files: Copy non-model files from local directory

Merge Risk: ⚪ Minimal · up to 4a208

The supplied changes do not establish an outstanding issue that should block merging after normal checks.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 28 functions across 12 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No prohibited security anti-pattern was introduced. The changed modelopt/examples code adds no torch.load(weights_only=False), numpy.load/np.load(allow_pickle=True), eval/exec, or # nosec usage. All a…
Title check		The title clearly and concisely summarizes the main changes: loading Hugging Face models from local disk and copying non-model files during Megatron export.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

No prohibited security anti-pattern was introduced. The changed modelopt/examples code adds no torch.load(weights_only=False), numpy.load/np.load(allow_pickle=True), eval/exec, or # nosec usage. All added trust_remote_code uses pass caller-controlled values; the example CLI flags and exporter parameters default to False. The PR adds no dependency changes in pyproject.toml or requirements files.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

Commit to this branch
Create a new PR

🧪 Generate unit tests (beta)

Commit to this branch
Create a new PR

Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

codecov · 2026-10-01T19:11:42Z

Codecov Report

❌ Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 78.81%. Comparing base (92a3cd5) to head (e1fc2c7).

Files with missing lines	Patch %	Lines
modelopt/torch/export/unified_export_megatron.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2610       +/-   ##
===========================================
+ Coverage   68.26%   78.81%   +10.55%     
===========================================
  Files         659      659               
  Lines       74089    74072       -17     
===========================================
+ Hits        50578    58381     +7803     
+ Misses      23511    15691     -7820

Flag	Coverage Δ
examples-diffusers	`20.04% <0.00%> (-0.01%)`	⬇️
examples-gpt-oss	`13.51% <0.00%> (+<0.01%)`	⬆️
examples-hf_ptq	`23.19% <0.00%> (-0.05%)`	⬇️
examples-llm_distill	`13.59% <0.00%> (-0.01%)`	⬇️
examples-llm_eval	`17.24% <0.00%> (-0.01%)`	⬇️
examples-llm_qat	`17.48% <0.00%> (-0.01%)`	⬇️
examples-llm_sparsity	`15.76% <0.00%> (+<0.01%)`	⬆️
examples-megatron_bridge	`26.43% <50.00%> (-0.15%)`	⬇️
examples-specdec_bench	`13.29% <0.00%> (+<0.01%)`	⬆️
examples-speculative_decoding	`17.47% <0.00%> (-0.07%)`	⬇️
examples-torch_onnx	`21.40% <0.00%> (-0.01%)`	⬇️
examples-torch_trt	`15.24% <0.00%> (+<0.01%)`	⬆️
examples-vllm_serve	`15.09% <0.00%> (+<0.01%)`	⬆️
gpu	`59.21% <50.00%> (+43.69%)`	⬆️
regression	`15.10% <0.00%> (-0.01%)`	⬇️
unit	`60.51% <25.00%> (+0.02%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2026-10-01T20:41:46Z

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-21 06:07 UTC

… copy_non_model_files (#2607) ### What does this PR do? Type of change: new feature ### Summary of the series Every pipeline that quantizes, prunes, sparsifies or distills a Hugging Face model now starts the same way: **make the checkpoint local, then load from it**. `ensure_local_checkpoint` returns a local directory as-is and downloads a Hub model ID in full if it is not cached yet -- once, on rank 0 -- and the model, tokenizer and processor are then loaded from that local directory. Before, a Hub ID went straight to `from_pretrained`, which fetched only the files loading needed, so every later step that needed the source as files found its own way back to it: `hf_ptq` through a hand-rolled cache lookup (`_resolve_model_path`, whose `TRANSFORMERS_CACHE` glob picked the lexicographically largest commit hash rather than `refs/main`) and a hardcoded download allowlist, regenerating tokenizer files along the way; the Megatron export by fetching only `*.py` for a Hub ID; other callers of `export_hf_checkpoint` not at all. With the source on local disk, exports follow one rule: an export writes the **model files** -- weights in any format and the metadata describing them -- and carries every other source file (**non-model files**: tokenizer, processor, remote code, chat templates, README, ...) over verbatim, from **local disk only**. Merge in order: 1. #2607 — Add `ensure_local_checkpoint` to load models from local disk, and `copy_non_model_files` ← **this PR** 2. #2608 — HF exporters write off-index safetensors and non-model files 3. #2609 — `hf_ptq`: load the model from local disk, downloading it first if needed 4. #2610 — `megatron_bridge`: load the HF model from local disk; Megatron export copies its non-model files 5. #2611 — `llm_sparsity`: load the model from local disk, downloading it first if needed ### This PR [1/5] Adds the two building blocks to `modelopt/torch/utils/plugins/hf_checkpoint_utils.py`, re-exported from `modelopt.torch.export`. Nothing calls them yet; #2609–#2611 make `ensure_local_checkpoint` the first step of each pipeline. - **`ensure_local_checkpoint(model_name_or_path, group=None) -> (hub_model_id, local_checkpoint_path)`** makes the whole checkpoint local before anything loads it: a local directory as-is, a Hub ID downloaded in full if not cached yet (guessing which files are redundant would risk dropping ones the checkpoint needs). Load the model, tokenizer and processor from `local_checkpoint_path`; `hub_model_id` is the Hub ID, or `None` for a local path, for records that must stay valid on other hosts. - Under `torch.distributed`, rank 0 of `group` resolves it and broadcasts the result -- or its error, so a failed download raises everywhere rather than stranding the waiting ranks. Rank 0 decides for local paths too, since ranks checking for themselves could disagree about a node-local directory. Without a `group`, every rank waits on a temporary gloo group whose 24 h timeout outlasts a large download; with one, the caller's group carries the wait. A rank that cannot see the resolved path raises, asking for shared storage. - `snapshot_download` silently returns a possibly incomplete cached snapshot when the Hub is unreachable. With `HF_HUB_OFFLINE` the cache is used quietly (the user's choice); with the Hub unexpectedly unreachable it is used with a warning, or the Hub's error is raised if nothing is cached. An input that is neither a local directory nor a Hub model ID the Hub can access -- a mistyped local path, say -- raises a `ValueError` saying so, chained to the Hub's error; a gated repo still raises `GatedRepoError`. - **`copy_non_model_files(source, export_dir)`** copies every file of a local checkpoint, recursively, except weight files in any format, top-level export metadata (`config.json`, `hf_quant_config.json`, stale quant configs, `.experiment.json`), files the export already wrote, hidden directories (`.git`, `.cache`), and links that resolve outside the checkpoint or its Hub `blobs/` (checked with the existing `resolve_checkpoint_file`). It never writes through a link already in the export, removes a partial copy when one fails, and skips `export_dir` when it lies inside the source (e.g. an earlier export). It returns `/`-separated paths on every platform. **What happens to the existing copy helpers.** `copy_non_model_files` replaces both copy rules the series' callers use today. #2609 moves `hf_ptq` off `copy_non_safetensor_files_from_ckpt`. #2610 moves the Megatron export onto `copy_non_model_files` (dropping its Hub-fetching `copy_hf_ckpt_remote_code` branch), moves `prune_minitron` and `export_distilled_megatron_to_hf` off `copy_hf_ckpt_remote_code`, and deprecates `copy_non_safetensor_files_from_ckpt`. `copy_hf_ckpt_remote_code` stays for now: puzzletron's `init_child_from_parent` still uses it, and migrating puzzletron is out of scope for this series. Also: module constants move to the top next to a note defining model versus non-model files; `off_index_safetensors_files` / `copy_off_index_safetensors` accept a `None` source and report what they copy; comments stop calling off-index weight files "sidecars". ### Usage ```python from modelopt.torch.export import copy_non_model_files, ensure_local_checkpoint hub_model_id, local_checkpoint_path = ensure_local_checkpoint("Qwen/Qwen3-8B") # every rank calls it model = AutoModelForCausalLM.from_pretrained(local_checkpoint_path) ... copy_non_model_files(local_checkpoint_path, export_dir) # after the export wrote its own files ``` ### Testing `pytest tests/unit/torch/utils tests/unit/torch/export tests/examples/hf_ptq` (GPU end-to-end and `test_dataset_utils` excluded): 684 passed. New tests cover the copy rule (weights, owned metadata, subdirectories, hidden dirs, links inside/outside/dangling, Hub `blobs/`), online/offline/unreachable/missing-repo resolution, and -- over two real gloo processes -- rank-0-only download, error propagation, an explicit group, a subgroup without global rank 0, and a local path hidden from one rank. The resolution paths were also checked live against the Hub. CPU only (torch 2.12, transformers 5.9, huggingface_hub 1.19); `transformer_engine` cannot be imported in this environment, so GPU and Megatron tests are left to CI. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ (additive) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (entries land with the behaviour changes in #2608–#2610) - Did you get Claude approval on this PR?: ❌ ### Additional Information Stacked series; this is the base.  ## Summary by CodeRabbit * **New Features** * Checkpoint exports can resolve local and Hub-hosted checkpoints, using cached snapshots when offline or Hub access is unavailable. Missing offline cache entries are reported as errors. * Exports can copy eligible non-model checkpoint files without overwriting existing export files. Hidden directories and unsafe or dangling links are excluded. * Distributed exports coordinate checkpoint resolution across processes and report failures, including when other processes cannot access the resolved path. * Safetensors files not listed in an index are treated as standalone weight files and can be copied into exports when eligible.  --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu · 2026-10-19T20:07:28Z

/ok to test b8de537

cjluo-nv

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Nudge: prior correctness and coverage concerns are resolved, but embedded agent directives in earlier CodeRabbit comments prevent automated approval.

Needs action:

Confirm human sign-off on this review; earlier CodeRabbit comments contain embedded agent prompts, which were ignored as untrusted input.

No action needed:

✔️ Resolved since the last review: local-source distillation exports, cached-module monkeypatching, quantize/prune Hub-ID and saved-tokenizer regressions, and Hub-ID API documentation.
The PR body justifies shared checkpoint utilities over direct partial Hub downloads and Transformers artifact regeneration; no second production subsystem is introduced.
Existing deprecation test edits are justified and retain all copying assertions.
GPU regressions were inspected but could not be executed with the available tools.

shengliangxu · 2026-10-19T20:23:21Z

/ok to test fcbacc6

kevalmorabia97 · 2026-10-19T20:24:27Z

can we fold this case into any existing tests instead of adding new ones since these tests are too slow to run?

Folded in 6723df2, and the standalone distill test is removed. test_distill_llm_hf_export now checks that both its inline export and the standalone re-export read the student from the copy ensure_local_checkpoint resolved; the re-export is passed the Hub ID. test_distill_vlm passes the Hub ID to its export step. A shared fake_hub_checkpoint fixture in conftest.py does the patching; it applies to single-rank steps, the only ones running in the pytest process. With more ranks the tests run on the local path as before.

shengliangxu · 2026-10-20T00:54:50Z

/ok to test 6723df2

The Megatron unified HF export copied every top-level non-safetensors file of a local source -- *.bin weights included -- but only *.py of a Hub ID, then regenerated the generation config, tokenizer and image processor. Now it writes its config and copies the source's non-model files with copy_non_model_files, from a local source only; a Hub ID gets a warning, as in the HF exporters. That is a breaking change for callers passing a Hub ID, who no longer get those files. The megatron_bridge quantize, export, prune and distillation-export scripts therefore make the HF model local first with ensure_local_checkpoint, keeping their argument as given and reading from hf_model_path. The tokenizer path recorded in saved Megatron checkpoints stays the Hub ID, which remains valid on other hosts. prune_minitron and the distillation VLM export -- which every rank reaches -- copy the non-model files from the main rank only, and the distillation LLM export now copies them too. copy_non_safetensor_files_from_ckpt lost its last internal user and is deprecated in favour of copy_non_model_files. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…D as hub_model_id As in hf_ptq: name the local checkpoint directory and the Hub identifier for what they are, with a student_ prefix in the distillation export. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…model ID Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

export_llm_to_hf and save_vlm_to_hf now copy the source's non-model files, which reads the student from local disk only, but distill.py still passed --student_hf_path unresolved: a Hub ID would fail the export after training. Resolve it with ensure_local_checkpoint at the start of main() and build, probe and export the student from that local copy; --student_hf_path stays as given. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

0.48 has branched (release/0.48.0), so this lands in 0.49: restore the 0.48 entry for the HF exporters to its released wording and file the Megatron export's feature, breaking-change and deprecation entries under 0.49. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

…es for a Hub model ID Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

The example tests are slow, so check the Hub-ID handling inside the runs that already exist instead of adding new ones: quantize/export and prune take the Hub ID, distillation's inline export and both standalone distillation exports read the resolved copy, and the quantized checkpoint records the Hub ID as its tokenizer. fake_hub_checkpoint only patches a single-rank step, which runs in the pytest process; with more ranks the tests run on the local path as before. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Saving a Megatron checkpoint loads the tokenizer it records, here the Hub ID, which failed for the fake org/tiny-model in CI. Map AutoTokenizer.from_pretrained(hub_model_id) to the resolved copy, as the Hugging Face cache does for a real Hub ID. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu · 2026-10-20T22:42:03Z

/ok to test 4f05440

Megatron-LM calls AutoTokenizer.from_pretrained with pretrained_model_name_or_path= as a keyword, which the wrapper only took positionally. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu · 2026-10-20T23:45:54Z

/ok to test e1fc2c7

cjluo-nv

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: folding the regressions into existing tests dropped coverage of the previously broken inline VLM export caller.

Needs action:

💬 Author folded the regressions into existing tests — restore inline VLM resolved-source coverage in test_distill_vlm as described in the inline comment, without adding another training run.
Confirm human sign-off; earlier CodeRabbit comments contain embedded agent directives, which were ignored as untrusted input.

No action needed:

✔️ Resolved: local-source distillation callers, cached-module patching, quantize/prune Hub-ID and saved-tokenizer assertions, and Hub-ID API documentation.
Reusing existing checkpoint utilities is justified over direct huggingface_hub.snapshot_download, Transformers artifact regeneration, and custom filesystem copying; no second production subsystem is introduced.
Deprecation-test edits retain all copying assertions. Folding slow tests is reasonable, but the VLM inline branch still needs equivalent coverage.
GPU tests were inspected, not executed.

cjluo-nv · 2026-10-21T00:51:46Z


-    # Separately convert the distilled Megatron checkpoint to HF with the standalone export script
+    # Separately convert the distilled Megatron checkpoint to HF with the standalone export script,
+    # passing the student as a Hub ID (see ``fake_hub_checkpoint``).


Bot comment.

The folded VLM regression now exercises only export_distilled_megatron_to_hf.py: the fixture is installed after training, and the preceding distill.py command does not pass hf_export_path. Consequently reverting distill.py's save_vlm_to_hf caller to args.student_hf_path would still pass this test, although that was one of the original post-training export failures covered by the removed parameterized regression. Preserve that coverage without another training run: serve the checkpoint and patch both scripts before training, enable inline HF export to a separate directory in the existing training command, and assert its resolved-copy marker before running the existing standalone export. Retain the standalone export and vision-weight assertions.

cjluo-nv

approve to unblock. @ChenhanYu probably need to approve as well.

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from a9eaef0 to 9d11a30 Compare October 1, 2026 01:32

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from a6cda6d to b1d1c21 Compare October 1, 2026 01:32

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from 9d11a30 to 9ce3570 Compare October 1, 2026 18:37

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from b1d1c21 to d765177 Compare October 1, 2026 18:37

shengliangxu added this pull request to stack #2621 October 1, 2026 18:53

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from 9ce3570 to b7e35e3 Compare October 1, 2026 20:32

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from d765177 to 36aa3ad Compare October 1, 2026 20:32

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from 36aa3ad to bf67146 Compare October 1, 2026 22:02

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from b7e35e3 to 27857e6 Compare October 1, 2026 22:02

shengliangxu changed the title ~~[4/5] Megatron export: copy non-model files from local disk only~~ [4/5] megatron_bridge: load the HF model from local disk; Megatron export copies its non-model files Oct 1, 2026

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from bf67146 to ee77aa3 Compare October 2, 2026 01:41

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch 2 times, most recently from 213ed0e to 0b6a3f4 Compare October 5, 2026 21:26

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch 2 times, most recently from eec4bd0 to 96b8404 Compare October 5, 2026 22:00

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from 0b6a3f4 to 4095979 Compare October 5, 2026 22:00

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch 2 times, most recently from 7fd6119 to fb32c1d Compare October 6, 2026 21:02

shengliangxu force-pushed the shengliangx/model-download-3-hf-ptq branch from 2ce59dd to 2d21f50 Compare October 6, 2026 21:02

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from fb32c1d to 6de984f Compare October 7, 2026 00:17

shengliangxu marked this pull request as ready for review October 7, 2026 05:59

cjluo-nv reviewed Oct 19, 2026

View reviewed changes

kevalmorabia97 requested review from kevalmorabia97 and removed request for a team and kaix-nv October 19, 2026 20:21

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from b8de537 to fcbacc6 Compare October 19, 2026 20:23

kevalmorabia97 reviewed Oct 19, 2026

View reviewed changes

Comment thread tests/examples/megatron_bridge/test_prune_minitron.py

kevalmorabia97 reviewed Oct 19, 2026

View reviewed changes

Comment thread tests/examples/megatron_bridge/test_quantize_export.py Outdated

shengliangxu added 10 commits October 20, 2026 22:41

Document that export_mcore_gpt_to_hf skips non-model files for a Hub …

d294c7d

…model ID Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Patch the runner's cached distill module in the resolved-checkpoint test

fb3be2e

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Test that quantize, export and prune run a Hub ID from its local copy

a5d9502

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Document that the vLLM fake-quant Megatron export skips non-model fil…

d1bd813

…es for a Hub model ID Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

shengliangxu force-pushed the shengliangx/model-download-4-megatron branch from 6723df2 to 4f05440 Compare October 20, 2026 22:41

Accept the keyword form in the fake Hub ID's tokenizer loader

e1fc2c7

Megatron-LM calls AutoTokenizer.from_pretrained with pretrained_model_name_or_path= as a keyword, which the wrapper only took positionally. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

cjluo-nv reviewed Oct 21, 2026

View reviewed changes

shengliangxu requested a review from meenchen October 21, 2026 00:52

cjluo-nv approved these changes Oct 21, 2026

View reviewed changes

shengliangxu merged commit 800d88a into main Oct 21, 2026
75 of 77 checks passed

shengliangxu deleted the shengliangx/model-download-4-megatron branch October 21, 2026 06:07

Conversation

shengliangxu commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Summary of the series

This PR

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Sep 30, 2026

Uh oh!

coderabbitai Bot commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Sequence Diagram(s)

Uh oh!

codecov Bot commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shengliangxu commented Oct 19, 2026

Uh oh!

cjluo-nv left a comment • edited by shengliangxu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shengliangxu commented Oct 19, 2026

Uh oh!

kevalmorabia97 Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

shengliangxu Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

shengliangxu commented Oct 20, 2026

Uh oh!

shengliangxu commented Oct 20, 2026

Uh oh!

shengliangxu commented Oct 20, 2026

Uh oh!

cjluo-nv left a comment

Choose a reason for hiding this comment

Uh oh!

cjluo-nv Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

cjluo-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

shengliangxu commented Sep 30, 2026 •

edited

Loading

coderabbitai Bot commented Sep 30, 2026 •

edited

Loading

codecov Bot commented Oct 1, 2026 •

edited

Loading

github-actions Bot commented Oct 1, 2026 •

edited

Loading

cjluo-nv left a comment •

edited by shengliangxu

Loading