From a31ee72cc925112f7733bc81eee7d2b2b3e6cd34 Mon Sep 17 00:00:00 2001 From: xxyux <1650459510@qq.com> Date: Wed, 15 Jul 2026 16:44:27 +0800 Subject: [PATCH] update group_NS --- python/paddle/optimizer/muon.py | 236 ++++++++++++------ .../hybrid_parallel_sharding_muon_model.py | 68 +++-- .../fleet/test_parallel_dygraph_muon.py | 4 +- 3 files changed, 210 insertions(+), 98 deletions(-) diff --git a/python/paddle/optimizer/muon.py b/python/paddle/optimizer/muon.py index 1251abac5f506..d707647373478 100644 --- a/python/paddle/optimizer/muon.py +++ b/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,12 +532,10 @@ 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, @@ -532,39 +543,53 @@ def _muon_update( 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,53 +603,69 @@ 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) + 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 + ) - find_master = param.name in self._master_weights - master_weight = ( - self._master_weights[param.name] if find_master else None - ) + 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 - lr * weight_decay) + else: + param.scale_(1.0 - lr * weight_decay) + + final_step = orthogonal_update * 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 - lr * weight_decay) + master_weight.subtract_(final_step) + paddle.assign(master_weight.astype(param.dtype), param) else: - param.scale_(1.0 - lr * weight_decay) - - final_step = orthogonal_update * 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 @@ -636,6 +677,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) @@ -655,10 +700,18 @@ 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: @@ -666,15 +719,44 @@ def _apply_optimize(self, loss, startup_program, params_grads): 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), @@ -683,7 +765,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, diff --git a/test/collective/fleet/hybrid_parallel_sharding_muon_model.py b/test/collective/fleet/hybrid_parallel_sharding_muon_model.py index ca41f773f7bd6..1d8da60195a0d 100644 --- a/test/collective/fleet/hybrid_parallel_sharding_muon_model.py +++ b/test/collective/fleet/hybrid_parallel_sharding_muon_model.py @@ -18,6 +18,11 @@ # Topology: sharding_degree=2, mp_degree=1 (2 ranks total) import os + +# Enable MUON_DEBUG to cover the debug logging branch in _muon_update_group. +# Must be set BEFORE `import paddle`: muon.py reads it at import time. +os.environ["MUON_DEBUG"] = "1" + import random import unittest from dataclasses import dataclass @@ -36,9 +41,6 @@ _default_should_use_muon, ) -# Enable MUON_DEBUG to cover the debug logging branch (muon.py L532-539) -os.environ["MUON_DEBUG"] = "1" - # Test-controlled flags (set via need_envs from test_parallel_dygraph_muon.py) g_enable_fuse_optimizer_states = int( os.environ.get("ENABLE_FUSE_OPTIMIZER_STATES", "0") @@ -95,6 +97,15 @@ def _ffn_split(matrix_2d, ortho_fn, intermediate_size=None): return paddle.concat([ortho_fn(gate), ortho_fn(up)], axis=1) +def _whole_matrix(matrix_2d, ortho_fn): + """Plain-function slice strategy: orthogonalise the whole matrix. + + Registered as a plain function (NOT functools.partial) to cover the + non-partial func_key branch in Muon._apply_optimize grouping. + """ + return ortho_fn(matrix_2d) + + @dataclass class TestConfig: """Test model config.""" @@ -189,6 +200,17 @@ def __init__( default_initializer=paddle.nn.initializer.Assign(np_batched_proj), ) + # Second 3D param with the SAME shape: the two 3D params fall into + # one group, covering the multi-param 3D concat/split path in + # _muon_update_group. + self.batched_proj2 = paddle.create_parameter( + shape=list(batched_proj_shape), + dtype='float32', + default_initializer=paddle.nn.initializer.Assign( + np_batched_proj * 0.5 + ), + ) + def forward(self, x): # Embedding h = self.embed_tokens(x) # [batch, seq, hidden] @@ -213,7 +235,8 @@ def forward(self, x): # Batched projection (3D param: triggers batched NS + transpose path) bp_scale = paddle.einsum('bsh,kmn->', out, self.batched_proj) - out = out + bp_scale * 1e-4 + bp_scale2 = paddle.einsum('bsh,kmn->', out, self.batched_proj2) + out = out + (bp_scale + bp_scale2) * 1e-4 # LM head logits = self.lm_head(out) # [batch, seq, vocab] @@ -294,6 +317,10 @@ def _build_split_concat_func_map(self, model): _ffn_split, intermediate_size=intermediate_size, ) + elif "down_proj" in name: + # Plain function (no partial): covers the `func_key = func` + # branch in Muon._apply_optimize group-key construction. + slice_map[param.name] = _whole_matrix return slice_map def build_optimizer( @@ -312,13 +339,14 @@ def build_optimizer( model: The model to optimize. ns_coeff: Newton-Schulz coefficient type. split_concat_func_map: Optional dict {param.name: split_concat_func}. - Covers muon.py L529 (split_concat_func call) and L535 (debug log). + Covers the split_concat_func call and MUON_DEBUG log in + _muon_update_group. ns_matmul_dtype: Optional explicit dtype for NS matmul. - Covers muon.py L283 (explicit ns_matmul_dtype branch). + Covers the explicit ns_matmul_dtype branch in Muon.__init__. multi_precision: If True, enable FP32 master weights. - Covers muon.py L560-564, L574-575, L582-583. + Covers the find_master branches in _muon_update_group. apply_decay_param_fun: Optional callable(param_name) -> bool. - Covers muon.py L443-446, L568-572. + Covers with_decay=False in _adamw_update / _muon_update_group. ns_coeffs: Optional custom NS coefficient list. """ muon_param_info_map = {} @@ -394,14 +422,16 @@ def _run_single_test( ns_coeff: Newton-Schulz coefficient type. color_params: Optional list of param name substrings to assign custom color group (covers muon_sharding_optimizer L388-394). - use_slice: If True, build split_concat_func_map for QKV/FFN params - (covers muon.py L529, L535). - explicit_dtype: If True, pass ns_matmul_dtype=paddle.float32 explicitly - (covers muon.py L283). + use_slice: If True, build split_concat_func_map for QKV/FFN/down + params (covers split_concat_func + MUON_DEBUG log in + _muon_update_group, and both partial / plain-function + func_key branches in _apply_optimize). + explicit_dtype: If True, pass ns_matmul_dtype=paddle.float32 + explicitly (covers the explicit dtype branch in Muon.__init__). multi_precision: If True, enable FP32 master weights - (covers muon.py L560-564, L574-575, L582-583). + (covers the find_master branches in _muon_update_group). apply_decay_param_fun: Optional callable(param_name) -> bool - (covers muon.py L443-446, L568-572). + (covers with_decay=False in _adamw_update / _muon_update_group). ns_coeffs: Optional custom NS coefficient list. """ # Allow env var to force multi_precision on @@ -499,11 +529,11 @@ def test_sharding_muon(self): Covers: - muon_sharding_optimizer.py L388-394: custom color from param.color dict - muon_sharding_optimizer.py L627-635, L665-667: fused gradient comm buffers - - muon.py L283: explicit ns_matmul_dtype=paddle.float32 - - muon.py L443-446: apply_decay_param_fun with_decay=False (AdamW path) - - muon.py L529: split_concat_func call - - muon.py L535: MUON_DEBUG logging (via MUON_DEBUG=1 env) - - muon.py L560-564, L574-575, L582-583: find_master=True (Muon path) + - muon.py: explicit ns_matmul_dtype=paddle.float32 + - muon.py: apply_decay_param_fun with_decay=False (AdamW + Muon paths) + - muon.py: split_concat_func call in _muon_update_group + - muon.py: MUON_DEBUG logging (via MUON_DEBUG=1 env) + - muon.py: find_master=True branches in _muon_update_group """ total = ( len(NS_COEFF_TYPES) + 4 diff --git a/test/collective/fleet/test_parallel_dygraph_muon.py b/test/collective/fleet/test_parallel_dygraph_muon.py index 13b6d990a8fc3..ba937c19657f8 100644 --- a/test/collective/fleet/test_parallel_dygraph_muon.py +++ b/test/collective/fleet/test_parallel_dygraph_muon.py @@ -76,8 +76,8 @@ def test_muon_sharding_release_grads_fused(self): def test_muon_sharding_multi_precision(self): """MuonSharding test with multi_precision=True. - Covers muon.py L575 (master_weight.scale_ with weight_decay), - L582-583 (master_weight.subtract_ + assign back to param). + Covers master_weight.scale_ with weight_decay and + master_weight.subtract_ + assign back to param in _muon_update_group. """ self.run_mnist_2accelerators( 'hybrid_parallel_sharding_muon_model.py',