当我们写几行提示与大型语言模型交互时,我们几乎立即收到十四行诗、调试建议或复杂分析。

这种以软件为中心的观点可能掩盖了一个基本现实:人工智能不仅仅是软件问题。它是一个涉及电子通过硅移动的物理问题,以及在内存和计算单元之间移动大量数据的挑战。

然而,像 LLM 这样的复杂 AI 工具不能仅使用 CPU 构建。这是因为 CPU 设计用于逻辑、分支决策和串行执行。另一方面,深度学习需要线性代数、大规模并行和概率操作。

在本文中,我们将看看 GPU 和 TPU 如何帮助构建现代 LLM 以及它们背后的架构。

矩阵乘法是核心

在其核心,每个神经网络执行一个基本操作数十亿次:矩阵乘法。

当我们问 LLM 问题时,我们的词语被转换为数字,流过数千亿次乘加运算。单个 700 亿参数模型的前向传递需要超过 140 万亿次浮点运算。

数学结构很简单。每个层执行 Y = W * X + B,其中 X 表示输入数据,W 包含学习参数,B 是偏置向量,Y 是输出。当我们将其扩展到数十亿参数时,我们正在执行数万亿次简单乘法和加法。

使这个工作负载特殊的是它对并行计算的依赖。矩阵运算中的每个乘法完全独立。计算第 1 行乘以第 1 列不需要等待第 2 行乘以第 2 列。我们可以将工作分割到数千个处理器上,计算期间零通信开销。

Transformer 架构放大了这种并行性。自注意力机制计算每个令牌与每个其他令牌之间的关系分数。作为参考,对于 4,096 令牌上下文窗口,这创建超过 1600 万注意力对。每个 transformer 层执行几个主要矩阵乘法,700 亿参数模型每个前向传递可以执行数百万这些操作。

CPU 的局限性

CPU 擅长需要复杂逻辑和分支决策的任务。现代 CPU 包含为不可预测代码路径设计的复杂机制,但神经网络不需要这些功能。

分支预测机制在预测条件语句结果方面达到 93-97% 准确率,消耗大量硅面积。神经网络几乎没有分支。它们以可预测模式数十亿次执行相同操作。

乱序执行重新排序指令以保持处理器在等待数据时忙碌。矩阵乘法具有不需要这种复杂性的完美可预测访问模式。大缓存层次结构(L1、L2、L3)隐藏随机访问的内存延迟,但神经网络数据顺序流过内存。

这意味着 CPU die 只有一小部分专用于算术。大多数晶体管预算用于控制单元管理乱序执行、分支预测和缓存一致性。当运行 LLM 时,这些数十亿晶体管空闲,消耗功率并占据本可用于算术单元的空间。

除了计算低效,CPU 面临更根本的限制:内存墙。这个术语描述处理器速度和内存访问速度之间日益扩大的差距。大型语言模型巨大。700 亿参数模型以 16 位精度存储占用约 140GB 内存。生成单个令牌,处理器必须从内存读取每个参数以执行必要的矩阵乘法。

传统计算机遵循冯·诺依曼架构,处理器和内存通过共享总线通信。执行任何计算,CPU 必须获取指令、从内存检索数据、执行操作,并写回结果。处理器和内存之间这种持续的信息传输创建计算机科学家称为冯·诺依曼瓶颈的东西。

无论增加多少核心数或时钟速度都无法解决这个问题。瓶颈不是算术运算,而是数据交付到处理器的速率。这就是为什么内存带宽而不是计算能力通常决定 LLM 性能。

GPU 架构

图形处理单元最初设计用于渲染视频游戏。渲染数百万像素的数学要求与深度学习惊人地相似,因为两者都需要大规模并行和高吞吐浮点算术。

NVIDIA 的 GPU 架构使用 SIMT(单指令多线程)。基本单元是称为 warp 的 32 线程组。warp 中的所有线程共享单个指令解码器,同时执行相同指令。这种共享控制单元节省大量硅面积,用数千个算术单元填充。

虽然现代 CPU 有 16 到 64 个复杂核心,NVIDIA H100 包含近 17,000 个更简单的核心。它们以更低时钟速度运行(1-2 GHz 对比 3-6 GHz),但大规模并行补偿更慢的单个操作。

