Amazon laat PyTorch-code zonder herschrijven op Trainium draaien
3 mins read

Amazon laat PyTorch-code zonder herschrijven op Trainium draaien

Amazon Web Services zet vandaag TorchNeuron live als publieke bèta. Met die software draait gewone PyTorch-code op Trainium, de eigen AI-chips van Amazon, zonder dat ontwikkelaars hun model hoeven om te bouwen. De lancering valt samen met de PyTorch Conference in San Jose.

Eén regel code

Wie tot nu toe PyTorch op Trainium wilde draaien, moest de code aanpassen aan interfaces die alleen voor Amazons chips bestaan. Volgens de aankondiging op de PyTorch-blog is dat nu grotendeels voorbij. Een model gaat met .to('neuron') naar de chip. Code die al het apparaat-neutrale torch.accelerator gebruikt, draait zelfs helemaal ongewijzigd.

Technisch haakt TorchNeuron in op PrivateUse1, het officiële aansluitpunt dat PyTorch heeft voor hardware van derden. Daardoor blijven vertrouwde onderdelen als torch.distributed en torch.profiler gewoon werken. De eager-modus (handig bij onderzoek en debuggen) wordt ondersteund, net als torch.compile voor productie. Het geheel werkt op Trainium2 en Trainium3 en verschijnt onder de Apache 2.0-licentie op GitHub. Installeren gaat via pip.

Technicus tussen rijen verlichte serverkasten in een datacenter

Eerste klanten zien lagere kosten

Amazon laat in de blogpost een paar vroege gebruikers aan het woord. Adam Sypniewski, CTO van spraakbedrijf Deepgram, zegt dat Trainium in hun eigen benchmarks ongeveer 60 procent meer seconden audio per dollar verwerkt dan hun bestaande GPU-opstelling. Arjun Desai, medeoprichter van Cartesia, schrijft dat zijn team de inferentiepijplijn met minimale wijzigingen kon overzetten. Hugging Face maakt Trainium een standaard backend voor Transformers en Diffusers.

Ook videostartup Decart staat op de partnerlijst. Dat is opvallend, want Amazon zou volgens eerdere berichten praten over een overname van Decart voor 7 miljard dollar.

Versie één heeft nog gaten

Amazon is open over wat nog ontbreekt. In de testronde van augustus werden ruim 22.000 tests uit de officiële PyTorch-testset overgeslagen omdat de betreffende operators nog geen eigen Neuron-implementatie hebben. Dynamische vormen van invoer (dynamic shapes) worden in deze release niet ondersteund, en hetzelfde geldt voor sparse tensors en NestedTensor. FP8 en de kwantisatiebibliotheek TorchAO zitten nog in een besloten bèta en moeten bij de volgende release algemeen beschikbaar komen.

Gericht op het CUDA-voordeel van Nvidia

De grootste troef van Nvidia is al jaren de software. Veel AI-code is geschreven met CUDA in het achterhoofd, en dat maakt overstappen naar andere chips duur. GeekWire ziet TorchNeuron dan ook als een directe aanval op die voorsprong. Volgens The Next Web zei Amazon-topman Andy Jassy in april dat de eigen chiptak 50 miljard dollar waard kan worden.

Tegelijk blijft Amazon zelf een grote Nvidia-klant: AWS bestelde eerder nog 2 miljoen extra GPU’s. Andere clouds bouwen ook aan eigen silicium, zoals Microsoft met de Maia 300, terwijl AWS ook met Qualcomm aan inferentiechips werkt.

Wat betekent dit

Voor teams die op AWS draaien wordt het een stuk goedkoper om Trainium gewoon eens te proberen. Een test kost geen maanden herschrijfwerk meer, hooguit een middag. Of de beloofde besparingen standhouden, hangt af van je model: wie leunt op dynamic shapes of exotische operators loopt in deze bèta snel tegen grenzen aan. Nvidia hoeft zich op korte termijn weinig zorgen te maken. De grote clouds slopen wel stap voor stap de software-drempel rond hun eigen chips, en dat zet op termijn druk op de prijzen van GPU-rekenkracht.