Hugging Face kondigt native ondersteuning aan voor Trainium, de chip van AWS, dankzij TorchNeuron, een nieuwe backend voor PyTorch. Transformers, Diffusers en Kernels kunnen nu rechtstreeks op deze chip draaien, wat het uitvoeren van modellen op Trainium en het delen ervan op de Hugging Face Hub moet vereenvoudigen. Volgens de bron zal een openbare bètaversie beschikbaar zijn vanaf 20 oktober 2026.
TorchNeuron is een extensie van PyTorch, de referentiebibliotheek voor de ontwikkeling van machine learning-modellen. Het voegt een backend toe aan PyTorch die specifiek is voor Trainium, met de « eager » uitvoeringsmodus, torch.compile (een tool voor code-optimalisatie) en native gedistribueerde API’s zoals FSDP en TP. Onderzoekers en ontwikkelaars kunnen het installeren met het commando pip install torch-neuronx.
Een verandering van aanpak voor de Hugging Face-bibliotheken
Tot nu toe verliep het gebruik van Trainium en Inferentia, twee chips ontworpen door AWS, de ene voor economische training van modellen en de andere voor grootschalige inferentie, via de optimum-neuron-bibliotheek. Die was gebaseerd op het torch-neuronx-pakket, zelf gebaseerd op PyTorch XLA, wat betekende dat je moest denken « als een compiler »: luie tensors, alleen statische vormen en hercompilatie bij elke graafbreuk. Hugging Face geeft aan dat optimum-neuron diende om deze complexiteit te verbergen via specifiek onderhouden code.
TorchNeuron verandert deze aanpak door zich rechtstreeks te integreren met PyTorch, dat Neuron ziet als een native apparaat. Trainium wordt dan een accelerator zoals elke andere voor Transformers, Diffusers, Accelerate, PEFT en TRL, met dezelfde klassen en scripts. Hugging Face raadt gebruikers van Neuron nu aan om via deze bibliotheken te gaan in plaats van via Optimum Neuron.
Voor tekstgeneratie is de ondersteuning breed: volgens Hugging Face zijn 89 % van de 100 meest gebruikte tekstgeneratiemodelarchitecturen meteen compatibel, tegenover 46 % met Optimum Neuron v0.4.5. Het volstaat om het argument device_map=neuron door te geven om het model naar Neuron-apparaten te verplaatsen.
Training, diffusie en rekenkernen
Voor training werkt een standaard fine-tuning-script met TRL op een Trainium-instantie, mits enkele Neuron-specifieke aanpassingen. Hugging Face preciseert dat het gepubliceerde voorbeeld Qwen/Qwen3-8B en de dataset HuggingFaceH4/ultrachat_200k gebruikt, en dat het werd uitgevoerd op 8 Trainium2-chips, oftewel 32 NeuronCores, dat wil zeggen de helft van een trn2.48xlarge-instantie. Het geproduceerde checkpoint is een standaard Transformers-checkpoint, bruikbaar op andere hardware.
Verschillende pipelines van Diffusers zijn compatibel met TorchNeuron, zowel in « eager » als gecompileerde modus. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 en H3 ondersteunen ook tensorparallelisme, waarbij elke laag van het model over meerdere kernen wordt verdeeld. Hugging Face geeft aan dat andere pipelines geleidelijk compatibel zullen worden en dat een handleiding het mogelijk maakt om deze ondersteuning zelf toe te voegen.
Ten slotte zal de Kernels-bibliotheek rekenkernen verwelkomen die zijn geoptimaliseerd voor Neuron, geschreven met de Neuron Kernel Interface (NKI). Een kernel voor vision-transformers, gepubliceerd op de Hub, maakt volgens de bron een versnelling « tot x2,4 » mogelijk op Trainium. Hugging Face bedankt het AWS-team voor zijn werk aan TorchNeuron en zijn reactievermogen op feedback van de community.
