News · · 8 min read

AWS TorchNeuron: Trainium Becomes a Native PyTorch Device, With No Trn2 in Seoul

In this post
  1. The old XLA path and the new code
  2. What the first release leaves out
  3. One customer, three numbers
  4. No Trn2 in Seoul
  5. What teams in Korea should weigh
  6. Little reaction so far
  7. Limitations
  8. Sources

AWS announced TorchNeuron on October 19, 2026. It is an open-source backend that registers Trainium, Amazon's AI chip, as a first-class PyTorch device named neuron, and it goes live as an Apache 2.0 public beta at PyTorch Conference North America on October 20. What changes is the code. Running PyTorch on Trainium used to mean the PyTorch/XLA pattern: an XLA device, a mark_step() at the end of each step, and XLA-specific optimizer helpers. With TorchNeuron, .to('neuron') takes their place. The catch is hardware: TorchNeuron runs only on Trn2 and Trn3 instances, and the AWS EC2 per-region instance list shows no Trn instances at all in the Seoul region.

This post covers the old and new code side by side, what the first release leaves out, three different numbers attached to the same customer story, and what the region gap means for teams in Korea.

The old XLA path and the new code

The previous Trainium PyTorch package (older torch-neuronx versions) was built on PyTorch/XLA. The Neuron programming guide calls this "lazy mode": tensors are placeholders that build a graph, and the graph compiles and runs when xm.mark_step() is called. Every training step needed a mark_step(), and calling loss.item() or print(loss) mid-step triggered an extra compile-and-execute. For training, that path has already ended. According to the Neuron docs, PyTorch/XLA is no longer supported for training on Trainium starting with Neuron 2.30, and PyTorch 2.9 was the last XLA-based version for training.

Putting the same single-device training loop in both styles shows the difference. We assembled both snippets from the Neuron programming guide and the PyTorch blog post so they run the same loop. We did not run them on a Trainium instance.

# Before: PyTorch/XLA-based torch-neuronx
import torch_xla.core.xla_model as xm

device = xm.xla_device()          # XLA device
model = MyModel().to(device)
for batch, targets in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch.to(device)), targets.to(device))
    loss.backward()
    optimizer.step()
    xm.mark_step()                # graph compiles and runs here

# TorchNeuron
model = MyModel().to('neuron')    # the line that changes
for batch, targets in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch), targets)
    loss.backward()
    optimizer.step()

The second version has no torch_xla import and no mark_step(). The PyTorch blog annotates that line as "the only change." Its example does skip moving the inputs to the device, though, so real code still needs the batches placed on neuron somewhere in the data pipeline.

The remaining differences are a list, so here they are as a table. The left column follows the archived Neuron docs; the right follows the PyTorch blog and the GitHub README.

ItemOld PyTorch/XLA pathTorchNeuron
Devicexm.xla_device() or 'xla''neuron'
ExecutionLazy; compiles and runs at mark_step()Eager by default; torch.compile(backend="neuron") optional
Data-parallel trainingxm.optimizer_step(), MpDeviceLoader, 'xla' backendStandard torch.distributed (DDP, FSDP2, and so on)
InstancesTrn1, Trn2Trn2, Trn3

The reason it works this way is how it registers. PyTorch has an official extension point for out-of-tree hardware, PrivateUse1, and TorchNeuron uses it to register Trainium. AWS says that is why torch.nn, torch.autocast, torch.distributed, and torch.profiler work on the neuron device as written. In Hugging Face Transformers, device_map="neuron" or device_map="auto" loads a model onto Trainium.

What the first release leaves out

"One line" holds only when a model fits inside the first release. The limit most likely to bite, from the README's known-limitations list, is static shapes. torch.compile supports only dynamic=False, and every new input shape triggers a recompile. The README suggests padding or bucketing inputs to keep the number of shapes down. For LLM serving, where text length varies on every request, that looks close to mandatory.

Numeric precision has conditions too. float64 runs as float32, and int64 support depends on the operation. The two sources describe low-precision training differently. The PyTorch blog says TorchAO and FP8 are in private beta and will be generally available in the next release. The README says TorchAO integration is not available in this release.

AWS says it measured compatibility in two ways: broad model coverage across Hugging Face and TorchTitan models, including a continuous sweep of 100,000 Hugging Face models that cover 99.9% of real-world usage by download count, and the PyTorch upstream test suites. Figure 1 shows how the upstream test outcomes break down.

Figure 1. TorchNeuron PyTorch upstream test outcomes, compiled from AWS figures in the PyTorch blog

The August 2026 run produced about 154,000 test outcomes, of which about 69,000 were skipped. About 30,000 were skipped by PyTorch's own eligibility checks, such as CUDA-only or CPU-only tests. The other roughly 39,000 were deliberate Neuron skips, and the largest group, about 22,000, covers operators that do not have a native Neuron lowering yet. The roughly 85,000 tests that ran is our own subtraction of 69,000 from 154,000. The post's text does not state a pass rate as a number, so the figure does not show one.

