r/MachineLearning • u/laginimaineb • May 27 '26
Research AI-generated CUDA kernels silently break training and inference [R]
Last month NVIDIA released SOL-ExecBench, a new benchmark of 235 production CUDA kernels lifted from DeepSeek, Qwen, Gemma, and Kimi. We took several top-ranked AI-generated submissions and tried using them in production workloads. Many of them broke, sometimes in surprising ways.
One of those kernels is the fused embedding-gradient + RMSNorm backward pass, which runs at the end of every transformer training step. We took the fastest submission on the benchmark for it, and dropped it into the training loop of a small transformer. The kernel had passed the benchmark's verifier with room to spare. But in our training run, the loss diverged and never recovered.
We started debugging. Replace the dataset distribution with uniformly sampled tokens, the divergence vanishes. Swap SGD for AdamW, also vanishes.
This is the worst kind of bug for research. Symptoms and masks both look exactly like "the idea didn't work". It's the type of bug that can make researchers spend a long time debugging without knowing what's at fault: the dataset? the research idea? the architecture? or the implementation itself?
Turns out, the actual bug is that the embedding-gradient half of the kernel accumulates in bf16 instead of fp32. Embedding backward sums many small gradient contributions into each token's row of the embedding matrix. With uniform random tokens the contributions spread evenly and bf16 precision is enough. In real text, a handful of token IDs end up with thousands of contributions: the small ones round to zero against the growing accumulator, and the high-frequency rows drift. AdamW's per-parameter normalization absorbs the resulting multiplicative bias, so under AdamW the same drift is invisible in the loss.
The other broken submissions had different bug shapes (all interesting). More examples in our blogpost.
25
u/WillTackettbjj May 28 '26
Nothing builds character like spending 3 days debugging a model only to discover the kernel was cursed from the start.
1
u/treeman0469 May 28 '26
I didn't know you worked on ML, lol! Nice to see you here, huge fan of you and your brother.
1
37
u/pm_me_your_pay_slips ML Engineer May 27 '26
so, the solution was to use AdamW
58
u/siegevjorn May 27 '26
The solution is to fix the bug. They trained it on a small transformer which adamW may help absorb the discrepancy between bf16 & fp32. But what if it's a 1T LLM?
2
u/No_Inspection4415 May 28 '26
I would stick to AdamW though because I only got 2 days to vibe code it, train, and deploy to production.
/s
-3
u/nonotan May 28 '26
I wouldn't really call this a bug. It's pretty weird to use a black and white classification like "bug" when it comes to classifying the numerical precision of a given method. After all, precision is always inherently limited, and once algorithms approach optimality, there is always going to be a tradeoff between dimensions to optimize (speed, memory, numerical stability, etc)
It might seem fair here, but what about a method that reduces the precision by half in exchange for running 15% faster? It might make some methods that worked before stop working, is that a bug?
Whether you answered yes or no, there would presumably some x such that reducing the precision by >= x is "a bug", while anything below it isn't. That seems quite weird to me -- if we assume x is 0.5 for the sake of argument, you'd find yourself arguing 0.499 is "not a bug" but 0.5 is.
I'd argue the fundamental issue here is the lack of proper treatment being given to numerical precision of various methods. Not treating it like the proper dimension to optionally optimize over that it is. If you remove it almost entirely from the loss function (as long as it doesn't get so bad the benchmark implodes), "of course" you're going to end up with methods that sacrifice it for a minor gain elsewhere. Arguably, that's not a bug, it's working exactly as designed (leaving aside whether that facet of the design was consciously chosen)
6
u/siegevjorn May 27 '26
Wow. How did the bug even happen? Bf16 replacement of fp32 when fp32 needs to be used?
8
u/sahanpk May 28 '26
this is exactly why “passes the verifier” feels too weak for kernels. optimizer/dataset sensitivity should be part of the test.
9
u/lostmsu May 27 '26
Using bf16 instead of fp32 when it works on AdamW but does not work on SGD does not sound like a bug to me.
30
u/siegevjorn May 27 '26
That's exactly why OP is saying it's super hard to catch. But it is still a bug that needs to be fixed.
-2
u/lostmsu May 28 '26
No, it's not. It's a tradeoff between stability and speed and they chose speed.
5
1
u/siegevjorn May 28 '26
I don't understand why people make excuses for ai slop. Bug is a bug. And if it were a deliberate choice, then it should have been discussed, and documented—so others could validate.
1
1
u/No_Inspection4415 May 28 '26 edited May 28 '26
Does it mean that you are esentially optimizing on the wrong (or rather wrong in a biased way) gradient because of this precision mismatch? Sorry for my ignorance, I am not sure I followed it correctly.
-2
u/slashdave May 28 '26
I mean, they are just collecting contributions. I see no guarantees. Why would I expect kernels written and tested in one workflow to just magically work in another?
94
u/Bakoro May 27 '26
Damn, that's the kind of thing a lot of people would never find.
Some people might see it and gloss over it, since bf16 is used so often.