Amazon targets Nvidia’s software edge with new AI tool

U.S.-based technology company Amazon has launched an open-source software tool that allows developers to run existing AI code on its own AI chips with minimal changes, making it easier to switch hardware.

The initiative targets a software advantage that has helped Nvidia maintain its position in AI computing.

In an announcement published on its official news website on October 20, Amazon said its cloud-computing arm, Amazon Web Services (AWS), had developed TorchNeuron, a tool that allows developers to run AI workloads built using PyTorch on its Trainium2 and Trainium3 chips with minimal code changes.

PyTorch is an open-source machine-learning framework widely used to develop and train AI models. Amazon estimates that 85% of AI researchers and developers use it. TorchNeuron is designed to let developers use familiar PyTorch code and tools on Trainium, potentially reducing the engineering work involved in moving workloads to alternative hardware.

“85% of AI developers write their code in PyTorch. But until now, trying a new chip meant rewriting and revalidating work that took months to perfect. That’s an enormous tax on teams who should be building products, not porting code,” AWS chief executive Matt Garman said.

“Today we’re changing that. TorchNeuron runs your existing PyTorch code natively on Trainium with a single config change. No rewrite. No revalidation. We’re open sourcing it under Apache 2.0, and we’re contributing the changes upstream to PyTorch itself, so the whole framework gets more hardware-agnostic for everyone.

“Developers should pick hardware on price, performance, and availability. Not because they’re locked in by code they already wrote,” he said.

The move addresses a key challenge for companies developing alternatives to Nvidia’s graphics processing units (GPUs).

NVIDIA’s CUDA platform provides the software tools and libraries that allow developers to run AI workloads on its GPUs. Its integration with frameworks such as PyTorch has helped build a broad developer ecosystem. Amazon says moving workloads to alternative chips can require substantial rewriting and revalidation, a barrier TorchNeuron is designed to reduce.

TorchNeuron is being released under the Apache 2.0 open-source licence as a public beta. Amazon is also contributing changes to PyTorch to reduce hardware-specific code in the framework, with the aim of making it easier to run AI workloads across different chips.

The unveiling has drawn support from Meta’s PyTorch engineering team. Jana Van Greunen, director of PyTorch engineering at Meta, said native support for Trainium would give researchers and engineers greater choice while allowing them to continue using familiar execution and distributed-computing tools.

Hugging Face, whose libraries are widely used to access and run AI models, has integrated support for Amazon’s Trainium chips into its Transformers and Diffusers libraries. The company said the new backend would let developers use familiar model definitions on Trainium, reducing the need to maintain separate, hardware-specific implementations.

Amazon said more than 30 early customers across areas including robotics, voice AI and world models were using the software.