Implement "NorMuon" Optimizer (Spectral Balancing)

[vukrosic]

Difficulty: Hard

Description:
We should investigate replacing the standard Muon optimizer with NorMuon. This variant modifies the Newton-Schulz iteration steps to effectively balance the spectral norm of the weight matrices. This has shown significant convergence speedups in later speedrun records.

Task:

1. Port the NorMuon optimization logic (specifically the update step and axis handling).
2. (optional) Run a baseline training run with standard Muon vs. NorMuon.
3. (optional) Perform a hyperparameter sweep on the learning rate, as NorMuon often requires different tuning than standard Muon.

References:

• Arxiv Paper: NorMuon
• Implementation PR (Original)
• Refinements & Fixes PR

---

🛠️ General Instructions

1. Environment Setup
To get started with development, follow these steps:

1. Fork this repository - Click the "Fork" button at the top right.
2. Clone your fork:

   git clone FORK_URL_HERE
   cd 5-dollar-llm

   (You may also clone it with our coding IDE)

Note: If you have already forked/cloned, please ensure you sync your fork with this repo & pull the latest changes to your local before starting - we make frequent changes)*

3. Install dependencies:

   pip install -r requirements.txt

2. Write your code

3. Verification & Testing

• Debug Mode: To quickly check if your code runs without errors (on CPU or GPU), use the debug script:

  python debug_moe.py

• Performance Test (optional): We will run the experiments anyways, but you may also run it (specify new name so you don't overwrite the baseline):

  python train_moe.py --experiment_name amp_speed_test

  (Note: This will use the GPU24GBMoEModelConfig by default)

4. Submission
Once finished, please create a Pull Request into the development branch, preferrably notify us on Discord as well.

No experiment is guaranteed to yield improvement, however, you will be credited for you work in any case.

[marc-shade]

🧠 NorMuon Research Complete

Key Findings

Paper: arXiv:2510.05491 - "NorMuon: Making Muon more efficient and scalable"

Performance

• 21.74% better than AdamW, 11.31% better than Muon (1.1B model)
• 2.9% computational overhead only
• 50% memory reduction vs AdamW

The Problem NorMuon Solves

Standard Muon produces non-uniform neuron norms after orthogonalization, causing some neurons to dominate. NorMuon adds per-neuron adaptive learning rates.

Implementation

Add ~7 lines after orthogonalization:

# After: update = orthogonalize(momentum)

################ NorMuon Addition ###################
vnorm = update.norm(dim=(-2,-1), keepdim=True)
v_mean = torch.mean(update * update, dim=-1, keepdim=True)
second_momentum.lerp_(v_mean, 1 - beta2)  # Track per-neuron stats
step_size = 1 / second_momentum.sqrt().add_(1e-10)
update.mul_(step_size)
update.mul_(vnorm / update.norm(dim=(-2,-1), keepdim=True).add_(1e-10))  # Preserve norm
#####################################################

# Continue: update *= lr_scale

Hyperparameters

lr = 0.02       # May need tuning vs standard Muon
momentum = 0.95 # β₁ 
beta2 = 0.95    # β₂ for per-neuron second-order stats

Memory

Only adds m scalars per weight matrix (one per neuron) - negligible overhead.

References

• Official repo: zichongli5/NorMuon
• modded-nanogpt PR #144
• modded-nanogpt PR #168 (fixes)

Happy to implement this as a PR!

[marc-shade]

✅ Implemented in PR #51

NorMuon is now implemented in optimizers/muon.py:

• NorMuon optimizer class
• Per-neuron second-order statistics tracking
• EMA with configurable beta2 parameter

Key implementation details:

• Saves Frobenius norm before normalization
• Computes row-wise variance: v_mean = mean(update² along dim=-1)
• Updates EMA buffer: second_buf.lerp_(v_mean, 1 - beta2)
• Per-neuron step size: 1 / (sqrt(second_buf) + eps)
• Restores original norm after normalization