SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and
  Composition of Experts
Authors: name: Raghu Prabhakar; name: Ram Sivaramakrishnan; name: Darshan Gandhi; name: Yun Du; name: Mingran Wang; name: Xiangyu Song; name: Kejie Zhang; name: Tianren Gao; name: Angela Wang; name: Karen Li; name: Yongning Sheng; name: Joshua Brot; name: Denis Sokolov; name: Apurv Vivek; name: Calvin Leung; name: Arjun Sabnis; name: Jiayu Bai; name: Tuowen Zhao; name: Mark Gottscho; name: David Jackson; name: Mark Luttrell; name: Manish K. Shah; name: Edison Chen; name: Kaizhao Liang; name: Swayambhoo Jain; name: Urmish Thakker; name: Dawei Huang; name: Sumti Jairath; name: Kevin J. Brown; name: Kunle Olukotun
Published: 2024-05-13 07:32:45

Abstract: 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.