Skip to content

Enable optional batched Muon NS iteration #79495

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xxyux wants to merge 1 commit into
base: release/3.4
Choose a base branch
from
+210 −98
Open

Enable optional batched Muon NS iteration #79495

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 159 additions & 77 deletions python/paddle/optimizer/muon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

if TYPE_CHECKING:
from collections.abc import Callable
from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING

Expand All @@ -40,6 +41,9 @@
# Debug logging for Muon optimizer
_logger = logging.getLogger(__name__)
MUON_DEBUG = os.environ.get("MUON_DEBUG", "0") == "1"
g_shard_bypass_dygraph_optimizer = int(
os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0)
)

__all__ = []

Expand Down Expand Up @@ -150,6 +154,15 @@ def _default_should_use_muon(name, shape, exclude_patterns):
return True


def _to_hashable(value):
"""Recursively convert lists/tuples/dicts into hashable tuples."""
if isinstance(value, (list, tuple)):
return tuple(_to_hashable(v) for v in value)
if isinstance(value, dict):
return tuple((k, _to_hashable(v)) for k, v in sorted(value.items()))
return value


class Muon(Optimizer):
r"""
Muon optimizer for MuonShardingOptimizer (Sharding Stage1 V3) usage.
Expand Down Expand Up @@ -519,52 +532,64 @@ def _adamw_update(
False, # amsgrad
)

def _muon_update(
def _muon_update_group(
self,
param,
grad,
group_params_grads,
lr,
momentum_buffer,
momentum_beta,
ns_steps,
nesterov,
epsilon,
weight_decay,
version,
):
"""In-place Muon update for a 2D parameter tensor.
"""Batched Muon update for a group of parameters with identical shape and split_concat_func."""
# Because shape is identical, use the first param's properties
param0, _ = group_params_grads[0]
param_shape = getattr(param0, "original_shape", param0.shape)
is_3d = len(param_shape) == 3

