Does Your Model Run on AWS Trainium? What to Check Before Committing

If your model is a Transformers text-generation or image-text-to-text model small enough to fit on four Trainium2 devices, you can now check whether it runs unmodified on AWS Trainium with a lookup instead of a porting spike. Hugging Face has published the torch-neuronx model support Space, a public pass/fail tracker for popular Hub models that reports, per model, how many Trainium cores it is sharded over, how long the run takes including compilation, and how much device memory it peaks at. One caveat belongs up front: every result in this tracker is a test run by Hugging Face itself on trn2 hardware, with no independent replication in the available evidence, so treat the numbers as vendor-side measurements rather than audited benchmarks. Even with that caveat, the tracker changes the first step of the GPU-versus-Trainium decision from “schedule a multi-week evaluation” to “search a list,” for the model classes it covers.

What it does not change is everything after the first step. A PASS certifies that a specific generation script compiled and produced a checked 32-token output. It says nothing about throughput, latency, or long generations, and the tracker’s own documentation is explicit about that. This article walks through what the tracker actually tests, how to read a row, how to reproduce a result yourself, and which verification work remains yours before you commit inference capacity to Trainium.

The question the tracker answers, and the one it doesn’t

The tracker’s stated purpose, per the Hugging Face announcement, is to answer a narrow question: “whether a model runs unmodified on Trainium,” with the script to run it yourself attached. That framing matters because “runs” has a precise, bounded meaning here. A green row means a model loaded, compiled through AWS’s neuronx-cc compiler, and generated tokens that passed a check. It does not mean the model meets any serving objective.

This distinction is where teams have historically burned time. Evaluating a new accelerator usually starts with a porting spike: engineers stand up an instance, fight the SDK, patch model code, and only then learn whether the model works at all. The tracker inverts that. Compatibility discovery happens before you spend anything, and the reproduction script means the discovery is auditable rather than anecdotal.

The question the tracker does not answer is the one that actually decides production deployments: will this model hold your latency and throughput SLOs at your batch sizes and sequence lengths, at a price that beats your GPU alternative? Answering that still requires your own profiling. What has changed is that you now reach that profiling stage only for models already known to compile and run, rather than discovering incompatibilities mid-evaluation.

What is actually tested: task scope, trn2 setup, and the 16-core ceiling

The tested scope is narrower than the phrase “popular models on the Hub” suggests, and reading it correctly is the first decision gate. According to the announcement post, only Transformers models are tested, and the selection is limited to text-generation and image-text-to-text tasks. Embedding models, rerankers, speech models, diffusion models, and non-Transformers architectures are outside the current coverage. If your workload lives there, the tracker has nothing to tell you yet, and the old evaluation cost applies in full.

The second gate is size. The tracker tests “only models that fit on 4 Trainium2 devices… which means 16 cores at most.” Each Trainium chip integrates two NeuronCores, per the NeuronMLP paper, so four devices is eight chips. A model that needs more than 16 cores to hold its weights is simply not in the test population, regardless of whether it might work with more aggressive sharding.

The test setup itself is documented in unusual detail. Each test loads the model on a trn2 instance, sharded over the smallest number of Trainium cores that can hold its weights, per the announcement. Most models are sharded with tensor parallelism; mixture-of-experts models without a tensor-parallel plan use expert parallelism only. Then the test generates 32 tokens and checks them.

Gate	Covered	Not covered
Task type	Text-generation, image-text-to-text	Embeddings, reranking, speech, diffusion, other tasks
Architecture	Transformers models	Non-Transformers architectures
Size	Fits on 4 Trainium2 devices (16 cores)	Anything needing more than 16 cores
What PASS certifies	A checked 32-token generation compiles and completes	Throughput, latency, long generations, SLO compliance

Reading a row: cores, compile-inclusive wall time, and peak device memory

Each row in the Space reports four things, per the announcement: the result, the number of Trainium cores the model is sharded over, the wall time of the run with compilation included, and the peak device memory. Each of those fields carries a specific operational reading.

The core count is a sizing signal. Since sharding uses the minimum cores needed to hold the weights, the number tells you roughly how much trn2 capacity the model occupies at rest, which feeds directly into instance selection and cost estimation. A model sharded over 2 cores leaves you different fleet options than one spread over 12.

The wall time needs careful interpretation because compilation is included. Neuron-compiled models pay a one-time compile cost before the first token, and that cost is folded into the reported figure. A row with a long wall time is not necessarily a slow model at inference; it may be a fast model behind a slow compile. For a production deployment where compilation happens once at build time, the compile-inclusive number overstates steady-state cost. For a workload that recompiles often, it understates nothing. The tracker does not separate the two components, so you cannot infer steady-state latency from the row.

Peak device memory is the closest thing to a capacity-planning number in the row. It tells you how much headroom the model leaves on the cores it occupies, which matters when you consider batch size, longer contexts, or co-locating models.

How the test works: static cache, torch.compile, and the 32-token check

The reproduction script follows a specific recipe, described in the announcement: it loads the model with from_pretrained, shards when needed, then generates with a static cache and torch.compile(backend="neuron"). Three details in that sentence are worth understanding before you treat a PASS as evidence.

The static cache matters because it pre-allocates key-value cache memory for a fixed generation length rather than growing it dynamically. Dynamic cache growth is a common source of graph breaks and recompilations on compiler-driven accelerators; a static cache sidesteps that, which is part of why the test can run unmodified models at all. The torch.compile(backend="neuron") call routes PyTorch’s compiler infrastructure through neuronx-cc, AWS’s production compiler for Trainium.

