推理阶段解析​

大多数主流的仅解码大语言模型(LLM),如 GPT-4、Qwen 系列等,本质上是下一个词预测器。其推理过程可分为预填充阶段和解码阶段。在预填充阶段,LLM 处理输入 token 以计算中间状态(keys 和 value),用于生成 “第一个” token 。由于输入全部已知,此阶段运算多为高度并行化矩阵运算,能充分利用 GPU。而解码阶段,LLM 一次自回归生成一个输出 token,直至满足停止条件。每个输出 token 都依赖之前迭代的所有输出状态(keys 和 values),此阶段类似矩阵向量运算,GPU 计算能力未被充分利用,数据从内存传输到 GPU 的速度决定了延迟,属于内存限制操作。​

批处理与 KV 缓存​

批处理是提高 GPU 利用率和有效吞吐量的常用方法。多个请求共享同一模型,权重内存成本得以分散,大批量数据一次性传输到 GPU 处理,提升了资源利用率。但批大小增加到一定程度可能导致内存溢出。传统的静态批处理并非最佳,因为不同请求生成 token 数量和执行时间不同,会造成资源浪费。​

解码阶段的常见优化手段是 KV 缓存。在每个时间步生成单个 token 时,为避免重新计算所有 token 的键和值张量,可将其缓存在 GPU 内存中。每次迭代生成新 token 时,将其添加到缓存中供下次迭代使用,一些实现中模型的每一层都有 KV 缓存。​

大模型内存需求​

LLM 对 GPU 显存的需求主要源于模型权重和 KV 缓存。以具有 70 亿个参数的 Llama 2-7B 模型为例,以 16 位精度(FP16 或 BF16)加载,模型权重约占用 14GB 内存。而 KV 缓存方面,每个 Token KV 缓存大小计算公式为:2 * (num_layers) * (num_heads * dim_head) * precision_in_bytes ,总的 Token KV 缓存(字节)= (batch_size) * (sequence_length) * 2 * (num_layers) * (hidden_size) * sizeof (FP16) 。如 16 位精度的 Llama 2-7B 模型,批量大小为 1,序列长度为 4096 时,KV 缓存约为 2GB。内存需求随批量大小和序列长度线性增长,这对推理服务的吞吐量和长上下文输入处理构成挑战。​

模型并行技术​

当模型权重占用显存过大时,可将模型分布在多个 GPU 上,即模型并行化。它能分散内存和计算,使更大模型或更大批量输入得以运行。与数据并行不同,数据并行在推理中不太常用,其将模型权重复制到多个设备,输入全局批量大小在每个设备上分成微批次,主要用于训练时间优化。​

Pipeline 并行​

Pipeline 并行化将模型垂直分片为块,每个块包含在单独设备上执行的层子集。如四路 Pipeline,模型按顺序分区,每层的四分之一子集在每个设备上执行,一个设备操作输出传递到下一个设备继续执行后续块。此方法缺点是由于处理的顺序性,某些设备或层在等待前一层输出(激活、梯度)时可能空闲。​

提升推理能力的创新方法​

思维链(CoT)​

思维链是一种提示方法,通过在输入中展示从问题到答案的逐步推理过程,引导大语言模型生成中间推理步骤,进而得出最终答案。它将复杂任务分解为一系列简单、逻辑连贯的子任务,使模型模仿人类思维过程逐步推导。例如在数学问题解决中,不仅给出答案,还展示每步计算依据和逻辑。​

CoT 主要通过提示工程实现,Few-shot CoT 在提示中提供少量包含推理步骤的示例,模型学习示例推理模式解决新问题;Zero-shot CoT 通过特定引导语句,如 “Let's think step by step”,促使模型自发生成推理步骤。在数学推理任务中,使用 CoT 提示的大语言模型准确率大幅提升,如 PaLM-540B 模型在 GSM8K 数学问题数据集上,应用 CoT 提示后准确率从 17% 提升至 58%。同时,CoT 增强了模型输出的可解释性,便于用户理解和优化模型。​

