Google, Meta and AWS launch 'native' PyTorch backends for TPUs and Trainium

PyTorch developers can now target two more families of AI chips through the framework's own device system, according to posts on the PyTorch blog this week. Google and Meta published TorchTPU, a backend for Google's TPUs, on Tuesday 20 October; AWS announced TorchNeuron, a backend for its Trainium chips, on Monday 19 October and said it would go live as a public beta on 20 October, at PyTorch Conference North America in San Jose.

What changes in your code

Both backends register through PrivateUse1, PyTorch's extension point for custom hardware, so the device becomes a string like "cuda".

• TorchTPU: the team says porting a model should take a change from model.to("cuda") to model.to("tpu"). It adds that "One cannot achieve peak TPU performance with that simple change", because shapes and sharding may be tuned for GPUs. Distributed code also swaps the nccl process-group backend for tpu_dist, and multi-chip runs need two TPU-specific environment variables (topology and interconnect addresses) that standard torchrun does not set. Compiled mode is torch.compile(model, backend="tpu").
• TorchNeuron: the device is 'neuron'. AWS says code written against torch.accelerator "runs unmodified", while compiled models need the backend named once with torch.compiler.set_default_backend("neuron") or per call.

The Hugging Face and AWS posts show Transformers loading straight onto the device: Hugging Face's with device_map="tpu", AWS's with device_map="neuron". In its own post, Hugging Face says its earlier Optimum TPU layer "is not required anymore to run Hugging Face models with Transformers and Diffusers".

How mature each one is

TorchTPU. Google and Meta describe it as "an open source collaboration between Google and Meta" that is "in active development". They say it is "used in production for both inference and training workloads for some of the largest Google Cloud customers", without naming them. By the team's count it supports both eager and compiled modes and about 450 ATen operators. Its model-coverage table reports 4,821 of 4,917 Hugging Face models running without a crash (98%) and 638 of 657 representative models matching GPU numerics (97%), tested on TPU generations from v5e onwards. Caveats it lists itself:

• torch.compile currently sets dynamic=False; support for dynamic shapes ("bounded dynamism") is still under development.
• FlexAttention is "not yet supported".
• By default each collective op is compiled on its own, which the team says can cost "30-40% end-to-end" performance against the optional faster paths.
• Because only one process can access a TPU device, reinforcement-learning colocation "can’t be supported out of the box".
• Helion kernel authoring is marked "(Experimental)".

Hugging Face reports its own testing. Running the Transformers library test suite, it writes: "It's not all green yet, and we're fine with that. Of the tests that run on TPU, 96.5% pass." The post does not say how many tests were run or which do not run on TPU. Hugging Face says it has filed 34 issues on the TorchTPU repo, 19 of them already fixed.

TorchNeuron. AWS says it was "going live as public Beta on Oct. 20", supporting Trainium2 and Trainium3, with eager mode, torch.compile and its NKI custom-kernel interface. Its "adaptive eager mode" is labelled experimental, and TorchAO and FP8 are "currently supported under private Beta". AWS charts its Hugging Face sweep, dated September 2026, at 96,618 of 101,076 models passing in eager mode (95.6%) and 95,097 of 98,420 in compiled mode (96.6%). On PyTorch's upstream test suites, in an August 2026 run, it charts pass rates of 94.7% (ATen), 96.6% (Dynamo), 94.7% (distributed) and 100% (profiler), measured against enabled tests. AWS says about 39,000 tests were deliberate Neuron skips, roughly 22,000 of them for operators "without a native Neuron lowering at this point". It also lists "Features TorchNeuron explicitly chose not to support in the first release", including sparse tensors, NestedTensor, forward-mode autograd, and dynamic and symbolic shapes.

Customer claims are attributed to the companies concerned. Deepgram's CTO, quoted by AWS, says that in benchmarks it evaluated, Trainium processes "about 60% more input audio seconds per dollar than our existing GPU baselines".

Licences

• TorchNeuron (aws-neuron/torch-neuronx): the GitHub API reports Apache-2.0, and the repository's LICENSE file is the Apache License 2.0, matching AWS's post.
• TorchTPU (google-pytorch/torch_tpu): the GitHub API reports "NOASSERTION". The LICENSE file itself applies the Apache License 2.0 to ./*, with a separate MIT licence (Andrej Karpathy, 2020) for examples/mingpt/impl/*. On our reading, that mixed file is why GitHub's detector returns no single licence; the core code is Apache 2.0.

Why it matters

On our reading, the practical change is that teams who wanted TPUs or Trainium previously went through separate layers (PyTorch/XLA, Optimum TPU, Neuron-specific tooling, by the sources' own accounts) and can now try the standard PyTorch device path instead. The vendors' coverage and performance figures are their own and have not been independently checked here. Both projects list unsupported features and beta components, so on our reading developers should treat them as early releases and test their own models.