From 594848e13af21aefd697886f73d43593238fbd5a Mon Sep 17 00:00:00 2001 From: xxyux <1650459510@qq.com> Date: Sat, 18 Apr 2026 23:43:36 +0800 Subject: [PATCH] feat: generalize MuonShardingOptimizer color handling and add V100 compatibility - muon_sharding_optimizer: replace hardcoded None/moe_expert color paths with generic _rank2params_2d_by_color iteration in step() and __init__ Step4; replace static _build_color_to_group_info(hcg) with dynamic _build_color_to_group_info_from_params(parameter_list, default_group) that scans param.color dicts at runtime; generalize reduce_gradients and _sharding_sync_parameters similarly; clean up comments (fix errors, translate Chinese to English, remove dead code and debug prints) - muon: remove built-in QKV/FFN split logic (QKVInfo, qkv_info, intermediate_size, muon_qkv_update_mode, muon_ffn_split) from the optimizer; callers now pass slice strategies via MuonParamInfo.slice_func, keeping model-specific split logic out of the optimizer core; add ns_matmul_dtype parameter to Muon.__init__ and _zeropower_via_newtonschulz5 with auto-detect (bfloat16 on Ampere+, float32 on V100 and older) to enable CI on V100 - optimizer: allow Muon class to skip incompatible base-class checks - test: update hybrid_parallel_sharding_muon_model and test_parallel_dygraph_muon to use current MuonParamInfo API (slice_func instead of deprecated qkv_info/ intermediate_size); remove GPU capability >= 8 skipIf guard so tests run on V100 Co-Authored-By: Claude Sonnet 4.6 --- .../muon_sharding_optimizer.py | 256 ++++++------- python/paddle/optimizer/muon.py | 357 +++--------------- python/paddle/optimizer/optimizer.py | 5 +- .../hybrid_parallel_sharding_muon_model.py | 265 ++++++++----- .../fleet/test_parallel_dygraph_muon.py | 11 +- 5 files changed, 352 insertions(+), 542 deletions(-) diff --git a/python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py b/python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py index 15b52bee738ef8..8113311b5a37c9 100644 --- a/python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py +++ b/python/paddle/distributed/fleet/meta_optimizers/muon_sharding_optimizer.py @@ -29,8 +29,8 @@ Parameters are grouped by their `color` attribute, which specifies the communication group to use: - color=None or -1: default sharding_group - - color='moe_expert': moe_sharding_group - - color=: hcg.get__parallel_group() (extensible design) + - color={'color': , 'group': }: custom group read directly + from the param, no code changes needed to add new color groups. """ import math @@ -118,6 +118,11 @@ def __init__(self, optimizer, hcg=None): self.comm_overlap = sharding_configs.comm_overlap self.comm_buffer_size_MB = sharding_configs.comm_buffer_size_MB self.use_reduce_avg = sharding_configs.use_reduce_avg + self.enable_fuse_optimizer_states = ( + sharding_configs.enable_fuse_optimizer_states + ) + if self.enable_fuse_optimizer_states: + self._inner_opt.use_fusion_storage() if self.use_reduce_avg and (not is_avg_reduce_op_supported()): self.use_reduce_avg = False @@ -135,8 +140,11 @@ def __init__(self, optimizer, hcg=None): self._parameter_list = list(optimizer._parameter_list) self._origin_parameter_list = list(optimizer._parameter_list) - # Build color -> group_info mapping - self._color_to_group_info = self._build_color_to_group_info(hcg) + # Build color -> group_info mapping dynamically from param.color attributes + sharding_group = hcg.get_sharding_parallel_group() + self._color_to_group_info = self._build_color_to_group_info_from_params( + self._parameter_list, sharding_group + ) # Extract MoE group info from color_to_group_info for backward compatibility moe_info = self._color_to_group_info.get('moe_expert', {}) @@ -144,28 +152,24 @@ def __init__(self, optimizer, hcg=None): self._moe_sharding_rank = moe_info.get('rank', 0) self._moe_sharding_group = moe_info.get('group', None) - # Get muon_param_info_map from Muon optimizer - # This map contains use_muon field for each parameter, determined by Trainer + # Get muon_param_info_map from the inner Muon optimizer. + # Each entry has use_muon=True/False, set by the Trainer before construction. self._muon_param_info_map = getattr( optimizer, '_muon_param_info_map', {} ) - # ---- Step 1: Separate params into categories by color ---- + # ---- Step 1: Separate params into 2D (Muon) and 1D (AdamW) by color ---- # Parameters are grouped by their `color` attribute: # - color=None or -1: default sharding_group (key: None) - # - color='moe_expert': moe_sharding_group (key: 'moe_expert') - # - color=: corresponding parallel group (key: ) + # - color={'color': , 'group': }: custom comm group # # For each color group: # - 2D (Muon) params: whole tensor, assigned to ranks via tensor-wise partition # - non-2D (AdamW) params: element-wise split via FusedCommBuffer - # - # This design is extensible: adding a new communication group only requires - # setting the `color` attribute on parameters, no code changes needed here. self._params_2d_by_color = defaultdict( list - ) # color -> list of 2D params - self._params_1d = [] # All non-2D params (single list, sharding_group only) + ) # color_key -> list of 2D params + self._params_1d = [] # all non-2D params self.clear_color = set() self._color_to_comm_buffer_list = {} for p in self._parameter_list: @@ -185,9 +189,6 @@ def __init__(self, optimizer, hcg=None): else: color_key = color_val - # Check if this color group supports 2D tensor-wise partition - group_info = self._color_to_group_info.get(color_key) - param_info = self._muon_param_info_map.get(p.name) assert param_info is not None, ( f"Parameter {p.name!r} (shape={list(p.shape)}) has no muon_param_info. " @@ -199,7 +200,7 @@ def __init__(self, optimizer, hcg=None): if use_muon: self._params_2d_by_color[color_key].append(p) else: - # Non-2D params always go to 1D element-wise split (sharding_group only) + # Non-2D params use element-wise split via FusedCommBuffer self._params_1d.append(p) # ---- Step 2: Partition 2D params for each color group ---- @@ -232,7 +233,7 @@ def __init__(self, optimizer, hcg=None): for p in params: self._param2rank_2d_by_color[color_key][p.name] = rank - # add sort 2d params + # Sort params within each color by owner rank for deterministic ordering for color_key, params_2d in self._params_2d_by_color.items(): params_2d.sort( key=lambda p: self._param2rank_2d_by_color[color_key][p.name] @@ -270,30 +271,22 @@ def __init__(self, optimizer, hcg=None): strategy.hybrid_configs['pp_configs'].release_gradients or sharding_configs.release_gradients ) + self._build_1d_comm_buffers() # ---- Step 4: Build the optimizer's parameter list ---- # The optimizer should see: - # - Non-MoE 2D params assigned to this rank (as whole tensors) - # - MoE expert 2D params assigned to this rank in moe_sharding_group + # - All 2D params assigned to this rank (all colors, as whole tensors) # - 1D slice_params for all non-2D params (element-wise shards) - local_2d_params = list( - self._rank2params_2d.get(self._sharding_rank, []) - ) - - if self._moe_sharding_world_size > 1: - local_2d_moe_params = list( - self._rank2params_2d_moe.get(self._moe_sharding_rank, []) - ) - else: - # moe_sharding_degree=1: this rank owns all its MoE expert params - local_2d_moe_params = list(self._rank2params_2d_moe.get(0, [])) + local_2d_params = [] + for color_key, rank2params in self._rank2params_2d_by_color.items(): + group_info = self._color_to_group_info.get(color_key, {}) + color_rank = group_info.get('rank', 0) + world_size = group_info.get('world_size', 1) + rank_key = color_rank if world_size > 1 else 0 + local_2d_params.extend(rank2params.get(rank_key, [])) - local_opt_params = ( - local_2d_params - + local_2d_moe_params - + list(self._local_parameter_list_1d) - ) + local_opt_params = local_2d_params + list(self._local_parameter_list_1d) self._set_inner_opt_attr('_parameter_list', local_opt_params) self._set_inner_opt_attr('_param_groups', local_opt_params) @@ -309,16 +302,16 @@ def __init__(self, optimizer, hcg=None): timer.set_timers() self.timers = timer.get_timers() - # --- [SLICE SIZE SUMMARY] Per-rank slice param sizes within this PP stage --- + # --- Per-rank parameter size summary (for load balancing diagnostics) --- _sg_group = hcg.get_sharding_parallel_group() _N = self._sharding_world_size - # 2D (non-MoE) params owned by this rank + # 2D params owned by this sharding rank (default color, via legacy alias) _local_2d_numel = sum( int(functools_reduce(lambda x, y: x * y, p.shape, 1)) for p in self._rank2params_2d.get(self._sharding_rank, []) ) - # 2D (MoE) params owned by this rank + # 2D MoE-expert params owned by this rank (moe_expert color, via legacy alias) _moe_rank_key = ( self._moe_sharding_rank if self._moe_sharding_world_size > 1 else 0 ) @@ -326,8 +319,7 @@ def __init__(self, optimizer, hcg=None): int(functools_reduce(lambda x, y: x * y, p.shape, 1)) for p in self._rank2params_2d_moe.get(_moe_rank_key, []) ) - # 1D (AdamW) slice: each rank owns ceil(param.numel / world_size) elements per param. - # Sum over all 1D params in this sharding group (same color). + # 1D (AdamW) slice: each rank holds ceil(numel / sharding_world_size) elements. _local_1d_numel = sum( math.ceil( int(functools_reduce(lambda x, y: x * y, p.shape, 1)) / _N @@ -370,50 +362,40 @@ def __init__(self, optimizer, hcg=None): # ------------------------------------------------------------------ @staticmethod - def _build_color_to_group_info(hcg): - """Build a mapping from color to communication group info. + def _build_color_to_group_info_from_params(parameter_list, default_group): + """Build color->group_info mapping dynamically from param.color attributes. + + When param.color is a dict containing 'group', the comm group is read + directly from the param — no hcg-specific method registration required. Returns: dict: { - None: {'group': sharding_group, 'world_size': N, 'rank': r}, - 'moe_expert': {'group': moe_sharding_group, 'world_size': M, 'rank': s}, - # Future colors can be added here + None: {'group': default_group, 'world_size': N, 'rank': r}, + '': {'group': group, 'world_size': M, 'rank': s}, + # additional entries auto-populated from param.color dicts } """ - color_to_info = {} - - # Default sharding group - sharding_world_size = hcg.get_sharding_parallel_world_size() - sharding_group = hcg.get_sharding_parallel_group() - color_to_info[None] = { - 'group': sharding_group, - 'world_size': sharding_world_size, - 'rank': sharding_group.rank if sharding_group else 0, + color_to_info = { + None: { + 'group': default_group, + 'world_size': len(default_group.ranks) if default_group else 1, + 'rank': default_group.rank if default_group else 0, + } } - - # MoE sharding group (if available) - if hasattr(hcg, "get_moe_sharding_parallel_world_size"): - moe_world_size = hcg.get_moe_sharding_parallel_world_size() - if moe_world_size > 0: - moe_group = hcg.get_moe_sharding_parallel_group() - color_to_info['moe_expert'] = { - 'group': moe_group, - 'world_size': moe_world_size, - 'rank': moe_group.rank if moe_group else 0, - } - - # Future: Add more color -> group mappings here as needed - # Example: - # if hasattr(hcg, "get_custom_parallel_world_size"): - # custom_world_size = hcg.get_custom_parallel_world_size() - # if custom_world_size > 0: - # custom_group = hcg.get_custom_parallel_group() - # color_to_info['custom'] = { - # 'group': custom_group, - # 'world_size': custom_world_size, - # 'rank': custom_group.rank if custom_group else 0, - # } - + for p in parameter_list: + color = getattr(p, 'color', -1) + if isinstance(color, dict): + color_key = color.get('color', -1) + if ( + color_key not in (-1, None) + and color_key not in color_to_info + ): + group = color.get('group', default_group) + color_to_info[color_key] = { + 'group': group, + 'world_size': len(group.ranks) if group else 1, + 'rank': group.rank if group else 0, + } return color_to_info def _partition_2d_parameters(self, params, world_size, label=""): @@ -515,7 +497,7 @@ def _build_1d_comm_buffers(self): else 256 * 1024 * 1024 ) - # Group 1D params by color (for MoE compatibility) + # Group 1D params by (color, comm_group) so each group uses its own FusedCommBuffer color_dict = defaultdict(list) for param in self._params_1d: color = getattr(param, 'color', -1) @@ -641,26 +623,23 @@ def reduce_gradients(self, parameter_list, hcg): paddle.device.synchronize() with framework.no_grad(): - # --- 2D params: reduce via comm buffers | per tensors --- if self._use_fuse_gradients: for comm_buffer in self.comm_buffer_2d: comm_buffer._comm_grads() else: - # --- Non-MoE 2D params: reduce to owner rank via sharding_group --- + # --- 2D params: reduce to owner rank via each color's group --- sharding_group = hcg.get_sharding_parallel_group() - self._reduce_2d_grads( - self._params_2d, self._param2rank_2d, sharding_group - ) - - # --- MoE expert 2D params: reduce to owner rank via moe_sharding_group --- - if self._params_2d_moe and self._moe_sharding_group is not None: - if self._moe_sharding_world_size > 1: - self._reduce_2d_grads( - self._params_2d_moe, - self._param2rank_2d_moe, - self._moe_sharding_group, + for color_key, params_2d in self._params_2d_by_color.items(): + if not params_2d: + continue + group_info = self._color_to_group_info.get(color_key, {}) + group = group_info.get('group', sharding_group) + world_size = group_info.get('world_size', 1) + if world_size > 1: + param2rank = self._param2rank_2d_by_color.get( + color_key, {} ) - # When moe_sharding_degree=1, no reduce needed (single rank group) + self._reduce_2d_grads(params_2d, param2rank, group) # --- 1D params: reduce-scatter via comm buffers --- for comm_buffer in self._comm_buffer_list: @@ -676,6 +655,7 @@ def reduce_gradients(self, parameter_list, hcg): if self._use_fuse_gradients: for comm_buffer in self.comm_buffer_2d: comm_buffer.scale_grads() + for comm_buffer in self._comm_buffer_list: comm_buffer.scale_grads() @@ -735,20 +715,16 @@ def _sharding_sync_parameters(self): with framework.no_grad(): all_tasks = [] - # --- Non-MoE 2D params: broadcast from owner via sharding_group --- - all_tasks.extend( - self._broadcast_2d_params(self._rank2params_2d, comm_group) - ) - - # --- MoE expert 2D params: broadcast from owner via moe_sharding_group --- - if self._params_2d_moe and self._moe_sharding_group is not None: - if self._moe_sharding_world_size > 1: + # --- 2D params: broadcast from owner via each color's group --- + for color_key, rank2params in self._rank2params_2d_by_color.items(): + group_info = self._color_to_group_info.get(color_key, {}) + group = group_info.get('group', comm_group) + world_size = group_info.get('world_size', 1) + if world_size > 1: all_tasks.extend( - self._broadcast_2d_params( - self._rank2params_2d_moe, self._moe_sharding_group - ) + self._broadcast_2d_params(rank2params, group) ) - # When moe_sharding_degree=1, no broadcast needed (single rank group) + # world_size=1: single rank group, no broadcast needed for task in all_tasks: task.wait() @@ -840,41 +816,30 @@ def step(self): if not isinstance(self._origin_parameter_list[0], dict): params_grads = [] - # --- Non-MoE 2D params on this rank: full tensors --- - local_2d = self._rank2params_2d.get(self._sharding_rank, []) - for param in local_2d: - if param.stop_gradient: - continue - grad_var = param._grad_ivar() - if hasattr(param, "main_grad") and param.main_grad is not None: - grad_var = param.main_grad - if grad_var is not None: - params_grads.append((param, grad_var)) - - # --- MoE expert params on this rank --- - # Pass the original param (2D or 3D) directly to the optimizer. - # _muon_update already handles both shapes: + # --- All 2D params on this rank (all colors): full tensors --- + # Pass the original param directly to the optimizer. + # _muon_update handles both shapes: # - 2D [H, I]: standard Newton-Schulz - # - 3D [n_experts, H, I]: per-expert Newton-Schulz loop (Step 4) - # Keeping the original name avoids registering _expert_N accumulator - # keys that are absent from model_sharded_state_dict, which would - # break sharded_state_dict (checkpoint save). - if self._moe_sharding_world_size > 1: - local_2d_moe = self._rank2params_2d_moe.get( - self._moe_sharding_rank, [] - ) - else: - local_2d_moe = self._rank2params_2d_moe.get(0, []) - - for param in local_2d_moe: - if param.stop_gradient: - continue - grad_var = param._grad_ivar() - if hasattr(param, "main_grad") and param.main_grad is not None: - grad_var = param.main_grad - if grad_var is None: - continue - params_grads.append((param, grad_var)) + # - 3D [n_experts, H, I]: per-expert Newton-Schulz loop + # Keeping the original param name avoids registering _expert_N + # accumulator keys absent from model_sharded_state_dict, which + # would break sharded_state_dict (checkpoint save). + for color_key, rank2params in self._rank2params_2d_by_color.items(): + group_info = self._color_to_group_info.get(color_key, {}) + color_rank = group_info.get('rank', 0) + world_size = group_info.get('world_size', 1) + rank_key = color_rank if world_size > 1 else 0 + for param in rank2params.get(rank_key, []): + if param.stop_gradient: + continue + grad_var = param._grad_ivar() + if ( + hasattr(param, "main_grad") + and param.main_grad is not None + ): + grad_var = param.main_grad + if grad_var is not None: + params_grads.append((param, grad_var)) # --- 1D params: slice params (element-wise shards) --- for param in self._params_1d: @@ -908,7 +873,8 @@ def step(self): @framework.dygraph_only def set_state_dict(self, state_dict): inner_state = {} - # Local parameters = local 2D + local MoE 2D + 1D slice params + # Collect local parameters: 2D whole-tensor params + 1D original params + # (set_state_dict uses legacy aliases; covers default and moe_expert colors) local_2d = list(self._rank2params_2d.get(self._sharding_rank, [])) if self._moe_sharding_world_size > 1: local_2d_moe = list( @@ -962,9 +928,9 @@ def sharded_state_dict(self, model_sharded_state_dict): Overrides the inner Muon optimizer's sharded_state_dict to handle V3's hybrid sharding scheme: - - 2D Muon params (non-MoE and MoE): whole tensor, shape matches - model's local_shape. Handled by delegating to the inner Muon's - sharded_state_dict after filtering out 1D param states. + - 2D Muon params: whole tensor, shape matches model's local_shape. + Handled by delegating to the inner Muon's sharded_state_dict after + filtering out 1D param states. - 1D AdamW params: accumulators are 1D shards (from reduce-scatter); wrapped with is_flattened=True + flattened_range, like V2. """ diff --git a/python/paddle/optimizer/muon.py b/python/paddle/optimizer/muon.py index 5d447204f35fc4..1eaabfef161cc8 100644 --- a/python/paddle/optimizer/muon.py +++ b/python/paddle/optimizer/muon.py @@ -16,6 +16,10 @@ import logging import os from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable import paddle from paddle.base import framework @@ -40,34 +44,6 @@ # ------------------------------------------------------------------ -@dataclass -class QKVInfo: - """Metadata for QKV weight matrices (GQA). - - Attributes: - head_num: Number of attention heads (Q heads). - kv_head_num: Number of key-value heads (for GQA). - num_key_value_groups: Number of Q heads per KV head. - """ - - head_num: int - kv_head_num: int - num_key_value_groups: int - - -@dataclass -class MLAInfo: - """Metadata for MLA weight matrices needed for head-split. - - Attributes: - param_name: Name of the parameter (q_b_proj, kv_b_proj, o_proj). - head_num: Number of attention heads. - """ - - param_name: str - head_num: int - - @dataclass class MuonParamInfo: """Muon update metadata for a single parameter. @@ -77,29 +53,13 @@ class MuonParamInfo: Attributes: use_muon: If True, use Muon (orthogonal) updates; otherwise AdamW. - qkv_info: Required for QKV weight matrices. - intermediate_size: Required for FFN gate_up weights when muon_ffn_split is True. + split_concat_func: Optional callable that implements the slice strategy. + Signature: split_concat_func(matrix, ortho_fn, **kwargs) -> sliced_matrix + If None, whole-matrix orthogonalisation is used. """ use_muon: bool = True - qkv_info: QKVInfo | None = None - mla_info: MLAInfo | None = None - intermediate_size: int | None = None - - @property - def is_qkv(self) -> bool: - """True if this is a QKV weight matrix.""" - return self.qkv_info is not None - - @property - def is_mla(self) -> bool: - """True if this is an MLA weight matrix.""" - return self.mla_info is not None - - @property - def is_ffn_gate_up(self) -> bool: - """True if this is an FFN gate_up weight matrix.""" - return self.intermediate_size is not None + split_concat_func: Callable | None = None # Type alias for the parameter info mapping @@ -213,21 +173,22 @@ class Muon(Optimizer): any of these substrings will use AdamW instead of Muon. Example: ``['embed', 'bias', 'lm_head', 'mlp.gate']``. Default: ``None``. - muon_qkv_update_mode (str): Strategy for QKV fused weight matrices. - ``"split_head"`` orthogonalises each Q/K/V head independently; - ``"split_qkv"`` treats Q, K, V as three separate matrices; - ``"fused_qkv"`` treats the entire QKV matrix as one. - Default: ``"split_head"``. - muon_ffn_split (bool): If True, split FFN gate_up fused weights into - gate and up projections and orthogonalise them independently. - Default: ``False``. muon_extra_scale_factor (float): Extra multiplicative scale applied after the dimension-dependent scaling in ``_scaling_fn``. Default: ``0.2``. muon_param_info_map (MuonParamInfoMap | None): Per-parameter metadata dict mapping param name to :class:`MuonParamInfo` (use_muon, - qkv_info, intermediate_size). Built by Trainer and passed in. + split_concat_func). Built by Trainer and passed in. Default: ``None``. + muon_slice_config (dict | None): Declarative slice configuration mapping + param name to (slice_fn, slice_kwargs) tuple. Built by model's + _build_muon_slice_config method and passed in. This allows slice + strategies to be defined in the model configuration rather than + hard-coded in the optimizer. Default: ``None``. + ns_matmul_dtype (paddle.dtype | None): Dtype for Newton-Schulz matmul + iterations. ``None`` = auto-detect: bfloat16 on Ampere+ (capability + >= 8.0), float32 on V100 and older. Pass ``paddle.float32`` + explicitly to force float32. Default: ``None``. multi_precision (bool): Maintain FP32 master weights when training in BF16/FP16. Default: ``False``. name (str | None): Optional name for the optimizer instance. @@ -254,10 +215,9 @@ def __init__( apply_decay_param_fun=None, muon_version=1, muon_exclude_patterns=None, - muon_qkv_update_mode="split_head", - muon_ffn_split=False, muon_extra_scale_factor=0.2, muon_param_info_map: MuonParamInfoMap | None = None, + ns_matmul_dtype=None, multi_precision=False, name=None, **kwargs, @@ -305,11 +265,22 @@ def __init__( self._apply_decay_param_fun = apply_decay_param_fun self._muon_split_logged = False self._muon_exclude_patterns = muon_exclude_patterns - self._muon_qkv_update_mode = muon_qkv_update_mode - self._muon_ffn_split = muon_ffn_split self._muon_extra_scale_factor = muon_extra_scale_factor self._ns_coeff_type = ns_coeff_type self._muon_param_info_map = muon_param_info_map or {} + # Dtype for Newton-Schulz matmul. + # None = auto: bfloat16 on Ampere+ (capability >= 8.0), float32 on older. + if ns_matmul_dtype is None: + cap = ( + paddle.device.cuda.get_device_capability() + if paddle.is_compiled_with_cuda() + else (0, 0) + ) + self._ns_matmul_dtype = ( + paddle.bfloat16 if cap[0] >= 8 else paddle.float32 + ) + else: + self._ns_matmul_dtype = ns_matmul_dtype self._default_dict.update(defaults) # ------------------------------------------------------------------ @@ -383,7 +354,11 @@ def _create_accumulators(self, block, parameters): @staticmethod def _zeropower_via_newtonschulz5( - X, steps=5, eps=1e-9, ns_coeff_type="simple" + X, + steps=5, + eps=1e-9, + ns_coeff_type="simple", + ns_matmul_dtype=paddle.bfloat16, ): """Approximate the matrix sign function via Newton-Schulz iteration. @@ -393,6 +368,8 @@ def _zeropower_via_newtonschulz5( eps: Small constant for numerical stability. ns_coeff_type: Type of coefficient set to use. Options: "simple", "quintic", "polar_express", "aol". + ns_matmul_dtype: Dtype for matmul iterations. Defaults to + bfloat16. Pass paddle.float32 for V100 compatibility. """ # Get coefficient set coeff_sets = _NS_COEFFICIENT_SETS.get( @@ -410,7 +387,7 @@ def _zeropower_via_newtonschulz5( X_flat = paddle.nn.functional.normalize( X_flat, p=2, axis=-1, epsilon=eps ) - X = X_flat.reshape(orig_shape).astype(paddle.bfloat16) + X = X_flat.reshape(orig_shape).astype(ns_matmul_dtype) # Iterate with cycling coefficients for i in range(steps): @@ -433,157 +410,6 @@ def _scaling_fn(orthogonal_update, version, extra_scale_factor=1.0): scale = max(dout, din) ** 0.5 return orthogonal_update * scale * extra_scale_factor - @staticmethod - def _ortho_qkv_per_head( - matrix_2d_global, - kv_head_num, - num_key_value_groups, - ortho_fn, - ): - """Orthogonalise each Q/K/V head independently (interleaved layout). - - Args: - matrix_2d_global: QKV weight matrix [hidden, (num_key_value_groups + 2)*kv_head_num*head_dim]. - kv_head_num: Number of K/V heads. - num_key_value_groups: Number of Q heads per KV head. - ortho_fn: Callable (2d_matrix) -> 2d_matrix applying NS + scaling. - - Returns: - orthogonal_update: Same shape as input, each head orthogonalised. - """ - # Interleaved layout: [Q_kv0, K0, V0, Q_kv1, K1, V1, ...] - head_dim = matrix_2d_global.shape[1] // ( - num_key_value_groups * kv_head_num + 2 * kv_head_num - ) - groups = paddle.split(matrix_2d_global, kv_head_num, axis=1) - - processed_groups = [] - for group in groups: - q_part, k_head, v_head = paddle.split( - group, - [num_key_value_groups * head_dim, head_dim, head_dim], - axis=1, - ) - q_heads = paddle.split(q_part, num_key_value_groups, axis=1) - q_ortho = paddle.concat([ortho_fn(h) for h in q_heads], axis=1) - processed_groups.append( - paddle.concat( - [q_ortho, ortho_fn(k_head), ortho_fn(v_head)], axis=1 - ) - ) - - return paddle.concat(processed_groups, axis=1) - - @staticmethod - def _ortho_qkv_sep( - matrix_2d, - kv_head_num, - num_key_value_groups, - ortho_fn, - ): - """Orthogonalise Q, K, V as three separate whole matrices (interleaved layout). - - Gathers all Q heads into one block, all K heads into one block, all V heads - into one block (across kv_groups), orthogonalises each block as a whole with - one NS call, then scatters back to interleaved order. - - Args: - matrix_2d: QKV weight matrix [hidden, (num_key_value_groups + 2)*kv_head_num*head_dim]. - kv_head_num: Number of K/V heads. - num_key_value_groups: Number of Q heads per KV head. - ortho_fn: Callable (2d_matrix) -> 2d_matrix applying NS + scaling. - - Returns: - orthogonal_update: Same shape as input, Q/K/V each orthogonalised as whole. - """ - # Interleaved layout: [Q_kv0, K0, V0, Q_kv1, K1, V1, ...] - head_dim = matrix_2d.shape[1] // ( - num_key_value_groups * kv_head_num + 2 * kv_head_num - ) - q_group_size = num_key_value_groups * head_dim - - # Step 1: gather Q / K / V parts from each kv_group - groups = paddle.split(matrix_2d, kv_head_num, axis=1) - q_parts, k_parts, v_parts = [], [], [] - for group in groups: - q_p, k_p, v_p = paddle.split( - group, [q_group_size, head_dim, head_dim], axis=1 - ) - q_parts.append(q_p) - k_parts.append(k_p) - v_parts.append(v_p) - - # Step 2: orthogonalise each projection as one whole matrix - q_ortho = ortho_fn(paddle.concat(q_parts, axis=1)) - k_ortho = ortho_fn(paddle.concat(k_parts, axis=1)) - v_ortho = ortho_fn(paddle.concat(v_parts, axis=1)) - - # Step 3: split back and restore interleaved layout - q_groups = paddle.split(q_ortho, kv_head_num, axis=1) - k_groups = paddle.split(k_ortho, kv_head_num, axis=1) - v_groups = paddle.split(v_ortho, kv_head_num, axis=1) - - return paddle.concat( - [ - paddle.concat([q_groups[i], k_groups[i], v_groups[i]], axis=1) - for i in range(kv_head_num) - ], - axis=1, - ) - - @staticmethod - def _ortho_ffn_gate_up(matrix, intermediate_size, ortho_fn): - """Orthogonalise gate and up projections independently for FFN. - - Args: - matrix: FFN weight tensor. - - 2D: [hidden, 2*intermediate_size] for standard FFN - - 3D: [num_experts, hidden, 2*intermediate_size] for MoE FFN - intermediate_size: Size of each of gate/up projections. - ortho_fn: Callable (2d_matrix) -> 2d_matrix applying NS + scaling. - - Returns: - orthogonal_update: Tensor with gate and up orthogonalised separately. - """ - if matrix.ndim == 2: - gate, up = paddle.split( - matrix, [intermediate_size, intermediate_size], axis=1 - ) - return paddle.concat([ortho_fn(gate), ortho_fn(up)], axis=1) - - elif matrix.ndim == 3: - # MoE FFN: [n_experts, hidden, 2*intermediate_size] - expert_updates = [] - for ei in range(matrix.shape[0]): - gate, up = paddle.split( - matrix[ei], [intermediate_size, intermediate_size], axis=1 - ) - expert_updates.append( - paddle.concat([ortho_fn(gate), ortho_fn(up)], axis=1) - ) - return paddle.stack(expert_updates, axis=0) - - else: - raise ValueError( - f"FFN gate_up split expects 2D or 3D tensor, got shape {matrix.shape}" - ) - - @staticmethod - def _ortho_mla_per_head( - matrix_2d_global, - head_num, - ortho_fn, - axis, - ): - """Orthogonalise each MLA head independently.""" - groups = paddle.split(matrix_2d_global, head_num, axis=axis) - - processed_groups = [] - for group in groups: - processed_groups.append(ortho_fn(group)) - - return paddle.concat(processed_groups, axis=axis) - # ------------------------------------------------------------------ # Per-parameter update rules # ------------------------------------------------------------------ @@ -658,9 +484,6 @@ def _muon_update( """ param_shape = getattr(param, "original_shape", param.shape) param_info = self._muon_param_info_map.get(param.name) - is_qkv = param_info is not None and param_info.is_qkv - is_mla: bool = param_info is not None and param_info.is_mla - is_ffn_gate_up = param_info is not None and param_info.is_ffn_gate_up with paddle.no_grad(): grad_f32 = ( @@ -692,6 +515,7 @@ def ortho_fn(m): steps=ns_steps, eps=epsilon, ns_coeff_type=self._ns_coeff_type, + ns_matmul_dtype=self._ns_matmul_dtype, ) scaled = Muon._scaling_fn( ns_out, version, self._muon_extra_scale_factor @@ -699,96 +523,26 @@ def ortho_fn(m): return scaled # Step 3: Newton-Schulz orthogonalisation - if is_ffn_gate_up and self._muon_ffn_split: - # FFN gate_up split: orthogonalise gate and up projections independently. - intermediate_size = param_info.intermediate_size - if MUON_DEBUG: - _global_rank = paddle.distributed.get_rank() - if _global_rank == 0: - _logger.info( - f"[Muon] FFN split: param={param.name}, " - f"shape={matrix_2d_global.shape}, " - f"intermediate_size={intermediate_size}" - ) - - orthogonal_update = Muon._ortho_ffn_gate_up( - matrix_2d_global, intermediate_size, ortho_fn - ) - elif matrix_2d_global.ndim == 3: - # 3D fused MoE expert tensor [n_experts, H, I]. - # Apply Newton-Schulz independently to each expert's 2D slice. - n_experts = matrix_2d_global.shape[0] - orthogonal_update = paddle.stack( - [ortho_fn(matrix_2d_global[ei]) for ei in range(n_experts)], - axis=0, - ) - elif is_qkv and self._muon_qkv_update_mode in ( - "split_head", - "split_qkv", + # 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 ): - # Read QKV head info from param_info - qkv_info = param_info.qkv_info - kv_head_num = qkv_info.kv_head_num - num_key_value_groups = qkv_info.num_key_value_groups - - if self._muon_qkv_update_mode == "split_head": - # split_head update: each Q/K/V head orthogonalised independently. - if MUON_DEBUG: - _global_rank = paddle.distributed.get_rank() - if _global_rank == 0: - _logger.info( - f"[Muon] QKV split_head: param={param.name}, " - f"shape={matrix_2d_global.shape}, " - f"heads={qkv_info.head_num}/{kv_head_num}, " - f"num_key_value_groups={num_key_value_groups}" - ) - orthogonal_update = Muon._ortho_qkv_per_head( - matrix_2d_global, - kv_head_num, - num_key_value_groups, - ortho_fn, - ) - else: - # split_qkv: Q, K, V each as a whole matrix, one NS call each. - if MUON_DEBUG: - _global_rank = paddle.distributed.get_rank() - if _global_rank == 0: - _logger.info( - f"[Muon] QKV split_qkv: param={param.name}, " - f"shape={matrix_2d_global.shape}, " - f"head_num={qkv_info.head_num}, kv_head_num={kv_head_num}, " - f"num_key_value_groups={num_key_value_groups}" - ) - orthogonal_update = Muon._ortho_qkv_sep( - matrix_2d_global, - kv_head_num, - num_key_value_groups, - ortho_fn, - ) - elif is_mla and self._muon_qkv_update_mode == "split_head": - # MLA split_head update: each head of [q_b_proj, kv_b_proj, o_proj] orthogonalised independently. - mla_info = param_info.mla_info - param_name: str = mla_info.param_name - head_num = mla_info.head_num + # Use slice function defined in model configuration + orthogonal_update = param_info.split_concat_func( + matrix_2d_global, ortho_fn + ) if MUON_DEBUG: _global_rank = paddle.distributed.get_rank() if _global_rank == 0: + _sf = param_info.split_concat_func _logger.info( - f"[Muon] MLA split_head: param={param.name}, param_name={param_name}, " - f"shape={matrix_2d_global.shape}, " - f"head_num={head_num}" + f"[Muon] Using split_concat_func: param={param.name}, " + f"split_concat_func={_sf.func.__name__}, " + f"args={_sf.args}, kwargs={_sf.keywords}" ) - assert param_name in ("q_b_proj", "kv_b_proj", "o_proj"), ( - f"Unsupported MLA param name: {param_name}" - ) - orthogonal_update = Muon._ortho_mla_per_head( - matrix_2d_global, - head_num, - ortho_fn, - 0 if param_name == "o_proj" else 1, - ) else: - # Standard 2D update: entire matrix as one Newton-Schulz call. + # Default: whole matrix orthogonalisation orthogonal_update = ortho_fn(matrix_2d_global) # Step 4: Apply update with optional weight decay @@ -817,6 +571,9 @@ def _apply_optimize(self, loss, startup_program, params_grads): if self._grad_clip is not None: params_grads = self._grad_clip(params_grads) + # apply for zcc + self._maybe_refuse() + group = self._default_dict lr = self._learning_rate if isinstance(lr, paddle.optimizer.lr.LRScheduler): diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py index e5c7dcd9733201..0b88ec31ce53d8 100644 --- a/python/paddle/optimizer/optimizer.py +++ b/python/paddle/optimizer/optimizer.py @@ -365,7 +365,10 @@ def _maybe_refuse(self): return # TODO(@gexiao): support other optimizer if needed - if self.__class__.__name__ != "AdamW": + if ( + self.__class__.__name__ != "AdamW" + and self.__class__.__name__ != "Muon" + ): return # add buffer check diff --git a/test/collective/fleet/hybrid_parallel_sharding_muon_model.py b/test/collective/fleet/hybrid_parallel_sharding_muon_model.py index a478e8e155cdb8..b0df5dabc7bb95 100644 --- a/test/collective/fleet/hybrid_parallel_sharding_muon_model.py +++ b/test/collective/fleet/hybrid_parallel_sharding_muon_model.py @@ -14,12 +14,14 @@ # # Validates Muon optimizer with MuonShardingOptimizer. # Muon requires whole 2D tensors for orthogonalization, so split_param is disabled. -# Tests all combinations of QKV/FFN/ns_coeff_type modes. +# Tests ns_coeff_type modes, custom color groups, and split_concat_func. # Topology: sharding_degree=2, mp_degree=1 (2 ranks total) +import os import random import unittest from dataclasses import dataclass +from functools import partial import numpy as np @@ -28,13 +30,13 @@ from paddle.distributed.fleet.utils import mix_precision_utils from paddle.optimizer.muon import ( MuonParamInfo, - QKVInfo, _default_should_use_muon, ) +# Enable MUON_DEBUG to cover the debug logging branch (muon.py L532-539) +os.environ["MUON_DEBUG"] = "1" + # Parameter combinations -QKV_UPDATE_MODES = ["split_head", "split_qkv", "fused_qkv"] -FFN_SPLITS = [True, False] NS_COEFF_TYPES = ["simple", "quintic", "polar_express", "aol"] # Model config @@ -52,6 +54,32 @@ sharding_degree = 2 +# ------------------------------------------------------------------ +# Slice functions (called as split_concat_func(matrix_2d_global, ortho_fn)) +# ------------------------------------------------------------------ + + +def _qkv_sep(matrix_2d, ortho_fn, kv_head_num=None, num_key_value_groups=None): + """Slice QKV into Q, K, V blocks, orthogonalise each as whole.""" + head_dim_local = matrix_2d.shape[1] // ( + num_key_value_groups * kv_head_num + 2 * kv_head_num + ) + q_dim = num_key_value_groups * kv_head_num * head_dim_local + k_dim = kv_head_num * head_dim_local + v_dim = kv_head_num * head_dim_local + + q, k, v = paddle.split(matrix_2d, [q_dim, k_dim, v_dim], axis=1) + return paddle.concat([ortho_fn(q), ortho_fn(k), ortho_fn(v)], axis=1) + + +def _ffn_split(matrix_2d, ortho_fn, intermediate_size=None): + """Split gate_up into gate and up, orthogonalise each.""" + gate, up = paddle.split( + matrix_2d, [intermediate_size, intermediate_size], axis=1 + ) + return paddle.concat([ortho_fn(gate), ortho_fn(up)], axis=1) + + @dataclass class TestConfig: """Test model config.""" @@ -199,9 +227,48 @@ def train_batch(self, batch, model, optimizer): optimizer.clear_grad() return loss - def build_optimizer(self, model, qkv_mode, ffn_split, ns_coeff): - """Build Muon optimizer, ref: PaddleFormers trainer.py L3122-3173.""" + def _init_weights(self): + """Create shared numpy weight arrays.""" + return ( + np.random.random_sample((vocab_size, hidden_size)), + np.random.random_sample((hidden_size, qkv_dim)), + np.random.random_sample((head_num * head_dim, hidden_size)), + np.random.random_sample((hidden_size, 2 * intermediate_size)), + np.random.random_sample((hidden_size, hidden_size)), + np.random.random_sample((hidden_size, vocab_size)), + ) + + def _build_split_concat_func_map(self, model): + """Build split_concat_func_map with QKV sep and FFN split for applicable params.""" + num_key_value_groups = head_num // kv_head_num + slice_map = {} + for name, param in model.named_parameters(): + if "qkv_proj" in name: + slice_map[param.name] = partial( + _qkv_sep, + kv_head_num=kv_head_num, + num_key_value_groups=num_key_value_groups, + ) + elif "up_gate_proj" in name: + slice_map[param.name] = partial( + _ffn_split, + intermediate_size=intermediate_size, + ) + return slice_map + def build_optimizer( + self, model, ns_coeff, split_concat_func_map=None, ns_matmul_dtype=None + ): + """Build Muon optimizer. + + Args: + 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). + ns_matmul_dtype: Optional explicit dtype for NS matmul. + Covers muon.py L283 (explicit ns_matmul_dtype branch). + """ muon_param_info_map = {} exclude_patterns = ["embed", "bias", "lm_head"] @@ -209,53 +276,32 @@ def build_optimizer(self, model, qkv_mode, ffn_split, ns_coeff): use_muon = _default_should_use_muon( name, param.shape, exclude_patterns ) + sf = None + if split_concat_func_map and param.name in split_concat_func_map: + sf = split_concat_func_map[param.name] + param_info = MuonParamInfo(use_muon=use_muon, split_concat_func=sf) + muon_param_info_map[param.name] = param_info - # QKV params: set QKVInfo - if "qkv_proj.weight" in name and len(param.shape) == 2: - param_info = MuonParamInfo( - use_muon=use_muon, - qkv_info=QKVInfo( - head_num=self.config.head_num, - kv_head_num=self.config.kv_head_num, - num_key_value_groups=self.config.head_num - // self.config.kv_head_num, - ), - ) - # FFN gate_up params: set intermediate_size - elif "up_gate_proj.weight" in name and ffn_split: - param_info = MuonParamInfo( - use_muon=use_muon, - intermediate_size=self.config.intermediate_size, - ) - else: - param_info = MuonParamInfo(use_muon=use_muon) + kwargs = {} + if ns_matmul_dtype is not None: + kwargs['ns_matmul_dtype'] = ns_matmul_dtype - muon_param_info_map[param.name] = param_info return paddle.optimizer.Muon( parameters=model.parameters(), learning_rate=0.001, weight_decay=0.00001, grad_clip=paddle.nn.ClipGradByGlobalNorm(0.5), muon_param_info_map=muon_param_info_map, - muon_qkv_update_mode=qkv_mode, - muon_ffn_split=ffn_split, ns_coeff_type=ns_coeff, + **kwargs, ) - def _run_single_test(self, qkv_mode, ffn_split, ns_coeff): - """Run single test combination.""" - # Init weights - np_embed = np.random.random_sample((vocab_size, hidden_size)) - np_qkv = np.random.random_sample((hidden_size, qkv_dim)) - np_o_proj = np.random.random_sample((head_num * head_dim, hidden_size)) - np_up_gate = np.random.random_sample( - (hidden_size, 2 * intermediate_size) + def _build_model(self, weights): + """Build a single model instance from weights.""" + np_embed, np_qkv, np_o_proj, np_up_gate, np_down_proj, np_lm_head = ( + weights ) - np_down_proj = np.random.random_sample((hidden_size, hidden_size)) - np_lm_head = np.random.random_sample((hidden_size, vocab_size)) - - # Distributed model - model_a = QKVFFNNet( + model = QKVFFNNet( self.config, np_embed, np_qkv, @@ -264,34 +310,65 @@ def _run_single_test(self, qkv_mode, ffn_split, ns_coeff): np_down_proj, np_lm_head, ) - model_a = mix_precision_utils.MixPrecisionLayer( - model_a, dtype="bfloat16" - ) - model_a = paddle.amp.decorate( - models=model_a, level='O2', dtype='bfloat16' + model = mix_precision_utils.MixPrecisionLayer(model, dtype="bfloat16") + model = paddle.amp.decorate(models=model, level='O2', dtype='bfloat16') + return model + + def _run_single_test( + self, ns_coeff, color_params=None, use_slice=False, explicit_dtype=False + ): + """Run single test combination. + + Args: + 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). + """ + weights = self._init_weights() + + # --- Distributed model (model_a) --- + model_a = self._build_model(weights) + + # Assign custom color before optimizer construction + if color_params: + hcg = fleet.get_hybrid_communicate_group() + sharding_group = hcg.get_sharding_parallel_group() + for name, p in model_a.named_parameters(): + for pattern in color_params: + if pattern in name: + p.color = { + 'color': 'test_color', + 'group': sharding_group, + } + break + + split_concat_func_map = ( + self._build_split_concat_func_map(model_a) if use_slice else None ) + ns_dtype = paddle.float32 if explicit_dtype else None + optimizer_a = self.build_optimizer( - model_a, qkv_mode, ffn_split, ns_coeff + model_a, + ns_coeff, + split_concat_func_map=split_concat_func_map, + ns_matmul_dtype=ns_dtype, ) - # Single-GPU reference model (same MixPrecisionLayer pattern for consistency) - model_b = QKVFFNNet( - self.config, - np_embed, - np_qkv, - np_o_proj, - np_up_gate, - np_down_proj, - np_lm_head, - ) - model_b = mix_precision_utils.MixPrecisionLayer( - model_b, dtype="bfloat16" - ) - model_b = paddle.amp.decorate( - models=model_b, level='O2', dtype='bfloat16' + # --- Reference model (model_b, single-GPU) --- + model_b = self._build_model(weights) + + split_concat_func_map_b = ( + self._build_split_concat_func_map(model_b) if use_slice else None ) optimizer_b = self.build_optimizer( - model_b, qkv_mode, ffn_split, ns_coeff + model_b, + ns_coeff, + split_concat_func_map=split_concat_func_map_b, + ns_matmul_dtype=ns_dtype, ) optimizer_b = mix_precision_utils.MixPrecisionOptimizer(optimizer_b) @@ -326,39 +403,53 @@ def _run_single_test(self, qkv_mode, ffn_split, ns_coeff): err_msg=f"Param {param_a.name} mismatch at step {idx}!", ) - @unittest.skipIf( - not paddle.is_compiled_with_cuda() - or paddle.device.cuda.get_device_capability()[0] < 8, - "BF16 matmul requires GPU compute capability >= 80 (Ampere+)", - ) def test_sharding_muon(self): - """Test all 24 parameter combinations.""" - total = len(QKV_UPDATE_MODES) * len(FFN_SPLITS) * len(NS_COEFF_TYPES) + """Test ns_coeff_type combinations + color/slice/dtype coverage. + + Phase 1: iterate all ns_coeff_types (basic, no slice, no color). + Phase 2: custom color group + split_concat_func + explicit fp32 dtype. + Covers: + - muon_sharding_optimizer.py L388-394: custom color from param.color dict + - muon.py L283: explicit ns_matmul_dtype=paddle.float32 + - muon.py L529: split_concat_func call + - muon.py L535: MUON_DEBUG logging (via MUON_DEBUG=1 env) + """ + total = len(NS_COEFF_TYPES) + 1 # +1 for color/slice/dtype test passed = 0 failed = [] - for qkv_mode in QKV_UPDATE_MODES: - for ffn_split in FFN_SPLITS: - for ns_coeff in NS_COEFF_TYPES: - print( - f"\n[Muon Test] qkv_mode={qkv_mode}, ffn_split={ffn_split}, ns_coeff={ns_coeff}" - ) - try: - self._run_single_test(qkv_mode, ffn_split, ns_coeff) - passed += 1 - print(f"[PASS] {qkv_mode}, {ffn_split}, {ns_coeff}") - except Exception as e: - failed.append((qkv_mode, ffn_split, ns_coeff, str(e))) - print( - f"[FAIL] {qkv_mode}, {ffn_split}, {ns_coeff}: {e}" - ) + # Phase 1: ns_coeff_type combinations + for ns_coeff in NS_COEFF_TYPES: + print(f"\n[Muon Test] ns_coeff={ns_coeff}") + try: + self._run_single_test(ns_coeff) + passed += 1 + print(f"[PASS] {ns_coeff}") + except Exception as e: + failed.append((ns_coeff, str(e))) + print(f"[FAIL] {ns_coeff}: {e}") + + # Phase 2: color + split_concat_func + explicit fp32 dtype + print("\n[Muon Test] color + split_concat_func + explicit fp32 dtype") + try: + self._run_single_test( + "simple", + color_params=["down_proj"], + use_slice=True, + explicit_dtype=True, + ) + passed += 1 + print("[PASS] color + slice + dtype") + except Exception as e: + failed.append(("color+slice+dtype", str(e))) + print(f"[FAIL] color + slice + dtype: {e}") print(f"\n{'=' * 60}") print(f"Muon Sharding Test Summary: {passed}/{total} passed") if failed: print("Failed combinations:") - for qkv, ffn, ns, err in failed: - print(f" - {qkv}, {ffn}, {ns}: {err[:100]}...") + for ns, err in failed: + print(f" - {ns}: {err[:100]}...") print(f"{'=' * 60}") if failed: diff --git a/test/collective/fleet/test_parallel_dygraph_muon.py b/test/collective/fleet/test_parallel_dygraph_muon.py index c908087b465fe7..983e0ccad49f57 100644 --- a/test/collective/fleet/test_parallel_dygraph_muon.py +++ b/test/collective/fleet/test_parallel_dygraph_muon.py @@ -18,20 +18,13 @@ TestMultipleAccelerators, ) -import paddle - class TestMuonParallel(TestMultipleAccelerators): - @unittest.skipIf( - not paddle.is_compiled_with_cuda() - or paddle.device.cuda.get_device_capability()[0] < 8, - "BF16 matmul requires GPU compute capability >= 80 (Ampere+)", - ) def test_muon_sharding_optimizer(self): - """MuonSharding test: iterate all QKV/FFN/ns_coeff_type combinations. + """MuonSharding test: iterate ns_coeff_type combinations. Test logic is in hybrid_parallel_sharding_muon_model.py, - iterating 24 combinations (3 qkv_modes * 2 ffn_splits * 4 ns_coeff_types). + iterating 4 ns_coeff_types. fp32 matmul is auto-selected on V100. """ self.run_mnist_2accelerators('hybrid_parallel_sharding_muon_model.py')