2026-10-10
In September I entered the AWS Trainium Frontier, a competition where you get one Trainium chip and 30 minutes to train the best language model you can. The score is bits per byte on held-out web text. I went from the baseline's 1.259 down to 0.959 over 106 versions of the training script, and most of what I learned along the way came from the times I was confidently wrong.
I had never used Trainium. It's AWS's own ML chip, with a 128 by 128 systolic array, a software-managed on-chip scratchpad, and a Python kernel language called NKI. I also didn't write this one by hand the way I did my diffusion model. The competition explicitly welcomes AI-assisted research, so I ran it as an experiment in that: I came up with the ideas and decided what to try next, and Claude Code agents wrote the code, ran the jobs on the dev box, and kept notes. By the end there were about 750 logged runs and a folder of findings longer than anything I've written myself.
That split turned out to be a skill of its own. The agents were fast and tireless, but they were also very good at producing a clean explanation for a number, and a clean explanation is not the same as a correct one.
The first big lesson was that the shape of the network almost doesn't matter in a 30-minute budget. We measured every width from 512 to 896 and every depth from 7 to 16, and they all landed within about 0.002 bits/byte of each other. A bigger model learns more per step, but it also takes longer per step, and on this chip those two effects nearly cancel. The math came out to about 0.0024 bits/byte per e-fold of parameters at a fixed time budget, which is basically a flat ridge.
What did matter was making steps cheaper. Getting torch.compile working on Neuron gave a 1.5x speedup on day one. Moving the hot path into NKI kernels (RMSNorm, RoPE, the MLP, the output head, the optimizer steps) made steps cheap enough that adding layers finally paid off, which it hadn't at any earlier speed. The single cleanest win of the whole project was finding that the embedding update was spending almost 80% of its work on padding, even though a comment in the code claimed otherwise. Fixing it was worth about 0.0009, and two runs agreed to five decimal places.
Two runs of the exact same file could differ by 0.0003 bits/byte. A lot of the changes I cared about were smaller than that. Twice I got excited about a hyperparameter change from a single run, and both times the second run cut the gain roughly in half. After that the rule was simple: anything under about 0.0005 needs at least two runs on each side before it counts.
My favorite mistake: one architecture change looked seven times slower, and the agent wrote a careful analysis of why the extra operations between the fused layers were too expensive to ship. That analysis then got used to rule out a couple of other ideas. Days later it turned out the change wasn't slow at all. Two NKI kernels had silently failed to load and fallen back to slow PyTorch code, and nobody had checked. We withdrew the conclusion and everything built on it, and "grep the log for fallbacks" became the first step for any slow result.
Submissions go through an automated reviewer that checks the code against the rules. One of mine got rejected, and the reason was embarrassing. Over 26 versions the file had grown from 5 environment-variable reads to 50, one convenient knob at a time, so the architecture of the scored model could in principle change without the code changing. The file also had a leftover branch that scaled the logits only at eval time, which directly contradicted a note at the top of the same file saying it didn't do that. Neither was an attempt to cheat, but a reviewer can't tell that from the code. The fix was to hard-code everything and delete the eval-only branch, and the next version went through.
The biggest late gain came from somewhere I didn't expect. With a few days left, we tested how much a plain n-gram model, built only from the training data, could add on top of the trained network. Offline, mixing it in was worth 0.0085 bits/byte, which was more than everything else from the previous week combined. A lot of it came from near-duplicate web pages, things like press releases and plant databases, where an order-16 context has been seen before and the n-gram is far more confident than a small network can be.
The version that shipped builds the n-gram tables from the training shards in a background thread while the network trains, then adds the n-gram's top prediction to the logits as a learned bias. On the leaderboard's private shard about 90% of the offline gain carried over. The final submission was a 14-layer network with that n-gram mix, at 0.959 on the public shard.
I'd look for the n-gram earlier, since it was sitting there the whole time. I'd also be stricter with the agents from the start: every claim backed by two runs, every slow result checked for fallbacks, and the submission file kept boring on purpose. The agents made it possible to try far more things than I could have alone, but the useful part of my job was saying "I don't believe that yet."
In the end I finished in the top 15 of 100 teams. Only the top 10 advanced to Phase 2, and I sat in 10th for a while before three teams showed up late and pushed me just below the cut. That stung a bit, but for a month on hardware I'd never touched, I'll take it.