TorchTPU and TorchNeuron make TPUs and Trainium first-class PyTorch devices

Two device strings landed in the same conference window. On 20 October 2026, Google and Meta's TorchTPU and AWS's TorchNeuron both promised the migration path PyTorch teams actually type: change the device, keep the model code.

Native TPU and Trainium backends register beside CUDA instead of wrapping JAX workflows

TorchTPU is an open-source PyTorch backend for Google TPUs, developed at Google with the PyTorch team at Meta. The public pitch is principled PyTorch citizenship: no JAX runtime dependency, solid eager mode, torch.compile via TorchDynamo, and TPU collectives exposed through PyTorch distributed APIs. Usability examples swap torch.device("cuda") for torch.device("tpu") and nccl for tpu_dist. Architecture notes describe a PrivateUse1 C++ backend, Strict and DeferAndFuse eager modalities, a multi-tier compilation cache, and roughly 450 ATen ops with forward and backward coverage for core deep-learning paths. The project reports out-of-box checks on thousands of popular Hugging Face models and validated paths for vLLM, SGLang, TorchTitan, Ray, and TorchRec. PyPI ships torch_tpu (0.2.0 cited in the README); source lives at google-pytorch/torch_tpu under an Apache-2.0-style project licence file.

TorchNeuron registers AWS Trainium through the same PrivateUse1 extension point so model.to("neuron") works on Trainium2 and Trainium3. Eager mode, an experimental Adaptive Eager fusion path, and torch.compile(..., backend="neuron") are first-class. Distributed APIs named in the launch post include DDP, FSDP2, tensor and pipeline parallelism, DTensor, and DeviceMesh, with collectives on Trainium's collective cores and EFA between nodes. Custom kernels use the Neuron Kernel Interface (NKI), with NKI-lib operators for FlashAttention and MoE blockwise paths. Neuron Explorer hooks the familiar torch.profiler. Inference plugins for vLLM Neuron and vLLM Omni Neuron sit on the same compile support. Code is Apache-2.0 at aws-neuron/torch-neuronx, announced as public beta at the conference.

PyTorch Conference turns accelerator portability into a device-string problem

PyTorch/XLA taught a generation of TPU users to mark compile regions and live with lazy-tensor friction. TorchTPU's authors say that stack lagged eager mode, compile latency, and native distributed primitives, and that TorchTPU was built so existing workloads run with trivial device changes even when peak performance still needs TPU-friendly shapes. AWS makes the parallel argument for Trainium: research and production libraries are written to PyTorch APIs, so Trainium must show up as a standard device rather than a side compiler workflow.

Shipping both backends in the same week reframes the non-CUDA story. The question becomes which native device string your cluster can feed, and how mature the eager and compile paths are on day two, rather than whether PyTorch can see the chip at all.

Open code on both sides; TorchTPU cites production Cloud customers, TorchNeuron opens as public beta

TorchTPU: open GitHub under google-pytorch/torch_tpu, pip wheels on PyPI, docs at torchtpu.readthedocs.io. The launch post says the stack is already used in production for some Google Cloud customers for training and inference while remaining in active development. Star count on the public mirror was still modest relative to inference frameworks when checked for this draft.

TorchNeuron: Apache-2.0 at aws-neuron/torch-neuronx, public beta as of 20 October 2026, with Neuron SDK release notes in the 2.33 line documenting the native backend direction. FP8 and TorchAO paths are called private beta in the launch post.

Maturity read: usable for teams with TPU VMs or Trainium instances who can own device-specific profiling. Production-shaped after you prove sharding, precision mode, and serving plugins on your model family. Demo if you only change the device string on a toy tensor.

Reach for TorchTPU or TorchNeuron when the chip is fixed and the code must stay PyTorch

TorchTPU fits teams whose capacity is already TPU slices on Google Cloud and who want eager debugging plus torch.compile without rebuilding the job as JAX. It is also the path when vLLM or SGLang on TPU is the serving target the Google/Meta post validates.

TorchNeuron fits when Trainium2 or Trainium3 is the purchased accelerator and you need .to("neuron") with NKI escape hatches for custom kernels. Adaptive Eager helps while iterating; compile is for stable production graphs.

CUDA PyTorch remains the default when your models, kernels, and serving stack are already tuned there and you lack TPU or Trainium quota. Keep PyTorch/XLA only for legacy TPU jobs you are not ready to re-validate; TorchTPU is the documented direction for new native work in the launch framing. JAX or native Neuron SDK stacks still make sense when you are already deep in those toolchains and PyTorch APIs are not the constraint.

Device strings run; peak performance still demands shapes, precision, and plugin glue

TorchTPU warns that .to("tpu") should run out of the box while peak performance still needs shapes and sharding suited to systolic MXUs. Default matmul precision maps to BF16-friendly medium mode; high and highest multi-pass modes cost roughly 3× or 6× MXU compute. One user-space process per TPU device complicates some RL colocation patterns the post flags as under investigation. DeferAndFuse and compile-cache behaviour will surprise teams who expect pure op-by-op CUDA semantics.

TorchNeuron's beta edges include Adaptive Eager still experimental, FP8/TorchAO under private beta, and serving depending on Neuron-specific vLLM plugins. Host-collective and stream multiplexing details differ from CUDA mental models. Independent cross-cloud benchmarks against H100 or GB200 class GPUs were not the point of either launch post.

Match the backend to the chips on the floor, then re-bench one training step and one serve path

If the cluster is TPU, install TorchTPU and move one training step plus one vLLM or SGLang smoke test before rewriting sharding. If the cluster is Trainium, install TorchNeuron, set the Neuron compile backend once, and prove eager debug then compiled train on a single node. Keep CUDA as the control. The Monday decision is which native device backend matches the chips on the floor, with PyTorch/XLA demoted for new TPU work unless a legacy path still blocks migration.