Enable optional batched Muon NS iteration

python/paddle/optimizer/muon.py

@@ -19,6 +19,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -40,6 +41,9 @@
 # Debug logging for Muon optimizer
 _logger = logging.getLogger(__name__)
 MUON_DEBUG = os.environ.get("MUON_DEBUG", "0") == "1"
+g_shard_bypass_dygraph_optimizer = int(
+    os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0)
+)
 
 __all__ = []
 
@@ -150,6 +154,15 @@ def _default_should_use_muon(name, shape, exclude_patterns):
     return True
 
 
+def _to_hashable(value):
+    """Recursively convert lists/tuples/dicts into hashable tuples."""
+    if isinstance(value, (list, tuple)):
+        return tuple(_to_hashable(v) for v in value)
+    if isinstance(value, dict):
+        return tuple((k, _to_hashable(v)) for k, v in sorted(value.items()))
+    return value
+
+
 class Muon(Optimizer):
     r"""
     Muon optimizer for MuonShardingOptimizer (Sharding Stage1 V3) usage.
@@ -519,52 +532,64 @@ def _adamw_update(
             False,  # amsgrad
         )
 
-    def _muon_update(
+    def _muon_update_group(
         self,
-        param,
-        grad,
+        group_params_grads,
         lr,
-        momentum_buffer,
         momentum_beta,
         ns_steps,
         nesterov,
         epsilon,
         weight_decay,
         version,
     ):
-        """In-place Muon update for a 2D parameter tensor.
+        """Batched Muon update for a group of parameters with identical shape and split_concat_func."""
+        # Because shape is identical, use the first param's properties
+        param0, _ = group_params_grads[0]
+        param_shape = getattr(param0, "original_shape", param0.shape)
+        is_3d = len(param_shape) == 3
 
