SambaNova SN40L: Überwindung der KI-Speichermauer mit Datenfluss und Zusammensetzung von Experten

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

May 13, 2024

papers.authors: 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

papers.abstract

Monolithische große Sprachmodelle (LLMs) wie GPT-4 haben den Weg für moderne generative KI-Anwendungen geebnet. Das Training, die Bereitstellung und die Wartung monolithischer LLMs im großen Maßstab bleiben jedoch prohibitiv teuer und herausfordernd. Das unverhältnismäßige Wachstum des Rechen-zu-Speicher-Verhältnisses moderner KI-Beschleuniger hat eine Speichermauer geschaffen, die neue Methoden zur Bereitstellung von KI erforderlich macht. Die Komposition von Experten (CoE) ist ein alternativer modulare Ansatz, der die Kosten und Komplexität des Trainings und der Bereitstellung senkt. Allerdings stellt dieser Ansatz zwei Hauptprobleme bei der Verwendung herkömmlicher Hardware dar: (1) Ohne verschmolzene Operationen haben kleinere Modelle eine niedrigere operationale Intensität, was es schwieriger macht, eine hohe Auslastung zu erreichen; und (2) das Hosting einer großen Anzahl von Modellen kann entweder prohibitiv teuer sein oder langsam, wenn dynamisch zwischen ihnen gewechselt wird. In diesem Paper beschreiben wir, wie die Kombination von CoE, Streaming-Datenfluss und einem dreistufigen Speichersystem die KI-Speichermauer skaliert. Wir beschreiben Samba-CoE, ein CoE-System mit 150 Experten und insgesamt einer Billion Parametern. Wir setzen Samba-CoE auf der SambaNova SN40L Reconfigurable Dataflow Unit (RDU) ein - einer kommerziellen Datenfluss-Beschleunigerarchitektur, die für Unternehmensinferenz- und Trainingsanwendungen gemeinsam entwickelt wurde. Der Chip führt ein neues dreistufiges Speichersystem mit verteiltem SRAM auf dem Chip, HBM auf dem Package und DDR DRAM außerhalb des Packages ein. Ein dediziertes Inter-RDU-Netzwerk ermöglicht das Skalieren über mehrere Sockets hinweg. Wir zeigen Beschleunigungen von 2x bis 13x bei verschiedenen Benchmarks, die auf acht RDU-Sockets im Vergleich zu einem nicht verschmolzenen Baseline-System ausgeführt werden. Wir zeigen, dass für CoE-Inferenzbereitstellungen der 8-Socket-RDU-Knoten den Maschinen-Footprint um bis zu 19x reduziert, die Modellwechselzeit um das 15- bis 31-fache beschleunigt und insgesamt eine Beschleunigung von 3,7x gegenüber einem DGX H100 und 6,6x gegenüber einem DGX A100 erreicht.

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.