Repository navigation
[OMNIML-6028] [3/5] hf_ptq: load the model from local disk, downloading it first if needed - #2609
Merged
Merged
Conversation
|
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. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (1)
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 6 remain after this review. 📝 WalkthroughWalkthroughHF PTQ resolves input checkpoints to local paths for model and tokenizer loading. Export handling uses the local checkpoint for source-file operations. Export calls no longer restore or save tokenizer state separately. ChangesHF PTQ checkpoint and export flow
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant PTQMain
participant ensure_local_checkpoint
participant ModelAndTokenizer
PTQMain->>ensure_local_checkpoint: Resolve input checkpoint
ensure_local_checkpoint-->>PTQMain: Return Hub ID and local path
PTQMain->>ModelAndTokenizer: Load model and tokenizer from local path
Merge Risk: ⚪ Minimal · up to No merge-blocking issue is established for these test changes; they are ready for normal checks. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
This was referenced Sep 30, 2026
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 1, 2026 01:32
a9eaef0 to
9d11a30
Compare
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 1, 2026 18:37
9d11a30 to
9ce3570
Compare
shengliangxu
added this pull request to stack #2621
October 1, 2026 18:53
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2609 +/- ##
==========================================
- Coverage 72.16% 72.15% -0.01%
==========================================
Files 635 635
Lines 70643 70643
==========================================
- Hits 50977 50970 -7
- Misses 19666 19673 +7
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
2 times, most recently
from
October 1, 2026 22:02
b7e35e3 to
27857e6
Compare
shengliangxu
changed the title
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 2, 2026 01:41
27857e6 to
213ed0e
Compare
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 5, 2026 21:26
213ed0e to
0b6a3f4
Compare
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
3 times, most recently
from
October 6, 2026 21:02
2ce59dd to
2d21f50
Compare
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 7, 2026 00:17
2d21f50 to
ebf6f65
Compare
shengliangxu
added a commit
that referenced
this pull request
Oct 7, 2026
… 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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
marked this pull request as ready for review
October 7, 2026 05:59
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 7, 2026 06:05
ebf6f65 to
a01879b
Compare
Base automatically changed from
shengliangx/model-download-2-hf-exporters
to
main
October 7, 2026 07:59
shengliangxu
added a commit
that referenced
this pull request
Oct 7, 2026
…2608) ### What does this PR do? Type of change: bug fix, 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` 2. #2608 — HF exporters write off-index safetensors and non-model files ← **this PR** 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 [2/5] `export_hf_checkpoint` put the tensors of off-index safetensors (GLM-4.7's `mtp.safetensors`) into `exclude_modules` but left copying the files to the `hf_ptq` example, so any other caller exported exclusions for weights the checkpoint lacked. It likewise left the source's non-model files to the caller, while its streaming paths copied only top-level `*.py`. Every HF export path now does both itself, from the model's source checkpoint on local disk: the regular, streaming and FSDP2 paths; the layerwise path, into the exporter's own directory (a recipe's `layerwise.export_dir` can set it apart from `export_dir`); and the vLLM fake-quant export. `save_non_weight_artifacts` copies the non-model files in place of its `*.py` glob. A Hub-ID source gets a warning instead of a fetch, since its cache holds only what loading needed. Until #2609, `hf_ptq` still runs its own copy afterwards -- harmless, as files already present are skipped -- and prints that warning for a Hub-ID `--pyt_ckpt_path`. ### Usage No API change; `export_hf_checkpoint(model, export_dir=...)` on a model loaded from a local directory now also writes the source's non-model and off-index files. ### Testing Same suite as #2607: 688 passed. New tests: the regular path copies non-model files (tiny Llama loaded from a local directory), a Hub-ID source warns without fetching, and the layerwise path writes `mtp.safetensors` into the exporter's directory even when `export_dir` differs. A CPU end-to-end FP8 export of `hf-internal-testing/tiny-random-LlamaForCausalLM` carried the tokenizer and README. 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?: ✅ (exports gain files; nothing is removed) - 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)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information Stacked on #2607. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Hugging Face and vLLM exports now preserve off-index safetensors and other non-model files from local source checkpoints, including files not listed in the checkpoint index. * Layerwise exports place copied safetensors in the correct output directory. * Exports from Hub sources warn when source files cannot be copied locally, without downloading them. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A Hub model ID reached the loader as-is, so every later step that needed the source checkpoint as files had its own way back to a directory: _resolve_model_path (whose TRANSFORMERS_CACHE glob picked the lexicographically largest commit hash rather than refs/main), a hardcoded sidecar allowlist, _resolved_local_dir, and hf_hub_download fallbacks. Now main() calls ensure_local_checkpoint once, leaving --pyt_ckpt_path as given and storing hf_model_path (the local copy every step reads) and hf_model_name (the Hub ID, or None). The layerwise resume directory is still named after the Hub ID. The unified exporters copy the source's non-model files themselves, so the example no longer does -- except after the deprecated TensorRT-LLM export, which does not. It also stops regenerating tokenizer and processor files: the source's are carried over verbatim. That fixes exports shipping pad_token set to the EOS token, which get_tokenizer sets for calibration, and removes the default padding side and pad token threaded through six functions only to undo that before saving. The notebooks do the same. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…model_id hf_model_path read as a Hub path when it is the local directory, and hf_model_name did not say it is the Hub identifier; name both for what they are. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…t call test_hf_ptq_model_type, new on main, still passed default_padding_side and default_pad_token, which this branch removes from post_quantize. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
force-pushed
the
shengliangx/model-download-3-hf-ptq
branch
from
October 7, 2026 07:59
a01879b to
c61b877
Compare
cjluo-nv
reviewed
Oct 7, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Nudge: the shared helpers are well tested, but the new hf_ptq orchestration lacks a regression test.
Needs action:
- Add a mocked Hub-ID regression in
tests/examples/hf_ptq/test_hf_ptq_args.pyprovingmain()resolves the checkpoint before loading, preservespyt_ckpt_path, and passes the local path through loading and TensorRT-LLM sidecar copying.
No action needed:
- Existing test edits are justified: removed tests cover deleted helpers, the library tests cover their replacements, and argument updates follow the changed signatures.
- The change reuses the checkpoint utilities and unified exporters rather than introducing duplicate logic. Tests were inspected, not executed.
…LM export copies from it Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
@cjluo-nv bot review: added the regression tests in 265631c (
Each test fails when the corresponding wiring is broken: copying from |
cjluo-nv
reviewed
Oct 7, 2026
cjluo-nv
left a comment
cjluo-nv
left a comment
Collaborator
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
LGTM: the prior orchestration-test concern is resolved; human sign-off remains required for the justified existing-test removals and updates.
Needs action:
- Confirm the existing-test edits: removed cases cover deleted helpers, shared-library tests cover their replacements, and signature updates follow the changed behavior.
No action needed:
- ✔️ Resolved since the last review: new mocked regressions verify checkpoint resolution before loading, preservation of the Hub ID, and TensorRT-LLM sidecar copying from the local directory.
- Tests inspected, not executed.
cjluo-nv
reviewed
Oct 7, 2026
| # --pyt_ckpt_path stays as given. Every later step reads the local copy of the whole | ||
| # checkpoint in local_checkpoint_path; hub_model_id is the Hub ID, or None for a local | ||
| # path. | ||
| args.hub_model_id, args.local_checkpoint_path = ensure_local_checkpoint( |
Collaborator
There was a problem hiding this comment.
ensure_local_checkpoint: will it point to the local HF downloaded cache or download the model to a particular location like /tmp?
Collaborator
Author
There was a problem hiding this comment.
It uses the local Hugging Face cache. It calls snapshot_download with no cache_dir/local_dir, so files land in the standard hub cache (HF_HUB_CACHE, default ~/.cache/huggingface/hub, following HF_HOME), and it returns the snapshot directory there (.../models--org--name/snapshots/<commit>). Nothing goes to /tmp, and a local directory is returned as-is.
If the model is already in the cache, it is not downloaded again: the cached snapshot is used directly, and only files missing from it are fetched.
The download behavior is the same as before this PR: from_pretrained("org/model") already downloaded into this same cache. The only difference is that this is now an explicit first step of the pipeline, before any loading. It makes the whole checkpoint local up front, instead of each later step (model, tokenizer, processor, export) fetching what it needs on its own.
Under torch.distributed, rank 0 does the download and the other ranks use the same path. The cache therefore has to be on storage every node can see; if a rank can't see the path, it raises with a hint.
cjluo-nv
approved these changes
Oct 7, 2026
Contributor
|
shengliangxu
added a commit
that referenced
this pull request
Oct 7, 2026
…rst if needed (#2611) ### What does this PR do? Type of change: bug fix ### 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` 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. #2611 — `llm_sparsity`: load the model from local disk, downloading it first if needed ← **this PR** 5. #2610 — `megatron_bridge`: load the HF model from local disk; Megatron export copies its non-model files ### This PR [4/5] **Loading.** `hf_sa.py` and `export_hf_ckpt.py` now make the model local first with `ensure_local_checkpoint` -- a local directory as-is, a Hub ID downloaded once -- and load the model and tokenizer from that copy (`args.local_checkpoint_path`), keeping their argument as given. That is also what lets `export_hf_checkpoint`, which carries the source's non-model files only from local disk (#2608), include them. **Exporting.** `hf_sa.py` no longer saves its tokenizer into the export: it sets `pad_token` to EOS only for generation, so the source's files are the right ones. `export_hf_ckpt.py` keeps its save, after the export, because it can add a pad token and resize the embeddings. ### Usage No flag change. ### Testing Both scripts byte-compile; their end-to-end tests need a GPU and are left to CI. No library code changes in this PR. ### 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?: ✅ - 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (no user-facing change beyond the series' entries in #2608 and #2610) - Did you get Claude approval on this PR?: ❌ ### Additional Information Based on `main` (#2607–#2609 are merged); #2610 is stacked on it. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Model export examples load checkpoints from a local copy while retaining the original model reference and, when available, its Hub ID. * Weight-sparsity exports save tokenizer updates, including padding-token changes. Attention-sparsity exports no longer save the tokenizer separately. * Attention-sparsity exports run the model in inference mode. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
shengliangxu
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: bug fix
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_checkpointreturns 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_ptqthrough a hand-rolled cache lookup (_resolve_model_path, whoseTRANSFORMERS_CACHEglob picked the lexicographically largest commit hash rather thanrefs/main) and a hardcoded download allowlist, regenerating tokenizer files along the way; the Megatron export by fetching only*.pyfor a Hub ID; other callers ofexport_hf_checkpointnot 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:
ensure_local_checkpointto load models from local disk, andcopy_non_model_fileshf_ptq: load the model from local disk, downloading it first if needed ← this PRmegatron_bridge: load the HF model from local disk; Megatron export copies its non-model filesllm_sparsity: load the model from local disk, downloading it first if neededThis PR [3/5]
Loading.
main()now makes the model local before anything else:ensure_local_checkpoint(--pyt_ckpt_path)returns a local directory as-is or downloads a Hub ID once, and every later step -- loading the model, tokenizer and processor, the MXFP4 cast, and the export's copy of the source's files -- reads that local copy.--pyt_ckpt_pathstays as given; the local directory isargs.local_checkpoint_path, andargs.hub_model_id(the Hub ID, orNone) still names the layerwise resume directory. This replaces the per-step ways back to the source --_resolve_model_path, a hardcoded sidecar allowlist,_resolved_local_dir,hf_hub_downloadfallbacks.Exporting. The unified exporters copy the source's non-model files themselves (#2608), so the example no longer does -- except after the deprecated TensorRT-LLM export, which does not. It also stops regenerating tokenizer and processor files: the source's are carried over verbatim. So a
pad_tokenthatget_tokenizersets for calibration (the EOS token, for a model without one) no longer leaks into the exported tokenizer; #2597 already stopped it overwriting an existing pad token. It also removes the default padding side and pad token threaded through six functions only to undo that before saving.The three notebooks do the same: make the model local first, load from it, and leave the tokenizer files to the export. Mostly deletions (+84 / −466).
Usage
No flag change.
--pyt_ckpt_path org/modeldownloads the whole repo once before loading; to skip files, download it yourself (e.g.hf download org/model --exclude 'original/*') and pass the local path.Testing
Same suite as #2607: 685 passed (the
hf_ptqtests moved to the library in #2607 or were dropped with the code they covered). Notebooks validated as JSON. CPU only (torch 2.12, transformers 5.9, huggingface_hub 1.19);transformer_enginecannot 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 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.).--pyt_ckpt_pathunchanged; exports carry the source's files as-is)CONTRIBUTING.md: N/Apad_tokenfix is [OMNIML-5935] Fix hf_ptq overwriting valid pad_token with eos_token #2597's entry; the verbatim copy is [OMNIML-6031] [2/5] HF exporters write off-index safetensors and non-model files #2608's)Additional Information
Stacked on #2608.
Summary by CodeRabbit