-        Applies Newton-Schulz orthogonalisation to the 2D weight matrix and
-        updates the parameter in-place. MuonShardingOptimizer assigns whole
-        2D tensors to ranks, so no sharding gather or TP communication is needed.
-        """
-        param_shape = getattr(param, "original_shape", param.shape)
-        param_info = self._muon_param_info_map.get(param.name)
+        param_info = self._muon_param_info_map.get(param0.name)
+        split_concat_func = param_info.split_concat_func if param_info else None
+
+        matrix_2d_list = []
+        find_master_list = []
 
         with paddle.no_grad():
-            grad_f32 = (
-                grad.astype(momentum_buffer.dtype)
-                if grad.dtype != momentum_buffer.dtype
-                else grad
-            )
+            # --- Pass 1: Sequential momentum update & preparation ---
+            for param, grad in group_params_grads:
+                momentum_buffer = self._get_accumulator(
+                    self._moment_acc_str, param
+                )
+                grad_f32 = (
+                    grad.astype(momentum_buffer.dtype)
+                    if grad.dtype != momentum_buffer.dtype
+                    else grad
+                )
 
-            # Step 1: Momentum update
-            new_momentum = paddle.lerp(
-                momentum_buffer, grad_f32, 1.0 - momentum_beta
-            )
-            paddle.assign(new_momentum, momentum_buffer)
-            update_buffer = (
-                paddle.lerp(grad_f32, momentum_buffer, momentum_beta)
-                if nesterov
-                else momentum_buffer
-            )
+                new_momentum = paddle.lerp(
+                    momentum_buffer, grad_f32, 1.0 - momentum_beta
+                )
+                paddle.assign(new_momentum, momentum_buffer)
+                update_buffer = (
+                    paddle.lerp(grad_f32, momentum_buffer, momentum_beta)
+                    if nesterov
+                    else momentum_buffer
+                )
+
+                matrix_2d_global = update_buffer.reshape(param_shape)
+                matrix_2d_list.append(matrix_2d_global)
+                find_master_list.append(param.name in self._master_weights)
 
-            # Step 2: Reshape update buffer to 2D matrix.
-            # MuonShardingOptimizer assigns whole 2D tensors to ranks, so params
-            # are already 2D/3D (no sharding gather needed).
-            matrix_2d_global = update_buffer.reshape(param_shape)
+            # --- Pass 2: Batched Newton-Schulz iteration ---
+            if len(matrix_2d_list) > 1:
+                if is_3d:
+                    batched_matrix = paddle.concat(matrix_2d_list, axis=0)
+                else:
+                    batched_matrix = paddle.stack(matrix_2d_list, axis=0)
+            else:
+                batched_matrix = matrix_2d_list[0]
 
-            # Shared NS + scaling closure (captures ns_steps, epsilon, version, ns_coeffs)
             def ortho_fn(m):
                 ns_out = Muon._zeropower_via_newtonschulz5(
                     m,
@@ -578,56 +603,74 @@ def ortho_fn(m):
                 )
                 return scaled
 
-            # Step 3: Newton-Schulz orthogonalisation
-            # Use split_concat_func from param_info if provided, otherwise default to whole matrix
-            if (
-                param_info is not None
-                and param_info.split_concat_func is not None
-            ):
-                # Use slice function defined in model configuration
-                orthogonal_update = param_info.split_concat_func(
-                    matrix_2d_global, ortho_fn
+            if split_concat_func is not None:
+                batched_orthogonal_update = split_concat_func(
+                    batched_matrix, ortho_fn
                 )
                 if MUON_DEBUG:
                     _global_rank = paddle.distributed.get_rank()
                     if _global_rank == 0:
-                        _sf = param_info.split_concat_func
+                        _sf = split_concat_func
+                        # split_concat_func may be a plain function or a
+                        # functools.partial; unwrap for readable logging.
+                        _fn = getattr(_sf, "func", _sf)
                         _logger.info(
-                            f"[Muon] Using split_concat_func: param={param.name}, "
-                            f"split_concat_func={_sf.func.__name__}, "
-                            f"args={_sf.args}, kwargs={_sf.keywords}"
+                            f"[Muon] Using split_concat_func: "
+                            f"params={[p.name for p, _ in group_params_grads]}, "
+                            f"split_concat_func={_fn.__name__}, "
+                            f"args={getattr(_sf, 'args', ())}, "
+                            f"kwargs={getattr(_sf, 'keywords', None)}"
                         )
             else:
-                # Default: whole matrix orthogonalisation
-                orthogonal_update = ortho_fn(matrix_2d_global)
-
-            find_master = param.name in self._master_weights
-            master_weight = (
-                self._master_weights[param.name] if find_master else None
-            )
+                batched_orthogonal_update = ortho_fn(batched_matrix)
+
+            # Unbatch results
+            if len(matrix_2d_list) > 1:
+                if is_3d:
+                    orthogonal_updates = paddle.split(
+                        batched_orthogonal_update,
+                        num_or_sections=len(matrix_2d_list),
+                        axis=0,
+                    )
+                else:
+                    orthogonal_updates = paddle.unbind(
+                        batched_orthogonal_update, axis=0
+                    )
+            else:
+                orthogonal_updates = [batched_orthogonal_update]
+
+            # --- Pass 3: Sequential apply weight decay & step ---
+            for i, (param, _) in enumerate(group_params_grads):
+                orthogonal_update = orthogonal_updates[i]
+                find_master = find_master_list[i]
+                master_weight = (
+                    self._master_weights[param.name] if find_master else None
+                )
 
-            lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param)
-            effective_lr = lr * lr_ratio
+                lr_ratio = (
+                    1.0 if self._lr_ratio is None else self._lr_ratio(param)
+                )
+                effective_lr = lr * lr_ratio
+
+                with_decay = True
+                if (
+                    self._apply_decay_param_fun is not None
+                    and not self._apply_decay_param_fun(param.name)
+                ):
+                    with_decay = False
+                if with_decay and weight_decay > 0:
+                    if find_master:
+                        master_weight.scale_(1.0 - effective_lr * weight_decay)
+                    else:
+                        param.scale_(1.0 - effective_lr * weight_decay)
+
+                final_step = orthogonal_update * effective_lr
 
-            with_decay = True
-            if (
-                self._apply_decay_param_fun is not None
-                and not self._apply_decay_param_fun(param.name)
-            ):
-                with_decay = False
-            if with_decay and weight_decay > 0:
                 if find_master:
-                    master_weight.scale_(1.0 - effective_lr * weight_decay)
+                    master_weight.subtract_(final_step)
+                    paddle.assign(master_weight.astype(param.dtype), param)
                 else:
-                    param.scale_(1.0 - effective_lr * weight_decay)
-
-            final_step = orthogonal_update * effective_lr
-
-            if find_master:
-                master_weight.subtract_(final_step)
-                paddle.assign(master_weight.astype(param.dtype), param)
-            else:
-                param.subtract_(final_step.astype(param.dtype))
+                    param.subtract_(final_step.astype(param.dtype))
 
     # ------------------------------------------------------------------
     # Core optimization step
@@ -639,6 +682,10 @@ def _apply_optimize(self, loss, startup_program, params_grads):
                 "Muon optimizer only supports dygraph mode."
             )
 
+        # Same bypass as the base Optimizer._apply_optimize: skip all parameter updates.
+        if g_shard_bypass_dygraph_optimizer:
+            return
+
         if self._grad_clip is not None:
             params_grads = self._grad_clip(params_grads)
 
@@ -658,26 +705,63 @@ def _apply_optimize(self, loss, startup_program, params_grads):
                 continue
 
             param_info = self._muon_param_info_map.get(param.name)
-            assert param_info is not None, (
-                f"muon_param_info_map does not have {param.name}"
-            )
-            use_muon = param_info.use_muon
+            if param_info is not None:
+                use_muon = param_info.use_muon
+            else:
+                use_muon = _default_should_use_muon(
+                    param.name,
+                    getattr(param, "original_shape", param.shape),
+                    self._muon_exclude_patterns,
+                )
+                _logger.warning(
+                    f"muon_param_info_map does not have {param.name}, "
+                    f"falling back to default rule: use_muon={use_muon}"
+                )
 
             self._ensure_accumulators(param, use_muon, group)
             if use_muon:
                 muon_params.append((param, grad))
             else:
                 adamw_params.append((param, grad))
 
-        # --- Pass 1: Muon updates (large temporary tensors) ---
+        # --- Pass 2: Muon updates ---
         lr_tensor = paddle.to_tensor(lr, dtype=paddle.float32)
         lr_tensor_f64 = paddle.to_tensor(lr, dtype=paddle.float64)
+
+        # Group parameters by split_concat_func and shape, then update
+        # each group in one batched call.
+        muon_groups = defaultdict(list)
         for param, grad in muon_params:
-            self._muon_update(
-                param,
-                grad,
+            assert len(param.shape) == 2 or len(param.shape) == 3, (
+                "Muon only supports 2D or 3D parameters."
+            )
+            param_shape = getattr(param, "original_shape", param.shape)
+            param_info = self._muon_param_info_map.get(param.name)
+
+            func_key = None
+            if param_info and param_info.split_concat_func is not None:
+                func = param_info.split_concat_func
+                if hasattr(func, "func"):
+                    # functools.partial: group by the underlying function
+                    # object plus all bound args/kwargs.
+                    func_key = (
+                        func.func,
+                        tuple(_to_hashable(v) for v in func.args),
+                        tuple(
+                            (k, _to_hashable(v))
+                            for k, v in sorted((func.keywords or {}).items())
+                        ),
+                    )
+                else:
+                    func_key = func
+
+            key = (func_key, tuple(param_shape))
+            muon_groups[key].append((param, grad))
+
+        for key, group_params in muon_groups.items():
+            self._muon_update_group(
+                group_params,
                 lr_tensor,
-                self._get_accumulator(self._moment_acc_str, param),
                 group.get("momentum", 0.95),
                 group.get("ns_steps", 5),
                 group.get("nesterov", True),
@@ -686,7 +770,7 @@ def _apply_optimize(self, loss, startup_program, params_grads):
                 version=group.get("muon_version", 3),
             )
 
-        # --- Pass 2: AdamW updates ---
+        # --- Pass 3: AdamW updates ---
         for param, grad in adamw_params:
             self._adamw_update(
                 param,