feat: add Muon optimizer with distributed sharding support

Add Muon optimizer implementation with Newton-Schulz orthogonalization
for distributed training:

- Muon optimizer (python/paddle/optimizer/muon.py):
  - Newton-Schulz iteration for orthogonal gradient updates
  - QKV split modes: per_head, qkv_sep, full
  - FFN gate_up split support
  - Multiple NS coefficient types: simple, quintic, polar_express, aol

- MuonShardingOptimizer:
  - Whole-tensor assignment for 2D parameters (Muon)
  - Element-wise sharding for non-2D parameters (AdamW)
  - Hybrid memory balancing across ranks

- Test coverage:
  - All 24 parameter combinations tested
  - 2-GPU sharding validation against single-GPU reference

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: xxyux

python/paddle/distributed/fleet/base/distributed_strategy.py

@@ -337,6 +337,8 @@ def __init__(self) -> None:
         ]
         self.sync_param_name: list[str] = ["embedding", "layer_norm", ".b_"]
 
+        self.use_muon_sharding: bool = False
+
         self.__lock_attr = True
         logger.info("distributed strategy initialized")
 

python/paddle/distributed/fleet/meta_optimizers/__init__.py

@@ -38,3 +38,4 @@
 from .recompute_optimizer import RecomputeOptimizer  # noqa: F401
 from .sharding_optimizer import ShardingOptimizer  # noqa: F401
 from .tensor_parallel_optimizer import TensorParallelOptimizer  # noqa: F401
+from .muon_sharding_optimizer import MuonShardingOptimizer  # noqa: F401

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py

@@ -1252,17 +1252,8 @@ def step(self):
         self._collect_comm_buffers()
         self._assign_slice_grad()
 
-        # Detect Muon by walking the wrapper chain; use name comparison to avoid
-        # a hard circular import.
-        core_opt = self._inner_opt
-        while hasattr(core_opt, '_inner_opt'):
-            core_opt = core_opt._inner_opt
-        is_muon = type(core_opt).__name__ == 'Muon'
-
         if not isinstance(self._parameter_list[0], dict):
             params_grads = []
-            # Build name→original-param map so Muon can recover full 2-D shape.
-            global_param_map = {p.name: p for p in self._parameter_list}
             for param in self._parameter_list:
                 if (
                     hasattr(param, "regularizer")
@@ -1280,25 +1271,8 @@ def step(self):
                 if hasattr(param, "main_grad") and param.main_grad is not None:
                     grad_var = param.main_grad
                 if grad_var is not None:
-                    if is_muon:
-                        from .muon_sharding_annotations import (
-                            annotate_muon_params,
-                        )
-
-                        original_p = global_param_map[param.name]
-                        if not annotate_muon_params(
-                            param, original_p, self._hcg, self.param2bucket
-                        ):
-                            continue
-
                     params_grads.append((param, grad_var))
 
-            if is_muon and params_grads:
-                from .muon_sharding_annotations import (
-                    sort_muon_params_grads,
-                )
-
-                sort_muon_params_grads(params_grads)
             if self._enable_timer:
                 self.timers("apply-optimize").start()
 

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/hybrid_parallel_optimizer.py

@@ -23,8 +23,8 @@
     DygraphShardingOptimizer,
     DygraphShardingOptimizerV2,
 )
-from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer_v3 import (
-    DygraphShardingOptimizerV3,
+from paddle.distributed.fleet.meta_optimizers.muon_sharding_optimizer import (
+    MuonShardingOptimizer,
 )
 from paddle.distributed.fleet.utils.hybrid_parallel_util import (
     obtain_optimizer_parameters_list,
@@ -287,9 +287,9 @@ def __init__(self, optimizer, hcg, strategy):
             split_param = strategy.hybrid_configs[
                 'sharding_configs'
             ].split_param
-            use_sharding_v3 = os.environ.get("FLAGS_sharding_v3", "0") == "1"
-            if use_sharding_v3 and split_param:
-                ShardingOptimizer = DygraphShardingOptimizerV3
+            use_muon_sharding = getattr(strategy, "use_muon_sharding", False)
+            if use_muon_sharding:
+                ShardingOptimizer = MuonShardingOptimizer
             elif split_param:
                 ShardingOptimizer = DygraphShardingOptimizerV2
             else:
@@ -340,7 +340,7 @@ def __init__(self, optimizer, hcg, strategy):
                     MixPrecisionOptimizer,
                     DygraphShardingOptimizer,
                     DygraphShardingOptimizerV2,
-                    DygraphShardingOptimizerV3,
+                    MuonShardingOptimizer,
                 ),
             )
 
@@ -637,7 +637,7 @@ def _hybrid_sync_grad(self, parameter_list):
                 (
                     DygraphShardingOptimizer,
                     DygraphShardingOptimizerV2,
-                    DygraphShardingOptimizerV3,
+                    MuonShardingOptimizer,
                 ),
             )
             self._inner_opt.reduce_gradients(parameter_list, self._hcg)