NAS export refactor + skip conversion on minitron restore #424
Conversation
…f mtp.plugins.* Signed-off-by: Keval Morabia <[email protected]>
WalkthroughReplaces the legacy "export" path with a new "export_nas" export workflow, adds NAS export primitives (ExportConfig, export_searchspace, restore_export), rewires mode descriptors, conditionally loads the mcore_minitron plugin, expands plugin exports, updates docs, and adapts tests to the new "export_nas" name and behavior. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant App
participant ModeMgr as ModeDescriptor
participant NASConv as nas.conversion
participant Model
Note over ModeMgr: Legacy flow (before change)
User->>App: apply_mode(..., mode="export")
App->>ModeMgr: select "export"
ModeMgr->>Model: legacy export/convert
Model-->>User: exported model
Note over ModeMgr: New NAS export flow (this change)
User->>App: apply_mode(..., mode="export_nas")
App->>ModeMgr: select "export_nas"
ModeMgr->>NASConv: export_searchspace(model, ExportConfig)
NASConv->>Model: in-place subnet export + optional BN calibration
NASConv-->>ModeMgr: metadata (subnet_config, patches)
ModeMgr-->>User: exported model + metadata
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
examples/megatron-lm/README.md (1)
136-139: Consider showing a complete usage example.The TIP explains uneven PP configuration but doesn't demonstrate how to integrate
MLM_EXTRA_ARGSinto the pruning command shown above (lines 129-134). Consider adding a concrete example:+ +Example with uneven PP: +```sh +PP=4 \ +TARGET_NUM_LAYERS=24 \ +HF_MODEL_CKPT=<pretrained_model_name_or_path> \ +MLM_MODEL_SAVE=Qwen3-8B-Pruned \ +MLM_EXTRA_ARGS="--decoder-first-pipeline-num-layers 7 --decoder-last-pipeline-num-layers 5" \ +bash megatron-lm/examples/post_training/modelopt/prune.sh qwen/Qwen3-8B +```
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
examples/megatron-lm/README.md(2 hunks)modelopt/torch/prune/__init__.py(1 hunks)modelopt/torch/prune/plugins/mcore_minitron.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
modelopt/torch/prune/plugins/mcore_minitron.py (1)
modelopt/torch/nas/plugins/megatron.py (1)
drop_mcore_language_model_layers(1392-1456)
⏰ 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: build-docs
- GitHub Check: code-quality
🔇 Additional comments (6)
examples/megatron-lm/README.md (2)
113-113: LGTM!Good addition of links to the pruning getting started section and guidelines. This helps users find the detailed documentation they need.
126-134: Verify the path change is intentional.The example now uses a relative path
Qwen3-8B-Prunedinstead of the absolute path/tmp/Qwen3-8B-DPruned. This changes the behavior:
- Relative path: saves to
$PWD/Qwen3-8B-Pruned- Absolute path: saves to
/tmp/Qwen3-8B-PrunedEnsure this change aligns with user expectations and is consistent with other examples in the file (e.g., lines 61, 67, 88, 94 use
/tmp/prefix).modelopt/torch/prune/__init__.py (2)
24-24: LGTM!Good addition of the
import_pluginutility to support conditional plugin loading.
29-30: LGTM!The conditional plugin loading pattern is appropriate for the
mcore_minitronplugin. This allows the plugin to be loaded only when its dependencies are available, preventing import errors when Megatron-Core is not installed. Theverbose=Falseflag suppresses unnecessary logging during import.modelopt/torch/prune/plugins/mcore_minitron.py (2)
40-40: LGTM!Good addition of the
drop_mcore_language_model_layersimport. This helper function is now available for re-export, making it accessible to users at the plugin level.
74-80: LGTM!The expanded
__all__makes the plugin's public API clearer and aligns with the PR objective to bring minitron imports one level up. Users can now access:
SUPPORTED_HPARAMS: For discovering supported pruning hyperparameters- Configuration and descriptor classes for the pruning mode
drop_mcore_language_model_layers: Helper function for manual layer droppingThis enables imports like
from modelopt.torch.prune.plugins.mcore_minitron import drop_mcore_language_model_layersor through the parent package asfrom modelopt.torch.prune import mcore_minitron(withmcore_minitron.drop_mcore_language_model_layers).
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #424 +/- ##
==========================================
+ Coverage 73.36% 73.38% +0.01%
==========================================
Files 180 180
Lines 17919 17934 +15
==========================================
+ Hits 13147 13160 +13
- Misses 4772 4774 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
kevalmorabia97
changed the title
Signed-off-by: Keval Morabia <[email protected]>
kevalmorabia97
force-pushed
the
kmorabia/pruning-doc-update-2
branch
from
1016d82 to
7a4394e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
modelopt/torch/nas/autonas.py(3 hunks)modelopt/torch/nas/conversion.py(3 hunks)modelopt/torch/nas/registry.py(1 hunks)modelopt/torch/nas/utils.py(0 hunks)modelopt/torch/prune/fastnas.py(1 hunks)modelopt/torch/prune/plugins/mcore_minitron.py(4 hunks)tests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py(3 hunks)tests/unit/torch/nas/plugins/test_hf_nas_save_restore.py(1 hunks)tests/unit/torch/nas/test_nas.py(3 hunks)tests/unit/torch/opt/test_chaining.py(3 hunks)
💤 Files with no reviewable changes (1)
- modelopt/torch/nas/utils.py
✅ Files skipped from review due to trivial changes (1)
- tests/unit/torch/nas/plugins/test_hf_nas_save_restore.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/prune/fastnas.py
- modelopt/torch/nas/registry.py
- tests/unit/torch/nas/test_nas.py
🧰 Additional context used
🧬 Code graph analysis (4)
modelopt/torch/nas/autonas.py (5)
modelopt/torch/opt/config.py (1)
get_kwargs_for_create_model_with_rules(322-383)modelopt/torch/nas/search_space.py (1)
generate_search_space(199-260)modelopt/torch/nas/utils.py (3)
get_subnet_config(160-170)sample(131-142)select(145-157)modelopt/torch/prune/fastnas.py (2)
sample(141-144)export_mode(349-351)modelopt/torch/prune/plugins/mcore_minitron.py (1)
export_mode(305-307)
modelopt/torch/prune/plugins/mcore_minitron.py (3)
modelopt/torch/nas/plugins/megatron.py (1)
drop_mcore_language_model_layers(1392-1456)modelopt/torch/nas/autonas.py (1)
export_mode(680-682)modelopt/torch/prune/fastnas.py (1)
export_mode(349-351)
modelopt/torch/nas/conversion.py (6)
modelopt/torch/opt/config.py (2)
ModeloptBaseConfig(59-147)ModeloptField(50-53)modelopt/torch/opt/conversion.py (2)
ApplyModeError(314-315)apply_mode(342-429)modelopt/torch/opt/mode.py (2)
ModeDescriptor(56-259)_ModeRegistryCls(267-344)modelopt/torch/utils/network.py (2)
compare_dict(423-427)unwrap_model(430-454)modelopt/torch/nas/search_space.py (1)
SearchSpace(38-196)modelopt/torch/nas/utils.py (2)
get_subnet_config(160-170)select(145-157)
tests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py (1)
tests/_test_utils/torch_dist/plugins/megatron_common.py (1)
run_mcore_inference(326-379)
⏰ 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: build-docs
- GitHub Check: code-quality
What does this PR do?
exporttoexport_nasmtp.plugins.mcore_minitron.*we can domtp.mcore_minitron.*which looks more neaterSummary by CodeRabbit
New Features
Documentation
Tests