chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,837 @@
|
||||
# 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 os
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
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"
|
||||
|
||||
__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
|
||||
|
||||
|
||||
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,
|
||||
**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 {}
|
||||
# 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 _muon_update(
|
||||
self,
|
||||
param,
|
||||
grad,
|
||||
lr,
|
||||
momentum_buffer,
|
||||
momentum_beta,
|
||||
ns_steps,
|
||||
nesterov,
|
||||
epsilon,
|
||||
weight_decay,
|
||||
version,
|
||||
):
|
||||
"""In-place Muon update for a 2D parameter tensor.
|
||||
|
||||
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)
|
||||
|
||||
with paddle.no_grad():
|
||||
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
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
# Shared NS + scaling closure (captures ns_steps, epsilon, version, ns_coeffs)
|
||||
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
|
||||
|
||||
# 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 MUON_DEBUG:
|
||||
_global_rank = paddle.distributed.get_rank()
|
||||
if _global_rank == 0:
|
||||
_sf = param_info.split_concat_func
|
||||
_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}"
|
||||
)
|
||||
else:
|
||||
# Default: whole matrix orthogonalisation
|
||||
orthogonal_update = ortho_fn(matrix_2d_global)
|
||||
|
||||
find_master = param.name in self._master_weights
|
||||
master_weight = (
|
||||
self._master_weights[param.name] if find_master else None
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
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)
|
||||
assert param_info is not None, (
|
||||
f"muon_param_info_map does not have {param.name}"
|
||||
)
|
||||
use_muon = param_info.use_muon
|
||||
|
||||
self._ensure_accumulators(param, use_muon, group)
|
||||
if use_muon:
|
||||
muon_params.append((param, grad))
|
||||
else:
|
||||
adamw_params.append((param, grad))
|
||||
|
||||
# --- Pass 1: Muon updates (large temporary tensors) ---
|
||||
lr_tensor = paddle.to_tensor(lr, dtype=paddle.float32)
|
||||
lr_tensor_f64 = paddle.to_tensor(lr, dtype=paddle.float64)
|
||||
for param, grad in muon_params:
|
||||
self._muon_update(
|
||||
param,
|
||||
grad,
|
||||
lr_tensor,
|
||||
self._get_accumulator(self._moment_acc_str, param),
|
||||
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 2: 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
|
||||
Reference in New Issue
Block a user