mcore_gpt_minitron
Module implementing top-level mcore_gpt_minitron
pruning handler for NVIDIA Megatron-Core / NeMo models.
Minitron pruning algorithm uses activation magnitudes to estimate importance of neurons / attention heads in the model. More details on Minitron pruning algorithm can be found here: https://arxiv.org/pdf/2407.14679
Actual implementation is at modelopt.torch.nas.plugins.megatron
.
Classes
Searcher for Minitron pruning algorithm. |
Functions
Get supported models (inside function to avoid circular imports). |
- class MCoreGPTMinitronSearcher
Bases:
BaseSearcher
Searcher for Minitron pruning algorithm.
- before_search()
Optional pre-processing steps before the search.
- Return type:
None
- property default_search_config: dict[str, Any]
Get the default config for the searcher.
- property default_state_dict: dict[str, Any]
Return default state dict.
- run_search()
Run actual search.
- Return type:
None
- sanitize_search_config(config)
Sanitize the search config dict.
- Parameters:
config (dict[str, Any] | None) –
- Return type:
dict[str, Any]
- get_supported_model_config_map()
Get supported models (inside function to avoid circular imports).
- Return type:
dict[type, str]