New historic NanoGPT record at 39.9s (-27.7s) from @DevenPzak , obliterating the prior record of 67.6s! This record introduces a new paradigm of thinking to NanoGPT: instead of optimizing matmuls or adding more expressive operations, optimize at the individual flop level with incredibly clever engineering and ML judgement. If a flop is low value on a particular step, skip it. Specifically: -(~8s) Sampled softmax. If a token doesn’t appear in a batch, skip its lm_head fwd/bwd some fraction of the time. -Sparse values. Only run an optimizer step for ngram embeddings that occurred in the batch. Set beta1 to zero to enable this. Beta2 is applied retroactively when the row is later used. -Sparse updates. Only update ngram and value embeddings once every 4 steps instead of once every 2. -Sparse communication. Shard the n-gram table across GPUs, and only pass the rows receiving updates on each step. -Sparse optimizer states. For the n-gram table, reduce from 2 floats in Adam optimizer per param, to 1 float per 768 params. -Hand-rolled flash attention for 64 dim heads. There are several additions that add accuracy too: -(~4s) EMA during last 300 steps, combined with lifting final_lr to 0.3 instead of 0.15. -(~1s) A new optimizer, Anvil2, which expands muon via a second tracked momentum buffer, improves the ortho coefficients, and modifies the cautious weight decay application. -A couple additional dynamic skip connections in the network. The most striking consequence of the ‘flop aware paradigm’ is you can grow parameters arbitrarily large, only limited by the available memory, since you can selectively choose how to expend flops on those parameters on each step. NanoGPT has kept active parameters below 124M, but total is unbounded, and has grown to 640M through embedding sparsity over the last year. This PR takes that to its logical conclusion on the 8xH100, scaling up to 65B sparse embedding parameters, which accounts for 25% of the PR’s gains. At frontier scale, where one is not bounded by an 8xH100, one could imagine where this paradigm could lead. github.com/KellerJordan/modd… As this was a very notable PR, I spoke with Deven for an hour to learn how he did it. Here’s his story on the changes: hyperstition.cc/training-nan…

There’s an enormous amount of open training data on Hugging Face. What does it take to make it work together 🤗? For Marin’s 535B run, we built on 25T tokens from 152 datasets with licenses permitting training. Here’s the work between downloading those and training a model 🧵

New NanoGPT Speedrun WR at 67.6s (-0.4s) from Jan Varho (jvarho on GitHub). Simple idea: mask impossible continuations. EG: "Hyperparameter" tokenizes to ["Hyper", "param", "eter"]. "eter" can never follow "Hyper" in the tokenized dataset since "Hypereter" tokenizes to ["H", "ype", "re", "ter"]. Yet, during multi-token prediction, "Hyper" learns to predict both "param" and "eter". Explicitly masking "eter" helps improve val loss, and during inference prevents a model from generating token pairs its never seen in training. github.com/KellerJordan/modd…

New NanoGPT Speedrun WR at 68.0 (-5.8s) from @theonlyglitch_ , with a decrease of QK dim from 128 to 96, a kernel fusion of [QK Norm, RoPE, KeyOffset, paired head layout] into Triton, and moving MLP bwk, QKV fwd/bwk to FP8. This is a heavily involved PR with over 1k lines of triton, 4 different FP8 scaling protocols, and fancy register aware epilogue placement. Two takeaways: 1) From a model perspective, if QK dims smaller than 128 works better at nano scale, then perhaps larger than 128 works better at Hero scale. 2) This Nvidia Engr is super legit. github.com/KellerJordan/modd…

