https://github.com/KellerJordan/modded-nanogpt Skip to content Navigation Menu Toggle navigation Sign in * Product + GitHub Copilot Write better code with AI + Security Find and fix vulnerabilities + Actions Automate any workflow + Codespaces Instant dev environments + Issues Plan and track work + Code Review Manage code changes + Discussions Collaborate outside of code + Code Search Find more, search less Explore + All features + Documentation + GitHub Skills + Blog * Solutions By size + Enterprise + Teams + Startups By industry + Healthcare + Financial services + Manufacturing By use case + CI/CD & Automation + DevOps + DevSecOps * Resources Topics + AI + DevOps + Security + Software Development + View all Explore + Learning Pathways + White papers, Ebooks, Webinars + Customer Stories + Partners * Open Source + GitHub Sponsors Fund open source developers + The ReadME Project GitHub community articles Repositories + Topics + Trending + Collections * Enterprise + Enterprise platform AI-powered developer platform Available add-ons + Advanced Security Enterprise-grade security features + GitHub Copilot Enterprise-grade AI features + Premium Support Enterprise-grade 24/7 support * Pricing Search or jump to... Search code, repositories, users, issues, pull requests... Search [ ] Clear Search syntax tips Provide feedback We read every piece of feedback, and take your input very seriously. [ ] [ ] Include my email address so I can be contacted Cancel Submit feedback Saved searches Use saved searches to filter your results more quickly Name [ ] Query [ ] To see all available qualifiers, see our documentation. Cancel Create saved search Sign in Sign up Reseting focus You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session. Dismiss alert {{ message }} KellerJordan / modded-nanogpt Public * Notifications You must be signed in to change notification settings * Fork 25 * Star 355 NanoGPT (124M) quality in 3.25B tokens 355 stars 25 forks Branches Tags Activity Star Notifications You must be signed in to change notification settings * Code * Issues 2 * Pull requests 0 * Actions * Projects 0 * Security * Insights Additional navigation options * Code * Issues * Pull requests * Actions * Projects * Security * Insights KellerJordan/modded-nanogpt This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository. master BranchesTags Go to file Code Folders and files Name Name Last commit Last commit message date Latest commit History 139 Commits data data img img records records .gitignore .gitignore README.md README.md requirements.txt requirements.txt run.sh run.sh train_gpt2.py train_gpt2.py View all files Repository files navigation * README Modded-NanoGPT This is a variant of the PyTorch GPT-2 trainer from Andrej Karpathy's llm.c repo. It: * Trains 3x more efficiently (taking only 3.15B tokens instead of 10B to reach the same validation loss). * Has shorter code (524 lines instead of 860). * Implements architectural modernizations (rotary embeddings, RMSNorm, ReLU^2). * Implements a new optimizer (Muon - Momentum Orthogonalized by Newton-schulz). To execute the training, run the following three commands on an 8xA100 or 8xH100 node. They complete in <45min on an 8xH100 with decent internet connection. pip install -r requirements.txt python data/cached_fineweb10B.py 35 # downloads the first 3.5B tokens ./run.sh This will train a 124M-parameter transformer for 6000 steps on 3.15B tokens of Fineweb [1], achieving ~3.275 validation loss. For comparison, the default llm.c PyTorch trainer yields >3.28 validation loss after training for 10B tokens. --------------------------------------------------------------------- Figures Figure 1. Proposed optimizer vs. a well-tuned AdamW. [fig_optimi] --------------------------------------------------------------------- Proposed optimizer For this training scenario, the proposed optimizer has the following properties: * Half the memory usage of Adam * 1.5x faster training * <9% wallclock overhead (which can be further brought down by distributing the overhead; it's currently performed redundantly on all 8 GPUs) The optimizer is defined as follows: [algo_optim] Where NewtonSchulz5 is the following Newton-Schulz iteration [2, 3]: @torch.compile def zeroth_power_via_newtonschulz5(G, steps=5, eps=1e-7): assert len(G.shape) == 2 a, b, c = (3.4445, -4.7750, 2.0315) X = G.bfloat16() / (G.norm() + eps) if G.size(0) > G.size(1): X = X.T for _ in range(steps): A = X @ X.T B = A @ X X = a * X + b * B + c * A @ B if G.size(0) > G.size(1): X = X.T return X.to(G.dtype) Provenance Many of the choices made to generate this optimizer were obtained experimentally by our pursuit of CIFAR-10 speedrunning. In particular, we experimentally obtained the following practices: * Using Nesterov momentum inside the update, with orthogonalization applied after momentum. * Using a specifically quintic Newton-Schulz iteration as the method of orthogonalization. * Using non-convergent coefficients for the quintic polynomial in order to maximize slope at zero, and thereby minimize the number of necessary Newton-Schulz iterations. * Running the Newton-Schulz iteration in bfloat16 (whereas Shampoo implementations often compute the preconditioners via inverse-pth-roots in fp32 or fp64). Our use of a Newton-Schulz iteration for orthogonalization traces to Bernstein & Newhouse (2024), who suggested it as a way to compute Shampoo [5, 6] preconditioners, and theoretically explored Shampoo without preconditioner accumulation. In particular, Jeremy Bernstein @jxbz sent us the draft, which caused us to experiment with various Newton-Schulz iterations as the orthogonalization method for this optimizer. If we had used SVD instead of a Newton-Schulz iteration, this optimizer would have been too slow to be useful. Bernstein & Newhouse also pointed out that Shampoo without preconditioner accumulation is equivalent to steepest descent in the spectral norm, and therefore Shampoo can be thought of as a way to smooth out spectral steepest descent. The proposed optimizer can be thought of as a second way of smoothing spectral steepest descent, with a different set of memory and runtime tradeoffs compared to Shampoo. --------------------------------------------------------------------- Other general differences between this codebase and NanoGPT To simplify the code, some features have been removed, including text generation. And to obtain a training speed improvement, we have diverged from being a strict reproduction of the GPT-2 paper. The speedup is due to the following changes: * Increased learning rate by 3x * Switched to trapezoidal learning rate schedule following [7] * Switched to rotary embeddings and ReLU^2 activation * Removed the special initialization for linear layers before residuals. Instead, just scale down the output of the attention block by a fixed scalar. * Removed all affine scale and bias parameters from the architecture, and switched to RMSNorm (actually this causes a slight slowdown, and I just did it to reduce code complexity) * Switched from AdamW to new optimizer, and removed learning rate warmup --------------------------------------------------------------------- References 1. Penedo, Guilherme, et al. "The fineweb datasets: Decanting the web for the finest text data at scale." arXiv preprint arXiv:2406.17557 (2024). 2. Nicholas J. Higham. Functions of Matrices. Society for Industrial and Applied Mathematics, 2008. Equation 5.22. 3. Gunther Schulz. Iterative Berechnung der reziproken Matrix. Z. Angew. Math. Mech., 13:57-59, 1933. 4. Jeremy Bernstein and Laker Newhouse. "Old Optimizer, New Norm: An Anthology." arxiv preprint arXiv:2409.20325 (2024). 5. Vineet Gupta, Tomer Koren, and Yoram Singer. "Shampoo: Preconditioned stochastic tensor optimization." International Conference on Machine Learning. PMLR, 2018. 6. Anil, Rohan, et al. "Scalable second order optimization for deep learning." arXiv preprint arXiv:2002.09018 (2020). 7. Hagele, Alexander, et al. "Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations." arXiv preprint arXiv:2405.18392 (2024). itsover_wereback About NanoGPT (124M) quality in 3.25B tokens Resources Readme Activity Stars 355 stars Watchers 4 watching Forks 25 forks Report repository Releases No releases published Packages 0 No packages published Languages * Python 99.9% * Shell 0.1% Footer (c) 2024 GitHub, Inc. Footer navigation * Terms * Privacy * Security * Status * Docs * Contact * Manage cookies * Do not share my personal information You can't perform that action at this time.