Build on Trainium

Build on Trainium è un programma di investimento da 110 milioni di dollari incentrato sulla ricerca sull’IA e sulla formazione universitaria per supportare la prossima generazione di innovazione e sviluppo su AWS Trainium. AWS Trainium è un chip di array sistolico di IA progettato appositamente per promuovere idee e applicazioni di IA all’avanguardia. Build on Trainium finanzia nuove ricerche sull’IA su Trainium, investendo in team accademici leader per creare innovazioni in aree critiche, tra cui nuove architetture di modelli, librerie di ML, ottimizzazioni, sistemi distribuiti su larga scala e altro ancora. Questa iniziativa pluriennale pone le basi per il futuro dell’IA ispirando la comunità accademica ad avvalersi di Trainium, investendo nella soluzione e contribuendo alla community open source che ruota intorno a essa. Unendo questi vantaggi con il Software Development Kit (SDK) Neuron e il recente lancio della Neuron Kernel Interface (NKI), i clienti di Trainium possono ora innovare su larga scala nel cloud.

Abbiamo creato un cluster di ricerca Trainium dedicato che può includere fino a 40.000 chip Trainium. Questi saranno accessibili tramite istanze Amazon EC2 Trn1 connesse su un’unica rete nell’ordine di grandezza dei peta-bit non bloccante utilizzando gli UltraClusters Amazon EC2. I team di ricerca e gli studenti possono accedere a questi chip tramite le prenotazioni autogestite di blocchi di capacità utilizzando Blocchi di capacità Amazon EC2 per ML.

Stiamo conducendo diversi cicli di inviti a presentare proposte (CFP) per gli Amazon Research Awards (ARA) rivolti all’ampia comunità di ricerca: le proposte selezionate ricevono crediti AWS Trainium e l’accesso al cluster di ricerca Trainium. Build on Trainium accoglie le proposte di ricerca che intendono impiegare le librerie e i framework di ML open source più diffusi e possono contribuire all’open source per migliorare le risorse per la community di sviluppatori di ML.

Neuron Kernel Interface (NKI) è una nuova interfaccia di programmazione per i chip di IA di AWS, Trainium e Inferentia. NKI fornisce l’accesso diretto alle primitive hardware e alle istruzioni disponibili su AWS Trainium e Inferentia, consentendo ai ricercatori di creare e ottimizzare i kernel di calcolo per prestazioni ottimali. È un ambiente di programmazione basato su Python che adotta la diffusa sintassi in stile Triton e la semantica a livello di riquadro. I ricercatori possono utilizzare NKI per migliorare i modelli di deep learning con nuove funzionalità, ottimizzazioni e innovazioni scientifiche. Per ulteriori informazioni, consulta la pagina della documentazione di NKI.

Presso il Device Realization Lab del MIT, stiamo utilizzando AWS Trainium per superare i limiti della ricerca sull’intelligenza artificiale nel settore sanitario. I nostri modelli di segmentazione ecografica 3D e di stima della velocità del suono si addestrano in modo estremamente rapido ed efficiente, riducendo i periodi sperimentali di oltre la metà e ottenendo livelli di precisione all’avanguardia. AWS Trainium ci ha permesso di scalare la nostra ricerca in modi non possibili con i tradizionali sistemi GPU. Addestrando le nostre reti neurali convoluzionali 3D su AWS Trainium (trn.32xlarge), abbiamo raggiunto prestazioni all’avanguardia con un throughput superiore del 50% e costi inferiori rispetto alle istanze NVIDIA A100. Utilizzando un cluster Trainium a 32 nodi, abbiamo condotto in modo efficiente oltre 180 esperimenti di ablazione, riducendo il tempo totale di addestramento da mesi a settimane, e accelerando l’innovazione dell’IA in campo medico al MIT. In futuro, intendiamo utilizzare Trainium per addestrare modelli di agenti IA in grado di gestire e automatizzare il flusso di lavoro dell’ecografia digitale, consentendo ai medici di risparmiare molto tempo e di fornire un’assistenza migliore ai pazienti.”

