$ show microsoft/dion
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.
The LLM identified 4 files worth highlighting for reuse:
Provides AsyncTask and AsyncRuntime classes that form a lightweight event loop for overlapping asynchronous collective operations with GPU computation. This pattern is general-purpose: any distributed optimizer or training step that requires non-blocking all-reduce, all-gather, or all-to-all can reuse this framework to improve GPU utilization. The code is self-contained (only depends on torch) and exemplifies clean async orchestration with work queues and progress tracking.
CudaGraphOptimizer is a drop-in wrapper that captures the entire optimizer step (including NCCL collectives) into a CUDA graph and replays it on subsequent iterations. It demonstrates a reusable pattern for graph-capturing any PyTorch optimizer that uses in-place updates on persistent tensors. The handling of hyperparameter tensors (LR, weight decay) as device-side scalars enables live updates without host synchronization, making it suitable for dynamic schedules. Although it references dion internals, the pattern is general and could be adapted to other optimizer implementations.
Implements the base class DistributedOrthoBase with megabatch_orthogonalize_async, which groups parameters of the same shape into a single communication round (all-to-all or all-gather) to reduce the number of collective calls and kernel launches. This pattern is foundational for any distributed optimizer that needs to perform per-parameter orthogonalization efficiently. While it uses internal utilities from dion.opt_utils, the grouping logic and asynchronous ordering of communication vs. computation can be extracted as a general-purpose distributed optimizer framework.
Provides optimized in-place foreach implementations of AdamW and Lion updates that operate on lists of parameter tensors. These kernels are capturable for CUDA graphs and handle mixed precision (stochastic rounding) and cautious weight decay. The pattern of writing foreach loops that avoid Python overhead and device synchronization is reusable for any custom optimizer that needs efficient per-parameter updates without per-parameter kernel launch overhead.