Code cleanup

Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>

Author: kinjalpatel27

modelopt/torch/quantization/plugins/megatron.py

@@ -611,13 +611,9 @@ class _MegatronTEGroupedMLP(_MegatronMLP):
     def _setup(self):
         if not hasattr(self, "parallel_state") or self.parallel_state is None:
             self.parallel_state = ParallelState(
-                mcore_parallel.get_expert_data_parallel_group(check_initialized=False),
-                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(
-                    check_initialized=False
-                ),
-                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(
-                    check_initialized=False
-                ),
+                mcore_parallel.get_expert_data_parallel_group(),
+                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
+                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
             )
         # initialize parallel state for submodules linear_fc1 and linear_fc2
         self.linear_fc1.parallel_state = self.parallel_state
@@ -630,13 +626,9 @@ class _MegatronSequentialMLP(_MegatronMLP):
     def _setup(self):
         if not hasattr(self, "parallel_state") or self.parallel_state is None:
             self.parallel_state = ParallelState(
-                mcore_parallel.get_expert_data_parallel_group(check_initialized=False),
-                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(
-                    check_initialized=False
-                ),
-                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(
-                    check_initialized=False
-                ),
+                mcore_parallel.get_expert_data_parallel_group(),
+                tensor_parallel_group=mcore_parallel.get_expert_tensor_parallel_group(),
+                expert_model_parallel_group=mcore_parallel.get_expert_model_parallel_group(),
             )
 
         # Initialize parallel state for submodules local_experts.*.linear_fc1 and local_experts.*.linear_fc2

tests/_test_utils/torch_dist/plugins/megatron_common.py

@@ -517,16 +517,16 @@ def copy_weights_from_grouped_to_non_grouped(te_grouped_moe_model, sequential_mo
     weight_mapping = {}
     sequential_key_template = "decoder.layers.{}.mlp.experts.local_experts.{}.linear_fc{}.weight"
     for key, value in te_grouped_state.items():
-        if "experts.linear_fc" in key and "weight" in key:
+        if "experts.linear_fc" in key and any(param in key for param in ("weight", "bias")):
             # Extract expert index from grouped weight name
             # Format: decoder.layers.X.mlp.experts.linear_fcY.weightZ
             parts = key.split(".")
             layer_idx = parts[2]  # X
             fc_idx = parts[5]  # Y (linear_fc1 or linear_fc2)
-            weight_idx = parts[6]  # Z (weight0, weight1, etc.)
-
+            param_idx = parts[6]  # weight0 / bias0 / etc.
+            match = re.search(r"\d+", param_idx)
+            expert_idx = match.group(0) if match else "0"  # Z for expert index
             # Map to sequential format: decoder.layers.X.mlp.experts.local_experts.Y.linear_fcZ.weight
-            expert_idx = weight_idx.replace("weight", "")
             sequential_key = sequential_key_template.format(layer_idx, expert_idx, fc_idx[-1])
             weight_mapping[sequential_key] = value
         elif isinstance(value, torch.Tensor):