Amazon har tillkännagivit TorchNeuron, en inbyggd PyTorch-backend med öppen källkod för sin Trainium-familj av AI-chips, i ett drag som tar bort ett av de största praktiska hindren för att använda anpassat kisel för arbetsbelastningar för maskininlärning. Enligt ett officiellt inlägg på PyTorch-bloggen som publicerades inför PyTorch Conference North America den 20 oktober, kan utvecklare nu köra befintlig PyTorch-modellkod på Trainium med en enkelradsändring, snarare än att skriva om sina arbetsflöden kring proprietära verktyg.

Tillkännagivandet väckte omedelbar uppmärksamhet eftersom det riktar sig till vad analytiker länge har beskrivit som Nvidias mest hållbara fördel: inte råchipprestanda, utan CUDA-programvaruekosystemet som binder utvecklare till sina GPU:er. Reuters inramade releasen som Amazon och sa att dess AI-chips kommer att fungera med det mest populära utvecklingsramverket inom AI, medan TechCrunch kallade TorchNeuron för företagets största skott på Nvidia hittills. För att få en bredare bild av hur acceleratormarknaden förändras, spårar vår pågående AI-hårdvarunyheter denna utveckling när den utvecklas.

Vad TorchNeuron faktiskt förändras

Kärnlöftet med TorchNeuron är enkelhet. I exemplet publicerat på PyTorch-bloggen körs en befintlig träningsslinga på Trainium efter inget annat än att flytta modellen till den nya enheten med `.to('neuron')`. Kod skriven mot `torch.accelerator`, PyTorchs rekommenderade device-agnostic API, körs oförändrad, enligt inlägget. För kompilerade arbetsbelastningar ställer utvecklare in standardbackend en gång med `torch.compiler.set_default_backend("neuron")` och anropar sedan `model.compile()` utan ytterligare ändringar.

Under huven registrerar TorchNeuron Trainium genom PrivateUse1, PyTorchs officiella förlängningspunkt för anpassade hårdvarubackends, vilket innebär att chippet beter sig som en vanlig PyTorch-enhet snarare än ett tillägg som är fastskruvat på ramverket. AWS beskriver integrationen som spänner över flera exekveringslägen:

  • Ivrigt utförande för forskning och felsökning, med ett experimentellt "adaptivt ivrigt läge" som minskar lanseringskostnader
  • torch.compile för produktionsprestanda, fånga modeller till grafer som Neuron-kompilatorn optimerar, med inbyggt stöd för FlexAttention
  • The Neuron Kernel Interface (NKI) för forskare som behöver anpassade kärnor på låg nivå

Backend fungerar på både Trainium2 och de nyare Trainium3-chippen, enligt tillkännagivandet.

Ett direkt strejk mot Nvidias Software Moat

Nvidias grepp om AI-datorer har aldrig vilat enbart på snabbare hårdvara. CUDA, företagets två decennium gamla mjukvaruplattform, är där det mesta av AI-forskningskoden skrivs, ställs in och distribueras, och varje utmanarchip – från AMD:s Instinct-linje till Googles TPU:er och Amazons Trainium – har fått ta itu med kostnaden för att be utvecklare att lämna den. Bevakningen av TorchNeuron-lanseringen i GeekWire, Benzinga och GuruFocus inramade på ett enhetligt sätt tillkännagivandet som ett angrepp på den vallgraven.

Amazon försöker inte dra ut utvecklare från PyTorch för att vinna över dem. Istället tar det Trainium in i PyTorch, så ramverkets enorma användarbas kan ta till sig chippet utan att ändra vanor. PyTorch-blogginlägget inkluderade ett stöd från Meta, som upprätthåller ramverket: "Inbyggt stöd för AWS Trainium utökar det valet, vilket gör det möjligt för forskare och ingenjörer att experimentera och iterera med det ivriga utförandet och distribuerade API:er de redan känner till", säger Jana van Greunen, chef för PyTorch Engineering på Meta, i inlägget.

Den inramningen är viktig för AWS molnverksamhet. Trainium-instanser har marknadsförts som ett billigare alternativ till Nvidia GPU:er för träning och slutledning, och chipsen har redan lockat marquee AI-kunder. Att få chipsen att kännas infödda för PyTorch tar direkt upp invändningen som har bromsat en bredare användning: porteringsansträngning.

Tekniskt omfång, från bärbara datorer till utbildning i klusterskala

Möjligheterna som listas på PyTorch-bloggen tyder på att Amazon siktar långt bortom hobbyexperiment. TorchNeuron täcker operatörsuppsättningarna sett över text-, multimodala och visionspråksmodeller samt bild- och videogenerering, statistiska rymdmodeller och världsmodeller, med BF16-, FP32- och FP16-precision; TorchAO- och FP8-stöd är i privat beta före allmän tillgänglighet i nästa utgåva.

Distribuerad utbildning, disciplinen som förvandlar enkel-GPU-kod till körningar i klusterskala, stöds genom PyTorchs standard-API:er – DDP, FSDP2, SimpleFSDP, tensor och pipeline-parallellism, och ZeRO-liknande skärning, tillsammans med DTensor och DeviceMesh för flerdimensionella processgrupper. Strömmar och händelser kartlägger PyTorchs semantik på Trainiums hårdvaruköer, sa inlägget.

För slutsatsen bygger Amazons Neuron-stackar på samma grund: vLLM Neuron tjänar stora språkmodeller, medan vLLM Omni Neuron tjänar diffusionsmodeller för bild- och videogenerering. Neuron Kernel Interface låter forskare skriva skräddarsydda operatorer för både ivrig och kompilerad exekvering, och AWS är open-sourcing NKI-lib, en uppsättning Trainium-optimerade kärnor som täcker FlashAttention och MoE blockvisa beräkningar för vanliga LLM och multimodala scenarier.

Prestandaverktyg ingår också. Neuron Explorer utökar standarden "torch.profiler" för att spåra exekvering från ramverk till hårdvara, länka varje händelse tillbaka till dess källlinje, och kan exportera parkettformaterade mätvärden som postanteckningarna kan konsumeras av agenter för automatiserad flaskhalsanalys.

Tillgänglighet och vad som kommer härnäst

TorchNeuron är inställd på sin öppen källkodslansering på PyTorch Conference North America den 20 oktober, enligt blogginlägget, med kompatibilitet som sträcker sig till Hugging Face Transformers v5. AWS beskriver Trainium som en portfölj av chip designade från grunden för molnskala AI-träning och slutledning.

Det strategiska sammanhanget är en marknad där varje större molnleverantör samtidigt är den största köparen av Nvidia-acceleratorer och en potentiell konkurrent. Microsoft, Google och Meta utvecklar alla anpassat kisel, och kinesiska företag inklusive Huawei bygger sina egna mjukvarustackar för att undkomma CUDA:s gravitation. Genom att träffa utvecklare inuti PyTorch istället för att be dem migrera, har Amazon valt den minst störande vägen in i den tävlingen – och den som mest sannolikt, om utvecklarupplevelsen håller i sig, flyttar verkliga träningsbelastningar från Nvidias hårdvara.

Huruvida TorchNeuron verkligen urholkar CUDA:s fördel kommer att bero på prestanda i skala och den långa svansen av verkliga modeller, frågor som lanseringsmaterialet inte besvarar. Men barriären som hindrade de flesta PyTorch-utvecklare från att ens prova Trainium – en omskrivning – har nu sänkts till en enda kodrad.

---

Stay ahead of AI

Få de senaste AI-nyheterna, analyserna och genombrotten – allt på ett ställe.

Läs mer AI-nyheter →