Hugging Face's language model training tool now supports Amazon's dedicated training chips and reduces redundant computation on repeated inputs in asynchronous reinforcement learning. These changes are relevant to developers expanding their training environments and teams training on multiple answers generated from the same question, respectively.
1. Use existing training code on Trainium
TRL v1.16.0, released on October 15, 2026, adds AWS Trainium support and a usage guide. The release notes say it fills remaining support gaps, building on the unified language model output head from the previous version.
The guide says that installing torch-neuronx and then TRL in the same environment lets existing trainer code run unchanged. Below is the guide's quick-start code for supervised fine-tuning, which trains a model to follow example answers.
from datasets import load_dataset
from trl import SFTTrainer
trainer = SFTTrainer(
model="Qwen/Qwen3-0.6B",
train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()2. Align lengths to reduce recompilation
Neuron compiles a computation graph whenever the input shape changes, taking about 1 minute each time. This version changes the default for pad_to_multiple_of to automatic selection. On Neuron, it pads input lengths to multiples of 128 to reduce distinct shapes, while preserving existing behavior on GPUs.
Supervised fine-tuning of Qwen3-0.6B on Trainium2 with 1 Neuron core and a batch size of 4 took 130 seconds for 20 steps with defaults, versus 260 seconds with length alignment disabled. This comparison applies to that example and does not guarantee training speeds for other models.
Model compilation is optional. Using torch_compile=True requires disabling Dynamo's automatic dynamic shapes. With the same model, batch size, and core count in the guide, initial compilation takes about 2 minutes. Subsequent time per step drops from 1.71 seconds without compilation to 1.15 seconds.
3. Compute shared prefixes only once
AsyncGRPO tree attention organizes inputs in a small training batch into a prefix tree. It connects answers branching from the same question like tree branches, computing their shared prefix once rather than repeating it for every answer.
Internally, non-branching portions are merged into segments, each gathering and attending to keys and values along its path. Non-attention layers run once per unique token, and only keys and values are duplicated along shared paths.
The release's Qwen3-1.7B measurements combine forward and backward computation for 32K tokens per generated output. With shared prompts and a packing ratio of 3.37, conventional sequence packing took 8179ms versus 2755ms for tree packing, a reported 2.97-fold speedup. Peak memory fell from 35.2GiB to 13.6GiB. The measurement hardware was not specified.
There is no separate packing selection option. Without shared portions, it takes the same form as conventional sequence packing, with overhead described in the release as about 1%.
4. Check caches and data types for online generation
Online trainers that generate answers during training in this environment cannot use vLLM, so they generate with Transformers. Since the default dynamic cache changes input shapes with every token, the guide recommends a static cache to keep shapes fixed during generation.
The static cache currently requires a bf16 model. A key/value data type mismatch has been reported with the default fp32 model. Prompt lengths are also aligned before generation, but a new cache size still requires recompilation.
Core allocation also needs checking for data-parallel execution. The default configuration of trn2.3xlarge, which contains 1 Trainium2 chip, groups 8 physical cores into 4 logical cores. The guide says supervised fine-tuning with 4 Accelerate processes in this configuration is 3.6 times faster than with 1 core. When selecting only some cores, follow the combinations allowed by the guide.
5. Assess computation savings and setup costs separately
The 2 changes reduce different costs. Tree attention reduces redundant computation and memory for shared inputs, while Trainium length alignment reduces shapes requiring recompilation at the cost of processing extra padding tokens. Choosing model compilation trades initial waiting time for faster subsequent iterations.
Tree attention benchmarks and Trainium usage examples were presented separately. No performance figures combining the 2 features were provided, so their gains should not be added together.
This update expands training hardware choices and improves shared-input efficiency, but adoption decisions require separate assessments of input sharing and compilation setup costs.