Skip to content

Update init group for adaptive muon #132

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

Merged
FDecaYed merged 6 commits into from
Mar 18, 2026
Merged

Update init group for adaptive muon #132

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c62fddf
add _init_group to adaptive muon
skyw Mar 17, 2026
33782da
add experimental warning
skyw Mar 17, 2026
9e57870
Merge branch 'main' into skyw/update_init_group_for_adaptive_muon
skyw Mar 18, 2026
69627ba
Update emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py
skyw Mar 18, 2026
4dee92b
fix AI format error
skyw Mar 18, 2026
e06ba85
fix AI format error
skyw Mar 18, 2026
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
69 changes: 41 additions & 28 deletions emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,9 @@ class AdaptiveMuon(muon.Muon):
descent for deep learning.* In Advances in neural information processing systems 28 (2015).
The step() method is overridden to include second moment normalization logic.

Warning:
This optimizer is experimental and may change in future versions.

Args:
params: Iterable of parameters to optimize or dicts defining parameter groups.
lr: Learning rate.
Expand Down Expand Up @@ -98,37 +101,49 @@ def __init__(
group.setdefault("beta2", beta2)
group.setdefault("eps", eps)

def _initialize_moment2(
@torch.no_grad() # type: ignore[misc]
@override
def _init_group(
self,
state: dict[str, torch.Tensor],
grad: torch.Tensor,
group: dict,
skip_non_grad_params: bool = True,
) -> None:
"""Initialize the second moment buffer if it doesn't exist.
"""Performs lazy state initialization for parameters.

The shape of the buffer depends on the moment2_method:
- "adamuon": Full elementwise buffer with same shape as grad
Extends the base class to also initialize the second moment buffer.
The shape of the moment2 buffer depends on the moment2_method:
- "adamuon": Full elementwise buffer with same shape as parameter
- "normuon": Reduced shape buffer (averaged along -1 if shape[-2] >= shape[-1], else -2)

Args:
state: The optimizer state dict for a parameter.
grad: The gradient tensor (used for shape/dtype).
group: Parameter group dictionary.
skip_non_grad_params: If True, skip parameters without gradients.
"""
if "moment2_buffer" not in state:
if self.moment2_method == "adamuon":
# Full elementwise second moment
moment2 = torch.zeros_like(grad)
elif self.moment2_method == "normuon":
# Row/column-wise second moment - reduced along one dimension
# Determine which dimension to reduce based on parameter shape
avg_dim = -1 if grad.shape[-2] >= grad.shape[-1] else -2
# Specify the shape with reduced dimension
moment2_shape = list(grad.shape)
moment2_shape[avg_dim] = 1
moment2 = torch.zeros(moment2_shape, dtype=grad.dtype, device=grad.device)
else:
raise TypeError(f"Invalid second moment method: {self.moment2_method}")

state["moment2_buffer"] = moment2
for p in group["params"]:
if skip_non_grad_params and p.grad is None:
continue
state = self.state[p]

if len(state) == 0:
state["momentum_buffer"] = torch.zeros_like(p.data)

if self.moment2_method == "adamuon":
# Full elementwise second moment
state["moment2_buffer"] = torch.zeros_like(p.data)
elif self.moment2_method == "normuon":
# Row/column-wise second moment - reduced along one dimension
# Determine which dimension to reduce based on parameter shape
if p.data.ndim != 2:
raise ValueError(
f"{self.__class__.__name__} only supports 2D parameters, got shape {tuple(p.data.shape)}"
)
avg_dim = -1 if p.data.shape[-2] >= p.data.shape[-1] else -2
# Specify the shape with reduced dimension
moment2_shape = list(p.data.shape)
moment2_shape[avg_dim] = 1
state["moment2_buffer"] = torch.zeros(moment2_shape, dtype=p.data.dtype, device=p.data.device)
Comment thread
skyw marked this conversation as resolved.
else:
raise TypeError(f"Invalid second moment method: {self.moment2_method}")
Comment thread
skyw marked this conversation as resolved.

def _apply_moment2_normalization(
self,
Expand Down Expand Up @@ -204,6 +219,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
loss = closure()

for group in self.param_groups:
self._init_group(group)

for p in group["params"]:
if p.dim() != 2:
raise ValueError(f"{self.__class__.__name__} only supports 2D parameters")
Expand All @@ -212,10 +229,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
continue
state = self.state[p]

if "momentum_buffer" not in state:
state["momentum_buffer"] = torch.zeros_like(grad)
self._initialize_moment2(state, grad)

exp_avg = state["momentum_buffer"]

self._apply_weight_decay_inplace(
Expand Down
Loading