· ai-devtools
AWS Launches TorchNeuron to Run PyTorch Models on AWS Trainium Without Modifications
Beyond reducing engineering effort, TorchNeuron allows enterprises to share hardware capacity between training and inference workloads, reallocating resources from inference to training during off-peak hours.
Amazon says that a single configuration change is enough to switch hardware, eliminating the need to rewrite and revalidate models when moving to Trainium. Today, 85% of AI researchers and developers write their code in PyTorch. Until now, the hardware that code could run on was far more limited than the code itself, and expanding to new chips meant rewriting and revalidating work that took months to perfect.
According to the TorchNeuron documentations, when leveraging PyTorch's automatic device detection through torch.accelerator, the only required change is replacing .to('cuda') to .to('neuron'). The backend also natively supports distributed training and PyTorch parallelism APIs, including DDP, FSDP2, DTensor, and others. In addition, TorchNeuron provides native support for Torch streams and events through torch_neuronx.Stream and torch_neuronx.Event. Amazon Annapurna Labs vice president Colin Brace, who leads the team behind TorchNeuron, emphasized that his team validated the 100,000 most popular models and more than 800 ML architectures to make sure a smooth developer experience when adopting the new backend.
