SambaNova SN40L: Superare il Muro della Memoria nell'IA con il Dataflow e la Composizione di Esperti

SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts

May 13, 2024

Autori: Raghu Prabhakar, Ram Sivaramakrishnan, Darshan Gandhi, Yun Du, Mingran Wang, Xiangyu Song, Kejie Zhang, Tianren Gao, Angela Wang, Karen Li, Yongning Sheng, Joshua Brot, Denis Sokolov, Apurv Vivek, Calvin Leung, Arjun Sabnis, Jiayu Bai, Tuowen Zhao, Mark Gottscho, David Jackson, Mark Luttrell, Manish K. Shah, Edison Chen, Kaizhao Liang, Swayambhoo Jain, Urmish Thakker, Dawei Huang, Sumti Jairath, Kevin J. Brown, Kunle Olukotun

cs.AI

Abstract

I modelli linguistici monolitici di grandi dimensioni (LLM) come GPT-4 hanno aperto la strada alle moderne applicazioni di intelligenza artificiale generativa. Tuttavia, l'addestramento, il servizio e la manutenzione di LLM monolitici su larga scala rimangono proibitivamente costosi e complessi. L'aumento sproporzionato del rapporto tra capacità di calcolo e memoria negli acceleratori AI moderni ha creato un "muro della memoria", rendendo necessari nuovi metodi per implementare l'AI. La Composizione di Esperti (CoE) è un approccio modulare alternativo che riduce i costi e la complessità dell'addestramento e del servizio. Tuttavia, questo approccio presenta due sfide principali quando si utilizza hardware convenzionale: (1) senza operazioni fuse, i modelli più piccoli hanno un'intensità operativa inferiore, rendendo più difficile raggiungere un'elevata utilizzazione; e (2) ospitare un gran numero di modelli può essere proibitivamente costoso o lento quando si passa dinamicamente tra di essi. In questo articolo, descriviamo come la combinazione di CoE, flusso di dati in streaming e un sistema di memoria a tre livelli affronti il muro della memoria dell'AI. Presentiamo Samba-CoE, un sistema CoE con 150 esperti e un totale di mille miliardi di parametri. Implementiamo Samba-CoE sull'Unità di Flusso di Dati Riconfigurabile (RDU) SambaNova SN40L, un'architettura commerciale di acceleratore a flusso di dati co-progettata per applicazioni di inferenza e addestramento aziendali. Il chip introduce un nuovo sistema di memoria a tre livelli con SRAM distribuita on-chip, HBM on-package e DRAM DDR off-package. Una rete dedicata inter-RDU consente di scalare verticalmente e orizzontalmente su più socket. Dimostriamo accelerazioni che vanno da 2x a 13x su vari benchmark eseguiti su otto socket RDU rispetto a una baseline non fusa. Mostriamo che, per le implementazioni di inferenza CoE, il nodo RDU a 8 socket riduce l'ingombro delle macchine fino a 19x, accelera il tempo di commutazione del modello da 15x a 31x e raggiunge un'accelerazione complessiva di 3,7x rispetto a un DGX H100 e di 6,6x rispetto a un DGX A100.

English

Monolithic large language models (LLMs) like GPT-4 have paved the way for modern generative AI applications. Training, serving, and maintaining monolithic LLMs at scale, however, remains prohibitively expensive and challenging. The disproportionate increase in compute-to-memory ratio of modern AI accelerators have created a memory wall, necessitating new methods to deploy AI. Composition of Experts (CoE) is an alternative modular approach that lowers the cost and complexity of training and serving. However, this approach presents two key challenges when using conventional hardware: (1) without fused operations, smaller models have lower operational intensity, which makes high utilization more challenging to achieve; and (2) hosting a large number of models can be either prohibitively expensive or slow when dynamically switching between them. In this paper, we describe how combining CoE, streaming dataflow, and a three-tier memory system scales the AI memory wall. We describe Samba-CoE, a CoE system with 150 experts and a trillion total parameters. We deploy Samba-CoE on the SambaNova SN40L Reconfigurable Dataflow Unit (RDU) - a commercial dataflow accelerator architecture that has been co-designed for enterprise inference and training applications. The chip introduces a new three-tier memory system with on-chip distributed SRAM, on-package HBM, and off-package DDR DRAM. A dedicated inter-RDU network enables scaling up and out over multiple sockets. We demonstrate speedups ranging from 2x to 13x on various benchmarks running on eight RDU sockets compared with an unfused baseline. We show that for CoE inference deployments, the 8-socket RDU Node reduces machine footprint by up to 19x, speeds up model switching time by 15x to 31x, and achieves an overall speedup of 3.7x over a DGX H100 and 6.6x over a DGX A100.