The blog also lists features the first release deliberately leaves out, including sparse tensors, NestedTensor, forward-mode autograd, dynamic and symbolic shapes, and per-channel quantization. AWS says customers did not ask for them.

One customer, three numbers

The headline customer is Deepgram, a voice AI company. Three pieces published the same day describe its results with three different numbers.

Who said itWhereNumberBaseline and conditions
Deepgram CTOPyTorch blogAbout 60%In the benchmarks Deepgram evaluated, versus existing GPU baselines; input audio seconds per dollar
AWS solutions architects and othersAWS Machine Learning blog11–70%Versus compute Deepgram had been running; offline batch; trn2.3xlarge; instance pricing at evaluation time
Colin Brace, AWS VPTechCrunch interview70%No conditions given beyond "production-grade accuracy"

The three numbers look like different summaries of the same measurement. The AWS blog's 11–70% is the full range, and 70% is its top end. The test was throughput-oriented offline batch work, not real-time serving. All three numbers come from AWS or its customer. We found no third-party measurement of TorchNeuron performance at the time of writing.

According to the AWS blog, this case did not end with a one-line change either. Deepgram and AWS replaced some layers with pre-built kernels from NKI-lib and modified kernel source. For serving, they removed Python and PyTorch from the path and had a Rust runtime call ahead-of-time compiled executables directly. Deepgram's CTO, in the PyTorch blog, called the port a multi-month engagement. The numbers above come after kernel-level tuning, not from a single setting.

No Trn2 in Seoul

The README requires a Trn2 or Trn3 instance. Trn1, which ran the older package, is not supported. That narrows TorchNeuron to regions with Trn2 or Trn3. We pulled the Trn entries from the AWS EC2 per-region instance list, alongside regions a team in Korea might consider.

RegionTrn2Other Trn/Inf types listed
US East (Ohio)YesTrn1, Trn1n, Inf1, Inf2
Asia Pacific (Hyderabad)YesTrn2u
Asia Pacific (Melbourne)YesTrn1
South America (São Paulo)YesInf1, Inf2
Asia Pacific (Seoul)NoInf1, Inf2
Asia Pacific (Tokyo)NoInf1, Inf2

In that list, these four are the only regions with Trn2. Trn3 UltraServers were announced on December 2, 2025, but they do not yet appear in this region list. Among Trn and Inf types, Seoul and Tokyo list only the inference chips Inf1 and Inf2. The list may not match what the console offers, so check the console before planning around it.

What teams in Korea should weigh

What TorchNeuron changes is the cost of trying Trainium. Before, a trial meant rewriting the training loop for XLA and learning lazy execution. Now a team can leave its PyTorch code mostly as it is, run it eagerly, and first check which operations fall back to the CPU. For teams under GPU cost pressure, the barrier to a comparison experiment is clearly lower.

Training and serving look like separate decisions. Training can run in an overseas region such as Ohio if the data is allowed to leave Korea. Real-time serving is different: the missing Trn2 in Seoul affects both latency and reviews of cross-border personal data transfer. For services that must stay in the Seoul region, Inf2 and the existing Neuron inference stack appear to remain the practical option for now.

The static-shape limit also weighs more on serving. Fixed-size models such as image classifiers are barely affected. Variable-length text generation and speech recognition need padding and bucketing designed up front. That Deepgram's results came from throughput-oriented offline batch work points the same way.

The partner list may include a Korean company. The PyTorch blog lists Nota AI among partners building with TorchNeuron to grow the Trainium ecosystem. Nota, an AI model compression company that listed on KOSDAQ in November 2025, goes by the same English name. The announcement does not say where this Nota AI is based or what it is doing with TorchNeuron, so we cannot confirm they are the same company.

Little reaction so far

The announcement is a day old, and developer communities have barely responded. The PyTorch blog post was submitted to Hacker News, but when we checked it had 2 points and no comments. We found no meaningful discussion on Reddit or X either.

Most opinions so far come from the announcing parties and the press. Meta's Jana van Greunen, who leads PyTorch engineering, said in the PyTorch blog that "developers should have a consistent experience across hardware platforms." Hugging Face's Jeff Boudier said TorchNeuron makes Trainium a native backend for Transformers and Diffusers. Both are launch partners.

The skeptical take comes from TechCrunch. The reporter noted that one of the points of PyTorch is to hide the difficulty of running AI code on hardware, and wrote: "We'll have to wait and see if pushing Nvidia out that far gives developers more choice and more affordable options, as Amazon promises, or just more configuration headaches." The same piece points out that AMD GPUs and Apple chips can already run PyTorch with a configuration change.

Limitations

ItemDetail
Hands-on runsWe did not run code on a Trainium instance. The code comparison is assembled from official examples
Performance numbersAll from AWS and its customer. No third-party measurement
Pass rateThe blog text gives no upstream pass rate as a number, so none is reported
Region dataBased on the EC2 docs list. Console availability may differ
TimingAs of the first day of the public beta. Coverage changes with each release

The code comparison here shows where the lines change, and no more. Whether a given model runs after a few edits in the first release depends on the operators and input shapes it uses, and this post did not measure that.

Sources

1 / 1