Move parallel state init and warnings to Quant DynamicModule

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.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.
 

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

@@ -21,6 +21,7 @@
 import torch
 
 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 +36,53 @@
 
 
 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
+
+    def convert(self, *args, **kwargs):
+        """Convert the module to a dynamic module."""
+        module = super().convert(*args, **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 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 modelopt_post_restore(self, prefix: str = ""):
         """Post-restore to correctly configure the TensorQuantizer states.