TorchNeuron is a PyTorch plugin that adds a native Neuron backend for AWS Trainium. It supports eager mode, torch.compile, and native distributed APIs such as FSDP and tensor parallelism, and it will be available as a public beta starting Oct. 20, 2026 via pip install torch-neuronx. With this backend, Hugging Face libraries no longer need optimum-neuron: Transformers, Diffusers, Accelerate, PEFT, and TRL can use the same classes and scripts as on other accelerators, treating Neuron as a regular device.

The article contrasts this with the previous path: optimum-neuron wrapped torch-neuronx, which was based on PyTorch XLA, meaning lazy tensors, static shapes, recompilation on graph breaks, and custom modeling code. TorchNeuron removes that layer by integrating into PyTorch, exposing neuron as a native device with eager execution, PyTorch distributed primitives (DTensor, FSDP2), and torch.compile. The recommendation is to use the Hugging Face libraries directly instead of optimum-neuron.

For text generation, Transformers now natively supports a wide range of models. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, compared with 46% for Optimum Neuron v0.4.5. Users pass device_map=neuron and can use tensor parallelism via DeviceMesh and DistributedConfig. The article’s example loads Qwen/Qwen3.8-27B with TP across all ranks, uses a static cache and prefill chunking for generation, and compiles with the Neuron backend.

Training is also unmodified from standard Hugging Face workflows. A TRL SFTTrainer example runs on Trainium with a few Neuron-specific settings: disabling torch._dynamo automatic dynamic shapes, disabling Transformers’ async weight loading with HF_DEACTIVATE_ASYNC_LOAD=1, and enabling –torch_compile true. Parallelism is configured purely through DistributedConfig with tp_size and fsdp_size; the trainer sees the model is already distributed and does not wrap it again. The example ran on 8 Trainium2 chips (32 NeuronCores), half a trn2.48xlarge instance, with a 4×8 TP/FSDP grid. The saved checkpoint is a standard Transformers checkpoint that can be pushed to the Hub and loaded on any hardware.

Diffusers inference also works: many pipelines run in eager and compiled modes, and FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1, and H3 support tensor parallelism. The sample uses Qwen-Image-2.1 with a TP=8 DeviceMesh, keeps the text encoder on CPU, and compiles the transformer with the Neuron backend. Launching uses NEURON_LOGICAL_NC_CONFIG=2 torchrun –nproc_per_node=8. More pipelines will gain TP support, and unsupported ones can be extended by following the included guide.

The post also points to Neuron kernels written with the Neuron Kernel Interface (NKI), which gives lower-level control for custom operations; one example kernel speeds up vision transformers by up to x2.4 on Trainium. The AWS team is thanked for the TorchNeuron work and for responding to feedback on the Hugging Face integration.