SambaNova SN40L: Superando la barrera de memoria en IA con flujo de datos y composición de expertos

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

May 13, 2024

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

Resumen

Los modelos de lenguaje monolíticos de gran escala (LLMs, por sus siglas en inglés) como GPT-4 han allanado el camino para las aplicaciones modernas de IA generativa. Sin embargo, entrenar, servir y mantener LLMs monolíticos a gran escala sigue siendo prohibitivamente costoso y desafiante. El aumento desproporcionado en la relación computación-memoria de los aceleradores de IA modernos ha creado una "barrera de memoria", lo que ha hecho necesario desarrollar nuevos métodos para implementar la IA. La Composición de Expertos (CoE, por sus siglas en inglés) es un enfoque modular alternativo que reduce el costo y la complejidad del entrenamiento y la implementación. No obstante, este enfoque presenta dos desafíos clave cuando se utiliza hardware convencional: (1) sin operaciones fusionadas, los modelos más pequeños tienen una menor intensidad operativa, lo que dificulta alcanzar una alta utilización; y (2) alojar un gran número de modelos puede ser prohibitivamente costoso o lento al cambiar dinámicamente entre ellos. En este artículo, describimos cómo la combinación de CoE, flujo de datos en streaming y un sistema de memoria de tres niveles escala la barrera de memoria de la IA. Presentamos Samba-CoE, un sistema CoE con 150 expertos y un billón de parámetros en total. Implementamos Samba-CoE en la Unidad de Flujo de Datos Reconfigurable (RDU, por sus siglas en inglés) SN40L de SambaNova, una arquitectura de acelerador de flujo de datos comercial que ha sido co-diseñada para aplicaciones empresariales de inferencia y entrenamiento. El chip introduce un nuevo sistema de memoria de tres niveles con SRAM distribuida en el chip, HBM en el paquete y DRAM DDR fuera del paquete. Una red inter-RDU dedicada permite escalar vertical y horizontalmente a través de múltiples sockets. Demostramos aceleraciones que van de 2x a 13x en varios benchmarks ejecutados en ocho sockets RDU en comparación con una línea base no fusionada. Mostramos que, para implementaciones de inferencia de CoE, el nodo RDU de 8 sockets reduce la huella de la máquina hasta 19x, acelera el tiempo de cambio de modelo entre 15x y 31x, y logra una aceleración general de 3.7x sobre un DGX H100 y 6.6x sobre 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.