PTQ reliability fixes: checkpoint resume + exclude_modules sentinel corruption - #2129
Open
wyattearp wants to merge 14 commits into
Open
PTQ reliability fixes: checkpoint resume + exclude_modules sentinel corruption#2129wyattearp wants to merge 14 commits into
wyattearp wants to merge 14 commits into
Conversation
…orruption Three fixes, found while quantizing nvidia/Nemotron-Cascade-2-30B-A3B to NVFP4: 1. examples/hf_ptq/hf_ptq.py: --save_quantized_state/--restore_quantized_state flags. Lets a failed *export* be retried without redoing calibration. 2. modelopt/torch/utils/dataset_utils.py: checkpoint_every/checkpoint_fn hook in _forward_loop/create_forward_loop. Periodic-checkpoint mechanism (not yet wired into hf_ptq.py's calibration call -- follow-up, not done here). 3. modelopt/torch/export/quant_aware_conversion.py: build_reverse_name_mapper's sentinel-strip logic (_map) silently accepted a mangled sentinel instead of raising, when a reverse-rename rule's `.` (matched as "any path separator char") ate into the appended sentinel before it could be stripped. In our export this corrupted 77/77 exclude_modules entries with an unsubstituted placeholder (`\x00backbone.pt_name_sentinel`, not matching the real sentinel `\x00modelopt_name_sentinel` -- direct evidence a rename rule rewrote it). vLLM couldn't match any excluded layer, defaulted every quant-aware fused linear (e.g. Mamba's in_proj, a MergedColumnParallelLinear) to quantized-width allocation, and crashed loading the correctly-unquantized-but-unlabeled weight. Extracted the strip logic into _strip_sentinel_or_raise(), which now raises QuantConversionUnsupportedError on a failed strip -- exactly the exception build_reverse_name_mapper's own docstring already promised for this case. No caller changes needed: both call sites in unified_export_hf.py already wrap this in a broad try/except that falls back to safe in-memory names with a warning: this fix just makes that safety net actually trigger, converting a silent, deployment-time failure (bare AssertionError in vLLM, no useful diagnostic) into a loud, immediate one at export time. All three built test-first (Red -> Green). 26 total tests passing across the two touched test files (2 new for the checkpoint flags, 3 new for the sentinel fix), no regressions on pre-existing tests. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZby2EsEkVQKwwNKFDpwBm
Contributor
📝 WalkthroughWalkthroughThe PR adds quantized-state persistence to HF PTQ, NemotronH configuration restoration, strict reverse-name sentinel validation, and periodic checkpoint callbacks for dataset forward loops. ChangesQuantized-state persistence
NemotronH configuration compatibility
Reverse name mapping validation
Forward-loop checkpoint callbacks
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant hf_ptq
participant ModelOpt
participant Export
CLI->>hf_ptq: Select save or restore option
hf_ptq->>ModelOpt: Restore state or quantize model
ModelOpt-->>hf_ptq: Return restored or calibrated state
hf_ptq->>ModelOpt: Save calibrated state when requested
hf_ptq->>Export: Export model
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…on export Root cause: transformers now ships native NemotronH support (transformers/models/nemotron_h/configuration_nemotron_h.py) with a newer config schema -- layers_block_type as a real stored field, with hybrid_override_pattern/num_hidden_layers as read-only properties derived from it. Many NemotronH checkpoints on the Hub (e.g. nvidia/Nemotron-Cascade-2-30B-A3B) still ship their own older, bundled trust_remote_code configuration_nemotron_h.py written against the *opposite* schema: hybrid_override_pattern/num_hidden_layers as real fields, with layers_block_type as a computed property that has no setter. When ModelOpt loads such a checkpoint for PTQ, transformers' AutoConfig resolves the model via its own newer native class (no trust_remote_code needed once a model_type is natively registered), silently converting the config to the new schema in memory. save_pretrained's plain to_dict() only serializes __dict__, so neither hybrid_override_pattern nor num_hidden_layers (both properties on the native class) makes it into config.json -- but layers_block_type does, since it's a real field there. The result: an exported config.json in the new schema, sitting next to a copied configuration_nemotron_h.py file that only understands the old one. Loading the export back through its own bundled class then fails outright (AttributeError: property 'layers_block_type' has no setter) or silently loses the pattern/layer-count metadata. Reproduced on three separate nvidia/Nemotron-Cascade-2-30B-A3B NVFP4 exports in a row; each one required hand-patching config.json afterward. Fix: sanitize_hf_config_for_deployment now reconstructs hybrid_override_pattern from layers_block_type using the exact inverse of transformers' own NemotronHConfig._list_to_pattern mapping (linear_attention -> M, moe -> E, full_attention -> *, mlp -> -), restores num_hidden_layers as len(layers_block_type), and drops the now-redundant layers_block_type / mtp_layers_block_type fields the old schema's class doesn't expect. No-op for non-NemotronH exports, no-op if hybrid_override_pattern is already present (nothing was dropped), and warns without guessing if layers_block_type contains an unrecognized value. Built test-first (Red -> Green): 4 new tests (restores correctly, ignores other model types, no-ops when already legacy schema, warns on unrecognized layer type), 19/19 passing in the touched test file, 165/165 passing across the full export test suite -- no regressions. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BZby2EsEkVQKwwNKFDpwBm
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/hf_ptq/hf_ptq.py (1)
1352-1365: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftMove restore handling before calibration preparation.
At Line 1352, restore handling occurs after batch-size probing, calibration dataloader creation, and
pre_quantize. A retry with--restore_quantized_statestill depends on calibration inputs and performs model work before export. This can fail when the original dataset is unavailable and defeats the export-retry path. Resolve restore mode before calibration setup, or skip setup that is not required by export.The PR objective is to retry export without repeating calibration.
🤖 Prompt for 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. In `@examples/hf_ptq/hf_ptq.py` around lines 1352 - 1365, Move the `_restore_quantized_state_if_requested` handling before batch-size probing, calibration dataloader creation, and `pre_quantize`; in restore mode, bypass calibration-only setup and proceed directly to export. Preserve the existing `quant_cfg` calibration path and quantized-state save behavior for non-restore runs.
🧹 Nitpick comments (2)
tests/examples/hf_ptq/test_hf_ptq_args.py (1)
73-88: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise CLI validation through
parse_args().This test bypasses the CLI boundary by calling
_restore_quantized_state_if_requesteddirectly. It cannot detect whetherparse_args()rejects both flags, and it does not cover state flags supplied with AutoQuantize. Add parser-level tests for both invalid combinations.As per path instructions, validate CLI/API inputs at boundaries, including incompatible save/restore and checkpoint settings.
🤖 Prompt for 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. In `@tests/examples/hf_ptq/test_hf_ptq_args.py` around lines 73 - 88, Update the tests around _parse_hf_ptq_args to exercise parse_args() validation instead of calling _restore_quantized_state_if_requested directly. Add parser-level coverage for both incompatible save/restore and checkpoint/AutoQuantize flag combinations, asserting each raises ValueError with the appropriate validation message before any restore or processing occurs.Source: Path instructions
modelopt/torch/utils/dataset_utils.py (1)
1188-1189: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winDocument
checkpoint_fnrank semantics.No production caller currently supplies
checkpoint_everyorcheckpoint_fn. The callback runs on every process that invokes the loop. Document that distributed callbacks must be rank-safe or collective before they write shared checkpoints.Source: Path instructions
🤖 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 `@examples/hf_ptq/hf_ptq.py`:
- Around line 1594-1614: Update examples/hf_ptq/hf_ptq.py lines 1594-1614 to
validate --save_quantized_state and --restore_quantized_state during
parse_args(), rejecting their simultaneous use and rejecting either flag with an
AutoQuantize recipe before calibration begins. Update
tests/examples/hf_ptq/test_hf_ptq_args.py lines 73-88 to exercise these invalid
combinations through parse_args() rather than only the private validation
helper.
In `@modelopt/torch/export/plugins/hf_checkpoint_utils.py`:
- Around line 144-156: Update the pattern reconstruction around
_NEMOTRON_H_PATTERN_CHAR_BY_LAYER_TYPE in hf_checkpoint_utils.py to handle
non-string, non-hashable layer_types entries without aborting export: validate
entries before lookup or handle TypeError alongside KeyError, then issue the
existing warning and preserve the native schema. Add a regression test covering
a list or object layers_block_type entry.
In `@modelopt/torch/utils/dataset_utils.py`:
- Around line 1203-1204: Validate checkpoint settings at the public API boundary
before constructing the dataloader or closure: reject checkpoint_every values
below zero, and require checkpoint_fn to be callable whenever checkpoint_every
is positive. Keep zero as the disabled-checkpoint case, let internal checkpoint
logic trust these validated invariants, and add regression tests covering both
invalid combinations.
In `@tests/examples/hf_ptq/test_hf_ptq_args.py`:
- Around line 58-70: Add a focused end-to-end persistence test alongside the
existing mocked tests, using a small local model and pytest tmp_path with the
real mto.save and mto.restore implementations. Verify state is written, restored
into the model, and remains usable by the export retry path; retain the current
delegation and rejection-flow unit tests unchanged.
In `@tests/unit/torch/utils/test_dataset_utils.py`:
- Around line 306-315: Update test_forward_loop_checkpoints_every_n_steps to
record the completed forward count in the checkpoint callback rather than
appending a constant value, then assert calls == [2, 4]. Preserve the existing
five-batch setup and checkpoint_every=2 configuration.
- Around line 293-296: Update the nested _Model class inside _tiny_loader to
define an __init__ method that calls torch.nn.Module’s initializer via
super().__init__(), while preserving its existing forward behavior.
---
Outside diff comments:
In `@examples/hf_ptq/hf_ptq.py`:
- Around line 1352-1365: Move the `_restore_quantized_state_if_requested`
handling before batch-size probing, calibration dataloader creation, and
`pre_quantize`; in restore mode, bypass calibration-only setup and proceed
directly to export. Preserve the existing `quant_cfg` calibration path and
quantized-state save behavior for non-restore runs.
---
Nitpick comments:
In `@tests/examples/hf_ptq/test_hf_ptq_args.py`:
- Around line 73-88: Update the tests around _parse_hf_ptq_args to exercise
parse_args() validation instead of calling _restore_quantized_state_if_requested
directly. Add parser-level coverage for both incompatible save/restore and
checkpoint/AutoQuantize flag combinations, asserting each raises ValueError with
the appropriate validation message before any restore or processing occurs.
🪄 Autofix
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: c1a685ac-2da7-42ca-8b30-083eb4111425
📒 Files selected for processing (8)
examples/hf_ptq/hf_ptq.pymodelopt/torch/export/plugins/hf_checkpoint_utils.pymodelopt/torch/export/quant_aware_conversion.pymodelopt/torch/utils/dataset_utils.pytests/examples/hf_ptq/test_hf_ptq_args.pytests/unit/torch/export/test_hf_checkpoint_utils.pytests/unit/torch/export/test_quant_aware_conversion.pytests/unit/torch/utils/test_dataset_utils.py
Reject --save_quantized_state and --restore_quantized_state together, and reject either with an AutoQuantize recipe (or the deprecated --auto_quantize_bits CLI path), at the CLI boundary instead of only at runtime inside _restore_quantized_state_if_requested. AutoQuantize's search state is not a single quantized model state, so persistence flags for it must fail before calibration work begins, not mid-run. Also exercises the save/restore rejection through parse_args() instead of only the private helper, per review feedback on PR NVIDIA#2129. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JjBHecV84Y6tcMJWKiBRzn Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
…store A malformed config.json (e.g. layers_block_type containing a nested list or dict) raised TypeError from the dict lookup, which the existing KeyError handler did not catch. That aborted export instead of warning and leaving the schema untouched, same as the already-handled unrecognized-string case. Catch TypeError alongside KeyError and add a regression test. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
…ard_loop boundary A negative checkpoint_every stays truthy under the modulo check inside _forward_loop, and a positive checkpoint_every with checkpoint_fn=None would call None once the first interval completes -- both failures only surface deep inside the loop, mid-run. Validate both at the public create_forward_loop boundary instead, before the dataloader/closure are built, and document that checkpoint_fn runs on every process and must be rank-safe or collective if it persists shared state. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
Matches the other nn.Module test double in this file, which does call super().__init__() explicitly rather than relying on the implicit default. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
…ount len(calls) == 2 also passes if checkpoint_fn fires at the wrong batches (e.g. after 1 and 3 instead of 2 and 4). Record the completed forward count in the callback and assert calls == [2, 4] so the test actually pins down where the callback fires. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
--restore_quantized_state retries a failed or interrupted export from a previously calibrated state, so it should not depend on calibration inputs at all. Previously the restore check ran deep inside the mono-quantization branch, after batch-size probing, calibration dataloader construction, and the pre-quantize generation preview had already run -- defeating the export-retry path when the original calibration dataset is unavailable. Check for a requested restore at the top of quantize_main instead, and skip straight to post_quantize/export when it happens. The generation-preview args passed to post_quantize are None in this path, which disables the before/after generation comparison (matching --skip_generate's existing None-preview behavior) without needing calib_dataloader, which post_quantize never actually used. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
The existing save/restore tests mock both mto.save and mto.restore, so they verify delegation and rejection flow but never that ModelOpt state is actually written, restored, or usable afterward. Add a focused test that quantizes a small local Linear layer, saves its state through hf_ptq's --save_quantized_state helper, restores it into a fresh model via --restore_quantized_state, and confirms the restored quantizer state matches and the model runs a forward pass. Also switches the existing mocked save test off a hardcoded /tmp path onto tmp_path. Signed-off-by: Wyatt Neal <wyatt.neal+git@gmail.com>
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
1 participant
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.
Summary
Three fixes found while quantizing
nvidia/Nemotron-Cascade-2-30B-A3B(NemotronH hybrid Mamba-Attention MoE) to NVFP4 viahf_ptq.py.examples/hf_ptq/hf_ptq.py: adds--save_quantized_state/--restore_quantized_state. Lets a failed export step be retried without redoing calibration (calibration on a 30B MoE is the expensive part; export/serialization bugs shouldn't force re-running it).modelopt/torch/utils/dataset_utils.py: adds an optionalcheckpoint_every/checkpoint_fnhook to_forward_loop/create_forward_loop, so long calibration runs can checkpoint periodically. (Not yet wired intohf_ptq.py's calibration call — that's a natural follow-up, intentionally left out of this PR to keep it scoped.)modelopt/torch/export/quant_aware_conversion.py:build_reverse_name_mapper's sentinel-strip logic (_map) silently accepted a mangled sentinel instead of raising. In our export, a reverse-rename rule matched past its intended boundary and corrupted the appended sentinel before it could be stripped cleanly (observed output was an unrelated, unsubstituted placeholder token, not the real sentinel — direct evidence of a rule mismatch, not a calibration gap). Effect: 100% ofexclude_modulesentries (77/77) were corrupted in our exportedconfig.json. vLLM then couldn't match any excluded layer against its weight-loading logic, defaulted every quant-aware fused linear (e.g. Mamba'sin_proj, aMergedColumnParallelLinear) to quantized-width allocation, and crashed loading the correctly-unquantized-but-mislabeled weight — a bareAssertionErrordeep in vLLM with no useful diagnostic pointing back at the real cause.Fix extracts the strip logic into
_strip_sentinel_or_raise(), which now raisesQuantConversionUnsupportedErroron a failed strip — exactly the exceptionbuild_reverse_name_mapper's own docstring already documents for this case, just not previously implemented. No caller changes needed: both call sites inunified_export_hf.pyalready wrap this in a broadtry/exceptthat falls back to safe in-memory names with a warning — this fix just makes that existing safety net actually trigger, turning a silent, deployment-time failure into a loud, immediate one at export time.All three built test-first (Red → Green). 26 tests total across the two touched test files (2 new for the checkpoint flags, 3 new for the calibration checkpoint hook, 3 new for the sentinel fix), no regressions on pre-existing tests.
Test plan
pytest tests/examples/hf_ptq/test_hf_ptq_args.pypytest tests/unit/torch/utils/test_dataset_utils.pypytest tests/unit/torch/export/test_quant_aware_conversion.pynvidia/Nemotron-Cascade-2-30B-A3BNVFP4 export without the fix, confirmed it's caught at export time with the fix applied.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests