导出设置

预览

Stanford CS336 Language Modeling from Scratch | Spring 2025 | Inference

类型: Detailed Summary 模型: gemini-2.5-pro-preview-06-05 创建时间: 2025-06-08 10:32

好的,这是优化后的内容:

本讲核心观点

本讲座聚焦于语言模型推理(Inference)的效率优化,探讨了“推理为什么慢”以及“如何让它变快”两大核心问题。

推理之所以慢,根源在于其与训练不同的特性:它通常是内存受限的,而非计算受限,并且需要处理动态变化的工作负载。尤其对于Transformer架构,自回归式的逐个token生成机制,特别是其中的Attention层,导致计算强度极低,成为性能瓶颈。

为提升效率,讲座介绍了四大类优化策略:
1. 架构内部优化:通过改进模型架构来减少内存消耗的“大户”——KV缓存。具体技术包括分组查询注意力(GQA)、多头潜在注意力(MLA)、跨层注意力(CLA)和局部注意力。
2. 探索替代架构:跳出Transformer框架,采用状态空间模型(SSM,如Mamba)、线性注意力、扩散模型等从根本上改变计算范式,以实现效率的飞跃。
3. 模型压缩技术:通过量化(降低数值精度)和模型剪枝(移除冗余部分并蒸馏)来减小模型体积,从而降低内存带宽压力。
4. 算法与系统优化:采用推测采样等无损加速算法,利用“检查比生成快”的原理提升速度;同时,借鉴操作系统思想,通过连续批处理和PagedAttention等技术,高效管理内存以应对动态工作负载。

所有优化的最终目标都是在准确性与效率之间找到最佳平衡,通过架构创新、算法优化和系统工程的结合,实现推理速度和吞吐量的最大化。

理解推理工作负载

概览

推理,即使用预训练好的模型根据提示生成响应,是语言模型价值变现的核心环节。它广泛应用于以下场景:
* 实际应用:聊天机器人、代码补全、批量数据处理等。
* 模型评估:评估模型的指令遵循等能力。
* 测试时计算:在输出最终答案前进行更多的“思考”(生成中间步骤)。
* 强化学习训练:生成样本以供奖励模型评分。

效率的重要性:训练是一次性投入,而推理则需频繁、大量地重复执行。据统计,OpenAI和Cursor等公司每天处理的生成任务量已达百亿甚至千亿级别,推理成本日益凸显。

关键性能指标

  • 首个token生成时间 (Time To First Token, TTFT):用户从输入到看到第一个输出token的等待时间,对交互式应用体验至关重要。
  • 延迟 (Latency):后续token的生成速度,通常以“秒/token”衡量,同样影响交互体验。
  • 吞吐量 (Throughput):系统单位时间内生成的总token数(tokens/秒),是衡量批量处理能力的核心指标。

效率的关键瓶颈:为何推理是内存受限的?

  • 训练 (监督式):可以一次性看到所有token,因此能在整个序列上并行计算(如Transformer中的矩阵乘法),充分利用硬件的计算能力,是计算受限的。
  • 推理 (自回归):必须逐个生成token,因为下一个token依赖于之前所有token。这种串行特性使得并行化变得困难,无法充分利用GPU的计算单元,导致大部分时间花在从内存中读写数据上,因此是内存受限 (memory limited) 的。

行业现状

无论是提供闭源模型的巨头(OpenAI, Google),还是服务于开源模型的公司(Together, Fireworks),都投入大量精力优化推理效率。社区也涌现出VLLM、TensorRT、TGI等优秀的开源推理框架。

计算强度分析:量化推理瓶颈

计算强度 (Arithmetic Intensity) 是衡量计算效率的核心指标,定义为“计算量(FLOPs)”与“内存访问量(Bytes)”之比。高计算强度意味着硬件在忙于计算(计算受限),低则意味着硬件在等待数据(内存受限)。

  • 矩阵乘法示例 ($X (B \times D) \cdot W (D \times F)$):其计算强度约等于批处理大小 B
  • 以H100 GPU为例:其硬件计算强度约为295。这意味着,要让H100这样的高端硬件“忙起来”(计算受限),而不是“闲着等数据”(内存受限),我们一次至少需要处理约295个序列的矩阵运算。
  • 推理的窘境:在生成阶段,我们一次只处理一个token,相当于批处理大小 B=1。其计算强度远小于295,差距悬殊,因此必然是内存受限。

推理的两个阶段与KV缓存

为了避免每次生成新token都重复计算整个历史序列,推理过程引入了KV缓存——存储并复用所有先前token的键(Key)和值(Value)向量。这使得推理分为两个不同特性的阶段:

  1. Prefill (预填充):并行处理用户输入的提示(prompt),计算并填充初始的KV缓存。此阶段类似于训练,可以并行处理,通常是计算受限的。
  2. Generation (生成):逐个生成新token。每生成一个token,只需计算当前token的Q、K、V,并与缓存中历史K、V进行Attention计算。此阶段是串行的,是内存受限的。

MLP层与Attention层的计算强度

  • MLP层
    • 生成阶段 ($T=1$):计算强度约为 B(并发请求数)。需要足够大的并发量才能提升效率。
  • Attention层
    • 生成阶段 ($T=1$):计算强度约为 S/(S+1) ≈ 1(S为上下文长度)。
    • 结论:Attention层的计算强度在生成阶段始终接近1,且无法通过批处理(增大B)改善,这是其成为内存瓶颈的根本原因。因为MLP权重对所有请求是共享的,而KV缓存是每个请求独有的,导致Attention层的内存访问量随并发数线性增长,抵消了计算量的增长。

减少KV缓存大小:最直接的优化路径

