New training speed record for @karpathy’s 124M-parameter NanoGPT setup: 3.28 Fineweb validation loss in 3.7B training tokens Previous record: 5B tokens Changelog: new optimizer 1/8
The new optimizer is defined as follows. It is based on orthogonalizing the update given by SGD-Nesterov-momentum in an efficient way
It uses half the memory of AdamW and takes 3% extra wallclock time per step for this setup Here's code to reproduce the result: https://github.com/KellerJorda...
A variant of this optimizer also yielded a training speed record for CIFAR-10. Scaling it up to NanoGPT was surprisingly simple 4/8 https://x.com/kellerjordan0/st...
The optimizer's core design decisions are: - For each parameter, orthogonalize the update *after* applying Nesterov momentum - Use a fast iterative method for the orthogonalization. Quintic Newton-Schulz works well - Run the iteration in bfloat16 5/8
And - To minimize the number of Newton-Schulz steps, select the quintic coefficients aggressively to maximize slope at zero, even at the cost of some variance / non-convergence. There's an interesting tradeoff space here 6/8
I would like to thank & acknowledge @jxbz for sending me his recent paper, which is where I learned about the crucial Newton-Schulz iteration method. He also had the insight that my initial quintic coefficients could be improved. 7/8 https://arxiv.org/abs/2409.203...
I'm still thinking about names for the optimizer, & there may be a full arXiv writeup at some point. I hope the community finds this intermediate result interesting; critiques are welcomed 8/8
@kellerjordan0 @karpathy This is very interesting, gonna try this weekend!
@Yuchenj_UW @karpathy Awesome
@kellerjordan0 @karpathy This is great! While speed is great a lower loss for the same amount of tokens is even greater. Can you go below the baseline if you continue training?
@kellerjordan0 @karpathy What a nice hyperparameter optimum you have here! Addding weight decay -> worse! Addding early dropout -> worse!
@bozavlado @karpathy Yeah, and wait til you hear that we can also completely remove warmup!
@kellerjordan0 @karpathy Tried it on my own diffusion model. Seems it's quite aggressive - the model quickly forgets what it has learned (10% drop on a metric). I set LR to be 1/10 of AdmW, and separated the params to two sets: 2D (maybe most are attention components?) and non-2D ( optimized by AdamW).
@kellerjordan0 @karpathy you guys should compare this to the AdEMAMix optimizer
@kellerjordan0 @karpathy Quite cool! Makes me want to test it, too. The real question for LLM pretraining is, does it work well with weight decay & comparison on 2T+ token baseline?
@pa_pfeiffer @karpathy good question
@kellerjordan0 @karpathy This looks great! Probably a silly question, but what's the difference between "tuned AdamW" and nanoGPT's AdamW? (which hparams did you tweak?)
@kellerjordan0 @karpathy Awesome research!!
@kellerjordan0 @karpathy Will this have arbitrary shape tensor version? (not only for convolution, since convolution can be taken as kind of linear after im2col algo)
@kellerjordan0 @karpathy Rather than solely optimizing for abstract metrics like loss in transformer models, the focus should be on improving tasks that deliver value to the user, such as accurately answering questions or ensuring the latent space captures meaningful structures.
@kellerjordan0 @karpathy How is the generalization?
@kellerjordan0 @karpathy Bestimizer
@kellerjordan0 @karpathy Beautiful
@kellerjordan0 @karpathy Wow a Twitter thread that's worth a paper
@kellerjordan0 @karpathy This is impressive!
@kellerjordan0 @karpathy Nice results! To compare other optimizers could you tell what the final val loss for the new optimizer is? On the 7k step run?
@kellerjordan0 @karpathy I feel like @jeremyphoward would be all over this
@kellerjordan0 @karpathy One question, the note says it should not be used for 1D parameters, so why not also for the last linear layer? Because the output of it is 1D? Any tipps for Networks that are more diverse in terms of Parameter dims? Do i have to collect all 1d layers beforehand?
@kellerjordan0 @karpathy What’s the scheduler? Loss shape looks unusual to me (after 4B tokens)