AI21 Labs addestra un modello di linguaggio basato su 178 miliardi di parametri utilizzando le istanze P4d di Amazon EC2, PyTorch

AI21 Labs utilizza il machine learning per sviluppare modelli linguistici incentrati sulla comprensione del significato: nel 2021 ha fissato l'obiettivo di addestrare il Jurassic-1 Jumbo recentemente rilasciato, un modello linguistico autoregressivo con 178 miliardi di parametri. Gli sviluppatori che si registrano per il beta test avranno accesso a Jurassic-1 Jumbo e potranno iniziare immediatamente a personalizzare il modello in base al loro caso d'uso. La startup software desiderava addestrare il modello in modo efficiente, quindi si è rivolta ad Amazon Web Services (AWS) e ha creato una soluzione utilizzando Amazon Elastic Compute Cloud (Amazon EC2), un servizio Web che fornisce capacità di elaborazione sicure e ridimensionabili nel cloud. La scelta di Amazon EC2 ha dato all'azienda il controllo sul processo di addestramento, inclusa l'allocazione dei nodi.

Per potenti funzionalità di elaborazione e di rete, l'azienda ha scelto le istanze P4d di Amazon EC2, che offrono una velocità di trasmissione effettiva elevata e una rete a bassa latenza per l'addestramento del machine learning e per le applicazioni di calcolo ad alte prestazioni nel cloud. Utilizzando le istanze P4d di Amazon EC2, AI21 Labs ha ottenuto le prestazioni e la memoria richieste distribuendo l'addestramento dei modelli su centinaia di GPU per fornire l'elaborazione del linguaggio naturale come servizio attraverso il suo modello Jurassic-1 Jumbo. Poiché ora l'azienda addestra e controlla il proprio modello su larga scala, può lavorare allo sviluppo di nuovi modelli sulla stessa scala e innovare con maggiore facilità.

Fondata nel 2017, AI21 Labs persegue una missione ibrida: condurre ricerche sull'elaborazione del linguaggio naturale e sviluppare prodotti per la lettura e la scrittura basati sull'intelligenza artificiale. Il suo prodotto di punta, Wordtune, è un assistente di scrittura e modifica intelligente lanciato nell'ottobre 2020 ed è cresciuto fino a supportare quasi un milione di utenti. L'altro prodotto principale, AI21 Studio, offre l'accesso tramite API ai modelli in linguaggio Jurassic-1 dell'azienda e lo sviluppo di modelli personalizzati. "Facciamo parte di un piccolo gruppo di aziende che offrono modelli linguistici come servizio, consentendo a chiunque, dagli sviluppatori indipendenti alle multinazionali, di creare app e servizi basati su una tecnologia di elaborazione del linguaggio naturale avanzata", afferma Yoav Shoham, cofounder e co-CEO di AI21 Labs. "Stiamo inoltre perseguendo innovazioni scientifiche e affrontando le sfide in materia di ingegneria del software poste da modelli di queste dimensioni e complessità".

Per addestrare in modo efficiente il suo primo megamodello di deep learning e supportare le relative esigenze di alta scalabilità e prestazioni, AI21 Labs aveva bisogno di potenza di elaborazione, una velocità di rete efficiente e accesso ad assistenza e supporto tecnico. Per questi motivi, all'inizio del 2021 l'azienda ha iniziato a implementare una soluzione su AWS, scegliendo di addestrare il modello utilizzando le istanze P4d di Amazon EC2. Queste istanze sono distribuite in cluster hyperscale denominati Amazon EC2 UltraCluster, che forniscono oltre 4.000 GPU NVIDIA A100, un'infrastruttura di rete non bloccante nell'ordine dei petabit e storage ad alta velocità di trasmissione effettiva e bassa latenza. 

