Amazon announced this Monday (19) the TorchNeuron, a tool that allows running code developed with PyTorch on AWS's Trainium artificial intelligence accelerators without needing to rewrite the applications. According to the company, migration can be done with a simple configuration change, reducing one of the main difficulties faced by developers who want to use alternative chips to Nvidia's.
The software will be made available for free and as open source on October 20, during the PyTorch Conference North America. Amazon said the tool will have an Apache 2.0 license and offer support for Trainium2 and Trainium3 accelerators.
PyTorch is a machine learning framework maintained by Meta and widely used in developing AI models. Although it is not exclusive to Nvidia, its integration with the CUDA ecosystem represents an important advantage for the manufacturer, which has invested for more than a decade in tools to run artificial intelligence applications on its GPUs.
With TorchNeuron, Amazon aims to reduce the need for software adaptations to use its own infrastructure. The tool also offers integration with compilation, distributed processing and custom operations development features for Trainium chips.
In an interview with Reuters, Colin Brace, vice president of software at Annapurna Labs, Amazon's chip division, summarized the strategy: “The idea has been to meet customers where they already are.”
According to the executive, more than 30 early customers already use TorchNeuron, including companies that develop robotics applications, world models and artificial intelligence for voice. Brace said that migration to Amazon's chips can provide significant savings, but did not present independent results that would allow comparing costs between the platforms.
TorchNeuron had already appeared in AWS technical documentation in restricted-access versions. The October announcement therefore marks a new stage of public distribution, not the emergence of the technology.
The promise of running code without rewriting also does not mean unrestricted compatibility with all PyTorch applications. Effective support and performance remain subject to the characteristics of the models and the features implemented by AWS.
