Amazon is releasing TorchNeuron as open-source software today. This allows existing PyTorch code to run natively on AWS Trainium chips. According to Amazon, only a single configuration change is needed, with no rewriting required. This makes it easier for developers to switch to an alternative to GPU hardware.
TorchNeuron is released under the Apache 2.0 license and supports Trainium2 and Trainium3. According to Amazon, 85 percent of AI researchers and developers write code in PyTorch. Until now, anyone who wanted to run that code on other chips often spent months rewriting and revalidating it.
Technically, TorchNeuron leverages the mechanism that already allows new hardware to register with PyTorch. As a result, developers can use the same tools for testing, debugging, and scaling on Trainium without modification. There is also integration with TorchTitan, PyTorch’s native distributed training framework, including Fully Sharded Data Parallelism (FSDP).
Amazon has made a major contribution to vLLM and has open-sourced its complete set of performance kernels. Hugging Face now natively supports Trainium in its development tools, without any Amazon-specific code. In addition, Amazon contributes changes upstream to PyTorch to remove code that ties the framework to specific chips. Every build is tested against hundreds of thousands of AI models.
From speech AI to drug research
Among the first users is Cartesia, which is working on conversational AI. The company ran its standard inference pipeline on Trainium with minimal modifications to the PyTorch code. The Lieber Institute for Brain Development, which conducts research on schizophrenia and other conditions, uses TorchNeuron to explore AI models that can accelerate or replace computationally intensive molecular dynamics simulations.
TorchNeuron follows a decade-long investment in proprietary chips. Since Amazon acquired chip designer Annapurna Labs in 2015, the company has developed four chip families, each designed for a specific cloud workload.