Applies Newton-Schulz orthogonalisation to the 2D weight matrix and
updates the parameter in-place. MuonShardingOptimizer assigns whole
2D tensors to ranks, so no sharding gather or TP communication is needed.
"""
param_shape = getattr(param, "original_shape", param.shape)
param_info = self._muon_param_info_map.get(param.name)
param_info = self._muon_param_info_map.get(param0.name)
split_concat_func = param_info.split_concat_func if param_info else None

matrix_2d_list = []
find_master_list = []

with paddle.no_grad():
grad_f32 = (
grad.astype(momentum_buffer.dtype)
if grad.dtype != momentum_buffer.dtype
else grad
)
# --- Pass 1: Sequential momentum update & preparation ---
for param, grad in group_params_grads:
momentum_buffer = self._get_accumulator(
self._moment_acc_str, param
)
grad_f32 = (
grad.astype(momentum_buffer.dtype)
if grad.dtype != momentum_buffer.dtype
else grad
)

# Step 1: Momentum update
new_momentum = paddle.lerp(
momentum_buffer, grad_f32, 1.0 - momentum_beta
)
paddle.assign(new_momentum, momentum_buffer)
update_buffer = (
paddle.lerp(grad_f32, momentum_buffer, momentum_beta)
if nesterov
else momentum_buffer
)
new_momentum = paddle.lerp(
momentum_buffer, grad_f32, 1.0 - momentum_beta
)
paddle.assign(new_momentum, momentum_buffer)
update_buffer = (
paddle.lerp(grad_f32, momentum_buffer, momentum_beta)
if nesterov
else momentum_buffer
)

matrix_2d_global = update_buffer.reshape(param_shape)
matrix_2d_list.append(matrix_2d_global)
find_master_list.append(param.name in self._master_weights)

# Step 2: Reshape update buffer to 2D matrix.
# MuonShardingOptimizer assigns whole 2D tensors to ranks, so params
# are already 2D/3D (no sharding gather needed).
matrix_2d_global = update_buffer.reshape(param_shape)
# --- Pass 2: Batched Newton-Schulz iteration ---
if len(matrix_2d_list) > 1:
if is_3d:
batched_matrix = paddle.concat(matrix_2d_list, axis=0)
else:
batched_matrix = paddle.stack(matrix_2d_list, axis=0)
else:
batched_matrix = matrix_2d_list[0]

# Shared NS + scaling closure (captures ns_steps, epsilon, version, ns_coeffs)
def ortho_fn(m):
ns_out = Muon._zeropower_via_newtonschulz5(
m,
Expand All @@ -578,53 +603,69 @@ def ortho_fn(m):
)
return scaled

# Step 3: Newton-Schulz orthogonalisation
# Use split_concat_func from param_info if provided, otherwise default to whole matrix
if (
param_info is not None
and param_info.split_concat_func is not None
):
# Use slice function defined in model configuration
orthogonal_update = param_info.split_concat_func(
matrix_2d_global, ortho_fn
if split_concat_func is not None:
batched_orthogonal_update = split_concat_func(
batched_matrix, ortho_fn
)
if MUON_DEBUG:
_global_rank = paddle.distributed.get_rank()
if _global_rank == 0:
_sf = param_info.split_concat_func
_sf = split_concat_func
# split_concat_func may be a plain function or a
# functools.partial; unwrap for readable logging.
_fn = getattr(_sf, "func", _sf)
_logger.info(
f"[Muon] Using split_concat_func: param={param.name}, "
f"split_concat_func={_sf.func.__name__}, "
f"args={_sf.args}, kwargs={_sf.keywords}"
f"[Muon] Using split_concat_func: "
f"params={[p.name for p, _ in group_params_grads]}, "
f"split_concat_func={_fn.__name__}, "
f"args={getattr(_sf, 'args', ())}, "
f"kwargs={getattr(_sf, 'keywords', None)}"
)
else:
# Default: whole matrix orthogonalisation
orthogonal_update = ortho_fn(matrix_2d_global)
batched_orthogonal_update = ortho_fn(batched_matrix)

# Unbatch results
if len(matrix_2d_list) > 1:
if is_3d:
orthogonal_updates = paddle.split(
batched_orthogonal_update,
num_or_sections=len(matrix_2d_list),
axis=0,
)
else:
orthogonal_updates = paddle.unbind(
batched_orthogonal_update, axis=0
)
else:
orthogonal_updates = [batched_orthogonal_update]

# --- Pass 3: Sequential apply weight decay & step ---
for i, (param, _) in enumerate(group_params_grads):
orthogonal_update = orthogonal_updates[i]
find_master = find_master_list[i]
master_weight = (
self._master_weights[param.name] if find_master else None
)

find_master = param.name in self._master_weights
master_weight = (
self._master_weights[param.name] if find_master else None
)
with_decay = True
if (
self._apply_decay_param_fun is not None
and not self._apply_decay_param_fun(param.name)
):
with_decay = False
if with_decay and weight_decay > 0:
if find_master:
master_weight.scale_(1.0 - lr * weight_decay)
else:
param.scale_(1.0 - lr * weight_decay)

final_step = orthogonal_update * lr

with_decay = True
if (
self._apply_decay_param_fun is not None
and not self._apply_decay_param_fun(param.name)
):
with_decay = False
if with_decay and weight_decay > 0:
if find_master:
master_weight.scale_(1.0 - lr * weight_decay)
master_weight.subtract_(final_step)
paddle.assign(master_weight.astype(param.dtype), param)
else:
param.scale_(1.0 - lr * weight_decay)

final_step = orthogonal_update * lr

if find_master:
master_weight.subtract_(final_step)
paddle.assign(master_weight.astype(param.dtype), param)
else:
param.subtract_(final_step.astype(param.dtype))
param.subtract_(final_step.astype(param.dtype))

# ------------------------------------------------------------------
# Core optimization step
Expand All @@ -636,6 +677,10 @@ def _apply_optimize(self, loss, startup_program, params_grads):
"Muon optimizer only supports dygraph mode."
)

# Same bypass as the base Optimizer._apply_optimize: skip all parameter updates.
if g_shard_bypass_dygraph_optimizer:
return

if self._grad_clip is not None:
params_grads = self._grad_clip(params_grads)

Expand All @@ -655,26 +700,63 @@ def _apply_optimize(self, loss, startup_program, params_grads):
continue

param_info = self._muon_param_info_map.get(param.name)
assert param_info is not None, (
f"muon_param_info_map does not have {param.name}"
)
use_muon = param_info.use_muon
if param_info is not None:
use_muon = param_info.use_muon
else:
use_muon = _default_should_use_muon(
param.name,
getattr(param, "original_shape", param.shape),
self._muon_exclude_patterns,
)
_logger.warning(
f"muon_param_info_map does not have {param.name}, "
f"falling back to default rule: use_muon={use_muon}"
)

self._ensure_accumulators(param, use_muon, group)
if use_muon:
muon_params.append((param, grad))
else:
adamw_params.append((param, grad))

# --- Pass 1: Muon updates (large temporary tensors) ---
# --- Pass 2: Muon updates ---
lr_tensor = paddle.to_tensor(lr, dtype=paddle.float32)
lr_tensor_f64 = paddle.to_tensor(lr, dtype=paddle.float64)

# Group parameters by split_concat_func and shape, then update
# each group in one batched call.
muon_groups = defaultdict(list)
for param, grad in muon_params:
self._muon_update(
param,
grad,
assert len(param.shape) == 2 or len(param.shape) == 3, (
"Muon only supports 2D or 3D parameters."
)
param_shape = getattr(param, "original_shape", param.shape)
param_info = self._muon_param_info_map.get(param.name)

func_key = None
if param_info and param_info.split_concat_func is not None:
func = param_info.split_concat_func
if hasattr(func, "func"):
# functools.partial: group by the underlying function
# object plus all bound args/kwargs.
func_key = (
func.func,
tuple(_to_hashable(v) for v in func.args),
tuple(
(k, _to_hashable(v))
for k, v in sorted((func.keywords or {}).items())
),
)
else:
func_key = func

key = (func_key, tuple(param_shape))
muon_groups[key].append((param, grad))

for key, group_params in muon_groups.items():
self._muon_update_group(
group_params,
lr_tensor,
self._get_accumulator(self._moment_acc_str, param),
group.get("momentum", 0.95),
group.get("ns_steps", 5),
group.get("nesterov", True),
Expand All @@ -683,7 +765,7 @@ def _apply_optimize(self, loss, startup_program, params_grads):
version=group.get("muon_version", 3),
)

# --- Pass 2: AdamW updates ---
# --- Pass 3: AdamW updates ---
for param, grad in adamw_params:
self._adamw_update(
param,
Expand Down
Loading
Loading