Abstract
Trainium of Thought is my entry in the AWS Trainium Frontier Competition, which is run by AWS Annapurna Labs on Devpost. Phase 1 fixes the hardware (a single AWS Trainium2 chip), the dataset and the wall-clock time. Entries are ranked by the validation bits per byte of the language model that exists when 30 minutes of training have elapsed. Under this rule the number of optimizer steps is not fixed, so each design decision affects both how much the model learns per step and how many steps fit in the budget.
This project reduced validation bits per byte from 1.3951 to 0.9656 in full-budget local runs. The best submitted model scored 0.9848 and held first place on the public leaderboard at its peak. Phase 1 is still open, and the entry is currently second on the public leaderboard (September 2026). Most of the improvement came from two changes that depend on each other: a wider and shallower model, and hand-fused kernels that shorten each training step. Several modifications that improved loss at a fixed step count produced a worse final model, and these negative results are reported.
01Problem statement
The competition has two phases. In Phase 1, which closes on September 30, 2026, each team of one to four people trains a language model on a single Trainium2 chip within a strict 30-minute wall-clock budget. Submissions are ranked on a public leaderboard by validation bits per byte, which the organisers recompute from the model's logits with a trusted evaluation function. The ten highest-ranked teams advance to Phase 2, where they receive full Trn2 servers for scaled training. The prize pool is $40,000, and the winners are invited to present at NeurIPS 2026.
The usual comparison between training recipes asks which one reaches the lowest loss after a fixed number of steps. Here the fixed quantity is time. A modification that lowers the loss at step 1,000 can still lose if it adds 20 ms to every step, because the run ends before the step at which the modification would have paid off. The objective is therefore the validation loss at the end of a complete 30-minute run, measured through the official scoring path. Short runs were used only to reject clearly unsuccessful ideas.
02Experimental protocol
- Profile one training step and attribute its time to operators.
- Change one thing, either the architecture or the system, and never both in the same experiment.
- Run a short gate that checks loss and throughput.
- If the change passes the gate, run the full 30 minutes. Only this result is recorded as the outcome.
- Build the submission artifact and confirm that its score matches the local run.
Every run was recorded in a ledger with its architecture, batch size, optimizer settings, kernel set, seconds per step, steps completed and final score. Kernels were first developed as isolated microbenchmarks and were always measured again inside the full trainer, because an operator that is faster in isolation can still increase step time once synchronisation and host overhead are included.
03Final configuration
| Aspect | Choice |
|---|---|
| Model | Width of about 1024, eight layers, attention in a subset of layers, U-Net-style skip connections |
| Parallelism | Four processes, 262k tokens per effective batch |
| Optimizers | Muon for the large matrices; AdamW for embeddings, the output head and small parameters; warmup and cooldown set to the number of steps that fit |
| Precision | BF16 state; FP8 on the projection paths where the numerical error remained acceptable |
| Fused kernels | Normalisation, RoPE, cross entropy, ReLU², flash attention, embedding scatter and parts of the optimizer step, written against the AWS Neuron Kernel Interface |
| Launch overhead | Small table preparations batched; residual operations fused where an end-to-end profile showed a benefit |
04Results
| Checkpoint | val bpb | Change | Measurement |
|---|---|---|---|
| Initial submission | 1.3951 | Reference recipe | Starting point |
| Submission 5 | 0.9848 | Wider and shallower model, first fused kernels | Leaderboard; first place at its peak |
| D62 | 0.9679 | Embedding scatter kernel, batched Muon tail | Local |
| D70 | 0.9669 | FP8 projection path | Local |
| D72 | 0.9656 | Batched FP8 tables, fused RoPE | Local |
Local validation, the built artifact and the public leaderboard are three different measurements. Packaging and runtime differences produced measurable offsets between them, so they are reported separately.
05Negative results
A faster kernel that slowed the trainer
A BF16 matrix-multiplication kernel that kept weights resident was 6 to 18% faster than the stock operator in microbenchmarks. Inside the trainer it increased step time from about 0.715 s to 0.870 s, because synchronisation and integration costs exceeded the saving.
Gating improved loss per step and reduced the number of steps
Attention gates and smear gates gave the best loss at a matched step count. They also added enough small operations per step that the run completed fewer updates and ended with a worse model. Moving the multiplications into kernels did not recover the loss, because slicing small gate inputs and launching many small operations consumed the gain.
More capacity, worse final model
A 203M-parameter model with value embeddings learned faster per step, completed 11% fewer steps and finished behind. A ten-layer model showed the same pattern.
Packaging errors
One built candidate ran 7% slower because the builder omitted a flash-attention flag. An earlier artifact passed locally and failed in the scorer because a module alias recursed and fell back without reporting an error. The submission builder now has its own parity tests.
06Conclusions
- Under a time budget, the quantity to optimise is the score at the end of the budget. Proxies at a fixed step count are misleading.
- Learning per step and time per step must be reported together for every experiment.
- An end-to-end profile of the trainer is more reliable than any microbenchmark.
- The code that builds the submission requires the same testing as the training code.