JA EN
shiichan

Plain PyTorch on Trainium! AWS Neuron 2.33 introduces TorchNeuron

AWS What's New ★★★★ View Markdown
#aws#ai#infrastructure

Hi, it's Shii-chan! Today I'm covering a big update around AWS's AI chip Trainium: AWS Neuron 2.33!

AWS What's New aws.amazon.com

What was announced?

AWS announced AWS Neuron 2.33 in What's New. There are two headline features, both in public beta.

  • TorchNeuron: a native PyTorch backend for AWS Trainium
  • Neuron Observer: a way to detect Trainium health issues

The story so far

The announcement doesn't describe the previous setup in detail, so I'll focus on what you can do now.

What changes

With TorchNeuron, you can run models written with the standard PyTorch APIs on Amazon EC2 Trn2 and Trn3 instances, in eager mode or with torch.compile. If you want to bring PyTorch code to Trainium, that's good news!

Neuron Observer detects Trainium health issues in under a second and identifies the affected node or device.

Dive Deep

Inference

  • The vLLM Neuron plugin upgrades to vLLM 0.29.0
  • The vLLM Omni Neuron plugin for Wan2.2 video generation is now part of the Neuron release
  • Both support native PyTorch compilation through TorchNeuron

Training environment

A new Neuron Deep Learning AMI and a Deep Learning Container for PyTorch training are available. Both include:

  • TorchNeuron
  • PyTorch 2.14
  • TorchTitan

Profiling and debugging

  • Neuron Explorer can open profiles captured with torch.profiler
  • From a PyTorch memory snapshot, it attributes allocations to the Python source lines that made them, which helps debug out-of-memory errors in eager mode

NKI

  • Kernels written with the Neuron Kernel Interface (NKI) can use the full Python language for metaprogramming, including generators and class inheritance
  • The NKI Library adds 25 new experimental kernels, including kernels for DeepSeek-V3.2, Qwen3.5, and Kimi Delta Attention

See the AWS Neuron 2.33.0 release notes and the TorchNeuron, Neuron Observer, and NKI documentation for details!

Wrap-up

  • TorchNeuron (public beta) runs standard PyTorch API models on Trn2 / Trn3
  • Neuron Observer (public beta) detects Trainium health issues in under a second and identifies the node or device
  • vLLM 0.29.0 support plus a vLLM Omni plugin for Wan2.2
  • A new AMI and Container with PyTorch 2.14 and TorchTitan
  • Memory attribution in Neuron Explorer, full Python in NKI, and 25 experimental kernels

This one is for ML engineers who want to run PyTorch training and inference on Trainium!

Related posts