SambaNova SN40L：データフローとエキスパートの構成によるAIメモリ壁のスケーリング

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

May 13, 2024

著者: 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

要旨

GPT-4のようなモノリシックな大規模言語モデル（LLM）は、現代の生成AIアプリケーションの道を切り開いてきました。しかし、モノリシックなLLMを大規模にトレーニング、提供、維持することは、依然として非常に高コストで困難です。現代のAIアクセラレータにおける計算能力とメモリ容量の不均衡な増加は、メモリの壁を生み出し、AIを展開するための新しい方法を必要としています。エキスパートの合成（CoE）は、トレーニングと提供のコストと複雑さを低減するモジュール型の代替アプローチです。しかし、このアプローチは従来のハードウェアを使用する際に2つの主要な課題を提示します：（1）融合操作がない場合、小さなモデルは演算密度が低く、高い利用率を達成することがより困難になります；（2）多数のモデルをホストすることは、それらを動的に切り替える際に非常に高コストまたは遅くなる可能性があります。 本論文では、CoE、ストリーミングデータフロー、および3層メモリシステムを組み合わせることで、AIのメモリの壁をどのようにスケーリングするかを説明します。150のエキスパートと合計1兆のパラメータを持つCoEシステムであるSamba-CoEについて説明します。Samba-CoEを、エンタープライズ向けの推論およびトレーニングアプリケーションのために共同設計された商用データフローアクセラレータアーキテクチャであるSambaNova SN40L再構成可能データフローユニット（RDU）に展開します。このチップは、オンチップ分散SRAM、オンパッケージHBM、およびオフパッケージDDR DRAMを備えた新しい3層メモリシステムを導入します。専用のRDU間ネットワークにより、複数のソケットにわたるスケールアップとスケールアウトが可能になります。8つのRDUソケットで実行されるさまざまなベンチマークにおいて、非融合ベースラインと比較して2倍から13倍の高速化を実証します。CoE推論展開において、8ソケットRDUノードはマシンのフットプリントを最大19倍削減し、モデル切り替え時間を15倍から31倍高速化し、DGX H100に対して全体で3.7倍、DGX A100に対して6.6倍の高速化を達成することを示します。

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.