WeHub Mirror
962 lines
36 KiB
Python
962 lines
36 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import math
|
||
import os
|
||
from typing import TYPE_CHECKING
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Callable
|
||
from collections import defaultdict
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING
|
||
|
||
if TYPE_CHECKING:
|
||
from paddle import Tensor
|
||
|
||
import paddle
|
||
from paddle import _C_ops
|
||
from paddle.base import framework
|
||
from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
|
||
ShardedStateDict,
|
||
ShardedWeight,
|
||
create_sharded_weight_with_new_local,
|
||
)
|
||
|
||
from ..nn.clip import GradientClipBase
|
||
from .optimizer import Optimizer
|
||
|
||
# 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__ = []
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# Parameter metadata dataclasses
|
||
# ------------------------------------------------------------------
|
||
|
||
|
||
@dataclass
|
||
class MuonParamInfo:
|
||
"""Muon update metadata for a single parameter.
|
||
|
||
This replaces the previous approach of setting dynamic attributes
|
||
directly on param objects.
|
||
|
||
Attributes:
|
||
use_muon: If True, use Muon (orthogonal) updates; otherwise AdamW.
|
||
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
|
||
split_concat_func: Callable | None = None
|
||
|
||
|
||
# Type alias for the parameter info mapping
|
||
MuonParamInfoMap = dict[str, MuonParamInfo]
|
||
|
||
# ------------------------------------------------------------------
|
||
# Newton-Schulz coefficient sets
|
||
# ------------------------------------------------------------------
|
||
|
||
_NS_COEFFICIENT_SETS = {
|
||
# Simple coefficient set (original)
|
||
"simple": [
|
||
(3.4445, -4.7750, 2.0315),
|
||
],
|
||
# Quintic iteration with optimized coefficients
|
||
# Source: https://leloykun.github.io/ponder/muon-opt-coeffs/
|
||
"quintic": [
|
||
(4.0848, -6.8946, 2.9270),
|
||
(3.9505, -6.3029, 2.6377),
|
||
(3.7418, -5.5913, 2.3037),
|
||
(2.8769, -3.1427, 1.2046),
|
||
(2.8366, -3.0525, 1.2012),
|
||
],
|
||
# Polar Express iteration from https://arxiv.org/abs/2505.16932
|
||
"polar_express": [
|
||
(8.2051, -22.9019, 16.4607),
|
||
(4.0664, -2.8612, 0.5184),
|
||
(3.9096, -2.8234, 0.5250),
|
||
(3.2856, -2.4153, 0.4853),
|
||
(2.2779, -1.6198, 0.3985),
|
||
(1.8726, -1.2307, 0.3585),
|
||
(1.8564, -1.2132, 0.3568),
|
||
(1.8750, -1.2500, 0.3750),
|
||
],
|
||
# AOL coefficients from https://github.com/thib-s/flash-newton-schulz
|
||
"aol": [
|
||
(4.0098, -7.0585, 2.4635),
|
||
(3.4585, -5.5479, 2.5959),
|
||
(2.7573, -3.2939, 1.4254),
|
||
(2.7215, -3.0494, 1.3169),
|
||
],
|
||
"deepseekv4":
|
||
# From DeepSeekV4: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/resolve/main/DeepSeek_V4.pdf
|
||
[(3.4445, -4.7750, 2.0315)] * 8 + [(2.0, -1.5, 0.5)] * 2,
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
# Default parameter classification
|
||
# ------------------------------------------------------------------
|
||
|
||
|
||
def _default_should_use_muon(name, shape, exclude_patterns):
|
||
"""Default fallback logic for determining if a parameter should use Muon.
|
||
|
||
This is only used when param.is_muon is not set. The actual exclusion
|
||
patterns must be configured via training_args.muon_exclude_patterns in yaml.
|
||
|
||
Args:
|
||
name: Parameter name.
|
||
shape: Parameter shape.
|
||
exclude_patterns: List of substrings to exclude from Muon updates.
|
||
Must be provided (e.g., ['embed', 'bias', 'lm_head', 'mlp.gate']).
|
||
|
||
Returns:
|
||
True if the parameter should use Muon (orthogonal) updates.
|
||
|
||
Raises:
|
||
ValueError: If exclude_patterns is None.
|
||
"""
|
||
if exclude_patterns is None:
|
||
raise ValueError(
|
||
"muon_exclude_patterns must be set in yaml config. "
|
||
"Example: muon_exclude_patterns: ['embed', 'bias', 'lm_head', 'mlp.gate']"
|
||
)
|
||
|
||
if len(shape) not in (2, 3):
|
||
return False
|
||
|
||
name_lower = name.lower()
|
||
for pattern in exclude_patterns:
|
||
if pattern.lower() in name_lower:
|
||
return False
|
||
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.
|
||
|
||
For 2-D weight matrices (identified by :func:`_default_should_use_muon`), Muon
|
||
applies orthogonal gradient updates via Newton-Schulz iteration. For all
|
||
other parameters (embeddings, biases, expert weights, …) it falls back to
|
||
a standard AdamW update.
|
||
|
||
Designed for ``MuonShardingOptimizer`` (Sharding Stage1 V3), where 2D parameters are
|
||
assigned as whole tensors to ranks. Currently we do not support TP=1, no sharding gather
|
||
or TP communication is needed during the optimizer step.
|
||
|
||
Args:
|
||
learning_rate (float | LRScheduler): Learning rate. Default: ``0.02``.
|
||
parameters (list[Tensor]): Flat list of parameters to optimize.
|
||
momentum (float): Momentum coefficient for the Muon update. Default: ``0.95``.
|
||
adam_beta1 (float): β₁ for the AdamW fallback. Default: ``0.9``.
|
||
adam_beta2 (float): β₂ for the AdamW fallback. Default: ``0.95``.
|
||
weight_decay (float): Decoupled weight decay. Default: ``0.01``.
|
||
ns_steps (int): Newton-Schulz iteration steps. Default: ``5``.
|
||
ns_coeff_type (str): Preset name for Newton-Schulz coefficients.
|
||
Options: ``"simple"``, ``"quintic"``, ``"polar_express"``,
|
||
``"aol"``, ``"deepseekv4"``, ``"custom"``. Default: ``"simple"``.
|
||
ns_coeffs (list[tuple[float, float, float]] | None): Custom
|
||
Newton-Schulz coefficient set. Each tuple is ``(a, b, c)``
|
||
for one iteration step. Default: ``None``.
|
||
Only used when ns_coeff_type=``custom``.
|
||
nesterov (bool): Use Nesterov momentum in Muon. Default: ``True``.
|
||
adam_epsilon (float): ε for numerical stability in AdamW. Default: ``1e-9``.
|
||
grad_clip (GradientClipBase | None): Gradient clipping. Default: ``None``.
|
||
apply_decay_param_fun (callable | None): Function to select which
|
||
parameters receive weight decay. Default: ``None``.
|
||
muon_version (int): Scaling-function version (1/2/3). Default: ``1``.
|
||
muon_exclude_patterns (list[str] | None): Parameter names containing
|
||
any of these substrings will use AdamW instead of Muon.
|
||
Example: ``['embed', 'bias', 'lm_head', 'mlp.gate']``.
|
||
Default: ``None``.
|
||
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,
|
||
split_concat_func). Built by Trainer and passed in.
|
||
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.
|
||
"""
|
||
|
||
_moment_acc_str = "moment1"
|
||
_moment2_acc_str = "moment2"
|
||
_beta1_pow_acc_str = "beta1_pow_acc"
|
||
_beta2_pow_acc_str = "beta2_pow_acc"
|
||
|
||
def __init__(
|
||
self,
|
||
learning_rate=0.02,
|
||
parameters=None,
|
||
momentum=0.95,
|
||
adam_beta1=0.9,
|
||
adam_beta2=0.95,
|
||
weight_decay=0.01,
|
||
ns_steps=5,
|
||
ns_coeff_type="simple",
|
||
ns_coeffs=None,
|
||
nesterov=True,
|
||
adam_epsilon=1e-9,
|
||
grad_clip=None,
|
||
lr_ratio: Callable[[Tensor], float] | None = None,
|
||
apply_decay_param_fun: Callable[[str], bool] | None = None,
|
||
muon_version=1,
|
||
muon_exclude_patterns=None,
|
||
muon_extra_scale_factor=0.2,
|
||
muon_param_info_map: MuonParamInfoMap | None = None,
|
||
ns_matmul_dtype=None,
|
||
multi_precision=False,
|
||
name=None,
|
||
muon_max_group_bytes=2 * 1024 * 1024 * 1024, # Max 2GB per group
|
||
**kwargs,
|
||
):
|
||
if parameters is None:
|
||
raise ValueError(
|
||
"parameters argument given to the Optimizer should not be None."
|
||
)
|
||
if not isinstance(parameters, list):
|
||
raise TypeError("parameters must be a list.")
|
||
if len(parameters) > 0 and isinstance(parameters[0], dict):
|
||
raise TypeError(
|
||
"Muon optimizer only supports a flat list of parameters, "
|
||
"not a list of parameter groups."
|
||
)
|
||
if grad_clip is not None and not isinstance(
|
||
grad_clip, GradientClipBase
|
||
):
|
||
raise TypeError(
|
||
"'grad_clip' should be an instance of GradientClipBase's derived class"
|
||
)
|
||
|
||
defaults = {
|
||
"momentum": momentum,
|
||
"adam_beta1": adam_beta1,
|
||
"adam_beta2": adam_beta2,
|
||
"weight_decay": weight_decay,
|
||
"ns_steps": ns_steps,
|
||
"nesterov": nesterov,
|
||
"epsilon": adam_epsilon,
|
||
"muon_version": muon_version,
|
||
"ns_coeff_type": ns_coeff_type,
|
||
}
|
||
|
||
super().__init__(
|
||
learning_rate=learning_rate,
|
||
parameters=parameters,
|
||
weight_decay=weight_decay,
|
||
grad_clip=grad_clip,
|
||
name=name,
|
||
)
|
||
|
||
self._multi_precision = multi_precision
|
||
self._master_weights = {}
|
||
self._lr_ratio = lr_ratio
|
||
self._apply_decay_param_fun = apply_decay_param_fun
|
||
self._muon_split_logged = False
|
||
self._muon_exclude_patterns = muon_exclude_patterns
|
||
self._muon_extra_scale_factor = muon_extra_scale_factor
|
||
self._ns_coeff_type = ns_coeff_type
|
||
if ns_coeff_type == "custom":
|
||
assert ns_coeffs is not None, (
|
||
"ns_coeffs must be provided when ns_coeff_type is 'custom'."
|
||
)
|
||
self._ns_coeffs = ns_coeffs
|
||
else:
|
||
assert ns_coeff_type in _NS_COEFFICIENT_SETS, (
|
||
f"Invalid ns_coeff_type: {ns_coeff_type}"
|
||
)
|
||
self._ns_coeffs = _NS_COEFFICIENT_SETS[ns_coeff_type]
|
||
self._muon_param_info_map = muon_param_info_map or {}
|
||
# Optional cap on the total size (in bytes) of a batched Muon group.
|
||
# When set, a group whose parameters sum to more than this many bytes
|
||
# is further split into sub-groups so each batched Newton-Schulz call
|
||
# stays under the limit, avoiding OOM. None disables splitting.
|
||
self._muon_max_group_bytes = muon_max_group_bytes
|
||
# 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)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Accumulator management
|
||
# ------------------------------------------------------------------
|
||
|
||
def _ensure_accumulators(self, param, use_muon, group):
|
||
"""Create optimizer accumulators for *param* if they do not exist yet."""
|
||
if (
|
||
self._moment_acc_str in self._accumulators
|
||
and param.name in self._accumulators[self._moment_acc_str]
|
||
):
|
||
return
|
||
|
||
# FP32 master weight for mixed-precision training
|
||
if self._multi_precision and self._is_dtype_fp16_or_bf16(param.dtype):
|
||
if param.name not in self._master_weights:
|
||
self._create_master_weight(param)
|
||
|
||
self._add_accumulator(
|
||
self._moment_acc_str,
|
||
param,
|
||
dtype=paddle.float32,
|
||
fill_value=0.0,
|
||
shape=param.shape,
|
||
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
|
||
)
|
||
|
||
if not use_muon:
|
||
# AdamW-specific states
|
||
self._add_accumulator(
|
||
self._moment2_acc_str,
|
||
param,
|
||
dtype=paddle.float32,
|
||
fill_value=0.0,
|
||
shape=param.shape,
|
||
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
|
||
)
|
||
for acc_name, init_val in [
|
||
(self._beta1_pow_acc_str, group.get("adam_beta1", 0.9)),
|
||
(self._beta2_pow_acc_str, group.get("adam_beta2", 0.95)),
|
||
]:
|
||
self._add_accumulator(
|
||
acc_name,
|
||
param,
|
||
dtype=paddle.float32,
|
||
fill_value=init_val,
|
||
shape=[1],
|
||
type=framework.core.VarDesc.VarType.DENSE_TENSOR,
|
||
)
|
||
|
||
def _create_accumulators(self, block, parameters):
|
||
"""Standard entry-point used by checkpoint-resume infrastructure."""
|
||
if isinstance(parameters, dict):
|
||
parameters = self._update_param_group(parameters)
|
||
for p in parameters:
|
||
param_info = self._muon_param_info_map.get(p.name)
|
||
if param_info is not None:
|
||
use_muon = param_info.use_muon
|
||
else:
|
||
use_muon = _default_should_use_muon(
|
||
p.name,
|
||
getattr(p, "original_shape", p.shape),
|
||
self._muon_exclude_patterns,
|
||
)
|
||
self._ensure_accumulators(p, use_muon, self._default_dict)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Newton-Schulz orthogonalisation
|
||
# ------------------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _zeropower_via_newtonschulz5(
|
||
X,
|
||
steps=5,
|
||
eps=1e-9,
|
||
ns_coeffs=None,
|
||
ns_matmul_dtype=paddle.bfloat16,
|
||
):
|
||
"""Approximate the matrix sign function via Newton-Schulz iteration.
|
||
|
||
Args:
|
||
X: Input tensor to orthogonalize. Must be 2D (M, N) or
|
||
3D (B, M, N) for batched operation.
|
||
steps: Number of Newton-Schulz iterations.
|
||
eps: Small constant for numerical stability.
|
||
ns_coeffs: List of (a, b, c) coefficient tuples for iteration.
|
||
If None, uses the "simple" preset.
|
||
ns_matmul_dtype: Dtype for matmul iterations. Defaults to
|
||
bfloat16. Pass paddle.float32 for V100 compatibility.
|
||
"""
|
||
if X.ndim < 2 or X.ndim > 3:
|
||
raise ValueError(
|
||
f"Input tensor X must be 2D or 3D (batched), got {X.ndim}D"
|
||
)
|
||
|
||
coeff_sets = (
|
||
ns_coeffs
|
||
if ns_coeffs is not None
|
||
else _NS_COEFFICIENT_SETS["simple"]
|
||
)
|
||
|
||
if X.shape[-2] > X.shape[-1]:
|
||
X = paddle.transpose(
|
||
X,
|
||
perm=[1, 0] if X.ndim == 2 else [0, 2, 1],
|
||
)
|
||
transpose = True
|
||
else:
|
||
transpose = False
|
||
|
||
orig_shape = X.shape
|
||
X_flat = X.flatten(start_axis=-2)
|
||
X_flat = paddle.nn.functional.normalize(
|
||
X_flat, p=2, axis=-1, epsilon=eps
|
||
)
|
||
X = X_flat.reshape(orig_shape).astype(ns_matmul_dtype)
|
||
|
||
if X.ndim == 3:
|
||
ns_step_fn = Muon._batched_newton_schulz_step
|
||
else:
|
||
ns_step_fn = Muon._newton_schulz_step
|
||
|
||
for i in range(steps):
|
||
a, b, c = coeff_sets[i % len(coeff_sets)]
|
||
X = ns_step_fn(X, a, b, c)
|
||
|
||
if transpose:
|
||
X = paddle.transpose(X, perm=[1, 0] if X.ndim == 2 else [0, 2, 1])
|
||
return X
|
||
|
||
@staticmethod
|
||
def _newton_schulz_step(X, a, b, c):
|
||
"""Single Newton-Schulz iteration step for 2D input."""
|
||
A = paddle.matmul(X, X, transpose_y=True)
|
||
B = paddle.addmm(input=A, x=A, y=A, beta=b, alpha=c)
|
||
X = paddle.addmm(input=X, x=B, y=X, beta=a, alpha=1.0)
|
||
return X
|
||
|
||
@staticmethod
|
||
def _batched_newton_schulz_step(X, a, b, c):
|
||
"""Single Newton-Schulz iteration step for 3D batched input."""
|
||
A = paddle.matmul(X, X, transpose_y=True)
|
||
B = paddle.baddbmm(A, A, A, beta=b, alpha=c)
|
||
X = paddle.baddbmm(X, B, X, beta=a, alpha=1.0)
|
||
return X
|
||
|
||
@staticmethod
|
||
def _scaling_fn(orthogonal_update, version, extra_scale_factor=1.0):
|
||
"""Apply dimension-dependent scaling to the orthogonal update."""
|
||
din, dout = orthogonal_update.shape[-2], orthogonal_update.shape[-1]
|
||
if version == 1:
|
||
scale = max(1, dout / din) ** 0.5
|
||
elif version == 2:
|
||
scale = (dout / din) ** 0.5
|
||
else: # version == 3 (default)
|
||
scale = max(dout, din) ** 0.5
|
||
return orthogonal_update * scale * extra_scale_factor
|
||
|
||
# ------------------------------------------------------------------
|
||
# Per-parameter update rules
|
||
# ------------------------------------------------------------------
|
||
|
||
def _adamw_update(
|
||
self,
|
||
param,
|
||
grad,
|
||
lr,
|
||
moment1,
|
||
moment2,
|
||
beta1_pow,
|
||
beta2_pow,
|
||
beta1,
|
||
beta2,
|
||
epsilon,
|
||
weight_decay,
|
||
):
|
||
"""In-place AdamW update for 1-D sharded parameters."""
|
||
|
||
lr_ratio = 1.0 if self._lr_ratio is None else self._lr_ratio(param)
|
||
with_decay = True
|
||
if (
|
||
self._apply_decay_param_fun is not None
|
||
and not self._apply_decay_param_fun(param.name)
|
||
):
|
||
with_decay = False
|
||
|
||
find_master = param.name in self._master_weights
|
||
master_weight = (
|
||
self._master_weights[param.name] if find_master else None
|
||
)
|
||
_, _, _, _, _, _, _ = _C_ops.adamw_(
|
||
param,
|
||
grad,
|
||
lr,
|
||
moment1,
|
||
moment2,
|
||
None, # moment2_max
|
||
beta1_pow,
|
||
beta2_pow,
|
||
master_weight,
|
||
None, # found_inf
|
||
beta1,
|
||
beta2,
|
||
epsilon,
|
||
lr_ratio,
|
||
weight_decay,
|
||
with_decay,
|
||
False, # lazy_mode
|
||
1000,
|
||
find_master,
|
||
False,
|
||
False, # amsgrad
|
||
)
|
||
|
||
def _split_group_by_bytes(self, group_params_grads):
|
||
"""Split a batched Muon group so each sub-group stays under the byte cap.
|
||
|
||
The batched Newton-Schulz update stacks every parameter's momentum
|
||
buffer (float32) into one tensor, so peak memory scales with the total
|
||
size of the group. When ``self._muon_max_group_bytes`` is set, greedily
|
||
pack parameters into sub-groups whose cumulative size does not exceed
|
||
the cap. A single parameter larger than the cap forms its own sub-group.
|
||
"""
|
||
max_bytes = self._muon_max_group_bytes
|
||
if not max_bytes:
|
||
yield group_params_grads
|
||
return
|
||
|
||
# Momentum buffer / batched matrix are float32 (see _ensure_accumulators).
|
||
elem_size = 4
|
||
sub_group = []
|
||
sub_bytes = 0
|
||
for param, grad in group_params_grads:
|
||
param_bytes = math.prod(param.shape) * elem_size
|
||
if sub_group and sub_bytes + param_bytes > max_bytes:
|
||
yield sub_group
|
||
sub_group = []
|
||
sub_bytes = 0
|
||
sub_group.append((param, grad))
|
||
sub_bytes += param_bytes
|
||
if sub_group:
|
||
yield sub_group
|
||
|
||
def _muon_update_group(
|
||
self,
|
||
group_params_grads,
|
||
lr,
|
||
momentum_beta,
|
||
ns_steps,
|
||
nesterov,
|
||
epsilon,
|
||
weight_decay,
|
||
version,
|
||
):
|
||
"""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
|
||
|
||
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():
|
||
# --- 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
|
||
)
|
||
|
||
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)
|
||
|
||
# --- 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]
|
||
|
||
def ortho_fn(m):
|
||
ns_out = Muon._zeropower_via_newtonschulz5(
|
||
m,
|
||
steps=ns_steps,
|
||
eps=epsilon,
|
||
ns_coeffs=self._ns_coeffs,
|
||
ns_matmul_dtype=self._ns_matmul_dtype,
|
||
)
|
||
scaled = Muon._scaling_fn(
|
||
ns_out, version, self._muon_extra_scale_factor
|
||
)
|
||
return scaled
|
||
|
||
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 = 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: "
|
||
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:
|
||
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
|
||
|
||
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
|
||
|
||
if find_master:
|
||
master_weight.subtract_(final_step)
|
||
paddle.assign(master_weight.astype(param.dtype), param)
|
||
else:
|
||
param.subtract_(final_step.astype(param.dtype))
|
||
|
||
# ------------------------------------------------------------------
|
||
# Core optimization step
|
||
# ------------------------------------------------------------------
|
||
|
||
def _apply_optimize(self, loss, startup_program, params_grads):
|
||
if not framework.in_dygraph_mode():
|
||
raise NotImplementedError(
|
||
"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)
|
||
|
||
# apply for zcc
|
||
self._maybe_refuse()
|
||
|
||
group = self._default_dict
|
||
lr = self._learning_rate
|
||
if isinstance(lr, paddle.optimizer.lr.LRScheduler):
|
||
lr = lr()
|
||
wd = group.get("weight_decay", 0.0)
|
||
|
||
muon_params = []
|
||
adamw_params = []
|
||
for param, grad in params_grads:
|
||
if grad is None:
|
||
continue
|
||
|
||
param_info = self._muon_param_info_map.get(param.name)
|
||
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 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:
|
||
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():
|
||
for sub_group in self._split_group_by_bytes(group_params):
|
||
self._muon_update_group(
|
||
sub_group,
|
||
lr_tensor,
|
||
group.get("momentum", 0.95),
|
||
group.get("ns_steps", 5),
|
||
group.get("nesterov", True),
|
||
group.get("epsilon", 1e-9),
|
||
wd,
|
||
version=group.get("muon_version", 3),
|
||
)
|
||
|
||
# --- Pass 3: AdamW updates ---
|
||
for param, grad in adamw_params:
|
||
self._adamw_update(
|
||
param,
|
||
grad,
|
||
lr_tensor_f64,
|
||
self._get_accumulator(self._moment_acc_str, param),
|
||
self._get_accumulator(self._moment2_acc_str, param),
|
||
self._get_accumulator(self._beta1_pow_acc_str, param),
|
||
self._get_accumulator(self._beta2_pow_acc_str, param),
|
||
group.get("adam_beta1", 0.9),
|
||
group.get("adam_beta2", 0.95),
|
||
group.get("epsilon", 1e-9),
|
||
wd,
|
||
)
|
||
|
||
@framework.dygraph_only
|
||
def step(self) -> None:
|
||
params_grads = [
|
||
(param, param._grad_ivar())
|
||
for param in self._parameter_list
|
||
if not param.stop_gradient and param._grad_ivar() is not None
|
||
]
|
||
self._apply_optimize(
|
||
loss=None, startup_program=None, params_grads=params_grads
|
||
)
|
||
|
||
def sharded_state_dict(
|
||
self,
|
||
model_sharded_state_dict: ShardedStateDict,
|
||
) -> ShardedStateDict:
|
||
"""Build a sharded optimizer state dict for flex checkpoint save/load.
|
||
|
||
The layout mirrors :class:`paddle.optimizer.AdamW`'s implementation so
|
||
that the same ``dist.save_state_dict`` / ``dist.load_state_dict`` path
|
||
works for Muon checkpoints.
|
||
|
||
Args:
|
||
model_sharded_state_dict: Sharded model state dict produced by
|
||
``model.sharded_state_dict()``.
|
||
|
||
Returns:
|
||
A dict mapping ``"<struct_name>.<state_type>"`` keys to
|
||
:class:`ShardedWeight` objects.
|
||
"""
|
||
_FP32_MASTER = "fp32_master_0"
|
||
_optimizer_scalar_names = [
|
||
"beta1_pow_acc_0",
|
||
"beta2_pow_acc_0",
|
||
]
|
||
_optimizer_vector_names = [
|
||
"moment1_0",
|
||
"moment2_0",
|
||
]
|
||
|
||
def _split_state_name(vname):
|
||
if _FP32_MASTER in vname:
|
||
return tuple(vname.split("_" + _FP32_MASTER + "_", 1))
|
||
for suffix in _optimizer_scalar_names + _optimizer_vector_names:
|
||
if vname.endswith(suffix):
|
||
return vname[: -(len(suffix) + 1)], suffix
|
||
raise ValueError(
|
||
f"Cannot parse optimizer state variable name: {vname!r}"
|
||
)
|
||
|
||
model_sharded_state_dict = dict(
|
||
sorted(model_sharded_state_dict.items())
|
||
)
|
||
|
||
# Build static-name → struct-name mapping (handles shared weights)
|
||
static_to_struct = {}
|
||
for struct_name, sw in model_sharded_state_dict.items():
|
||
local_name = sw.local_tensor.name
|
||
if local_name not in static_to_struct:
|
||
static_to_struct[local_name] = struct_name
|
||
|
||
optimizer_state_dict = self.state_dict()
|
||
master_weights = optimizer_state_dict.pop("master_weights", None)
|
||
optimizer_state_dict.pop("LR_Scheduler", None)
|
||
|
||
sharded_state: ShardedStateDict = {}
|
||
|
||
# Optimizer states (moment1, moment2, beta_pow scalars)
|
||
for key, tensor in optimizer_state_dict.items():
|
||
static_name, state_type = _split_state_name(key)
|
||
struct_name = static_to_struct[static_name]
|
||
sharded_param = model_sharded_state_dict[struct_name]
|
||
unified_name = f"{struct_name}.{state_type}"
|
||
|
||
if state_type in _optimizer_vector_names:
|
||
# Vector states share the same sharding layout as the parameter
|
||
if tensor.is_dist():
|
||
sharded_state[unified_name] = ShardedWeight(
|
||
key=unified_name,
|
||
local_tensor=tensor,
|
||
local_shape=tensor.shape,
|
||
global_shape=tensor.shape,
|
||
global_offset=sharded_param.global_offset,
|
||
)
|
||
else:
|
||
# Reshape accumulator if numel matches but shape differs.
|
||
# MoE: grouped_gemm_experts param.shape is 3D
|
||
# [n_experts, H, I] but model.state_dict() returns actual
|
||
# C++ storage shape 2D [n_experts*H, I]. moment1 was
|
||
# created with 3D shape, so we need to reshape here.
|
||
# V2 is unaffected: its moments are always 1D shards,
|
||
# so shape always matches and reshape is never triggered.
|
||
target_shape = sharded_param.local_shape
|
||
if (
|
||
tuple(tensor.shape) != tuple(target_shape)
|
||
and tensor.numel()
|
||
== paddle.to_tensor(list(target_shape)).prod().item()
|
||
):
|
||
tensor = tensor.reshape(target_shape)
|
||
sharded_state[unified_name] = (
|
||
create_sharded_weight_with_new_local(
|
||
unified_name, tensor, sharded_param
|
||
)
|
||
)
|
||
else:
|
||
# Scalar states (beta_pow) are replicated – save as-is
|
||
sharded_state[unified_name] = ShardedWeight(
|
||
key=unified_name,
|
||
local_tensor=tensor,
|
||
local_shape=(1,),
|
||
global_shape=(1,),
|
||
global_offset=(0,),
|
||
)
|
||
|
||
# FP32 master weights
|
||
if master_weights:
|
||
for weight_key, tensor in master_weights.items():
|
||
struct_name = static_to_struct[weight_key]
|
||
sharded_param = model_sharded_state_dict[struct_name]
|
||
unified_name = f"{struct_name}.w_0"
|
||
|
||
if tensor.is_dist():
|
||
sharded_state[unified_name] = ShardedWeight(
|
||
key=unified_name,
|
||
local_tensor=tensor,
|
||
local_shape=tensor.shape,
|
||
global_shape=tensor.shape,
|
||
global_offset=sharded_param.global_offset,
|
||
)
|
||
else:
|
||
sharded_state[unified_name] = (
|
||
create_sharded_weight_with_new_local(
|
||
unified_name, tensor, sharded_param
|
||
)
|
||
)
|
||
|
||
return sharded_state
|