SambaNova SN40L : Surmonter le mur de la mémoire en IA grâce au flux de données et à la composition d'experts

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

May 13, 2024

Auteurs: 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

Résumé

Les modèles de langage monolithiques de grande taille (LLMs) comme GPT-4 ont ouvert la voie aux applications modernes d'IA générative. Cependant, l'entraînement, le déploiement et la maintenance de ces LLMs monolithiques à grande échelle restent extrêmement coûteux et complexes. L'augmentation disproportionnée du ratio calcul-mémoire des accélérateurs d'IA modernes a créé un "mur de mémoire", nécessitant de nouvelles méthodes pour déployer l'IA. La Composition d'Experts (CoE) est une approche modulaire alternative qui réduit les coûts et la complexité de l'entraînement et du déploiement. Cependant, cette approche présente deux défis majeurs lors de l'utilisation de matériel conventionnel : (1) sans opérations fusionnées, les modèles plus petits ont une intensité opérationnelle réduite, ce qui rend plus difficile d'atteindre une utilisation élevée ; et (2) héberger un grand nombre de modèles peut être soit prohibitivement coûteux, soit lent lors de la commutation dynamique entre eux. Dans cet article, nous décrivons comment la combinaison de la CoE, du flux de données en continu et d'un système mémoire à trois niveaux permet de surmonter le mur de mémoire de l'IA. Nous présentons Samba-CoE, un système CoE comprenant 150 experts et un total de mille milliards de paramètres. Nous déployons Samba-CoE sur le SambaNova SN40L Reconfigurable Dataflow Unit (RDU) - une architecture d'accélérateur de flux de données conçue conjointement pour les applications d'inférence et d'entraînement en entreprise. Le processeur introduit un nouveau système mémoire à trois niveaux avec de la SRAM distribuée sur puce, de la HBM sur package et de la DRAM DDR hors package. Un réseau inter-RDU dédié permet une montée en puissance et en échelle sur plusieurs sockets. Nous démontrons des accélérations allant de 2x à 13x sur divers benchmarks exécutés sur huit sockets RDU par rapport à une base de référence non fusionnée. Nous montrons que pour les déploiements d'inférence CoE, le nœud RDU à 8 sockets réduit l'empreinte machine jusqu'à 19x, accélère le temps de commutation de modèle de 15x à 31x, et obtient une accélération globale de 3,7x par rapport à un DGX H100 et de 6,6x par rapport à 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.