Amazon brengt vandaag TorchNeuron uit als open-source software. Daarmee draait bestaande PyTorch-code native op de AWS Trainium-chips. Volgens Amazon is één configuratiewijziging genoeg en hoeft er niets herschreven te worden. Ontwikkelaars kunnen zo makkelijker overstappen naar een alternatief voor GPU-hardware.
TorchNeuron wordt uitgebracht onder de Apache 2.0-licentie en ondersteunt Trainium2 en Trainium3. Volgens Amazon schrijft 85 procent van de AI-onderzoekers en -ontwikkelaars code in PyTorch. Wie die code op andere chips wilde draaien, was tot nu toe vaak maanden bezig met herschrijven en opnieuw valideren.
Technisch maakt TorchNeuron gebruik van het mechanisme waarmee nieuwe hardware zich al in PyTorch kan registreren. Daardoor werken de tools die ontwikkelaars gebruiken om te testen, te debuggen en te schalen zonder aanpassingen op Trainium. Er is ook integratie met TorchTitan, het native distributed training-framework van PyTorch, inclusief Fully Sharded Data Parallelism (FSDP).
Amazon heeft een grote bijdrage geleverd aan vLLM en zijn complete set performance-kernels open-source gemaakt. Hugging Face ondersteunt Trainium inmiddels native in zijn ontwikkeltools, zonder Amazon-specifieke code. Daarnaast draagt Amazon wijzigingen upstream bij aan PyTorch om code te verwijderen die het framework aan specifieke chips bindt. Elke build wordt getest tegen honderdduizenden AI-modellen.
Van spraak-AI tot medicijnonderzoek
Tot de eerste gebruikers behoort Cartesia, dat werkt aan conversationele AI. Het bedrijf draaide zijn standaard inference-pipeline op Trainium met minimale aanpassingen aan de PyTorch-code. Het Lieber Institute for Brain Development, dat onderzoek doet naar onder meer schizofrenie, gebruikt TorchNeuron om AI-modellen te verkennen die rekenintensieve moleculaire dynamica-simulaties kunnen versnellen of vervangen.
TorchNeuron volgt op een investering in eigen chips die al tien jaar loopt. Sinds Amazon in 2015 chipontwerper Annapurna Labs overnam, heeft het bedrijf vier chipfamilies ontwikkeld, elk voor een specifieke cloudworkload.