Controlled LLM Training on Spectral Sphere

Scaling large models requires optimization strategies that ensure rapid convergence grounded in stability. Maximal Update Parametrization ($\boldsymbol{\mu}$P) provides a theoretical safeguard for width-invariant $\Theta(1)$ activation control, whereas emerging optimizers like Muon are only "half-aligned" with these constraints: they control updates but allow weights to drift. To address this limitation, we introduce the **Spectral Sphere Optimizer (SSO)**, which enforces strict module-wise spectral constraints on both weights and their updates. By deriving the steepest descent direction on the spectral sphere, SSO realizes a fully $\boldsymbol{\mu}$P-aligned optimization process. To enable large‑scale training, we implement SSO as an efficient parallel algorithm within Megatron. Through extensive pretraining on diverse architectures, including Dense 1.7B, MoE 8B-A1B, and 200-layer DeepNet models, SSO consistently outperforms AdamW and Muon. Furthermore, we observe significant practical stability benefits, including improved MoE router load balancing, suppressed outliers, and strictly bounded activations.

Training Large Language Models requires a delicate balance: the algorithms must learn as fast as possible without becoming unstable. Today's fastest training methods only partially follow the fundamental rules of stability. Over time, the model's internal metrics slowly drift out of safe bounds, causing numerical signals to explode and forcing engineers to rely on messy, trial-and-error patches. To solve this, we built the Spectral Sphere Optimizer (SSO). Instead of letting parameters drift, our algorithm confines the model's learning process to a strict mathematical geometry—a "spectral sphere." SSO calculates the absolute fastest path to improve the model while guaranteeing its internal values never cross into unstable territory. We also engineered a way to run this highly complex math efficiently across massive supercomputers. When tested on models with up to 8 billion parameters, SSO consistently outperformed popular industry methods. By naturally preventing massive numerical errors and keeping training smooth, our research provides developers with a fast, reliable, and theoretically sound tool to train the next generation of massive AI systems.