Frontier models span language, video and audio, on architectures from dense and mixture-of-experts transformers to diffusion, state-space, world and structure-prediction models. PyTorch is where AI research happens, where the models are defined, and increasingly where production happens too. Researchers want to experiment and innovate, and customers want the performance AWS Trainium offers to enable them to achieve their performance and scale goals. The training and serving libraries are written against PyTorch’s own APIs, and new libraries and new versions ship continuously. Over the last several years, PyTorch made those APIs accelerator-agnostic, with official extension points for registering custom hardware.

Today, to address the growing demand from the developers and research community on Trainium, we are excited to share an even simpler approach to program the chip as a standard PyTorch device, ahead of its open source launch at PyTorch Conference North America on Oct. 20. AWS Trainium is a portfolio of AI chips designed from the ground up for next generation cloud scale AI training and inference. We are announcing TorchNeuron, an open source native PyTorch backend for Trainium, so the model code, the libraries, and the tools you already use run on Trainium. This integration builds directly on the accelerator-agnostic extension points, including the PrivateUse1 device backend and torch.accelerator APIs, created by the PyTorch community.