Projects
My work is available primarily through the GitHub and 🤗 Hugging Face Hub. These projects range from first-principles implementations of modern AI systems to original research on optimization, memory, and training dynamics.
Scaling Orthogonal Optimizers: Muon & Dion3 (PyTorch Con 2026)
An end-to-end integration of $O(n^3)$ orthogonal optimizers into PyTorch’s FSDP2 distributed training stack, and the subject of my PyTorch Con 2026 presentation. This project bridged the gap between academic theory and production deployment by resolving severe framework-level API collisions, including diagnosing and patching a silent tensor-aliasing memory leak during Distributed Checkpoint (DCP) recovery.
The resulting integration for AI2’s olmo-core (PR #919) successfully scales megabatched, fractional Gram Newton-Schulz updates across HSDP meshes. It proves that orthogonal optimizers can survive aggressive FP8 mixed-precision quantization and cluster node failures without losing their convergence advantage over standard AdamW. ››
Muon-Based Optimization for GRPO
I am developing a Muon-based optimizer variant for GRPO-style reinforcement learning. The project examines whether Muon’s matrix-aware updates can improve stability and optimization behavior in reinforcement-learning settings, where noisy rewards and long-horizon generation create dynamics very different from pretraining.
The work includes optimizer design, integration into the training loop, instrumentation of gradient behavior, and controlled comparisons against conventional optimization baselines.››
Engram Conditional Memory in olmo-core
I extended olmo-core with Engram conditional memory to study how external memory interacts with attention-based and hybrid architectures.
The project became a broader investigation of architecture-specific learning dynamics. I traced anomalous loss and gradient behavior to implementation defects, corrected the integration, and reran the experiments to separate genuine memory effects from artifacts of the original code. ››
Distributed Muon
A reverse-engineered and annotated implementation of Moonshot AI’s distributed Muon optimizer, with an emphasis on multi-node communication, CPU offloading, and the systems decisions that make Muon practical at scale. ››
Understanding the Muon Optimizer
A theory-to-code guide to Muon, covering Newton–Schulz orthogonalization, matrix-shaped parameter updates, and the differences between Muon and conventional adaptive optimizers. ››
Reinforcement Learning for Tool Use
A hands-on implementation of ReTool, exploring how reinforcement learning can teach language models when and how to invoke external tools across multi-step reasoning trajectories. ››
Model Architecture from First Principles
Implementations of DeepSeek’s Multi-Head Latent Attention and Mixture-of-Experts architectures, along with a complete Transformer-from-scratch tutorial.
Post-Training Methods, Visualized
An interactive guide to distillation, preference optimization, RLHF, DPO, and related post-training methods. ››