[nvbug 6289151, nvbug 6301817] Fix exported Step layer type and RoPE metadata#1693
[nvbug 6289151, nvbug 6301817] Fix exported Step layer type and RoPE metadata#1693meenchen wants to merge 1 commit into
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a sanitizer that trims Hugging Face ChangesLayer-types sanitization for HF export
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1693 +/- ##
==========================================
+ Coverage 67.72% 76.11% +8.39%
==========================================
Files 511 511
Lines 56168 57764 +1596
==========================================
+ Hits 38037 43969 +5932
+ Misses 18131 13795 -4336
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:
|
meenchen
force-pushed
the
fix-step35-flash-trtllm-layer-types
branch
2 times, most recently
from
0bccd7b to
97687b3
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Small, focused bug fix (+121/-5, 4 files) for Step-3.5-Flash HF export where layer_types includes trailing MTP/next-token-prediction entries, breaking Transformers 5.x's len(layer_types) == num_hidden_layers validation in TRT-LLM.
Correctness:
sanitize_hf_config_for_deploymentis conservative: it only trims trailing entries when the mismatch is exactly explained bynum_nextn_predict_layers, and leaves unexplained mismatches untouched._as_nonnegative_intcorrectly excludesbool(which is anintsubclass) before the integer check._get_num_nextn_predict_layersfalls back config_data → model.config →_mtp_layer_prefixeslength; when it returnsNone,None == <int>isFalse, so no trim happens. Anum_nextn_predict_layers == 0value can never trigger a trim because the equal-length case already returns early. All edges consistent.- Integration point in
export_hf_checkpointis correct: called aftersave_pretrainedwritesconfig.jsonand before the config is re-written to disk; trimminglayer_types[:num_hidden_layers]keeps the head decoder entries, matching the documented MTP-entries-come-last assumption. - The
plugins/__init__.pyreorder (movinghf_checkpoint_utilsstar-import ahead ofvllm_fakequant_hf) is benign; both are simple plugin star-imports with no cross-dependency. Private helpers (_as_nonnegative_int,_get_num_nextn_predict_layers) won't leak viaimport *.
Tests: Three unit tests cover the config-nextn-count trim, the model.config-nextn-count trim, and the keep-unexplained-mismatch case, including warning emission. Meaningful coverage of the core branches.
No licensing changes, no new subsystem/abstraction, no prompt-injection attempts in the PR metadata. Full TRT-LLM serve wasn't run (checkpoint not visible from cluster), which is an acceptable limitation noted by the author for a config-only sanitizer with unit coverage.
Complex PR: 1 existing test file modified or removed. Looping in a human for approval.
meenchen
added
the
cherry-pick-0.45.0
meenchen
force-pushed
the
fix-step35-flash-trtllm-layer-types
branch
from
97687b3 to
914cf1e
Compare
| from .model_utils import get_language_model_from_vl, is_multimodal_model | ||
| from .moe_utils import _export_fused_experts | ||
| from .plugins import SpeculativeDecodingExporter, has_spec_opt | ||
| from .plugins import SpeculativeDecodingExporter, has_spec_opt, sanitize_hf_config_for_deployment |
There was a problem hiding this comment.
Hard import of a symbol from an optionally-guarded plugin.
This now hard-imports sanitize_hf_config_for_deployment, which lives in hf_checkpoint_utils. In plugins/__init__.py that module is imported inside import_plugin(...) — a guard whose whole purpose is to tolerate the module failing to import. By contrast, SpeculativeDecodingExporter / has_spec_opt come from hf_spec_export, which is imported unguarded.
So if hf_checkpoint_utils ever fails under the guard, the symbol silently won't exist and this line raises ImportError, breaking all HF export — not just the Step path. In practice the deps (huggingface_hub, safetensors, tqdm) are always present for export, so the risk is low, but the import_plugin guard is partially defeated by taking a hard dependency on one of its symbols.
Note the import reorder in plugins/__init__.py doesn't close this gap — it only changes binding order; it doesn't make the symbol exist if the guarded import itself fails. Suggested fixes:
- guard this import defensively (
try/except ImportError→ no-op fallback), or - if
hf_checkpoint_utilsis genuinely required for export, drop theimport_pluginguard around it so failures are loud rather than silently amputating a symbol.
There was a problem hiding this comment.
Addressed in d83dcb8. I removed the import_plugin guard around hf_checkpoint_utils in plugins/init.py, so sanitize_hf_config_for_deployment is no longer a hard import from an optionally amputated plugin symbol. This matches the existing direct hf_checkpoint_utils imports in export code: failures are now loud instead of silently dropping the symbol. Validated with ruff-format, ruff-check, py_compile, and the focused hf_checkpoint_utils unit tests.
There was a problem hiding this comment.
Updated in 43a9abc. I changed the resolution from making hf_checkpoint_utils unguarded to preserving the import_plugin guard and adding a no-op fallback for sanitize_hf_config_for_deployment when the guarded import does not provide it. That keeps modelopt.torch.export import paths resilient to missing optional deps while keeping the direct unified_export_hf import stable. Validated with ruff-format, ruff-check, py_compile, and the focused hf_checkpoint_utils unit tests.
meenchen
force-pushed
the
fix-step35-flash-trtllm-layer-types
branch
from
914cf1e to
d83dcb8
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re-review of Step-3.5-Flash HF export layer_types fix (now +159/-5, 4 files).
Previous critical comment (Edwardf0t1): RESOLVED. 💬 Author replied: removed the import_plugin guard around hf_checkpoint_utils in plugins/__init__.py, so sanitize_hf_config_for_deployment is now an unguarded import matching the existing hf_spec_export pattern; failures are loud rather than silently amputating the symbol. Verified in the current tree — the guard is gone and the import is now plain from .hf_checkpoint_utils import *. This matches the reviewer's suggested fix exactly, and the hard import in unified_export_hf.py:93 is now safe.
Correctness: The sanitizer is conservative — it only trims trailing layer_types when the surplus is exactly num_nextn_predict_layers, otherwise leaves the config untouched. _as_nonnegative_int correctly excludes bool. _get_num_nextn_predict_layers falls back config_data → model.config → _mtp_layer_prefixes count; the new _count_mtp_layer_prefixes correctly excludes broad prefixes ("mtp" → not counted, "mtp.layers.0" → counted). Integration point (after save_pretrained, before re-writing config.json) is correct.
Tests: Five unit tests now cover config-count trim, model.config-count trim, MTP-prefix counting, broad-prefix-only no-op, and the keep-unexplained-mismatch case — meaningful branch coverage.
No licensing changes; no prompt-injection in PR metadata.
Why nudge rather than approve: the resolution looks correct but warrants a quick human confirmation, and the author notes the full TRT-LLM serve path was not run (Step-3.5 checkpoint not visible from the cluster login node), so the end-to-end deployment fix is unverified — only unit coverage exists for the config-only sanitizer.
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/export/plugins/__init__.py`:
- Line 23: Restore guarded import semantics: replace the unconditional "from
.hf_checkpoint_utils import *" with a guarded import that uses the package's
import_plugin mechanism (or a try/except ImportError around importing
hf_checkpoint_utils) so third‑party optional extras don't cause hard import
failures; specifically, use
import_plugin("modelopt.torch.export.plugins.hf_checkpoint_utils") or wrap the
import of hf_checkpoint_utils in try/except, emit a warning when the extra is
missing, and ensure any names expected from hf_checkpoint_utils are either
imported when available or left undefined/placeholder so package import paths
remain resilient.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 74637461-2fc7-43aa-83c9-862d6ca48c18
📒 Files selected for processing (4)
modelopt/torch/export/plugins/__init__.pymodelopt/torch/export/plugins/hf_checkpoint_utils.pymodelopt/torch/export/unified_export_hf.pytests/unit/torch/export/test_hf_checkpoint_utils.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/export/plugins/hf_checkpoint_utils.py
- tests/unit/torch/export/test_hf_checkpoint_utils.py
- modelopt/torch/export/unified_export_hf.py
meenchen
force-pushed
the
fix-step35-flash-trtllm-layer-types
branch
2 times, most recently
from
43a9abc to
bdf0fe4
Compare
### What does this PR do? Type of change: Bug fix Extends the calibration/memory-probe `use_cache` guard to Step 3.7-style nested text configs. Step 3.7 remote code reads the language config under `model.config.text_config` directly and raises `AttributeError` when `use_cache` is absent during PTQ calibration with Transformers >5. This keeps the existing Step 3.5 behavior and applies the same temporary set/restore logic to the nested text config. ### Usage No API change. PTQ calibration continues to use the existing forward-loop path. ### Testing - `pre-commit run ruff-format --files modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `pre-commit run ruff-check --files modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `python -m py_compile modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `python -m pytest tests/unit/torch/utils/test_dataset_utils.py -k "disable_use_cache or iter_use_cache_configs or forward_loop_runs_under_disabled" -vv` ### Before your PR is "*Ready for review*" - 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?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information This is separate from PR #1693. Step 3.7 needs both fixes if both failure paths are exercised: this PR fixes PTQ calibration-time `use_cache` handling, while PR #1693 fixes exported config `layer_types` metadata for deployment config loading. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved handling of cache flags stored in nested model configuration objects: cache is reliably disabled during dataset operations and restored or removed afterward. * **Tests** * Added unit tests covering nested-config disabling, restoration/removal of cache flags post-operation, and deduplication when nested configs reference the same object. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re-review of the Step-3.5-Flash HF export layer_types fix (+171/-5, 4 files). Small, focused, correct config-sanitizer change.
Previous critical comment (Edwardf0t1 / CodeRabbit, line 93 & line 23 — hard import of an optionally-guarded plugin symbol): RESOLVED.
- 💬 Author replied (commit
43a9abca6): kept theimport_plugin("hf_checkpoint_utils")guard and added a conditional no-op fallback forsanitize_hf_config_for_deploymentwhen the guarded import doesn't provide it. Verified in the current tree:plugins/__init__.pydefinesdef sanitize_hf_config_for_deployment(...) -> None: return Noneonly when the symbol is absent fromglobals(). I confirmedimport_pluginswallowsModuleNotFoundError(and other exceptions with a warning), so the hard import inunified_export_hf.py:93is now always resolvable — real function when deps present, no-op when not — while optional-plugin semantics are preserved. This is a correct, clean resolution and both reviewers confirmed it.
Correctness: Sanitizer is conservative — trims trailing layer_types only when the surplus is exactly num_nextn_predict_layers (None == int → False, so missing metadata never trims). _as_nonnegative_int correctly excludes bool; _count_mtp_layer_prefixes correctly counts only *.layers.<n> prefixes (broad mtp excluded). Integration point (after save_pretrained writes config.json, before re-write) is correct.
Tests: Five unit tests cover config-count trim, model.config-count trim, MTP-prefix counting, broad-prefix-only no-op, and keep-unexplained-mismatch — meaningful branch coverage. Licensing clean (standard NVIDIA header). No prompt-injection in PR metadata.
Why nudge rather than approve: the resolution is correct, but the author notes the full TRT-LLM serve path was not run (Step-3.5 checkpoint not visible from the cluster login node), so the end-to-end deployment fix remains verified only by unit tests on the config-only sanitizer. A human with the checkpoint can give final sign-off.
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
meenchen
force-pushed
the
fix-step35-flash-trtllm-layer-types
branch
from
bdf0fe4 to
4f1e560
Compare
meenchen
changed the title
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re-review of the Step-3.5-Flash HF export config sanitizer (+283/-5, 4 files), now extended with the RoPE-metadata fix (nvbug 6301817) folded in from closed PR #1745.
Previous critical comment (Edwardf0t1 / CodeRabbit — hard import of an optionally-guarded plugin symbol): RESOLVED. 💬 Author replied (commit 43a9abca6): kept the import_plugin("hf_checkpoint_utils") guard in plugins/__init__.py and added a conditional no-op fallback def sanitize_hf_config_for_deployment(...) -> None: return None only when the symbol is absent from globals(). Verified in the current tree; since import_plugin swallows ModuleNotFoundError, the hard import in unified_export_hf.py:93 is always resolvable (real function when deps present, no-op otherwise). CodeRabbit also confirmed this as addressed. Clean resolution — dropping from concerns.
New content this round — RoPE fix (nvbug 6301817): _sanitize_llama3_rope_config fills missing rope_theta into llama3 rope_parameters/rope_scaling from config.json["rope_theta"] or model.config.rope_theta. Conservative: only touches dicts whose rope_type/type is "llama3", only when rope_theta is absent, and _get_rope_theta uses is not None so a real 0 would be respected. Non-llama3 and pre-existing values are left untouched. Runs before the existing layer_types trim in the same sanitize_hf_config_for_deployment call, correctly placed after save_pretrained and before re-writing config.json.
Layer_types fix (nvbug 6289151): unchanged from prior round and still correct/conservative — trims trailing entries only when surplus is exactly num_nextn_predict_layers (None == int → False), _as_nonnegative_int excludes bool, _count_mtp_layer_prefixes counts only *.layers.<n>.
Tests: Five layer_types tests plus six new RoPE tests (top-level fill, model.config fallback, legacy rope_scaling, existing-value-not-overwritten, non-llama3 ignored) — meaningful branch coverage for both fixes.
Licensing clean (existing standard NVIDIA Apache-2.0 headers, no license-file/SPDX changes). No prompt-injection in PR metadata.
Why nudge rather than approve: the prior critical is resolved and the new RoPE logic is well-unit-tested, but the author notes the full B200 PTQ + vLLM/TRT-LLM serve path was not run end-to-end (Step-3.5 checkpoint not visible from the cluster login node), so the actual deployment config-load fix is verified only by unit tests on the config-only sanitizer. A human with the checkpoint should give final sign-off. Minor (non-blocking): sanitize_hf_config_for_deployment's docstring still only documents the layer_types behavior and doesn't mention the new RoPE handling, and _get_rope_theta/_sanitize_llama3_rope_config lack docstrings unlike the other helpers.
3 participants
What does this PR do?
Type of change: Bug fix
Fixes two Step-3.5-Flash HF checkpoint export config issues that make deployment stacks fail during config loading before weights are loaded:
nvbug 6289151: exportedlayer_typescan include the main decoder layers plus MTP / next-token-prediction entries. Transformers 5.x validateslen(layer_types) == num_hidden_layers, so the sanitizer trims only trailing next-token-prediction entries when the mismatch is exactly explained bynum_nextn_predict_layersor supported MTP layer-prefix metadata.nvbug 6301817: Transformers 5.10.2 / vLLM validatesllama3RoPE metadata and requiresrope_thetainsiderope_parameters. The sanitizer fills missingrope_thetaintollama3rope_parameters/rope_scalingfromconfig.json["rope_theta"]ormodel.config.rope_theta.Unexplained
layer_typesmismatches and non-llama3RoPE metadata are left unchanged.Usage
N/A. Export Step-3.5 HF checkpoints as usual with
hf_ptq.py; the exportedconfig.jsonis sanitized before it is written back.Testing
/Users/weimingc/miniconda3/envs/modelopt/bin/pre-commit run ruff-format --files modelopt/torch/export/plugins/hf_checkpoint_utils.py tests/unit/torch/export/test_hf_checkpoint_utils.py/Users/weimingc/miniconda3/envs/modelopt/bin/pre-commit run ruff-check --files modelopt/torch/export/plugins/hf_checkpoint_utils.py tests/unit/torch/export/test_hf_checkpoint_utils.py/Users/weimingc/miniconda3/envs/modelopt/bin/pre-commit run mypy --files modelopt/torch/export/plugins/hf_checkpoint_utils.py tests/unit/torch/export/test_hf_checkpoint_utils.py/Users/weimingc/miniconda3/envs/modelopt/bin/pre-commit run mypy --all-filesPYTHONPATH=/Users/weimingc/workspace/nvbug/Model-Optimizer /Users/weimingc/miniconda3/envs/modelopt/bin/python -m pytest tests/unit/torch/export/test_hf_checkpoint_utils.py -vvMissing required keys in rope_parameters ... {'rope_theta'}rope_parameters["rope_theta"] == 500000Full B200 PTQ + vLLM serve was not run locally because the reported Step-3.5 checkpoint/export paths are not visible from the configured cluster login node.
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.).CONTRIBUTING.md: N/AAdditional Information
nvbug 6289151
nvbug 6301817
PR #1745 was closed because its RoPE metadata fix is folded into this PR.
Related MTP handling: #1532.