Minitron pruning fixes for Nemotron-3.5-Lightning-30B-A3B and Deepseek - #2159
Conversation
Enables pruning and HF-checkpoint export of native NemotronHForCausalLM models that ship without remote code and carry MTP heads (e.g. NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16) on the modern Megatron-Bridge / transformers stack: - Clear mtp_hybrid_override_pattern in the calibration provider overrides. On newer Megatron-LM, mtp_process is derived from the hybrid pattern, so mtp_num_layers=0 alone left mtp_process=True and calibration tripped `assert self.config.mtp_num_layers > 0`. - Stream pruned weights through a config-only bridge (AutoBridge.from_hf_config(...).save_hf_pretrained(...)) when available. Its hf_keys=None, so the embedding conversion task is not dropped when transformers' saved key differs from the bridge mapping (backbone.embedding vs backbone.embeddings), which previously left an owner-less PP placeholder and crashed save_hf_weights. VLMs and older Megatron-Bridge builds keep the dummy-model path. - Only write hybrid_override_pattern for older remote-code configs that lack layer_types; on native configs it is a read-only property. - Preserve n_shared_experts (a fixed count) instead of re-deriving it from moe_shared_expert_intermediate_size // moe_ffn_hidden_size, which is DeepSeek-style logic that corrupts NemotronH's count. Extends the nemotron_h prune test to export to HF and reload, with an n_shared_experts assertion, and drops the now-dead megatron_format branch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe pruning flow adds MoE candidate filtering, validates DeepSeek-style shared-expert configurations, and uses configuration-only HF export when supported. Distributed scripts now abort peers before cleanup on uncaught failures. ChangesPruning and HF export
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant prune_minitron
participant AutoBridge
participant HF_output
prune_minitron->>AutoBridge: from_hf_config(hf_config)
AutoBridge-->>prune_minitron: config-only bridge
prune_minitron->>AutoBridge: save_hf_pretrained(...)
AutoBridge->>HF_output: stream HF weights
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — findings
Scope: full review (the trigger comment carried no scoping instructions). All 4 changed files reviewed (+70/-42): examples/megatron_bridge/prune_minitron.py, examples/megatron_bridge/README.md, tests/_test_utils/examples/megatron_bridge.py, tests/examples/megatron_bridge/test_prune_minitron.py. No modelopt/ source changed, so no mode-registration or modelopt_state schema surface is affected.
Findings: CRITICAL: 0 · IMPORTANT: 3 · SUGGESTION: 2
What holds up
All three fixes are well-targeted and I could not fault the logic:
- Fix 1 (MTP calibration) — clearing
mtp_hybrid_override_patternalongsidemtp_num_layers=0addresses the actual root cause (the providerfinalize()re-appending the MTP suffix somtp_processstays true), not the assertion symptom. Consistent with the coding standard on fixing causes rather than side effects. - Fix 3a (
n_shared_experts) — thehas_explicit_shared_sizenarrowing is right in both directions: NemotronH carries both the count and an explicitmoe_shared_expert_intermediate_size, so the count must be preserved; DeepSeek-V3 has no explicit shared-size field, so it still needs the derivation. Neither Qwen3.5-MoE (shared_expert_intermediate_size, no count) nor dense configs are affected. - Fix 3b (
hybrid_override_pattern) — skipping the write whenlayer_typesis present is consistent with the pre-existinglayer_typesslicing just above it, which already carries the pruned cadence; writing both would be two sources of truth for the same thing.
Most impactful findings
-
[IMPORTANT] The
nemotron_htest case is now skipped in CI entirely, not just its export assertions.config_only_hf_export_supported()requires nemo:26.08+, but both workflows pinnvcr.io/nvidia/nemo:26.06(example_tests.yml:156,gpu_tests.yml:53) — as the adjacent "until nemo:26.08 container" TODO confirms. Before this PR the case ran the full prune and asserted the Megatron checkpoint, which covered fix 1 (the crash this PR exists to fix) and fix 3. After it, nothing runs until the container bumps. Only the export leg needs 26.08+, so gating the whole parametrization gives up coverage that does not depend on it. Suggested restructuring is in the inline comment. -
[IMPORTANT] The config-only path is gated more broadly than the bug it fixes, and drops
trust_remote_code. The condition is "API present AND not VLM", so on 26.08+ every non-VLM export (remote-code Nemotron, DeepSeek, Qwen3.5-MoE, dense Qwen3) switches mechanism — which qualifies the PR description's "existing … exports are unaffected"; that holds on 26.06 only. Separately,args.trust_remote_codeis threaded through every other HF call in this block (line 596 and all three fallback calls) but not intofrom_hf_config/save_hf_pretrained, whosesource_pathpoints back at the original repo. -
[IMPORTANT] The fallback is a silent path to a known-broken outcome, and its README workaround was deleted. Your new helper documents the dummy-model path as not round-tripping the pruned config, yet a 26.06 + transformers>=5 user now takes it with no log line and no README note — the removed
transformers<5hint was exactly the guidance for that combination. Awarn_rank_0on fallback plus a note scoped to older containers restores discoverability.
The two SUGGESTIONs (unused from_auto_config in the gate; silent floor-division truncation in the DeepSeek n_shared_experts derivation) are inline and non-blocking.
Risk assessment
Low-to-moderate, and confined to the example script — no library code, no checkpoint/state schema, no public API. The end-to-end validation on the 30B target model is convincing for the path it exercised. The residual risk is entirely about unexercised paths: the primary MTP fix has no CI coverage on the container CI actually uses (finding 1), and the export-mechanism switch reaches every non-VLM model on 26.08+ with no test behind it (finding 2). Neither is a correctness defect in what you ran; both are gaps that would let a future regression through silently.
Findings 1 and 3 are cheap to address. Finding 2 may be resolvable with a sentence confirming the config-only bridge never loads remote code.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2159 +/- ##
==========================================
- Coverage 78.74% 78.16% -0.59%
==========================================
Files 522 522
Lines 60368 60392 +24
==========================================
- Hits 47538 47205 -333
- Misses 12830 13187 +357
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:
|
Follow-up to the NemotronH Lightning prune + HF export fix. mcore_minitron: add an optional candidate_filter search config option (default None). Candidates come from a Cartesian product of independent per-hparam choice lists, so a per-hparam restriction cannot express a constraint *between* two hparams. The filter runs before the metric computation, and receives every supported hparam with non-searched ones filled in from the model config. It is also re-applied to candidates restored from a search checkpoint, since all_candidates_per_constraint is cached by constraints alone and would otherwise silently reuse candidates computed under a different filter. prune_minitron.py wires this up for DeepSeek-style MoE configs, which size the shared expert as n_shared_experts * moe_intermediate_size and so can only represent a shared size that is a multiple of the routed one. A --prune_export_config never reaches the filter, so the export path raises instead of writing a config that disagrees with the saved weights. Review feedback on the export path: - Set trust_remote_code on the config-only bridge. save_hf_pretrained takes no such argument and reads it off the bridge to fetch source artifacts, and from_hf_config cannot infer it because AutoConfig consumes the kwarg. - Warn when falling back to the dummy-model path, and keep the README transformers<5 note scoped to containers without config-only export. - Drop the unused from_auto_config from the capability check. Tests: nas_memory_mb now exercises candidate_filter (exact rejected count, surviving candidate validity, regenerated top-k goldens) and covers the model-config fallback for hparams outside the search space. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
changed the title
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: 2
🤖 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/prune/plugins/mcore_minitron.py`:
- Around line 575-586: Update the candidate-cache flow around
_compute_candidate_metrics and the cache-removal logic at lines 595-602 to
retain every metric-qualified CandidateSubnet regardless of candidate_filter.
Build a temporary filtered view using the current filter only when performing
top-k selection, so resumed searches with no or less restrictive filters can
recover previously rejected candidates; add a checkpoint-resume test covering
that filter change.
In `@tests/_test_utils/examples/megatron_bridge.py`:
- Around line 34-39: Add a brief comment immediately before the local AutoBridge
import in the helper, documenting that Megatron Bridge is optional and
unavailable installations cause the helper to return False.
🪄 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: 90368826-3115-47f6-935b-2c30740677cb
📒 Files selected for processing (6)
examples/megatron_bridge/README.mdexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/_test_utils/examples/megatron_bridge.pytests/examples/megatron_bridge/test_prune_minitron.pytests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/megatron_bridge/README.md
- tests/examples/megatron_bridge/test_prune_minitron.py
|
Summary of changes since reviewThree groups of work landed in 1. Review feedback
2.
|
|
/claude review |
672e982 to
2ae34c4
Compare
There was a problem hiding this comment.
Claude review — 6 findings (CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 4)
Scope: full review (trigger comment carried no scoping instructions). All 11 changed files opened — modelopt/ (2) first, then examples/ (6), then tests/ (3). Nothing deprioritized.
Most impactful
1. mtp_hybrid_override_pattern is an unconditional provider override with a hard hasattr assert (examples/megatron_bridge/prune_minitron.py:424) — IMPORTANT Compatibility
load_mbridge_model_from_hf asserts hasattr(provider, key) for every override key (modelopt/torch/utils/plugins/mbridge.py:95). mtp_hybrid_override_pattern belongs to the modern Megatron-LM / HybridModelProvider era that mbridge.py already guards behind HAS_HYBRID. If nemo:26.06 providers lack it, this assert fires before any model is built, breaking prune_minitron.py for every architecture — including the qwen3 CI case, which (unlike the new nemotron_h case) is not skipped on 26.06. All the end-to-end validation in the PR body was on nemo:26.08.rc6, so the container CI actually runs is untested for this line. Please either gate the key on hasattr or confirm it exists on 26.06.
2. The config-only export path has no weights-vs-config verification (prune_minitron.py:718-731) — IMPORTANT Export
The ~60 lines of field-by-field hf_cfg write-back above are handed to AutoBridge.from_hf_config(hf_cfg), and whatever save_hf_pretrained emits becomes config.json — unchecked. The PR description documents this exact failure already occurring here (pruned Qwen3: hidden_size=24 weights, config.json saying 32, checkpoint unloadable). The chosen mitigation is an architecture allowlist, but _HYBRID_PROVIDER_TYPES admits every hybrid provider while only native NemotronH was validated, so an unvalidated hybrid gets the same silent mismatch — surfacing only when a user later loads the checkpoint. An assert on the written config after save turns a class of silent corruption into a loud export-time error, and would also make the isinstance(...) restriction safe to relax later.
Suggestions (non-blocking)
prune_minitron.py:711— comment says "Preferred path for all non-VLMs" but the code restricts to hybrids, contradicting its own next paragraph.prune_minitron.py:652-663— the newValueErrorfires aftersave_artifactsalready wrote aconfig.json, which the early-exit guard inmain()then reads as "Pruned model already exists… Exiting…" on the next run — unpruned checkpoint, exit code 0. The constraint is checkable fromargsup front.modelopt/torch/utils/distributed.py:240-241—abort()'s docstring says a re-raisedSystemExitshuts the PG down "normally", but the subsequentfinally: dist.cleanup()sees non-Nonesys.exc_info()and skips the barrier. Low practical risk (all ranks reachsys.exittogether), but the doc and code disagree. Also suggests an explicitcleanup(barrier=...)escape hatch, since theexc_infoheuristic makes behavior depend on ambient interpreter state.tests/_test_utils/examples/megatron_bridge.py:17-18— module-levelfrom megatron.bridge import AutoBridgemakes this shared helper import-fatal without Megatron-Bridge (breakingtest_qad.pycollection, not skipping) and renders the new function'stry/exceptdead. The siblingqwen35_moe_bridge_supported()deliberately imports inside thetry.modelopt/torch/prune/plugins/mcore_minitron.py:589— "No subnets found fitting the constraints!" is misleading whencandidate_filterrejected the whole grid;num_filteredis already in scope.
Verified as correct
candidate_filterplumbing — defaultNonekeeps existing searches byte-identical;SUPPORTED_HPARAMSfallback viaself.model.configcorrectly covers skipped/single-choice hparams; the filter runs before_compute_candidate_metrics, so it makes the search cheaper; re-application to checkpoint-restored candidates closes the real hole created by caching on constraints alone, and the docstring correctly documents the narrowing-only resume semantics. No lambda reachesstate_dict()(BaseSearcher.state_dictonly serializesdefault_state_dictkeys), so no pickling hazard from theprune_minitron.pylambda.dist.abort()design — not tearing down the process group at all (rather than just skipping the barrier) is the right call given peers parked inrecv_from_prev_pipeline_rank_; flushing both streams beforeos._exitis necessary sinceos._exitskips buffer flushing; re-raisingSystemExitcorrectly preserves the--score_lower_boundexit code. Keeping it out ofcleanup()so no library caller gets a surprise process exit is the right boundary._is_deepseek_style_moenarrowing — correctly restricts then_shared_expertsre-derivation to configs with no explicit shared-size field; NemotronH keeps its fixed count (matching the new{"n_shared_experts": 1}test assertion) while the_SHARED_EXPERT_SIZE_FIELDSloop handles its pruned size. Source-config detection correctly unwrapstext_configfor VLMs.hybrid_override_patternwrite-back gate — skipping it whenlayer_typesis present correctly avoids the read-only property on native configs.- Mid-
main()dist.cleanup()indistill.py/export_distilled_megatron_to_hf.py— composes safely with the new outerexcept BaseException: dist.abort(); the PG is already gone, so a rank-0 export failure cannot stall. - No
modelopt_stateschema change, no mode registration change, no public API signature change — restore/backward compatibility is unaffected.
Risk assessment
Moderate, concentrated in one line. The library-side changes (candidate_filter, dist.abort/cleanup) are additive, default-off, and sound. Finding 1 is the only thing that could break currently-green CI, and it is a one-line guard away from resolved. Finding 2 is about hardening a path whose failure mode this PR already experienced once — worth adding before the allowlist inevitably widens.
Nothing here duplicates CodeRabbit's style/security surface.
2ae34c4 to
30d78d2
Compare
A rank raising inside a distributed entrypoint left the job stalled until the process group timed out, with no diagnostic output: the failing rank blocked in cleanup's barrier while its peers blocked in a pipeline recv, and Python only prints a traceback once the enclosing `finally` returns. - dist.cleanup(): skip the barrier when unwinding from an exception, since peers may be blocked in a collective this rank will never reach. - dist.abort(): new helper that prints the traceback, flushes and exits immediately. Skipping the barrier alone is not enough -- destroy_process_group blocks for the same reason (confirmed by stack dump) -- so the error path must not tear the process group down at all. SystemExit is re-raised rather than aborted: an intentional exit (e.g. the --score_lower_bound gate) is reached by every rank, so it needs no traceback and keeps its exit code. Kept out of cleanup() so no library caller gets a surprise process exit. Measured on a 2-GPU PP run whose rank 0 raises during calibration: 10 min timeout kill with no visible error -> 31s, exit 1, real traceback. Also narrow the config-only HF export to hybrid models. Routing every non-VLM export through it regressed dense models: a pruned Qwen3 exported with hidden_size 24 weights but a config.json still saying 32, so the checkpoint would not reload. The dummy-model path round-trips the pruned config for non-hybrids and is what CI has always exercised. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
30d78d2 to
91d6361
Compare
Update since last review (
|
| Test | Result |
|---|---|
test_prune_minitron[qwen3] |
PASSED |
test_prune_minitron[nemotron_h] |
PASSED (58s), via the config-only path, no fallback warning |
Re-confirmed after deleting the skip mark: 2 passed, exit 0. Also still passing on nemo:26.08.rc6, along with nas_memory_mb and nas_params.
Other changes from this review round
cleanup()now treatsSystemExitas the clean path, so the--score_lower_boundgate keeps barrier-then-destroy (the docstring had promised this; the code did not).candidate_filteris no longer re-applied to checkpoint-restored candidates — it matchesscore_func, whose cached scores are reused without re-validation.mtp_hybrid_override_patternverified present on nemo:26.06.00/26.06.01 for bothMambaModelProviderandGPTModelProvider, so the unconditional provider override cannot trip thehasattrassert.- Declined the post-save config assertion:
test_prune_minitronalready reloads the exported checkpoint, which is how the Qwen3 mismatch surfaced, so it is covered by a test rather than by runtime checks over a field list that will drift.
|
/claude review |
There was a problem hiding this comment.
Claude review — 4 findings (CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 4)
Claude review passed — no blocking issues found. LGTM
Scope: full review (trigger comment carried no scoping instructions). All 10 changed files reviewed — modelopt/ (2) first, then examples/ (6), then tests/ (2). Nothing deprioritized.
Prior-round findings: both resolved
mtp_hybrid_override_patternas an unconditional provider override (previously IMPORTANT) —load_mbridge_model_from_hfhard-assertshasattr(provider, key)(modelopt/torch/utils/plugins/mbridge.py:95), so a provider missing this key would breakprune_minitron.pyfor every architecture on the container CI actually runs. The test table now reportstest_prune_minitron[qwen3]and[nemotron_h]passing on nemo:26.06.01 as well as 26.08.rc6; both exercise this override on the dense and the hybrid provider respectively, so the attribute is present on both. Considered settled empirically.- No weights-vs-config verification on the config-only export path (previously IMPORTANT) — the
nemotron_hcase now ends inAutoModelForCausalLM.from_pretrained(pruned_path), a default-strict load that fails on exactly the mismatch class the PR description documents (pruned weights vs. staleconfig.json). That puts the shipped path under CI. Residual, non-blocking:_HYBRID_PROVIDER_TYPESadmits every hybrid provider on 26.08+, so a hybrid other than NemotronH (e.g. a non-VLM Qwen3.5-MoE — only the VL variant is in CI, and VLMs are excluded from this branch) would still surface a mismatch at user load time rather than at export. Worth an export-time config assert whenever the allowlist widens, but the allowlist plus the new load test is a legitimate mitigation for what ships here. tests/_test_utils/examples/megatron_bridge.pyno longer appears in the diff — the import-fatal shared-helper concern is gone with it.
This round (all non-blocking)
_is_deepseek_style_moekeys on attribute presence, not value (prune_minitron.py:121) —n_shared_experts=Noneis a valid no-shared-expert setting in DeepSeek-V2/V3-style configs, and it currently routes into both the divisibility filter (line 573) and the new export guard (line 653), each of which then raisesTypeError: ... %: 'NoneType' and 'int'instead of taking the no-shared-expert path. One-line fix suggested inline.base_configcan handcandidate_filteraNone(mcore_minitron.py:560-564) — this searcher's own export path documents thatmoe_ffn_hidden_size,kv_channelsandnum_query_groupsmay be unset, but it normalizes them only after the search. The new docstring promises non-searched hparams are 'filled in from the model config' without noting the fill-in may beNone, so an arithmetic filter dies inside thetqdmloop with no hint the value was a fallback. Suggested either reusing the existing normalization or droppingNones so the failure names the field.cleanup()'s docstring overstates what skipping the barrier buys (distributed.py:218-221) —abort()'s own docstring correctly notes thatdestroy_process_groupstalls for the same reason as the barrier, andcleanup()still calls it unconditionally on the error path. Two entrypoints have the sametry/finally: dist.cleanup()shape as the five that were fixed but did not get theabort()line:modelopt/torch/puzzletron/scoring.py:85-88andexamples/hf_ptq/example_utils.py:96-99. Either extendabort()to those two or tighten the docstring.pruned_bridge.trust_remote_code = ...is a silent no-op if the field moves (prune_minitron.py:723) — the only affected models are remote-code hybrids, which is exactly the set neither CI (locally-built tiny checkpoint) nor the manual validation (native NemotronH, no remote code) covers.mbridge.py:95already usesassert hasattr(...)for this same hazard; mirroring it makes a future rename loud.
Verified as correct
candidate_filterplumbing — defaultNoneleaves existing searches byte-identical; the filter runs before_compute_candidate_metrics, which is analytically pure, so short-circuiting has no side effects and makes the search strictly cheaper;base_configis computed once outside the loop from global (PP-invariant) config fields, so all ranks filter identically andselectedcannot diverge. TheSUPPORTED_HPARAMSfallback is confirmed by the new GPU test, wheremoe_shared_expert_intermediate_sizeis in the skip list and is supplied from the model config. No lambda reachesstate_dict()— confirmed empirically by the new test passing while writing a search checkpoint with a nested-function filter. Narrowing-only resume semantics are documented on the config key.dist.abort()design — not tearing the PG down at all is the right call given peers parked inrecv_from_prev_pipeline_rank_; flushing both streams beforeos._exitis necessary sinceos._exitskips buffer flushing; re-raisingSystemExitpreserves the--score_lower_boundexit code, and since thatsys.exit(1)(line 756) is not rank-guarded,cleanup()'s 'every rank reaches it' premise holds and the barrier is correctly still taken. Keeping the process exit out ofcleanup()is the right library boundary.n_shared_experts/_SHARED_EXPERT_SIZE_FIELDSordering — the_SHARED_EXPERT_SIZE_FIELDSwrite-back loop runs before_is_deepseek_style_moe(text_cfg), but by construction the predicate is only true when neither field exists, so the earliersetattrcannot flip it. DeepSeek semantics are right:moe_intermediate_sizetakes the pruned routed size andn_shared_experts = shared // routedagainst that same pruned size.hybrid_override_patternwrite-back gate — skipping it whenlayer_typesis present avoids the read-only property on native configs and matches the pre-existinglayer_typesslicing that already carries the pruned cadence.- Export-branch CI coverage —
from_hf_configexists on 26.06+, soqwen3covers the dummy-model branch andnemotron_hthe config-only branch on both containers; the fallbackwarn_rank_0is correctly unreachable for VLMs. - Test changes —
contextlib/io/reare all already imported intest_mcore_mamba_minitron_pruning.py, so the new capture block is sound; the 256-of-512 assertion is exact and consistent withmoe_ffn_hidden_sizein [12, 16] against an unpruned shared size of 16. - No
modelopt_stateschema change, no mode registration change, no public API signature change (candidate_filteris additive with aNonedefault) — restore and backward compatibility are unaffected.
Risk assessment
Low. The two library-side changes are additive and default-off (candidate_filter) or strictly reduce a stall (cleanup/abort). Both of the previous round's IMPORTANT findings are addressed — one by cross-container verification, one by a strict-reload assertion in CI. Everything remaining is defensive hardening against configs and future dependency versions this PR does not target: None-valued MoE fields (findings 1-2), an incomplete rollout of the fail-fast fix to two non-Megatron-Bridge entrypoints (finding 3), and a silent-drop risk on a duck-typed bridge attribute (finding 4).
Nothing here overlaps CodeRabbit's style/security surface.
Pruning any MLA model with q_lora_rank set died during calibration with `AttributeError: 'tuple' object has no attribute 'view'`. Root cause: hidden_size importance estimation blanket-patches every TELayerNormColumnParallelLinear with return_layernorm_output=True to capture post-layernorm activations. When q_lora_rank is set, MCore builds linear_q_up_proj as a TELayerNormColumnParallelLinear (the Q-LoRA layernorm is fused into it, which is why q_layernorm is IdentityOp), so it was patched too -- even though its layernorm is over the latent rank, not hidden_size. TE then returns ((out, ln_out), bias) and MCore's `q, _ = self.linear_q_up_proj(...)` leaves q a tuple. Isolated by probing the module before/after conversion: conversion alone keeps the (Tensor, None) return and the forward passes; registering the importance hooks is what nests it. Fix: exclude MLA's Q/KV up-projections from both the patch and unpatch loops. test_mcore_mla_pruning did not catch this because it builds MLA without q_lora_rank, where MCore uses a plain linear_q_proj and nothing is patched. Adds the deepseek_v3 case to test_prune_minitron, which covers this path plus candidate_filter end-to-end (n_group=1 keeps num_moe_experts divisible after expert pruning). Verified on 2 GPUs: prune + HF export + reload, 52s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
changed the title
Folded in: MLA Q-LoRA fix — DeepSeek-V3 now prunes end-to-end (
|
| Setup | linear_q_up_proj returns |
Forward |
|---|---|---|
| Before conversion | tuple(Tensor, NoneType) |
— |
| After conversion, no hooks | tuple(Tensor, NoneType) |
OK |
| After conversion + importance hooks | tuple(tuple(Tensor, Tensor), NoneType) |
AttributeError |
Fix: exclude MLA's Q/KV up-projections from the patch and unpatch loops (~20 lines). test_mcore_mla_pruning never caught this because it builds MLA without q_lora_rank, where MCore uses a plain linear_q_proj and nothing is patched.
Verified on 2 GPUs (nemo:26.08.rc6):
| Check | Result |
|---|---|
test_prune_minitron[deepseek_v3] (restored) |
PASSED, 52s |
test_prune_minitron[qwen3] / [nemotron_h] |
PASSED |
tests/gpu_megatron/torch/prune/plugins/ |
18 passed, incl. test_mcore_mla_pruning |
| Pruned DeepSeek checkpoint reload | 246,032 params; hidden_size=112, moe_intermediate_size=56, n_shared_experts=1, q_lora_rank=32 preserved |
The restored deepseek_v3 case now covers both the MLA Q-LoRA path and candidate_filter end-to-end: without the filter the search picks a shared size that is not a multiple of the routed one, and the reload fails on the shape mismatch.
One aside worth recording: an early DeepSeek run printed Rejected 546 candidates via candidate_filter immediately followed by No subnets found fitting the constraints!, and I spent time chasing the params target before realising the filter had emptied the grid — concrete support for the suggestion to name the rejected count in that assertion.
The pruned checkpoint intermittently reloaded with the *source* config against pruned weights. Every rank ran bridge.hf_pretrained.save_artifacts() (writing the source config.json), while save_hf_pretrained writes the pruned config from rank 0 alone -- so a late write from another rank could land last. This only bites the config-only path; the dummy-model path writes the pruned config from every rank, so per-rank ordering makes it self-correcting. That also explains the earlier Qwen3 failure, which had the same symptom and was on the config-only path at the time -- so narrowing the export gate moved Qwen3 off the racy path rather than fixing a config-only export defect. The gate stays narrow (hybrids are still the only models that need it), but the comment no longer claims config-only export drops pruned fields for dense models. Guard the artifact write to rank 0 and barrier before reading it back. Verified: nemotron_h failed intermittently in a 3-test session before, 3/3 clean sessions after. Review follow-ups: - _is_deepseek_style_moe keyed on the value, not hasattr: n_shared_experts=None is a valid "no shared expert" setting, and it made both the filter lambda and the export guard raise TypeError on `%` instead of taking the no-shared path. - base_config omits hparams MCore leaves unset, so a filter touching one raises KeyError naming the field rather than TypeError from inside the search loop. - The "no subnets found" assert names the rejected count, so an emptied grid is not mistaken for too-tight constraints (hit while debugging DeepSeek). - cleanup() docstring: skipping the barrier is not sufficient on its own, since destroy_process_group blocks for the same reason. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Race fix on HF export (
|
kevalmorabia97
changed the title
kevalmorabia97
added
the
cherry-pick-done
#2159 #2112 (#2179) ## Cherry-picked PRs - #1975 - #2076 - #2071 - #2093 - #2084 - #2115 - #2133 - #2146 - #2064 - #2159 - #2112 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added PTQ recipes for Nemotron model families and NVFP4 mixer-MLP quantization. * Added Qwen3-VL multimodal speculative-decoding support, including video inputs. * Added compatibility with multiple vLLM KV-cache layouts and newer Transformers versions. * Added a Nemotron 3.5 Lightning quantization-aware distillation workflow. * **Bug Fixes** * Improved Hugging Face, QLoRA, and PEFT checkpoint exports. * Improved distributed job shutdown when a process encounters an error. * Improved pruning validation and candidate selection for MoE models. * **Documentation** * Updated supported-model lists and recipe paths. * Removed Phi-3 Vision and Phi-4 Multimodal quantization support. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Slawomir Kierat <skierat@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: James Shen <yueshen@nvidia.com> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com> Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Co-authored-by: skierat <skierat@nvidia.com> Co-authored-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-authored-by: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Co-authored-by: Zhiyu <zhiyuc@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: kinjalpatel27 <31936134+kinjalpatel27@users.noreply.github.com> Co-authored-by: sychen52 <41452870+sychen52@users.noreply.github.com> Co-authored-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Co-authored-by: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Co-authored-by: Jenny Chen <jennifchen@nvidia.com> Co-authored-by: sugunav14 <178320438+sugunav14@users.noreply.github.com>
2 participants
What does this PR do?
Type of change: Bug fix + new feature
Two model families that could not be pruned end-to-end now can:
nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16) — a nativeNemotronHForCausalLMthat ships without remote code and carries MTP heads. Fixes a calibration crash and HF-export failures on the modern Megatron-Bridge / transformers stack.candidate_filtersearch option tomcore_minitronso its MoE-FFN dimensions stay prunable while remaining representable in HF.Also makes a rank-local failure under pipeline parallelism fail fast instead of stalling.
1. Nemotron Lightning: prune + HF export (
examples/megatron_bridge/prune_minitron.py)MTP calibration crash. On newer Megatron-LM,
mtp_processis derived from the hybrid pattern, not frommtp_num_layers. Setting onlymtp_num_layers=0in the calibration provider overrides was insufficient: the provider'sfinalize()re-appended the MTP suffix tohybrid_layer_pattern(becausemtp_hybrid_override_patternwas still set andmtp_use_repeated_layer=True), somtp_process=Truewhilemtp_num_layers=0and the calibration forward hitassert self.config.mtp_num_layers > 0. Fix: also clearmtp_hybrid_override_patternin the calibration overrides so MTP is fully disabled (MTP heads are dropped from the pruned model, as before).HF export via a config-only bridge (hybrid models only). The old export built a dummy HF model to obtain the bridge, then streamed weights. This breaks on native NemotronH because (a) native
NemotronHConfigmakeshybrid_override_patterna read-only property, and (b) transformers 5.12 saves the input embedding under a different key than the bridge mapping expects (backbone.embeddingvsbackbone.embeddings); the mismatch madebuild_conversion_tasksdrop the embedding task on its owning rank, leaving an owner-less PP placeholder that crashedsave_hf_weightswithObject must exist on at least one PP rank. Fix: stream weights through a config-only bridge (AutoBridge.from_hf_config(hf_cfg).save_hf_pretrained(...)), available since Megatron-Bridge 0.5.0 (nemo:26.06). A config-only bridge hashf_keys=None, so the embedding task is never dropped, no dummy model is built, and the output uses the canonical HF key names.This is restricted to hybrid providers, which are the only models that need it; non-hybrids keep the dummy-model path that CI has always exercised.
Writing the source artifacts is now rank-0-only. Every rank used to write the source
config.json, which races with the prunedconfig.jsonthatsave_hf_pretrainedwrites from rank 0 alone: a late write from another rank leaves a checkpoint whose config does not match its weights. This reproduced intermittently on both Qwen3 and NemotronH before the fix, and 3/3 clean runs after.save_hf_pretrainedtakes notrust_remote_codeargument — it reads the flag off the bridge to fetch the source checkpoint's artifacts, andfrom_hf_configcannot infer it becauseAutoConfig.from_pretrainedconsumes the kwarg rather than storing it on the config. So the flag is set explicitly on the bridge instance; otherwise remote-code models would silently lose it.Config write-back correctness:
hybrid_override_patternis only written for older remote-code configs that lacklayer_types; native configs carry the cadence inlayer_types(read-onlyhybrid_override_patternis skipped).n_shared_expertsis preserved (a fixed count) instead of being re-derived bymoe_shared_expert_intermediate_size // moe_ffn_hidden_size, which is DeepSeek-style logic that would corrupt NemotronH's count.Non-hybrids, VLMs, and Megatron-Bridge builds without config-only export keep the dummy-model path, with a
warn_rank_0when a hybrid has to fall back. The README'stransformers<5workaround is removed: it existed because the dummy-model path broke on transformers 5, and the config-only path handles NemotronH on every supported container.2.
candidate_filterformcore_minitron(modelopt/torch/prune/plugins/mcore_minitron.py)DeepSeek-style MoE configs have no explicit shared-expert-size field: they size the shared expert as
n_shared_experts * moe_intermediate_size, wheremoe_intermediate_sizeis the (also prunable) routed expert size. So only candidates withmoe_shared_expert_intermediate_size % moe_ffn_hidden_size == 0can be written back to HF at all.Candidates come from a Cartesian
product()of independent per-hparam choice lists, so no per-hparam restriction can express a constraint between two hparams. New optionalcandidate_filtersearch-config key (defaultNone, so existing behaviour is unchanged): a callable that rejects candidate configs before the metric computation, making the search cheaper rather than more expensive. It receives every supported hparam, with non-searched ones filled in from the model config, so a filter still works when one of its hparams was skipped or had a single choice.Rejected candidates are not cached, so — like
score_func, whose cached scores are reused without re-validation — the filter is assumed unchanged when resuming from acheckpoint.prune_minitron.pywires this up for DeepSeek-style configs, so bothmoe_ffn_hidden_sizeandmoe_shared_expert_intermediate_sizestay prunable (the search then only picks shared sizes that are a multiple of the routed one). A--prune_export_configthat violates the constraint never reaches the filter, so the export path now raisesValueErrorinstead of writing a checkpoint whose config disagrees with its weights.3. MLA Q-LoRA pruning (
modelopt/torch/prune/plugins/mcore_minitron.py)Pruning any MLA model with
q_lora_rankset died during calibration withAttributeError: 'tuple' object has no attribute 'view'.hidden_sizeimportance estimation blanket-patches everyTELayerNormColumnParallelLinearwithreturn_layernorm_output=Trueto capture post-layernorm activations. Whenq_lora_rankis set, MCore buildslinear_q_up_projas aTELayerNormColumnParallelLinear— the Q-LoRA layernorm is fused into it, which is whyq_layernormisIdentityOp— so it was patched too, even though its layernorm is over the latent rank, nothidden_size. TE then returns((out, ln_out), bias)and MCore'sq, _ = self.linear_q_up_proj(...)leavesqa tuple.Isolated by probing the module before and after dynamic conversion:
linear_q_up_projreturnstuple(Tensor, NoneType)tuple(Tensor, NoneType)tuple(tuple(Tensor, Tensor), NoneType)So conversion is innocent; registering the importance hooks is the trigger. Fix: exclude MLA's Q/KV up-projections from both the patch and unpatch loops.
test_mcore_mla_pruningdid not catch this because it builds MLA withoutq_lora_rank, where MCore uses a plainlinear_q_projand nothing is patched.4. Fail fast instead of stalling on a rank-local error under PP (
modelopt/torch/utils/distributed.py)A rank raising inside a distributed entrypoint left the whole job stalled until the process group timed out, with no diagnostic output at all: the failing rank blocked in
cleanup()'s barrier while its peers blocked inrecv_from_prev_pipeline_rank_, and Python only prints a traceback once the enclosingfinallyreturns. A crash on one rank was indistinguishable from a slow job.dist.cleanup()skips the barrier when unwinding from an exception.dist.abort()prints the traceback, flushes and exits immediately. Skipping the barrier alone is not enough — a stack dump showed the failing rank then blocking indestroy_process_groupfor the same reason — so the error path must not tear the process group down at all.SystemExitis re-raised rather than aborted, so an intentional exit (e.g. the--score_lower_boundgate) keeps its exit code and prints no traceback. Kept out ofcleanup()so no library caller gets a surprise process exit.main()intry/finally: the fiveexamples/megatron_bridgescripts.Measured on a 2-GPU PP run whose rank 0 raises during calibration: 10 min timeout kill with no visible error → 31s, exit 1, real traceback. This is a latent, pre-existing issue (the
try/finallypredates this PR); it only surfaces on a failing PP run, which is why CI never hit it.Usage
torchrun --nproc_per_node 4 examples/megatron_bridge/prune_minitron.py \ --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 \ --pp_size 4 \ --prune_target_active_params 3e9 \ --output_hf_path /path/to/Nemotron-3.5-Lightning-30B-A3B-Pruned-A3.0BTesting
End-to-end on nemo:26.08.rc6 (4× GB300, transformers 5.12.1, Megatron-Bridge with config-only export): pruning + export complete (
EXIT=0, "Saved pruned model … Done!"). The exported checkpoint has canonical pluralbackbone.embeddings.weightkeys, 0 MTP tensors, and a config that reloads correctly (num_hidden_layers=52fromlayers_block_type,n_shared_experts=1,num_nextn_predict_layers=0, prunedhidden_size/mamba_*/MoE dims, reconstructedhybrid_override_pattern).Pruning search log (
--prune_target_active_params 3e9)tests/examples/megatron_bridge/test_prune_minitron.py—nemotron_hnow exports to HF and reloads (previously it stopped at a Megatron checkpoint, since the dummy-model path neededtransformers<5), plus ann_shared_expertsconfig assertion; the deadmegatron_formatbranch is gone. It runs on the CI container: verified on nemo:26.06.01 (transformers 5.8.1) and nemo:26.08.rc6.tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py— thenas_memory_mbsearch test now passes acandidate_filterand asserts the exact number of rejected candidates (256 of the 512-combo grid) plus the surviving candidates' validity; itsexpected_top_kgoldens are regenerated accordingly. Becausemoe_shared_expert_intermediate_sizeis in that test's skip list, this also covers the model-config fallback for hparams that are not in the search space.Verified on 2 GPUs, on both the CI container (nemo:26.06.01) and nemo:26.08.rc6:
test_prune_minitron[qwen3]test_prune_minitron[deepseek_v3]candidate_filterend-to-endtest_prune_minitron[nemotron_h]test_mcore_mamba_hybrid_pruning_nas_memory_mbtest_mcore_mamba_hybrid_pruning_nas_paramsBefore your PR is "Ready for review"
candidate_filterdefaults toNone(existing searches unchanged), and the config-only export is limited to hybrid providers on nemo:26.08+, so dense / MoE / VLM exports keep the path they use today.CONTRIBUTING.md: N/AAdditional Information
Enables the Prune + Distill workflow for
NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16(native, no-remote-codeNemotronHForCausalLMwith MTP heads). Pruning-time MTP support was scoped and intentionally deferred — MTP heads are dropped and can be re-derived via a short SFT withmtp_num_layers=1on the pruned+distilled model.