Nyhet

Hugging Face och AWS: TorchNeuron gör Trainium till en vanlig PyTorch-enhet

Hugging Face har presenterat TorchNeuron, ett PyTorch-plugin från AWS som ger Trainium-chippen en inbyggd backend. Bibliotek som Transformers, Diffusers och TRL kan då köras på Trainium utan särskilda Neuron-versioner av koden. TorchNeuron blir tillgängligt som publik beta den 20 oktober 2026.

Det viktigaste i korthet

• TorchNeuron är ett PyTorch-plugin med en ny backend för AWS Trainium och stöd för eager mode, torch.compile samt FSDP och tensorparallellism (TP).

• Publik beta startar den 20 oktober 2026 och installeras med pip install torch-neuronx.

• Enligt Hugging Face stöds 89 procent av de 100 mest använda arkitekturerna för textgenerering direkt, mot 46 procent med Optimum Neuron v0.4.5.

• Transformers, Diffusers, Accelerate, PEFT och TRL körs direkt på Trainium. Hugging Face rekommenderar nu dessa i stället för Optimum Neuron.

• I träningsexemplet finjusterades Qwen3-8B med TRL på 8 Trainium2-chip (32 NeuronCores), alltså en halv trn2.48xlarge-instans.

• Diffusers-pipelines som FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 och H3 fungerar med tensorparallellism.

Vad har hänt

Hugging Face beskriver TorchNeuron som en ny, inbyggd backend för PyTorch på AWS Trainium. Pluginet gör att ”neuron” fungerar som en vanlig enhet i PyTorch, med eager-körning, PyTorchs egna distribuerade verktyg (DTensor, FSDP2) och torch.compile.

Tidigare har Hugging Face-upplevelsen på Trainium och Inferentia levererats via biblioteket Optimum Neuron. Det bygger på paketet torch-neuronx och PyTorch XLA, vilket enligt Hugging Face innebar att man var tvungen att tänka som en kompilator: lazy tensors i stället för eager-körning, enbart statiska former och omkompilering vid varje grafbrott. Optimum Neuron dolde den komplexiteten med egen kod, bland annat egna modellklasser och tränare. Med TorchNeuron behövs inte det, och Hugging Face rekommenderar att Neuron-användare från nu använder bibliotekens vanliga klasser och skript.

För inferens räcker det att skicka argumentet device_map=neuron för att flytta en modell till Neuron-enheterna. Källan visar ett exempel med tensorparallellism och en statisk cache, vilket rekommenderas starkt för att undvika täta omkompileringar.

För träning körs samma bibliotek som på andra acceleratorer: Transformers, Accelerate, PEFT och TRL. Ingen särskild tränare eller Neuron-specifik kodgaffel behövs. Hugging Face visar ett skript med TRL:s SFTTrainer som startas med torchrun över 32 processer, en per NeuronCore, i ett rutnät där TP_SIZE=4 och FSDP_SIZE=8. Parallelliseringen sköts av Transformers: modellen delas upp när vikterna laddas, och ingen --fsdp-flagga eller Accelerate-FSDP-konfiguration krävs. Träningen kördes klart, och den sparade checkpointen är en vanlig Transformers-checkpoint som kan laddas upp till Hugging Face Hub och användas på annan hårdvara.

Hugging Face nämner några förbehåll. Eftersom formerna ska vara statiska stängs automatiska dynamiska former av i torch.compile, och med TRL 1.16 eller senare fylls batcher på Neuron ut till en multipel av 128 tokens. Den inbyggda asynkrona viktladdningen i Transformers stöds inte på Neuron och stängs av med HF_DEACTIVATE_ASYNC_LOAD=1. Neuron-backenden kan dessutom inte ännu kompilera Dynamos automatiska dynamiska former.

Inom Diffusers fungerar många pipelines i både eager- och kompilerat läge. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 och H3 stöder även tensorparallellism. Fler pipelines ska få stödet, och användare kan själva lägga till det med en guide. I bildexemplet ligger textkodaren kvar på CPU.

Slutligen lyfter Hugging Face fram Kernels-biblioteket. Neuron-teamet publicerar kernels på Hub, och med Neuron Kernel Interface (NKI) kan användare skriva och dela egna operationer. Enligt Hugging Face kan en av dem ge upp till 2,4 gånger snabbare vision transformers på Trainium. AWS får tack för arbetet med TorchNeuron.

Vad det betyder för dig i Sverige

Källan tar inte upp Sverige specifikt. Svenska forskare och utvecklare som använder eller överväger AWS Trainium kan enligt Hugging Face köra samma skript och bibliotek som på andra acceleratorer. Det kan minska behovet av Neuron-specifik kod. Den som i dag använder Optimum Neuron uppmanas av Hugging Face att gå över till Transformers, TRL och övriga bibliotek direkt. Betan kan installeras med pip från den 20 oktober 2026.

Det här vet vi inte än

Källan anger inga priser och inga regioner där Trainium-instanser finns. Den säger inte heller när betan blir en stabil version. Det framgår inte vilka 11 procent av de populära textgenereringsarkitekturerna som saknar stöd, och inga jämförande prestandasiffror mot andra acceleratorer redovisas, förutom påståendet om upp till 2,4 gånger snabbare vision transformers. Det är också oklart vad som händer med Optimum Neuron på sikt, och om TorchNeuron omfattar Inferentia eller bara Trainium.