“Il gruppo di ricerca CMU Catalyst lavora sull’ottimizzazione dei sistemi di ML. Il nostro progetto aspira a semplificare l’ottimizzazione tra diversi sistemi di ML. Trainium è l’unica soluzione che offre un controllo di basso livello e un’interfaccia di programmazione accessibile tramite la Neuron Kernel Interface (NKI).

con il supporto di AWS attraverso il programma Build on Trainium, il nostro ricercatore è stato in grado di esplorare ottimizzazioni avanzate su un kernel critico: FlashAttention. Quello che ci ha stupito maggiormente è stata la velocità con cui abbiamo potuto iterare: abbiamo ottenuto miglioramenti significativi rispetto allo stato dell’arte precedente in appena una settimana utilizzando NKI, Neuron profiler e documentazione dell’architettura accessibili pubblicamente. La combinazione di strumenti potenti e informazioni chiare sull’hardware ha reso accessibile al nostro team un’ottimizzazione sofisticata e di basso livello.

AWS Trainium e Neuron Kernel Interface (NKI) consentono ai ricercatori come noi di innovare più velocemente, eliminando gli ostacoli che in genere rallentano il lavoro di ottimizzazione specifico dell’hardware.”

 

“Grazie al programma Build on Trainium, il suo team ha ottenuto l’accesso completo al nuovo stack del compilatore open source NKI di AWS Neuron, inclusa la visibilità diretta su ISA e sulle API di Trainium per una pianificazione e un’allocazione della memoria precise. Questo livello di visibilità e controllo consente ai suoi studenti di analizzare più facilmente le opportunità di ottimizzazione e scoprire in modo più efficace implementazioni performanti.”

Christopher Fletcher, Professore associato di informatica presso la University of California, Berkeley
 

“L’accesso ad AWS Trainium e Inferentia è stato fondamentale per far progredire la nostra ricerca e formazione su sistemi di intelligenza artificiale efficienti e su larga scala. Utilizziamo queste piattaforme per l’addestramento “mixture of experts” e l’ottimizzazione delle inferenze, al fine di prototipare nuove tecniche di runtime e pianificazione che migliorano la scalabilità, l’efficienza e la portabilità sulle architetture degli acceleratori emergenti. Grazie allo stack Neuron Developer, i ricercatori della UIUC stanno sviluppando nuovo runtime e nuove tecniche di pianificazione che migliorano l’efficienza e la portabilità dei carichi di lavoro dell’IA. Il team è rimasto particolarmente colpito dall’apertura dello stack Neuron Developer, che rende queste piattaforme preziose per la ricerca sul runtime e consente innovazioni in termini di sparsità, gerarchie di memoria ed efficienza comunicativa che vanno oltre le tradizionali architetture GPU.”

“Grazie ad AWS Trainium e al programma Build on Trainium, io e i miei studenti siamo stati in grado di accelerare in modo significativo le nostre simulazioni di circuiti quantistici. Il progetto ha riunito un forte gruppo di studenti che hanno sviluppato in collaborazione un simulatore ad alte prestazioni, consentendo una sperimentazione più approfondita e un apprendimento pratico in una scala che, semplicemente, non era possibile in precedenza.”

“Il nostro team di ricerca presso la UTS sta esplorando l’integrazione di algoritmi di watermarking tree-ring sviluppando kernel Neuron NKI personalizzati. Avere accesso allo stack open source Neuron tramite il programma Build on Trainium è qualcosa di rivoluzionario. Ci offre una visibilità senza precedenti sull’architettura di Trainium e la capacità di lavorare direttamente a livello hardware. L’accesso a Trainium ha consentito al nostro team di accelerare in modo significativo i nostri carichi di lavoro di watermarking, riducendo i cicli di iterazione e consentendoci di esplorare modelli e tecniche più complessi. Questa profondità di accesso consente ai nostri ricercatori di prototipare nuove idee, sperimentare ottimizzazioni di basso livello e superare i limiti di ciò che i sistemi di watermarking possono ottenere con i moderni acceleratori di intelligenza artificiale.”