AWS Trainium Frontier
Train a language model from scratch in 30 minutes on one AWS Trainium chip. PyTorch, NKI kernels, AI agents.
The AWS Trainium Frontier is a competition run on AWS's own ML chips. Phase 1 gives every team a single Trn2 chip and exactly 30 minutes of training, then scores the model on held-out web text in bits per byte, where lower is better. You can change anything in the training script: the architecture, the optimizer, the data sampling, and even the hardware kernels, which on Trainium are written in NKI, a Python kernel language. I'd never touched a Trainium chip before this.
Between August 31 and September 30 I took the organizers' baseline from 1.259 to 0.959 bits/byte on the public validation shard, across 106 submission versions and about 750 logged training runs, and finished in the top 15 of 100 teams. The top 10 advanced to Phase 2; I was 10th for a while before three late teams pushed me just below the cut.
The final model. The n-gram tables are counted from the training data during the run, so nothing outside the training loop goes into it.
how I worked
The competition treats AI-assisted research as a first-class way to enter, and that is how I did it. I came up with the ideas and made the calls about what to try next, and AI coding agents (Claude Code) wrote most of the code, ran the experiments on the dev box, and kept a written record of every result. The interesting part of the project was learning how to run that loop well: how to tell a real gain from noise, and how to catch an agent confidently explaining a result that turned out to be a bug.
what's in the final model
- A 14-layer, 512-wide transformer with gated attention, rotary embeddings, and a hashed bigram embedding table with 2^19 rows. It is small on purpose: in a 30-minute budget, step count matters more than size.
- Custom NKI kernels for most of the forward and backward pass, including fused RMSNorm, RoPE, MLP, the output head, and the Muon and AdamW optimizer steps. The vendor's attention kernel stayed, because my own attention kernels were faster on their own but produced NaNs inside the full model.
- A Muon optimizer for the matrices, AdamW for everything else, and an exponential moving average of the weights over the last stretch of training.
- An n-gram model built from the training data while the network trains, with context orders up to 16. Its predictions are mixed into the network's logits as a learned bias. This was the single biggest gain late in the competition, worth about 0.007 bits/byte.
where the gains came from
Roughly in order: getting torch.compile working on Neuron, moving the hot path into NKI kernels so steps got cheap enough to afford more layers, fixing a slow embedding update that was mostly padding, tuning the weight averaging, and finally the n-gram mix. Almost every architecture change I tried was a wash. Width and depth anywhere in a wide range landed within 0.002 of each other, because a bigger model costs nearly as much in lost steps as it gains in quality.
challenges
- Noise. Two runs of the same file differed by up to 0.0003 bits/byte, and many real changes were smaller than that. I learned the hard way to run anything below about 0.0005 at least twice before believing it.
- Silent kernel fallbacks. One architecture change looked seven times slower and got written off as too expensive. It turned out two NKI kernels had quietly failed and fallen back to slow PyTorch code. After that, every slow result got its logs checked for fallbacks first.
- A rejected submission. Submissions are screened by an automated rule reviewer, and one of mine was rejected because the file had slowly grown 50 environment-variable knobs, which meant the scored architecture could change without the code changing. The fix was hard-coding everything again.
- Compiler limits. A DeltaNet-style hybrid layer and several fused kernels hit internal compiler errors or miscompiled to NaN, and some ideas were closed simply because the toolchain couldn't build them yet.
The longer story, with the mistakes, is in the write-up.