This repository provides a PyTorch implementation of Preconditioned Stochastic Gradient Descent (PSGD), a general-purpose second-order optimizer. It implements various preconditioner structures including Kronecker product, low-rank approximation, dense matrix, affine group, and X-matrix preconditioners. The code offers both Newton-type and whitening preconditioners with multiple update rules for the preconditioner matrices. Key features include support for different geometries (Q0.5EQ1.5, EQ, QEP, etc.), momentum handling, gradient clipping, and numerical stability mechanisms. The library includes optimizers suitable for various deep learning applications as demonstrated in the provided examples (vision transformers, GPT2, LSTMs, CNNs, etc.).
This repository is a specialized library of distributed orthonormal optimizers for large-scale language model training, implementing the Dion family of algorithms (Dion, Dion2, Dion3, NorMuon, Muon) alongside traditional optimizers (AdamW, Lion). It provides a complete training pipeline including a GPT model, data loaders for FineWeb datasets, and configurable parallelism strategies (FSDP2, TP, DDP). The core innovation lies in communication-efficient orthogonal gradient updates via submatrix selection, Newton-Schulz iterations, and polar decomposition, all wrapped in an async megabatch framework that overlaps collectives with computation. The codebase is well-structured with modular optimizer implementations, extensive tests (including distributed process spawning), and CUDA graph capture support. It is primarily an application library for ML practitioners, but retains a research character through its multiple algorithm variants and academic references. The high engineering quality and active maintenance make it suitable for both production use and further research extensions.