Llama 3.1 8B on AWS Trainium and Inferentia.
One trn1.2xlarge fine-tunes the model, one inf2.xlarge serves it, one trn2.3xlarge repeats the fine-tune, and every number has a committed log in torchneuronx. The GPU points come from MI300X-vs-H200, which used the same request shape and metric schema. The conditions that differ are listed under the first figure. The talk covers the same material.
Serving one 8B model, and what a token costs
$ per 1M output tokens, hourly rate divided by tokens per hour. On-demand list prices.
Inferentia2 stops at 32 streams because 32 GB of device memory holds about 48 resident sequences at 2048 context, and the KV budget is spent. Its decode latency stays flat across the sweep, 63 to 71 ms per token, while aggregate throughput grows 26x. Compared at 32 streams, an Inferentia2 token cost 2.4x an MI300X token and 1.8x an H200 token. At each side's best point the gap was 5.7x and 4.6x. No GPU ran on AWS in either study, so the like-for-like rental (g6e.xlarge, one L40S, $1.86/hr) is the open comparison.
| What differs | Neuron study | GPU study |
|---|---|---|
| vLLM | 0.16 on Neuron, the newest DLAMI that boots on NeuronCore-v2 | 0.26 on the GPUs |
| PyTorch | 2.9.1 | 2.11.0 |
| Cloud | AWS us-west-2, on-demand list price | DigitalOcean (MI300X), Nebius (H200), published GPU-hour rates |
| Device memory | 32 GB | 192 GB (MI300X), 141 GB (H200) |
| Run date | 2026-07 | 2026-08 |
Training on one Trainium1 and one Trainium2
Steady-state tok/s. Hover a bar for MFU.
| Objective, one Trainium1 | Status | Measured |
|---|---|---|
| SFT, LoRA Llama 3.1 8B | works | 2,952 tok/s, 68.3% MFU at seq 2048 |
| ORPO Llama 3.1 8B | works | 1,181 tok/s, 30.2% MFU at length 1024 |
| Pretraining 362M, SmolLM2 shape | runs, unresolved | 4,573 tok/s, 7.0% MFU; the hand-written XLA loop recompiles every step |
| DPO Llama 3.1 8B | unresolved | the reference forward compiles outside the step and the lane dies in a host transfer |
| GRPO, RLVR Qwen3 1.7B | blocked | the training model class has no generate() |
Speculative decoding on one Trainium1
Compile cost, quality and stability
Boot time to first token, inf2
Every graph is compiled ahead of time. Same graphs, new weights: nothing recompiled.
Held-out loss after the fine-tune
| Trainium1 | 2.149 to 1.251 |
| Trainium2 | 2.148 to 1.265 |
Byte-identical held-out rows on both chips. Starting losses agree to three decimals, which is what makes the two runs comparable.
30 minutes of sustained load, inf2
- Retention vs first iteration
- 100.4%
- Throughput, first to last
- 118.8 to 119.3 tok/s
- First-token p99 spread
- 1.6%
Concurrency 8, seven iterations. No thermal or stability droop at this load.
Where it fit, on this evidence
| LoRA or full SFT, supported architecture, static shapes | strong fit |
| Preference optimisation without a reference model (ORPO) | works |
| Pretraining a small model from scratch | works to about 400M parameters on one small instance |
| Online RL of any kind (GRPO, PPO, RLVR) | blocked |
| Cost-optimised 8B serving | a single GPU was 1.8x to 4.6x cheaper per token here |
| Serving where capacity or data residency decides | reasonable, inside the concurrency ceiling |
| Architecture outside the exporter list | a wall, not a tuning problem |
Most of the study's walls were toolchain, not silicon: an exporter allowlist, ahead-of-time compilation per tensor shape, and a training class without generate(). The instances were terminated on 2026-08-26; the analysis re-runs from the committed results with no AWS account.