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

📄 arXiv: 2405.07518v2 📥 PDF

作者: 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.AR, cs.AI

发布日期: 2024-05-13 (更新: 2024-11-05)

备注: 2024 57th IEEE/ACM International Symposium on Microarchitecture (MICRO)

DOI: 10.1109/MICRO61859.2024.00100

💡 一句话要点

SambaNova SN40L：利用数据流和专家组合扩展AI内存墙

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 专家组合 数据流架构 三层存储 大型语言模型 AI加速器

📋 核心要点

  1. 现有单体LLM训练和部署成本高昂，且现代AI加速器面临计算与内存比率失衡的“内存墙”问题。
  2. 论文提出结合专家组合（CoE）、流式数据流和三层存储系统，以降低成本和复杂性，并扩展AI内存墙。
  3. 实验表明，SambaNova SN40L RDU在CoE推理部署中，显著降低了机器占用空间，加快了模型切换速度，并实现了总体加速。

📝 摘要（中文）

像GPT-4这样的大型单体语言模型（LLM）为现代生成式AI应用铺平了道路。然而，大规模训练、服务和维护单体LLM仍然非常昂贵且具有挑战性。现代AI加速器计算与内存比率的失衡增长造成了内存墙，因此需要新的AI部署方法。专家组合（CoE）是一种替代的模块化方法，可以降低训练和服务的成本和复杂性。然而，当使用传统硬件时，这种方法存在两个关键挑战：（1）在没有融合操作的情况下，较小的模型具有较低的运算强度，这使得实现高利用率更具挑战性；（2）托管大量模型在动态切换时，要么成本过高，要么速度过慢。本文介绍了如何结合CoE、流式数据流和三层存储系统来扩展AI内存墙。我们描述了Samba-CoE，一个拥有150个专家和万亿总参数的CoE系统。我们在SambaNova SN40L可重构数据流单元（RDU）上部署Samba-CoE，这是一种为企业推理和训练应用共同设计的商业数据流加速器架构。该芯片引入了一种新的三层存储系统，包括片上分布式SRAM、片上封装HBM和片外封装DDR DRAM。专用的RDU间网络实现了在多个socket上的横向扩展。我们展示了在八个RDU socket上运行的各种基准测试相比未融合的基线，速度提高了2倍到13倍。我们表明，对于CoE推理部署，8-socket RDU节点将机器占用空间减少了高达19倍，模型切换时间加快了15倍到31倍，并且与DGX H100相比实现了3.7倍的总体加速，与DGX A100相比实现了6.6倍的总体加速。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在训练和推理过程中面临的内存瓶颈问题，即“AI内存墙”。现有方法，特别是基于单体模型的方案，在扩展性和成本效益方面存在局限性。传统硬件在处理专家组合（CoE）模型时，由于运算强度低和模型切换开销大，难以实现高效利用。

核心思路：论文的核心思路是结合CoE的模块化特性、流式数据流的并行处理能力以及三层存储系统的分层存储优势，从而突破内存墙的限制。通过CoE将大型模型分解为多个小型专家模型，降低单个模型的内存需求；利用数据流架构实现专家模型之间的快速切换和并行计算；采用三层存储系统优化数据访问速度，提高整体性能。

技术框架：Samba-CoE系统包含150个专家模型，总参数量达到万亿级别。该系统部署在SambaNova SN40L RDU上，RDU采用三层存储结构：片上SRAM用于存储激活值和权重，片上HBM用于存储模型权重，片外DDR DRAM用于存储更大的模型权重和数据。RDU之间通过专用网络互连，实现多节点扩展。数据在RDU之间以流式方式传输，实现并行计算。

关键创新：论文的关键创新在于将CoE与数据流架构和三层存储系统相结合，从而有效解决了传统硬件在处理大型CoE模型时面临的挑战。这种组合方式不仅降低了内存需求，还提高了计算效率和模型切换速度。此外，SambaNova SN40L RDU本身也是一个创新点，其可重构的数据流架构和三层存储系统为CoE的部署提供了硬件基础。

关键设计：SambaNova SN40L RDU的关键设计包括：可重构的数据流架构，允许根据不同的CoE模型动态调整计算流程；三层存储系统，优化数据访问速度；专用的RDU间网络，实现多节点扩展；以及针对CoE优化的编译器和运行时系统。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在CoE推理部署中，8-socket RDU节点相比DGX H100，实现了3.7倍的总体加速，相比DGX A100，实现了6.6倍的总体加速。同时，机器占用空间减少了高达19倍，模型切换时间加快了15倍到31倍。这些数据表明，SambaNova SN40L RDU在处理大型CoE模型时具有显著的性能优势和成本效益。

🎯 应用场景

该研究成果可广泛应用于需要处理大规模AI模型的场景，如自然语言处理、计算机视觉和推荐系统。特别是在生成式AI领域，可以降低大型语言模型的部署成本和复杂性，加速AI技术的落地和应用。此外，该技术还有潜力应用于其他需要高性能计算和大规模数据处理的领域，如科学计算和金融分析。

📄 摘要（原文）

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 2$\times$ to 13$\times$ 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 19$\times$, speeds up model switching time by 15$\times$ to 31$\times$, and achieves an overall speedup of 3.7$\times$ over a DGX H100 and 6.6$\times$ over a DGX A100.