Update tp support in muon (#2385)

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

megatron/core/optimizer/muon.py

@@ -8,7 +8,6 @@
 import torch
 from torch.optim.optimizer import ParamsT
 
-from megatron.core import parallel_state
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_pg_size, log_single_rank
@@ -76,7 +75,7 @@ def scaled_orthogonalize_fn(
                 f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}',
             )
             size = [grad.size(-2), grad.size(-1)]
-            if partition_dim:
+            if partition_dim is not None:
                 size[partition_dim] *= get_pg_size(tp_group)
             orth_grad = newton_schulz_tp(
                 grad,
@@ -130,8 +129,7 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t
             tp_group = None
         partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None)
         if partition_dim == -1:
-            # llm-shower use different default value for partition_dim than TE.
-            # Because -1 is a valid index for ndarray, we decided to not overload it.
+            # emerging-optimizers use None instead of -1 to indicate no tensor parallel
             partition_dim = None
 
         if self.split_qkv and self.is_qkv_fn(p):  # type: ignore[misc]
@@ -201,8 +199,6 @@ def get_megatron_muon_optimizer(
     # before this function receive properly created collection
     if pg_collection is None:
         pg_collection = ProcessGroupCollection.use_mpu_process_groups()
-        pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True)
-        pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group()
 
     log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}')