既然推理的主要瓶颈是内存,而KV缓存是内存消耗的“大户”,那么最直接的优化思路就是用各种方法在不严重影响模型效果的前提下,把KV缓存变小。

分组查询注意力 (Grouped-Query Attention, GQA)

  • 问题:标准多头注意力(MHA)中,每个查询头(Query Head)都有一套独立的键/值头(Key/Value Head),导致KV缓存较大。
  • 思想:让多组查询头共享同一套键/值头。这是在MHA(N个Q配N个K/V)和多查询注意力MQA(N个Q配1个K/V)之间的一种折中。
  • 效果:将KV缓存大小减少 N/K 倍(N为查询头数,K为键/值头数),显著提升吞吐量。Llama 2 (70B) 和 Llama 3 等模型均已采用。

多头潜在注意力 (Multi-Head Latent Attention, MLA)

  • 思想:不减少键/值头的数量,而是将每个键/值向量投影到一个更低的维度空间,从而“压缩”KV缓存。
  • 示例 (DeepSeek v2):将16384维的K/V向量压缩至512维。
  • 效果:同样能显著减少KV缓存,带来可观的性能提升,且有研究声称其准确性甚至优于MHA。

跨层注意力 (Cross-Layer Attention, CLA)

  • 思想:不仅在头之间共享,还在不同层之间共享KV缓存。
  • 效果:进一步优化了准确性与KV缓存大小之间的平衡关系。

局部注意力 (Local Attention)

  • 思想:只让token关注其邻近的上下文(滑动窗口注意力),使KV缓存大小与序列长度无关。
  • 问题与方案:会损害长距离依赖建模能力。因此通常与全局注意力层交错使用(混合层),如Character AI每6层中设1个全局层。

Transformer的替代方案:跳出框架的革命

除了在Transformer内部修修补补,我们还可以探索全新的架构,从根本上改变计算范式。

状态空间模型 (State-Space Models, SSM)

  • 思想:借鉴信号处理领域的思想,用循环的方式(RNN-like)处理序列,将Attention的二次复杂度降为线性。
  • 代表
    • S4:早期模型,在特定长上下文任务上表现优异,但语言能力不足。
    • Mamba:通过使模型参数依赖于输入,解决了SSM在语言任务上的短板,性能可与同规模Transformer媲美。
    • Jamba:混合Transformer与Mamba层,取长补短。
  • 优势:将随序列长度增长的 $O(T)$ KV缓存,变为固定大小的 $O(1)$ 状态,极大提升长序列推理效率。

扩散模型 (Diffusion Models)

  • 思想:颠覆自回归范式,并行生成所有token,然后通过多步迭代的方式对整个序列进行“去噪”和“修正”。
  • 优势:并行生成,有望实现极高的token输出速度。
  • 示例 (Mercury Coder):在代码生成任务上,其速度远超Transformer类模型,同时保持了有竞争力的准确性。

模型压缩:让模型更“轻”

量化 (Quantization)

  • 思想:降低模型权重和激活值的数值精度(如从16位浮点数降至8位甚至4位整数),直接减少模型大小和内存传输量。
  • 关键技术
    • LLM.int8():为解决异常值(outliers)对量化的干扰,该技术将异常值分离出来用高精度处理,其余大部分值进行int8量化。
    • 激活感知量化:根据激活值动态决定哪些权重需要保持高精度。

模型剪枝 (Model Pruning)

  • 思想:识别并移除模型中不重要的部分(如层、头、神经元),然后通过知识蒸馏等方法“修复”剪枝后模型的性能。
  • 效果:可以得到一个参数量更少、速度更快,但性能与原始大模型相当甚至更好的新模型。例如,NVIDIA通过此方法从15B模型剪枝得到性能优于Mistral 7B的8B模型。

推测采样:无损加速的巧妙算法

该方法利用了“检查已生成的token序列”比“从头生成单个token”快得多的特性。

  • 思想
    1. 使用一个更小、更快的草稿模型 (draft model) 预先生成一段候选token。
    2. 使用目标大模型 (target model) 并行地、一次性地验证这些候选token。
    3. 根据一套数学上严谨的接受/拒绝规则,采纳草稿模型“猜对”的token,对“猜错”的部分由大模型重新生成。
  • 类比:这个过程好比让一个反应快、知识稍逊的“学霸”(草稿模型)先快速写出答案草稿,再由知识渊博但批阅速度很快的“教授”(目标大模型)来并行检查和修正,从而在保证最终答案质量的同时,大大缩短了等待时间。
  • 关键特性无损。数学上保证最终生成的序列与直接用大模型生成完全一致。
  • 效果:通常能带来2倍左右的加速,且草稿模型的设计有很大的创新空间(如Medusa, EAGLE等改进方案)。

处理动态工作负载:系统层面的智慧

现实世界的推理服务面临请求到达时间不一、序列长度各异等挑战,传统的静态批处理效率低下。

连续批处理 (Continuous Batching)

  • 思想:采用迭代级调度。每生成一步token,调度器就检查是否有新请求到达,并动态地将其加入或移出批处理队列,从而消除不必要的等待,最大化GPU利用率。

PagedAttention (VLLM)

  • 思想:借鉴操作系统中的虚拟内存分页机制来管理KV缓存。
  • 解决的问题:传统方式为每个请求预分配连续的大块内存,导致严重的内部碎片(空间浪费)和外部碎片(空间不可用)。
  • 解决方案
    • 将KV缓存分割成非连续的、固定大小的块 (blocks)
    • 通过一个逻辑地址到物理地址的映射表来管理这些块,消除了内存碎片。
    • 通过写时复制 (copy-on-write) 技术,高效实现多序列间的KV缓存共享(如共享系统提示),进一步节省内存。