The science behind TorchNeuron, Amazon’s new PyTorch backend for AWS Trainium chips
Captured source
source ↗The science behind TorchNeuron, Amazon’s new PyTorch backend for AWS Trainium chips - Amazon Science
Close
Close
Social
bluesky
threads
youtube
github
rss
Menu
Research
Research areas
Automated reasoning
Cloud and systems
Computer vision
Conversational AI
Economics
Information and knowledge management
Machine learning
Operations research and optimization
Quantum technologies
Robotics
Search and information retrieval
Security, privacy, and abuse prevention
Sustainability
Our scientific contributions
Publications
Research from our scientists and collaborators.
Conferences
Our experts present and discuss cutting-edge research at scientific meetings globally.
Research areas
Automated reasoning
Cloud and systems
Computer vision
Conversational AI
Economics
Information and knowledge management
Machine learning
Operations research and optimization
Quantum technologies
Robotics
Search and information retrieval
Security, privacy, and abuse prevention
Sustainability
Our scientific contributions
Publications
Research from our scientists and collaborators.
Conferences
Our experts present and discuss cutting-edge research at scientific meetings globally.
News & blog
The latest from Amazon researchers
Amazon Science Blog
Technical deep-dives and perspectives from our scientists.
News
Research milestones and recent achievements.
The latest from Amazon researchers
Amazon Science Blog
Technical deep-dives and perspectives from our scientists.
News
Research milestones and recent achievements.
Collaborations
Amazon Research Awards
Overview
Call for proposals
Latest news
Research stories
Recipients
Amazon Nova AI Challenge
Overview
Rules
FAQs
Teams
Research collaborations
Overview
Carnegie Mellon University
Columbia University
Hampton University
Howard University
IIT Bombay
Johns Hopkins University
Max Planck Society
MIT
Tennessee State University
University of California, Los Angeles
University of Illinois Urbana-Champaign
University of Southern California
University of Texas at Austin
Virginia Tech
University of Washington
Amazon Research Awards
Overview
Call for proposals
Latest news
Research stories
Recipients
Amazon Nova AI Challenge
Overview
Rules
FAQs
Teams
Research collaborations
Overview
Carnegie Mellon University
Columbia University
Hampton University
Howard University
IIT Bombay
Johns Hopkins University
Max Planck Society
MIT
Tennessee State University
University of California, Los Angeles
University of Illinois Urbana-Champaign
University of Southern California
University of Texas at Austin
Virginia Tech
University of Washington
Resources
Code and datasets
Amazon Nova
Try Amazon’s frontier foundation models.
Code and datasets
Amazon Nova
Try Amazon’s frontier foundation models.
Careers
Careers
Explore our open roles.
Amazon Scholars
Faculty research opportunities on industry-scale technical challenges.
Postdoctoral Science Program
Early-career research opportunities alongside experienced industry scientists.
Careers
Explore our open roles.
Amazon Scholars
Faculty research opportunities on industry-scale technical challenges.
Postdoctoral Science Program
Early-career research opportunities alongside experienced industry scientists.
Search
Submit Search
Cloud and systems
The science behind TorchNeuron, Amazon’s new PyTorch backend for AWS Trainium chips
Runtime optimizations and a thorough evaluation process driven by agentic AI help ensure that models written in PyTorch will execute efficiently and reliably on Trainium systems.
By Robert Hundt
October 22, 2026
7 min read
Share
Share
Copy link
X
Line
QZone
Sina Weibo
分享到微信
x
PyTorch is the one framework that most AI developers use to build their models. It’s a Python library of operations performed on tensors , higher-dimensional analogues of a matrix, or 2-D grid of numbers. Almost all of the crucial features of contemporary AI models, such as input data, model weights, and attention mechanisms, are represented by tensors, and PyTorch drastically simplifies writing code to manipulate them. This week, we announced the availability of TorchNeuron, an open-source backend that translates PyTorch code into operations executable on AWS Trainium systems. PyTorch code that already runs on other accelerators will run on Trainium with minimal or no code changes. TorchNeuron bridges PyTorch and Trainium, and both sides gain from it. PyTorch developers can now run their code on Trainium, which was built from the ground up for large-tensor workloads. They can also tap the AWS Neuron SDK, which includes a profiler that traces each instruction with cycle-level accuracy and the Neuron Kernel Interface (NKI, pronounced "Nikki"), which enables control almost at the assembly level.
TorchNeuron enables seamless integration between PyTorch models, the AWS Neuron SDK, and AWS Trainium chips.
TorchNeuron includes several optimizations that enable PyTorch code to run efficiently on Trainium, including caching at multiple levels and an optional “adaptive eager” mode that doesn’t execute individual operations one by one, as is typical in plain eager mode, but instead groups and compiles operations for more efficient execution, approaching fully compiled performance. Much of the innovation behind TorchNeuron, however, is the way its performance was evaluated. AI agents searched through hundreds of thousands of open-source models and categorized them according to their architectural characteristics. This allowed the Annapurna team around Khaled ElGalaind and Sarang Ozarde to identify 841 representative models covering the full spectrum of model architectures. From those models, they also distilled tens of thousands of practically important tensor dimensions, or “shapes”, which were then used to tune and benchmark TorchNeuron’s performance. The result is that PyTorch programmers can trust TorchNeuron to execute their code accurately, efficiently, and most importantly, right out of the box.
Runtime optimizations
Adaptive eager mode
Standard PyTorch has two native execution modes. One is called “eager”, meaning that PyTorch operations are executed immediately, in program order. The other mode is compiled mode, where the programmer selects blocks of code that should be compiled in the traditional sense, meaning that a...
Excerpt shown — open the source for the full document.
Notability
notability 7.0/10New PyTorch backend for AWS Trainium, notable infrastructure.
Amazon (Nova) has a writing signal matching evals and quality, infrastructure.