Megaprop
A cross-repository control project integrating Newton-Muon and Locoprop-style preconditioned optimizers, with FSDP and MuP support, into Megatron-LM and TransformerEngine for training mixture-of-experts models at scale.
Megaprop adds preconditioned stochastic gradient descent primitives (FOOF, KFAC, Locoprop-S and Newton-Muon) to Megatron-LM, by threading a feature gram matrix alongside the usual gradient computation through changes to TransformerEngine, Megatron-LM and a companion Emerging-Optimizers project, with out-of-the-box support for FSDP and MuP.
In a Megatron GPT sweep on FineWeb-Edu, with the feature gram matrix refreshed every 8 steps, Locoprop-S beats AdamW on wallclock time; Newton-Muon currently trails, which the project attributes to the cost of its polar iteration. A diagonal approximation of the gram matrix, chosen to save memory, performs well in early experiments.
The project pins specific fork branches of Megatron-LM, Emerging-Optimizers and TransformerEngine, and documents the size of its patch against each as several hundred to just over a thousand added lines.