Minor pruning update

[kevalmorabia97]

Minor cleanup

Summary by CodeRabbit

• New Features

  • Added an optional parameter to the model conversion API to allow selecting a custom registry for applying modes.
  • Pruning now performs model conversion automatically as part of the workflow, simplifying setup.

• Documentation

  • Updated pruning example: revised loop signature and evaluation helper usage to reflect the new workflow and parameters.

• Chores

  • General cleanup and alignment of internal flows with the updated conversion mechanism.

[coderabbitai]

Walkthrough

Updates adjust example code signatures, introduce a registry parameter to NAS conversion for mode application, and switch pruning to convert models before search. Public API changes: convert(model, mode, registry=NASModeRegistry). Other flows and error handling remain unchanged.

Changes

Cohort / File(s)	Summary
Examples: pruning usage examples/pruning/README.md	Forward loop signature changed to def forward_loop(_):. Updated evaluation call to evaluate_and_print_results(prefix, forward_step, train_iterator, model, ...).
NAS conversion registry support modelopt/torch/nas/conversion.py	convert gains registry: _ModeRegistryCls = NASModeRegistry and forwards it to apply_mode. Export function unchanged.
Pruning flow conversion modelopt/torch/prune/pruning.py	prune now calls mtn.convert (instead of apply_mode) to convert model to DynamicModule before mtn.search. No signature changes.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor Client
  participant Conversion as convert(model, mode, registry)
  participant Registry as registry.apply_mode
  Client->>Conversion: call convert(model, mode, registry?)
  Note over Conversion: Default registry = NASModeRegistry
  Conversion->>Registry: apply_mode(model, mode)
  Registry-->>Conversion: returns converted model
  Conversion-->>Client: converted model

Loading

sequenceDiagram
  autonumber
  actor User
  participant Prune as prune(model, ...)
  participant MTN as mtn.convert
  participant Search as mtn.search
  User->>Prune: invoke prune(model, ...)
  Prune->>MTN: convert(model, mode)
  MTN-->>Prune: DynamicModule model
  Prune->>Search: search(DynamicModule, ...)
  Search-->>Prune: search results
  Prune-->>User: results

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

I twitch my whiskers at refactors neat,
A registry hopping to set the beat.
Convert, then prune—search down the trail,
Dynamic burrows, no chance to fail.
In README fields I leave a mark,
A lighter pawprint—onward, hark! 🐇✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title Check	⚠️ Warning	The title “Minor pruning update” is overly generic and fails to capture the core changes introduced by this pull request, including the addition of a registry parameter to the conversion.convert function and the integration of mtn.convert into the pruning pipeline, as well as the accompanying README example updates. It also omits the broader context of enhancing the registry mechanism and replacing apply_mode calls, and using “minor” adds ambiguity about the significance of these modifications.	Please revise the title to succinctly reflect the primary modifications, for example “Add registry parameter to conversion.convert and integrate mtn.convert into pruning pipeline,” so that it clearly indicates both the signature change and the pruning integration for reviewers.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

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

---

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.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

examples/pruning/README.md (1)

58-59: Verify the undefined variables in the example code.

The example shows evaluate_and_print_results(prefix, forward_step, train_iterator, model, ...) but the variables prefix, forward_step, and train_iterator are not defined in the example snippet. Consider adding brief comments or placeholder definitions to make the example more complete and avoid confusion for users following the example.

Example addition:

+# Define the evaluation parameters
+prefix = "train"
+forward_step = ...  # Your forward step function
+train_iterator = ...  # Your training data iterator
+
 def forward_loop(_):
     evaluate_and_print_results(prefix, forward_step, train_iterator, model, ...)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 17439e6 and 706a3af.

📒 Files selected for processing (3)

• examples/pruning/README.md (3 hunks)
• modelopt/torch/nas/conversion.py (2 hunks)
• modelopt/torch/prune/pruning.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

modelopt/torch/nas/conversion.py (2)

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

• _ModeRegistryCls (267-344)

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

• apply_mode (342-429)

modelopt/torch/prune/pruning.py (1)

modelopt/torch/nas/conversion.py (1)

• convert (29-89)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (5)

• GitHub Check: linux
• GitHub Check: wait-checks / wait
• GitHub Check: wait-checks / wait
• GitHub Check: code-quality
• GitHub Check: build-docs

🔇 Additional comments (4)

modelopt/torch/prune/pruning.py (1)

201-202: LGTM! Clean refactoring to use the NAS conversion API.

The change from apply_mode to mtn.convert with PruneModeRegistry is consistent with the updated convert API that now accepts a registry parameter. This provides better separation of concerns by delegating mode application and DynamicModule conversion to the NAS layer.

examples/pruning/README.md (2)

69-69: Good clarification for handling model lists.

The comment helpfully clarifies that if the model is a list, users should pass model[0] to the prune API.

---

84-84: LGTM! More specific tutorial reference.

The updated link points to a specific NeMo tutorial for pruning and distillation, which is more helpful than a generic reference.

modelopt/torch/nas/conversion.py (1)

89-89: LGTM! Correct forwarding of registry parameter.

The registry parameter is properly forwarded to apply_mode, enabling callers to specify custom mode registries.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.78%. Comparing base (17439e6) to head (706a3af).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #397      +/-   ##
==========================================
- Coverage   73.79%   73.78%   -0.01%     
==========================================
  Files         171      171              
  Lines       17583    17582       -1     
==========================================
- Hits        12975    12973       -2     
- Misses       4608     4609       +1     

☔ 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.