Was reading "Lost in Backpropagation: The LM Head is a Gradient Bottleneck" by @nthngdy and @yoavartzi the other day. Their ideas seemed very clever, but my ability to follow their intuition was lacking. So I decided to imagine "What would it be like to work as a neuron", to better process their ideas. The main audience when writing this was just myself. Sharing in case others find it a fun read. A Day in the Life of a Neuron Associate at Grub Zoo. Welcome to Grub Zoo. Proud home of 1000 animals. Each animal has its own designated 5 minute feeding window during the day where it eats one serving of grub. The feeding schedule looks something like: -Feed the horse and the Giraffe from 4:00-4:05. -Feed the rhino from 4:15-4:20. -Feed the 8 alligators from 7:10-7:15. Initially the Zoo ran off a simple lookup chart, where one could query a 5 minute window, and get returned a list of all animals that needed to be fed. Well, Bob from Marketing had a great (terrible) idea to replace our lookup table with AI. But instead of using AI, we could just become AI (wtf?). He explained that we’d each get some role in a mega human neural network, then use the power of ‘backpropagation’ to tell each other what to do. I think this is his promo project or something, not sure. His idea is that the optimal feeding window should depend on all sorts of properties of the weather, and once we get this learning machine running, it can plug in a bunch of these advanced features. The first part of the plan was to designate 10 employees as feeders. Each feeder had to write down how they would divvy up any grub they were given between the 1,000 animals. One feeder, Jeremy, decided he’d give all grub he was given to the sloth. Weird guy. Another one, Anna, decided she’d split up the grub equally across all 1,000 animals. (How’s she gonna manage that, lol?). My job was the supplier. Every 5 minutes I’d show up with some servings of grub. Then I’d have to divvy it up between the feeders, who’d then divvy it up between the animals based on their prearranged fractions. The weird part here is that I was the only one who was aware of the time of day- the feeders always stuck to their fractions. Bob said that during each 5 minute interval we’d have two phases: an execution phase, and a feedback phase. Well for the first 5 minute execution phase I frankly had no idea what was going on, so I just gave an equal amount of grub to each of the 10 feeders, who then distributed it out. Somehow Anna was able to get grub out to all 1,000 animals, impressive. Only the horse and the giraffe were supposed to get fed in this window, but the small amount of grub ended up spread out across all animals, and nobody was happy. I voted to drop Bob’s idea, but he figured we could salvage things. His insight was “Most windows only have a couple animals who need getting fed, so after everybody distributes their grub, we should apply some sort of bunching action towards the ones with the most grub. That way at least somebody gets a decent meal. Sometimes multiple animals need grub and we won’t be sure we got it right, so just do soft bunching.” He called it softmax. We applied his technique, and sure enough one animal got a good meal, but it was the wrong animal! And the animals at Grub Zoo are extremely picky about when they get fed! Well, Bob said this was ok at first (what?), we just needed the feedback phase. The feedback phase started with the animals. The ones that didn’t get fed but were supposed to, complained loudly. The feeders were the only ones who could observe and hear the animals. Each feeder would first go observe all 1000 animals to see which ones were complaining. Then they’d update their fractions slightly towards giving them more grub. Now if the supplier had given the feeder little or no grub that window, the feeder would take little effort in updating their fraction as it wouldn’t matter anyway. But if they had been given a large amount of grub, they’d take tremendous honor in updating their fraction quite heavily. But they had to be careful here! They were going to use that fraction for every time window, so they didn’t want to over-correct for one 5 minute interval. The next phase of feedback was for me as the supplier. Each of the 10 feeders came and told me either ‘Man I coulda made a difference I had the fractions right just give me more grub’, or ‘I was off on this one, didn’t need grub here’. Then my job was to update ‘Ok for this 5 minute interval, give grub to feeder X’, because they told me their fraction is calibrated. I figured none of the feeders had it exactly right, but if I gave more grub to the feeders that seemed decent at the right times, their errors would average out and softmax would help bunch it up on the right animals. To me the process felt pretty dumb, as the feeders were just like mindless cogs, trying to find the one right fraction. They could see all 1000 animals. But all I could see were 10 people shouting ‘give me more/less grub’. And somehow I was supposed to manage all the windows from that? Over time I developed a massive lookup table. For each 5 minute interval, how much grub to give to each of the 10 feeders. Each feeder started to get more sophisticated specialization. For instance, one guy fed just the 8 alligators. I knew to only give him grub from 7:10-7:15, because that’s when he complained. Well, Bob got wind that I was using a lookup table for each 5 minute window and freaked out. Said it wasn’t true AI. Said I was getting demoted to a neuron. 9 other people got hired as neurons too, working in parallel to me. Now instead of a unique fraction for each window, I had a simpler role. I had to learn a fixed ratio across the 10 feeders, and a simple rule for when to switch on. Initially I just put the fraction on for the morning. And then tuned the fraction a bit to give grub to the feeders that were clamoring for more grub in the morning. I figured if I helped focus on some of the morning intervals, other ‘neuron associates’ could manage other windows of the day. I never got to see what these other neuron associates were up to, but I figured if I listened to the feeders I could help compensate for the gaps. Well, we got some new funding, and Bob decided to hire some more ‘neuron associates’, so we are up to 20 now. Communication is becoming a real problem. We never talk to each other, and it’s getting kind of weird only listening to the 10 feeders and trying to compensate for what all my coworkers are doing. Now Bob had the craziest idea yet. He decided to break up our 20 neuron associates into two 10 member squads: Alpha and Bravo. I got put on Bravo. During each interval I’d get to see both the time of day, and how much the Alpha Squad was planning to give to each feeder. I decided to take on a targeted role: IF the time was between 10-11am AND I saw that Alpha Squad was giving grub to feeder 2, then I’d give grub to feeder 1 and 2. I learned that in the morning whenever Alpha Squad gave grub to feeder 2, feeder 2 would still complain and want even more, and feeder 1 complained as well. I noticed that sometimes Alpha Squad would slip up and not give grub to feeder 2, which meant I didn’t trigger on feeder 1 or 2. When this happened, I’d complain loudly. Hey Alpha! Give more grub to feeder 2! Problem is, team Alpha was getting overwhelmed during the feedback cycle. They had to listen to both the feeders directly, and every single member on squad Bravo. The problem got even worse when we doubled the size of Alpha and Bravo to 20 people each. That’s when Bob started hiring a new role, called ‘transporters’ and ‘messengers’. There was one messenger and transporter for each feeder. Everyone only got to receive feedback from the messengers, and give grub to the transporters. During the execution phase, the transporters would collect the grub from squad Alpha, then I could make my decision based on if transporter 2 had grub. All of us in squad Bravo would then add on to the grub the transporters were carrying. During the feedback phase, I’d tell messenger 2 if I wanted more grub sent to feeder 2. This way, Alpha squad didn’t have to talk to both the feeders and Bravo squad. The messengers just added up the feedback from both the feeders and Bravo. We ran this operation for several weeks, with Alpha and Bravo squads distributing grub to the transporters, the messengers giving feedback, and the feeders divvying it out and observing the animals. It seemed like the dumbest thing I’d ever seen! A lookup table could have done the same thing! Well, that’s when Bob made a change that actually blew my mind. He called it the ‘magic of generalization’. The first thing he did was give all neuron associates access to the temperature, in addition to the time of day. Now in the old world, we’d have to manually make updates to all 1000 animals in the lookup table to figure out how they responded to temperature. In the new world, I only had to change my allocation to the 10 transporters. The feeders had already built fractions that grouped the animals into correlated sets. I developed a new trigger condition. IF the time was between 10-11am AND transporter 2 had grub AND the temperature was above 60F, then give grub to transporter 1, 2, and 3. I came to learn that in very cold weather, these messengers didn’t complain as much. Over time Bob added new features I could trigger on. Rain. Wind. Sunshine. It was starting to be quite a lot to handle. Bob always had to go around explaining to every neuron associate how this new measurement would work. Eventually Bob got pretty tired of this. He decided that only squad Alpha could see things like time, temperature, rain, wind, and sunshine. At first this seemed completely crazy to me. How could I know when to trigger, without these key features? We concocted a pretty funny plan. I noticed that messenger 9 NEVER complained. My guess is that feeder 9 was just completely checked-out. Never gave any grub to anyone. Team Alpha had picked up on this. They also knew that whenever it was between 10-11am and above 60F, messenger 2 would complain to them. So they decided to give grub to transporter 9 whenever it got hot in the morning- not because transporter 9 needed the grub. But because team Bravo could then use that signal to make better decisions. I developed a new heuristic: IF transporter 2 AND transporter 9 have grub, then we give grub to transporter 1,2,3. AND we also take that grub away from transporter 9 and add it to the pile, since feeder 9 never complains. So that became my job. I’d wake up every morning, drive in to work as a neuron associate, then pay very close attention to which 10 transporters had grub, and how that correlated to the complaints I’d get from the 10 messengers. I’d update my fractions among the 10 transporters, and update my feedback to the 10 messengers. Later that month my team name got updated to Squad Bravo-46. What the heck? Turns out Bob had really scaled the operation. We now had 50 Bravo squads in series. My job was always the same, talk to 10 messengers, and observe 10 transporters. I had no idea what measurements got fed into Alpha Squad. And honestly, I had no idea how many animals even lived at the Zoo anymore. I didn’t even know how many team members were on Bravo-46, since I never talked to them. Alpha Squad was the only team that got to see the pure, raw measurements. And the 10 feeders were the only team that got to see the actual animals. In particular, it seemed like the feeders functioned like quite the bottleneck. I'd love to get to observe all 1,000 animals, instead of just 10 messengers and transporters. Anyways, that’s my job. Hope the giraffe is doing alright.

