Fix saving pruned Nemotron-3-Nano hybrid_override_pattern with MTP or Pipe symbols

[kevalmorabia97]

Saving pruned Nemotron-3-Nano (with MTP) to HF format raised an assertion which is fixed here

Tested on nemo:26.04 with transformers 4.57

Summary by CodeRabbit

• Bug Fixes

  • Improved exported hybrid layer patterns for pruned models by removing MTP and pipeline-parallel markers.
  • Ensured exported configurations accurately represent the model’s main layers.

• Tests

  • Added coverage for pruning models with an MTP prediction layer and hybrid override patterns.

[coderabbitai]

📝 Walkthrough

Walkthrough

The pruning export reconstructs hybrid_override_pattern from main-layer characters. The NemotronH test fixture enables one MTP prediction layer and configures the MTP hybrid pattern.

Changes

MTP-aware hybrid pattern export

Layer / File(s)	Summary
Normalize hybrid patterns during HF export examples/megatron_bridge/prune_minitron.py, tests/examples/megatron_bridge/test_prune_minitron.py	The export parses main-layer characters before writing the HF hybrid pattern. The NemotronH fixture enables one MTP prediction layer with mtp_hybrid_override_pattern="*E".

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: aanoosheh, chenhanyu, jenchen13

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The commit only changes hybrid_override_pattern assignment in prune_minitron.py; it adds no banned security patterns, # nosec comments, or dependencies.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: fixing the saving of hybrid_override_pattern when pruning Nemotron-3-Nano models with MTP or Pipe symbols.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch kmorabia/nemotron3-mtp-prune-fix

---

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

[coderabbitai]

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.

👉 Steps to fix this

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/examples/megatron_bridge/test_prune_minitron.py (1)

43-51: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Assert the normalized pattern in the HF output.

The fixture enables one MTP prediction layer, but the HF branch only checks that config.json exists and that the parameter count is below the target. Add assertions on the reloaded configuration to verify that hybrid_override_pattern contains only main-layer characters, has num_hidden_layers entries, and has MTP disabled. This prevents regressions that save the old MTP suffix or PP boundaries.

As per coding guidelines, tests must exercise the behavior they claim to validate. As per path instructions, integration tests must protect regressions and document expected behavior.

🤖 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/megatron_bridge/test_prune_minitron.py` around lines 43 - 51,
Update the HF-output assertions in the test case using
create_tiny_nemotron_h_dir to reload config.json and validate
hybrid_override_pattern contains only main-layer characters, has exactly
num_hidden_layers entries, and confirms MTP is disabled; retain the existing
file-existence and parameter-count checks.

Sources: Coding guidelines, 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/megatron_bridge/prune_minitron.py`:
- Around line 664-667: Update the hybrid pattern normalization block around
hybrid_override_pattern to write the normalized value to
text_cfg.hybrid_override_pattern when provider is a VLM (_HYBRID_PROVIDER_TYPES)
and that field exists. Preserve the existing normalization via
parse_main_layer_chars and continue using the root hf_cfg field for non-VLM
configurations.

---

Nitpick comments:
In `@tests/examples/megatron_bridge/test_prune_minitron.py`:
- Around line 43-51: Update the HF-output assertions in the test case using
create_tiny_nemotron_h_dir to reload config.json and validate
hybrid_override_pattern contains only main-layer characters, has exactly
num_hidden_layers entries, and confirms MTP is disabled; retain the existing
file-existence and parameter-count checks.

🪄 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: a4609c66-e8af-4634-97e2-28d37dec6bbe

📥 Commits

Reviewing files that changed from the base of the PR and between 77dbeb1 and a74ce3d.

📒 Files selected for processing (2)

• examples/megatron_bridge/prune_minitron.py
• tests/examples/megatron_bridge/test_prune_minitron.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.97%. Comparing base (77dbeb1) to head (1449998).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2061      +/-   ##
==========================================
+ Coverage   65.51%   66.97%   +1.45%     
==========================================
  Files         521      521              
  Lines       59812    59812              
==========================================
+ Hits        39185    40058     +873     
+ Misses      20627    19754     -873     

Flag	Coverage Δ
examples	43.03% <ø> (+2.07%)	⬆️
unit	55.42% <ø> (ø)

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

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

🚀 New features to boost your workflow:

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-04 14:45 UTC