[Feature] support ep and enhance step latency for muon

- Remove distributed_mesh parameter; extract device mesh and process
group
  directly from DTensor metadata to support heterogeneous meshes (ViT 1D
+ LM 2D).
- Pre-compute adjust_lr ratios in __init__ based on global (unsharded)
shape,
  avoiding incorrect shape references inside async tasks after
communication.
- Add MoE expert-parallel support: per-expert Newton-Schulz
orthogonalization,
  requires n_experts % ep_size == 0 on the EP dimension. On the FSDP
dimension,
  skip communication when n_experts % fsdp_size == 0 (each rank holds
complete
  experts), use sub-group all-gather when fsdp_size % n_experts == 0,
otherwise
  fall back to batched all-to-all.
- Add AGRS (All-Gather + Reduce-Scatter) path for remainder batches to
avoid
  zero-padding overhead, with even-sharding guard to prevent deadlock.
- Refactor shared utilities
(group_tensors_by_device_mesh_and_placements,
  cal_total_norm) from grad_norm.py to dtensor.py.
- Remove `# type: ignore` from file head, fix lint, add full type
annotations.

Author: nil0x9

xtuner/v1/config/optim.py

@@ -185,14 +185,9 @@ def build(self, model):
                     f"Muon params: {num_muon_regular / 1e6:.2f}M, AdamW params: {num_adamw / 1e6:.2f}M (counts by numel)"
                 )
             logger.info(f"Untrainable parameters names: {untrainable_names}")
-            logger.info(
-                f"using Muon optimizer distributed_mesh_size: {model.fsdp_mesh.size()}, "
-                f"distributed_mesh: {model.fsdp_mesh}"
-            )
 
         optimizer = Muon(
             param_groups,
-            distributed_mesh=model.language_model.fsdp_mesh,  # TODO: 暂不支持 EP>1; maybe rm device_mesh dependency?
             lr=self.lr,
             mu=self.momentum,
             betas=self.betas,