TorchNeuron Is Public Beta. Your CUDA Pin Did Not Move

The PyTorch blog introduced TorchNeuron as a native backend for AWS Trainium, with a public beta on October 20 and Apache 2.0 code at github.com/aws-neuron/torch-neuronx. Eager mode, torch.compile, and the Neuron Kernel Interface sit on Trainium2 and Trainium3. That is a backend. It is not your CUDA pin.

GeekWire says the stack was in private beta since late last year, that Amazon said Monday it is opening to all developers, and that more than 30 companies plus internal Amazon teams already used it. Colin Brace at Annapurna Labs: previously you called Amazon-specific APIs; now the same PyTorch runs after a basic configuration change. “If you know PyTorch, then you know Trainium.” About Amazon, by Isaac Schultz, puts the open-source drop at PyTorch Conference North America on October 20 and names Hugging Face plus TorchTitan/FSDP as the compatibility story.

We already filed PyTorch 2.14’s clamp gradient and vLLM against the speed of light. This week the conference news is a device string.

Public beta is a date, not an SLA

PyTorch.org: public Beta October 20. About Amazon: Apache 2.0 at the North America conference the same day. GeekWire: open source Tuesday in San Jose, public beta announced Monday. Those timestamps agree enough to use. They do not agree with older Neuron docs that still called the native backend a closed beta you join through a support rep. Believe the launch posts on the day they launch. Recheck the docs before you promise a customer.

Private beta, GeekWire: more than 30 companies and internal teams. Deepgram, speech AI, San Francisco: about 60% more audio per dollar than its existing GPU setup, now production on Trainium, per Brace. Prime Video: Prime 360 rebuilds NBA plays in 3D so you can watch a replay from an angle a camera did not occupy; Trainium-backed; select broadcasts this NBA season. Those are named workloads. They are not your BERT finetune.

Brace’s quote is a slogan with a footnote. You know PyTorch. You do not yet know Trainium quotas, instance types, or whether your fused kernel exists in NKI-lib. The footnote is the rest of the PyTorch blog: FlashAttention and MoE blockwise kernels in NKI-lib, custom ops registered in the register_autograd style, Neuron Explorer tying torch.profiler events back to source lines.

Hugging Face, in both the PyTorch post and About Amazon: Transformers, Diffusers, trl, kernels. Transformers v5 for text and audio training and inference. Diffusers for image and video. trl for RL on large models. NKI kernels through the Hugging Face kernels library. AWS and Hugging Face added the support so existing HF workflows run without device-specific code, in that telling. About Amazon adds TorchTitan and FSDP as the scaling path you already know from other devices.

About Amazon also says Amazon made a major open-source contribution to vLLM and open-sourced its complete set of performance kernels. If you serve on vLLM, that sentence is the one to test. It is not a promise that your current vLLM pin takes --device neuron today. Check the version. We have been burned by image tags before.

A config change is not a free lunch on silicon

The migration story is: stop calling Neuron-specific APIs, set the backend, run. Brace says that out loud. The hardware story is still Trainium2 and Trainium3. If your bill is GPU on a different cloud, you are changing more than a flag. You are changing the building.

GeekWire frames this as Amazon challenging Nvidia’s software advantage. Fine as a headline. Brace, asked about Nvidia, said Amazon has “always been about choice,” and pointed at Intel next to Graviton. Choice is a catalog. Your runtime is one row. CUDA wheels you already vendor-pinned will not compile themselves into NKI.

Custom kernels are the escape hatch and the trap. NKI exists so you can write what the library lacks. The blog shows registering a custom operator the way PyTorch wants. That is research and performance engineering, not pip install and go. If your model is dense Llama-shaped and FlashAttention is in NKI-lib, you might be in the happy path. If your model is a one-off attention variant, you are writing NKI whether the slogan says you know Trainium or not.

Eager mode, including an “adaptive eager” with better performance in the PyTorch post, is for debug. torch.compile is the production claim. Two modes means two failure modes. A graph that compiles on CUDA can still bounce on Neuron. Public beta is when you find that out. Do not freeze a Trainium compile error into an incident on day one unless you already had a private-beta relationship.

Neuron Explorer using torch.profiler is the right kind of boring. Same tool, events mapped to lines. If you already profile PyTorch, you are not learning a new GUI from scratch. You are learning what the traces mean on this chip.

Transformers loading GGUF while llama.cpp owns the no-Python box was a similar split: the high-level API is not the runtime. TorchNeuron is the runtime for people who already wrote nn.Module. GGUF is not this post.

Named customers are not your capacity plan

