SambaNova SN40L: Het overkomen van de AI-geheugenbarrière met dataflow en compositie van 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

Samenvatting

Monolithische grote taalmodellen (LLM's) zoals GPT-4 hebben de weg geëffend voor moderne generatieve AI-toepassingen. Het trainen, serveren en onderhouden van monolithische LLM's op schaal blijft echter buitengewoon kostbaar en uitdagend. De onevenredige toename in de compute-to-memory-ratio van moderne AI-accelerators heeft een geheugenmuur gecreëerd, wat nieuwe methoden vereist om AI te implementeren. Composition of Experts (CoE) is een alternatieve modulaire aanpak die de kosten en complexiteit van trainen en serveren verlaagt. Deze aanpak brengt echter twee belangrijke uitdagingen met zich mee bij het gebruik van conventionele hardware: (1) zonder gefuseerde operaties hebben kleinere modellen een lagere operationele intensiteit, wat het bereiken van een hoge benutting moeilijker maakt; en (2) het hosten van een groot aantal modellen kan buitengewoon kostbaar of traag zijn bij het dynamisch schakelen tussen hen. In dit artikel beschrijven we hoe de combinatie van CoE, streaming dataflow en een drie-lagen geheugensysteem de AI-geheugenmuur schaalbaar maakt. We beschrijven Samba-CoE, een CoE-systeem met 150 experts en een totaal van een biljoen parameters. We implementeren Samba-CoE op de SambaNova SN40L Reconfigurable Dataflow Unit (RDU) – een commerciële dataflow-acceleratorarchitectuur die mede is ontworpen voor enterprise-inferentie- en trainingsapplicaties. De chip introduceert een nieuw drie-lagen geheugensysteem met on-chip gedistribueerde SRAM, on-package HBM en off-package DDR DRAM. Een toegewijd inter-RDU-netwerk maakt schaalbaarheid mogelijk over meerdere sockets. We demonstreren snelheidswinsten variërend van 2x tot 13x op verschillende benchmarks die draaien op acht RDU-sockets in vergelijking met een niet-gefuseerde baseline. We laten zien dat voor CoE-inferentie-implementaties de 8-socket RDU Node de machinefootprint tot 19x verkleint, de modelwisseltijd met 15x tot 31x versnelt en een algehele snelheidswinst van 3.7x behaalt ten opzichte van een DGX H100 en 6.6x ten opzichte van een 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.