The Pulse
Hugging Face Adds Native PyTorch Training on AWS Trainium
Hugging Face says its libraries can now use AWS Trainium through TorchNeuron, a native PyTorch backend entering public beta on October 20, 2026. The change brings familiar training and distributed APIs to the hardware, while the beta still
AI.info Team · Written with AI · fact-checked ·
Hugging Face says developers can train models on AWS Trainium through familiar PyTorch libraries, with the new path arriving first as a public beta and still requiring workarounds for static input shapes. Hugging Face published its integration announcement on October 19, 2026; TorchNeuron is scheduled to become publicly available the following day.
Trainium becomes a native PyTorch device
TorchNeuron is a PyTorch plugin that registers Trainium as a native device, enabling eager execution, torch.compile and distributed tools. Hugging Face lists Transformers, Diffusers, Accelerate, PEFT and TRL as libraries that can use Trainium directly, with the same classes and training scripts developers use on other accelerators. The post also describes support for its Kernels library.
The new route replaces a more specialized workflow built around Optimum Neuron and PyTorch XLA. Hugging Face says that earlier setup involved lazy tensors, static shapes and recompilation when a graph broke, with Optimum Neuron maintaining custom model classes and trainers to hide some of that complexity. The company now recommends Neuron users work directly with standard Hugging Face libraries instead.
TorchNeuron builds on PyTorch’s extension points for custom hardware. The PyTorch announcement says the backend supports eager execution and torch.compile on Trainium2 and Trainium3, as well as standard distributed APIs. Developers can move a model to the neuron device, while compiled models can select the Neuron backend.
“Native support for AWS Trainium expands that choice, enabling researchers and engineers to experiment and iterate using the eager execution and distributed APIs they already know.”
Jana van Greunen, Director of PyTorch Engineering at Meta
The demo runs on eight Trainium2 chips
Hugging Face’s training example fine-tunes Qwen3-8B on the UltraChat dataset using TRL’s standard SFTTrainer. It uses tensor parallelism and Fully Sharded Data Parallelism (FSDP), with torch.compile enabled. The company ran the example on eight Trainium2 chips, or 32 NeuronCores—half of a trn2.48xlarge instance.
The training script loads the model and dataset, configures distributed training, and saves the resulting checkpoint through the trainer. Hugging Face says that checkpoint is a standard Transformers model: it can be shared on the Hugging Face Hub and loaded on other hardware, rather than being tied to Trainium.
The post also reports a separate inference comparison: 89% of the 100 most-used text-generation model architectures work out of the box with TorchNeuron, compared with 46% using Optimum Neuron v0.4.5. Those numbers describe architecture support for text-generation inference; they are not measurements of training speed or model quality. For inference, the example uses a static cache, which Hugging Face recommends to avoid frequent recompilations.
Beta users still face shape limits
The move to native PyTorch does not remove every hardware-specific constraint. Hugging Face’s training example disables automatic dynamic shapes because the Neuron backend cannot yet compile them. It uses wrapped packing to produce sequences up to 2,048 tokens; with TRL version 1.16 or later, batches are padded to multiples of 128 tokens to limit the number of input shapes. The post says compiled graphs can then be reused for later steps with the same shapes.
The example also disables Transformers’ asynchronous weight loading, which Hugging Face says Neuron does not yet support. TorchNeuron’s public beta is scheduled to begin October 20, 2026, and can be installed with pip install torch-neuronx. PyTorch says the release will be open source under the Apache 2.0 license and support Trainium2 and Trainium3. The familiar library interfaces simplify the programming path, while shape handling and weight-loading limits remain practical considerations for users adopting the beta.