SambaNova SN40L: Масштабирование стены памяти искусственного интеллекта с помощью потоков данных и композиции экспертов

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

Аннотация

Монолитные большие языковые модели (LLM), такие как GPT-4, проложили путь для современных генеративных приложений искусственного интеллекта. Однако обучение, обслуживание и поддержка монолитных LLM в масштабе остаются чрезмерно дорогими и сложными. Неравномерное увеличение соотношения вычислений к памяти современных ускорителей искусственного интеллекта создало "стену памяти", требуя новых методов развертывания ИИ. Композиция экспертов (CoE) представляет собой альтернативный модульный подход, снижающий стоимость и сложность обучения и обслуживания. Однако этот подход сталкивается с двумя ключевыми проблемами при использовании обычного оборудования: (1) без объединенных операций у менее крупных моделей ниже операционная интенсивность, что затрудняет достижение высокой утилизации; и (2) размещение большого количества моделей может быть либо чрезмерно дорогим, либо медленным при динамическом переключении между ними. В данной статье мы описываем, как комбинирование CoE, потоковой модели данных и трехуровневой системы памяти преодолевает "стену памяти" искусственного интеллекта. Мы описываем Samba-CoE, систему CoE с 150 экспертами и общим числом параметров в триллиона. Мы развертываем Samba-CoE на устройстве SambaNova SN40L Reconfigurable Dataflow Unit (RDU) - коммерческой архитектуре ускорителя потоков данных, спроектированной для предприятий для применения в выводе и обучении. Этот чип вводит новую трехуровневую систему памяти с распределенной SRAM на кристалле, HBM на пакете и DDR DRAM вне пакета. Выделенная сеть между RDU позволяет масштабироваться вверх и вширь по нескольким сокетам. Мы демонстрируем ускорение от 2 до 13 раз на различных тестах на восемь сокетах RDU по сравнению с базовой моделью без объединенных операций. Мы показываем, что для развертывания вывода CoE узел RDU с восемью сокетами сокращает машинное пространство до 19 раз, ускоряет время переключения модели на 15-31 раз и достигает общего ускорения в 3,7 раза по сравнению с DGX H100 и в 6,6 раза по сравнению с 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.