Google lancia uno strumento in grado di scalare e parallelizzare le reti neurali
GSPMD separa la programmazione di un modello ML dalla parallelizzazione ed è in grado di scalare la maggior parte delle architetture di rete di deep learning
Google AI ha lanciato GSPMD – Parallelizzazione generale e scalabile per grafici di calcolo ML, per affrontare le sfide di scalabilità. GSPMD è in grado di scalare la maggior parte delle architetture di rete di deep learning ed è stato applicato a molti modelli di deep learning che includono GShard-M4, BigSSL, LaMDA , ViT e MetNet-2. GSPMD è stato inoltre integrato in più framework ML, inclusi TensorFlow e JAX, che utilizzano XLA come compilatore condiviso.
La soluzione separa il compito di programmare un modello ML dalla sfida della parallelizzazione. Consente agli sviluppatori di modelli di scrivere programmi come se fossero eseguiti su un singolo dispositivo con memoria e capacità di calcolo molto elevate. L’utente deve solo aggiungere poche righe di codice di annotazione a un sottoinsieme di tensori critici nel codice del modello per indicare come partizionare i tensori. Con GSPMD, gli sviluppatori possono impiegare diversi algoritmi di parallelismo per diversi casi d’uso senza la necessità di reimplementare il modello.
La separazione tra programmazione del modello e parallelismo consente agli sviluppatori di ridurre al minimo la duplicazione del codice. GSPMD è progettato per supportare un’ampia varietà di algoritmi di parallelismo con un’astrazione e un’implementazione uniformi. Supporta anche modelli nidificati di parallelismo. La soluzione facilita l’innovazione sugli algoritmi di parallelismo consentendo agli esperti di prestazioni di concentrarsi sugli algoritmi che utilizzano al meglio l’ hardware , invece dell’implementazione che implica molte comunicazioni tra dispositivi.
Nel recente set di benchmark delle prestazioni MLPerf , un modello di solo encoder simile a BERT con circa 500 miliardi di parametri a cui il team ha applicato GSPMD per la parallelizzazione su 2048 chip TPU-V4, ha prodotto risultati altamente competitivi, utilizzando fino al 63% del picco FLOPS offerti dai TPU-V4. Come meccanismo condiviso e robusto per diverse modalità di parallelismo, GSPMD consente agli utenti di passare comodamente da una modalità all’altra in diverse parti di un modello. Ciò è particolarmente utile per i modelli che possono avere componenti diversi con caratteristiche prestazionali distinte, come i modelli multimodali che gestiscono sia le immagini che l’audio.
“Poiché ciò richiede spesso la costruzione di modelli più grandi e ancora più complessi, siamo lieti di condividere il documento GSPMD e la corrispondente libreria open source con la più ampia comunità di ricerca e speriamo che sia utile per un training efficiente di reti neurali profonde su larga scala. “, hanno scritto Yuanzhong Xu e Yanping Huang, ingegneri del software; Google Research, Brain Team, nel post del blog .