L'approccio dell'azienda è stato ulteriormente ottimizzato su GPUDirectRDMA a bassa latenza e larghezza di banda elevata, insieme a Elastic Fabric Adapter (EFA), un'interfaccia di rete per istanze Amazon EC2 che consente ai clienti di eseguire applicazioni che richiedono comunicazioni di alto livello tra nodi su larga scala su AWS. A causa delle dimensioni del modello, il team aveva bisogno di utilizzare l'elaborazione parallela per ottenere tempi di addestramento efficienti, quindi si è rivolto alle funzionalità di rete di AWS per supportare l'addestramento distribuito e il parallelismo dei modelli. "Le istanze P4d di Amazon EC2 offrono reti ad alte prestazioni da 400 Gbps su EFA", afferma Opher Lieber, Technical Lead per Jurassic in AI21 Labs. "La velocità di rete da GPU a GPU influisce direttamente sulla capacità di scalare in modo efficiente e mantenere l'efficienza nei costi quando si passa a centinaia di GPU".

AI21 Labs ha iniziato implementando la sua base di codice sulle istanze P4d di Amazon EC2 attivate per EFA. Quindi ha testato e verificato le prestazioni e la scalabilità efficiente del suo approccio di addestramento multinodo. Successivamente, il team ha avviato un rapido addestramento sul modello completo, che utilizza centinaia di GPU, per verificare funzionalità e prestazioni. Da lì è stata in grado di iniziare ad addestrare il suo modello Jurassic-1 Jumbo su AWS. Per l'orchestrazione, l'azienda ha scelto una soluzione interna che alloca le istanze utilizzando un kit di sviluppo software AWS, l'AWS SDK per Python (Boto3), che semplifica l'integrazione dell'applicazione, della libreria o dello script Python di un cliente con diversi servizi AWS.

Per lo storage, AI21 Labs ha scelto Amazon Simple Storage Service (Amazon S3), che offre scalabilità, disponibilità dei dati, sicurezza e prestazioni all'avanguardia nel settore. "Con l'aiuto del team AWS siamo riusciti a ottenere ottime prestazioni su Amazon S3, quindi è stata una scelta facile sia in termini di prestazioni che di prezzo", afferma Lieber. Il team utilizza i bucket Amazon S3 per archiviare e caricare i checkpoint in modo efficiente e distribuito. Per registrare i progressi e gli eventi dell'addestramento, il team usa Amazon CloudWatch, un servizio di monitoraggio e osservabilità. 

Durante l'implementazione della sua soluzione, AI21 Labs ha sfruttato il supporto di AWS. Il suo team ha consultato gli specialisti di AWS, i quali hanno fornito indicazioni in risposta alle domande e ai dubbi relativi al livello di servizio, all'architettura e all'hardware. Inoltre, l'azienda ha migliorato le prestazioni di Jurassic-1 Jumbo utilizzando PyTorch in AWS, un framework di deep learning open source che semplifica lo sviluppo di modelli di machine learning e la loro implementazione in produzione. 

AI21 Labs ha completato l'addestramento in diversi mesi, terminandolo a giugno 2021. Il nuovo megamodello, un modello linguistico autoregressivo, ha 178 miliardi di parametri, paragonabili all'offerta della concorrenza dell'azienda. Offre anche un vocabolario differenziato di 256.000 elementi che fornisce funzionalità estese di rappresentazione del testo e supporto per le entità denominate. L'azienda offre ora Jurassic-1 Jumbo (insieme alla sua controparte, Jurassic-1 Large, che ha 7 miliardi di parametri) in open beta tramite l'offerta AI21 Studio dell'azienda. Il servizio consente a un'ampia gamma di sviluppatori di creare prodotti sul modello Jurassic-1 Jumbo, e AI21 Labs è già stato adottato in molti settori, dal marketing alla creazione di contenuti, ai giochi, alla ricerca medica, all'automotive, alle telecomunicazioni e alla finanza.

Possedendo e avendo accesso diretto al proprio modello, AI21 può adattarsi e innovare senza dipendere da terze parti e può esplorare obiettivi di innovazione continui, che sono una parte fondamentale della sua missione. AI21 Labs è attualmente impegnata nella prototipazione di altri modelli, che prevede di addestrare su larga scala. "L'addestramento e la proprietà dei nostri megamodelli continueranno a essere un fattore di differenziazione fondamentale per le nostre offerte Wordtune e AI21 Studio", afferma Shoham.