PyTorch conference extends training and inference to more chips: TorchTPU and TorchNeuron can make migration a device change
At the PyTorch Conference North America in San Jose, the foundation's projects announced several efforts to run PyTorch natively on more hardware. TorchTPU, a Google and Meta collaboration, gives Google TPUs a native backend supporting eager execution and torch.compile; TorchNeuron for AWS Trainium shipped as an Apache 2.0 public beta that runs PyTorch workloads with a single configuration change. Workload migration can require only a device change, while peak performance can require additional changes to shapes, sharding and operations.