MBridge pruning minor fix for saving pruned NemotronH

[kevalmorabia97]

What does this PR do?

Type of change: Bug fix

Testing

Nemotron Nano v2 pruned can be saved

Summary by CodeRabbit

• Bug Fixes

  • Fixed Hugging Face model loading to properly respect the trust_remote_code parameter during model instantiation.

• Improvements

  • Enhanced distributed training logging with rank-0 aware warning and logging mechanisms for cleaner, non-redundant output in multi-GPU and multi-node scenarios.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes introduce rank-0 aware logging improvements in a searcher utility, update a method signature to respect a trust_remote_code flag during model loading, and add documentation comments in example files regarding future validation configuration refactoring.

Changes

Cohort / File(s)	Summary
Validation Config Documentation examples/megatron_bridge/distill.py, modelopt/torch/utils/plugins/mbridge.py	Added TODO comments and commented-out validation config lines indicating planned migration to Nemo 26.04 validation configuration pattern.
Model Loading API Update examples/megatron_bridge/prune_minitron.py	Updated AutoBridge.from_hf_pretrained() call to include trust_remote_code parameter, enabling the method to respect this flag during bridge instantiation.
Rank-0 Aware Logging modelopt/torch/opt/searcher.py	Replaced generic warnings and print statements with rank-0 specific helpers (warn_rank_0, print_rank_0) for checkpoint validation and state loading operations.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~15 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: a minor fix for MBridge pruning to enable saving of pruned NemotronH models, which aligns with the primary functional change in prune_minitron.py.
Merge Conflict Detection	✅ Passed	✅ No merge conflicts detected when merging into main

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch kmorabia/mbridge-minor-pruning

---

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments

modelopt/torch/opt/searcher.py (1)

266-267: Nit: save_search_checkpoint still uses bare print instead of print_rank_0.

Line 267 uses print(...) while the rest of this file now uses print_rank_0. The dist.is_master() guard on line 262 makes this safe, but switching to print_rank_0 would be more consistent and would let you drop the redundant master check for the verbose path.

Proposed fix

         if verbose:
-            print(f"Saving searcher state to {checkpoint}...")
+            print_rank_0(f"Saving searcher state to {checkpoint}...")

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.74%. Comparing base (95511a0) to head (5a25091).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #887   +/-   ##
=======================================
  Coverage   73.73%   73.74%           
=======================================
  Files         199      199           
  Lines       21165    21163    -2     
=======================================
  Hits        15606    15606           
+ Misses       5559     5557    -2     

☔ View full report in Codecov by Sentry.
📢 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.

[jenchen13]

will you implement this in the current PR or a later PR?

[kevalmorabia97]

Later. It wont be merged into 26.02 container so will raise error if we do it now but wanted to keep it for reference incase we try to mount latest m-bridge and run the script