SambaNova SN40L: Superando a Barreira de Memória em IA com Fluxo de Dados e Composição de Especialistas

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

Resumo

Modelos monolíticos de linguagem de grande escala (LLMs, na sigla em inglês), como o GPT-4, abriram caminho para aplicações modernas de IA generativa. No entanto, treinar, servir e manter LLMs monolíticos em escala continua sendo proibitivamente caro e desafiador. O aumento desproporcional na relação computação-memória dos aceleradores de IA modernos criou uma "parede de memória", exigindo novos métodos para implantação de IA. A Composição de Especialistas (CoE, na sigla em inglês) é uma abordagem modular alternativa que reduz o custo e a complexidade do treinamento e da execução. No entanto, essa abordagem apresenta dois desafios principais ao usar hardware convencional: (1) sem operações fundidas, modelos menores têm menor intensidade operacional, o que torna mais difícil alcançar alta utilização; e (2) hospedar um grande número de modelos pode ser proibitivamente caro ou lento ao alternar dinamicamente entre eles. Neste artigo, descrevemos como a combinação de CoE, fluxo de dados em streaming e um sistema de memória de três camadas escala a parede de memória da IA. Apresentamos o Samba-CoE, um sistema CoE com 150 especialistas e um trilhão de parâmetros totais. Implantamos o Samba-CoE na Unidade de Fluxo de Dados Reconfigurável (RDU, na sigla em inglês) SN40L da SambaNova – uma arquitetura comercial de acelerador de fluxo de dados co-projetada para aplicações empresariais de inferência e treinamento. O chip introduz um novo sistema de memória de três camadas com SRAM distribuída no chip, HBM no pacote e DRAM DDR fora do pacote. Uma rede dedicada entre RDUs permite escalar vertical e horizontalmente em múltiplos sockets. Demonstramos acelerações que variam de 2x a 13x em vários benchmarks executados em oito sockets RDU em comparação com uma linha de base não fundida. Mostramos que, para implantações de inferência CoE, o Nó RDU de 8 sockets reduz a pegada de máquina em até 19x, acelera o tempo de troca de modelo em 15x a 31x e alcança uma aceleração geral de 3,7x em relação a um DGX H100 e 6,6x em relação a um 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.