[论文解读] Chiplet Cloud: Building AI Supercomputers for Serving Large Generative Language Models

Huwan Peng, Scott Davidson+3 more|arXiv (Cornell University)|Jul 5, 2023

Ferroelectric and Negative Capacitance Devices被引用 5

一句话总结

Chiplet Cloud 提出了一种基于小芯片的 ASIC 人工智能超算架构，将所有大型语言模型参数存储在片上 SRAM 中，以消除带宽瓶颈，从而实现每 token 的总拥有成本（TCO）前所未有的降低。通过在数千个芯片上协同设计硬件与软件映射，该架构在 A100 GPU 和 TPUv4 云平台之上分别实现了最高 94 倍和 15 倍的 TCO/Token 降低。

ABSTRACT

Large language models (LLMs) such as OpenAI's ChatGPT and Google's Gemini have demonstrated unprecedented capabilities of autoregressive AI models across multiple tasks triggering disruptive technology innovations around the world. However, as models continue to grow the cost to serve these models also continues to grow threatening the democratization of LLMs. To address this issue, we propose Chiplet Cloud, a chiplet-based ASIC LLM-supercomputer architecture whose goal is to optimize the total cost of ownership (TCO) per generated token. This architecture is a highly parameterizable ASIC and server-level architecture leveraging thousands of replicated accelerator modules collaborating to scale-up the performance of LLMs at cloud-scale. To determine specific parameterizations of the Chiplet Cloud architecture, we implemented a two-phase hardware-software co-design methodology that can search the massive design space and fine tune the architecture across a collection of LLMs based on an accurate inference simulation. A common bottleneck for LLMs is the memory access performance therefore we introduce CC-MEM, a scalable on-chip memory system for Chiplet Cloud architectures. Using the CC-MEM, Chiplet Clouds can be built using only SRAMs for design points where the power and performance of memory access is critical. The CC-MEM also includes a compression decoder module to add support for sparse models without impacting the compute units using a Store-as-Compressed, Load-as-Dense mechanism. We evaluate Chiplet Cloud architectures across eight popular LLMs. Using fine tuned Chiplet Cloud servers we are able to achieve $97 imes$ and $18 imes$ improvement in TCO/Token over rented GPU and TPU clouds, or a $8.3 imes$ and $3.7 imes$ improvement over fabricated GPU and TPU clouds respectively. Chiplet Cloud can also support $1.7 imes$ larger models with a sparsity of 60\%.

研究动机与目标

• 解决现有 GPU 和 TPU 平台在部署大型生成式语言模型（LLMs）时面临的硬件成本上升和能效低下问题。
• 克服由于摩尔定律放缓以及 LLM 推理中内存带宽需求增加所导致的单片芯片可扩展性限制。
• 将每生成一个 token 的总拥有成本（TCO）降至最低，以实现真实应用场景中大规模 LLM 的普惠化访问。
• 设计一种可扩展、成本优化的人工智能超算架构，支持不同模型尺寸和上下文长度的多样化 LLM 工作负载。

提出的方法

• 设计一种基于小芯片的 ASIC 架构，其中数千个复制的芯片小模块协同工作，以并行方式执行 token 生成。
• 将所有模型参数存储在单个小芯片的片上 SRAM 中，以消除片外内存带宽限制，这是大规模 LLM 推理中的关键瓶颈。
• 通过软硬件协同设计，将模型计算映射到各个小芯片上，利用并行性来缓解芯片间通信开销。
• 实施一种全面的设计方法论，探索软硬件联合设计空间，生成所有有效配置下的性能-成本权衡分析。
• 应用帕累托前沿分析，识别在多种 LLM 和上下文长度下最小化 TCO/Token 的最优系统设计。
• 在四个真实世界的大模型（如 GPT-3、PaLM 540B 等）上评估该架构，使用 Lambda 和 Google Cloud 的真实成本模型。

Figure 1 . Compared to A100 GPU and TPUv4, Chiplet Cloud can achieve up to $94\times$ and $15\times$ improvement in TCO/Token on GPT-3 and PaLM 540B, respectively.

实验结果

研究问题

• RQ1 基于小芯片的 ASIC 架构是否能在服务大型生成式语言模型方面，显著降低 TCO/Token，相比现有的 GPU 和 TPU 平台？
• RQ2 通过将所有模型参数存储在片上 SRAM 中，如何消除大规模 LLM 推理中的带宽瓶颈？
• RQ3 在可扩展的人工智能超算设计中，小芯片晶粒尺寸、芯片间通信与系统成本之间的最优权衡是什么？
• RQ4 单一通用的 Chiplet Cloud 服务器设计是否能高效支持多个 LLM 及不同上下文长度，且性能损失最小？
• RQ5 在大规模并行小芯片架构中，软件感知的硬件映射能在多大程度上提升系统效率并降低 TCO？

主要发现

• Chiplet Cloud 在服务 GPT-3 和 PaLM 540B 时，相比 A100 GPU 云平台，实现了最高 94.4 倍的 TCO/Token 改进；相比 TPUv4 云平台，实现了 15.2 倍的改进。
• 通过将所有模型参数存储在片上 SRAM 中，该架构消除了片外内存带宽作为性能瓶颈，从而在更低成本下实现高吞吐量。
• 该设计方法论成功识别出软硬件设计空间中的帕累托最优配置，实现了准确的性能-成本权衡分析。
• 单一通用的 Chiplet Cloud 服务器在支持多个模型和不同上下文长度时，分别仅产生 9% 和 12% 的聚合 TCO/Token 性能损失，相较于单独优化的服务器。
• 该方法通过适度减小芯片尺寸来降低系统成本和能耗，同时借助智能软件映射保持了高计算效率。
• 结果表明，基于小芯片的 ASIC 可通过大幅降低推理的资本和运营成本，实现 LLM 的全球规模部署。

Figure 2 . A generative language model is constructed of multiple decoder layers with the same architecture. It generates outputs in an autoreggressive fashion. FC layers often dominate the computation.