🚢 Marin 535B-A23B started training this week! As usual, the whole process is open. Voyage plan: pretraining (80%) + midtraining (20%) on 18.75T tokens on 11 x GB200 NVL72 for ~3 months (2.7e24 FLOPs). Post-training will follow. Before kicking off the run, we trained a 4-rung scaling ladder from 1.6B-A61M (48B tokens) to 27.7B-A1.2B (926B tokens) to debug issues, and to make a forecast of our hero run. This is by far our biggest run, so definitely expecting the unexpected.

At Marin we’ve found that scaling laws let us simulate the entire training trajectory. This became clear in our last 67B MoE run. The run hit the loss target within 1%, but perhaps more interestingly, arbitrary points during training could also be predicted to within 1%. This model is finishing up long context and will soon enter RL. Based on this finding, we’ve used a scaling ladder to simulate the training trajectory of the 535B MoE we kicked off yesterday. Will see how it goes. Now if someone can figure out how to fit a scaling law to each parameter we can just predict our final model and skip this training business. Something else I've learned from this process is that getting a training run off the ground is way more than just fiddling with the architecture. Lot of herculean efforts from teammates to get trillions of tokens curated, get hardware running on a new GPU stack, improve MFU, and propose great ideas that were tested and integrated. Given that this is our largest run to date, I anticipate lots of learning and adapting along the way. But thats part of the fun. Details: github.com/marin-community/m…

