One-line change lets PyTorch run on Amazon Trainium chips in new open-source beta

Amazon has released TorchNeuron, an open-source software tool that lets developers run existing PyTorch code on Amazon’s Trainium AI chips with a single configuration change, according to The Next Web. The framework launches into public beta on October 20, offering a pathway for teams to shift workloads away from Nvidia hardware without extensive re-engineering.

The risk

The primary challenge for developers is the time and complexity involved in migrating AI models to new hardware. Historically, moving code to a different chip required months of rewriting and testing to ensure compatibility. TorchNeuron addresses this by plugging into PyTorch’s built-in system for adding new hardware, allowing standard testing and scaling tools to function as usual. However, the initial release has limitations. It does not yet support certain PyTorch features, such as dynamic input shapes. Approximately 22,000 PyTorch tests are skipped in this first version, indicating that not all functionalities are immediately available.

Who is exposed

This update affects software engineers and data scientists using PyTorch for AI model training and inference. Amazon reports that early users of TorchNeuron include voice AI startups Cartesia and Deepgram, as well as the Lieber Institute for Brain Development, which utilizes the tool for drug research. Hugging Face has also integrated Trainium support into its AI tools. Additionally, Uber recently signed an agreement to use these chips as part of an AWS expansion deal, signaling broader industry interest in Amazon’s custom silicon solutions.

Cost and compliance

Amazon’s custom chip business has reached a $25 billion annual revenue run rate, according to company statements. Chief executive Andy Jassy previously suggested the division could be worth $50 billion. The launch of TorchNeuron is part of Amazon’s strategy to compete with Nvidia, whose chips and CUDA software currently dominate the AI computing market. TorchNeuron is released under the Apache 2.0 license, allowing for widespread adoption and modification. Decart, an AI chip software startup reportedly in talks to be acquired by Amazon, is listed as a launch partner.

What should you do?

Developers should evaluate whether their current PyTorch workloads rely on unsupported features like dynamic input shapes before migrating. Teams can test the public beta version to compare performance and compatibility with their existing models. Monitor Amazon’s updates for feature completion, as the initial release skips a significant number of standard tests.

The background in plain language

When moving software to new hardware, teams usually face compatibility hurdles that require code adjustments. This tool aims to hide those differences, letting the software talk to the chip without manual translation. Check whether your project uses advanced features that might not work yet. Teams should verify that their testing pipelines run smoothly on the new backend. Ensure that performance metrics remain consistent after the switch. This approach reduces the risk of errors during migration.

What to do and how to stay safe: Amazon

• Audit your PyTorch code to identify usage of features not yet supported by TorchNeuron, such as dynamic input shapes.
• Run a pilot test of your model on the Trainium beta environment to verify that training rates and results match your current hardware.
• Monitor official documentation for updates on the 22,000 skipped tests to determine when full feature parity is achieved.
• Consult with your infrastructure team to assess the long-term cost and performance benefits of switching to Amazon Trainium chips.

Step-by-step guide: Browser security settings are the architectural walls of your digital life

General security guidance from the Vector Tribune newsroom. It is not confirmed advice from the organisations named in this story.

Frequently asked questions

When does TorchNeuron become available for public testing?

TorchNeuron enters public beta on October 20, according to a PyTorch blog post by Amazon engineers.

Which Trainium chips are supported by this release?

The tool supports Amazon’s Trainium2 and Trainium3 AI chips, as stated by The Next Web.

Is TorchNeuron free to use?

Yes, it is open-source and released under the Apache 2.0 license, allowing developers to use and modify the code freely.