Fix param count calculation + account depth pruning

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

Author: kevalmorabia97

.vscode/settings.json

@@ -40,7 +40,7 @@
         "--no-cov",
     ],
     "evenBetterToml.schema.enabled": false, // disable toml/json schema since we have custom fields
-    "python.analysis.extraPaths": [
+    "cursorpyright.analysis.extraPaths": [
         "./tests/" // add tests to python path just like pytest does in pyproject.toml
     ],
     "git.alwaysSignOff": true,

modelopt/torch/nas/plugins/megatron.py

@@ -1040,6 +1040,7 @@ def modify(
         ffn_hidden_size_divisor: int = 1,
         mamba_head_dim_divisor: int = 1,
         num_moe_experts_divisor: int = 1,
+        num_layers_divisor: int = 1,
     ):
         """Modify the dynamic choices of the module according to provided keyword arguments.
 
@@ -1048,10 +1049,15 @@ def modify(
             ffn_hidden_size_divisor: The divisor of the mlp ffn_hidden_size.
             mamba_head_dim_divisor: The divisor of the mamba head_dim.
             num_moe_experts_divisor: The divisor of the number of MoE experts.
+            num_layers_divisor: The divisor of the number of layers.
         """
-        hp = self.get_hparam("hidden_size")
-        choices = {int(make_divisible(c, hidden_size_divisor)) for c in hp.choices}  # type: ignore[arg-type]
-        hp.choices = list(set(hp.choices) & choices | {hp.original})
+        for hp_name, divisor in [
+            ("hidden_size", hidden_size_divisor),
+            ("num_layers", num_layers_divisor),
+        ]:
+            hp = self.get_hparam(hp_name)
+            choices = {int(make_divisible(c, divisor)) for c in hp.choices}  # type: ignore[arg-type]
+            hp.choices = list(set(hp.choices) & choices | {hp.original})
 
         for layer in self.decoder.layers:
             layer.modify(

modelopt/torch/prune/plugins/mcore_minitron.py

@@ -78,7 +78,7 @@ def prune(
                     constraints = {"params": "60%"}
 
                     # Specify export_config with pruned hyperparameters
-                    # This is supported and required if the model is converted via ``mcore_minitron`` mode.
+                    # This is supported only if the model is converted via ``mcore_minitron`` mode.
                     constraints = {
                         "export_config": {
                             "ffn_hidden_size": 128,

modelopt/torch/prune/pruning.py

@@ -46,7 +46,7 @@
 def num2hrb(num: float, suffix="") -> str:
     """Convert big floating number to human readable string."""
     step = 1000  # step between units is 1000
-    units = ["", "K", "M", "G", "T", "P", "E"]
+    units = ["", "K", "M", "B", "T", "P", "E"]
     while abs(num) >= step and len(units) > 1:
         num /= step
         units.pop(0)

modelopt/torch/utils/logging.py

@@ -112,6 +112,13 @@ def param_num(network: nn.Module, trainable_only: bool = False, unit=1e6) -> flo
     Returns:
         The number of parameters in the model in the given unit.
     """
+    from modelopt.torch.opt.dynamic import DynamicModule
+
+    if isinstance(network, DynamicModule):
+        # NOTE: model.parameters() doesnt consider active_slice so we dont get sorted or trimmed parameters!
+        raise NotImplementedError(
+            "param_num doesn't support DynamicModule. Please use param_num_from_forward instead."
+        )
     return (
         sum(
             p.numel() if not trainable_only or p.requires_grad else 0

modelopt/torch/utils/network.py

@@ -46,9 +46,14 @@ def param_num_megatron(
     Returns:
         The number of parameters in the model (reduced across TP and PP ranks).
     """
+    from modelopt.torch.opt.dynamic import DynamicModule
+
     if from_forward:
         assert args is not None, "args must be provided if from_forward is True"
         params = int(param_num_from_forward(model, args, unit=1.0))
+    elif isinstance(model, DynamicModule):
+        # NOTE: model.parameters() doesnt consider active_slice so we dont get sorted or trimmed parameters!
+        raise NotImplementedError("DynamicModule input is not supported without from_forward.")
     else:
         params = sum(p.numel() for p in model.parameters())
     reduced_params = torch.Tensor([params]).to(device=next(model.parameters()).device)

modelopt/torch/utils/plugins/megatron_model.py

@@ -26,6 +26,7 @@ def prune_minitron(model, export_config, config, channel_divisor=64):
                     channel_divisor=channel_divisor,
                     mamba_head_dim_divisor=4,
                     num_moe_experts_divisor=1,
+                    num_layers_divisor=1,
                 ),
             )
         ],

tests/_test_utils/torch/nas_prune/minitron_common.py

@@ -83,7 +83,15 @@ def _test_gpt_search_space(
         normalization=normalization,
     ).cuda()
 
-    mtn.convert(model, [("mcore_minitron", get_mcore_minitron_config(channel_divisor))])
+    mtn.convert(
+        model,
+        [
+            (
+                "mcore_minitron",
+                get_mcore_minitron_config(channel_divisor=channel_divisor, num_layers_divisor=1),
+            )
+        ],
+    )
 
     assert isinstance(model, _DynamicMCoreLanguageModel)
     for m in model.modules():
@@ -255,7 +263,14 @@ def _test_gpt_moe_search_space(rank, size):
 
     mtn.convert(
         model,
-        [("mcore_minitron", get_mcore_minitron_config(channel_divisor, num_moe_experts_divisor=1))],
+        [
+            (
+                "mcore_minitron",
+                get_mcore_minitron_config(
+                    channel_divisor=channel_divisor, num_moe_experts_divisor=1, num_layers_divisor=1
+                ),
+            )
+        ],
     )
 
     moe = model.decoder.layers[0].mlp

tests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py

@@ -77,7 +77,16 @@ def _test_mamba_search_space(rank, size):
 
     mtn.convert(
         model,
-        [("mcore_minitron", get_mcore_minitron_config(channel_divisor, mamba_head_dim_divisor))],
+        [
+            (
+                "mcore_minitron",
+                get_mcore_minitron_config(
+                    channel_divisor=channel_divisor,
+                    mamba_head_dim_divisor=mamba_head_dim_divisor,
+                    num_layers_divisor=1,
+                ),
+            )
+        ],
     )
 
     assert isinstance(model, _DynamicMCoreLanguageModel)