Add RMNP (Row-Momentum Normalized Preconditioning) optimizer (arXiv:2603.20527)

[JohnLangford]

What

Adds RMNP (Row-Momentum Normalized Preconditioning), the optimizer from
Deng et al., RMNP: Row-Momentum Normalized Preconditioning for Scalable
Matrix-Based Optimization, arXiv:2603.20527,
as an opt-in member of the Muon/NorMuon family.

Method

RMNP replaces Muon's Newton-Schulz orthogonalization with a single row-wise
(input-dimension, d_in) ℓ₂ normalization of the momentum update. For an
m×n weight (m = d_out, n = d_in), the update is

V ← μ·V + G                      # momentum (Muon convention; see note below)
D ← row_normalize(V) = (diag(V Vᵀ))^{-1/2} V     # each row / its ℓ₂ norm
W ← W − η·D

row_normalize divides each row (length d_in) by its own ℓ₂ norm. This is
exactly Muon's orthogonalization (V Vᵀ)^{-1/2} V with the full inverse-square-root
replaced by its diagonal. The two coincide when V Vᵀ is (block-)diagonally
dominant — the empirically observed structure of the Transformer layerwise
Hessian that motivates the method; the paper shows orthogonalization and row-wise
d_in ℓ₂ normalization are asymptotically equivalent for the Transformer.
This drops the per-iteration cost from O(mn·min(m,n)) to O(mn) while, in
the paper's experiments, matching Muon-level optimization quality.

How it's wired into dion

RMNP is a standalone DistributedOrthoBase subclass (dion/rmnp.py) that
mirrors Muon/NorMuon. The row normalization is a module-level row_normalize
primitive passed as the newton_schulz_func, so it drops straight into the
existing, tested megabatch assembly path — literally "replace Newton-Schulz with
row normalization." Consequently RMNP:

• reuses Muon's momentum (muon_update_pre_orthogonalize) and weight update
  (muon_update_post_orthogonalize), and muon_update_megabatch_async;
• inherits Muon's distributed support unchanged (single device, DDP, FSDP2, and
  the num_heads / split_sizes param-group options), since the matrix is
  assembled the same way before the preconditioner runs;
• carries no preconditioner state beyond momentum (no Newton-Schulz iteration,
  so no use_triton / use_polar_express / use_gram_newton_schulz options).

Interpretation notes (flagged, since the paper leaves small choices open)

• Momentum convention. The paper writes the EMA V ← β·V + (1−β)·G; dion's
  shared Muon momentum is V ← μ·V + G. Because row normalization is
  scale-invariant, these differ only by the constant factor 1−β, which cancels
  — the normalized update is identical — so RMNP reuses Muon's momentum
  unchanged. (Covered by test_scale_invariant.)
• Learning-rate scaling. The paper applies no post-normalization rescaling
  (W ← W − η·D) and tunes the matrix LR directly, so adjust_lr defaults to
  None. spectral_norm / rms_norm remain available for consistency with
  the family / LR transfer, but were not studied in the RMNP paper.
• epsilon is added to each row norm (X / (‖X‖ + ε)), matching both the
  paper's formula and dion's Newton-Schulz stabilization idiom.

Tests

• tests/test_rmnp.py (CPU-runnable primitive + equivalence checks):
  per-row d_in normalization, shape preservation, scale invariance, the
  input-dimension axis, epsilon stability, and the asymptotic-equivalence
  claim — row_normalize equals the SVD polar factor (orthogonalization) to
  machine precision on a diagonal-Gram matrix V = diag(d)·Q (orthonormal rows),
  and diverges when rows are correlated (so the equivalence is regime-specific,
  not trivial). A CUDA check also matches the shipped Newton-Schulz kernel within
  its bf16 tolerance. Plus end-to-end: the applied update has unit-norm rows, and
  a toy loss decreases.
• tests/test_optimizers.py: a TestRMNP class mirroring TestMuon
  (basic / determinism / params-change / nesterov / cautious_wd / adjust_lr /
  megabatch / mixed shapes / state), plus RMNP entries in the shared num_heads
  and split_sizes fused-vs-separate parity tests, mixed matrix+scalar groups,
  and hyperparameter validation.

Full test suite passes locally (CUDA, 4×GPU). Sharded correctness was also
verified out-of-band: RMNP under FSDP2 row-sharding (dim 0) and column-sharding
(dim 1) reproduces the single-device result exactly.

Docs

README (intro list, parallelization table, file list, 1D-sharding section) and
CHANGELOG updated.

[Copilot]

Pull request overview

Adds the RMNP (Row-Momentum Normalized Preconditioning) optimizer as an opt-in member of the Muon/NorMuon family, implemented by reusing the existing Muon megabatch/distributed update path while swapping Newton–Schulz orthogonalization for a row-wise ℓ₂ normalization primitive.

Changes:

• Introduces dion.RMNP and the row_normalize preconditioner (Muon-compatible megabatch + distributed wiring).
• Adds targeted RMNP primitive/equivalence tests plus optimizer-level parity and validation coverage.
• Updates README and CHANGELOG to document RMNP and its parallelization/sharding support.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_rmnp.py	New unit + CUDA-gated end-to-end tests for row_normalize and RMNP behavior.
tests/test_optimizers.py	Integrates RMNP into the shared optimizer test matrix (parity, state, validation, mixed groups).
README.md	Documents RMNP and adds it to the optimizer/parallelization tables and sharding section.
dion/rmnp.py	Implements row_normalize and the RMNP optimizer by reusing Muon’s megabatch update path.
dion/init.py	Exposes RMNP at the package top level.
CHANGELOG.md	Notes RMNP addition and its intended behavior/defaults.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[JohnLangford]

Thanks for the review — both addressed in 3b816c5:

1. row_normalize docstring (rmnp.py): reworded so it is RMNP that replaces Muon's full (X Xᵀ)^{-1/2} orthogonalization factor with its diagonal (diag(X Xᵀ))^{-1/2} (previously phrased backwards).
2. test_normalizes_input_dimension_not_output (test_rmnp.py): axis labels corrected to be precise — rows are indexed by d_out (dim -2) and normalization is over d_in (dim -1); the prose now matches the assertions.

Also added test_sharded_update_matches_single_device, a multi-GPU test pinning that the FSDP2-sharded update (both row- and column-sharded) reproduces the single-device update exactly.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.