Repository navigation
[Tracking] TRL on AWS Trainium (Neuron) #7658
Description
This tracks TRL support on AWS Trainium with TorchNeuron, the native PyTorch backend (device neuron, torch.compile(backend="neuron")). transformers and accelerate already handle the device (huggingface/transformers#44302 sets it up in TrainingArguments, huggingface/transformers#47035 makes neuron the default compile backend, and accelerate has MULTI_NEURON).
Setup used for everything below: a trn2 instance exposing 1 Trainium2 device (4 logical NeuronCores, 96 GB). torch 2.14.0, torch-neuronx 2.14.3, neuronx-cc 2.28, Python 3.14, transformers 5.19.0, accelerate 1.15.0, peft 0.21.2. Each trainer ran 3 steps on a tiny model with each option below; SFT and GRPO also ran on Qwen3-0.6B.
Status
| TRL 1.14.2 | TRL 1.15 | main | |
|---|---|---|---|
| Reward | ✅ | ✅ | ✅ |
| SFT, DPO, KTO, GRPO, RLOO, Distillation | ✅ | ❌ step 1 (#7659) | ✅ |
| + LoRA, packing, GC off | ✅ | ❌ (#7659) | ✅ |
| + eval, checkpoint save, GA, loss types (IPO, DFT), conversational and prompt-completion data | not tested | ❌ (#7659) | ✅ |
torch_compile=True |
❌ step 2 (#7660) | ❌ | ❌ step 2 (#7660) |
torch_compile=True + automatic_dynamic_shapes=False |
✅ | ❌ (#7659) | ✅ |
activation_offloading=True |
❌ (#7662) | ❌ (#7662) | ✅ |
GRPO/RLOO cache_implementation="static" |
❌ (#7664) | ❌ (#7664) | ✅ with a bf16 model |
| Model families (SFT): Qwen2, Qwen3, Llama 3.2, Mistral, Gemma 2, Qwen3-MoE, GPT-OSS | ✅ (Llama, Gemma 2, Qwen3-MoE, GPT-OSS tested) | ❌ (#7659) | ✅ all 7 |
| VLM (Qwen2-VL): SFT, SFT + LoRA, DPO, GRPO | not tested | ❌ (#7659) | ✅ |
torch_compile + static shapes on other families |
✅ Gemma 2, Qwen3-MoE, SmolLM2-135M (Llama arch); ❌ the tiny Llama 3.2 test model (neuronx-cc NCC_ITEN406, specific to its tiny dimensions) |
||
| Data parallel, 2 NeuronCores (cores 0,1) | ✅ SFT | ❌ (#7659) | ✅ SFT, DPO, GRPO, SFT + LoRA |
| Data parallel, cores 2,3 and all 4 cores | ✅ SFT | ||
| Data parallel, cores 1,2 | ❌ collective init fails in the Neuron runtime (no_hier no_mesh); pairs must be (0,1)/(2,3) |
Sub-issues
- Fused LM head requires Triton: SFT, DPO, KTO, GRPO, RLOO and Distillation fail on Trainium (and any non-CUDA device) since v1.15 #7659 Fused LM head requires Triton: SFT, DPO, KTO, GRPO, RLOO and Distillation fail on Trainium since v1.15
-
torch_compile=Truefails on the second step on Trainium: the Neuron dynamo backend can't run the dynamic-shape recompile #7660torch_compile=Truefails on the second step (dynamic shapes in the Neuron dynamo backend; upstream, with workaround). Fix: [Trainer] Compile with static shapes when the backend is neuron transformers#49617 - Docs: a Trainium (Neuron) page #7661 Docs: a Trainium page (bucket shapes with
pad_to_multiple_of, compile with static shapes, generation) -
activation_offloading=Truecrashes on Trainium:torch.cuda.Streamis used for every non-XPU/NPU accelerator #7662activation_offloading=Truecrashes on Trainium -
trl envreportsaccelerator(s): cpuon Trainium #7663trl envreportsaccelerator(s): cpuon Trainium - GRPO/RLOO
cache_implementation="static"fails with the default fp32 model: the static cache is fp32, the keys are bf16 #7664 GRPO/RLOOcache_implementation="static"fails with the default fp32 model (not Neuron-specific). Fix: Allocate the StaticCache values with the value states' dtype transformers#49554 - GRPO/RLOO/Distillation:
pad_to_multiple_ofis not applied to the prompts passed togenerate#7665 GRPO/RLOO/Distillation:pad_to_multiple_ofis not applied to the prompts passed togenerate(on Trainium, a GRPO step whose prompt bucket has been seen takes 25-32 s, against 112-325 s with the static cache alone) - Default
pad_to_multiple_ofto"auto": pad to a multiple of 128 on Neuron #7724 Defaultpad_to_multiple_ofto"auto", which pads to a multiple of 128 on Neuron and staysNoneelsewhere - Read
is_triton_availableonce in the fused LM head kernels so Dynamo does not warn #7731 Dynamo warns aboutis_triton_availableundertorch_compile - GRPO/RLOO on Trainium: each new prompt length bucket recompiles generation for ~12 minutes #7733 GRPO/RLOO on Trainium: each new prompt length bucket recompiles generation for ~12 minutes
Not started
- Neuron CI. TRL CI runs on
aws-g6e-4xlarge/aws-g5-12xlarge-cache. We need a Trn2 runner group reachable from huggingface/trl; optimum-neuron usesaws-inf2-8xlargeand forgeforge-use2. Until then, GPU CI covers the torch tile path by running the kernels on CPU tensors and comparing them to Triton (Run the fused LM head and distillation kernels on devices without Triton #7669). - vLLM on Neuron for GRPO/RLOO generation
- An NKI kernel for the fused LM head, if profiling shows the torch tile path matters
Activity
torch_compile=Truefails on the second step on Trainium: the Neuron dynamo backend can't run the dynamic-shape recompile #7660activation_offloading=Truecrashes on Trainium:torch.cuda.Streamis used for every non-XPU/NPU accelerator #7662trl envreportsaccelerator(s): cpuon Trainium #7663cache_implementation="static"fails with the default fp32 model: the static cache is fp32, the keys are bf16 #7664pad_to_multiple_ofis not applied to the prompts passed togenerate#7665