PyTorch conference extends training and inference to more chips: TorchTPU and TorchNeuron can make migration a device change

The PyTorch ecosystem gained two native hardware backends, so the same PyTorch code can run directly on Google TPUs and AWS Trainium instead of needing a separate implementation per chip. Previously, landing PyTorch workloads on these accelerators required extra adaptation work; TorchTPU and TorchNeuron make that layer a native backend, and TorchNeuron runs with a single configuration change. TorchTPU is an open source Google and Meta collaboration supporting eager execution and torch.compile, with validated integrations covering TorchTitan for training, vLLM and SGLang for inference, and Ray for orchestration; TorchNeuron launched as an Apache 2.0 public beta at the conference on October 20, supporting eager execution, torch.compile and native APIs such as DTensor and torch.distributed.

Moving to a new chip can require only a device change to run, but reaching peak performance can still mean changing shapes, sharding and operations. The distinction separates running from running fast: device substitution addresses runnability, while performance tuning remains a separate step. The TorchTPU description states that workload migration can require only a device change while peak performance can require additional changes to shapes, sharding and operations; TorchNeuron targets Trn2 and Trn3 instances and covers workflows from research experimentation to production deployment.

The backends target researchers and engineers using PyTorch for training and inference, covering Google TPUs and AWS Trainium Trn2 and Trn3 instances; TorchNeuron is released as an Apache 2.0 public beta.

The material states explicitly that migration can require only a device change, but peak performance can require additional changes to shapes, sharding and operations, so a device swap does not by itself mean matching performance.