The 32-token generation is a minimal end-to-end check: long enough to exercise the generation loop, short enough to keep the test fleet affordable. It is also, by the tracker’s own admission, the boundary of what a PASS certifies.

The compiler itself deserves a word, because a PASS is fundamentally a statement that neuronx-cc handled your model’s graph. That compiler is a production system with known rough edges: researchers who built an auto-generated backend matching its performance “reported these results to the AWS Trainium team, who have acknowledged these issues with neuronx-cc,” per the ACT paper, and AWS adopted that tooling. This is not a reason to distrust the tracker; it is a reason to expect that some models fail for compiler reasons rather than fundamental incompatibility, and that per-version results will move as the compiler improves. The tracker anticipates this by versioning its data, which is covered below.

What a PASS certifies, and why throughput is still on you

The tracker’s own warning is unambiguous, quoted from the announcement: “A PASS says nothing about throughput, or about longer generations.” Take that literally. A model that generates 32 correct tokens may still fall over at 4,000-token generations, may deliver tokens at a rate your product cannot tolerate, and may cost more per request than the GPU instance you were trying to escape.

So the verification work after a PASS looks like this: reproduce the tracker’s run to confirm the result on your own instance (cheap, scripted, hours at most), then run your own serving benchmarks at production-realistic batch sizes, prompt lengths, and generation lengths, and profile where the time goes. The TorchNeuron stack includes system-level profiling tooling intended for exactly this stage; the fetched evidence does not document that workflow in detail, so I will not pretend to walk through it here, but budget for it as a real work item rather than an afternoon.

One inference worth flagging as inference: because compilation is included in the tracker’s wall time and the generation is short, nothing in a row constrains steady-state per-token latency in either direction. If someone on your team reads a PASS row as a performance benchmark, correct them before that reading reaches a capacity plan.

The cost case for checking, with caveats

Why bother with Trainium at all? The economic argument recurs across the research literature, but its numbers are older and narrower than the headlines suggest. The NeuronMLP paper describes each NeuronCore as delivering up to 95 TFLOPS of FP16/BF16 compute, “comparable to NVIDIA A100 GPUs at roughly 60% of the cost.” On the training side, the HLAT paper reports that “the training cost on trn1 is approximately 60% of GPU,” based on pre-training a 7B and 70B model family on 4096 Trainium chips over 1.8 trillion tokens. Both figures come from research papers, not current price sheets: the A100 comparison predates the current GPU generation, and the trn1 figure describes training, not trn2 inference. Treat the “roughly 60%” figures from NeuronMLP and HLAT as directional motivation to evaluate, not as a number to put in a business case. Your actual comparison is your negotiated trn2 price against your actual GPU availability, which for many teams is the binding constraint anyway: scarce GPU capacity is the reason this evaluation is happening at all.

The platform itself is past the experimental stage. The HLAT work shows Trainium operating at 4096-chip scale for serious pre-training, and the NinjaLLM paper demonstrates AWS Trainium and Inferentia2 hosting LLMs as “a compelling alternative to traditional Nvidia GPUs” in a RAG serving context. On the inference side specifically, NeuronMLP reports 1.74× faster inference on Trainium with a 0.03 average accuracy drop, per the paper, though that result is for its own SVD-compression technique, not for unmodified models, so it shows headroom rather than a default outcome.

Reproducing a result yourself

Everything the Space shows comes from the public aws-neuron/torch-neuronx-model-support dataset, with one Parquet file per pair of neuronx-cc and library versions, per the announcement. That versioning is more useful than it looks. Compiler behavior changes between releases, so a FAIL on one neuronx-cc version is not a permanent verdict, and the dataset structure lets you see whether a result moved between versions.

The reproduction path is: find your model in the Space, pull the provided run.py script, launch a trn2 instance with the matching library versions, and run it. Because the script is the same one that produced the row, a discrepancy between your run and the published row is itself informative, pointing at environment or version drift rather than model incompatibility. I would treat reproduction as mandatory rather than optional: the published results are Hugging Face-run with no independent replication, and running the script yourself is the cheapest available check on that basis. It also gives you the compiled artifact and a working environment as a starting point for your own latency and throughput benchmarking, which is where the real evaluation begins.

Verdict: when the lookup replaces the porting spike

If your model is a Transformers text-generation or image-text-to-text model that plausibly fits on 16 Trainium2 cores, check the Space before scheduling any porting work. A PASS plus the provided script converts compatibility discovery from a multi-week spike into a lookup and a scripted reproduction run. That is a genuine change in the economics of the GPU-versus-Trainium decision for covered model classes, because it removes the riskiest unknown at near-zero cost.

Then hold the line on what was certified. A PASS means a 32-token generation compiled and checked, on Hugging Face’s trn2 environment, with no independent replication. Your latency, throughput, and long-generation behavior remain unmeasured until you measure them, and the ~60%-of-GPU cost figures in the HLAT and NeuronMLP papers are research-paper anchors, not current inference pricing. Outside the tested scope, wrong task, wrong architecture, or too large for 16 cores, the tracker cannot help you, and the full evaluation cost returns exactly as before. The honest way to use this tool is as a gate that is cheap to walk through, not as a finish line.