标准 GPU 核心在单个线程上一次执行单个数字操作。认识到 AI 工作负载由矩阵运算主导,NVIDIA 从 Volta 架构开始引入 Tensor Core。Tensor Core 是专用硬件单元,在单个时钟周期执行整个矩阵乘积累加运算。标准核心每个周期完成一个浮点运算,Tensor Core 执行 4×4 矩阵乘法涉及 64 个单独操作(乘步骤 16 次乘法和 16 次加法,加上 16 次累加)瞬间。这代表矩阵运算吞吐量的 64 倍改进。

Tensor Core 还支持混合精度算术,这对实际 AI 部署至关重要。它们可以接受较低精度格式如 FP16 或 BF16 的输入(使用 FP32 一半内存),同时在更高精度 FP32 中累加结果以保持数值准确性。这种组合增加吞吐量并减少内存需求,而不牺牲稳定模型训练和准确推理所需的精度。

为了喂养这些数千个计算单元,GPU 使用高带宽内存(HBM)。与插在主板上的独立模块上的 DDR 内存不同,HBM 由使用硅通孔(微观垂直线)垂直堆叠在一起的 DRAM die 组成。这些堆叠放置在直接相邻 GPU die 的硅中介层上,最小化数据必须传输的物理距离。

这种架构允许 GPU 实现超过 3,350 GB/s 的内存带宽在 H100 上,比 CPU 快 20 多倍。有了这个带宽,H100 可以在约 0.04 秒加载 140GB 模型,实现每秒 20 或更多令牌的令牌生成速度。这是生硬、令人沮丧的交互和自然对话节奏之间的区别。

大规模并行计算和极端内存带宽的组合使 GPU 成为 AI 工作负载的主导平台。

TPU 架构

2013 年,Google 计算出如果每个用户每天仅使用语音搜索三分钟,他们将需要使用 CPU 加倍数据中心容量。这导致张量处理单元(TPU)。

定义特征是脉动阵列,互连算术单元的网格(256×256,总共 65,536 个处理器)。权重加载到阵列并保持固定,而输入数据水平流动。每个单元将其存储的权重乘以传入数据,加到垂直流动的_running sum_,并将两个值传递给邻居。

这种设计意味着中间值从不接触主内存。从 DRAM 读取消耗大约比乘法多 200 倍的能量。通过保持结果在相邻处理器之间流动,脉动阵列消除大部分内存访问开销,实现比 CPU 高 30 到 80 倍的性能每瓦。

Google 的 TPU 没有缓存、分支预测、乱序执行或推测预取。这种极端专门化意味着 TPU 不能运行通用代码,但对于矩阵运算,效率增益巨大。Google 还引入了 bfloat16,使用 8 位指数(匹配 FP32 范围)和 7 位尾数。神经网络容忍低精度但需要宽范围,使这种格式理想。

实践意义

理解硬件差异有直接实际影响。

训练和推理有根本不同的需求。

  • 训练存储参数、梯度和优化器状态。作为参考,总内存达到参数计数的 16 到 20 倍。例如,训练 4050 亿参数的 LLaMA 3.1 需要 16,000 个 80GB 的 H100 GPU。
  • 推理更宽容。它跳过后向传播,需要更少操作。这就是为什么我们可以在训练不足的消费者 GPU 上运行 70 亿参数模型。

批处理对效率重要。GPU 通过同时处理多个输入实现峰值性能。每个额外输入摊销加载权重的成本。单请求推理通常未充分利用并行硬件。

从 CPU 到 GPU 和 TPU 的转变代表计算哲学的根本转变。CPU 体现逻辑和顺序操作时代,优化为低延迟。GPU 和 TPU 代表通过概率操作进行数据转换时代。这些是线性代数的专用引擎,通过压倒性并行算术实现结果。

本文为学习目的的个人翻译,译文仅供参考。

原文链接:Why AI Needs GPUs and TPUs: The Hardware Behind LLMs

版权归原作者或原刊登方所有。本文为非官方译本;如有不妥,请联系删除。