New NanoGPT Speedrun WR at 73.8s (-0.8s) from Mister-dev-oss, CerovazS, MarioPaerle, GabrieleCirillo, and crisostomi on GitHub. They added an incredibly sophisticated fp8 implementation on the MLP down proj fwd pass. 280 lines that include putting the activation quantization step into the prior kernel, delayed amax scaling, and fusing the weight quantization and transpose together. github.com/KellerJordan/modd…

New NanoGPT Speedrun WR at 74.6s (−0.8s), by Jan Varho (jvarho on GitHub). The clever idea: if the target is the token " bananas", give partial credit for predicting " banana". More generally, during early training, add an auxiliary loss on the longest token-prefix of each target token. github.com/KellerJordan/modd…

New NanoGPT Speedrun WR at 76.0s (-0.3s) from aryavohra on Github, with an algebraic rewrite of XSA. Equivalent math, faster execution. 220 more of these and NanoGPT will be sub 10 seconds. github.com/KellerJordan/modd…

New NanoGPT Speedrun WR at 76.3 (-2.9s) from @Lisennlp , with a radical boost to the final attention layer that drops step count by 7%, called lightweight Dynamically Composable MHA. A secondary local window (112 tokens) is recomputed on the same q/k. After softmax it mixes the per-head attn maps into one unified map, then applies the shared map to each head's values with a per-head scale. Every mix and scale operation is computed per-token from a MUDD MLP. All in a custom kernel. Takes inspiration from Dynamically Composable MHA arxiv.org/abs/2405.08553. github.com/KellerJordan/modd…

