尧图网站建设 
Transformers 大模型推理优化实战低精度量化、Flash Attention 与架构创新【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers本文基于 Transformers 仓库中的官方优化指南docs/source/ar/llm_tutorial_optimization.md编写系统讲解如何从低精度量化、Flash Attention、架构创新相对位置编码 KV Cache 优化三个维度降低大语言模型LLM的推理内存占用并提升推理效率。读完后你将掌握 VRAM 需求估算方法、8-bit/4-bit 量化的完整实操代码、KV Cache 与多轮对话加速的实现原理以及 MQA/GQA 等模型架构选择的决策依据并能对照仓库源码如src/transformers/utils/quantization_config.py、src/transformers/cache_utils.py理解每个优化手段的落地位置。1. 为什么 LLM 部署难内存与长上下文的挑战GPT3/4、Falcon、Llama 这类大语言模型在文本理解与生成上接近人类水平但要在真实任务中部署它们存在两大障碍参数量巨大要表现出接近人类的文本理解与生成能力LLM 通常需要数十亿乃至数千亿参数见 Kaplan et al. 的 scaling 研究这直接推高了推理时的内存需求长上下文许多真实任务需要给模型提供大量背景信息要求模型在推理时能处理超长输入序列。这两点共同指向一个核心问题如何增强 LLM 的计算与内存能力尤其是在处理海量输入序列时。本指南围绕三类已被验证有效的技术展开低精度Lower Precision研究表明以 8-bit / 4-bit 等更低的数值精度运行可以在不明显损害模型效果的前提下获得计算与内存收益详见仓库文档docs/source/en/main_classes/quantization.mdFlash Attention对注意力算法的改造版本不仅内存更高效还通过优化 GPU 显存VRAM利用方式带来速度提升架构创新Architectural Innovations由于 LLM 在推理时总是以长输入上下文 自回归生成的方式运行研究者提出了专门的模型结构来提升推理效率代表性进展包括 ALiBi、旋转位置编码RoPE、多查询注意力MQA与分组查询注意力GQA。整个指南从张量视角分析自回归生成过程通过一系列可运行的实践示例逐一验证每种优化带来的实际收益。2. 低精度从 VRAM 估算到 8-bit / 4-bit 量化2.1 先算账加载一个 LLM 需要多少显存把 LLM 看作一组权重矩阵和向量把文本输入看作一串向量序列下文中权重指模型中所有权重矩阵与向量。每个参数都是一个如4.5689的十进制数通常以 float32、bfloat16 或 float16 格式存储由此可以估算把模型加载进内存所需的显存加载一个 X 十亿参数的模型权重在 float32 精度下大约需要 4 * X GB 显存。但如今模型很少以完整 float32 训练更多是 bfloat16较少为 float16因此经验法则变为加载一个 X 十亿参数的模型权重在 bfloat16/float16 精度下大约需要 2 * X GB 显存。对于较短的文本输入少于 1024 个 token推理的内存需求几乎完全由加载权重主导因此可以近似认为推理内存需求 ≈ 把模型装入 GPU 显存的内存需求。按 bfloat16 估算一些知名模型的显存需求模型参数量bfloat16 显存估算GPT3175B2 × 175 350 GBBLOOM176B2 × 176 352 GBLlama-2-70b70B2 × 70 140 GBFalcon-40b40B2 × 40 80 GBMPT-30b30B2 × 30 60 GBbigcode/starcoder15.5B2 × 15.5 31 GB写这篇文档时市场上最大的 GPU 芯片是提供 80 GB 显存的 A100 与 H100。上表多数模型仅加载就需要超过 80 GB因此必然需要张量并行tensor parallelism和/或流水线并行pipeline parallelism。指南原文指出当时 Transformers 不原生支持张量并行因为它要求以特定方式编写模型结构可参考 text-generation-inference 库中的自定义建模写法而朴素流水线并行开箱即用。当前仓库的英文版同篇文档已更新说明对于配置类中定义了base_tp_plan的受支持模型Transformers 现已支持张量并行——从这一演进可以看出仓库在此方向上的持续投入。朴素流水线并行的用法非常简单加载模型时传device_mapauto框架会自动把不同层放到可用的多块 GPU 上。注意朴素流水线并行虽然有效但并不解决 GPU 空闲idling问题需要更高级的流水线并行才能进一步利用硬件。如果你有一台 8 × 80GB A100 的节点可以这样加载 BLOOMpip install transformers accelerate bitsandbytes optimumfrom transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(bigscience/bloom, device_mapauto, pad_token_id0)使用device_mapauto后注意力层会被均匀分布到所有可用 GPU 上。2.2 实测基线bigcode/octocoder 在 bfloat16 下需要约 29 GB后续所有实验使用bigcode/octocoder约 15.5B 参数因为它能跑在单块 40 GB 的 A100 上。需要注意本节所有内存与速度优化手段同样适用于需要模型/张量并行的更大模型。模型以 bfloat16 加载按上面的经验法则预期推理显存需求约 31 GB。先加载模型与 tokenizer再传给 Transformers 的pipeline对象from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline import torch model AutoModelForCausalLM.from_pretrained(bigcode/octocoder, dtypetorch.bfloat16, device_mapauto, pad_token_id0) tokenizer AutoTokenizer.from_pretrained(bigcode/octocoder) pipe pipeline(text-generation, modelmodel, tokenizertokenizer)prompt Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\npython\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single生成的代码可以直接用来做字节到 GiB 的换算def bytes_to_giga_bytes(bytes): return bytes / 1024 / 1024 / 1024调用torch.cuda.memory.max_memory_allocated()测量 GPU 峰值显存分配bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出29.0260648727417与粗略计算非常接近。数字不是完全精确是因为字节到千字节换算应乘 1024 而非 1000所以该经验公式也可以理解为至多 X GB的估算。如果改用完整 float32 运行需要约 64 GB 显存。如今几乎所有模型都以 bfloat16 训练。如果你的 GPU 支持 bfloat16就没有理由用完整 float32 运行模型——float32 不会给出比训练精度更好的推理结果。如果不確定模型权重在 Hub 上以什么格式存储可以查看检查点的 config 文件中的dtype字段。建议用from_pretrained(..., dtype...)加载时设置与 config 中一致的精度类型除非原始类型是 float32此时推理可以用float16或bfloat16。为了准确测量后续实验的峰值显存定义一个flush(...)函数释放全部已分配内存del pipe del model import gc import torch def flush(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats()flush()Accelerate 库还提供了设备无关的辅助方法release_memory()可一并释放模型占用的内存from accelerate.utils import release_memory # ... release_memory(model)2.3 量化原理为什么 8-bit / 4-bit 能跑且不伤精度如果 GPU 显存不到 32 GB 怎么办研究发现模型权重可以被量化到 8-bit 甚至 4-bit 而几乎没有性能损失Dettmers et al. 的工作GPTQ 论文进一步表明量化到 3-bit / 2-bit 也能在可接受的损失下运行。不深入细节量化方案的共同目标是降低权重的存储精度同时尽量保持模型推理结果准确即尽可能接近 bfloat16 的推理结果。量化对文本生成尤其有效因为生成任务只关心选择最可能的下一批 token并不真正关心下一 tokenlogit分布的精确数值——只要logit分布大体不变argmax或topk的结果就不变。尽管具体技术各有不同所有量化方法的工作流程都可以概括为三步把所有权重量化到目标精度加载量化后的权重输入序列仍以 bfloat16 格式传入在计算时把权重动态反量化回 bfloat16再与 bfloat16 输入向量做运算。也就是说每个矩阵乘法$$ Y X * W $$其中 $X$ 是输入、$W$ 是权重矩阵、$Y$ 是输出都变成了$$ Y X * \text{dequantize}(W) $$反量化与重量化在所有权重矩阵上按输入穿过网络图的顺序依次执行。因此使用量化权重通常不会降低推理时间反而会使其变长——这一点在实测中会被反复验证。2.4 8-bit 量化实操显存从 29 GB 降到 15 GB用 Transformers 做量化需要安装bitsandbytes库pip install bitsandbytes8-bit 量化只需给from_pretrained加上量化配置。指南原文使用旧式写法load_in_8bitTrue在当前仓库中推荐也是当前代码路径实际消费的写法是传入BitsAndBytesConfig。该类的定义位于 src/transformers/utils/quantization_config.py从源码可以看到load_in_8bit启用 LLM.int8() 的 8-bit 量化默认Falseload_in_4bit用bitsandbytes的 FP4/NF4 层替换 Linear 层实现 4-bit 量化默认False两者互斥同时为True会抛出ValueErrorllm_int8_threshold默认 6.0LLM.int8() 论文中异常值检测的阈值超过该阈值的隐藏状态值会以 fp16 计算bnb_4bit_compute_dtype4-bit 反量化后参与计算的精度如可设torch.bfloat16提速bnb_4bit_quant_type4-bit 量化数据类型可选fp4或nf4bnb_4bit_use_double_quant双重量化对第一次量化的常数再做一次量化进一步省内存。model AutoModelForCausalLM.from_pretrained(bigcode/octocoder, quantization_configBitsAndBytesConfig(load_in_8bitTrue), pad_token_id0)重新运行示例并测量内存pipe pipeline(text-generation, modelmodel, tokenizertokenizer) result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\npython\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single与之前完全相同精度无损。再看内存bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出15.219234466552734显存从约 29 GB 降到 15 GB 出头这个模型因此可以跑在 4090 这类消费级 GPU 上。内存效率收益很大、输出几乎无退化但可以观察到推理有轻微变慢——这正是 2.3 节所说的反量化开销。删除模型并清空内存进入下一组实验del model del pipeflush()2.5 4-bit 量化实操15.5B 模型只占约 9.5 GB4-bit 量化使用同一套 API把load_in_8bitTrue换成load_in_4bitTruemodel AutoModelForCausalLM.from_pretrained(bigcode/octocoder, quantization_configBitsAndBytesConfig(load_in_4bitTrue), pad_token_id0) pipe pipeline(text-generation, modelmodel, tokenizertokenizer) result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\n\ndef bytes_to_gigabytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single argument输出与之前几乎一致——只是代码块前少了一个python标注。内存需求bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出9.543574333190918只有 9.5 GB对 15B 参数模型来说非常少。虽然本例中精度退化极小但实践中 4-bit 量化相比 8-bit 或完整 bfloat16 推理更容易产生不同结果需要你自行在目标任务上验证。同时注意这里的推理又比 8-bit 略慢原因是 4-bit 使用了更激进的量化方式quantize与dequantize在推理时耗时更长。del model del pipeflush()小结OctoCoder 以 8-bit 运行时所需显存从约 32 GB 降到约 15 GB以 4-bit 运行时进一步降到 9 GB 出头。4-bit 量化让模型可以跑在 RTX3090、V100、T4 等对大多数人更易获得的 GPU 上。如果想把显存压得比 4-bit 更低可以研究仓库量化文档docs/source/en/main_classes/quantization.md中介绍的 GPTQModel 等实现。结论性提醒模型量化是用有时还有推理时间换内存效率的交易。如果显存不是你的约束通常不必考虑量化但很多 GPU 不量化根本无法运行 LLM此时 4-bit / 8-bit 量化方案是极其有用的工具。3. Flash Attention把 $N^2$ 的显存压力降到线性当今高性能 LLM 大多共享同一套基础结构前馈层、激活层、层归一化层以及最关键的自注意力层。自注意力层让模型理解输入 token 间的上下文关系但其 GPU 峰值显存消耗在计算与内存两个维度上都随输入 token 数序列长度记为 $N$二次方增长。短序列1000 token 以内不明显长序列约 16000 token就成为严重问题。对长度为 $N$ 的输入 $\mathbf{X}$自注意力层输出 $\mathbf{O}$ 的计算公式为$$ \textbf{O} \text{Attn}(\mathbf{X}) \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} \mathbf{W}_q \mathbf{X}, \mathbf{V} \mathbf{W}_v \mathbf{X}, \mathbf{K} \mathbf{W}_k \mathbf{X} $$$\mathbf{X} (\mathbf{x}1, ... \mathbf{x}{N})$ 是输入序列。投影 $\mathbf{Q}$ 和 $\mathbf{K}$ 各含 $N$ 个向量于是 $\mathbf{QK}^T$ 的规模是 $N^2$。LLM 通常有多个注意力头并行做自注意力计算。假设 LLM 有 40 个注意力头、以 bfloat16 运行存储 $\mathbf{QK^T}$ 矩阵的显存需求是 $40 \times 2 \times N^2$ 字节$N 1000$只需约 50 MB 显存$N 16000$需要 19 GB 显存$N 100,000$仅存储 $\mathbf{QK}^T$ 矩阵就接近 1 TB。一句话默认自注意力算法在大输入上下文下迅速变得内存昂贵。随着 LLM 能力增强它们从处理几句话的翻译/摘要走向处理整个页面的任务长输入成为刚需。出路在于一种不需要显式构造 $\mathbf{QK}^T$ 矩阵的新注意力算法——Tri Dao 等人提出的Flash Attention。其核心思想是把 $\mathbf{V} \times \operatorname{Softmax}(\mathbf{QK}^T)$ 的计算拆开通过对多组 Softmax 分块步骤迭代逐块计算输出的小片段$$ \textbf{O}i \leftarrow s^a{ij} * \textbf{O}i s^b{ij} * \mathbf{V}{j} \times \operatorname{Softmax}(\mathbf{QK}^T{i,j}) \quad \text{for multiple } i, j \text{ iterations} $$其中 $s^a_{ij}$ 与 $s^b_{ij}$ 是需要在每一组 $i, j$ 上重新计算的 Softmax 归一化统计量。Flash Attention 的完整实现更复杂这里做了大幅简化深入细节请阅读 Flash Attention 原论文。关键要点通过追踪 Softmax 归一化统计量并配合巧妙的数学处理Flash Attention 在显存成本只随 $N$ 线性增长的前提下给出与默认自注意力层数值完全相同的输出。从公式直觉上会以为 Flash Attention 更慢——它确实需要更多 FLOPs因为 Softmax 归一化统计量必须不断重算。但推理中它远快于默认注意力Flash Attention 大幅降低了对 GPU 上更慢的高带宽显存VRAM的访问需求转而利用芯片上更快的 SRAM。本质上Flash Attention 保证所有中间读写都发生在片上高速 SRAM而不是访问慢速 VRAM 来计算输出向量 $\mathbf{O}$。实践建议很明确只要可用就没有任何理由不用 Flash Attention——它数学上输出相同且更快、更省内存。4. 架构创新为长文本任务优化的模型结构前面我们从两个方向改进了计算与内存效率把权重转成低精度格式用更高效版本替换自注意力算法。本节讨论第三个维度——让 LLM 的结构本身适配长文本任务例如检索增强问答、摘要、聊天。注意聊天不仅要求处理长文本输入还要求高效处理用户与助手之间的多轮对话如 ChatGPT 场景。模型一旦训练完成基础结构就很难改变因此在设计阶段就要预判任务并相应优化架构。结构中有两个组件对长输入序列最先成为内存和/或性能瓶颈位置编码positional embeddings键值缓存key-value cache4.1 位置编码从绝对编码到 RoPE / ALiBi自注意力把每个 token 与其他 token 建立关联。对输入序列Hello, I, love, you$\operatorname{Softmax}(\mathbf{QK}^T)$ 矩阵可以形如每个词 token 获得一组概率质量用以关注所有其他词 token从而与全部其他 token 建立关系。例如单词love以 5% 关注Hello、以 30% 关注I、以 65% 关注自身。一个基于自注意力但没有位置编码的 LLM会很难理解各输入文本的相对位置因为 $\mathbf{QK}^T$ 计算出的概率得分以 $O(1)$ 计算把每个词 token 关联到每个其他词 token与相对位置距离无关。对没有位置编码的 LLM 来说每个 token 到所有其他 token 的距离看起来都一样区分Hello I love you与You love I hello将非常困难。为了让 LLM 理解句子顺序需要额外线索通常以位置编码positional encodings也称positional embeddings的形式注入位置编码把每个 token 的位置编码为数值表示让 LLM 据此更好理解句子顺序。《Attention Is All You Need》的 authors 引入了正弦位置嵌入$\mathbf{P} \mathbf{p}_1, \ldots, \mathbf{p}_N$其中每个向量 $\mathbf{p}_i$ 是其位置 $i$ 的正弦函数。位置编码直接加到输入序列向量上$$ \mathbf{\hat{X}} \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N \mathbf{x}_1 \mathbf{p}_1, \ldots, \mathbf{x}_N \mathbf{p}_N $$从而提示模型更好地学习句子顺序。另一种路线如 BERT 的 Devlin et al.使用学习式位置嵌入即在训练中学习 $\mathbf{P}$。正弦与学习式位置嵌入曾是编码句子顺序的主流方法但发现了两个问题二者都是绝对位置嵌入即为每个位置 id$0, \ldots, N$编码唯一嵌入。研究表明绝对位置嵌入导致 LLM 在长文本输入上表现差——对长文本让模型学习 token 之间的相对位置距离比绝对位置更有优势使用学习式位置嵌入时LLM 必须在固定输入长度 $N$ 上训练难以外推到比训练长度更长的输入。近期能解决上述问题的相对位置嵌入更受青睐其中最突出的是旋转位置嵌入RoPE, Rotary Position EmbeddingALiBiRoPE 与 ALiBi 的共同观点是句子顺序信息最好直接注入自注意力算法——因为 token 正是在那里建立相互关系的具体而言应通过修改 $\mathbf{QK}^T$ 的计算来注入顺序信息。RoPE指出位置信息可以编码进 query-key 向量对如 $\mathbf{q}_i$ 与 $\mathbf{k}_j$中把每个向量分别旋转角度 $\theta \cdot i$ 与 $\theta \cdot j$$i, j$ 是各向量的句子位置$$ \mathbf{\hat{q}}_i^T \mathbf{\hat{k}}_j \mathbf{q}i^T \mathbf{R}{\theta, i - j} \mathbf{k}_j $$其中 $\mathbf{R}_{\theta, i - j}$ 是旋转矩阵。$\theta$不在训练中学习而是设为一个依赖于训练时最大输入序列长度的预定义值。这样处理后$\mathbf{q}_i$ 与 $\mathbf{k}_j$ 之间的概率得分仅在 $i \ne j$ 时受影响且只取决于相对距离 $i - j$与两个向量的具体位置 $i$、$j$ 无关。RoPE 被当今许多重要 LLM 采用如 Falcon、Llama、PaLM。仓库中对应实现集中在 src/transformers/modeling_rope_utils.pyRoPE 频率计算与配置解析各模型如 Qwen、Llama 系列的建模代码通过它构造旋转缓存。ALiBi则提出更简单的相对位置编码方案把输入 token 之间的相对距离作为负整数 × 预定义缩放值m在计算 softmax 前直接加到 $\mathbf{QK}^T$ 矩阵的每个 query-key 条目上。ALiBi 论文表明这种简单的相对位置编码让模型在超长文本输入序列上仍保持高性能。ALiBi 同样被 MPT、BLOOM 等重要模型采用。外推能力对比RoPE 与 ALiBi 都能外推到训练中未见过的输入长度但开箱即用的外推效果 ALiBi 明显优于 RoPE对 ALiBi只需把下三角位置矩阵的数值加大以匹配输入序列长度即可对 RoPE沿用训练时的 $\theta$ 会在远长于训练长度的输入上得到糟糕结果不过社区发现了一些调整 $\theta$ 的有效技巧如缩放 RoPE 基频让 RoPE 位置编码也能在extrapolated 输入上正常工作。RoPE 与 ALiBi 都是不在训练中学习的相对位置编码基于如下直觉关于文本输入的位置线索应直接给到自注意力层的 $\mathbf{QK}^T$ 矩阵应促使 LLM 学习 token 之间恒定的相对距离位置编码token 彼此距离越远query-key 概率应越低。RoPE 通过增大 query-key 向量夹角来降低向量点积ALiBi 通过给点积加大负数来降低。结论面向需要处理大文本输入任务的 LLM最好用 RoPE 或 ALiBi 这类相对位置编码训练。注意即使模型只在固定长度如 $N_1 2048$上训练借助位置编码外推实际可以用远大于 $N_1$ 的输入如 $N_2 8192$。4.2 KV Cache自回归生成的加速器及其代价LLM 的自回归文本生成方式是反复输入序列 → 采样下一个 token → 把该 token 追加到输入序列 → 直到模型生成结束标记 token。先用一段简码展示自回归生成的朴素形态用torch.argmax取最可能的下一 tokeninput_ids tokenizer(prompt, return_tensorspt)[input_ids].to(cuda) for _ in range(5): next_logits model(input_ids)[logits][:, -1:] next_token_id torch.argmax(next_logits, dim-1) input_ids torch.cat([input_ids, next_token_id], dim-1) print(shape of input_ids, input_ids.shape) generated_text tokenizer.batch_decode(input_ids[:, -5:]) generated_text输出shape of input_ids torch.Size([1, 21]) shape of input_ids torch.Size([1, 22]) shape of input_ids torch.Size([1, 23]) shape of input_ids torch.Size([1, 24]) shape of input_ids torch.Size([1, 25]) [ Here is a Python function]可以看到每一步都把刚采样的 token 追加进输入序列。除极少数例外LLM 都用因果语言建模目标训练因此注意力结果矩阵的上三角被掩码——token从不依赖后续 token$\mathbf{q}_i$ 永远不会与 $j i$ 的 $\mathbf{k}j, \mathbf{v}j$ 建立关系只关注此前的 $\mathbf{k}{m i}, \mathbf{v}{m i}$。为了减少不必要的计算可以缓存每一层在所有前序时间步的 key-value 向量。在 Transformers 中向forward传入use_cache标志即可取回 KV cache并在下一步连同当前 token 一起传回past_key_values None # past_key_values 就是 KV cache generated_tokens [] next_token_id tokenizer(prompt, return_tensorspt)[input_ids].to(cuda) for _ in range(5): next_logits, past_key_values model(next_token_id, past_key_valuespast_key_values, use_cacheTrue).to_tuple() next_logits next_logits[:, -1:] next_token_id torch.argmax(next_logits, dim-1) print(shape of input_ids, next_token_id.shape) print(length of key-value cache, past_key_values.get_seq_length()) # past_key_values 形状为 [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim] generated_tokens.append(next_token_id.item()) generated_text tokenizer.batch_decode(generated_tokens) generated_text输出shape of input_ids torch.Size([1, 1]) length of key-value cache 20 shape of input_ids torch.Size([1, 1]) length of key-value cache 21 shape of input_ids torch.Size([1, 1]) length of key-value cache 22 shape of input_ids torch.Size([1, 1]) length of key-value cache 23 shape of input_ids torch.Size([1, 1]) length of key-value cache 24 [ Here, is, a, Python, function]可以看到使用 KV cache 时输入 token 序列不再变长始终是一个输入向量而 KV cache 的长度在每步解码时增加 1。使用 KV cache 意味着 $\mathbf{QK}^T$ 本质上退化成了 $\mathbf{q}_c\mathbf{K}^T$其中 $\mathbf{q}_c$ 是当前传入 token 的 query 投影永远只是单个向量。KV cache 带来两大优势计算效率显著提升相比计算完整 $\mathbf{QK}^T$ 矩阵计算量更少推理速度更快峰值内存不再随生成 token 数二次增长而是线性增长。应始终使用 KV cache它结果相同且对长输入序列有显著加速。Transformers 在使用 text pipeline 或generate方法 时默认启用 KV cache。仓库中 src/transformers/cache_utils.py 里的DynamicCache是生成模型的默认缓存实现逐层存储 key/value 张量形状为[batch_size, num_heads, seq_len, head_dim]若传入 config还会根据滑动窗口/混合层结构把缓存张量压缩到[batch_size, num_heads, min(seq_len, sliding_window), head_dim]进一步降低内存需求。更完整的缓存体系说明见仓库文档 KV Cache 指南。⚠️ 注意尽管建议使用 KV cache启用它之后 LLM 输出可能与未启用时略有不同。这是矩阵乘法内核kernel本身的数值特性所致属于已知行为。4.2.1 多轮对话KV cache 复用历史上下文KV cache 对聊天这类需要多次自回归解码的应用尤其有用。看一个例子User: How many people live in France? Assistant: Roughly 75 million people live in France User: And how many are in Germany? Assistant: Germany has ca. 81 million inhabitants在这段对话中LLM 运行了两次自回归解码第一轮KV cache 为空输入 prompt 是User: How many people live in France?模型自回归生成Roughly 75 million people live in France每个解码步都在增长 KV cache第二轮输入 prompt 本应是User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?。得益于 cache前两句话的 key-value 向量已经计算过所以实际只需处理User: And how many in Germany?。处理这段缩短的 prompt 时其计算出的 key-value 向量直接拼接进第一轮解码留下的 cache随后基于完整三句话的 cache 自回归生成第二个回答Germany has ca. 81 million inhabitants。有两点值得注意保留完整上下文对聊天 LLM 至关重要否则模型无法理解对话前文。例如上例中模型需要理解用户问 And how many are in Germany 时指的是人口KV cache 对聊天极有价值它允许持续生长已编码的聊天历史而不必从头重新编码若用 encoder-decoder 结构就必须从头重编。在transformers中当传入return_dict_in_generateTrue叠加默认的use_cacheTrue时generate调用会返回past_key_values。注意该能力尚不通过pipeline接口暴露# Generation as usual prompt system_prompt Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here model_inputs tokenizer(prompt, return_tensorspt) generation_output model.generate(**model_inputs, max_new_tokens60, return_dict_in_generateTrue) decoded_output tokenizer.batch_decode(generation_output.sequences)[0] # Piping the returned past_key_values to speed up the next conversation round prompt decoded_output \nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here model_inputs tokenizer(prompt, return_tensorspt) generation_output model.generate( **model_inputs, past_key_valuesgeneration_output.past_key_values, max_new_tokens60, return_dict_in_generateTrue ) tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]输出is a modified version of the function that returns Mega bytes instead. def bytes_to_megabytes(bytes): return bytes / 1024 / 1024 Answer: The function takes a number of bytes as input and returns the number of很好不再为注意力层重复计算相同的 key 与 value。但要注意一个代价$\mathbf{QK}^T$ 矩阵的峰值内存需求被大幅降低了把 KV cache 留在内存里对长输入序列或多轮聊天可能非常吃内存——cache 需要为所有自注意力层、所有注意力头保存所有前序输入向量 $\mathbf{x}_i$$i \in {1, \ldots, c-1}$的 key-value 向量。以bigcode/octocoder为例计算 KV cache 需要存储的浮点数个数2k 和 v× 序列长度 × 注意力头数 × 每头维度 × 层数。在假想的 16000 输入序列长度下config model.config 2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head输出7864320000约80 亿个浮点数以float16存储这 80 亿个浮点数需要约 15 GB 内存——接近模型权重本身的一半研究者提出了两种显著降低 KV cache 内存成本的方法即下文两小节。4.2.2 多查询注意力MQANoam Shazeer 在Fast Transformer Decoding: One Write-Head is All You Need中提出MQA。如标题所示与其使用n_head组 key-value 投影权重只需用一组跨所有注意力头共享的 value 投影权重模型性能并不会显著退化。使用单组 value 投影权重后key-value 向量 $\mathbf{k}_i, \mathbf{v}_i$ 在所有注意力头中必然相同因此 cache 中只需存 1 组 key-value 投影而不是n_head组。由于大多数 LLM 使用 20 到 100 个注意力头MQA 能大幅降低 KV cache 的内存消耗。对本指南使用的 LLM可以把 16000 序列长度下的 cache 内存需求从 15 GB 降到不到 400 MB。除了省内存MQA 还提升计算效率。自回归解码中每步都需要把巨大的 key-value 向量重新读回、与当前 key-value 向量拼接再送入 $\mathbf{q}_c\mathbf{K}^T$ 计算——持续重载带来的内存带宽需求可能成为严重的时间瓶颈。减小 key-value 向量体积即减少内存访问从而缓解带宽瓶颈。要理解的关键点把 key-value 注意力头数降到 1只有在配合 KV cache 使用时才有意义。不使用 KV cache 的单个前向传播峰值内存不变——每个注意力头仍有唯一 query 向量每个头仍有不同的 $\mathbf{QK}^T$ 矩阵。MQA 已被社区广泛采用众多热门 LLM 都在使用包括 Falcon、PaLM、MPT、BLOOM本指南使用的bigcode/octocoder检查点同样采用 MQA。4.2.3 分组查询注意力GQA谷歌的 Ainslie et al. 提出的GQA发现相比使用常规的多个 key-value 头投影MQA 有时会带来质量退化。该工作认为可以通过更温和地减少 query 头投影权重数量来保住更多模型性能不用 1 组 key-value 投影权重而是用n n_head组。把n选为远小于n_head的值如 2、4、8既能保住 MQA 绝大部分内存与速度收益又只牺牲更少的模型容量因此性能损失更小。此外GQA 作者发现已存在的模型检查点可以用不到原始预训练 5% 的计算量uptrain成 GQA 结构。虽然 5% 的预训练计算量本身可能仍是天文数字但 GQA uptraining 让既有检查点可以服务于更长的输入序列。GQA 提出较晚撰写本文档时采用率相对更低其最著名的应用是 Llama-v2。总结性建议如果 LLM 以自回归解码方式部署、且需要处理大输入序列如聊天场景强烈建议使用 GQA 或 MQA。5. 结语把最好的算法与架构用足研究社区不断为越来越大的 LLM 提出加速推理的新方法。例如很有前景的**推测解码speculative decoding**方向由更小更快的语言模型生成简单 token只有困难 token才交给大 LLM 生成。GPT3/4、Llama-2-70b、PaLM 这类巨型 LLM 之所以能在 Hugging Face Chat、ChatGPT 等聊天界面中高速运行很大程度上正得益于本文讨论的**精度量化、算法Flash Attention、架构相对位置编码 KV Cache MQA/GQA**三个层面的优化。未来 GPU、TPU 等加速器只会更快、内存更大但无论如何都应确保用上了当前可用的最佳算法与架构把每一份算力花在刀刃上。快速决策参考场景首选手段预期收益显存装不下模型8-bit / 4-bit 量化BitsAndBytesConfig本例 29 GB → 15.2 GB8-bit→ 9.5 GB4-bit多卡加载超大模型device_mapauto流水线并行 / 张量并行突破单卡 80 GB 上限长上下文≥16k tokenFlash Attention注意力显存从 $O(N^2)$ 降到 $O(N)$长文本任务摘要/检索/聊天RoPE 或 ALiBi 相对位置编码支持外推到训练长度之外的输入自回归解码 多轮聊天始终启用 KV cache考虑 MQA/GQA 架构解码步数计算量下降16k 序列 cache 从约 15 GB 降到 400 MBMQA【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Transformers 大模型推理优化实战：低精度量化、Flash Attention 与架构创新
Transformers 大模型推理优化实战低精度量化、Flash Attention 与架构创新【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers本文基于 Transformers 仓库中的官方优化指南docs/source/ar/llm_tutorial_optimization.md编写系统讲解如何从低精度量化、Flash Attention、架构创新相对位置编码 KV Cache 优化三个维度降低大语言模型LLM的推理内存占用并提升推理效率。读完后你将掌握 VRAM 需求估算方法、8-bit/4-bit 量化的完整实操代码、KV Cache 与多轮对话加速的实现原理以及 MQA/GQA 等模型架构选择的决策依据并能对照仓库源码如src/transformers/utils/quantization_config.py、src/transformers/cache_utils.py理解每个优化手段的落地位置。1. 为什么 LLM 部署难内存与长上下文的挑战GPT3/4、Falcon、Llama 这类大语言模型在文本理解与生成上接近人类水平但要在真实任务中部署它们存在两大障碍参数量巨大要表现出接近人类的文本理解与生成能力LLM 通常需要数十亿乃至数千亿参数见 Kaplan et al. 的 scaling 研究这直接推高了推理时的内存需求长上下文许多真实任务需要给模型提供大量背景信息要求模型在推理时能处理超长输入序列。这两点共同指向一个核心问题如何增强 LLM 的计算与内存能力尤其是在处理海量输入序列时。本指南围绕三类已被验证有效的技术展开低精度Lower Precision研究表明以 8-bit / 4-bit 等更低的数值精度运行可以在不明显损害模型效果的前提下获得计算与内存收益详见仓库文档docs/source/en/main_classes/quantization.mdFlash Attention对注意力算法的改造版本不仅内存更高效还通过优化 GPU 显存VRAM利用方式带来速度提升架构创新Architectural Innovations由于 LLM 在推理时总是以长输入上下文 自回归生成的方式运行研究者提出了专门的模型结构来提升推理效率代表性进展包括 ALiBi、旋转位置编码RoPE、多查询注意力MQA与分组查询注意力GQA。整个指南从张量视角分析自回归生成过程通过一系列可运行的实践示例逐一验证每种优化带来的实际收益。2. 低精度从 VRAM 估算到 8-bit / 4-bit 量化2.1 先算账加载一个 LLM 需要多少显存把 LLM 看作一组权重矩阵和向量把文本输入看作一串向量序列下文中权重指模型中所有权重矩阵与向量。每个参数都是一个如4.5689的十进制数通常以 float32、bfloat16 或 float16 格式存储由此可以估算把模型加载进内存所需的显存加载一个 X 十亿参数的模型权重在 float32 精度下大约需要 4 * X GB 显存。但如今模型很少以完整 float32 训练更多是 bfloat16较少为 float16因此经验法则变为加载一个 X 十亿参数的模型权重在 bfloat16/float16 精度下大约需要 2 * X GB 显存。对于较短的文本输入少于 1024 个 token推理的内存需求几乎完全由加载权重主导因此可以近似认为推理内存需求 ≈ 把模型装入 GPU 显存的内存需求。按 bfloat16 估算一些知名模型的显存需求模型参数量bfloat16 显存估算GPT3175B2 × 175 350 GBBLOOM176B2 × 176 352 GBLlama-2-70b70B2 × 70 140 GBFalcon-40b40B2 × 40 80 GBMPT-30b30B2 × 30 60 GBbigcode/starcoder15.5B2 × 15.5 31 GB写这篇文档时市场上最大的 GPU 芯片是提供 80 GB 显存的 A100 与 H100。上表多数模型仅加载就需要超过 80 GB因此必然需要张量并行tensor parallelism和/或流水线并行pipeline parallelism。指南原文指出当时 Transformers 不原生支持张量并行因为它要求以特定方式编写模型结构可参考 text-generation-inference 库中的自定义建模写法而朴素流水线并行开箱即用。当前仓库的英文版同篇文档已更新说明对于配置类中定义了base_tp_plan的受支持模型Transformers 现已支持张量并行——从这一演进可以看出仓库在此方向上的持续投入。朴素流水线并行的用法非常简单加载模型时传device_mapauto框架会自动把不同层放到可用的多块 GPU 上。注意朴素流水线并行虽然有效但并不解决 GPU 空闲idling问题需要更高级的流水线并行才能进一步利用硬件。如果你有一台 8 × 80GB A100 的节点可以这样加载 BLOOMpip install transformers accelerate bitsandbytes optimumfrom transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(bigscience/bloom, device_mapauto, pad_token_id0)使用device_mapauto后注意力层会被均匀分布到所有可用 GPU 上。2.2 实测基线bigcode/octocoder 在 bfloat16 下需要约 29 GB后续所有实验使用bigcode/octocoder约 15.5B 参数因为它能跑在单块 40 GB 的 A100 上。需要注意本节所有内存与速度优化手段同样适用于需要模型/张量并行的更大模型。模型以 bfloat16 加载按上面的经验法则预期推理显存需求约 31 GB。先加载模型与 tokenizer再传给 Transformers 的pipeline对象from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline import torch model AutoModelForCausalLM.from_pretrained(bigcode/octocoder, dtypetorch.bfloat16, device_mapauto, pad_token_id0) tokenizer AutoTokenizer.from_pretrained(bigcode/octocoder) pipe pipeline(text-generation, modelmodel, tokenizertokenizer)prompt Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\npython\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single生成的代码可以直接用来做字节到 GiB 的换算def bytes_to_giga_bytes(bytes): return bytes / 1024 / 1024 / 1024调用torch.cuda.memory.max_memory_allocated()测量 GPU 峰值显存分配bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出29.0260648727417与粗略计算非常接近。数字不是完全精确是因为字节到千字节换算应乘 1024 而非 1000所以该经验公式也可以理解为至多 X GB的估算。如果改用完整 float32 运行需要约 64 GB 显存。如今几乎所有模型都以 bfloat16 训练。如果你的 GPU 支持 bfloat16就没有理由用完整 float32 运行模型——float32 不会给出比训练精度更好的推理结果。如果不確定模型权重在 Hub 上以什么格式存储可以查看检查点的 config 文件中的dtype字段。建议用from_pretrained(..., dtype...)加载时设置与 config 中一致的精度类型除非原始类型是 float32此时推理可以用float16或bfloat16。为了准确测量后续实验的峰值显存定义一个flush(...)函数释放全部已分配内存del pipe del model import gc import torch def flush(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats()flush()Accelerate 库还提供了设备无关的辅助方法release_memory()可一并释放模型占用的内存from accelerate.utils import release_memory # ... release_memory(model)2.3 量化原理为什么 8-bit / 4-bit 能跑且不伤精度如果 GPU 显存不到 32 GB 怎么办研究发现模型权重可以被量化到 8-bit 甚至 4-bit 而几乎没有性能损失Dettmers et al. 的工作GPTQ 论文进一步表明量化到 3-bit / 2-bit 也能在可接受的损失下运行。不深入细节量化方案的共同目标是降低权重的存储精度同时尽量保持模型推理结果准确即尽可能接近 bfloat16 的推理结果。量化对文本生成尤其有效因为生成任务只关心选择最可能的下一批 token并不真正关心下一 tokenlogit分布的精确数值——只要logit分布大体不变argmax或topk的结果就不变。尽管具体技术各有不同所有量化方法的工作流程都可以概括为三步把所有权重量化到目标精度加载量化后的权重输入序列仍以 bfloat16 格式传入在计算时把权重动态反量化回 bfloat16再与 bfloat16 输入向量做运算。也就是说每个矩阵乘法$$ Y X * W $$其中 $X$ 是输入、$W$ 是权重矩阵、$Y$ 是输出都变成了$$ Y X * \text{dequantize}(W) $$反量化与重量化在所有权重矩阵上按输入穿过网络图的顺序依次执行。因此使用量化权重通常不会降低推理时间反而会使其变长——这一点在实测中会被反复验证。2.4 8-bit 量化实操显存从 29 GB 降到 15 GB用 Transformers 做量化需要安装bitsandbytes库pip install bitsandbytes8-bit 量化只需给from_pretrained加上量化配置。指南原文使用旧式写法load_in_8bitTrue在当前仓库中推荐也是当前代码路径实际消费的写法是传入BitsAndBytesConfig。该类的定义位于 src/transformers/utils/quantization_config.py从源码可以看到load_in_8bit启用 LLM.int8() 的 8-bit 量化默认Falseload_in_4bit用bitsandbytes的 FP4/NF4 层替换 Linear 层实现 4-bit 量化默认False两者互斥同时为True会抛出ValueErrorllm_int8_threshold默认 6.0LLM.int8() 论文中异常值检测的阈值超过该阈值的隐藏状态值会以 fp16 计算bnb_4bit_compute_dtype4-bit 反量化后参与计算的精度如可设torch.bfloat16提速bnb_4bit_quant_type4-bit 量化数据类型可选fp4或nf4bnb_4bit_use_double_quant双重量化对第一次量化的常数再做一次量化进一步省内存。model AutoModelForCausalLM.from_pretrained(bigcode/octocoder, quantization_configBitsAndBytesConfig(load_in_8bitTrue), pad_token_id0)重新运行示例并测量内存pipe pipeline(text-generation, modelmodel, tokenizertokenizer) result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\npython\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single与之前完全相同精度无损。再看内存bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出15.219234466552734显存从约 29 GB 降到 15 GB 出头这个模型因此可以跑在 4090 这类消费级 GPU 上。内存效率收益很大、输出几乎无退化但可以观察到推理有轻微变慢——这正是 2.3 节所说的反量化开销。删除模型并清空内存进入下一组实验del model del pipeflush()2.5 4-bit 量化实操15.5B 模型只占约 9.5 GB4-bit 量化使用同一套 API把load_in_8bitTrue换成load_in_4bitTruemodel AutoModelForCausalLM.from_pretrained(bigcode/octocoder, quantization_configBitsAndBytesConfig(load_in_4bitTrue), pad_token_id0) pipe pipeline(text-generation, modelmodel, tokenizertokenizer) result pipe(prompt, max_new_tokens60)[0][generated_text][len(prompt):] result输出Here is a Python function that transforms bytes to Giga bytes:\n\n\ndef bytes_to_gigabytes(bytes):\n return bytes / 1024 / 1024 / 1024\n\n\nThis function takes a single argument输出与之前几乎一致——只是代码块前少了一个python标注。内存需求bytes_to_giga_bytes(torch.cuda.max_memory_allocated())输出9.543574333190918只有 9.5 GB对 15B 参数模型来说非常少。虽然本例中精度退化极小但实践中 4-bit 量化相比 8-bit 或完整 bfloat16 推理更容易产生不同结果需要你自行在目标任务上验证。同时注意这里的推理又比 8-bit 略慢原因是 4-bit 使用了更激进的量化方式quantize与dequantize在推理时耗时更长。del model del pipeflush()小结OctoCoder 以 8-bit 运行时所需显存从约 32 GB 降到约 15 GB以 4-bit 运行时进一步降到 9 GB 出头。4-bit 量化让模型可以跑在 RTX3090、V100、T4 等对大多数人更易获得的 GPU 上。如果想把显存压得比 4-bit 更低可以研究仓库量化文档docs/source/en/main_classes/quantization.md中介绍的 GPTQModel 等实现。结论性提醒模型量化是用有时还有推理时间换内存效率的交易。如果显存不是你的约束通常不必考虑量化但很多 GPU 不量化根本无法运行 LLM此时 4-bit / 8-bit 量化方案是极其有用的工具。3. Flash Attention把 $N^2$ 的显存压力降到线性当今高性能 LLM 大多共享同一套基础结构前馈层、激活层、层归一化层以及最关键的自注意力层。自注意力层让模型理解输入 token 间的上下文关系但其 GPU 峰值显存消耗在计算与内存两个维度上都随输入 token 数序列长度记为 $N$二次方增长。短序列1000 token 以内不明显长序列约 16000 token就成为严重问题。对长度为 $N$ 的输入 $\mathbf{X}$自注意力层输出 $\mathbf{O}$ 的计算公式为$$ \textbf{O} \text{Attn}(\mathbf{X}) \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} \mathbf{W}_q \mathbf{X}, \mathbf{V} \mathbf{W}_v \mathbf{X}, \mathbf{K} \mathbf{W}_k \mathbf{X} $$$\mathbf{X} (\mathbf{x}1, ... \mathbf{x}{N})$ 是输入序列。投影 $\mathbf{Q}$ 和 $\mathbf{K}$ 各含 $N$ 个向量于是 $\mathbf{QK}^T$ 的规模是 $N^2$。LLM 通常有多个注意力头并行做自注意力计算。假设 LLM 有 40 个注意力头、以 bfloat16 运行存储 $\mathbf{QK^T}$ 矩阵的显存需求是 $40 \times 2 \times N^2$ 字节$N 1000$只需约 50 MB 显存$N 16000$需要 19 GB 显存$N 100,000$仅存储 $\mathbf{QK}^T$ 矩阵就接近 1 TB。一句话默认自注意力算法在大输入上下文下迅速变得内存昂贵。随着 LLM 能力增强它们从处理几句话的翻译/摘要走向处理整个页面的任务长输入成为刚需。出路在于一种不需要显式构造 $\mathbf{QK}^T$ 矩阵的新注意力算法——Tri Dao 等人提出的Flash Attention。其核心思想是把 $\mathbf{V} \times \operatorname{Softmax}(\mathbf{QK}^T)$ 的计算拆开通过对多组 Softmax 分块步骤迭代逐块计算输出的小片段$$ \textbf{O}i \leftarrow s^a{ij} * \textbf{O}i s^b{ij} * \mathbf{V}{j} \times \operatorname{Softmax}(\mathbf{QK}^T{i,j}) \quad \text{for multiple } i, j \text{ iterations} $$其中 $s^a_{ij}$ 与 $s^b_{ij}$ 是需要在每一组 $i, j$ 上重新计算的 Softmax 归一化统计量。Flash Attention 的完整实现更复杂这里做了大幅简化深入细节请阅读 Flash Attention 原论文。关键要点通过追踪 Softmax 归一化统计量并配合巧妙的数学处理Flash Attention 在显存成本只随 $N$ 线性增长的前提下给出与默认自注意力层数值完全相同的输出。从公式直觉上会以为 Flash Attention 更慢——它确实需要更多 FLOPs因为 Softmax 归一化统计量必须不断重算。但推理中它远快于默认注意力Flash Attention 大幅降低了对 GPU 上更慢的高带宽显存VRAM的访问需求转而利用芯片上更快的 SRAM。本质上Flash Attention 保证所有中间读写都发生在片上高速 SRAM而不是访问慢速 VRAM 来计算输出向量 $\mathbf{O}$。实践建议很明确只要可用就没有任何理由不用 Flash Attention——它数学上输出相同且更快、更省内存。4. 架构创新为长文本任务优化的模型结构前面我们从两个方向改进了计算与内存效率把权重转成低精度格式用更高效版本替换自注意力算法。本节讨论第三个维度——让 LLM 的结构本身适配长文本任务例如检索增强问答、摘要、聊天。注意聊天不仅要求处理长文本输入还要求高效处理用户与助手之间的多轮对话如 ChatGPT 场景。模型一旦训练完成基础结构就很难改变因此在设计阶段就要预判任务并相应优化架构。结构中有两个组件对长输入序列最先成为内存和/或性能瓶颈位置编码positional embeddings键值缓存key-value cache4.1 位置编码从绝对编码到 RoPE / ALiBi自注意力把每个 token 与其他 token 建立关联。对输入序列Hello, I, love, you$\operatorname{Softmax}(\mathbf{QK}^T)$ 矩阵可以形如每个词 token 获得一组概率质量用以关注所有其他词 token从而与全部其他 token 建立关系。例如单词love以 5% 关注Hello、以 30% 关注I、以 65% 关注自身。一个基于自注意力但没有位置编码的 LLM会很难理解各输入文本的相对位置因为 $\mathbf{QK}^T$ 计算出的概率得分以 $O(1)$ 计算把每个词 token 关联到每个其他词 token与相对位置距离无关。对没有位置编码的 LLM 来说每个 token 到所有其他 token 的距离看起来都一样区分Hello I love you与You love I hello将非常困难。为了让 LLM 理解句子顺序需要额外线索通常以位置编码positional encodings也称positional embeddings的形式注入位置编码把每个 token 的位置编码为数值表示让 LLM 据此更好理解句子顺序。《Attention Is All You Need》的 authors 引入了正弦位置嵌入$\mathbf{P} \mathbf{p}_1, \ldots, \mathbf{p}_N$其中每个向量 $\mathbf{p}_i$ 是其位置 $i$ 的正弦函数。位置编码直接加到输入序列向量上$$ \mathbf{\hat{X}} \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N \mathbf{x}_1 \mathbf{p}_1, \ldots, \mathbf{x}_N \mathbf{p}_N $$从而提示模型更好地学习句子顺序。另一种路线如 BERT 的 Devlin et al.使用学习式位置嵌入即在训练中学习 $\mathbf{P}$。正弦与学习式位置嵌入曾是编码句子顺序的主流方法但发现了两个问题二者都是绝对位置嵌入即为每个位置 id$0, \ldots, N$编码唯一嵌入。研究表明绝对位置嵌入导致 LLM 在长文本输入上表现差——对长文本让模型学习 token 之间的相对位置距离比绝对位置更有优势使用学习式位置嵌入时LLM 必须在固定输入长度 $N$ 上训练难以外推到比训练长度更长的输入。近期能解决上述问题的相对位置嵌入更受青睐其中最突出的是旋转位置嵌入RoPE, Rotary Position EmbeddingALiBiRoPE 与 ALiBi 的共同观点是句子顺序信息最好直接注入自注意力算法——因为 token 正是在那里建立相互关系的具体而言应通过修改 $\mathbf{QK}^T$ 的计算来注入顺序信息。RoPE指出位置信息可以编码进 query-key 向量对如 $\mathbf{q}_i$ 与 $\mathbf{k}_j$中把每个向量分别旋转角度 $\theta \cdot i$ 与 $\theta \cdot j$$i, j$ 是各向量的句子位置$$ \mathbf{\hat{q}}_i^T \mathbf{\hat{k}}_j \mathbf{q}i^T \mathbf{R}{\theta, i - j} \mathbf{k}_j $$其中 $\mathbf{R}_{\theta, i - j}$ 是旋转矩阵。$\theta$不在训练中学习而是设为一个依赖于训练时最大输入序列长度的预定义值。这样处理后$\mathbf{q}_i$ 与 $\mathbf{k}_j$ 之间的概率得分仅在 $i \ne j$ 时受影响且只取决于相对距离 $i - j$与两个向量的具体位置 $i$、$j$ 无关。RoPE 被当今许多重要 LLM 采用如 Falcon、Llama、PaLM。仓库中对应实现集中在 src/transformers/modeling_rope_utils.pyRoPE 频率计算与配置解析各模型如 Qwen、Llama 系列的建模代码通过它构造旋转缓存。ALiBi则提出更简单的相对位置编码方案把输入 token 之间的相对距离作为负整数 × 预定义缩放值m在计算 softmax 前直接加到 $\mathbf{QK}^T$ 矩阵的每个 query-key 条目上。ALiBi 论文表明这种简单的相对位置编码让模型在超长文本输入序列上仍保持高性能。ALiBi 同样被 MPT、BLOOM 等重要模型采用。外推能力对比RoPE 与 ALiBi 都能外推到训练中未见过的输入长度但开箱即用的外推效果 ALiBi 明显优于 RoPE对 ALiBi只需把下三角位置矩阵的数值加大以匹配输入序列长度即可对 RoPE沿用训练时的 $\theta$ 会在远长于训练长度的输入上得到糟糕结果不过社区发现了一些调整 $\theta$ 的有效技巧如缩放 RoPE 基频让 RoPE 位置编码也能在extrapolated 输入上正常工作。RoPE 与 ALiBi 都是不在训练中学习的相对位置编码基于如下直觉关于文本输入的位置线索应直接给到自注意力层的 $\mathbf{QK}^T$ 矩阵应促使 LLM 学习 token 之间恒定的相对距离位置编码token 彼此距离越远query-key 概率应越低。RoPE 通过增大 query-key 向量夹角来降低向量点积ALiBi 通过给点积加大负数来降低。结论面向需要处理大文本输入任务的 LLM最好用 RoPE 或 ALiBi 这类相对位置编码训练。注意即使模型只在固定长度如 $N_1 2048$上训练借助位置编码外推实际可以用远大于 $N_1$ 的输入如 $N_2 8192$。4.2 KV Cache自回归生成的加速器及其代价LLM 的自回归文本生成方式是反复输入序列 → 采样下一个 token → 把该 token 追加到输入序列 → 直到模型生成结束标记 token。先用一段简码展示自回归生成的朴素形态用torch.argmax取最可能的下一 tokeninput_ids tokenizer(prompt, return_tensorspt)[input_ids].to(cuda) for _ in range(5): next_logits model(input_ids)[logits][:, -1:] next_token_id torch.argmax(next_logits, dim-1) input_ids torch.cat([input_ids, next_token_id], dim-1) print(shape of input_ids, input_ids.shape) generated_text tokenizer.batch_decode(input_ids[:, -5:]) generated_text输出shape of input_ids torch.Size([1, 21]) shape of input_ids torch.Size([1, 22]) shape of input_ids torch.Size([1, 23]) shape of input_ids torch.Size([1, 24]) shape of input_ids torch.Size([1, 25]) [ Here is a Python function]可以看到每一步都把刚采样的 token 追加进输入序列。除极少数例外LLM 都用因果语言建模目标训练因此注意力结果矩阵的上三角被掩码——token从不依赖后续 token$\mathbf{q}_i$ 永远不会与 $j i$ 的 $\mathbf{k}j, \mathbf{v}j$ 建立关系只关注此前的 $\mathbf{k}{m i}, \mathbf{v}{m i}$。为了减少不必要的计算可以缓存每一层在所有前序时间步的 key-value 向量。在 Transformers 中向forward传入use_cache标志即可取回 KV cache并在下一步连同当前 token 一起传回past_key_values None # past_key_values 就是 KV cache generated_tokens [] next_token_id tokenizer(prompt, return_tensorspt)[input_ids].to(cuda) for _ in range(5): next_logits, past_key_values model(next_token_id, past_key_valuespast_key_values, use_cacheTrue).to_tuple() next_logits next_logits[:, -1:] next_token_id torch.argmax(next_logits, dim-1) print(shape of input_ids, next_token_id.shape) print(length of key-value cache, past_key_values.get_seq_length()) # past_key_values 形状为 [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim] generated_tokens.append(next_token_id.item()) generated_text tokenizer.batch_decode(generated_tokens) generated_text输出shape of input_ids torch.Size([1, 1]) length of key-value cache 20 shape of input_ids torch.Size([1, 1]) length of key-value cache 21 shape of input_ids torch.Size([1, 1]) length of key-value cache 22 shape of input_ids torch.Size([1, 1]) length of key-value cache 23 shape of input_ids torch.Size([1, 1]) length of key-value cache 24 [ Here, is, a, Python, function]可以看到使用 KV cache 时输入 token 序列不再变长始终是一个输入向量而 KV cache 的长度在每步解码时增加 1。使用 KV cache 意味着 $\mathbf{QK}^T$ 本质上退化成了 $\mathbf{q}_c\mathbf{K}^T$其中 $\mathbf{q}_c$ 是当前传入 token 的 query 投影永远只是单个向量。KV cache 带来两大优势计算效率显著提升相比计算完整 $\mathbf{QK}^T$ 矩阵计算量更少推理速度更快峰值内存不再随生成 token 数二次增长而是线性增长。应始终使用 KV cache它结果相同且对长输入序列有显著加速。Transformers 在使用 text pipeline 或generate方法 时默认启用 KV cache。仓库中 src/transformers/cache_utils.py 里的DynamicCache是生成模型的默认缓存实现逐层存储 key/value 张量形状为[batch_size, num_heads, seq_len, head_dim]若传入 config还会根据滑动窗口/混合层结构把缓存张量压缩到[batch_size, num_heads, min(seq_len, sliding_window), head_dim]进一步降低内存需求。更完整的缓存体系说明见仓库文档 KV Cache 指南。⚠️ 注意尽管建议使用 KV cache启用它之后 LLM 输出可能与未启用时略有不同。这是矩阵乘法内核kernel本身的数值特性所致属于已知行为。4.2.1 多轮对话KV cache 复用历史上下文KV cache 对聊天这类需要多次自回归解码的应用尤其有用。看一个例子User: How many people live in France? Assistant: Roughly 75 million people live in France User: And how many are in Germany? Assistant: Germany has ca. 81 million inhabitants在这段对话中LLM 运行了两次自回归解码第一轮KV cache 为空输入 prompt 是User: How many people live in France?模型自回归生成Roughly 75 million people live in France每个解码步都在增长 KV cache第二轮输入 prompt 本应是User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?。得益于 cache前两句话的 key-value 向量已经计算过所以实际只需处理User: And how many in Germany?。处理这段缩短的 prompt 时其计算出的 key-value 向量直接拼接进第一轮解码留下的 cache随后基于完整三句话的 cache 自回归生成第二个回答Germany has ca. 81 million inhabitants。有两点值得注意保留完整上下文对聊天 LLM 至关重要否则模型无法理解对话前文。例如上例中模型需要理解用户问 And how many are in Germany 时指的是人口KV cache 对聊天极有价值它允许持续生长已编码的聊天历史而不必从头重新编码若用 encoder-decoder 结构就必须从头重编。在transformers中当传入return_dict_in_generateTrue叠加默认的use_cacheTrue时generate调用会返回past_key_values。注意该能力尚不通过pipeline接口暴露# Generation as usual prompt system_prompt Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here model_inputs tokenizer(prompt, return_tensorspt) generation_output model.generate(**model_inputs, max_new_tokens60, return_dict_in_generateTrue) decoded_output tokenizer.batch_decode(generation_output.sequences)[0] # Piping the returned past_key_values to speed up the next conversation round prompt decoded_output \nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here model_inputs tokenizer(prompt, return_tensorspt) generation_output model.generate( **model_inputs, past_key_valuesgeneration_output.past_key_values, max_new_tokens60, return_dict_in_generateTrue ) tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]输出is a modified version of the function that returns Mega bytes instead. def bytes_to_megabytes(bytes): return bytes / 1024 / 1024 Answer: The function takes a number of bytes as input and returns the number of很好不再为注意力层重复计算相同的 key 与 value。但要注意一个代价$\mathbf{QK}^T$ 矩阵的峰值内存需求被大幅降低了把 KV cache 留在内存里对长输入序列或多轮聊天可能非常吃内存——cache 需要为所有自注意力层、所有注意力头保存所有前序输入向量 $\mathbf{x}_i$$i \in {1, \ldots, c-1}$的 key-value 向量。以bigcode/octocoder为例计算 KV cache 需要存储的浮点数个数2k 和 v× 序列长度 × 注意力头数 × 每头维度 × 层数。在假想的 16000 输入序列长度下config model.config 2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head输出7864320000约80 亿个浮点数以float16存储这 80 亿个浮点数需要约 15 GB 内存——接近模型权重本身的一半研究者提出了两种显著降低 KV cache 内存成本的方法即下文两小节。4.2.2 多查询注意力MQANoam Shazeer 在Fast Transformer Decoding: One Write-Head is All You Need中提出MQA。如标题所示与其使用n_head组 key-value 投影权重只需用一组跨所有注意力头共享的 value 投影权重模型性能并不会显著退化。使用单组 value 投影权重后key-value 向量 $\mathbf{k}_i, \mathbf{v}_i$ 在所有注意力头中必然相同因此 cache 中只需存 1 组 key-value 投影而不是n_head组。由于大多数 LLM 使用 20 到 100 个注意力头MQA 能大幅降低 KV cache 的内存消耗。对本指南使用的 LLM可以把 16000 序列长度下的 cache 内存需求从 15 GB 降到不到 400 MB。除了省内存MQA 还提升计算效率。自回归解码中每步都需要把巨大的 key-value 向量重新读回、与当前 key-value 向量拼接再送入 $\mathbf{q}_c\mathbf{K}^T$ 计算——持续重载带来的内存带宽需求可能成为严重的时间瓶颈。减小 key-value 向量体积即减少内存访问从而缓解带宽瓶颈。要理解的关键点把 key-value 注意力头数降到 1只有在配合 KV cache 使用时才有意义。不使用 KV cache 的单个前向传播峰值内存不变——每个注意力头仍有唯一 query 向量每个头仍有不同的 $\mathbf{QK}^T$ 矩阵。MQA 已被社区广泛采用众多热门 LLM 都在使用包括 Falcon、PaLM、MPT、BLOOM本指南使用的bigcode/octocoder检查点同样采用 MQA。4.2.3 分组查询注意力GQA谷歌的 Ainslie et al. 提出的GQA发现相比使用常规的多个 key-value 头投影MQA 有时会带来质量退化。该工作认为可以通过更温和地减少 query 头投影权重数量来保住更多模型性能不用 1 组 key-value 投影权重而是用n n_head组。把n选为远小于n_head的值如 2、4、8既能保住 MQA 绝大部分内存与速度收益又只牺牲更少的模型容量因此性能损失更小。此外GQA 作者发现已存在的模型检查点可以用不到原始预训练 5% 的计算量uptrain成 GQA 结构。虽然 5% 的预训练计算量本身可能仍是天文数字但 GQA uptraining 让既有检查点可以服务于更长的输入序列。GQA 提出较晚撰写本文档时采用率相对更低其最著名的应用是 Llama-v2。总结性建议如果 LLM 以自回归解码方式部署、且需要处理大输入序列如聊天场景强烈建议使用 GQA 或 MQA。5. 结语把最好的算法与架构用足研究社区不断为越来越大的 LLM 提出加速推理的新方法。例如很有前景的**推测解码speculative decoding**方向由更小更快的语言模型生成简单 token只有困难 token才交给大 LLM 生成。GPT3/4、Llama-2-70b、PaLM 这类巨型 LLM 之所以能在 Hugging Face Chat、ChatGPT 等聊天界面中高速运行很大程度上正得益于本文讨论的**精度量化、算法Flash Attention、架构相对位置编码 KV Cache MQA/GQA**三个层面的优化。未来 GPU、TPU 等加速器只会更快、内存更大但无论如何都应确保用上了当前可用的最佳算法与架构把每一份算力花在刀刃上。快速决策参考场景首选手段预期收益显存装不下模型8-bit / 4-bit 量化BitsAndBytesConfig本例 29 GB → 15.2 GB8-bit→ 9.5 GB4-bit多卡加载超大模型device_mapauto流水线并行 / 张量并行突破单卡 80 GB 上限长上下文≥16k tokenFlash Attention注意力显存从 $O(N^2)$ 降到 $O(N)$长文本任务摘要/检索/聊天RoPE 或 ALiBi 相对位置编码支持外推到训练长度之外的输入自回归解码 多轮聊天始终启用 KV cache考虑 MQA/GQA 架构解码步数计算量下降16k 序列 cache 从约 15 GB 降到 400 MBMQA【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考