Add NAS to Minitron pruning for parameter based auto-pruning #720
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #720 +/- ##
==========================================
- Coverage 74.23% 74.18% -0.05%
==========================================
Files 192 192
Lines 19033 19240 +207
==========================================
+ Hits 14129 14274 +145
- Misses 4904 4966 +62 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
kevalmorabia97
force-pushed
the
kmorabia/minitron-auto
branch
from
621fcaf to
408a574
Compare
kevalmorabia97
force-pushed
the
kmorabia/minitron-auto
branch
3 times, most recently
from
11bc408 to
1cbe285
Compare
kevalmorabia97
changed the title
kevalmorabia97
requested review from
AAnoosheh,
ChenhanYu,
LianaMikael,
danielkorzekwa and
realAsma
kevalmorabia97
changed the title
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the 📝 WalkthroughWalkthroughAdds parameter-count–based automatic NAS pruning for Minitron alongside existing manual export_config pruning; updates configs/APIs with granular divisors, implements search-and-prune utilities, extends docs and examples, and updates tests and utilities to exercise both pruning modes. Changes
Sequence DiagramsequenceDiagram
participant User
participant MCoreMinitronSearcher
participant SearchSpace
participant Pruner
participant Model
participant Validator
User->>MCoreMinitronSearcher: run_search(model, constraints)
activate MCoreMinitronSearcher
MCoreMinitronSearcher->>SearchSpace: _generate_search_space_combos(constraints)
activate SearchSpace
SearchSpace-->>MCoreMinitronSearcher: candidate_configs
deactivate SearchSpace
MCoreMinitronSearcher->>Pruner: sort_and_select_top_k(candidate_configs)
activate Pruner
Pruner-->>MCoreMinitronSearcher: top_k_configs
deactivate Pruner
loop for each config in top_k_configs
MCoreMinitronSearcher->>Model: _prune(model, config)
activate Model
Model-->>MCoreMinitronSearcher: pruned_model
deactivate Model
MCoreMinitronSearcher->>Validator: evaluate(pruned_model)
activate Validator
Validator-->>MCoreMinitronSearcher: score
deactivate Validator
MCoreMinitronSearcher->>MCoreMinitronSearcher: record_candidate(config, score)
end
MCoreMinitronSearcher-->>User: best_architecture
deactivate MCoreMinitronSearcher
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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: 4
🤖 Fix all issues with AI agents
In @.vscode/settings.json:
- Around line 43-45: Add the missing VS Code Python extension setting by adding
"python.analysis.extraPaths": ["./tests/"] alongside the existing
"cursorpyright.analysis.extraPaths": ["./tests/"] in the settings JSON; ensure
the new key is placed at the same object level as
"cursorpyright.analysis.extraPaths" and that the resulting .vscode settings
remain valid JSON (no trailing comma errors) so both IDEs use ./tests/ on the
Python path.
In @examples/pruning/README.md:
- Line 32: Fix the typo in the README sentence by replacing "requisred" with
"required" in the line that reads "For GradNAS pruning for Hugging Face BERT /
GPT-J, no additional dependencies are requisred." so it reads "For GradNAS
pruning for Hugging Face BERT / GPT-J, no additional dependencies are required."
In @modelopt/torch/opt/searcher.py:
- Line 215: The print_rank_0("Running forward loop...") log bypasses the silence
control; move that call inside the existing context manager (with no_stdout() if
silent else nullcontext():) or wrap it with an if not silent check so it only
runs when silent is False; update the code around print_rank_0, the no_stdout()
context usage, and the silent parameter passed from self.config["verbose"] to
ensure the message respects the provided silence flag.
🧹 Nitpick comments (3)
modelopt/torch/utils/network.py (1)
152-152: Optional: Fix grammar in docstring.The phrase "This can helpful" is missing "be" and should read "This can be helpful".
✏️ Proposed grammar fix
- This can helpful for MoE or dynamic modules, where the state dict might contain extra parameters that + This can be helpful for MoE or dynamic modules, where the state dict might contain extra parameters thatmodelopt/torch/prune/plugins/mcore_minitron.py (1)
504-510: Input dictionary mutation may cause side effects.The method mutates the input
search_spacedictionary by callingpop()onhparams_to_skip. If the caller expects the dictionary to remain unchanged (e.g., for logging or reuse), this could cause unexpected behavior.♻️ Create a copy before mutation
+ # Create a copy to avoid mutating the input + search_space = dict(search_space) + if hparams_to_skip: print_rank_0(f"Skipping {hparams_to_skip=} during search space generation...") for hparam in hparams_to_skip: if hparam in search_space: search_space.pop(hparam) else: warn(f"Hparam {hparam} not found in search space! Skipping...")tests/gpu/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py (1)
288-289: Hardcoded parameter count may be fragile.The assertion
assert param_count == 31776.0is a hardcoded check. If model architecture or default configurations change, this will fail silently without indicating what changed.Consider adding a comment explaining how this value was computed, or computing it from model dimensions to make the test more self-documenting.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (20)
.vscode/settings.jsonCHANGELOG.rstexamples/pruning/README.mdmodelopt/torch/nas/plugins/megatron.pymodelopt/torch/nas/search_space.pymodelopt/torch/opt/searcher.pymodelopt/torch/prune/__init__.pymodelopt/torch/prune/plugins/mcore_minitron.pymodelopt/torch/prune/pruning.pymodelopt/torch/utils/logging.pymodelopt/torch/utils/network.pymodelopt/torch/utils/plugins/megatron_generate.pymodelopt/torch/utils/plugins/megatron_mmlu.pymodelopt/torch/utils/plugins/megatron_preprocess_data.pytests/_test_utils/torch/megatron/models.pytests/_test_utils/torch/nas_prune/minitron_common.pytests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.pytests/gpu/torch/nas/plugins/test_megatron_mamba_dynamic_modules.pytests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.pytests/gpu/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
💤 Files with no reviewable changes (1)
- modelopt/torch/prune/init.py
🧰 Additional context used
🧬 Code graph analysis (8)
modelopt/torch/utils/network.py (1)
modelopt/torch/opt/dynamic.py (1)
DynamicModule(338-896)
modelopt/torch/opt/searcher.py (2)
modelopt/torch/utils/logging.py (2)
no_stdout(99-103)print_rank_0(106-109)modelopt/torch/utils/network.py (1)
run_forward_loop(510-580)
tests/_test_utils/torch/nas_prune/minitron_common.py (2)
modelopt/torch/prune/pruning.py (1)
prune(31-210)modelopt/torch/prune/plugins/mcore_minitron.py (1)
get_mcore_minitron_config(630-657)
tests/gpu/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py (3)
modelopt/torch/prune/plugins/mcore_minitron.py (2)
get_mcore_param_count(539-545)get_mcore_minitron_config(630-657)tests/_test_utils/torch/nas_prune/minitron_common.py (1)
prune_minitron(19-37)tests/_test_utils/torch/megatron/utils.py (1)
run_mcore_inference_with_dummy_input(122-129)
modelopt/torch/prune/plugins/mcore_minitron.py (2)
modelopt/torch/nas/utils.py (3)
get_subnet_config(166-176)sample(137-148)sort_parameters(187-197)modelopt/torch/nas/search_space.py (2)
sample(95-128)sort_parameters(131-164)
modelopt/torch/nas/plugins/megatron.py (4)
modelopt/torch/opt/dynamic.py (2)
get_hparam(801-803)get_hparam(1217-1222)modelopt/torch/utils/network.py (1)
make_divisible(200-216)modelopt/torch/opt/utils.py (1)
get_hparam(74-76)modelopt/torch/opt/hparam.py (3)
choices(123-125)choices(128-141)original(157-159)
tests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py (2)
modelopt/torch/prune/plugins/mcore_minitron.py (2)
get_mcore_minitron_config(630-657)_generate_search_space_combos(472-536)tests/_test_utils/torch/nas_prune/minitron_common.py (1)
prune_minitron(19-37)
tests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py (2)
modelopt/torch/nas/modules/container.py (1)
DynamicModuleList(103-131)modelopt/torch/nas/plugins/megatron.py (9)
NumAttentionHeadsHp(256-295)choices(764-766)active_slice(279-295)active_slice(696-705)active_slice(732-748)active(725-729)_get_output_size_indices(317-379)expand_head_indices(248-253)_get_input_size_indices(415-425)
🪛 LanguageTool
examples/pruning/README.md
[grammar] ~32-~32: Ensure spelling is correct
Context: .../ GPT-J, no additional dependencies are requisred. ## Getting Started As part of the pr...
(QB_NEW_EN_ORTHOGRAPHY_ERROR_IDS_1)
⏰ 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 (37)
modelopt/torch/utils/logging.py (1)
46-53: LGTM! Aligns with ML community conventions.The change from "G" (giga) to "B" (billion) for the third magnitude unit is appropriate for this logging utility in an ML context. The ML/AI community standardly uses "B" notation for parameter counts (e.g., "GPT-3 175B", "LLaMA 7B"), and this PR specifically deals with parameter-based pruning where consistency with that convention improves user experience.
tests/_test_utils/torch/megatron/models.py (3)
316-316: LGTM! Backward-compatible parameter addition.The new
mamba_num_headsparameter is correctly defined as optional with a None default, maintaining backward compatibility. The type hint and placement among other Mamba-specific parameters are appropriate.
350-350: LGTM! Parameter correctly propagated to TransformerConfig.The
mamba_num_headsparameter is properly passed through to the TransformerConfig constructor, consistent with how other Mamba-specific parameters are handled.
362-362: LGTM! Helpful version context added.The comment clarifies that MoE support is available in Megatron-Core 0.16+, which is useful context for understanding the conditional logic.
modelopt/torch/utils/network.py (1)
115-121: LGTM! Proper defensive guard for DynamicModule.The check correctly prevents incorrect parameter counting for DynamicModule instances where
model.parameters()doesn't account for dynamic slicing. The error message clearly directs users to the appropriate alternative function.modelopt/torch/utils/plugins/megatron_mmlu.py (1)
50-51: LGTM! Clean public API declaration.The
__all__declaration correctly exportsmegatron_mmluas the module's public interface. This aligns with similar updates across other plugin modules and follows Python best practices for explicit API exposure.modelopt/torch/utils/plugins/megatron_generate.py (1)
27-28: LGTM! Correctly exports both public generation functions.The
__all__declaration appropriately exportsmegatron_generateandmegatron_prefillas the module's public API. Internal utilities likeget_current_memory_infoare correctly excluded from the export list.modelopt/torch/utils/plugins/megatron_preprocess_data.py (1)
45-46: LGTM! Proper public API declaration for the preprocessing utility.The
__all__declaration correctly exportsmegatron_preprocess_dataas the module's public function. Private implementation classes (_Encoder,_Partition) and the CLI-onlymain()function are appropriately excluded.modelopt/torch/opt/searcher.py (1)
38-38: LGTM!The import addition is correct and properly used in the new log statement.
modelopt/torch/nas/search_space.py (1)
138-138: LGTM - Simplified logging output.The removal of verbose summary printing streamlines the output while retaining the essential sorting notification.
modelopt/torch/prune/pruning.py (1)
81-81: LGTM - Documentation accurately reflects the updated API.The wording change correctly indicates that
export_configis now optional (for manual pruning) rather than required, aligning with the new NAS-based auto pruning mode that usesparamsconstraints instead.CHANGELOG.rst (2)
17-17: LGTM - Clear feature documentation.The changelog entry accurately describes the new NAS-based auto pruning capability and references the appropriate documentation.
84-84: LGTM - Improved documentation link.The direct anchor link provides better user experience by navigating directly to the pruning guidelines.
tests/gpu/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py (2)
54-54: LGTM - Helpful clarification.The comment clearly explains the pattern semantics for better test readability.
78-91: LGTM - Test updated to use new configuration API.The test correctly uses the new
get_mcore_minitron_confighelper with explicit divisor parameters, aligning with the broader API changes while preserving test semantics.examples/pruning/README.md (4)
28-28: LGTM - Updated prerequisites.The container version update and pip upgrade instruction ensure users have the latest dependencies for the new features.
40-167: LGTM - Excellent documentation structure.The comprehensive documentation clearly distinguishes between Manual Pruning and NAS-based Auto Pruning modes, with:
- Clear mode descriptions and use cases
- Well-structured code examples
- Detailed "Under the Hood" explanations
- Helpful tips and notes
- Advanced configuration guidance
This will significantly improve user understanding and adoption of the new NAS-based auto pruning feature.
172-175: LGTM - Clear support matrix.The updated matrix effectively distinguishes between manual (
export_config) and auto (paramswithscore_func) pruning modes for Minitron.
179-206: LGTM - Well-organized examples section.The reorganized examples section provides clear pathways for users to find relevant framework-specific documentation and examples, with good references to external resources and tutorials.
modelopt/torch/prune/plugins/mcore_minitron.py (7)
166-171: LGTM - Clean dataclass definition.The
CandidateSubnetdataclass is a clean way to encapsulate subnet candidate information. Consider adding__slots__for memory efficiency if many candidates are expected, but this is optional.
227-230: LGTM - Clear constraint validation.The assertion properly validates that exactly one constraint type is provided and that it's one of the supported types (
export_configorparams).
322-365: LGTM - Comprehensive pruning logic.The
_prunemethod correctly handles both width and depth pruning, updates model config, and reinitializes MoE token dispatcher. The kv_channels handling for None case is a good defensive measure.
630-657: LGTM - Well-structured configuration API.The keyword-only signature enforces explicit parameter naming, improving code clarity. The recursive
_set_divisorshelper correctly handles nested configuration dictionaries.
154-161: LGTM - Clean layer filtering and reindexing.The refactored code cleanly separates the filtering of kept layers from their reindexing, making the logic easier to follow.
875-875: LGTM - Helpful logging for checkpoint restoration.The added log message improves observability when activations and scores are loaded from a checkpoint.
583-588: Parameter counting logic for decoder layers is correct and properly handles selective layer counting.The condition
("decoder.layers." not in name or layer_numbers_to_count is None)correctly excludes decoder layer parameters from the initial sum whenlayer_numbers_to_count is not None, then adds them back selectively in the loop. This is tested indirectly through the pruning test suite, including uneven pipeline parallelism scenarios.modelopt/torch/nas/plugins/megatron.py (2)
27-48: LGTM - Cleaned up imports.The import changes remove unused parallel state utilities and add the required
make_divisibleutility.
1037-1043: LGTM - Cleaner divisor handling with loop.The refactored loop pattern for applying divisors to
hidden_sizeandnum_layersis cleaner and more maintainable than separate blocks of code. The preservation of the original value via| {hp.original}ensures the model can always restore to its original configuration.tests/_test_utils/torch/nas_prune/minitron_common.py (1)
19-36: LGTM - API updated for constraints-based pruning.The signature change from
export_configtoconstraintsproperly reflects the expanded pruning API that now supports bothexport_configandparamsconstraint types. The explicit divisor parameters inget_mcore_minitron_configprovide clear configuration.Note: The hardcoded divisors (
mamba_head_dim_divisor=4,num_moe_experts_divisor=1,num_layers_divisor=1) may need to be parameterized if tests require different configurations in the future.tests/gpu/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py (4)
86-91: LGTM - Updated config with explicit divisors.The explicit divisor parameters make the test configuration clearer and align with the updated
get_mcore_minitron_configAPI.
175-176: Verify reduced iterations don't affect test reliability.Forward loop iterations reduced from 5 to 2. While this speeds up tests, ensure it still collects sufficient activations for reliable importance estimation in the test scenarios.
345-373: Good handling of non-deterministic layer ordering.The test correctly handles the fact that layer ordering can vary based on PP configuration by checking
sorted_layersand providing expected values for known orderings. TheRuntimeErrorwith a clear message (line 372) helps identify when new orderings need to be handled.The hardcoded expected values are extensive but necessary for deterministic validation. Consider extracting them to a constant or helper function if they need to be reused or modified frequently.
385-390: LGTM - Appropriate test skipping conditions.The test correctly skips for configurations that aren't supported:
- More than 4 GPUs (line 388)
- MoE support in Mamba (line 390) with a TODO note
tests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py (2)
205-215: LGTM - Consistent constraint wrapper usage.The test correctly wraps the
export_configin a constraints dictionary and passes it toprune_minitron, aligning with the updated API.
510-534: LGTM - Thorough unit test for search space generation.The test validates the
_generate_search_space_combosmethod with:
- Width pruning (50%) correctly filters
hidden_sizeto[96, 128, 160]andnum_attention_headsto[24, 32]- Depth pruning (25%) correctly filters
num_layersto[7, 8]hparams_to_skipcorrectly removesffn_hidden_sizefrom the search space- Total combinations: 3 × 2 × 2 = 12 ✓
The explicit verification of all 12 combinations ensures deterministic behavior.
tests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py (2)
86-98: LGTM - Explicit config for search space test.The updated
mtn.convertcall with explicit divisor configuration makes the test setup clearer and more maintainable.
170-236: LGTM - Comprehensive test for attention head sorting and pruning.This test thoroughly validates the
NumAttentionHeadsHpclass:
- Verifies choices are multiples of
num_query_groups- Tests importance-based ranking with
enforce_order- Validates
active_slicereturns correct indices after pruning- Verifies QKV and projection layer index mappings are correct
The cleanup via
destroy_model_parallel()is important since this test doesn't use the spawned multiprocess pattern.
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
Signed-off-by: Keval Morabia <[email protected]>
kevalmorabia97
force-pushed
the
kmorabia/minitron-auto
branch
from
745e960 to
e41b108
Compare
| """ | ||
| from modelopt.torch.opt.dynamic import DynamicModule | ||
|
|
||
| if isinstance(network, DynamicModule): |
There was a problem hiding this comment.
this seems to be applicable only for pruning. Quantization have QuantModule base block - should we have something similar for pruning and do:
if isinstance(network, PruneModule):
or
if hasattr(network, "some_special_pruning_attr")
There was a problem hiding this comment.
Since quantization doesnt change param count, I dont see a reason to use this function. Besides, users can still use this on Quantized + Exported model since its not longer DynamicModule. Currently we will just raise an exception if Quantized DynamicModule but later if we need to support it, we can think of such a way. Wdyt?
Signed-off-by: Keval Morabia <[email protected]>
kevalmorabia97
force-pushed
the
kmorabia/minitron-auto
branch
from
69f57f4 to
c01ced9
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Around line 543-549: The code is mutating the input dict search_space when
hparams_to_skip is provided; instead create a local shallow copy (e.g.,
new_search_space = dict(search_space) or use search_space.copy()) before
performing the pop operations so the caller's dict isn't changed, then operate
on and return/use that local copy for the rest of the search-space generation
path; update the block that references hparams_to_skip, the pop calls, and any
subsequent uses of search_space in the surrounding function (keep print_rank_0
and warn calls as-is but point them at the copied dict).
♻️ Duplicate comments (3)
.vscode/settings.json (1)
43-45: Add VS Code Python extension extraPaths alongside Cursor.This drops
python.analysis.extraPaths, which likely breaks the documented VS Code setup. Please add it back alongsidecursorpyright.analysis.extraPaths.✅ Proposed fix
"cursorpyright.analysis.extraPaths": [ "./tests/" // add tests to python path just like pytest does in pyproject.toml - ], + ], + "python.analysis.extraPaths": [ + "./tests/" + ],examples/pruning/README.md (1)
32-32: Fix typo: "requisred" should be "required".📝 Proposed fix
-For GradNAS pruning for Hugging Face BERT / GPT-J, no additional dependencies are requisred. +For GradNAS pruning for Hugging Face BERT / GPT-J, no additional dependencies are required.modelopt/torch/prune/plugins/mcore_minitron.py (1)
479-491: Layer reversion logic may still leave model in inconsistent state.The reversion code saves
all_layersandstart_layer_numberbefore_prune, then attempts to restore after scoring. However,_prunecallsdrop_mcore_language_model_layerswhich modifiesmodel.config.num_layerspermanently (line 163 indrop_mcore_language_model_layers). Thesample(max)call only resets hyperparameters, notmodel.config.num_layers.This could cause assertions to fail at line 304 (
assert sorted(self.sorted_layers) == list(range(1, self.model.config.num_layers + 1))) on subsequent iterations sincemodel.config.num_layerswould be the pruned count, not the original.🛠️ Suggested fix
Save and restore
model.config.num_layersalongside the layer list:if candidate.score is None: # not restored from checkpoint all_layers = self.model.decoder.layers start_layer_number = all_layers[0].layer_number + original_num_layers = self.model.config.num_layers self._prune(candidate.ss_config, prune_depth=True) candidate.score = self.eval_score(silent=False) self.save_search_checkpoint(verbose=False) # reset to max subnet and revert dropped layers sample(self.model, sample_func=max) for layer in all_layers: layer.layer_number = start_layer_number start_layer_number += 1 self.model.decoder.layers = all_layers + self.model.config.num_layers = original_num_layers
🧹 Nitpick comments (3)
tests/_test_utils/torch/megatron/models.py (1)
316-352: Validatemamba_num_headsconstraints (if required by Megatron).If upstream expects
mamba_num_headsto be compatible withhidden_size/mamba_head_dim, add a fast‑fail check so tests fail early with a clear error. Please verify the current Megatron TransformerConfig/Mamba requirements before adding the guard.Possible guard (adjust per upstream rules)
def get_mcore_mamba_hybrid_model( @@ config = TransformerConfig( @@ mamba_state_dim=mamba_state_dim, mamba_num_heads=mamba_num_heads, mamba_head_dim=mamba_head_dim, mamba_num_groups=mamba_num_groups, @@ ) + if mamba_num_heads is not None: + assert mamba_num_heads > 0 + # If upstream expects divisibility or exact match, enforce it here: + # assert hidden_size % mamba_num_heads == 0 + # assert mamba_num_heads * mamba_head_dim == hidden_sizemodelopt/torch/prune/plugins/mcore_minitron.py (1)
166-171: Consider usingfrozen=Truefor the dataclass.
CandidateSubnetis used as a data container wheress_configandparamsshould be immutable after creation. Onlyscoregets updated. Consider usingfrozen=Trueor makingscoreexplicitly mutable via a different pattern to prevent accidental modifications.tests/gpu/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py (1)
510-534: Good test coverage for the new NAS search space generation.The test thoroughly verifies
_generate_search_space_combos:
- Validates correct filtering based on
max_width_pruningandmax_depth_pruningthresholds.- Confirms
hparams_to_skipproperly excludesffn_hidden_size.- Asserts both the count and exact expected combinations.
One minor note: This test doesn't require GPU and could potentially be placed in a non-GPU test module to speed up CI when GPU resources are constrained. Consider moving to a unit test file if reorganizing tests in the future.
Signed-off-by: Keval Morabia <[email protected]>
kevalmorabia97
force-pushed
the
kmorabia/minitron-auto
branch
from
246ffc8 to
c735e7b
Compare
2 participants
What does this PR do?
Type of change: New feature
hidden_sizemultiple of 256,ffn_hidden_sizemultiple of 512) to make the process efficient. Users can configure this if they want to.Usage
Pruning API is same as before:
To configure search space (shows defaults):
Testing
Qwen3-8B -> 6B (~2 hours on 8xA5000)
Nemotron-Nano-9B-v2 -> 7B (~2.5 hours on 8xA5000)
Before your PR is "Ready for review"
Additional Information
OMNIML-3043
Summary by CodeRabbit
Release Notes
New Features
Documentation
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.