Skip to content

Commit a0e376b

Browse files
mkhona-nvidiamkhona-nvidia
authored
MOSO: Combining SOAP and Muon (#222)
* Initial Implementation of ShMuon: SOAPified Muon Signed-off-by: mikail <mkhona@nvidia.com> * changed to adam in the shmuon eigenbasis Signed-off-by: mikail <mkhona@nvidia.com> * fix state init dtype Signed-off-by: mikail <mkhona@nvidia.com> * renamed shmuon to moso Signed-off-by: mikail <mkhona@nvidia.com> * Addressed MR comments Signed-off-by: mikail <mkhona@nvidia.com> * removed extra args from moso Signed-off-by: mikail <mkhona@nvidia.com> * remove muon scale factor since moso already has the adamW update Signed-off-by: mikail <mkhona@nvidia.com> * removed code for removed flags Signed-off-by: mikail <mkhona@nvidia.com> --------- Signed-off-by: mikail <mkhona@nvidia.com>
1 parent 7e20f92 commit a0e376b

3 files changed

Lines changed: 441 additions & 0 deletions

File tree

‎emerging_optimizers/soap/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,13 @@
1212
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
15+
from emerging_optimizers.soap.moso import MOSO
1516
from emerging_optimizers.soap.rekls import REKLS
1617
from emerging_optimizers.soap.soap import SOAP
1718

1819

1920
__all__ = [
21+
"MOSO",
2022
"REKLS",
2123
"SOAP",
2224
]

‎emerging_optimizers/soap/moso.py‎

Lines changed: 297 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,297 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
from typing import TYPE_CHECKING, Callable, override
16+
17+
18+
if TYPE_CHECKING:
19+
from typing import overload
20+
21+
import torch
22+
from torch import optim
23+
from torch.optim.optimizer import ParamsT
24+
25+
from emerging_optimizers import mixin as opt_mixin
26+
from emerging_optimizers import registry, utils
27+
from emerging_optimizers.scalar_optimizers import update_functions
28+
from emerging_optimizers.soap import soap_utils
29+
from emerging_optimizers.soap.soap import _clip_update_rms_in_place
30+
31+
32+
__all__ = ["MOSO"]
33+
34+
35+
@registry.register_optimizer("moso")
36+
class MOSO(opt_mixin.WeightDecayMixin, optim.Optimizer):
37+
r"""Momentum One-Sided SOAP.
38+
39+
MOSO tracks EMA momentum like Muon, accumulates a SOAP/Shampoo-style covariance of that momentum on the
40+
smaller matrix side, and applies an Adam update in the covariance eigenbasis.
41+
Conceptually, this is one-sided SOAP where ``G_t G_t^T`` is replaced by ``M_t M_t^T`` (or ``M_t^T M_t`` for the
42+
right side), and the update is computed by projecting the momentum into the eigenbasis, applying Adam there, and
43+
projecting back:
44+
45+
.. math::
46+
47+
C_t = \beta_s C_{t-1} + (1 - \beta_s) M_t M_t^T,\quad C_t = Q_M \Lambda_M Q_M^T
48+
49+
U_t = Q_M \operatorname{Adam}(Q_M^T M_t)
50+
51+
for the left-preconditioned case where ``M_t.shape[0] <= M_t.shape[1]``; the right-preconditioned case uses
52+
``C_t = M_t^T M_t`` and computes ``U_t = \operatorname{Adam}(M_t Q_M) Q_M^T``.
53+
54+
Args:
55+
params: Iterable of parameters to optimize or dicts defining parameter groups.
56+
lr: Learning rate.
57+
momentum: EMA coefficient for the Muon-style momentum.
58+
betas: Inner Adam beta parameters ``(beta1, beta2)``.
59+
shampoo_beta: EMA coefficient for the one-sided momentum covariance.
60+
eps: Inner Adam epsilon for numerical stability.
61+
weight_decay: Weight decay coefficient.
62+
max_update_rms: Clip the update RMS to this value (0 means no clipping).
63+
"""
64+
65+
def __init__(
66+
self,
67+
params: ParamsT,
68+
lr: float = 3e-4,
69+
momentum: float = 0.95,
70+
betas: tuple[float, float] = (0.9, 0.95),
71+
shampoo_beta: float = 0.95,
72+
eps: float = 1e-8,
73+
weight_decay: float = 0.01,
74+
*,
75+
max_update_rms: float = 0.0,
76+
) -> None:
77+
self.weight_decay_method = "decoupled"
78+
self.max_update_rms = max_update_rms
79+
80+
defaults = {
81+
"lr": lr,
82+
"momentum": momentum,
83+
"betas": betas,
84+
"shampoo_beta": shampoo_beta,
85+
"eps": eps,
86+
"weight_decay": weight_decay,
87+
}
88+
super().__init__(params, defaults)
89+
90+
@torch.no_grad() # type: ignore[misc]
91+
def _init_group(
92+
self,
93+
group: dict,
94+
skip_non_grad_params: bool = True,
95+
) -> None:
96+
"""Performs lazy state initialization for 2D parameters."""
97+
for p in group["params"]:
98+
if skip_non_grad_params and p.grad is None:
99+
continue
100+
101+
if p.dim() != 2:
102+
raise TypeError("MOSO is only supported for 2D tensors")
103+
104+
state = self.state[p]
105+
if len(state) == 0:
106+
rows, cols = p.shape
107+
preconditioner_size = min(rows, cols)
108+
state["step"] = 0
109+
state["momentum_buffer"] = torch.zeros_like(p.data, dtype=torch.float32)
110+
state["exp_avg"] = torch.zeros_like(p.data, dtype=torch.float32)
111+
state["exp_avg_sq"] = torch.zeros_like(p.data, dtype=torch.float32)
112+
state["M"] = torch.zeros(
113+
preconditioner_size,
114+
preconditioner_size,
115+
device=p.device,
116+
dtype=torch.float32,
117+
)
118+
state["Q_M"] = torch.eye(preconditioner_size, device=p.device, dtype=torch.float32)
119+
120+
if TYPE_CHECKING:
121+
122+
@overload
123+
def step(self, closure: None = ...) -> None: ...
124+
125+
@overload
126+
def step(self, closure: Callable[[], float]) -> float: ...
127+
128+
@torch.no_grad() # type: ignore[misc]
129+
@override
130+
def step(self, closure: Callable[[], float] | None = None) -> float | None:
131+
"""Performs a single optimization step."""
132+
if closure is not None:
133+
raise ValueError("closure is not supported")
134+
135+
for group in self.param_groups:
136+
self._init_group(group)
137+
138+
for p in group["params"]:
139+
if p.grad is None:
140+
continue # pragma: no cover
141+
142+
grad = p.grad.to(torch.float32)
143+
state = self.state[p]
144+
curr_iter_1_based = state["step"] + 1
145+
146+
self._apply_weight_decay_inplace(
147+
p,
148+
grad,
149+
group["lr"],
150+
group["weight_decay"],
151+
)
152+
153+
state["momentum_buffer"].lerp_(grad, 1 - group["momentum"])
154+
momentum = state["momentum_buffer"]
155+
156+
shampoo_beta = 1 - (1 - group["shampoo_beta"]) / (1 - group["shampoo_beta"] ** curr_iter_1_based)
157+
158+
with utils.fp32_matmul_precision("highest"):
159+
_update_one_sided_momentum_factor(
160+
momentum_factor=state["M"],
161+
momentum=momentum,
162+
shampoo_beta=shampoo_beta,
163+
)
164+
165+
left_preconditioned = momentum.shape[0] <= momentum.shape[1]
166+
with utils.fp32_matmul_precision("highest"):
167+
state["Q_M"], state["exp_avg"], state["exp_avg_sq"] = _update_eigenbasis_and_adam_exp_avgs(
168+
momentum_factor=state["M"],
169+
eigenbasis=state["Q_M"],
170+
exp_avg=state["exp_avg"],
171+
exp_avg_sq=state["exp_avg_sq"],
172+
left_preconditioned=left_preconditioned,
173+
use_eigh=state["step"] == 0,
174+
power_iter_steps=1,
175+
)
176+
177+
with utils.fp32_matmul_precision("highest"):
178+
momentum_projected = _project_to_one_sided_eigenbasis(
179+
x=momentum,
180+
eigenbasis=state["Q_M"],
181+
left_preconditioned=left_preconditioned,
182+
)
183+
adam_update = update_functions.calculate_adam_update(
184+
momentum_projected,
185+
state["exp_avg"],
186+
state["exp_avg_sq"],
187+
betas=group["betas"],
188+
eps=group["eps"],
189+
correct_bias=True,
190+
nesterov=False,
191+
step=curr_iter_1_based,
192+
)
193+
update = _project_from_one_sided_eigenbasis(
194+
x=adam_update,
195+
eigenbasis=state["Q_M"],
196+
left_preconditioned=left_preconditioned,
197+
)
198+
199+
_clip_update_rms_in_place(update, self.max_update_rms)
200+
p.add_(update, alpha=-group["lr"])
201+
202+
state["step"] += 1
203+
204+
return None
205+
206+
207+
@torch.no_grad() # type: ignore[misc]
208+
def _update_one_sided_momentum_factor(
209+
momentum_factor: torch.Tensor,
210+
momentum: torch.Tensor,
211+
shampoo_beta: float,
212+
) -> None:
213+
"""Update the smaller-side covariance of the Muon momentum."""
214+
left_preconditioned = momentum.shape[0] <= momentum.shape[1]
215+
maybe_transposed_momentum = momentum if left_preconditioned else momentum.T
216+
momentum_factor.lerp_(maybe_transposed_momentum @ maybe_transposed_momentum.T, 1 - shampoo_beta)
217+
218+
219+
@torch.no_grad() # type: ignore[misc]
220+
def _update_eigenbasis_and_adam_exp_avgs(
221+
momentum_factor: torch.Tensor,
222+
eigenbasis: torch.Tensor,
223+
exp_avg: torch.Tensor,
224+
exp_avg_sq: torch.Tensor,
225+
left_preconditioned: bool,
226+
*,
227+
use_eigh: bool,
228+
power_iter_steps: int,
229+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
230+
"""Update one eigenbasis and keep Adam state aligned with that basis."""
231+
exp_avg = _project_from_one_sided_eigenbasis(
232+
x=exp_avg,
233+
eigenbasis=eigenbasis,
234+
left_preconditioned=left_preconditioned,
235+
)
236+
237+
eigenbasis, exp_avg_sq = _sort_one_sided_eigenbasis_and_exp_avg_sq(
238+
momentum_factor=momentum_factor,
239+
eigenbasis=eigenbasis,
240+
exp_avg_sq=exp_avg_sq,
241+
left_preconditioned=left_preconditioned,
242+
)
243+
244+
if use_eigh:
245+
(updated_eigenbasis,) = soap_utils.get_eigenbasis_eigh([momentum_factor])
246+
else:
247+
(updated_eigenbasis,) = soap_utils.get_eigenbasis_qr(
248+
[momentum_factor],
249+
[eigenbasis],
250+
power_iter_steps=power_iter_steps,
251+
)
252+
253+
exp_avg = _project_to_one_sided_eigenbasis(
254+
x=exp_avg,
255+
eigenbasis=updated_eigenbasis,
256+
left_preconditioned=left_preconditioned,
257+
)
258+
return updated_eigenbasis, exp_avg, exp_avg_sq
259+
260+
261+
@torch.no_grad() # type: ignore[misc]
262+
def _sort_one_sided_eigenbasis_and_exp_avg_sq(
263+
momentum_factor: torch.Tensor,
264+
eigenbasis: torch.Tensor,
265+
exp_avg_sq: torch.Tensor,
266+
left_preconditioned: bool,
267+
) -> tuple[torch.Tensor, torch.Tensor]:
268+
"""Sort eigenbasis slots by approximate eigenvalue and permute Adam second moments."""
269+
approx_eigvals = utils.eig.conjugate(momentum_factor, eigenbasis, diag=True)
270+
sort_idx = torch.argsort(approx_eigvals, descending=True, stable=True)
271+
sorted_eigenbasis = eigenbasis[:, sort_idx]
272+
exp_avg_sq_dim = 0 if left_preconditioned else 1
273+
return sorted_eigenbasis, exp_avg_sq.index_select(exp_avg_sq_dim, sort_idx)
274+
275+
276+
@torch.no_grad() # type: ignore[misc]
277+
def _project_to_one_sided_eigenbasis(
278+
x: torch.Tensor,
279+
eigenbasis: torch.Tensor,
280+
left_preconditioned: bool,
281+
) -> torch.Tensor:
282+
"""Project a matrix into the smaller-side covariance eigenbasis."""
283+
if left_preconditioned:
284+
return eigenbasis.T @ x
285+
return x @ eigenbasis
286+
287+
288+
@torch.no_grad() # type: ignore[misc]
289+
def _project_from_one_sided_eigenbasis(
290+
x: torch.Tensor,
291+
eigenbasis: torch.Tensor,
292+
left_preconditioned: bool,
293+
) -> torch.Tensor:
294+
"""Project a matrix from the smaller-side covariance eigenbasis."""
295+
if left_preconditioned:
296+
return eigenbasis @ x
297+
return x @ eigenbasis.T

0 commit comments

Comments
 (0)