minor rename

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

nemo/collections/llm/api.py

@@ -42,7 +42,7 @@
     PruningConfig,
     QuantizationConfig,
     Quantizer,
-    prune_gpt_model,
+    prune_language_model,
     save_pruned_model,
     set_modelopt_spec_if_exists_in_ckpt,
     setup_trainer_and_restore_model_with_modelopt_spec,
@@ -377,7 +377,7 @@ def prune(
         trainer_kwargs={"max_steps": steps, "limit_val_batches": steps, "val_check_interval": steps},
         model_config_overrides={"sequence_parallel": False},
     )
-    prune_gpt_model(model, pruning_config, data, trainer)
+    prune_language_model(model, pruning_config, data, trainer)
     save_pruned_model(trainer, save_path)
 
     console = Console()

nemo/collections/llm/modelopt/prune/__init__.py

@@ -14,6 +14,6 @@
 
 """Prune utilities for using TensorRT Model Optimizer."""
 
-from .pruner import PruningConfig, prune_gpt_model, save_pruned_model
+from .pruner import PruningConfig, prune_language_model, save_pruned_model
 
-__all__ = ["PruningConfig", "prune_gpt_model", "save_pruned_model"]
+__all__ = ["PruningConfig", "prune_language_model", "save_pruned_model"]

nemo/collections/llm/modelopt/prune/pruner.py

@@ -81,13 +81,13 @@ def __post_init__(self):
                 raise ValueError("drop_layers cannot be used with other pruning parameters")
 
 
-def prune_gpt_model(
+def prune_language_model(
     model: llm.GPTModel,
     pruning_config: PruningConfig,
     data_module: pl.LightningDataModule | None = None,
     trainer: nl.Trainer | None = None,
 ) -> llm.GPTModel:
-    """Prune a GPT model in-place based on the provided pruning configuration.
+    """Prune a GPT / Mamba (sub-class of GPT) model in-place based on the provided pruning configuration.
 
     Args:
         model (llm.GPTModel): The model to prune.
@@ -101,7 +101,7 @@ def prune_gpt_model(
         llm.GPTModel: The pruned model.
     """
     if pruning_config.drop_layers:
-        mtp.plugins.drop_mcore_gpt_layers(model, layers_to_drop=pruning_config.drop_layers)
+        mtp.plugins.drop_mcore_language_model_layers(model, layers_to_drop=pruning_config.drop_layers)
     else:
         assert data_module is not None, "data_module is required to prune the model."
         assert trainer is not None, "trainer is required to prune the model."
@@ -116,7 +116,7 @@ def prune_gpt_model(
         }
         mtp.prune(
             model,
-            mode="mcore_gpt_minitron",
+            mode="mcore_minitron",
             constraints={"export_config": export_config},
             dummy_input=None,  # Not used
             config={"forward_loop": partial(llm.validate, data=data_module, trainer=trainer, tokenizer="model")},

requirements/requirements_nlp.txt

@@ -13,7 +13,7 @@ megatron_core
 multi-storage-client>=0.21.0
 nltk>=3.6.5
 numpy<2  # tensorstore has an implicit compiled dependency on numpy<2
-nvidia-modelopt[torch]==0.35.0
+nvidia-modelopt==0.35.0
 nvidia-resiliency-ext>=0.3.0,<1.0.0; platform_system != 'Darwin'
 nvtx
 opencc