updated parallel state for experts

Author: kinjalpatel27

modelopt/torch/quantization/plugins/megatron.py

@@ -22,6 +22,7 @@
 import megatron.core.parallel_state as mcore_parallel
 import megatron.core.tensor_parallel.layers as megatron_parallel
 import megatron.core.transformer.mlp as megatron_mlp
+import megatron.core.transformer.moe.experts as megatron_moe
 import torch
 import transformer_engine.pytorch.module.grouped_linear as te_grouped_linear
 from megatron.core.extensions import transformer_engine as megatron_te
@@ -38,7 +39,7 @@
 from modelopt.torch.utils.distributed import ParallelState
 
 from ..nn import QuantModuleRegistry, TensorQuantizer
-from ..nn.modules.quant_linear import RealQuantLinear, _QuantLinear
+from ..nn.modules.quant_linear import RealQuantLinear
 from ..qtensor import QTensorWrapper
 from .custom import CUSTOM_MODEL_PLUGINS, CUSTOM_POST_CALIBRATION_PLUGINS, _ParallelLinear
 
@@ -518,29 +519,18 @@ def forward(self, input, *args, **kwargs):
 
 
 # Register the public te.pytorch.GroupedLinear class
-@QuantModuleRegistry.register({te_grouped_linear.GroupedLinear: "te_GroupedLinear_public"})
+@QuantModuleRegistry.register({te_grouped_linear.GroupedLinear: "te_GroupedLinear"})
 class _QuantTEGroupedLinear(_MegatronParallelLinear):
     def _setup(self):
-        if not hasattr(self, "parallel_state") or self.parallel_state is None:
-            data_parallel_group = None
-            try:
-                data_parallel_group = get_data_parallel_group(with_context_parallel=True)
-            except AssertionError:
-                data_parallel_group = get_data_parallel_group()
-
-            self.parallel_state = ParallelState(
-                data_parallel_group,
-                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
-                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
-            )
-            self.input_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_input)
-            self.weight_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_weight)
-            self.output_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_output)
-            self.output_quantizer.disable()
-
+        # GroupedMLP stores the weights as weight0, weight1, etc. To run setup in order to
+        # initialize the quantizer states, self.weight is used to extract shape, dtype etc. Assigning
+        # self.weight0 to self.weight to run the quantizer states initialization.
+        self.weight = self.weight0
         # Memorize the original weight.dtype for modelopt_post_restore given that
         # the dtype can change later.
-        self.original_weight_dtype = None if self.weight0 is None else self.weight0.dtype
+        super()._setup()
+        # Revert the weight to None after setup.
+        self.weight = None
 
     @property
     def functionals_to_replace(self):
@@ -577,7 +567,7 @@ def modelopt_post_restore(self, prefix: str = ""):
         # self.weight0 to self.weight to run the quantizer states initialization.
         self.weight = self.weight0
         super().modelopt_post_restore(prefix=prefix)
-        # Revert the weight to None after post_restore to avoid the weight being None during forward pass.
+        # Revert the weight to None after post_restore.
         self.weight = None
 
     def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
@@ -611,3 +601,41 @@ class _QuantTEGroupedColumnParallelLinear(_QuantTEGroupedLinear, _MegatronColumn
 )
 class _QuantTEGroupedRowParallelLinear(_QuantTEGroupedLinear, _MegatronRowParallelLinear):
     _is_row_parallel = True
+
+
+# Register the public megatron_moe.TEGroupedMLP class
+@QuantModuleRegistry.register({megatron_moe.TEGroupedMLP: "megatron_moe_TEGroupedMLP"})
+class _QuantTEGroupedMLP(_MegatronMLP):
+    def _setup(self):
+        if not hasattr(self, "parallel_state") or self.parallel_state is None:
+            data_parallel_group = None
+            try:
+                data_parallel_group = get_data_parallel_group(with_context_parallel=True)
+            except AssertionError:
+                logger.warning(
+                    "Context parallel group is not initialized, using data parallel group"
+                )
+                data_parallel_group = get_data_parallel_group()
+
+            self.parallel_state = ParallelState(
+                data_parallel_group,
+                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
+                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
+            )
+
+
+# Register the public megatron_moe.SequentialMLP class
+@QuantModuleRegistry.register({megatron_moe.SequentialMLP: "megatron_moe_SequentialMLP"})
+class _QuantSequentialMLP(_MegatronMLP):
+    def _setup(self):
+        if not hasattr(self, "parallel_state") or self.parallel_state is None:
+            try:
+                data_parallel_group = mcore_parallel.get_expert_data_parallel_group()
+            except AssertionError:
+                data_parallel_group = None
+
+            self.parallel_state = ParallelState(
+                data_parallel_group,
+                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
+                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
+            )

tests/gpu/torch/quantization/plugins/test_megatron.py

@@ -516,17 +516,19 @@ def test_fp8_real_quantize():
         mtq.NVFP4_DEFAULT_CFG,
     ],
 )
-def test_moe_sharded_state_dict(need_8_gpus, tmp_path, config):
+@pytest.mark.parametrize("moe_grouped_gemm", [False, True])
+def test_moe_sharded_state_dict(tmp_path, config, moe_grouped_gemm):
     size = torch.cuda.device_count()
-    # TODO: Meta device doesn't work with TE
     # TODO: Add support for compress=True for TEGroupedMLP
+    if size < 4:
+        pytest.skip("Requires at least 4 GPUs for expert parallel test")
     moe_config = {
-        "tp_size": 2,
+        "tp_size": 1,
         "ep_size": 2,
         "etp_size": 2,
         "num_moe_experts": 4,
-        "moe_grouped_gemm": True,
-        "use_te": True,
+        "moe_grouped_gemm": moe_grouped_gemm,
+        "use_te": moe_grouped_gemm,
     }
     spawn_multiprocess_job(
         size=size,
@@ -627,10 +629,12 @@ def test_te_grouped_vs_sequential_quantize():
     )
 
 
-def _test_expert_model_parallel_amax_sync(ep_size, etp_size, moe_grouped_gemm, config, rank, size):
+def _test_expert_model_parallel_amax_sync(
+    tp_size, ep_size, etp_size, moe_grouped_gemm, config, rank, size
+):
     """Test expert parallel synchronization with different configurations."""
     initialize_for_megatron(
-        tensor_model_parallel_size=1,
+        tensor_model_parallel_size=tp_size,
         pipeline_model_parallel_size=1,
         expert_model_parallel_size=ep_size,
         expert_tensor_parallel_size=etp_size,
@@ -639,7 +643,7 @@ def _test_expert_model_parallel_amax_sync(ep_size, etp_size, moe_grouped_gemm, c
 
     # Create model with expert parallelism
     model = _gpt_model_provider(
-        tp_size=1,
+        tp_size=tp_size,
         ep_size=ep_size,
         etp_size=etp_size,
         hidden_size=256,
@@ -700,6 +704,7 @@ def test_expert_parallel_sync(ep_size, etp_size, moe_grouped_gemm):
         size=size,
         job=partial(
             _test_expert_model_parallel_amax_sync,
+            1,
             ep_size,
             etp_size,
             moe_grouped_gemm,