🚢 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:

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.

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.

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.

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

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:

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!

New NanoGPT Speedrun WR at 79.7 (-1.5s) from @TrianX , with a brilliant solution to hash collisions on the bigram hash embedding. Instead of every bigram in a bucket returning the same embed, a secondary hash gives each bigram its own ±1 sign pattern (one of 8192), applied element-wise, e.g. x·[1,−1,1,1,…]. Each bigram in the bucket then reads a different partial reflection of the one stored row.

Quoting @dlwh : we are at risk of losing the reputation of spiky loss runs! This run incorporates some stability techniques from my past projects: Hyperball, Gated Norm, and Gated Attention. Excited to see the next run from Marin!

Building momentum at Marin! Upgrading from Dense -> 129B parameter MoEs -> architecture improvements -> optimizer improvements gives our pretraining recipe an estimated 6x cumulative learning speedup, accounting for MFU. Includes community contributions.

AI Agent literature/web review can get much better. It was peculiar to see how under-the-radar the NanoGPT Speedrun was for agents during Parameter Golf. Many objective improvements, like faster RopE, were not copied. SmearGate was copied incorrectly, and only fixed after a month. Several others were copied in the last couple days, often by the original speedrun author. Even the attributions were not aware of the NanoGPT origins.

Researchers' brilliant ideas often get lost in the sea of endless SOTA claims on weak baselines. At Marin we battle-test ideas in an open arena, where anyone's idea can be promoted to the next hero run. One that recently rose up was @Jianlin_S MoE Quantile Balancing, used in our last 1e22 and ongoing 130B run. Animated visuals of how QB performed are available in the OpenAthena blog.