Move parallel_state init and warnings to Quant DynamicModule + MBridge pruning doc update (#854)

## What does this PR do?

**Type of change:** Minor improvement <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

Only quantization DynamicModules use the parallel_state attribute so for
all other model opt methods, we see a parallel state not initialized
warning which could be confusing hence moving it to QuantModule class
instead

Minor update to MBridge pruning docs

## Testing
<!-- Mention how have you tested your change if applicable. -->

N/A

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Distributed parallel state support is now available in quantization
workflows for multi-GPU training.

* **Bug Fixes**
* Improved resource cleanup in distributed training to ensure proper
environment finalization.

* **Documentation**
* Updated example paths and added new manual pruning configuration
examples.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: kevalmorabia97

CHANGELOG.rst

@@ -226,7 +226,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 - Add support for UNet ONNX quantization.
 - Enable ``concat_elimination`` pass by default to improve the performance of quantized ONNX models.
 - Enable Redundant Cast elimination pass by default in :meth:`moq.quantize <modelopt.onnx.quantization.quantize>`.
-- Add new attribute ``parallel_state`` to :class:`DynamicModule <modelopt.torch.opt.dynamic.DynamicModule>` to support distributed parallelism such as data parallel and tensor parallel.
+- Add new attribute ``parallel_state`` to :class:`QuantModule <modelopt.torch.quantization.nn.modules.quant_module.QuantModule>` to support distributed parallelism such as data parallel and tensor parallel.
 - Add MXFP8, NVFP4 quantized ONNX export support.
 - Add new example for torch quantization to ONNX for MXFP8, NVFP4 precision.
 

examples/megatron_bridge/README.md

@@ -18,7 +18,7 @@ This directory contains examples of using Model Optimizer with [NeMo Megatron-Br
 
 Running these examples requires many additional dependencies to be installed (e.g., Megatron-Bridge, Megatron-core, etc.), hence we strongly recommend directly using the NeMo container (e.g., `nvcr.io/nvidia/nemo:26.02`) which has all the dependencies installed.
 
-To get the latest ModelOpt features and examples, you can mount your latest ModelOpt cloned repository to the container at `/opt/Model-Optimizer` or pull the latest changes once inside the docker container (`cd /opt/Model-Optimizer && git checkout main && git pull`).
+To get the latest ModelOpt features and examples, you can mount your latest ModelOpt cloned repository to the container at `/opt/Megatron-Bridge/3rdparty/Model-Optimizer` or pull the latest changes once inside the docker container (`cd /opt/Megatron-Bridge/3rdparty/Model-Optimizer && git checkout main && git pull`).
 
 ## Pruning
 
@@ -30,17 +30,27 @@ Example usage to prune Qwen3-8B to 6B on 2-GPUs (Pipeline Parallelism = 2) while
     top-10 candidates are evaluated for MMLU score (5% sampled data) to select the best model.
 
 ```bash
-torchrun --nproc_per_node 2 /opt/Model-Optimizer/examples/megatron_bridge/prune_minitron.py \
+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_target_params 6e9 \
     --hparams_to_skip num_attention_heads \
     --output_hf_path /tmp/Qwen3-8B-Pruned-6B
 ```
 
+Example usage for manually pruning to a specific architecture using following defaults:
+    1024 samples from [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) for calibration.
+
+```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": 3584, "ffn_hidden_size": 9216}' \
+    --output_hf_path /tmp/Qwen3-8B-Pruned-6B-manual
+```
+
 To see the full usage for advanced configurations, run:
 
 ```bash
-python /opt/Model-Optimizer/examples/megatron_bridge/prune_minitron.py --help
+python /opt/Megatron-Bridge/3rdparty/Model-Optimizer/examples/megatron_bridge/prune_minitron.py --help
 ```
 
 > [!TIP]

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
@@ -377,8 +375,8 @@ def score_func_mmlu(m):
 
 
 if __name__ == "__main__":
-    args = get_args()
     dist.setup()
+    args = get_args()
     try:
         main(args)
     finally:

modelopt/torch/opt/dynamic.py

@@ -31,7 +31,6 @@
 from torch.nn.parameter import Parameter
 
 from modelopt.torch.utils import get_unwrapped_name, is_channels_last, unwrap_model
-from modelopt.torch.utils.distributed import ParallelState
 from modelopt.torch.utils.network import bind_forward_method
 
 from .config import ModeloptBaseRule, RulesDict
@@ -359,14 +358,10 @@ class DynamicModule(nn.Module):
     should ensure only to expose ``hparams`` in the outermost class and handle other ``hparams``
     internally including ``hparams`` of child modules that are exposed on their own usually
     (e.g. block module implementations containing DynamicLinear).
-
-    In addition, the class also provides ``parallel_state`` attribute that can be used to access
-    the parallel state of the module.
     """
 
     # this is needed to store the special attributes for dynamic modules
     _dm_attribute_manager: _DMAttributeManager
-    _parallel_state: ParallelState
 
     def __init__(self, *args, **kwargs):
         """Initializing a dynamic module is not allowed!"""
@@ -657,10 +652,6 @@ def bind_forward_method_if_needed(self):
         # setup new hparams and dynamic attributes
         module._setup(**setup_kwargs)
 
-        # setup parallel state now that the module is converted
-        if module.parallel_state is None:
-            module._initialize_parallel_state()
-
         return module
 
     def _setup(self, **setup_kwargs: Any):
@@ -867,36 +858,6 @@ def original_cls(self) -> type[nn.Module]:
         """
         return self._get_dm_attribute_manager().og_cls
 
-    @property
-    def parallel_state(self) -> ParallelState | None:
-        """Return the parallel state of the dynamic module."""
-        return getattr(self, "_parallel_state", None)
-
-    @parallel_state.setter
-    def parallel_state(self, parallel_state: ParallelState):
-        """Set the parallel state of the dynamic module."""
-        assert isinstance(parallel_state, ParallelState), (
-            "parallel_state must be a ParallelState object!"
-        )
-        self._parallel_state = parallel_state
-
-    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
-        after `_setup` is called.
-        """
-        if torch.distributed.is_initialized():
-            warnings.warn(
-                f"Distributed training is initialized but no parallel_state is set for {type(self)}. "
-                "Using default parallel_state which has data_parallel_group set to the default process group and "
-                "tensor_parallel_group is unspecified. "
-                "If you are using tensor parallelism for this module, you should set the parallel_state "
-                "in its `_setup` method."
-            )
-
-        self.parallel_state = ParallelState(data_parallel_group=None)
-
     def get_original_cls_by_level(self, level: int = -1) -> type[nn.Module]:
         """Return the original class of the dynamic module.
 

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

@@ -17,10 +17,13 @@
 
 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
 
 from ...tensor_quant import QUANT_DESC_8BIT_PER_TENSOR
 from ...utils import is_torch_export_mode
@@ -35,7 +38,55 @@
 
 
 class QuantModule(DynamicModule):
-    """A base class for quantized modules."""
+    """A base class for quantized modules.
+
+    In addition, the class also provides ``parallel_state`` attribute that can be used to access
+    the parallel state of the module.
+    """
+
+    _parallel_state: ParallelState
+
+    @classmethod
+    @torch.no_grad()
+    def convert(cls, module: nn.Module, **setup_kwargs: Any) -> "QuantModule":
+        """Convert the module to a dynamic module."""
+        module = super().convert(module, **setup_kwargs)
+
+        # setup parallel state now that the module is converted
+        if module.parallel_state is None:
+            module._initialize_parallel_state()
+
+        return module
+
+    @property
+    def parallel_state(self) -> ParallelState | None:
+        """Return the parallel state of the quant module."""
+        return getattr(self, "_parallel_state", None)
+
+    @parallel_state.setter
+    def parallel_state(self, parallel_state: ParallelState):
+        """Set the parallel state of the dynamic module."""
+        assert isinstance(parallel_state, ParallelState), (
+            "parallel_state must be a ParallelState object!"
+        )
+        self._parallel_state = parallel_state
+
+    def _initialize_parallel_state(self):
+        """Initialize the parallel state of the dynamic module.
+
+        This method is called only if the `QuantModule` does not have a `parallel_state` attribute
+        after `_setup` is called.
+        """
+        if torch.distributed.is_initialized():
+            warnings.warn(
+                f"Distributed training is initialized but no parallel_state is set for {type(self)}. "
+                "Using default parallel_state which has data_parallel_group set to the default process group and "
+                "tensor_parallel_group is unspecified. "
+                "If you are using tensor parallelism for this module, you should set the parallel_state "
+                "in its `_setup` method."
+            )
+
+        self.parallel_state = ParallelState(data_parallel_group=None)
 
     def modelopt_post_restore(self, prefix: str = ""):
         """Post-restore to correctly configure the TensorQuantizer states.