minor

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

Author: kevalmorabia97

examples/megatron_bridge/README.md

@@ -43,9 +43,8 @@ Example usage for manually pruning to a specific architecture using following de
 ```bash
 torchrun --nproc_per_node 2 /opt/Megatron-Bridge/3rdparty/Model-Optimizer/examples/megatron_bridge/prune_minitron.py \
     --hf_model_name_or_path Qwen/Qwen3-8B \
-    --prune_export_config '{"hidden_size": 3072, "ffn_hidden_size": 9216}' \
-    --hparams_to_skip num_attention_heads \
-    --output_hf_path /tmp/Qwen3-8B-Pruned-6B
+    --prune_export_config '{"hidden_size": 3584, "ffn_hidden_size": 9216}' \
+    --output_hf_path /tmp/Qwen3-8B-Pruned-6B-manual
 ```
 
 To see the full usage for advanced configurations, run:

examples/megatron_bridge/prune_minitron.py

@@ -91,7 +91,7 @@ def get_args() -> argparse.Namespace:
 
     # Pruning parameters
     parser.add_argument(
-        "--prune_intermediate_checkpoint",
+        "--prune_intermediate_ckpt",
         type=str,
         default=None,
         help=(
@@ -169,18 +169,16 @@ def get_args() -> argparse.Namespace:
     args = parser.parse_args()
 
     # Post-process arguments
-    if args.prune_intermediate_checkpoint is None:
+    if args.prune_intermediate_ckpt is None:
         if args.output_megatron_path:
-            args.prune_intermediate_checkpoint = (
+            args.prune_intermediate_ckpt = (
                 f"{args.output_megatron_path}/modelopt_pruning_scores.pth"
             )
         elif args.output_hf_path:
-            args.prune_intermediate_checkpoint = (
-                f"{args.output_hf_path}/modelopt_pruning_scores.pth"
-            )
+            args.prune_intermediate_ckpt = f"{args.output_hf_path}/modelopt_pruning_scores.pth"
         print_rank_0(
             "No checkpoint provided to cache intermediate pruning scores. "
-            f"Setting to: {args.prune_intermediate_checkpoint}"
+            f"Setting to: {args.prune_intermediate_ckpt}"
         )
 
     if args.prune_export_config:
@@ -247,7 +245,7 @@ def main(args: argparse.Namespace):
 
     pruning_config = {
         "forward_loop": forward_loop,
-        "checkpoint": args.prune_intermediate_checkpoint,
+        "checkpoint": args.prune_intermediate_ckpt,
     }
     if args.prune_target_params is not None:
         # Restrict search space to a smaller set of candidates

modelopt/torch/quantization/nn/modules/quant_module.py

@@ -17,8 +17,10 @@
 
 import contextlib
 import warnings
+from typing import Any
 
 import torch
+import torch.nn as nn
 
 from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls
 from modelopt.torch.utils.distributed import ParallelState
@@ -44,9 +46,11 @@ class QuantModule(DynamicModule):
 
     _parallel_state: ParallelState
 
-    def convert(self, *args, **kwargs):
+    @classmethod
+    @torch.no_grad()
+    def convert(cls, module: nn.Module, **setup_kwargs: Any) -> "QuantModule":
         """Convert the module to a dynamic module."""
-        module = super().convert(*args, **kwargs)
+        module = super().convert(module, **setup_kwargs)
 
         # setup parallel state now that the module is converted
         if module.parallel_state is None:
@@ -56,7 +60,7 @@ def convert(self, *args, **kwargs):
 
     @property
     def parallel_state(self) -> ParallelState | None:
-        """Return the parallel state of the dynamic module."""
+        """Return the parallel state of the quant module."""
         return getattr(self, "_parallel_state", None)
 
     @parallel_state.setter
@@ -70,7 +74,7 @@ def parallel_state(self, parallel_state: ParallelState):
     def _initialize_parallel_state(self):
         """Initialize the parallel state of the dynamic module.
 
-        This method is called only if the `DynamicModule` does not have a `parallel_state` attribute
+        This method is called only if the `QuantModule` does not have a `parallel_state` attribute
         after `_setup` is called.
         """
         if torch.distributed.is_initialized():