Deepgram’s 60% more audio per dollar is a customer benchmark Brace repeated. Audio is their job. If your job is images, do not steal 60%. If your job is audio, still ask which GPU they were on, which batch size, and whether the dollar is instance list price or a committed rate. GeekWire did not print those. Neither should you invent them.

Prime 360 is a product feature on a video site. 3D replay on Trainium is a existence proof that internal Amazon teams used the beta. It is not an API you import.

GeekWire’s other named Trainium users are the expensive ones. Anthropic, Amazon-backed, trains and runs Claude on the chips. OpenAI agreed in February to use 2 gigawatts of Trainium capacity as part of a deal in which Amazon is investing $50 billion, per that article. Those sentences are corporate. They do not put a trn2 instance in your VPC. They do explain why Amazon wants a PyTorch-shaped on-ramp: the labs already in the family should not be the only ones who can spell Neuron.

Startups and enterprises in robotics, video, world models, self-driving, speech, drug discovery, LLMs, and finance are “already building on it,” About Amazon says. That is a list of industries, not a list of GitHub stars. Treat it as marketing geography.

If you are on CUDA because the rest of the company is on CUDA, TorchNeuron does not end that meeting. It starts a meeting about whether Trainium capacity is even enabled. Public beta on a conference Tuesday is a good time to read the blog. It is a bad time to swap device="cuda" in prod because a VP forwarded GeekWire.

Trainium2 and Trainium3 are still the floor

The PyTorch post and About Amazon both name Trainium2 and Trainium3. There is no Trainium1 nostalgia tour in these write-ups. If your account only shows older Neuron instances, public beta on a blog does not conjure new SKUs. Capacity is an AWS conversation. The GitHub repo is a software conversation. People will mix them in Slack on Tuesday. Separate the threads.

TorchTitan plus FSDP is the scaling sentence meant for people who already sharded on CUDA. Same patterns, different device, in About Amazon’s telling. Fully Sharded Data Parallelism is not a new algorithm here. It is a claim that you will not rewrite your trainer to a vendor API. Test that claim with a small transformer, not with the 70B you cannot afford to debug on a new backend during a conference keynote.

Hugging Face’s “no Amazon-specific code” line is doing a lot of work. Transformers v5, Diffusers, trl, kernels: four products. Text and audio in Transformers. Image and video in Diffusers. RL in trl. Kernels as the NKI delivery path. If you are on an old Transformers 4.x pin, you are not in that sentence. Upgrade is a separate PR. Do not hide it inside “we enabled Trainium.”

NKI-lib’s FlashAttention and MoE blockwise kernels are the two workloads Amazon wanted on the slide. Dense and MoE, training and multimodal, in the blog’s language. Custom ops through register_autograd is how you fill gaps. If you cannot read a NKI kernel, you are a customer of NKI-lib, not an author. That is fine. It also means you wait when your attention variant is missing.

Neuron Explorer mapping profiler events to source lines is the adult tool in the announcement. Same torch.profiler you already run. The value is the link from a hardware event back to a Python line. Use it when compile mode is slow, not as a screenshot for the launch channel.

GeekWire’s Nvidia-challenge frame will dominate the secondary coverage. Brace’s “choice” answer is the corporate counter. Your lockfile does not care. device="cuda" stays until you change it, run it, and keep the metrics. A conference in San Jose is a good place to hear the talk. It is a bad place to merge unsigned kernel pins.

What to pin this week, and what to leave

If you already have Trainium access and a PyTorch training script that is not a pile of custom CUDA, try the backend on a throwaway job. Pin the TorchNeuron commit you actually ran. Pin Transformers if you are in the HF path. Write down eager vs compile.

If you serve with vLLM, read Amazon’s vLLM contribution notes, not just the sentence that they contributed. The kernel dump is useful when it lands in a version you can install. Until that version is the one in your lockfile, your serving path is still whatever compiled last week.

San Jose is a conference city. The repo path is github.com/aws-neuron/torch-neuronx. Bookmark the repo, not the airport.

If you write kernels, NKI-lib plus register_autograd is the homework. FlashAttention and MoE blockwise are the examples they wanted you to see.

Skip rewriting a working CUDA stack this week because the conference is in San Jose. Skip quoting 60% on a slide that is not about speech. Skip treating Brace’s slogan as a guarantee that every torch.distributed call maps. Skip the closed-beta docs if the launch post says public; skip the launch post if your account still cannot see the instances.

CUDA did not move. The pin in your lockfile is still the pin. TorchNeuron is another device, now with a public door and an Apache 2.0 repo. Walk through the door with a toy model. Leave production where it compiles today.