Question for the LLM Research Community: Is anyone aware of fully reproducible experimental results showing that MLA beats GQA under the same KV cache? I ask for two reasons: (1) I find it surprising that there is still a divide between Chinese labs using MLA, and Western labs using GQA + sliding window, when certainly many ablations have been run by many parties. (2) At Marin we are considering MLA for our next large scale run, but in early (!) smaller scale ablations it appears worse than heavily tuned feature-rich GQA, even after controlling for KV cache. I would like to run better experiments here. Below I cover my thoughts on general reproducibility, then specifics on MLA. Every empirical result in ML is only contextually true. Conditioned on the data distribution, optimizer settings, model width, model depth, finer architecture details, hardware, kernel engineering, initialization, token count, tokenizer, context length, and evaluation protocol, one can reach different conclusions. Contextual results are still useful. Typically if I see a promising method, I will first attempt a full 'context jump', where I apply it to my own context, hoping results transfer. Sometimes they do. If they don't, I can try 2 things: modify the implementation of the method, or modify the context. Ideally I have access to the full context of the original result. Then I can perform a 'context bridge', where I ablate one aspect of the context at a time, isolating exactly why a method performs differently. This lets me make an informed decision to either update my context to let the method shine, or stick with my context and leave the method out. MLA is tricky to assess at small scale. A core aspect of MLA is compressing hidden_dim->latent_dim. Then for each head, latent_dim->head_dim. Typically head_dim is fixed at 128, partially for hardware reasons, and partially for learning dynamics (head_dim of 8 wouldn't have sufficient representational capacity). To get MLA dynamics, you want hidden_dim>>latent_dim, and latent_dim>head_dim. This window closes at small scale. The degree of tuning can unfairly alter the scales. In GQA we have partial RoPE, QK Norm, Gated Attention, attention sharpening, sliding window, and other techniques that give a 30%+ training boost. They don't seem to give the same boost to MLA. On one hand, you want to compare techniques apples:apples with equal tuning. On the other hand, there is a finite amount of future tuning you can do, so prior tuning influences which approach is most pragmatic. Creating controlled tests between MLA and GQA is tricky. Several factors: kv_cache, quadratic attention flops, attention projection flops. kv_cache is controlled by scaling down kv_heads to match MLA, or scaling up kv_latent to match kv_heads. quadratic attention flops are controlled by scaling up GQA's query head count to match MLA head count, or scaling down MLA head count. Also scaling up GQA head_dim 128->192, or scaling down MLA head_dim to 192->128. In general, it's informative to context match to both option A's preferred context and option B's preferred context. Sliding window is another confounder. MLA is theoretically elegant, if we ignore RoPE. It replaces the 'replicate' op of kv_heads in GQA with a 'mix' op (pic below). Since the 'mix' can learn to 'replicate' if it wants, MLA is purely more expressive, and the cost of 'mix' is hidden at inference with absorb trick. Yet in practice, I find that at small scale this 'mix' op doesn't add much value and interacts poorly with the optimizer dynamics. And the change to RoPE hurts. My current plan is to first tune and ablate our model features around MLA, then run 3 scaling ladders: MLA, GQA with 2 kv_heads, and GQA with higher kv_heads. For each ladder, fit a loss vs compute projection. If MLA performs worse at our target compute compared to both GQA options, drop it. If MLA beats 2 kv_heads but loses to higher KV_heads, then it becomes a kv_cache tradeoff. Early results indicate MLA will perform worse than both feature-rich GQA ladders, but we will see. Any positive external reproducible results for MLA would help make sure I give it the best chance possible.

Also, the post has a math duel from a Renaissance beef over cubic equations. Anyway, here’s our latest architecture and data mix from a run we kicked off over the weekend. Looking like a pretty good run so far. Entire recipe is public! Wandb: wandb.ai/marin-community/mar…

New NanoGPT Speedrun WR at 79.2 (-0.5s) from @sisovicm, by moving the MLP up projection to FP8 on the forward pass, and increasing step count by 5 to offset the loss increase. The march towards lower precision continues! github.com/KellerJordan/modd…