思维树(ToT)​

思维树是在思维链基础上的拓展,将问题解决过程视为在一棵思维树上的搜索。每个节点代表问题解决过程中的一个状态,包含输入信息和到该节点为止的思维序列(部分解决方案)。ToT 允许模型探索多种推理路径,处理问题时需回答四个关键问题:如何分解问题为思维步骤(思维分解);如何基于当前树状态生成下一个思维步骤的多个候选(思维生成器);如何评估不同状态在解决问题上的进展(状态求值器);选择何种搜索算法在思维树中进行搜索(搜索算法)。​

在实现机制上,思维分解需根据问题性质和复杂程度划分合理粒度的思维步骤;思维生成器基于当前树状态,通过从思维链提示中采样思维或直接提出不同想法生成候选;状态求值器通过对每个状态独立赋值或跨状态投票等方式评估状态优劣;常用搜索算法包括广度优先搜索(BFS)和深度优先搜索(DFS),可根据问题特点选择。在 24 点游戏、创意写作、迷你填字游戏等复杂多步骤推理任务中,ToT 表现出独特优势,相比 CoT 能更全面地探索推理路径,解决复杂问题能力更强。​

推理优化技术​

内存管理优化​

  1. 页面注意力机制(Page attention):KV 缓存占用大量内存,且易产生内存碎片。Page Attention 机制借鉴操作系统内存管理思路,将数据按逻辑块组织,用 “页表” 监控,映射到物理内存。其特点为固定大小块,分配定长小内存块给 KV 缓存;块可共享,在多个请求间复用;按需分配,推理时边生成边分配内存,无需按最大序列长度预留。​
  1. Raddix tree KV 缓存:Raddix 树是一种空间优化的 trie 结构。Raddix tree KV 缓存技术将 KV 缓存组织成 radix 树结构,适合多个请求有相同前缀的场景,能高效回调 KV 数据并在请求间共享,操作复杂度为 O (nlogn),相比注意力机制的 O (n²) 有显著优势。​

注意力机制优化​

  1. Group Query Attention(GQA)和 Multi Query Attention(MQA):多头注意力是 Transformer 模型核心,但每个头需自己的 KV 对,导致处理实时或长文本时 KV 对庞大,占用大量内存。GQA 允许多个 query 共享 Key 和 Value 对,减少计算 attention 时 KV 对数量;MQA 则让多个 query 共享一组 KV,更为节省资源。​
  1. Flash Multi Latent Attention(Flash MLA):DeepSeek 推出的 Flash MLA 技术,不保存完整的 K 和 V,而是用低秩压缩将它们投影成维度更小的 latent 向量,降低缓存体积,在计算 attention 时再进行 “上投影”,且 “上投影” 权重矩阵会和 query 里的矩阵权重融合,加快 attention 计算速度。​

其他优化方法​

  1. 量化:量化是重要的优化手段,使用 8 比特或 4 比特整型替换 16 比特或 32 比特浮点类型,可降低内存使用和带宽需求,尤其在解码阶段效果显著。量化模型能够缓存更多 tokens,处理更多并发请求,提升计算速度。​
  1. 算子融合:将 LayerNorm、矩阵乘、激活函数等融合成一个算子,减少访存和 Kernel Launch 的开销,降低每个 token 生成的成本,提升解码阶段性能。​
  1. 连续批处理和混合批处理:连续批处理新请求加入时不中断正在处理的批样本,减少推理时延,最大化 GPU 和显存利用率,但要求复杂调度机制和高效 KV Cache 管理;混合批处理结合多种批处理方式,进一步优化解码阶段性能,更好地利用 GPU 资源。​

大模型推理技术涵盖了从基础概念到创新方法,再到多种优化手段的丰富内容。随着技术不断发展,这些技术将持续演进,推动大模型在更广泛领域高效应用。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