MOSO: use RMSProp directly (remove double momentum)

[mchrzanowski]

Summary

MOSO applies a Muon-style first-moment EMA (momentum) to the gradient before preconditioning. The inner Adam previously also exposed beta1, adding a second, redundant momentum on the projected update direction. This double-averaging over-damps the update and measurably hurts convergence — and it was reachable by default (betas=(0.9, 0.95)).

This PR removes the inner first moment entirely and makes double momentum structurally impossible: the inner optimizer is now RMSProp (calculate_rmsprop_update, matching torch.optim.RMSprop with momentum=0, centered=False), controlled by a single rms_beta. The now-dead exp_avg first-moment state and its eigenbasis reprojection are removed. There is no longer any API surface to add a second momentum.

-        betas: tuple[float, float] = (0.9, 0.95),
+        rms_beta: float = 0.95,
-        adam_update = update_functions.calculate_adam_update(..., betas=(0.0, rms_beta), correct_bias=True, ...)
+        rmsprop_update = update_functions.calculate_rmsprop_update(..., alpha=rms_beta, eps=...)

Evidence

Hybrid Mamba + 128-expert MoE, WSD schedule, per-config LR-tuned. Removing the inner first moment (beta1 0.9 → 0):

rung	old default (beta1=0.9)	fixed (no inner momentum)	Δ
1B (88B tok)	1.7238	1.6268	−0.097

• Holds across a 1B → 7B scaling ladder: the fixed optimizer beats the AdamW baseline at every scale (1B 1.627, 2B 1.512, 4B 1.433, 7B 1.359) and closes most of the gap to Muon+Lion / SOAP-KL; the old default was the worst of the four optimizers.
• The RMSProp form (this PR) was re-validated at 1B against the beta1=0 Adam form: the two are identical after the first ~200 steps — RMSProp drops the second-moment bias correction, which only affects the earliest steps (a benign larger-step transient).

Behavior / compatibility

• Breaking API change: betas is removed; use rms_beta. The old (0.9, 0.95) behavior is intentionally no longer reachable.
• The inner update is RMSProp rather than Adam-with-beta1=0; numerically identical after the bias-correction warmup.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR simplifies MOSO by replacing the inner Adam update with RMSProp-style normalization. The main changes are:

• Adds and exports calculate_rmsprop_update.
• Replaces betas with rms_beta in MOSO.
• Removes the inner first-moment optimizer state.
• Updates MOSO tests for the new parameter shape.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
emerging_optimizers/scalar_optimizers/update_functions/init.py	Adds the RMSProp update helper to the update-functions export surface.
emerging_optimizers/scalar_optimizers/update_functions/rmsprop.py	Introduces the RMSProp update helper used by MOSO.
emerging_optimizers/soap/moso.py	Switches MOSO from an inner Adam update to RMSProp-style second-moment normalization.
tests/test_moso.py	Updates MOSO tests for rms_beta and the removed betas argument.

_{Reviews (6): Last reviewed commit: "MOSO: use RMSProp directly instead of Ad..." | Re-trigger Greptile}

[mkhona-nvidia]

What is this test?

[skyw]

Looks it test wrong kwargs name.

Not a very strong test, doesn't hurt to keep it with proper naming, e.g. test_betas_raise_type_error.

[mchrzanowski]

Renamed to test_betas_raise_type_error and dropped the redundant construction line (per @skyw). It's intentionally light — it just guards that the removed betas kwarg now raises TypeError, so double momentum can't sneak back in through the old API. Done in bc7ba66.

[skyw]

Testing is weak, but should be OK.
Defer to @mkhona-nvidia for approval.

[skyw]

Looks it test wrong kwargs name.

Not a very strong test, doesn't hurt to keep it with proper naming, e.g. test_betas_raise_type_error.

[mkhona-nvidia]

@mchrzanowski can we replace the current scalar_optimizers.calculate_adam_update_function with a RMSProp update directly from pytorch via torch.optim.rmsprop.rmsprop(args,kwargs)? https://docs.pytorch.org/docs/2.13/generated/torch.optim.rmsprop.rmsprop_function.html#torch.optim.rmsprop.rmsprop

Basically we are using the adam function and zeroing out momentum to simulate an RMSProp, might as well do it directly.

[mchrzanowski]

@mkhona-nvidia done in bc7ba66 — MOSO now calls a dedicated calculate_rmsprop_update (added to scalar_optimizers.update_functions) instead of calculate_adam_update with a zeroed first-moment beta, and I removed the now-dead exp_avg state and its eigenbasis reprojection.

One note: I didn't wire in torch.optim.rmsprop.rmsprop directly because that functional updates parameters in place and can't return the eigenbasis-space update MOSO needs (it projects the update back out of the covariance eigenbasis and RMS-clips before applying lr). The helper is the same math — torch.optim.RMSprop with momentum=0, centered=False: v = alpha*v + (1-alpha)*g^2; update = g / (sqrt(v) + eps).

This does change numerics slightly vs. the old path, which bias-corrected the second moment (correct_bias=True); true RMSProp drops that. I re-validated at 1B (Hybrid-Mamba MoE, LR 2.5e-3): the RMSProp path tracks the previous beta1=0 run — the only difference is a benign larger-step transient in the first ~200 iters (bias correction ≈ 1 after that); the train-loss delta is already down to −0.02 by iter 160 and closing.

[skyw]

/ok to test bc7ba66