整体说明

  • FLOPS（Floating-Point Operations Per Second）和FLOPs（Floating-Point Operations）是衡量计算性能的两个相关但含义不同的术语
  • FLOPs 是浮点运算次数 ，是模型复杂度的评估指标，用于评估一个模型的复杂度
  • FLOPS 是每秒浮点运算次数 ，是硬件计算能力的单位，比如用于评估 GPU 性能

FLOPs（Floating-Point Operations）

  • FLOPs 是 浮点运算次数 ，即模型或算法执行的总浮点计算量（如加、减、乘、除等操作数量）
  • FLOPs 常用于衡量算法/模型的计算复杂度 ，举例：
    • 矩阵乘法中，两个 \( n \times n \) 矩阵相乘需要 \( 2n^3 \) FLOPs（\( n \times n \)个数，每个数需要 \(n\) 次乘法和 \(n\) 次加法操作）
    • 在深度学习中，FLOPs常用来估计模型的计算开销（如卷积层的计算量）
  • 1 GFLOPs = 10亿次运算（算法复杂度）
  • 注：”s” 为小写，表示复数（Operations）
  • 用法：ResNet-50 模型约需 3.8 GFLOPs（38亿次浮点运算）处理一张图像
    • 注意这里的十亿用的是 G，而不是 B，两者都有十亿的含义，但不同地方用不同的值

FLOPS（Floating-Point Operations Per Second）

  • FLOPS 是每秒浮点运算次数 ，是硬件计算能力的单位
  • FLOPS 常用于衡量处理器（如CPU/GPU）的理论峰值性能。例如：
    • 1 FLOPS = 1次浮点运算/秒
    • 1 TFLOPS（Tera-FLOPS）= \( 10^{12} \) 次浮点运算/秒
  • 1 GFLOPS = 10亿次运算/秒（硬件速度）
  • 注：”S” 为大写，代表 “Second”（每秒）
  • 用法：NVIDIA A100 GPU的峰值性能为 312 TFLOPS

一些常见错误表达

  • 错误的写法如 “Flops” 或 “flops” 可能导致歧义，建议严格区分大小写

注：本文包含 AI 辅助创作

Paper Summary

  • 整体总结：
    • 核心：DuoAttention 是一种通过区分 Retrieval Heads 和 Streaming Heads 来优化 LLM 内存和计算资源的框架
    • 具体：DuoAttention 可以显著减少了长上下文应用中解码和 Pre-filling 的内存使用和延迟
      • 因为 DuoAttention 对 Retrieval Heads 应用完整的 KV 缓存（Streaming Heads 仅缓存 Sink Token 和 Recent Token）
    • 对比之前 MHA 和 GQA 的效果（内存大幅减少、解码速度大幅提升）
      • MHA 模型内存减少高达 \(2.55\times\)，MHA 模型解码速度提升高达 \(2.18\times\)，Pre-filling 加速高达 \(1.73\times\)
      • GQA 模型内存减少高达 \(1.67\times\)，GQA 模型解码速度提升高达 \(1.50\times\)，Pre-filling 加速高达 \(1.63\times\)
      • 且与完全注意力相比准确率损失最小（minimal accuracy loss）
    • 当与量化结合时，DuoAttention 可以进一步提升 KV 缓存容量，在单个 A100 GPU 上支持高达 3.30M 个上下文 Token
  • 背景 & 问题提出：
    • 部署长上下文（long-context）LLM 至关重要，但长上下文带来了显著的计算和内存挑战
    • 跨所有注意力头缓存所有 Key 和 Value （KV）状态会消耗大量内存
    • 现有的 KV 缓存剪枝方法要么损害 LLM 的长上下文能力，要么仅提供有限的效率提升
  • 作者发现：
    • 只有一小部分注意力头（Retrieval Heads），对于处理长上下文至关重要，并且需要对所有 Token 进行完整的注意力计算
    • 而其他头（Streaming Heads），主要关注最近的 Token 和 Attention Sinks，不需要完整的注意力计算
  • 基于这一洞察，论文引入了 DuoAttention：
    • 该框架仅对 Retrieval Heads 应用完整的 KV 缓存，同时对 Streaming Heads 使用轻量级的、恒定长度的 KV 缓存
    • 从而在不损害其长上下文能力的情况下，减少 LLM 解码和 Pre-filling 的内存占用和延迟
  • DuoAttention 使用一种轻量级的、基于优化的算法以及合成数据来准确识别 Retrieval Heads
  • 内存方面：
    • 对于多头注意力模型最高减少 2.55\(\times\)
    • 对于分组 Query 注意力模型最高减少 1.67\(\times\)
  • 效率方面：
    • 对于多头注意力模型解码速度最高提升 2.18\(\times\)， Pre-filling 速度最高提升 1.73\(\times\)
    • 和分组 Query 注意力模型解码速度最高提升1.50\(\times\)， Pre-filling 速度最高提升 1.63\(\times\)
    • 与完整注意力相比，准确率损失最小（minimal accuracy loss）
  • 开源链接：github.com/mit-han-lab/duo-attention

Introduction and Discussion

  • LLM 处于人工智能革命的前沿，驱动着高级应用，如多轮对话、长文档摘要以及涉及混合模态的任务，如视觉和视频理解
    • 这些应用通常需要处理大量的上下文 Token ；
    • 例如，总结整个《哈利·波特》系列可能涉及分析约一百万个 Token
    • 对于视觉语言模型，挑战更加严峻，其中一张 224×224 的图像对应 256 个 Token ，而一段三分钟、24 FPS 的视频会生成约 1.1M 个 Token
  • 在此类应用中部署 LLM 的一个关键问题是长上下文推理问题
    • 完整的注意力机制要求所有 Token 关注所有先前的 Token 以获得准确的表示，这导致解码延迟线性增加， Pre-filling 延迟二次方增加
    • KV 缓存技术存储所有先前 Token 的 Key 和 Value ，导致内存使用量随上下文长度线性增长
    • 随着序列变长，内存越来越多地被 KV 缓存消耗，给注意力机制带来了显著的计算负担
      • 例如，在 Llama-3-8B 模型架构中，为 1M 个 Token 提供服务并使用 FP16 KV 缓存将需要至少 137 GB 的内存（这已经超过了单个 80GB GPU 的容量）
    • 而且使用如此大上下文进行 Pre-filling 和解码会有显著延迟，这对 LLM 在长上下文场景中的有效使用构成了重大挑战
  • 尽管有许多努力来克服注意力机制在长上下文推理中的挑战，但显著的计算和内存问题仍然存在
    • 架构修改，如分组 Query 注意力，需要模型预训练，并且无法降低计算成本
      • 线性注意力（Linear Attention）方法虽然在计算和内存需求上较低，但在长上下文场景下往往不如 Transformer 模型
      • 近似注意力（Approximative attention）方法，如 H\({}_{2}\)O、StreamingLLM、TOVA 和 FastGen，常常在长上下文应用中牺牲精度，并且与关键的 KV 缓存优化技术（如分组 Query 注意力）不兼容
    • KV 缓存量化虽然有用，但并未减少注意力机制的计算时间
      • 系统级优化，包括 FlashAttention、FlashDecoding 和 PagedAttention，虽然有效，但并未减少 KV 缓存大小，并且在扩展上下文时仍然需要大量计算
  • 论文引入了一个关键观察
    • LLM 中的注意力头可以分为两种不同的类型 ：Retrieval Heads 和 Streaming Heads ，如图 1 所示
    • Retrieval Heads 仅占总头数的一小部分，对于处理长上下文至关重要，并且需要对所有 Token 进行完整的注意力计算
    • Streaming Heads（大多数注意力头），主要关注最近的 Token 和 Attention Sinks ，并且可以在仅包含 Recent Token 和 Attention Sinks 的简化 KV 缓存下有效运行
  • 基于 Retrieval Heads 和 Streaming Heads 的二分法，论文提出了 DuoAttention
    • DuoAttention 是一种通用、直接且易于集成的方法，能显著加速 LLM 的解码和 Pre-filling ，并减少内存占用，尤其是在长上下文场景中
    • DuoAttention 的核心创新是一种轻量级的、基于优化的过程，它使用合成数据集来识别不可压缩的 Retrieval Heads
    • 与依赖注意力模式分析 (2024;) 的现有方法不同，DuoAttention 直接测量因 Token 丢弃而产生的输出偏差，从而实现更高的压缩率和改进的部署效率
  • DuoAttention 的设计注重简洁和高效：每个 Transformer 层有两个 KV 缓存
    • 一个用于关键 Retrieval Heads 的完整 KV 缓存
    • 一个用于 Streaming Heads 的恒定 KV 缓存，仅存储 Attention Sinks 和最近的 Token
  • 这种设计使得 DuoAttention 能够显著减少内存使用、提高模型的解码速度，且与完整注意力相比，精度损失最小
  • DuoAttention 与重要的优化技术（如分组 Query 注意力和量化）完全兼容
    • 当结合 8-bit 权重和 4-bit KV 缓存量化时，DuoAttention 使得 Llama-3-8B 模型能够在单个 A100 GPU 上处理高达 3.3M 上下文 Token
      • 与标准的完整注意力 FP16 部署相比，实现了 \(6.4\times\) 的容量提升
    • DuoAttention 为在需要百万级上下文处理的应用中部署 LLM 铺平了道路

DuoAttention

Retrieval Heads 和 Streaming Heads

Retrieval Heads
  • 在基于 Transformer 的 LLM 中，注意力头表现出独特且一致的模式，反映了它们的专门功能
  • 图 1 使用句子“最好的水果是橙子。什么是最好的水果？橙子。”可视化了 Llama-2-7B-32K-Instruct 模型中的两种注意力头
  • 左图突出显示了一个在解码过程中强调相关 Token 的注意力头；
    • 例如，在解码第二个“最好的水果”时，第一个“最好的水果”被加重；在推断第二个“橙子”时，初始的“橙子”被突出显示
    • 这些注意力头，论文称之为 Retrieval Heads ，对于上下文处理至关重要，因为它们捕获了上下文相关的 Token
    • 压缩 Retrieval Heads 的 KV 缓存将导致关键上下文信息的丢失，因此它们需要对所有 Token 进行完整的注意力计算
Streaming Heads
  • 图 1 中间图描绘的注意力头主要关注最近的 Token 和 Attention Sinks，不强调上下文中较早的相关 Token
    • 论文称这些为 Streaming Heads
  • 压缩 Streaming Heads 的 KV 缓存是可行的，因为丢弃未被关注的中国 Token 不会显著改变注意力输出
    • 可以通过仅保留 Attention Sinks 和 Recent Token 的 KV 状态来优化 Streaming Heads ，而不会损害模型管理长上下文的能力
Impact of Token Pruning on Retrieval and Streaming Heads
  • 图 1 的右图显示了一个初步的 Passkey 检索实验
    • 当 Retrieval Heads KV 缓存中的中间 Token 被剪枝时，模型的性能显著下降
    • 移除 Streaming Heads 的中间 Token 对 Passkey 检索精度没有显著影响
  • 这一观察表明，我们可以在不牺牲模型长上下文能力的情况下提高计算效率：
    • 通过丢弃 Streaming Heads 的中间 Token ，同时保持 Retrieval Heads 的完整注意力，将 Streaming Heads 的内存需求降低到 \(O(1)\)，从而提高了处理长上下文的效率

Optimization-Based Identification of Retrieval Heads

Definition of Retrieval Heads
  • 第 2.1 节定性地定义了 Retrieval Heads 和 Streaming Heads ，但为了精确识别，论文需要一个具体且量化的定义
  • 在论文中，论文将“Retrieval Heads”定义为：
    • 当被限制为仅关注 Recent Token 和 Attention Sinks 时，会显著改变模型输出的注意力头
  • 论文使用这个标准来区分 Retrieval Heads 和 Streaming Heads
    • 这个定义不同于现有工作 (2024; )，它们仅依赖注意力分数来识别 Retrieval Heads ，忽略了
      • 1）压缩特定注意力头 KV 缓存的端到端影响
      • 2）Value 状态的角色
      • 3）注意力分布在层和头之间的可变性
    • 论文的定义直接测量输出偏差 ，即使它们在注意力分数中不明显，论文也能够识别对长上下文处理至关重要的注意力头
    • 论文在第 3.5 节中提供的消融研究支持了这一论点
Optimization-based Identification
  • 论文采用一种基于优化的方法来识别 Retrieval Heads ，灵感来自先前在 CNN 滤波器剪枝方面的工作，如图 2 所示
    • 首先为 LLM 中的每个 KV 头分配一个门控值 \(\alpha_{i,j}\)
      • 这个值直观地表示了第 \(i\) 层第 \(j\) 个 KV 头在处理长上下文信息时的重要性
      • 在使用分组 Query 注意力的模型中，一个 KV 头可能与多个注意力头相关联，论文的方法考虑了对整个注意力头组的 KV 缓存压缩
  • 论文的基于优化的识别方法直接评估了仅使用 Sink Token 和 Recent Token 压缩每个 KV 头 KV 缓存的影响
    • 首先将每个头的门控值 \(\alpha_{i,j}\in[0,1]\) 初始化为 1，假设所有头最初都作为 Retrieval Heads
    • 然后优化这些门控值，同时保持 LLM 的参数固定，将可训练参数的数量限制在 \(N\times H\)，并防止影响模型的原始能力
  • 在前向传播过程中，论文结合每个 KV 头的完整注意力和流式注意力的输出，使用门控值作为混合权重：
    $$\texttt{attn}_{i,j}=\alpha_{i,j}\cdot\texttt{full_attn}+(1-\alpha_{i,j})\cdot\texttt{streaming_attn}$$
    • 其中注意力计算定义为：
      $$\texttt{full_attn} =\texttt{softmax}(\boldsymbol{Q}\boldsymbol{K}^{T}\odot\boldsymbol{M}_{\text{causal} })\boldsymbol{V}, \\
      \texttt{streaming_attn} =\texttt{softmax}(\boldsymbol{Q}\boldsymbol{K}^{T}\odot\boldsymbol{M}_{\text{streaming} })\boldsymbol{V},$$
    • 其中 \(\boldsymbol{M}_{\text{causal} }\) 是因果注意力掩码，而 \(\boldsymbol{M}_{\text{streaming} }\) 表示一个类 \(\Lambda\) 掩码，仅关注最近和初始的 Token
Synthetic Dataset for Identifying Retrieval Heads
  • 仅依赖自然语言建模目标不足以识别 Retrieval Heads
    • 自然文本中需要长跨度推理的监督信号是稀疏的，且大多数 Token 可以使用局部上下文进行推断
  • 论文设计了一个专门旨在增强模型长上下文检索能力的合成数据集，使论文能够有效地识别哪些 KV 头可以在不损害模型性能的情况下被压缩
  • 如图 3 所示，论文通过在一个非常长的上下文中，在十个随机位置嵌入十个随机生成的 \(s\) 个 Token 的 passkey sequences 来创建一个 passkey-retrieval 数据集
    • 模型的任务是在上下文末尾回忆这十个序列
Training and Loss Functions
  • 论文优化蒸馏损失，即完整注意力模型的最后一个隐藏状态与使用 DuoAttention 的模型的最后一个隐藏状态之间的 L2 差异，仅关注整个输入中最后 \(l\) 个 Passkey Token ：
    $$\mathcal{L}_{\text{distill} }=\frac{1}{N}\sum_{i=1}^{N}\sum_{j=\bar{T}-l+1}^{T}(\boldsymbol{H}_{\text{full} }^{(i)}[j]-\boldsymbol{H}_{\text{mixed} }^{(i)}[j])^{2}$$
  • 论文的合成数据集确保每个监督信号都与最终的压缩策略相关，使得该过程在信息检索精度方面是无损的
    • 事实证明，它比仅使用自然语言建模更有效
    • 论文使用 L1 正则化项来鼓励门控值的稀疏性：
      $$\mathcal{L}_{\text{reg} }=\sum_{i=1}^{L}\sum_{j=1}^{H}|\alpha_{i,j}|,.$$
  • 最终的训练损失是蒸馏损失和正则化损失的组合，由一个超参数 \(\lambda\) 加权，论文在实验中将其设置为 0.05：
    $$\mathcal{L}=\mathcal{L}_{\text{distill} }+\lambda\mathcal{L}_{\text{reg} }.$$
  • 由于可训练参数的总数仅为数千个浮点数，此优化过程相当快，仅需要 2,000 步
    • 论文论文中的所有训练实验都可以在 8×NVIDIA A100 GPU 服务器上进行

Deploying LLMs with DuoAttention

Binarizing Attention Implementations（二值化注意力）
  • 在推理时，论文仅对指定的 Retrieval Heads 应用完整注意力，这些 Retrieval Heads 是使用训练阶段优化的门控值识别的
  • 论文根据阈值 \(\tau\) 对每个头的注意力策略进行二值化，以区分 Retrieval Heads 和 Streaming Heads ：
    $$\text{attn}_{i,j}=\begin{cases}\text{full_attn}&\text{if }\alpha_{i,j}>\tau \\ \text{streaming_attn}&\text{otherwise}\ \end{cases}$$
Reordering Attention Heads（重排注意力头）
  • 在部署之前，论文通过根据注意力头分配重新排序 Query 、 Key 和 Value 投影权重的输出通道来预处理模型
  • 这种重新排序将 Retrieval Heads 和 Streaming Heads 分组为两个不同的、连续的簇，从而允许在层内管理这两种类型头的 KV 缓存时进行高效的切片和连接操作，而不是依赖 scattering 和 gathering 操作
Decoding
  • 如图 5 所示，论文在解码期间为 LLM 的每一层分配两个 KV 缓存 ：
    • 一个用于 Retrieval Heads ，存储所有过去的 Key 和 Value ；
    • 另一个用于 Streaming Heads ，仅存储 Attention Sinks 和最近的 Token ，保持恒定大小
  • 当处理一个新 Token 时，其 Query 、 Key 和 Value 向量沿头维度分割，以计算 Retrieval Heads 的完整注意力和 Streaming Heads 的流式注意力
    • 然后将结果沿头维度连接以进行输出投影
Chunked Pre-filling（分块 Pre-filling）
  • 论文使用 FlashAttention-2 来 Pre-fill Retrieval Heads 和 Streaming Heads 的 KV 缓存
    • 在长上下文 LLM 中，分块 Pre-filling 是一种常见做法，将提示分成固定长度的块来 Pre-filling KV 缓存
    • 这种技术通过将线性层中的峰值中间激活大小从序列长度降低到块大小，显著降低了峰值内存使用
  • DuoAttention 与分块 Pre-filling 完全兼容，并且 DuoAttention 中 Streaming Heads 的 Pre-filling 可以在线性时间和恒定内存复杂度下实现，无需专门的核
  • 如图 5 所示，计算了某一层的 KV 后，Streaming Heads 的 KV 缓存会立即被剪枝，仅保留 Sink Token 和最近的 Token
    • 下一个传入 Token 块在 Pre-filling 期间将仅关注恒定数量的上下文 Token
  • 令 \(L\) 表示序列长度，\(K\) 表示块大小（chunk size）
    • Streaming Heads 的 Pre-filling 时间复杂度从 \(O(L^{2})\) 优化到 \(O(LK)\)，内存复杂度从 \(O(L)\) 减少到 \(O(K)\)
  • 需要注意的是，DuoAttention 的设计非常适合批量操作，这可以在具有大批量大小的服务场景中进一步提高 LLM 的效率

Experiments

Setups

  • 模型、数据集和基线 (Models, Datasets, and Baselines)
    • 论文在长上下文和短上下文基准测试上评估 DuoAttention，证明论文的方法在保留模型处理长短上下文任务性能的同时，显著提高了效率
      • 对于长上下文评估
        • 论文使用 Needle-in-a-Haystack (NIAH) 基准测试 (Kamradt, 2024) 和 LongBench (2023)
      • 对于短上下文评估
        • 论文评估了在 MMLU (2021)、MBPP (2021) 和 MT-Bench (2023) 上的性能
    • 论文采用了 SOTA 开源模型，包括 Llama-2-7B-chat (2023b)（及其长上下文变体 Llama-2-7B-32K-Instruct (Together, 2023)）、Llama-3-[8,70]B-Instruct（及其长上下文变体 Llama-3-8B-Instruct-Gradient-1048k）以及 Mistral-7B-v0.2-Instruct (2023)
    • 论文将论文的方法与 KV 缓存压缩算法进行了比较，包括 H2O (2023b)、TOVA (2024)、FastGen (2024) 和 StreamingLLM (2023b)
  • Implementation details
    • 论文使用 PyTorch (2019) 和来自 FlashInfer (2024) 的 RoPE (2021) 和 RMSNorm 内核来实现 DuoAttention
    • 对于 Retrieval Heads 的识别
      • 论文使用批量大小为 1，将 10 个 32 词（Words）的 passkeys 插入到 BookSum (2021) 数据集中
      • 识别过程使用 128 个 Sink Token 和 256 个 Recent Token
      • 训练样本从范围为 1,000 个 Token 到模型特定的最大长度（间隔 50 intervals）中采样（问题：这里是指样本长度的采样）

        Training samples are drawn from 50 intervals ranging from 1,000 tokens to the model-specific maximum length

    • passkeys 在上下文中的 1000 个点处随机插入（更多细节包含在附录 A.1 节中）
    • 论文使用 AdamW (2015) 优化器优化门控值，初始学习率为 0.02，在前 400 步从 0.002 进行预热，并在最后 400 步降回 0.002
      • 所有实验在 NVIDIA A100 GPU 上运行 2,000 步

Long-Context Benchmarks

  • 使用 Needle-in-a-Haystack (NIAH) 基准测试和 LongBench (2023) 来评估 DuoAttention
  • 使用了两个长上下文模型：Llama-2-7B-32K-Instruct 和 Llama-3-8B-Instruct-Gradient-1048k
    • DuoAttention 配置：
      • Llama-2-7B-32K-Instruct 使用 25% 的 Retrieval Heads 比例
      • Llama-3-8B-Instruct-Gradient-1048k 使用 50% 的比例
    • 论文在相同的 KV 缓存预算下，将 DuoAttention 与 H2O、TOVA 和 StreamingLLM 进行比较
      • 论文为 DuoAttention 使用 64 个 Sink Token 、256 个 Recent Token 和 32,000 的 Pre-filling 块大小
    • 由于 H2O 和 TOVA 的原始设计不支持长上下文，论文修改了它们的算法，将 Pre-filling 阶段替换为 FlashAttention，并模拟输入最后 50 个 Token 的解码（遵循 Tang 等人 (2024b) 的方法）
    • FastGen 的算法不允许指定 KV 压缩比，因为它会随输入波动
      • 论文调整了注意力恢复比例，以确保在图 6 所示的实验中，KV 缓存预算平均高于 25% 或 50%
    • FastGen 在 Attention Profiling 阶段的二次内存成本限制了其处理长上下文样本的能力
      • 论文测量了 FastGen 在 NIAH 上对 Llama-2-7B 最高到 24K 上下文、对 Llama-3-8B 最高到 32K 上下文的性能；
      • 超过这些大小会导致内存不足错误
    • 详细的基线实现和理由在附录 A.3 节和 A.5 节中提供
  • Needle-in-a-Haystack (NIAH) 是一个具有挑战性的压力测试，旨在评估模型从冗长上下文中准确识别和检索相关信息的能力
    • 如图 6 所示，所有基线方法都无法从长序列的不同深度检索到正确答案，因为它们在生成过程中丢弃了包含必要信息的 KV 缓存
    • DuoAttention 保留了 Retrieval Heads 中的所有 KV 缓存，同时仅丢弃 Streaming Heads 中的缓存，从而保留了模型的检索能力
    • DuoAttention 在所有序列深度上都表现出强大的性能，有效处理高达 1048K Token 的长度
  • LongBench (2023) 是一个全面的长上下文数据集套件，涵盖多个任务和自然文本，旨在更全面地评估长上下文理解能力
    • 图 7 显示了在 14 个 LongBench 任务上的性能，比较了不同方法基于其 KV 缓存预算的表现
    • DuoAttention 在大多数任务上显示出 KV 预算和准确性之间的优越权衡，突显了其泛化能力
    • DuoAttention 在大多数任务上实现了与完全注意力相当的性能，对 MHA 使用 25% 的 KV 缓存预算，对 GQA 使用 50% 的 KV 缓存预算，这与在 needle-in-a-haystack 基准测试中观察到的结果一致
    • 论文在附录的表 5 和表 6 中将 DuoAttention 与 FastGen 进行了比较
    • 附录中的表 3 和表 4 提供了两个模型使用 25% 和 50% KV 缓存预算在所有 21 个 LongBench 任务上的完整结果，表明 DuoAttention 在大多数任务上始终优于基线，并取得了最高的平均分数

Short-Context Benchmarks

  • 为了确保 DuoAttention 不损害模型在短上下文任务上的性能，论文将其与所有基线一起在三个短上下文基准测试上进行了评估：MMLU、MBPP 和 MT-Bench
    • 这些基准测试评估模型的知识、编码能力和帮助性
    • 对 MMLU 使用 one-shot 提示，对 MBPP 和 MT-Bench 使用 zero-shot 提
    • 对于 DuoAttention，在 MMLU 上配置 32 个 Sink Token 和 128 个 Recent Token ，在 MBPP 和 MT-Bench 上配置 16 个 Sink Token 和 64 个 Recent Token
  • 如图 8 和表 1 所示
    • 在相同的 KV 缓存预算下，DuoAttention 在各种模型（包括 Llama-2-7B、Llama-3-8B 和 Llama-3-70B-Instruct）上始终优于所有基线
    • 在 50% KV 缓存预算下，DuoAttention 在大多数基准测试上实现了近乎无损的性能，表明它保留了模型的原始能力

Efficiency Results

  • 论文在单个 NVIDIA A100 GPU 上评估了 DuoAttention 在 Llama-2-7B 和 Llama-3-8B 模型上的解码延迟和内存使用情况
  • 论文为整个基准测试序列预分配 KV 缓存，以防止动态内存分配的额外开销
  • 权重和激活的默认数字格式为 BFloat16
  • 通过对 Llama-2-7B 采用 25% 的 Retrieval Heads 比例，对 Llama-3-8B 采用 50% 的比例，DuoAttention 在保持准确性的同时显著提高了效率
Decoding Efficiency
  • 如图 9 所示
    • DuoAttention 的解码速度呈线性缩放，但与完全注意力相比斜率更平缓，这反映了所选的 Retrieval Heads 比例
      • 这种高效的缩放带来了内存使用的显著减少和解码速度的显著提升
    • 这些改进随着上下文长度的增加而接近 Retrieval Heads 比例的倒数
  • 图 11 显示
    • 在固定上下文大小下，DuoAttention 在不同 KV 预算设置下的加速和内存节省
    • 随着部署配置中 Retrieval Heads 比例的降低，解码延迟和内存使用都线性下降
    • 在图 11 的设置下，DuoAttention 在 A100 GPU 上实现了最大改进：MHA 模型内存减少 2.55 倍，GQA 模型内存减少 1.67 倍；MHA 模型延迟减少 2.18 倍，GQA 模型延迟减少 1.50 倍
Pre-filling Efficiency
  • 如第 2.3 节所述，DuoAttention 也加速了 LLM 的长上下文 Pre-filling
  • 图 10 显示
    • DuoAttention 显著降低了 Pre-filling 延迟和内存使用，并且这些节省随着 Pre-filling 块大小的减小而增加
      • 这是因为 Streaming Heads 的时间和内存复杂度随着块大小的减小而降低
    • DuoAttention 实现了 MHA 模型延迟减少高达 1.73 倍，GQA 模型延迟减少高达 1.63 倍，同时 MHA 模型内存减少高达 2.38 倍，GQA 模型内存减少高达 1.53 倍
Combination with Quantization
  • 为了将更多 Token 装入有限的内存，我们可以将权重和 KV 缓存量化与 DuoAttention 结合，以最大化 KV 缓存容量
  • 先前的研究表明，权重量化 (2023a;) 和 4-bit KV 缓存量化 (2024;) 不会损害模型性能
  • 论文将 DuoAttention 与 QServe (2024) 量化方法和内核相结合，以实现 8-bit 权重和 4-bit KV 缓存的 LLM 推理
  • 测量结果如图 12 所示
    • 将量化技术与 DuoAttention 结合，使论文能够在单个 A100-80G GPU 上使用 Llama-3-8B 模型容纳高达 3.30M 个 Token ，与朴素的完全注意力 BF16 部署相比，容量增加了 \(6.4\times\)

Ablation Studies

  • 论文使用 Mistral-7B-Instruct-v0.2 在 passkeys 检索和 MMLU 数据集上进行了消融研究
  • 对于 passkeys 检索任务，论文将一个 8 词的 passkeys 嵌入到一个 30K 词的文本中，并在 100 个插入深度上进行线性扫描，报告精确匹配准确率
  • 基于优化与基于 Attention Profiling 的 Retrieval Heads 识别 (Optimization-based vs. Attention Profiling-based Retrieval Head Identification)
    • 论文评估了论文的基于优化的方法与 FastGen (2024) 和 RazorAttention (2024a) 中使用的 Attention Profiling 方法，两者使用相同的合成 passkeys 数据集
    • 图 13 (1) 中的结果表明，论文的方法显著优于 Attention Profiling ，后者难以识别 Retrieval Heads ，从而影响了模型的准确优化
  • 使用合成数据优化与语言建模 (Optimizing with Synthetic Data vs. Language Modeling)
    • 如图 13 (1) 所示，论文使用合成数据识别 Retrieval Heads 的方法比传统的 Language Modeling（在自然数据中的所有 Token 上计算损失）产生了明显更好的结果
  • 优化中结合 Sink 和 Recent 注意力的必要性 (Necessity of Sink+Recent Attention in Optimization)
    • 图 13 (2) 强调了在优化阶段结合 Sink 和 Recent 注意力的重要性
    • 仅依赖 Sink Token 或 Recent Token 注意力不足以有效识别 Retrieval Heads
  • 部署阶段配置 (Deployment Phase Configuration)
    • 论文分析了 Streaming Heads 中注意力 Sink 和 Recent Token 的部署配置
    • 论文的发现表明：
      • 性能在 16 个 Sink Token 和 64 个 Recent Token 时达到稳定（图 13 (3)）
      • 进一步增加只会带来边际改进
    • 问题：论文的发现跟 StreamingLLM 的 4 个 Token 足以的发现有矛盾！
  • 已有很多方法在扩展 LLM 并提高其处理长上下文的效率；这些方法可以分为四个主要类别：
    • 优化模型架构、使用近似注意力机制、应用 KV 缓存量化以及系统级优化
  • Model Architecture
    • MQA (2019) 和 GQA (2023) 通过在 Query 头之间共享 KV 头来减小 KV 缓存的大小
    • 但这些方法需要使用特定架构进行预训练，并且不会降低计算成本（但会降低显存）
    • 线性注意力 Transformer (2023) 减少了内存使用，但在需要长上下文处理的任务上往往表现不佳
  • Approximate Attention
    • 诸如 Sparse Transformer (2019) 和 LongFormer (2020) 等方法使用 Local Attention 或 Block Attention 模式来降低计算复杂度
    • BigBird (2020) 通过结合 Local Attention 和 Global Attention 实现线性复杂度，但其中许多方法需要定制的 GPU 内核或重新训练，限制了其实用性
    • H2O (2023b) 和 TOVA (2024) 基于 Query 模式丢弃 Token 来简化注意力
    • StreamingLLM (2023b) 识别了“注意力 Sink ”并提出始终保留 Initial Token 和 Recent Token 以维持恒定的解码延迟和内存使用，使模型能够处理比预训练序列长度多得多的输入 Token
    • FastGen (2024) 分析注意力头以在解码期间丢弃 Token
    • 论文的实验表明：
      • 这些方法会降低 LLM 的长上下文能力
      • 这些方法无法降低长上下文 LLM 的 Pre-filling 成本
  • KV Cache Quantization
    • 诸如 8-bit 和 4-bit 量化 (2024; 2024; 2024) 等技术减小了 KV 缓存的大小，但它们没有解决注意力内核的计算开销问题
    • 这些方法与 DuoAttention 是互补的，可以结合使用以进一步减少内存使用
  • System Optimizations
    • vLLM (2023) 和 FlashAttention (2022; 2023) 通过优化批处理（Batch Processing）和利用 GPU 内存层次结构来提高注意力计算效率
    • FlashDecoding (2024) 和 RingAttention (2023a) 在解码速度和序列级并行性方面引入了进一步的改进
    • 这些方法提高了计算性能，但它们没有解决 KV 缓存大小减少的问题，它们与 DuoAttention 互补，以实现额外的速度和内存优化
  • Recent Works
    • 一些近期工作与 DuoAttention 有相似的想法
    • Wu 等人 (2024) 引入了 Retrieval Heads 的概念来解释 LLM 的长上下文能力
      • 但他们的方法没有压缩非 Retrieval Heads 的 KV 缓存，仅关注准确性
    • MInference (2024) 通过使用稀疏注意力模式来加速长上下文 LLM 的 Pre-filling
      • 但没有优化解码期间的 KV 缓存存储或延迟
    • RazorAttention (2024a) 也将注意力头分为 Retrieval 和 Non-Retrieval 类别
      • 但 RazorAttention 使用 Attention Profiling 方法而不是 Optimization-based 方法区区分
        • 论文的实验表明，Attention Profiling-based 方法不如论文的 Optimization-based 的方法准确
      • 而且，RazorAttention 没有优化 Pre-filling
        • DuoAttention 提供了更有效的 KV 缓存管理和更高的压缩率，从而在长上下文应用中为 Pre-filling 和解码带来了更好的性能

Appendix A

A.1 Experimental Details

  • 论文使用 PyTorch (2019) 中的 FSDP 进行模型训练，并使用 DeepSpeed Ulysses (2023) 序列并行来支持长序列
  • 在训练期间，论文使用 Guo 等人 (2024) 实现的、如图 14 所示的高效块稀疏近似 \(\Lambda\) 类注意力来计算流式注意力
  • 不同模型的最大序列长度各不相同，详见表 2

A.2 Full LongBench Results

A.3 在长上下文基准测试上 H2O 和 TOVA 的实现 (Implementation of H2O and TOVA on Long-Context Benchmarks)

  • H2O (2023b) 和 TOVA (2024) 算法的原始设计与 Pre-filling 阶段的 FlashAttention (2022) 不兼容，因为它们依赖注意力分数来执行 Token Eviction（驱逐）
    • 由于 FlashAttention 中的注意力分数从未被具体化，这些算法无法用于 Pre-filling ，这是它们的主要缺陷之一
    • 因此，不可能在像“大海捞针”和 LongBench 这样的长上下文设置中评估这些算法，因为它们会在上下文 Pre-filling 期间导致内存不足（OOM）
  • 为了与这些策略进行比较，论文修改了算法：
    • 在 Pre-filling 期间，论文使用 FlashAttention 进行精确计算
    • 在解码阶段，论文根据生成 Token 对上下文 Token 的注意力分数执行 Token Eviction
  • 这种修改相比原始设计提高了性能，因为 Pre-filling 是精确的，并且 Token Eviction 仅发生在解码期间
    • 在极端情况下，如果答案中只有一个生成 Token （例如，多项选择题任务），论文实现的 H2O 和 TOVA 将与完全注意力一样精确，这并非它们的真实精度
    • 为了接近它们的真实性能，论文在长输入基准测试（“大海捞针”和 LongBench）中模拟最后 50 个 Token 作为生成 Token ，以足够长时间地执行它们的 Token Eviction 策略，论文的算法也是如此
  • 此实验设置也被 Tang 等人 (2024b) 使用
    • 实验结果表明论文的方法可以通过此压力测试，而 H2O 和 TOVA 则不能

A.4 Implementation of FastGen on Long-Context Benchmarks

  • 由于缺乏 FastGen (2024) 算法的官方实现，论文使用一个社区代码库 (2024) 对其进行了复现，该代码库被 FastGen 的官方仓库引用
    • 在 FastGen 算法中，剪枝比率不能直接配置；而是使用恢复比率 \(T\) 来控制稀疏度，如 FastGen 论文中所述
  • 为了量化稀疏度，论文计算了所有测试用例的平均 KV 缓存使用量作为整体稀疏度的度量
    • 对于 Llama-2-7B 模型，论文将恢复比率设置为 \(0.7\)，确保平均 KV 缓存预算超过完整 KV 缓存的 25%
    • 对于 Llama-3-8B 模型，论文将恢复比率设置为 \(0.87\)，确保平均 KV 缓存预算超过完整 KV 缓存的 50%
  • 由于 FastGen 使用用户提供提示的完整注意力图来分析不同头的类型，它会导致 \(O(n^{2})\) 的注意力图复杂度
    • 论文无法在长上下文中测试其性能
  • 对于长上下文基准测试，论文使用了 8 个 A100-80G GPU，对于 Llama-2-7B 模型实现了最高 24k Token 的序列长度，对于 Llama-3-8B 模型实现了最高 32k Token 的序列长度
  • 除了图 6 中显示的“大海捞针”基准测试结果外，论文还评估了FastGen 在两个模型上的 LongBench 表现
    • 但由于 FastGen 的二次内存消耗，论文仅报告了在 8x A100-80G GPU 上使用 FastGen 可以运行的数据集结果
    • 如表 5 和表 6 所示，DuoAttention 在 LongBench 数据集上 consistently 优于 FastGen

补充表格和图标

  • 图 15: NIAH result on the Mistral-7B-Instruct-v0.2 model
  • 图 16: NIAH result on the Mistral-7B-Instruct-v0.3 model

注：本文包含 AI 辅助创作

Paper Summary

  • 整体说明：
    • 作者发现：Window Attention 提供了一个部分解决方案，但当初始 Token 被排除时，其性能会急剧下降（这些 Token 作为“Attention Sink”的作用很重要）
    • 本文提出了一个简单而高效的框架 StreamingLLM，使 LLM 能够在无需微调的情况下处理无限长度的文本
      • 通过将 Attention Sink 与 Recent Token 结合，StreamingLLM 可以高效地对多达 4M Token 的文本进行建模
    • 论文还进一步通过实验证实，使用专用的 Sink Token 预训练模型可以改善流式性能
      • StreamingLLM 首次将 LLM 的预训练窗口大小与其实际文本生成长度解耦，为 LLM 的流式部署铺平了道路
  • 第一个问题提出：在流式（streaming）应用（如多轮对话）中部署 LLM 是迫切需要的，但面临两大挑战
    • 挑战一：在解码阶段，缓存先前 Token 的键和值状态（KV）会消耗大量内存
    • 挑战二：popular LLM 无法泛化到比训练序列长度更长的文本
  • 第二个问题提出：Window Attention 是一种自然的方法，仅缓存最近的 KV，但论文发现当文本长度超过缓存大小时，该方法会失效
  • Insight：作者观察到一个有趣的现象，即 Attention Sink ：
    • 保留（Keeping）初始 Token（initial Token）的 KV 会大幅恢复 Window Attention 的性能
    • 注：Window Attention 会丢失最初的 Token 信息（窗口外），而刻意保留初始 Token 的 KV 能大幅提升模型性能
    • 论文证明了 Attention Sink 的出现是由于对初始 Token 的 Strong 注意力分数 ，即使它们在语义上并不重要
    • 注：Attention Sink 的定义应该是 无论初始 Token 与语言建模任务的相关性如何，都有大量的注意力分数分配给了初始 Token
  • 方案：StreamingLLM
    • StreamingLLM 是一个高效的框架，使经过有限长度注意力窗口训练的 LLM 能够无需任何微调即可泛化到无限序列长度
    • StreamingLLM 能够使 Llama-2、MPT、Falcon 和 Pythia 在高达 4M Token 甚至更多的文本上实现稳定高效的语言建模
  • 论文的其他发现，在预训练期间添加一个占位符 Token 作为专用的 Attention Sink 可以进一步改善流式部署
    • 在流式设置中，StreamingLLM 相比滑动窗口重计算基线实现了高达 22.2 倍的加速

Introduction and Discussion

  • LLM (2018; 2020; 2022; OpenAI, 2023; 2023a, 2023b) 正变得无处不在，驱动着许多自然语言处理应用，如对话系统 (2022; 2023; 2023)、文档摘要 (2020; 2023a)、代码补全 (2021; 2023) 和问答 (2023)
    • 为了释放预训练 LLM 的全部潜力，它们应该能够高效且准确地执行长序列生成
      • 例如，一个理想的聊天机器人助手应该能够在长达数日的对话内容上稳定工作
    • 但对于 LLM 来说，泛化到比其预训练长度更长的序列是非常具有挑战性的，例如 Llama-2 的 4K (2023b)
  • 原因是 LLM 在预训练期间受到注意力窗口的限制
    • 尽管在扩展此窗口大小 (2023; 2023; 2023) 以及改进长输入的训练 (2022; 2023) 和推理 (2022; 2023; 2023; 2021; 2023b) 效率方面付出了大量努力，可接受的序列长度本质上仍然是有限的 ，这不允许持久部署
  • 论文介绍了 LLM 流式应用的概念，并提出了一个问题：论文能否在不牺牲效率和性能的情况下，为无限长度的输入部署 LLM？
  • 当将 LLM 应用于无限输入流时，会出现两个主要挑战：
    • 1）在解码阶段，基于 Transformer 的 LLM 会缓存所有先前 Token 的键和值状态（KV），如图 1 (a) 所示，这可能导致过多的内存使用和不断增加的解码延迟 (2022)
    • 2）现有模型的长度外推能力有限，即当序列长度超过预训练期间设置的注意力窗口大小时，它们的性能会下降 (2023; 2022)
  • 一种直观的方法，称为 Window Attention (2020)（图 1 b），仅维护最近 Token 的 KV 状态的固定大小滑动窗口
    • 虽然它在缓存初始填满后确保了恒定的内存使用和解码速度，但一旦序列长度超过缓存大小，模型就会崩溃，即即使只是驱逐第一个 Token 的 KV ，如图 3 所示
  • 另一种策略是带重计算的滑动窗口（如图 1 c 所示），它为每个生成的 Token 重建最近 Token 的 KV 状态
    • 虽然它提供了强大的性能，但由于需要在其窗口内计算二次注意力，这种方法明显更慢，使其对于现实世界的流式应用不切实际
  • 为了理解 Window Attention 的失败原因，论文发现了自回归 LLM 的一个有趣现象：无论初始 Token 与语言建模任务的相关性如何，都有大量的注意力分数分配给了初始 Token ，如图 2 所示
    • 论文称这些 Token 为“ Attention Sink ”
    • 尽管它们缺乏语义重要性，却收集了显著的注意力分数
    • 论文将原因归咎于 Softmax 操作，它要求所有上下文 Token 的注意力分数总和为一
    • 因此，即使当前查询在许多先前的 Token 中没有强匹配项，模型仍然需要将这些不需要的注意力值分配到某处，以便总和为一
    • 初始 Token 成为汇聚 Token 背后的原因是直观的：
      • 由于自回归语言建模的性质，初始 Token 对几乎所有后续 Token 都是可见的，这使得它们更容易被训练成 Attention Sink
  • 基于上述见解，论文提出了 StreamingLLM，一个简单高效的框架，使经过有限注意力窗口训练的 LLM 能够无需微调即可处理无限长度的文本
    • StreamingLLM 利用了 Attention Sink 具有高注意力值这一事实，保留它们可以保持注意力分数分布接近正常
    • 因此，StreamingLLM 只需保留 Attention Sink Token 的 KV（仅需 4 个初始 Token 就足够）以及滑动窗口的 KV，以锚定注意力计算并稳定模型的性能
  • 通过 StreamingLLM，包括 Llama-2 (2023b)、MPT (Team, 2023)、Falcon (2023) 和 Pythia (2023) 在内的模型可以可靠地对 4M Token 进行建模，甚至可能更多
    • 与唯一可行的基线——带重计算的滑动窗口相比，StreamingLLM 实现了高达 22.2 倍的加速，实现了 LLM 的流式使用
  • 图 1: StreamingLLM 与现有方法的示意图 在长度为 \(L\) 的文本上预训练的语言模型预测第 \(T\) 个 Token (\(T\gg L\))
    • (a) Dense Attention 具有 \(O(T^{2})\) 的时间复杂度和不断增长的缓存大小
      • 当文本长度超过预训练文本长度时，其性能下降
    • (b) Window Attention 缓存最近 \(L\) 个 Token 的 KV
      • 虽然在推理中高效，但一旦起始 Token 的键和值被驱逐，性能就会急剧下降
    • (c) 带重计算的滑动窗口（Sliding Window with Re-computation）为每个新 Token 从 \(L\) 个最近 Token 重建 KV 状态
      • 虽然它在长文本上表现良好，但其 \(O(TL^{2})\) 的复杂度（源于上下文重计算中的二次注意力）使其相当慢
    • (d) StreamingLLM 保留 Attention Sink （几个初始 Token ）以进行稳定的注意力计算，并结合了最近 Token
      • 在扩展文本上高效且提供稳定的性能
    • 困惑度是使用 Llama-2-13B 模型在 PG-19 测试集中第一本书（65K Token ）上测量的
  • 图 2: Llama-2-7B 在 256 个句子上的平均注意力对数概率可视化，每个句子长度为 16。观察包括：
    • (1) 前两层（第 0 层和第 1 层）的注意力图呈现出“局部”模式，最近 Token 获得更多注意力
    • (2) 在底部两层之上，模型在所有层和头中都严重关注（heavily attends）初始 Token

StreamingLLM

The Failure of Window Attention and Attention Sinks

  • 虽然 Window Attention 技术在推理过程中提供了效率，但它导致了极高的语言建模困惑度
    • 该模型的性能不适合部署在 流式应用 (streaming applications) 中
    • 论文使用 Attention Sink 的概念来解释 Window Attention 的失败，这为 StreamingLLM 提供了灵感
  • 识别困惑度激增点 (Identifying the Point of Perplexity Surge) 图 3 显示了在 20K Token 文本上的语言建模困惑度
    • 当文本长度超过 缓存 (cache) 大小时，由于排除了初始 Token ，困惑度会急剧上升
    • 这表明，初始 Token ，无论它们与预测 Token 的距离如何，对于维持 LLM 的稳定性都至关重要
Why do LLMs break when removing initial tokens’ KV?
  • 论文在图 2 中可视化了 Llama-2-7B 和模型所有层和头中的注意力图
  • 论文发现，除了底部两层之外（注：这里的底部两层是最开始的两层），模型在所有层和头上都持续关注初始 Token
    • 这意味着：移除这些初始 Token 的 KV 将移除 SoftMax 函数（公式 1）注意力计算中分母的相当一部分
  • 这种改变导致注意力分数的分布发生显著变化，偏离了正常推理环境下的预期
    $$\text{SoftMax}(x)_{i}=\frac{e^{x_{i} } }{e^{x_{1} }+\sum_{j=2}^{N}e^{x_{j} } },\quad x_{1}\gg x_{j},j\in 2,\ldots,N \tag{1}$$
  • 对于初始 Token 在语言建模中的重要性，有两种可能的解释：
    • (1) 它们的语义至关重要
    • (2) 模型学习到了对其绝对位置的偏好
  • 为了区分这两种可能性，论文进行了实验（表 1），其中前四个 Token 被替换为换行符 Token “\n”
    • 观察结果表明，模型仍然显著关注这些初始的换行符 Token
    • 此外，重新引入它们可以将语言建模困惑度恢复到与拥有原始初始 Token 相当的水平
    • 这表明起始 Token 的绝对位置（而非其语义价值）具有更重要的意义
LLMs attend to Initial Tokens as Attention Sinks（LLM 将初始 Token 视为 Attention Sink ）
  • 为了解释为什么无论它们与语言建模的语义相关性如何，模型都不成比例地关注初始 Token，论文引入了 “Attention Sink” 的概念
  • SoftMax 函数（公式 1）的性质阻止所有被关注的 Token 具有零值
    • 这要求在所有层的所有头中从其他 Token 聚合一些信息，即使当前的 Embedding 已经有足够的 self-contained 信息用于预测
    • 因此，模型倾向于将不必要的注意力值转储到特定的 Token 上
  • 在量化异常值领域也进行了类似的观察 (2023; 2023)，这导致了提出 SoftMax-Off-by-One (Miller, 2023) 作为潜在的补救措施
  • 图 3：各种 LLM 在 20K Token 文本上的语言建模困惑度。观察结果显示了一致的趋势：
    • (1) 一旦输入长度超过预训练注意力窗口大小， Dense Attention 就会失败
    • (2) 一旦输入长度超过缓存大小，即初始 Token 被逐出（evicted）， Window Attention 就会崩溃
    • (3) StreamingLLM 表现出稳定的性能，其困惑度几乎与带重计算的滑动窗口 (sliding window with re-computation) 基线相匹配
  • 为什么各种自回归 LLM，如 Llama-2、MPT、Falcon 和 Pythia，都一致地将初始 Token 作为它们的 Attention Sink，而不是其他 Token ？(Why do various autoregressive LLMs, such as Llama-2, MPT, Falcon, and Pythia, consistently focus on initial tokens as their attention sinks, rather than other tokens?)
    • 论文的解释很简单：由于自回归语言建模的顺序性质，初始 Token 对所有后续 Token 都是可见的，而后来的 Token 仅对有限的后续 Token 集合可见
    • 因此，初始 Token 更容易被训练成为 Attention Sink，捕获不必要的注意力
  • 论文注意到，LLM 通常被训练为使用多个初始 Token 作为 Attention Sink，而不仅仅是一个
    • 如图 2 所示，引入四个初始 Token 作为 Attention Sink，足以恢复 LLM 的性能
      • 只添加一个或两个则无法实现完全恢复
    • 作者认为这种模式的出现是因为这些模型在预训练期间没有在所有输入样本中包含一致的起始 Token
    • 尽管 Llama-2 确实在每个段落前加上一个 <s> Token ，但这发生在文本分块（text chunking）之前 ，导致第零个位置大多被随机 Token 占据
      • 问题：如何理解这里的 text chunking 会影响第一个 <s> Token ？
    • 这种缺乏统一起始 Token 的情况导致模型使用几个初始 Token 作为 Attention Sink
  • 论文假设，通过在所有训练样本的开头加入一个稳定的可学习 Token ，它可以单独作为一个专门的 Attention Sink，从而无需多个初始 Token 来确保一致的流式处理
    • 论文将在第 3.3 节验证这一假设

Rolling KV Cache with Attention Sinks

  • 为了在已训练的 LLM 中启用 LLM 流式处理，论文提出了一种简单的方法，可以在不进行任何模型微调的情况下恢复 Window Attention 的困惑度
    • 方法：除了当前的滑动窗口 Token 之外，论文在注意力计算中重新引入了几个起始 Token 的 KV
  • StreamingLLM 中的 KV 缓存概念上可以分为两部分，如图 4 所示：
    • (1) Attention Sink (四个初始 Token ) 稳定注意力计算；
    • (2) 滚动 KV 缓存 (Rolling KV Cache) 保留最近的 Token ，这对语言建模至关重要
  • StreamingLLM 的设计是通用的，可以无缝集成到任何使用相对位置编码的自回归语言模型中，例如 RoPE (2021) 和 ALiBi (2022)
  • 在确定相对距离并向 Token 添加位置信息时，StreamingLLM 关注的是 缓存内 的位置（within the cache），而不是 原始文本中 的位置
    • 这种区别对 StreamingLLM 的性能至关重要
    • 例如，如果当前缓存（图 4）有 Token [0, 1, 2, 3, 6, 7, 8] 并且正在解码第 9 个 Token ，则分配的位置是 [0, 1, 2, 3, 4, 5, 6, 7]，而不是原始文本中的位置 [0, 1, 2, 3, 6, 7, 8, 9]
      • 问题：直观上看，怎么觉得这样反而会出现问题？因为原始文本中的真实相对位置信息被修改了
  • 对于像 RoPE 这样的编码
    • 论文在引入旋转变换 之前 缓存 Token 的键 (Keys)
    • 然后在每个解码阶段，论文对滚动缓存中的键应用位置变换
  • 另一方面，与 ALiBi 集成更直接
    • 这里，对注意力分数应用连续线性偏置，而不是 ‘跳跃’ 偏置
    • 这种在缓存内分配位置 Embedding 的方法对 StreamingLLM 的功能至关重要，确保模型即使在其预训练注意力窗口大小之外也能高效运行
  • 表 1：
    • Window Attention 在长文本上表现不佳
    • 当论文重新引入最初的四个 Token 以及最近的 1020 个 Token (对应表中 4+1020) 时，困惑度得以恢复
    • 将原始的四个初始 Token 替换为换行符 Token “n” (对应表中 4”n”+1020) 实现了 comparable 困惑度恢复
    • 缓存配置 x+y 表示添加 x 个初始 Token 和 y 个最近 Token
    • 困惑度是在 PG19 测试集中第一本书（65K Token ）上测量的
  • 表 2：重新引入的初始 Token 数量对 StreamingLLM 的影响
    • (1) Window Attention (0+y) 的困惑度急剧增加
    • (2) 引入一个或两个初始 Token 不能完全恢复模型困惑度，表明模型不仅仅使用第一个 Token 作为 Attention Sink
    • (3) 引入四个初始 Token 通常就足够了；进一步添加收益递减
    • 缓存配置 x+y 表示将 x 个初始 Token 添加到 y 个最近 Token
    • 困惑度是在 concatenated PG19 测试集中的 400K Token 上评估的

Pre-Training LLMs with Attention Sinks

  • 如第 3.1 节所述，模型过度关注多个初始 Token 的一个重要原因是缺乏一个指定的 Sink Token 来卸载过多的注意力分数
    • 因此，模型无意中使用了全局可见的 Token ，主要是初始 Token ，作为 Attention Sink
    • 一个潜在的补救措施可以是故意加入一个全局可训练的 Attention Sink Token ，表示为 “Sink Token”，它将作为不必要注意力分数的储存库
    • 或者，用像 SoftMax-off-by-One (2023) 这样的变体替换传统的 SoftMax 函数，
      $$\text{SoftMax}_{1}(x)_{i}=\frac{e^{x_{i} } }{1+\sum_{j=1}^{N}e^{x_{j} } } \tag{2}$$
  • 它不要求所有上下文 Token 上的注意力分数总和为 1，可能也是有效的
    • 注意 SoftMax\(_1\) 相当于在注意力计算前添加一个具有全零键和值特征的 Token
      • 问题：因为全是零，所以各种 Attention 的加权平均后均等价于没有增加该 Token
    • 论文将此方法称为 “Zero Sink“ 以符合论文的框架
  • 为了验证，论文在相同设置下从头开始预训练了三个具有 160M 参数的语言模型
    • 第一个模型使用标准的 SoftMax 注意力 (Vanilla)
    • 第二个模型用 SoftMax\(_1\)（Zero Sink）替换了常规的注意力机制
    • 第三个模型在所有训练样本前添加了一个可学习的占位符 Token （Sink Token）
  • 如表 3 所示，Zero Sink 在某种程度上缓解了 Attention Sink 问题，但模型仍然依赖其他初始 Token 作为 Attention Sink
    • 引入 Sink Token 在稳定注意力机制方面非常有效
    • 只需将此 Sink Token 与最近的 Token 配对就足以稳定模型的性能，并且最终的评估困惑度甚至略有改善
    • 鉴于这些发现，论文建议在所有样本中使用 Sink Token 来训练未来的 LLM，以优化流式部署
  • 表 3：比较在预训练期间使用标准注意力、前置零 Token 和可学习 Sink Token
    • 为了确保稳定的流式困惑度，标准模型需要几个初始 Token
    • 虽然 Zero Sink 显示出轻微改进，但它仍然需要其他初始 Token
    • 若使用可学习 Sink Token 训练的模型，仅添加 Sink Token 就显示出稳定的流式困惑度
    • 缓存配置 \(x\)+\(y\) 表示添加 \(x\) 个初始 Token 和 \(y\) 个最近 Token
    • 困惑度是在 PG19 测试集中第一个样本上评估的（问题：1 个样本就够评估困惑度了？）

Experiments

  • 论文使用四个近期主流的模型家族来评估 StreamingLLM：Llama-2 (2023b)、MPT (2023)、Pythia (2023) 和 Falcon (2023)
    • Llama-2、Falcon 和 Pythia 采用了 RoPE (2021)
    • MPT 采用了 ALiBi (2022)
    • RoPE 和 ALiBi 是近期研究中两种最具影响力的位置编码技术
  • 论文多样化的模型选择确保了研究结果的有效性和鲁棒性
    • 论文将 StreamingLLM 与已建立的基线方法进行比较，例如 Dense Attention、 Window Attention 以及带重计算的滑动窗口方法（Sliding Window with Re-computation）
    • 在所有后续使用 StreamingLLM 的实验中，除非另有说明，论文默认使用四个初始 Token 作为 Attention Sink

Language Modeling on Long Texts Across LLM Families and Scales

  • 论文首先使用 Concatenated PG19 (2020) 测试集评估 StreamingLLM 的语言建模困惑度（Perplexity），该测试集包含 100 本长书籍
  • 对于 Llama-2 模型，缓存大小设置为 2048，而对于 Falcon、Pythia 和 MPT 模型，则设置为 1024
    • 这是预训练窗口大小的一半，选择此值是为了增强可视化清晰度
  • 图 3 表明
    • 在跨越 20K Token 的文本上，StreamingLLM 在困惑度方面可以与 Oracle 基线（带重计算的滑动窗口）相媲美
    • 当输入长度超过其预训练窗口时， Dense Attention 技术会失败；
    • 当输入长度超过缓存大小时， Window Attention 技术会因初始 Token 被逐出而表现不佳
  • 在图 5 中，论文进一步证实了 StreamingLLM 可以可靠地处理异常长的文本，涵盖超过 4M 个 Token ，跨越一系列模型家族和规模
    • 这包括 Llama-2 [7,13,70]B、Falcon [7,40]B、Pythia-[2,8,6,9,12]B 和 MPT-[7,30]B

Results of Pre-Training with a Sink Token

  • 为了验证论文在所有预训练样本中引入一个 Sink Token 可以改进流式 LLM 的建议，论文在相同条件下训练了两个语言模型，每个模型有 160M 参数
    • 一个模型遵循原始训练设置
    • 另一个在每个训练样本的开头加入了一个 Sink Token
  • 论文的实验使用了 Pythia-160M (2023) 代码库并遵循其训练方法
    • 论文在一个 8xA6000 NVIDIA GPU 服务器上使用去重后的 Pile (2020) 数据集训练模型
    • 除了将训练批大小减少到 256 之外，论文保留了所有 Pythia 训练配置，包括学习率调度、模型初始化和数据集排列
    • 两个模型都训练了 143,000 步
  • 收敛性与正常模型性能 (Convergence and Normal Model Performance)
    • 在预训练期间包含一个 Sink Token 对模型收敛性以及后续在一系列 NLP 基准测试中的性能没有负面影响
    • 如图 6 所示，原始模型与使用 Sink Token 训练的模型表现出相似的收敛动态
    • 论文在七个不同的 NLP 基准测试上评估这两个模型，包括 ARC-[Challenge, Easy] (2018)、HellaSwag (2019)、LAMBADA (2016)、OpenbookQA (2018)、PIQA (2020) 和 Winogrande (2019)
    • 如表 4 所示，使用 Sink Token 预训练的模型与使用原始方法训练的模型表现相似
  • 流式性能 (Streaming Performance)
    • 如表 3 所示，使用传统方法训练的模型与使用 Sink Token 增强的模型在流式困惑度上存在差异
      • 原始模型需要添加多个 Token 作为 Attention Sink 以维持稳定的流式困惑度
      • 使用 Sink Token 训练的模型仅使用该 Sink Token 就能达到令人满意的流式性能
      • 注：其实原始训练方法需要 4 个 Token 这个事情，表 3 不够明显，表 2 MPT-7B 更明显
  • 注意力可视化 (Attention Visualization)
    • 图 7 对比了使用和不使用 Sink Token 预训练的模型的注意力图
    • 没有 Sink Token 的模型，类似于 Llama-2-7B（图 2），在浅层显示局部注意力，在深层则关注初始 Token
    • 相比之下，使用 Sink Token 训练的模型在所有层和头中都持续关注 Sink ，表明存在有效的注意力卸载机制
    • 这种对 Sink 的 Strong 关注，加上对其他初始 Token 注意力的减少，解释了 Sink Token 在提升模型流式性能方面的有效性

Results on Streaming Question Answering with Instruction-tuned Models

  • 为了展示 StreamingLLM 在现实世界中的适用性，论文使用指令微调（Instruction-tuned）的 LLM 模拟多轮问答，这在现实场景中很常见
  • 论文首先将 ARC-[Challenge, Easy] 数据集中的所有问答对拼接起来，将连续的流输入到 Llama-2-[7,13,70]B-Chat 模型中，并使用精确匹配（Exact Match）准则评估每个答案位置上的模型补全情况
  • 如表 5 所示，
    • Dense Attention 会导致内存不足（Out-of-Memory, OOM）错误，表明它不适合此设置
    • Window Attention 方法虽然运行高效，但由于输入长度超过缓存大小时会产生随机输出，导致准确率低下
    • StreamingLLM 表现出色，能高效处理流式格式，其准确率与 One-shot、Sample-by-sample 的基线准确率相当
  • 为了突出一个更适合 StreamingLLM 的场景，论文引入了一个数据集 StreamEval，其灵感来源于 LongEval (2023) 基准测试
  • 如图 8 所示
    • 与 LongEval 在长跨度设置上使用单一查询不同，论文每提供 10 行新信息就查询一次模型
    • 每个查询的答案始终在 20 行之前，这反映了现实世界中问题通常与近期信息相关的实例
    • 问题：从图 8 中看，查询间隔是 20 行，一定会超过窗口吗？
  • 如图 9 所示
    • 采用 StreamingLLM 的 LLM 即使在输入长度接近 120K Token 时也能保持合理的准确率
    • Dense Attention 和 Window Attention 分别在达到预训练文本长度和 KV 缓存大小时失败
  • 论文使用了两个上下文扩展模型，LongChat-7b-v1.5-32k (2023) 和 Llama-2-7B-32K-Instruct (2023)，以表明 StreamingLLM 可以与上下文扩展技术互补
    • 在 StreamingLLM 中，上下文扩展意味着扩大流式 LLM 的最大缓存大小，从而能够捕获更广泛的局部信息

Ablation Studies

  • 初始 Token 数量 (Numbers of Initial Tokens)
    • 在表 2 中，论文通过消融实验研究了添加不同数量的初始 Token 与 Recent Token 对流式困惑度的影响
    • 结果表明，仅引入一个或两个初始 Token 是不够的，而四个初始 Token 的阈值似乎就足够了，后续增加 Token 数量带来的效果微乎其微
    • 这一结果证明了论文在 StreamingLLM 中引入 4 个初始 Token 作为 Attention Sink 的选择是合理的
  • 缓存大小 (Cache Sizes)
    • 在表 6 中，论文评估了缓存大小（Cache Size， Attention Window Size）对 StreamingLLM 困惑度的影响
    • 与直觉相反 ，增加缓存大小并不会持续降低语言建模的困惑度
    • 这种不一致性表明了一个潜在的局限性，即这些模型可能无法最大化利用它们接收到的整个上下文信息
    • 未来的研究工作应致力于增强这些模型更好利用广泛上下文的能力

Efficiency Results

  • 论文将 StreamingLLM 的解码延迟（Decoding Latency）和内存使用量与带重计算的滑动窗口基线进行了基准测试，带重计算的滑动窗口 是唯一具有可接受质量的基线
  • 两种方法均使用 Huggingface Transformers (2020) 库实现，并在单个 NVIDIA A6000 GPU 上使用 Llama-2-7B 和 Llama-2-13B 模型进行测试
  • 如图 10 所示
    • 随着缓存大小（Attention Window Size）的增加 ，StreamingLLM 的解码速度呈线性增长
    • 而带重计算的滑动窗口基线的解码延迟呈二次方增长
    • StreamingLLM 实现了令人印象深刻的加速，每个 Token 的加速比高达 \(22.2\times\)
    • 且 StreamingLLM 仍保持了与重计算基线一致的内存占用
  • 注意：这里仅仅考虑效率，具体模型性能指标见前面的其他图

补充：Related Work

  • 关于将 LLM 应用于长文本已经进行了广泛的研究，主要集中在三个领域：
    • 长度外推（Length Extrapolation）
    • 上下文窗口扩展（Context Window Extension）
    • 改进 LLM 对长文本的利用（Improving LLMs’ Utilization of Long Text）
  • 虽然看似相关，但值得注意的是，一个方向的进展并不一定导致另一个方向的进展
    • 例如，扩展 LLM 的上下文大小并不能提高模型在上下文大小之外的性能，而且这两种方法都不能确保有效利用长上下文
    • 论文的 StreamingLLM 框架主要属于第一类（长度外推），即 LLM 被应用于显著超过预训练窗口大小的文本，甚至可能是无限长度
    • 论文不扩展 LLM 的注意力窗口大小，也不增强模型对长文本的记忆和使用能力
    • 后两个类别与论文的重点正交，并且可以与论文的技术结合
  • 长度外推（第一类）旨在使在较短文本上训练的语言模型能够在测试时处理较长的文本
    • 一个主要的研究方向是针对 Transformer 模型开发相对位置编码方法，使其能够在训练窗口之外运行
    • 其中一项工作是 Rotary Position Embeddings (RoPE) (2021)，它在每个注意力层中转换查询和键以整合相对位置信息
      • 后续研究 (2022; 2023) 表明其在超过训练窗口的文本上表现不佳
    • 另一种方法 ALiBi (2022) 根据查询和键之间的距离对注意力分数进行偏置，从而引入相对位置信息
      • 虽然这显示出改进的外推能力，但论文在 MPT 模型上的测试突显了当文本长度远大于训练长度时会出现崩溃
    • 当前的方法尚未实现无限长度外推，导致没有现有的 LLM 适合流式应用
  • 上下文窗口扩展（第二类）侧重于扩展 LLM 的上下文窗口，使其能够在一个前向传递中处理更多 Token
    • 一条主要的工作线解决了训练效率问题
      • 考虑到训练期间注意力计算的二次复杂度，开发长上下文 LLM 既是计算挑战也是内存挑战
      • 解决方案范围从系统优化的 FlashAttention (2022; Dao, 2023)（加速注意力计算并减少内存占用）到近似注意力（Approximative attention）方法 (2020a; 2020; 2020; 2020)，这些方法以模型质量换取效率
    • 最近，关于使用 RoPE 扩展预训练 LLM 的工作激增 (2023;)，涉及位置插值和微调
    • 但所有上述技术仅将 LLM 的上下文窗口扩展到有限的程度 ，这未能达到论文处理无限输入的主要关注点
  • 改进 LLM 对长文本的利用（第三类）优化 LLM 以更好地捕获和使用上下文中的内容，而不是仅仅将它们作为输入
    • 正如 (2023) 和 (2023) 所强调的，前述两个方向的成功并不一定能转化为对长上下文的胜任利用
    • 解决 LLM 内部对长上下文的有效使用仍然是一个挑战
    • 论文的工作集中于稳定地利用最近 Token ，实现 LLM 的无缝流式应用

附录 A：Discussions

  • 应用 (Applications)
    • StreamingLLM 特别适合流式应用，例如多轮对话，其中持续运行而不严重依赖大量内存或历史数据至关重要
      • 例如，在 LLM-based 日常助手应用中，StreamingLLM 使模型能够在较长时间内无缝运行
    • 它基于最近的交互生成响应，从而避免了频繁刷新缓存的需要
    • 传统方法可能需要在对话长度超过训练长度时重置缓存，导致丢失最近的上下文，或者可能需要根据最近的文本历史重新计算键值状态，这可能效率低下
  • 局限性 (Limitations)
    • 虽然 StreamingLLM 提高了 LLM 在流式上下文中的效率，但它并没有扩展模型的上下文窗口或增强其长期记忆能力
    • 如章节 C 中详述，模型仅限于在其当前缓存的范围内运行
    • StreamingLLM 不适合需要长期记忆和广泛数据依赖性的任务，例如长文档问答和摘要
    • 但它在仅需要短期记忆的场景中表现出色 ，例如日常对话和短文档问答，其优势在于能够根据最近的上下文生成连贯的文本，而无需刷新缓存
  • 更广泛的社会影响 (Broader Societal Impacts)
    • StreamingLLM 显著提高了 LLM 的效率和可访问性，使其在各个部门的使用民主化
      • 通过在对话代理等应用中实现不间断的快速交互，StreamingLLM 改善了用户体验，尤其是在需要固定长度模型的场景中
      • 这一进步使得对话更加无缝和具有上下文感知能力，可能惠及教育、医疗保健和客户服务等行业
    • StreamingLLM 在处理过程中的效率降低了计算负载，符合对环境可持续 AI 技术的需求
      • 这一方面对于在技术资源有限的地区推广先进的 AI 工具至关重要
    • 但 StreamingLLM 的潜在负面影响与通用语言模型相关的风险类似，例如错误信息和生成有偏见内容的风险
      • 必须通过强有力的道德准则和保障措施来解决这些风险
    • 虽然 StreamingLLM 具有语言模型共有的一些风险，但其在提升用户体验、 democratizing AI 访问和促进可持续性方面的积极贡献是值得注意的
      • 这些好处强调了负责任地部署和合乎道德地使用该技术的重要性
  • 稀疏 Transformer (Sparse Transformers)
    • 关于高效 Transformer 模型的文献主要集中于降低自注意力机制的计算和内存复杂性
      • 一项相关的工作是通过将注意力范围限制在固定的、预定义的模式来稀疏化注意力矩阵，例如局部窗口或固定步长的块模式 (2022)
      • Sparse Transformer (2019) 引入了注意力矩阵的稀疏分解，将注意力的计算复杂度降低到 \(O(n\sqrt{n})\)
      • LongFormer (2020) 将扩张的局部 Window Attention 与任务驱动的全局注意力相结合
      • Extended Transformer Construction (ETC) Ainslie 等 (2020) 提出了一种新颖的全局-局部注意力机制，包含四种注意力模式：全局到全局、局部到局部、局部到全局和全局到局部
      • 基于 ETC，BigBird (2020a) 提出了另一种线性复杂度的注意力替代方案，利用全局 Token、局部滑动 Window Attention 和随机注意力
    • 但这些方法有几个局限性
      • 一：Sparse Transformer 和 ETC 需要为特定的块稀疏矩阵乘法变体定制 GPU 内核
      • 二：LongFormer、ETC 和 BigBird 都依赖于全局注意力模式，这不适合自回归语言模型
      • 三：这些方法与预训练模型不兼容，需要从头开始重新训练
    • 相比之下，论文的方法使用标准的 GPU 内核易于实现，并且与使用 Dense Attention 的预训练自回归语言模型兼容，这些模型在 NLP 社区中普遍存在
      • 这种兼容性提供了显著的优势，允许利用现有的预训练模型而无需任何微调
  • 同期工作 (Concurrent Works)
    • 论文的研究与 Han 等人的工作同时进行，他们对语言模型长度泛化失败进行了理论研究，确定了三个分布外因素
      • 受此分析启发，他们的方法采用“\(\Lambda\)”形注意力模式并重新配置位置编码距离以增强 LLM 中的长度泛化
      • 这种方法与论文的方法有相似之处
      • 但论文的工作揭示了“Attention Sink”现象，即 Transformer 模型倾向于将高注意力分数分配给语义较小的初始 Token
        • 这一现象超出了长度泛化失败的范围，表明 Transformer 模型中存在一个更普遍的问题
      • 论文不仅在自回归语言模型中观察到这种“Attention Sink”行为，而且在编码器 Transformer（如 BERT，见章节 H）和视觉 Transformer (ViTs) Darcet 等 (2023) 中也观察到，表明其在 Transformer 架构中更广泛地存在
      • 为了缓解“Attention Sink”现象，论文建议在预训练期间引入一个可学习的 Sink Token ，并通过广泛的消融研究支持论文的发现
    • 与此同时，Darcet 等人在视觉 Transformer 中观察到类似的注意力集中在随机背景 patch Token 上的现象，称为“寄存器（registers）”
      • 这些寄存器充当全局图像信息的存储库
      • 他们的解决方案是添加专用的“寄存器” Token ，旨在平衡注意力分布
      • “Attention Sink”与此概念类似
      • 在论文的论文中，“Attention Sink” 是初始 Token ，不成比例地吸引后续 Token 的注意力
      • 在预训练期间引入专用的 Sink Token 可以防止模型不适当地使用内容 Token 作为 Attention Sink ，从而实现更有效的注意力分布
      • 但存在一个关键区别：视觉 Transformer 中的“寄存器”在中间层充当全局信息持有者，而论文的“Attention Sink”在自回归模型中作为初始 Token 定位
      • 这种位置差异表明，注意力计算中的 softmax 函数可能在 Attention Sink 的出现中扮演更基本的角色

附录 C：Accuracy on StreamEval with Increasing Query-Answer Line Distance（行距增加时的精确率）

  • 为了评估 StreamingLLM 对扩展输入的处理能力，论文在 StreamEval 上评估了 Llama-2-7B-32K-Instruct 模型，重点关注不同缓存配置下不同的查询-答案行距
    • 在 StreamEval 中，每行包含 23 个 Token ，使得行距相当于 Token 距离的 \(23\times\) 行距
    • 准确率是通过对 100 个样本的结果取平均值计算的，每个样本包含 100 个查询
  • 表 7 说明
    • 当查询和答案之间的 Token 距离在缓存大小之内时，StreamingLLM 保持准确率
    • 但随着该距离增加，准确率会降低，并在最终超过缓存容量时降至零
  • 这些结果表明，虽然 StreamingLLM 在基于最近上下文生成连贯文本方面是有效的，但它不能扩展语言模型的上下文长度
    • 这些结果也强调了当前语言模型中一个更广泛的挑战：它们无法充分利用缓存中的上下文信息，这一发现与 Liu 等人的观察结果一致

附录 D：Long-Range Benchmark Evaluation

  • 论文使用 Llama-2-7B-chat 模型（最大上下文长度 4k）在 Long-Bench Bai 等 (2023) 上评估了 StreamingLLM，该基准包含三个关键 NLP 任务：
    • 单文档问答 NarrativeQA Kocisky 等 (2017) 和 Qasper Dasigi 等 (2021)
    • 多文档问答 HotpotQA Yang 等 (2018) 和 2WikiMQA Ho 等 (2020)
    • 摘要 GovReport Huang 等 (2021), MultiNews Fabbri 等 (2019)
  • LongBench 为 Llama-2-7B-chat 模型设置了默认的最大序列长度 3,500 个 Token ，从中间截断以保留开头和结尾信息（各 1,750 个 Token ）
  • 表 8 显示，使用 4+3496 缓存配置的 StreamingLLM 表现不如 truncation 基线，这可能是由于丢失了关键的初始输入提示信息
    • 但将 Attention Sink 数量调整为 1750 可以将性能恢复到文本截断基线的水平
    • 这些结果证实了章节 C 中的发现，表明 StreamingLLM 的有效性取决于其缓存中的信息，其缓存内性能与文本截断基线相当
  • 问题：这里的 truncation 基线是指直接保留 前后 1750 个 Token 吗？
  • 回答：是的，与 StreamingLLM 1750+1750 的最大区别在于，StreamingLLM 1750+1750 的位置信息是缓存窗口内部的，不是真实文本中的
    • 从表 8 中可知，两者的模型效果差不多

附录 E：在较长序列上 Llama-2-7B 的注意力可视化 (Llama-2-7B Attention Visualization on Longer Sequences)

  • 图 2 使用短序列（长度为 16）可视化了 Llama-2-7B 的注意力图，以便清晰展示
  • 论文在图 11 中进一步可视化了 Llama-2-7B 在较长序列（长度为 128）上的注意力
  • 论文发现短序列上的观察结果在较长序列上也成立
    • 即在大多数层中，无论初始 Token 与序列中其余 Token 之间的距离如何，初始 Token 的注意力分数远高于序列中其余 Token 的注意力分数
  • 因为序列越长， Attention Sink 的分数在热力图上的显示就越细
  • 论文在章节 F 中使用不同的方法进一步分析了较长序列（长度为 4096）上的注意力分布
  • 补充观察：从图上看，仍然是输入的浅层（低层）上关注局部注意力，深层关注 Sink Token

附录 F：Qualitative Analysis of Attention Sinks in Long Inputs

  • 图 2 和图 13 使用短序列说明了 Attention Sink 现象以便清晰展示
  • 扩展此分析，图 12 展示了在长输入（序列长度为 4096）中指向第一个 Token 的注意力分数（经过 SoftMax 后）的分布
  • 论文对 256 个序列的注意力分数取平均值，每个序列包含 4096 个 Token ，绘制数据表示第 4096 个 Token 在每个层中对初始 Token 的注意力分配
  • 第一个 Token 的注意力分数显著高，通常超过总注意力的一半，除了最底部的两个层（最浅的两层）
  • 这一观察经验性地证实了大多数层和头对第一个 Token 的偏好关注，无论序列中其他 Token 的距离如何
  • 这种趋势强调了序列中初始 Token 的关键作用，因为移除它们会由于 SoftMax 函数分母的大部分被移除而对语言模型性能产生巨大影响

附录 G：Llama-2-70B 注意力可视化 (Llama-2-70B Attention Visualization)

  • 图 2 展示了 Llama-2-7B 的注意力可视化，论文在图 13 中进一步可视化了 Llama-2-70B 的注意力
  • 论文发现对 Llama-2-7B 的观察结果在 Llama-2-70B 上也成立，
  • 其中在大多数层中，初始 Token 的注意力分数远高于其余 Token 的注意力分数

附录 H：Attention Sinks in Encoder Transformers

  • 在论文中，论文主要探讨了在自回归、 Decoder-only 语言模型（如 GPT 和 Llama）中观察到的 Attention Sink 现象
    • 基于章节 3.1 的见解，论文提出这一现象可能扩展到其他 Transformer 架构，包括编码器模型，如 BERT Devlin 等 (2019) 和 ViT Dosovitskiy 等 (2021)
  • 这一假设源于这些模型共享相似的 Transformer 结构并使用 SoftMax 注意力机制
    • 为了证实论文的假设，论文分析了 BERT-base-uncased 的注意力模式
  • 如图 14 所示
    • BERT-base-uncased 表现出 Attention Sink 现象，其特征是在大多数层中分配给 [SEP] Token 的注意力分数不成比例地高
      • 这表明模型始终依赖无处不在的 [SEP] Token 作为注意力的焦点
    • Darcet 等人的同期研究在视觉 Transformer 中识别出类似的注意力尖峰，归因于随机背景补丁 Token 充当全局图像信息的“寄存器”
    • 作者认为这些“寄存器”类似于论文观察到的 Attention Sink 现象，表明这是所有 Transformer 模型的普遍特征

附录 I：Using More Sink Tokens in the Pre-Training Stage

  • 章节 3.3 说明，在预训练阶段加入单个专用的 Sink Token 不会影响模型性能，可通过将 Attention Sink 集中到一个 Token 上来增强流式性能
    • 本节深入探讨在预训练期间添加额外的 Sink Token 是否能够进一步优化预训练语言模型的性能
  • 如图 15 所示，论文的实验表明，在预训练期间加入一个或两个 Sink Token，预训练损失曲线与基线（原始）模型非常相似
  • 但如表 9 详述，引入第二个 Sink Token 在大多数基准任务中并未产生实质性的性能改进
  • 进一步分析，如表 10 所示，显示包含额外的 Sink Token 并不会增强流式性能（理解：这里的额外指的主要是多余一个的部分？）
    • 模型似乎依赖两个 Sink Token 来维持稳定的流式性能
    • 这些发现表明，单个 Sink Token 足以改善流式性能，添加更多 Sink Token 并不会带来整体语言模型性能的进一步提升
    • 这与视觉 Transformer (ViT) Darcet 等 (2023) 中的发现形成对比，在 ViT 中发现多个“寄存器”是有益的
  • 表 10: 预训练期间添加零 Token 和可学习 Sink Token 与原始注意力的比较
    • 缓存配置 \(x\)+\(y\) 表示添加 \(x\) 个初始 Token 和 \(y\) 个最近 Token
    • 困惑度在 PG19 测试集的第一个样本上评估

整体介绍

  • DeepSpeed 是由微软开发的开源的深度学习优化库，旨在提高大规模模型训练的效率和可扩展性
  • DeepSpeed 框架的核心技术有：
    • ZeRO 冗余优化技术 ：通过分布式内存管理，将模型参数、梯度和优化器状态进行分区，大幅降低显存占用，是首次支持千亿级参数模型训练的框架
    • 3D 并行策略 ：支持数据并行、流水线并行和张量切片模型并行，并可灵活组合
    • 混合精度训练 ：自动混合精度训练（AMP）将单精度和半精度浮点数结合使用，降低内存需求的同时提升计算效率
    • 智能推理优化器 ：支持张量并行与异构内存技术，提供低延迟高吞吐的分布式推理服务，可将成本降低 70%
    • 全链路内存管理 ：集成 CPU 卸载与显存碎片整理技术，单卡即可训练百亿级模型，资源利用率提升 6 倍
  • DeepSpeed 框架的组件构成有：
    • Apis ：提供易用的 API 接口
    • 运行时组件 ：管理、执行和性能优化，基于 Python 语言实现，负责部署训练任务到分布式设备、数据分区、模型分区等
    • 底层内核 ：用 C++ 和 CUDA 实现，优化计算和通信
  • DeepSpeed 生态兼容性极好 ：原生兼容 PyTorch 与 Hugging Face 生态，通过简洁 API 可快速迁移项目，开发效率提升 300%

安装 DeepSpeed

  • 仅需一行安装命令即可：

    1
    pip install deepspeed

  • 暗转完成后，可以使用 ds_report 命令验证安装是否成功

    • 这个命令可以查看环境配置信息

DeepSpeed 的使用

  • 代码修改 ：仅仅需要非常少的代码修改即可将原始训练代码切换到 DeepSpeed 框架上（DeepSpeed 与 PyTorch 无缝集成，只需少量修改即可启用加速）

    • 第一步：通过 deepspeed.initialize 将模型包装为 DeepSpeed 引擎，自动应用优化
    • 第二步：使用 model_engine.backward 和 model_engine.step 替换PyTorch原本的 loss.backward() 和 optimizer.step()

  • 配置 DeepSpeed ：DeepSpeed 的优化行为通过 JSON 配置文件（ds_config.json）指定

    • 一般的项目都会自带一些配置好的 .json 文件示例，可直接修改使用

  • 运行训练 ：使用 DeepSpeed 的命令行工具启动训练

    • 单节点训练命令为：

      1
      deepspeed train.py --deepspeed_config ds_config.json

    • 对于多节点集群，使用下面的命令：

      1
      deepspeed --num_gpus 8 --num_nodes 2 train.py --deepspeed_config ds_config.json
      • 其中--num_gpus 指定每节点使用的 GPU 数量，--num_nodes 指定集群中的节点数

    • 其他常用参数：

      • 使用--log_dir 参数启用日志记录，监控内存使用、训练速度等

  • 特别说明：使用了 DeepSpeed 框架的代码需要使用 deepspeed 命令来启动

    • 补充：直接使用 python 命令启动时会出现 deepspeed.initialize 调用的错误
    • 原因：DeepSpeed 作为一个分布式训练库，需要特殊的启动器来管理多个进程和 GPU 之间的通信和资源分配

DeepSpeed 使用示例（基于 PyTorch）

  • DeepSpeed 使用示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader, Dataset
    import deepspeed

    class DiyModel(nn.Module):
        def __init__(self, input_dim, output_dim):
            super(DiyModel, self).__init__()
            self.fc = nn.Sequential(
                nn.Linear(input_dim, 128),
                nn.ReLU(),
                nn.Linear(128, output_dim)
            )
        
        def forward(self, x):
            return self.fc(x)

    class RandomDataset(Dataset):
        def __len__(self):
            return 1000
        
        def __getitem__(self, idx):
            x = torch.randn(32)
            y = torch.randint(0, 10, (1,)).item()
            return x, y

    # 初始化模型和数据（无需为 DeepSpeed 特别处理）
    model = DiyModel(32, 10)
    train_dataset = RandomDataset()
    train_loader = DataLoader(train_dataset, batch_size=32)

    # 初始化 DeepSpeed 引擎（仅需使用 deepspeed.initialize 初始化模型得到 model_engine 即可）
    # 注：这一步会自动将模型参数移动到 GPU 上，下面使用的数据也需要将数据移动到对应的 GPU 上才能运行，否则会出现 设备不一致的错误（CPU vs GPU）
    # zero_optimization 字段的 stage 键值对应如下效果
    ## stage: 0：不使用 ZeRO 优化（默认值是 0）
    ## stage: 1：优化器状态分片
    ## stage: 2：优化器状态和梯度分片
    ## stage: 3：优化器状态、梯度和参数分片（最高内存效率）
    model_engine, optimizer, _, _ = deepspeed.initialize(
        args=None,
        model=model,
        model_parameters=model.parameters(),
        config={
            "train_batch_size": 32,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 0.001,
                    "betas": [0.9, 0.999]
                }
            },
            "fp16": {
                "enabled": True,
                "loss_scale": 0,
                "loss_scale_window": 1000,
                "initial_scale_power": 16
            },
            "zero_optimization": {
                "stage": 2,  # 这里指定Zero层级（0、1、2、3）
                "offload_optimizer": {
                "device": "cpu"  # 可选：指定优化器卸载设备
                }
            }
        }
    )

    # 训练过程（训练时不再使用原来的模型，使用 model_engine）
    for epoch in range(10):
        for batch_idx, (data, target) in enumerate(train_loader):
            # 将数据挪到和模型相同的 GPU 上
            device = model_engine.device
            model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取dtype
            data = data.to(device, dtype=model_dtype) # 将数据移至模型所在设备并转换为与模型相同的dtype
            target = target.to(device)

            outputs = model_engine(data) # 使用（deepspeed.initialize 初始化得到的）model_engine 来进行前向过程
            loss = nn.CrossEntropyLoss()(outputs, target) 
            model_engine.backward(loss) # 使用 model_engine 来进行后向过程
            model_engine.step() # 使用 model_engine 来更新模型参数（注：此时不再需要显示调用 optimizer)
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item()}")

附录：如何指定目标 GPU

启动时指定 CUDA_VISIBLE_DEVICES 环境变量

  • 在运行命令前设置环境变量来限制 DeepSpeed 可见的 GPU：

    1
    CUDA_VISIBLE_DEVICES=0,1,2,3 deepspeed your_script.py --args ...

  • 该方案是最常用的方法 ，且适用于常见的很多框架

在 Python 脚本中设置环境变量

  • 也可以在Python脚本中通过os.environ设置这个环境变量：

    1
    2
    3
    4
    5
    6
    7
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  # 指定使用GPU 0和1

    import deepspeed
    import torch

    # 训练逻辑

  • 注：需要在导入DeepSpeed或PyTorch之前设置

  • 该方案同样适用于常见的很多框架，但因为需要修改代码，不常用

使用 deepspeed 命令的 --include 参数

  • 如果使用的是 DeepSpeed 的 launcher，也可以通过--include参数指定使用的 GPU： 
    1
    deepspeed --include localhost:0,1 your_script.py --args ...

多机多卡如何指定 GPU

  • 如果你使用的是单机多卡，以上方法都能很好地工作
  • 对于多机多卡训练，通常需要结合 deepspeed 命令的其他参数如 --hostfile 等一起使用

附录：数据加载位置管理

  • 由于 deepspeed 会自动将模型参数加载到指定 GPU 上，所以数据也要加载到指定 GPU，否则会出现设备不一致的错误
  • 加载命令如下（亲测解决方案）： 
    1
    2
    device = model_engine.device
    data = data.to(device)

附录：混合精度训练数据格式管理

  • 由于 deepspeed 在启动混合精度训练时，可能会按照指定格式来指定参数形式，此时数据也需要转换为指定类型
  • 解决方式如下（亲测遇到错误时的解决方案）： 
    1
    2
    model_dtype = next(model_engine.parameters()).dtype # 通过模型的第一个参数获取类型，注：写这么复杂的原因是，当前还不支持直接调用 model_engine.dtype()
    data = data.to(device, dtype=model_dtype)

本文主要介绍 DeepSeek-R1 相关的解读，笼统而简单的介绍，详情可查看本人的其他博客

Background

2025年01月20日，deepseek 正式发布 DeepSeek-R1，并同步开源模型权重

  • 开源 DeepSeek-R1 推理大模型，与 o1 性能相近。‍‍（冷启动 SFT -> RL -> COT + 通用数据SFT（80W）-> 全场景RL）
  • 开源 DeepSeek-R1-Zero，预训练模型直接 RL，不走 SFT。（纯强化学习）
  • 开源用 R1 数据蒸馏的 Qwen、Llama 系列小模型，蒸馏模型超过 o1-mini 和 QWQ。（直接使用80W数据进行SFT）

报告核心说明

  • 首次验证了纯 RL 也可以训练出大模型的推理能力
  • aha moment：顿悟时刻，主要指DeepSeek-R1-Zero模型训练过程中，模型在某个关键时刻突然学会自我反省的情况

DeepSeek-R1-Zero

  • 预训练后直接进入RL阶段

DeepSeek-R1-Zero奖励模型

  • 直接采用了一种基于规则的奖励系统，包括两种奖励模型作为评估指标，分别是准确率奖励模型和格式奖励模型
    • 准确率奖励模型 ：评估response是否准确
    • 格式奖励模型 ：评估格式是否准确，具体来说，格式奖励要求模型将思考过程放在“和”标签之间

DeepSeek-R1-Zero演化过程

DeepSeek-R1

DeepSeek-R1 使用了冷启动 + 多阶段训练的方式：

  • 阶段1：使用少量高质量的 CoT 数据进行冷启动，预热模型。（相较于直接RL，冷启动预热能让模型快速进入稳定训练阶段）
  • 阶段2：进行面向推理的强化学习，提升模型在推理任务上的性能
  • 阶段3：使用拒绝采样和监督微调，进一步提升模型的综合能力
  • 阶段4：再次进行强化学习，使模型在所有场景下都表现良好

DeepSeek-R1之MoE

普通的MoE

  • Mixture of Experts，混合专家模型。最早1991年的论文《Adaptive Mixtures of Local Experts》中提出了混合专家模型的雏形，架构图如下：

Switch Transformer中的MoE

DeepSeekMoE（DeepSeek-V1）

  • 原始论文：DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

  • 传统的Transformer和MoE

  • 问题：上图中说 \(\boldsymbol{e}_i^l\) 是每个专家的质心(Centroid)，但是未给出这个质心是怎么来的，是否可训练？

    • 回答：在上述原始论文中确实没有说清楚，但是一些文章中有提到，比如 GSHARD: SCALING GIANT MODELS WITH CONDITIONAL COMPUTATION AND AUTOMATIC SHARDING 中 Algorithm1 提到输入门控网络的该值是可训练的权重，上图中的乘法实际上也就是一个线性层，实现可以如下：
      1
      2
      3
      4
      5
      6
      7
      8
      9
      10
      11
      12
      13
      def __init__():
      	...
      	# 门控网络
      	self.gate = nn.Linear(input_dim, num_experts)

      def forward(self, x):
      	# 计算门控权重
      	gate_logits = self.gate(x)
      	gate_probs = F.softmax(gate_logits, dim=-1)

      	# 选择top-k专家
      	topk_values, topk_indices = torch.topk(gate_probs, self.k)
      	topk_gates = topk_values / topk_values.sum(dim=-1, keepdim=True)

  • 个人思考：以上被选中的FNN系数和不为1（小于1），但因为在每一个Transformer层中，FNN的结果和上一层的隐向量叠加后，都有LayerNorm存在（将每个token的隐向量分别归一化为均值为0，方差为1的向量），所以隐向量的值不会越来越小，在V3版本的公式中，会考虑在选择了 TopK FNN 后，再进行一次归一化

  • 改进一：Fine-Grained Expert Segmentation，更精细化的专家拆分

    • DeepSeekMoE的精细化MoE：可以看出DeepSeekMoE中将原始的 \(N\) 个FNN扩展为 \(mN\) 个（注意：只是拆分的更细，参数总量是相同的），选择的FNN数量也从 \(K\) 个扩展到 \(mK\) 个

  • 改进二：Shared Expert Isolation，独立的共享专家

    • 使用了 \(K_s\) 个固定的共享专家，需要路由的专家数量为 \(mN - K_s\)
    • 最终的MOE层输出由3部分组成，共享专家的输出结果 + Top_K个路由专家输出结果 + 残差

  • 改进三：Load Balance Consideration，负载均衡考量

    • Expert-Level Balance Loss（专家级别的负载均衡损失函数）：
    • Device-Level Balance Loss（设备级别的负载均衡损失函数）：

  • DeepSeekMoE结构图：

  • 其他说明：DeepSeek-R1共61个Transformer层，其中前三个层是正常的FNN层，后面的4-61层均用MoE取代FNN层

DeepSeek-V2

DeepSeek-V3

  • 参考链接：DeepSeek-V3 Technical Report

  • 门控函数从SoftMax优化为了Sigmoid

    • 个人理解：为什么用Sigmoid更好？因为本来选择了topK就还需要再做一次归一化的(这次归一化是直按线性权重分配，不使用Softmax)，使用Sigmoid速度更快，不影响选择topK且归一化后效果一致？

DeepSeek-R1之MLA

DeepSeek-R1之GRPO

  • 原始论文：DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
  • 核心思路：使用多次采样的归一化结果作为reward，放弃Critic Model，减少显存使用
  • 一个有趣的对比：DeepSeek GRPO在简单控制系统上和PPO的对比 - 王兴兴的文章 - 知乎

    • 对于整个系统中间过程和信息，比较清晰的问题（中间过程能被价值评价清晰），比如类似上面的控制系统（或者其他机器人系统），PPO还是最简单粗暴出效果很好的；但对于像DeepSeek用来搞数学RL推理，由于中间过程没法很好的描述和计算中间过程的价值，确实还是GRPO更快更方便（只看最终结果）；

  • 注意，GRPO中使用了KL散度的近似形式，Approximating KL Divergence —— 来自：Deepseek的RL算法GRPO解读 - AIQL的文章 - 知乎
    • 估计形式为（注意以下式子中右边是左边的无偏梯度的前提是 \(o_{i,t} \sim \pi_\theta(\cdot \vert q,\mathbf{o}_{i,<t})\) ）：
      $$
      \mathbb{D}_\text{KL}[\pi_\theta\Vert\pi_{\text{ref}}] \approx \frac{\pi_{\text{ref}}(o_{i,t}\vert q,\mathbf{o}_{i,<t})}{\pi_\theta(o_{i,t}\vert q,\mathbf{o}_{i,<t})} - \log \frac{\pi_{\text{ref}}(o_{i,t}\vert q,\mathbf{o}_{i,<t})}{\pi_\theta(o_{i,t}\vert q,\mathbf{o}_{i,<t})} - 1, \quad o_{i,t} \sim \pi_\theta(\cdot \vert q,\mathbf{o}_{i,<t})
      $$
      • 直观理解：上面的式子右边满足KL散度的基本特性
        • 当两个分布足够接近时，第一项趋近于1，第二项趋近于0，整体趋近于0；
        • 两个分布不相等时，上式右边取值总是大于0，可以通过求导证明：当 \(x>0\) 时，有 \(x - \log x - 1 \ge 0\)
  • 其他团队对GRPO的改进：阶跃&清华新论文：DeepSeek-R1的GRPO 可以更简洁 - 机器之心的文章 - 知乎

    阶跃星辰与清华大学近期的一项研究发现，只需使用带 GAE （λ= 1，γ= 1）的普通 PPO 以及基于规则的简单奖励函数，无需任何 KL 正则化，就足以扩展在推理任务上的响应长度和基准性能，类似于在 DeepSeek-R1-Zero 上观察到的现象
    使用这种极简方法，他们打造了 Open-Reasoner-Zero，这是首个面向大规模推理的强化学习训练的开源实现。并且该实现在 GPQA Diamond 基准上的表现优于 DeepSeek-R1-Zero-Qwen-32B，同时仅需使用 1/30 的训练步数。需要强调，该团队不仅开源了代码，还发布了参数设置、训练数据和模型权重

DeepSeek-R1之MTP

  • 参考链接：deepseek技术解读(2)-MTP（Multi-Token Prediction）的前世今生 - 姜富春的文章 - 知乎
  • 基本思想：
    • 预测阶段（Predict） ：通过 K 个头一次生成 K 个 token 的预测
    • 验证阶段（Verify） ：将 K 个 token 组装成 K 个 <input,label> 对，并行地利用输入 Main Model 作为评估验证，如果输出 label 与 Main Model 一致，则接受该 token
    • 接受阶段（Accept） ：最终接受满足 Main Model 的最大长度 tokens 作为输出

Deepseek MTP实现细节

  • 原始报告：DeepSeek-V3 Technical Report内容如下

  • 问题：上面图中设计的MTP中，无法做到整整的并行，比如，仅知道 \(t_1\) 时，只能预测得到 \(t_2\) ，无法得到 \(t_3\) ，因为在任意一个Module中， \(t_2\) 都依赖着 \(t_3\) 作为输入（只是输入后不用再过 \(L\) 层Transformer Block了，仅过一层就行）

  • 训练时：使用多个MTP Module，综合大家的损失共同更新梯度
    $$\mathcal{L}_\text{MTP} = \frac{\lambda}{D}\sum_{k=1}^D \mathcal{L}_\text{MTP}^k$$

  • 推断时：Deepseek直接丢弃了MTP Module，仅使用第一个（相当于跟普通不使用MTP的时候一致，只是吃到了训练的红利），部分文章中提到最多使用2个

    Our MTP strategy mainly aims to improve the performance of the main model, so during inference, we can directly discard the MTP modules and the main model can function independently and normally. Additionally, we can also repurpose these MTP modules for speculative decoding to further improve the generation latency.

DeepSeek-R1 API远程调用

附录：Sparse MoE实现Demo

  • Sparse MoE 的简单实现示例 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    import torch
    import torch.nn as nn

    class Expert(nn.Module):
    	def __init__(self, input_size, output_size):
    		super(Expert, self).__init__()
    		self.fc = nn.Linear(input_size, output_size)

    	def forward(self, x):
    		return self.fc(x)

    # Dense MoE实现，对每个Token，所有专家都参与计算
    class MoE(nn.Module):
    	def __init__(self, num_experts, input_size, output_size):
    		super(MoE, self).__init__()
    		self.experts = nn.ModuleList([Expert(input_size, output_size) for _ in range(num_experts)])
    		self.gate = nn.Linear(input_size, num_experts)

    	def forward(self, x):
    		gate_scores = torch.softmax(self.gate(x), dim=1)
    		expert_outputs = [expert(x) for expert in self.experts]
    		expert_outputs = torch.stack(expert_outputs, dim=1)
    		output = torch.sum(gate_scores.unsqueeze(-1) * expert_outputs, dim=1)
    		return output

    # SparseMoE实现，对每个Token，仅少量专家参与计算
    class SparseMoE(nn.Module):
    	def __init__(self, num_experts, input_size, output_size, k=2):
    		super(SparseMoE, self).__init__()
    		self.experts = nn.ModuleList([Expert(input_size, output_size) for _ in range(num_experts)])
    		self.gate = nn.Linear(input_size, num_experts)
    		self.k = k

    	def forward(self, x):
    		gate_scores = self.gate(x)
    		topk_scores, topk_indices = torch.topk(gate_scores, k=self.k, dim=1)

    		batch_size = x.size(0)
    		expert_outputs = []
    		for b in range(batch_size):
    			# 获取当前样本选中的K个专家的输出
    			selected_outputs = [self.experts[idx](x[b].unsqueeze(0)) for idx in topk_indices[b]]
    			# 堆叠输出，维度为 (1, k, output_size)
    			selected_outputs = torch.stack(selected_outputs, dim=1)
    			expert_outputs.append(selected_outputs)

    		# 合并所有样本，维度为 (batch_size, k, output_size)
    		expert_outputs = torch.cat(expert_outputs, dim=0)

    		# 计算选中专家的权重 (batch_size, k)
    		weights = torch.softmax(topk_scores, dim=1)

    		# 加权求和，维度为 (batch_size, output_size)
    		output = torch.sum(weights.unsqueeze(-1) * expert_outputs, dim=1)
    		return output

    # 使用Demo
    num_experts = 4
    input_size = 10
    output_size = 5
    batch_size = 32
    moe = SparseMoE(num_experts, input_size, output_size)
    input_data = torch.randn(batch_size, input_size)
    output = moe(input_data)
    print(output.shape)

    # torch.Size([32, 5])

本文主要介绍 LLM 的 Tokenization（一些研究文献也叫作 Tokenize）

Tokenization 技术分类

  • 分词（Tokenization），有些地方也称为词元化，用于将文本进行切分，以便于语言模型能够理解和使用
  • Tokenization 技术有三种粒度：Word-level Tokenization、Character-level Tokenization 和 Sub-word Tokenization

Word-level Tokenization（Word 粒度）

  • 按照单词为最小单位进行分词来进行分词，每个单词是一个 token
  • 优点：
    • 每个 token 都是完整的单词，能准确表达语义（理解：模型学到的embedding就是当前单词的），便于理解和处理
  • 缺点：
    • 词表一般会很大，比如英文中各种词语变化非常大，词表动辄超过 10W+
    • 容易遇到 OOV（未登录词）问题
    • 对新词语或者拼写错误词语敏感，比如“MAGA”等近期新出的词语就无法识别到

Character-level Tokenization（Char 粒度）

  • 按照字符来进行分词，每个字符是一个 token
  • 优点：
    • 无 OOV 问题，且任何的新词都能被拆解到字符粒度
    • 词表固定，且词表很小（就几十个很少的字符集合），占用存储少
  • 缺点：
    • 相同句子长度下 token 太多，导致模型训练和推理慢
    • 单个字符无法表达语义，学出来的 embedding 也难以对应到语义上，比如 “Queen” 和 “King” 就很难被联系起来

Sub-word Tokenization（Subword 粒度）

  • 介于 Word-level 和 Character-level 之间，可以灵活地将一个单词拆分成一个或多个有语义信息的部分，比如 “tokenization” 可以拆为 “token”+”ization”，而 “love” 则不需要拆分，能够在减少 OOV 和 表达语义之间做 Trade-off
  • Sub-word Tokenization 包括 BPE，WordPiece，Unigram 等
    • BPE（Byte-Pair Encoding） ，即字节对编码，是一种 Tokenize 算法。其核心思想在于将最常出现的子词对合并，直到词汇表达到预定的大小时停止（实际上该方法本质是一种压缩算法，最早1994年提出，被用于通用数据压缩）
      • 特别说明 ：Byte-Pair Encoding 中的 Byte 不是字节的意思，可以理解为编码的最小单位，一般是字符级别（这里很容易误解）
      • BPE 的缺点 ：由于 BPE 是字符级别，直接使用 BPE 可能导致词表过大
        • 理解：一个字符可能包含很多个字节，比如 UTF-8（Unicode编码的一种）是一种变长的编码方式，它可以使用 1 到 4 个字节来表示一个字符（理论上 4 个字节可以表示所有语言的所有字符）；对于英文，通常 1 个字符 1 个字节，但对于汉字，通常使用 2 到 3 个字节，由于中文的字符空间过大，所以直接使用 BPE 可能导致词表过大
    • BBPE（Byte-level BPE） ，BBPE 是 BPE 的一个变体，与 BPE 的唯一区别是，BPE 是字符级别，BBPE 是字节级别，特别适用于中文，日文等基础字符过于庞大的语言编码，BBPE 会使用 “_”作为每个字第一个 Byte 的前缀来指示单词边界，从而识别
    • WordPiece 是 Google 提出的一种方法，没有公开的实现，其实现逻辑较为复杂，与 BPE 类似，但是每次合并子词时不是按照出现频率最高的词对，而是其他策略（Google是使用模型来预估词对出现的概率），比如 HuggingFace 上的一个实现是使用一个计算公式计算优先级 \(Score_{<A,B>} = \frac{Count_{A}}{Count_{A} \times Count_{B}}\)，然后按照优先级选择合并的词对
      • 公开资料显示：WordPiece 的目标是寻找最大化（分词训练集数据）似然的分词组合
      • 大名鼎鼎的 BERT 使用的就是 WordPiece 分词方法
    • Unigram ，又名 ULM（Unigram Language Model） ：先初始化一个巨大的词汇表 ，再逐渐删除出现概率低的词汇（会从词汇表中挑出使得 loss 增长最小的 10%~20% 的词汇），直到词汇表打到预定的大小时停止
  • 缺点：
    • 需要提前使用一些特定的算法，按照不同的文本训练集/分词方法进行分词，可能导致不同模型的分词结果不同
    • 每个模型需要增加一个自己的分词词表和分词函数（不同词表分词流程不同），以确保其他开发者可以使用
  • 优点：
    • 能准确控制词表大小，可大可小
    • 相对 Char 粒度分词，每个 Token 的语义相对独立，在词表数量足够的情况下，单个token不会大规模重复语义
    • 相对 Word 粒度分词，出现 OOV 的概率大大降低

分词的一般流程

  • 一个分词算法的目标是实现编码和解码
  • 常见的分词算法包括三个函数功能：
    • 词表生成 ：对给定数据进行编码生成
    • 编码 ：根据词表，对给定句子进行编码，输出为 token 列表
    • 解码 ：根据词表，给定 token 列表，输出句子

BPE 词表生成过程

  • BPE 的词表生成函数的输入输出 ：
    • 输入 ：训练数据 \(\mathcal{D}\)
    • 输出 ：词表（BPE的词表是有序的合并规则集合）
  • BPE 生成词表的执行步骤如下：
    • 第一步：初始化词汇表 vacab ：将训练数据 \(\mathcal{D}\) 中的文本拆分为字符级别的词汇表并统计字符出现次数
      • 注：单词的结尾使用一个特殊字符来表示 /w，后续处理过程中该字符与其他字符等价处理
    • 第二步：统计频率 ：统计所有相邻词汇对的出现频数并记录（注意：这一步是所有相邻的词汇都要统计）
    • 第三步：合并最频繁的字符对 ：将频数最高的词汇对 \(<v_1,v_2>\) 合并为一个新的词汇 \(v_{\text{new}}\)，并更新词汇表
      • 假设：原始词汇表频次统计次数为 \(v_1:count_1\)，\(v_2:count_2\)，而 \(v_3 = <v_1,v_2>\) 出现的频次是 \(count_3\)，
      • 更新方式：增加一个新的词汇 \(v_3: count_3\)，减少 \(<v_1,v_2>\) 对应的次数，更新结果为：\(v_1:count_1-count_3\)，\(v_2:count_2-count_3\)，若更新后的频数值为0，则该词汇可以从词表中永久删除了
      • 问题：可能出现新的词汇已经在词汇表中了吗？此时如何处理？
        • 回答：是可能出现的，相当于不同路径生成了同一个词汇，此时一般可以选择跳过更新或者将频数增加到新的词汇上，同时减少子词汇的频数

          如果更注重词汇表的稳定性和避免不必要的合并，那么跳过合并可能是较好的选择；如果希望更准确地反映子词的出现频率，以便在后续处理中更好地权衡不同子词的重要性，那么更新词频则更为合适。在一些实际的 BPE 实现中，也可能会综合考虑这两种方式，或者根据一些特定的条件来动态决定采取哪种处理方式

    • 第四步：重复第二步和第三步 ，直到触发终止条件，终止条件一般有：
      • 词汇表达到预定的词汇表大小：迭代过程中词汇表会越来越大，可以给定确定的词汇表大小
      • 迭代次数达到最大值

BPE 编码过程

  • 根据词表，对输入文本进行编码，生成 token 序列
  • 具体流程
    • 根据词表和有序的合并规则集合，对输入文本进行编码
    • 由于合并规则是有序的，所以按照规则即可实现固定的编码，相同文本编码一定是相同的
  • 实际上，在词表中加入权重分后，不需要任何其他文件（规则隐含在权重分中），可以使用贪心的方式编码
    • 首先轮询一遍，搜索所有的相邻编码 A 和 B 并尝试合并，如果合并结果 AB 在词表中（这里的词表可以使用有序词表，会加快匹配效率），则取出 AB 的权重分（注意，这里先不合并）
    • 轮询一遍以后分两种情况
      • 如果没有任何候选 AB 在词表中，则结束，编码完成
      • 否则，对所有候选的匹配 AB 按照权重分进行倒序排列，选择得分最高的匹配 AB 执行合并操作（注意，一个句子可以同时匹配多个完全相同的地方）
  • 另一种不需要权重分的贪心编码方式：
    • 直接从第一个字符开始，贪心匹配从第一个字符为初始起点的，最大满足长度的匹配（不需要权重分，只需要词表即可）
  • 代码示例，来自github.com/boyu-ai/Hands-on-NLP
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    ordered_vocabulary = {key: x for x, key in enumerate(vocabulary)}
    sentence = "nanjing beijing"
    print(f"输入语句：{sentence}")
    tokens = sentence.split(' ')
    tokenized_string = []
    for token in tokens:
        key = token+'_'
        splits = list(key)
        #用于在没有更新的时候跳出
        flag = 1
        while flag:
            flag = 0
            split_dict = {}
            #遍历所有符号进行统计
            for i in range(len(splits)-1): 
                #组合两个符号作为新的符号
                current_group = splits[i]+splits[i+1] 
                if current_group not in ordered_vocabulary:
                    continue
                if current_group not in split_dict:
                    #判断当前组合是否在词表里，如果是的话加入split_dict
                    split_dict[current_group] = ordered_vocabulary[current_group] 
                    flag = 1
            if not flag:
                continue
                
            #对每个组合进行优先级的排序（此处为从小到大）
            group_hist=[(k, v) for k, v in sorted(split_dict.items(),\
                key=lambda item: item[1])] 
            #优先级最高的组合
            merge_key = group_hist[0][0] 
            new_splits = []
            i = 0
            # 根据优先级最高的组合产生新的分词
            while i < len(splits):
                if i+1>=len(splits):
                    new_splits.append(splits[i])
                    i+=1
                    continue
                if merge_key == splits[i]+splits[i+1]:
                    new_splits.append(merge_key)
                    i+=2
                else:
                    new_splits.append(splits[i])
                    i+=1
            splits=new_splits
        tokenized_string+=splits

    print(f"分词结果：{tokenized_string}")

BPE 解码过程

  • 根据词表，对输出token进行解码，生成文本
  • 具体流程
    • 解码比较简单，按照token直接从词表查询原始字符值，添加到文本后面即可

其他说明

  • BPE 在对中英文处理时，方式不同，英文时，需要先按照空格进行切分，然后再进行词表生成，中文则不需要进行切分
    • 理解：英文需要切分的原因是为了防止出现"e y"这种包含空格的 token 吗？
    • 问题：实际上也可以包含吧？感觉不一定会影响效果（注：目前是不包含的）
  • OpenAI 统计 GPT3 和 GPT4 分词数量的在线链接OpenAI Tokenizer
  • ChatGPT-BPE编码过程实例(from OpenAI Tokenizer)
    • 实例1
    • 实例2
    • 实例3
    • 以上实例1-3说明：
      • BPE编码不是从最长的 token 开始搜索，否则实例3的编码应该是2个 token 才对
      • BPE编码也不是为了保证将每一个句子编码成最少的 token 数量（但原始算法的基本思想是所有训练数据上，尽量压缩 token 数量，减少存储）

附录：扩展词表的方法

附录：分词器比较

  • 常用分词器有：SentencePiece 和 Tiktoken
    对比维度 SentencePiece Tiktoken
    公司 Google OpenAI
    设计目标和使用场景 通用分词工具 OpenAI 专为 GPT 系列模型设计
    分词算法 支持多种算法，如 BPE、Unigram 等 BPE 算法
    词表特点 词表可根据数据集和任务训练，可调整大小和内容 针对 GPT 系列模型预先训练好，用户无法直接训练或修改
    性能和效率 性能和效率受算法选择和词表大小影响，某些情况下训练和分词速度可能较慢 分词速度快，能高效计算文本token数量
  • SentencePiece 是 Google 开源的一个分词工具库，包含以下优点：
    • SentencePiece 使用特殊的符号来转义空格，可以精确区分 love you ! 和 love you! 两个句子中 ! 前的空格
    • 高效实现了如 BPE 和 ULM 等分词方法，且自动将语料先转换为 Unicode 编码（注意这里不是 UTF-8），再输入分词算法，这样可以解决中文没有空格的问题
      • 理解（待确定）：Unicode 编码中文时是以字为单位的，比如你对应 U+4F60，可以从 U+ 准确识别每个文字开头？
      • 一些文章中提到，中文字符是在进行 BBPE 算法时，需要先将中文经过 UTF-8 变成字节，然后再经过 Unicode 统一编码成字符，最后再进行分词表构造
  • Tiktoken 的使用讲解：NLP（五十四）tiktoken的使用
  • 注：Llama1 和 Llama2 使用 SentencePiece，Llama3 开始使用 Tiktoken 了
  • 两种分词器的使用 Demo，分别展示 SentencePiece 和 Tiktoken 从指定数据构造词表、加载词表、编码和解码的全过程： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    import sentencepiece as spm
    import tiktoken

    # 数据文件构造
    data = [
    	"This is a sample sentence.",
    	"Another sample sentence for testing."
    ]
    with open('data.txt', 'w', encoding='utf-8') as f:
    	for line in data:
    		f.write(line + '\n')

    # SentencePiece 使用 Demo
    # 构造 SentencePiece 词表
    spm.SentencePieceTrainer.train(input='data.txt', model_prefix='sp_model', vocab_size=100)

    # 加载 SentencePiece 词表
    sp = spm.SentencePieceProcessor()
    sp.load('sp_model.model')

    # 编码和解码
    text = "This is a test."
    sp_encoded = sp.encode(text)
    sp_decoded = sp.decode(sp_encoded)
    print("SentencePiece 编码结果:", sp_encoded)
    print("SentencePiece 解码结果:", sp_decoded)

    # Tiktoken 使用 Demo
    # 对于 Tiktoken，我们使用预训练的 cl100k_base 编码
    # 这里不需要构造词表，直接加载预训练编码
    encoding = tiktoken.get_encoding("cl100k_base") 
    # 还可以选择其他编码方式，不同模型使用不同的编码方式，可以调用函数对应抽取到编码方式
    # print(tiktoken.encoding_for_model('gpt-3.5-turbo'))

    # 编码和解码
    tiktoken_encoded = encoding.encode(text)
    tiktoken_decoded = encoding.decode(tiktoken_encoded)
    print("Tiktoken 编码结果:", tiktoken_encoded)
    print("Tiktoken 解码结果:", tiktoken_decoded)

附录：transformers.AutoTokenizer 的使用

tokenizer.tokenize() 处理单文本

  • tokenize 处理单文本示例 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    from transformers import AutoTokenizer

    # Tokenizer 加载，根据路径下的文件初始化自身，路径下会包含：
    # # tokenizer.json（包含词表） 
    # # tokenizer_config.json（包含 tokenizer_class（可以按照需要自己实现并提 PR 到 transformers 库） 和 special_token 等信息）
    # # special_tokens_map.json（通常与 tokenizer_config.json 中的内容重复，可选，可被类初始化时读取使用），special_token 包含 bos,eos,pad,unk 四个，是模型必须指定的，有时候也实现在 tokenizer_config.json 等文件中
    # # vocab.json (词表信息，可选，可被类初始化时读取使用)
    tokenizer = AutoTokenizer.from_pretrained("model_name")

    text = "Hello world! This is a test."
    # 单行分词
    tokens = tokenizer.tokenize(text)
    print("Tokens:", tokens) # 这里输出的可能是 类似 ['Hello', ',', 'Ġworld', '!', 'ĠThis', 'Ġis', 'Ġa', 'Ġtest', '.'] 的结果（分词器里面使用 'Ġ'（Unicode 字符 U+0120，带点空格）在这里表示普通空格（U+0020））
    # 注：这一步中，如果是BBPE 分词，且输入是中文甚至可能看起来像是乱码，因为常见的 BBPE 分词方式会先将中文转换成 UTF-8 编码后再进行分词

    # 转换为索引 IDs
    ids = tokenizer.convert_tokens_to_ids(tokens)
    print("IDs:", ids) # 这里输出具体索引（int 类型的数字列表）

    # 解码验证，将 IDs 解码为原始文本
    decoded = tokenizer.decode(ids)
    print("解码结果:", decoded) # 与原始输出文本 text 一致

tokenizer() 批量处理文本

  • tokenizer(text) 实际上是 tokenizer.__call__(text) 的简写，是最常用、最完整的文本编码方法
  • tokenizer(text) 函数包含以下步骤：
    • 1）分词 ：将文本拆分成 Tokens（注意： Token 不是 ID，不是整数）
    • 2）映射 ：将 Toke ns 转换为对应的 IDs
    • 3）添加特殊标记 ：如[CLS], [SEP], [PAD]等
    • 4）填充和截断 ：处理长度不一致的问题，需指定参数
    • 5）返回张量 ：转换为 PyTorch/TensorFlow 张量，注意返回的是 ID 列表，不是 Token 列表
  • tokenizer(text) 函数参数较多，重点需要关注的如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    def __call__(
    		self,
    		text: Union[str, List[str], List[str], List[List[str]], None] = None, # 待编码的单条文本或文本列表；若为 List[str] 类型，表示已按词切分
    		text_pair: Optional[Union[str, List[str], List[str], List[List[str]]]] = None, # 与 text 配对的第二条文本（或列表），用于需要“句子对”输入的任务，如 NLI、QA
    		text_target: Union[str, List[str], List[str], List[List[str]], None] = None, # 作为“标签/目标”的单条文本或文本列表，用于 seq2seq 等需对目标序列编码的场景
    		text_pair_target: Optional[
    			Union[str, List[str], List[str], List[List[str]]]
    		] = None, # 与 text_target 配对的第二条目标文本（或列表）
    		add_special_tokens: bool = True, # 是否在输出中自动添加特殊符号（如 [CLS]、[SEP]）
    		padding: Union[bool, str, PaddingStrategy] = False, # 是否/如何填充：True/False、'longest'、'max_length' 或 "do_not_pad"（PaddingStrategy 枚举）
    		truncation: Union[bool, str, TruncationStrategy, None] = None, # 是否/如何截断：True/False、'only_first'、'only_second'、"longest_first" 或  "do_not_truncate"（TruncationStrategy 枚举）
    		max_length: Optional[int] = None, # 允许的最大序列长度；超出则按 truncation 策略处理
    		stride: int = 0, # 滑动窗口截断时，相邻片段间的重叠 token 数
    		is_split_into_words: bool = False, # 表明输入 text 是否已经按词切分（List[str]），此时不重新分词
    		pad_to_multiple_of: Optional[int] = None, # 将序列长度填充到指定值的整数倍（常见于 TensorRT / 优化内核）
    		padding_side: Optional[str] = None, # 显式指定填充方向：'left' 或 'right'；默认采用 tokenizer 配置
    		return_tensors: Optional[Union[str, TensorType]] = None, # 返回张量类型：'np'、'pt'、'tf' 等 TensorType 枚举；None 则返回 Python 列表
    		return_token_type_ids: Optional[bool] = None, # 是否返回 token_type_ids（区分句子对）；None 时按模型默认
    		return_attention_mask: Optional[bool] = None, # 是否返回 attention_mask；None 时按模型默认
    		return_overflowing_tokens: bool = False, # 是否返回因截断被溢出的 token 列表
    		return_special_tokens_mask: bool = False, # 是否返回特殊符号掩码（1 表示特殊 token，0 表示普通 token）
    		return_offsets_mapping: bool = False, # 是否返回每个 token 对应原始文本的 (start, end) 字符偏移
    		return_length: bool = False, # 是否返回编码后序列长度
    		verbose: bool = True, # 遇到警告/异常时是否打印详细信息
    		**kwargs,
    ) -> BatchEncoding: # BatchEncoding 对象支持字典索引和属性访问
tokenizer() 参数说明（按常用程度排序）
  • text ：必填参数，要分词的文本数据
    • 支持类型：单个字符串（如 "Hello world"）、字符串列表（如 ["Hello, World", "你好，你好"]）、嵌套字符串列表（如 [["Hello, World", "你好，你好"], ["Hi，你好吗？", "谢谢你"]]，仅部分模型支持）
    • 说明：当传入列表时，会自动批量处理；嵌套列表通常用于处理“文本对”场景，需结合 padding 等参数使用
  • padding ：控制是否填充（补全）序列到相同长度，可选值如下：
    • False / None：不填充（默认值），每个序列保留原始长度
    • True / "longest"：填充到批量中最长序列的长度 ，此时 "max_length" 指定了也不会使用
    • "max_length"：填充到 max_length 参数指定的长度（需同时设置 max_length）
    • "do_not_pad"：等价于 False
    • 说明：填充时使用分词器的 pad_token（默认通常是 <pad>），填充位置由 padding_side 控制（默认右侧）
  • truncation ：控制是否截断过长的序列，可选值如下：
    • False / None：不截断（默认值），若序列长度超过 max_length 会（理解：不能截断，又不能超过 max_length，矛盾了，必须报错）
    • True / "longest_first"：截断到 max_length（需指定 max_length），优先截断最长的序列（批量场景）
    • "only_first"：仅截断输入文本对中的第一个文本（如 premise）
    • "only_second"：仅截断输入文本对中的第二个文本（如 hypothesis）
    • "do_not_truncate"：等价于 False
    • 说明：截断时默认从右侧截断（由 truncation_side 控制）
  • max_length ：指定序列的最大长度
    • 类型：整数（int）
    • 说明：需与 padding（设为 "max_length"）或 truncation（设为 True 等）配合使用；
      • 若不指定，默认使用分词器的 model_max_length（通常是模型支持的最大输入长度，如 512、1024 等）
  • return_tensors ：指定返回的张量类型
    • 可选值：None（默认）、"pt"（PyTorch 张量）、"tf"（TensorFlow 张量）、"np"（NumPy 数组）、"jax"（JAX 张量）
    • 说明：默认返回普通 Python 列表；指定后返回对应框架的张量，可直接传入模型训练/推理
  • return_attention_mask ：是否返回注意力掩码
    • 类型：布尔值（bool），默认 True
    • 说明：注意力掩码用于告诉模型哪些 token 是真实文本（值为 1），哪些是填充的 pad token（值为 0）；若设为 False，返回结果中不包含 attention_mask 键
  • return_token_type_ids ：是否返回 token 类型 ID（用于区分文本对）
    • 类型：布尔值（bool），默认值由模型决定（如 BERT 类模型默认 True，GPT 类模型默认 False）
    • 说明：文本对场景（如 text = (sent1, sent2)）中，token_type_ids 用 0 标识第一个文本的 token，1 标识第二个文本的 token；单文本场景中全为 0
    • 注意：亲测，很多现代 GPT 模型中，即使设置 return_token_type_ids=True，输出都全 0 了
  • add_special_tokens ：是否添加模型所需的特殊 token
    • 类型：布尔值（bool），默认 True
    • 说明：特殊 token 包括 cls_token（如 <cls>）、sep_token（如 <sep>）、bos_token（如 <s>）、eos_token（如 </s>）等，不同模型的特殊 token 不同；设为 False 则仅返回原始分词结果，不添加任何特殊 token
  • return_offsets_mapping ：是否返回 token 在原始文本中的偏移量（起始/结束索引）
    • 类型：布尔值（bool），默认 False
    • 说明：返回的 offsets_mapping 是一个列表，每个元素为 (start, end) 元组，对应每个 token 在原始文本中的字符位置（可用于实体标注、文本对齐等任务）
  • return_length ：是否返回每个序列的原始长度（未填充/未截断前）
    • 类型：布尔值（bool），默认 False
    • 说明：返回的 length 列表包含每个序列在处理前的实际长度，便于后续计算有效 token 数
  • padding_side ：填充方向
    • 可选值："right"（默认，右侧填充）、"left"（左侧填充）
    • 说明：部分模型（如 GPT 类）可能需要左侧填充，需根据模型要求设置
  • truncation_side ：截断方向
    • 可选值："right"（默认，右侧截断）、"left"（左侧截断）
    • 说明：根据任务需求调整，例如处理历史对话时可能需要左侧截断旧对话
  • stride ：截断时的重叠长度（用于长文本分段处理）
    • 类型：整数（int），默认 0
    • 说明：当文本长度超过 max_length 时，截断后保留前一段序列的 stride 个 token，避免丢失上下文（如长文档分类、问答任务）
  • return_overflowing_tokens ：是否返回截断后溢出的 token 组成的额外序列
    • 类型：布尔值（bool），默认 False
    • 说明：结合 stride 使用，长文本会被分成多个重叠的子序列，返回所有子序列及对应的 overflow_to_sample_mapping（映射子序列到原始样本的索引）
  • allow_multiple_sentences ：是否允许单个文本包含多个句子（通过标点/换行分隔）
    • 类型：布尔值（bool），默认 True
    • 说明：部分分词器支持自动拆分多句子，但通常不影响核心分词逻辑，保持默认即可
  • verbose ：是否输出详细日志
    • 类型：布尔值（bool），默认 True
    • 说明：设为 False 可关闭警告信息（如序列长度超过模型最大长度的警告）

简单的返回结果示例

  • 示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    from transformers import AutoTokenizer

    model_path = "path_to_model"

    # 初始化分词器（以 BERT 为例）
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    # 1. 基础用法（单文本）
    text = "Hello, world! This is a test."
    inputs = tokenizer(
    	text,
    	padding=True,
    	truncation=True,
    	max_length=10,
    	return_tensors="pt"
    )
    print(type(inputs)) # <class 'transformers.tokenization_utils_base.BatchEncoding'>
    print(inputs.keys())  # 输出：dict_keys(['input_ids', 'attention_mask'])
    print(inputs["input_ids"].shape)  # 输出：torch.Size([1, 9])

    # 文本对用法（如句子相似度），仅在 BERT RoBERTa 等模型中可以使用，其他模型会输出 inputs_pair["token_type_ids"]，但是包含的值都是 0
    text_pair = ("I like cats.", "I love dogs.")
    inputs_pair = tokenizer(
    	text_pair,
    	padding="max_length",
    	max_length=5,
    	return_token_type_ids=True,
    	return_tensors="np"
    )
    print(inputs_pair.keys()) # 输出：dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])
    print(inputs_pair["input_ids"].shape) # 输出：(2, 5)
    print(inputs_pair["token_type_ids"])  # 在 BERT RoBERTa 等模型中，0 对应第一个句子，1 对应第二个句子；现代大模型一般都是 0
    # 输出：
    # [[0 0 0 0]
    #  [0 0 0 0]]

    # 长文本分段（带重叠）
    long_text = "This is a very long text that exceeds the max length. We need to split it into chunks with overlap."
    inputs_long = tokenizer(
    	long_text,
    	max_length=10,
    	truncation=True,
    	stride=3,
    	return_overflowing_tokens=True,
    	return_offsets_mapping=True
    )
    print(inputs_long.keys()) # 输出：dict_keys(['input_ids', 'attention_mask', 'offset_mapping', 'overflow_to_sample_mapping'])
    print(type(inputs_long)) # <class 'transformers.tokenization_utils_base.BatchEncoding'>
    print(len(inputs_long["input_ids"]))  # 输出：3
    print(inputs_long["input_ids"])
    # 输出：
    # [
    #   [3031, 472, 358, 2481, 2093, 3402, 524, 159, 369, 156],
    #   [159, 369, 156, 5041, 112, 1769, 1838, 408, 11694, 574],
    #   [408, 11694, 574, 1394, 48782, 537, 30033, 112]
    # ]

    # 批量处理多条文本
    texts = ["你好", "世界", "深度学习"]
    encoded_batch = tokenizer(texts)
    print(encoded_batch.keys()) # 输出：dict_keys(['input_ids', 'attention_mask'])
    print(encoded_batch["input_ids"]) # [[135], [487], [18834]]

其他更多完整详细示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
from transformers import AutoTokenizer
MODEL = "path_to_model"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

# 故意构造 3 条长度差异很大的句子
raw_texts = [
	"Hello world!",
	"This is a much longer sentence, designed to exceed the usual 8 ~ 12 token budget.",
	"Hi"
]

def group_title(title):
	print(f"\n{title}")

def pretty_print(enc):
	print("input_ids       :", enc["input_ids"])
	if "attention_mask" in enc:
		print("attention_mask  :", enc["attention_mask"])
	if "token_type_ids" in enc:
		print("token_type_ids  :", enc["token_type_ids"])
	if "overflowing_tokens" in enc:
		print("overflowing_tokens:", enc["overflowing_tokens"])
	print("-" * 50)

# 默认：不 padding 也不 truncation -> 长度各异的 list
group_title("默认：padding=False, truncation=None")
enc = tokenizer(raw_texts)
pretty_print(enc)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# 只 padding 不 truncation -> 长度一致，无截断
group_title("只 padding（最长）不 truncation")
enc = tokenizer(raw_texts, padding=True)
pretty_print(enc)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# 只 truncation 不 padding -> 超过 max_length 被截断，仍返回 list
group_title("只 truncation=True 不 padding, max_length=12")
enc = tokenizer(raw_texts, truncation=True, max_length=12)
pretty_print(enc)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# padding + truncation -> 固定长度，超长截断，不足填充
group_title("padding + truncation, max_length=16")
enc = tokenizer(raw_texts, padding=True, truncation=True, max_length=16)
pretty_print(enc)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# 返回 PyTorch 张量（必须 padding，否则长度不同会抛错）
group_title("return_tensors='pt' 必须配合 padding")
enc = tokenizer(raw_texts, padding=True, truncation=True, max_length=16, return_tensors="pt")
print("input_ids 形状:", enc["input_ids"].shape)
print("attention_mask 形状:", enc["attention_mask"].shape)
print("类型:", type(enc["input_ids"]))

# 按照窗口拆分成多个 chunks，stride + return_overflowing_tokens 演示滑动窗口
group_title("stride=4 + return_overflowing_tokens 截断窗口")
long_text = "This is a very long sentence that will be split into overlapping chunks."
enc = tokenizer(long_text, truncation=True, max_length=10, stride=4, return_overflowing_tokens=True, return_offsets_mapping=True)
print("共返回片段数:", len(enc["input_ids"]))
for idx, (ids, off) in enumerate(zip(enc["input_ids"], enc["offset_mapping"])):
	print(f"片段 {idx}: tokens={ids}")
	print(f"offsets={off}")
	print("-" * 30)

# padding='max_length' 强制 pad 到 max_length
group_title("padding='max_length' 强制 32")
enc = tokenizer(raw_texts, padding='max_length', max_length=32)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# pad_to_multiple_of 演示 8 的倍数填充
group_title("pad_to_multiple_of=8, 填充到最小符合要求的 8 的倍数，比 21 大的倍数是 24")
enc = tokenizer(raw_texts, padding=True, pad_to_multiple_of=8)
print("每条长度:", [len(ids) for ids in enc["input_ids"]])

# tokenizer 类属性：默认最大长度 & 截断策略
group_title("tokenizer 默认属性")
print("tokenizer.model_max_length :", tokenizer.model_max_length) # 从原始模型文件中读取
print("tokenizer.truncation_side  :", tokenizer.truncation_side)
print("tokenizer.padding_side     :", tokenizer.padding_side)
# 其他：如 tokenizer.vocab_size 等

# 默认：padding=False, truncation=None
# input_ids       : [[20769, 3121, 224], [3156, 597, 483, 2799, 6991, 14941, 235, 7605, 533, 15890, 494, 15299, 444, 247, 8581, 444, 240, 241, 10539, 14517, 237], [23383]]
# attention_mask  : [[1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1]]
# --------------------------------------------------
# 每条长度: [3, 21, 1]
#
# 只 padding（最长）不 truncation
# input_ids       : [[20769, 3121, 224, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3], [3156, 597, 483, 2799, 6991, 14941, 235, 7605, 533, 15890, 494, 15299, 444, 247, 8581, 444, 240, 241, 10539, 14517, 237], [23383, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]
# attention_mask  : [[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
# --------------------------------------------------
# 每条长度: [21, 21, 21]
#
# 只 truncation=True 不 padding, max_length=12
# input_ids       : [[20769, 3121, 224], [3156, 597, 483, 2799, 6991, 14941, 235, 7605, 533, 15890, 494, 15299], [23383]]
# attention_mask  : [[1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1]]
# --------------------------------------------------
# 每条长度: [3, 12, 1]
#
# padding + truncation, max_length=16
# input_ids       : [[20769, 3121, 224, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3], [3156, 597, 483, 2799, 6991, 14941, 235, 7605, 533, 15890, 494, 15299, 444, 247, 8581, 444], [23383, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]]
# attention_mask  : [[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
# --------------------------------------------------
# 每条长度: [16, 16, 16]
#
# return_tensors='pt' 必须配合 padding
# input_ids 形状: torch.Size([3, 16])
# attention_mask 形状: torch.Size([3, 16])
# 类型: <class 'torch.Tensor'>
#
# stride=4 + return_overflowing_tokens 截断窗口
# 共返回片段数: 2
# 片段 0: tokens=[3156, 597, 483, 2606, 2218, 14941, 649, 1253, 621, 11819]
# offsets=[(0, 4), (4, 7), (7, 9), (9, 14), (14, 19), (19, 28), (28, 33), (33, 38), (38, 41), (41, 47)]
# ------------------------------
# 片段 1: tokens=[649, 1253, 621, 11819, 1519, 44984, 48907, 237]
# offsets=[(28, 33), (33, 38), (38, 41), (41, 47), (47, 52), (52, 64), (64, 71), (71, 72)]
# ------------------------------
#
# padding='max_length' 强制 32
# 每条长度: [32, 32, 32]
#
# pad_to_multiple_of=8
# 每条长度: [24, 24, 24]
#
# tokenizer 默认属性
# tokenizer.model_max_length : 131072
# tokenizer.truncation_side  : right
# tokenizer.padding_side     : right

常见注意事项

  • 建议批量处理，而不是单条处理
    • 大批量处理时注意内存使用
  • 不同模型有不同的最大长度限制，模型配置中会包含
  • 某些模型会自动添加特殊标记，也可以自己预设标记： 
    1
    tokenizer.pad_token = tokenizer.eos_token  # 设置填充token

附录：transformers.AutoTokenizer.decode 的使用

  • tokenizer.decode 用于将模型输出的 token ID 序列（或 token 索引序列）反向解码为自然语言文本

    • 会自动处理 token 间的拼接逻辑（如去除 subword 分隔符、恢复原始词汇），屏蔽底层 tokenization 细节（如 BPE、WordPiece 等子词切分的拼接）

  • 函数签名：

    1
    2
    3
    4
    5
    6
    7
    def decode(
        self,
        token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: Optional[bool] = None,
        **kwargs,
    ) -> str:

  • 需要注意：

    • 解码结果依赖于 tokenizer 的训练数据和切分规则（如 BPE 切分的词汇可能需要特殊拼接逻辑），需确保解码使用的 tokenizer 与编码（tokenizer.encode）时一致；
    • 对于包含填充 token（PAD）的 input_ids，建议设置 skip_special_tokens=True 或指定 padding_token_id，避免解码出无效的填充文本；

参数说明

  • token_ids：必选，对应编码时的 input_ids，需要解码的 token ID 序列，支持两种输入形式：
    • 单个整数（单个 token ID），例如 101、2023；
    • 整数列表/张量（多个 token ID 组成的序列），例如 [101, 2054, 2182, 102]、torch.tensor([101, 3845, 102])，也支持批量输入（如二维张量，shape 为 [batch_size, seq_len]）
      • 注意：这里的批量不是 list 形式，必须是一个张量才行
  • skip_special_tokens：布尔值，默认 False
    • 设为 True 时，解码过程中会自动跳过所有特殊 token（如 [CLS]、[SEP]、[PAD]、[MASK] 等，具体取决于 tokenizer 的配置）；
    • 设为 False 时，会保留所有特殊 token 原样输出（例如解码结果可能包含 "<s>" "<pad>" 等标记）
    • 实际使用中通常设为 True，以获取干净的自然语言文本
  • clean_up_tokenization_spaces：布尔值，默认 True
    • 设为 True 时，会自动清理解码后文本中多余的空格（如 subword 拼接后残留的空格、特殊 token 移除后的空字符）；
    • 设为 False 时，保留 token 拼接后的原始空格布局（可能出现连续空格或不合理空格）
    • 注：clean_up_tokenization_spaces 参数的必要性已大大降低，亲测这个参数为 True 或 False 的结果是一样的，多个特殊的字符串测试没有例外
  • padding_token_id：整数，默认使用 tokenizer 配置中的 padding_token_id（如 [PAD] 对应的 ID）
    • 当 input_ids 中包含填充 token ID 时，可通过该参数指定需要忽略的填充 ID，避免解码出 [PAD] 对应的文本标记
  • stop_at_eos：布尔值，默认 False（部分 tokenizer 版本默认 True，需结合具体库版本确认）
    • 设为 True 时，解码过程中遇到 eos_token_id 会立即终止，仅返回 eos_token_id 之前的文本；
    • 设为 False 时，会完整解码 input_ids 中的所有 token ID，包括 eos_token_id 之后的内容

使用示例

  • 使用示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    # 基础解码（跳过特殊 token）：
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    input_ids = [101, 2054, 2182, 2003, 1996, 3014, 102]  # 包含 [CLS]（101）和 [SEP]（102）
    text = tokenizer.decode(input_ids, skip_special_tokens=True)
    print(text)  # 输出："i love this movie"（自动去除特殊 token 并拼接子词）

    # 保留特殊 token：
    text = tokenizer.decode(input_ids, skip_special_tokens=False)
    print(text)  # 输出："[CLS] i love this movie [SEP]"

    # 解码批量输入（二维张量）：
    import torch
    batch_input_ids = torch.tensor([[101, 2054, 102], [101, 2182, 102]])
    batch_text = tokenizer.decode(batch_input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    print(batch_text)  # 输出批量文本（具体格式取决于 tokenizer，通常为列表或拼接字符串）

decode() vs batch_decode()

  • tokenizer.decode() 和 tokenizer.batch_decode() 的核心区别是 处理输入的维度和场景
    • tokenizer.decode() 用于解码「单个序列」的 input_ids，无批量优化
    • tokenizer.batch_decode() 用于批量解码「多个序列」的 input_ids，有批量优化，比循环调用 decode() 更高效
    • 本质是「单样本」与「多样本」的适配差异
  • 输入要求：
    • decode() 仅接受 1维输入 ：
      • 支持 list[int]（如 [101, 2054, 3110, ...]）、1D Tensor、1D numpy数组，若传入2D数据会直接报错
    • batch_decode() 仅接受 2维输入 ：
      • 支持 list[list[int]]（如 [[101, 2054], [101, 3110]]）、2D Tensor、2D numpy数组，自动处理批量样本
  • 输出情况：
    • decode() 输出单个字符串（str）；
    • batch_decode() 输出字符串列表（list[str]）
  • 两者用相同参数解码时，结果应该一致

附录：中英文 Tokenization 编码效率对比

  • 注：本文以 Qwen3 编码方式为例，且只是简单的粗糙对比，不严谨

  • 分别尝试了将文本从中文翻译为英文和从英文翻译为中文两种方式（翻译工作由 豆包完成），代码如下：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    from transformers import AutoTokenizer
    import sys
    model_name = "/Users/jiahong/Workspace/IdeaProjects/Torch/LLaMA-Factory/model/Qwen3-0.6B/"

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    print("Chinese to English")
    print("Chinese")
    prompt_c = """xxx"""
    print(f"{len(prompt_c)} 字, 存储 {sys.getsizeof(prompt_c)-49} B") # 49 是对象的固定存储部分
    token_c = tokenizer.tokenize(prompt_c)
    print(f"len of tokens: {len(token_c)}, {len(prompt_c)/len(token_c)} 字/token")

    print("English")
    prompt_e = """xxx"""
    print(f"{len(prompt_e.split(' '))} 单词, {sys.getsizeof(prompt_e)-49} B, len: {len(prompt_e)}")
    token_e = tokenizer.tokenize(prompt_e)
    # print(token_e)
    print(f"len of tokens: {len(token_e)}, {len(prompt_e)/len(token_e)} B/token, {len(prompt_e.split(' '))/len(token_e)} words/token")

    print("="*30)
    print("English to Chinese")
    print("Chinese")
    prompt_c = """xxx"""
    print(f"{len(prompt_c)} 字, 存储 {sys.getsizeof(prompt_c)-49} B")
    token_c = tokenizer.tokenize(prompt_c)
    # print(token_c)
    print(f"len of tokens: {len(token_c)}, {len(prompt_c)/len(token_c)} 字/token")

    print("English")
    prompt_e = """xxx"""
    print(f"{len(prompt_e.split(' '))} 单词, {sys.getsizeof(prompt_e)-49} B, len: {len(prompt_e)}")
    token_e = tokenizer.tokenize(prompt_e)
    print(f"len of tokens: {len(token_e)}, {len(prompt_e)/len(token_e)} B/token, {len(prompt_e.split(' '))/len(token_e)} words/token")

    # Chinese to English
    # Chinese
    # 2514 字, 存储 5053 B
    # len of tokens: 1784, 1.4091928251121075 字/token
    # English
    # 1611 单词, 21073 B, len: 10524
    # len of tokens: 2239, 4.700312639571237 B/token, 0.7195176418043769 words/token
    # ==============================
    # English to Chinese
    # Chinese
    # 2061 字, 存储 4147 B
    # len of tokens: 1421, 1.450387051372273 字/token
    # English
    # 814 单词, 10739 B, len: 5357
    # len of tokens: 1223, 4.3802125919869175 B/token, 0.6655764513491415 words/token

  • 结论（以 Qwen3 的编码方式为例）：

    • 中文：约 1.5 字/token
    • 英文：约 0.7 words/token，且大约合 4.5 B/token
    • 编码效率上看，汉译英时，中文优于英文(0.8:1)；英译汉时，则英文较优(1.16:1)

  • 有趣的是，汉译英和英译汉的效率差异很大：

    • 汉译英时，1 个单词大概对应 1.56 个中文字
    • 英译汉时，1 个单词大概对应 2.53 个中文字
    • 解释：不翻译别人，自由发挥的情况下，文字使用效率是最高的

附录：新增 Token 的方法

  • 可以对 HuggingFace 模型添加新的 Token，添加 Token 的类型可以有两种
    • 特殊 Token（Special Token）：
      • Special Token 是具有特定语义或结构作用的标记，常用于向模型指示文本的结构、任务类型或特殊信息
      • Special Token 会保证原子性 (Atomicity) ： Special Token 永远不会被分词器拆分成更小的 subword 或字符，它们作为一个整体被识别
      • 通常可通过两种方式添加：
        • 通过 tokenizer.add_special_tokens() 方法添加（推荐使用）
          • 接受一个字典，如 {'additional_special_tokens': ['<NEW_TOKEN>']}
        • 通过 tokenizer.add_tokens() 并设置参数 special_tokens=True 来添加
      • 在解码时，通常可以选择跳过这些 Special Token（使用 skip_special_tokens=True）
      • Special Token 添加后，词表中显示它的属性为 "special": true
    • 常规 Token (Regular Token)：
      • Regular Token 指的是模型训练过程中常见的词汇、词片段（subword）或字符，它们在自然语言文本中出现，主要用于表示文本内容
      • 主要通过 tokenizer.add_tokens() 方法添加
        • 新添加的 Token 会被加入到词汇表中，从现有词汇表的末尾开始分配新的索引
        • 如果添加的 Token 是一个完整的词汇，它通常不会被分词器拆分 ，但如果它本身就是 subword 的一部分，则会按照分词器的逻辑处理
        • 如果一个新词汇被添加，分词器在遇到该词时会优先使用该新词汇作为一个整体
      • 常规 Token 添加后，词表中显示它的属性为 "special": false
  • 此外，如果在模型预训练时已经增加了 <mask_id> 还可以手动添加（已测试没问题）：
    • 直接修改 tokenizer.json 中的字段
      • 分别修改 <mask_id_xx> 为指定 Token 文本，比如 <start_think>
      • 映射关系一般有两个地方
        • 一个是 "added_tokens" 中的 ID 映射
        • 另一个是词表中的映射关系（一般也在 "tokenizer.json" 的 "model" -> "vocab" 中）

新增 Token 对模型嵌入层的影响

  • 无论是哪种方式添加了 Token，都必须调用 model.resize_token_embeddings(len(tokenizer)) 来调整模型的词嵌入层 (Embedding Layer) 大小，使其与新的词汇表大小匹配，否则新增的 Token 无法被模型正确处理
  • Special Tokens 通常会被赋予特定的嵌入初始化方式（例如部分模型对特殊 tokens 有默认初始化逻辑），而普通 tokens 的嵌入则可能随机初始化

修改文件的方式添加 Special Token

  • 不建议使用这种方式，建议通过代码动态添加并 save_pretrained 的方式持久化到文件中
  • 手动修改文件容易出错，比如 Special Token 的添加，除了词表外，还需要修改 special_tokens_map.json 文件, 否则 tokenizer 不会将其识别为 Special Token
    • 注：部分模型也可以将 “special_tokens_map” 作为一个 key 放到 tokenizer.json 中？

tokenizer.add_special_tokens() 函数使用及对比

  • tokenizer.add_special_tokens() 接受一个字典对象作为参数，且字典的 key 必须在 tokenizer.SPECIAL_TOKENS_ATTRIBUTES 中
  • 在 Hugging Face 的 tokenizers 库中，tokenizer.SPECIAL_TOKENS_ATTRIBUTES 是分词器类（如 PreTrainedTokenizer）的一个内置属性，它定义了 “特殊 token 类型” 的标准名称集合
    • 注意：不同版本可能写作 SPECIAL_TOKENS_ATTRIBUTES 或类似属性，本质一致
    • 常见配置为：['bos_token', 'eos_token', 'unk_token', 'sep_token', 'pad_token', 'cls_token', 'mask_token', 'additional_special_tokens']
    • “additional_special_tokens” 用于添加自定义的 Special Token
  • 使用 tokenizer.add_special_tokens() 添加 与 tokenizer.add_tokens(..., special_tokens=True) 的区别：
    • Special Token 是否会出现在 special_tokens_map.json 文件 和 tokenizer_config.json 的 "additional_special_tokens" 字段中？
      • tokenizer.add_tokens(..., special_tokens=True) 方式添加的不会出现
      • tokenizer.add_special_tokens() 方式添加的则会出现
    • 两者都会出现在 tokenizer_config.json 的 "added_tokens_decoder" 中
    • 说明：add_special_tokens（标准用法）会将新添加的 Token 关联到分词器的任何标准特殊 Token 属性，而 add_tokens(..., special_tokens=True) 不会
      • add_tokens(..., special_tokens=True) 添加后只能在词表中找到索引，其他地方看不到
  • 添加 Special Token 的示例： 
    1
    2
    3
    4
    5
    print(tokenizer.special_tokens_map) # 可能输出为：{'bos_token': '<bos>', 'eos_token': '<eos>', 'unk_token': '<unk>', 'pad_token': '<pad>'}
    tokenizer.add_special_tokens({"additional_special_tokens": ["<llm_assistant>", "<llm_user>"]})
    print(tokenizer.special_tokens_map) # 可能输出为：{'bos_token': '<bos>', 'eos_token': '<eos>', 'unk_token': '<unk>', 'pad_token': '<pad>', 'additional_special_tokens': ['<llm_assistant>', '<llm_user>']}

    # 注意：tokenizer.add_tokens(["<longcat_assistant>", "<longcat_system>"], special_tokens=True) 方式添加的 Special Token 不会出现在 special_tokens_map 中

新增 Special Token 与解码注意事项

  • 特别注意：推理引擎解码时，一般会默认会自动跳过所有 Special Token
  • 在定义 Special token 时需要非常小心
    • 工具调用等场景中，会依赖工具调用相关的 Token 的明文文本做工具解析（例如 <tool_call> 等）
    • 因此需要设置为这部分 Token 为 Regular Token (配置方式为：tokenizer.json 中设置 "special": false)
  • 必须设置为 Regular Token 的 Token 有：
    • tool 调用回复相关的 Token
    • thinking 相关的 Token

新增 Special Token 的初始化

  • 新增 Token 的训练可能不太充足，可考虑使用已有 Token 的 embedding 来进行初始化
  • 常用方式：使用添加 Token 前的编码 Token（可能为多个）对应的 embedding 均值作为添加 Token 后的 embedding 初始值
    • 比如新增 '<think>' 作为 Special Token，这个 Token 在原来的场景里面可能是编码为 ['<', 'think', '>'] 这三个 Token 的，此时可以使用 这三个 Token embedding 的均值作为新 Token 的 embedding 初始值
    • 注：这种方式对于有语义的 Special Token 可能有一定帮助（比如 assistant: 或 <assistant:> 等），但是对于非常特殊的非语义 Token 理论上没有帮助，如 ACBD34 或 <ACBD34> 等）
    • 亲测：这种方法得到的结果可能持平甚至负向，原因是提前保留的 <mask_id_xx> 对应的 Embedding 其实已经被训练过了（压低出现的概率），这里替换可能反而会影响结果
  • 初始化流程（以 HF 格式的模型为例）：
    • 在 model.safetensors.index.json 中 找到 embedding 层变量和 LM Head 层变量所在的 safetensors 文件
    • 读取对应文件，检索到对应的 embedding 并修改（注意：一定要两个层的参数都修改）
      • 注：可能部分模型上该 embedding 参数和 LM Head 层参数是共享的？
  • 已有模型，修改初始化值的代码示例（一般来说续修改 embedding 层 和 LM Head 层）： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    # 以使用旧 Token ID embedding 的均值初始化为例
    from safetensors.torch import load_file, save_file
    from transformers import AutoTokenizer
    import torch

    # 新旧路径加载 Tokenizer
    tokenizer_base = AutoTokenizer.from_pretrained('/path_to_old_base_tokenizer/') # 不包含新 Special Token
    tokenizer_new = AutoTokenizer.from_pretrained('/path_to_new_tokenizer/') # 包含新 Special Token

    # 指定需要修改的文件
    # # 对于多文件存储的大模型参数，需要在 model.safetensors.index.json 中找到 Embedding 层参数或者 LM Head 层参数所在的路径
    st_path_base = '/path_to_old_base_safetensors/model_xx-of-xx.safetensors' # 随机初始化/原始初始化的模型权重
    st_path_new = '/path_to_new_safetensors/model_xx-of-xx.safetensors' # 待写入路径，写入后包含自定义初始化的 embedding 的结果

    # 从原来的 safetensors 文件路径读取数据，得到一个字典数据 dict[name, weight]
    safetensors_weights = load_file(st_path_base)

    # # print weights for debug
    # for name, weight in safetensors_weights.items():
    #     print(name, weight.shape)

    # 从上面的打印结果中查找对应的权重名称并抽取出对应的值，也可以写个 for 循环，逐个修改
    weight_name = '$LM_Head_name or $Embedding_name'
    embedding_layer_weight = safetensors_weights[weight_name]

    # 待替换的 Special Token
    # # 注意肯定是不存在于 tokenizer_base 中但新定义到 tokenizer_new 中的新 Token
    # # 这种 Special Token 使用第一个 tokenizer_base 编码一般是编码为几个 Token 的组合，但使用 tokenizer_new 时一般是编码为一个整体（跟新修改的 ID 对齐）
    new_special_tokens = [3,5,18] # 这里只是一个示例，需要修改为自己的；也可以使用 文本而不是 ID，下面的脚本对应修改即可

    for new_token_id in new_special_tokens:
    	token_text = tokenizer_new.decode([new_token_id]) # 抽取原始文本，也可提前定义
    	old_token_ids = tokenizer_base(token_text).input_ids # 使用旧的 tokenizer 编码新的样本
    	# print(token_text, new_token_id, old_token_ids) # 可打印看看 Token 对应情况
    	new_embedding = torch.mean(safetensors_weights[weight_name][old_token_ids], dim=0) # 按照旧 Token ID 检索向量并做平均
    	# print(new_token_id, "old == new: ", torch.equal(embedding_layer_weight[new_token_id], new_embedding)) # 打印日志，纯新的 Token 不应该相等，旧的则应该相等（因为旧的两者 ID 都只有一个且相同）
    	safetensors_weights[weight_name][new_token_id] = new_embedding # 核心代码，替换 Token embedding

    # 将 safetensors 对象写入新的目标路径
    save_file(safetensors_weights, st_path_new)

附录：新增 Special Token 初始化的讨论

  • 纯新增的 Special Token 确实可以考虑使用一些特殊的初始化操作，比如使用 Embedding 的均值
  • 非纯新增的 Special Token (使用预留的 <mask_1> 等转换而来)， 理论上在预训练中会被训练到：
    • 即使 <mask_1>等 token 永远没出现在语料库中，也会参与训练，因为训练时是整个词表参与的，Softmax 保证了每次都会训练到所有 Embedding
    • 理解：为了压低语料库中不存在的 Token （如 <mask_1>）出现的概率（监督学习目标输出其他正常 Token），也需要学习 <mask_1> 等 Token 的 Embedding 的
  • 思考：使用预留的 <mask_1> 等转换 Special Token，特别是 Special Token 没有太明确的明文语义时，不应该随便替换 Special Token 的 Embedding
    • 因为替换后得到的可能是很奇怪的无语义 Embedding，反而丢失了原始 <mask_1> 在预训练中学到的 Embedding（这个 Embedding 可以保证这个 Special Token 以较低的概率出现）
    • 进一步的理解：如果替换 Embedding，也不建议使用原始 Token 文字使用旧 词表编码的结果
      • 因为特意设计作为非常特殊用途的，非语义的 Special Token，理论上就是 Special 的，一般来说不应该含有语义（其文本形式应该可以为任意值），这时候使用 旧词表的编码结果来初始化是奇怪的，打破了这种特殊的 Special Token 没有语义的设定

本文主要介绍LLM的Tokenization，sentencepiece的使用

打印词表编码

  • 代码

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    import sentencepiece as spm

    def export_vocab_to_file(model_path, output_file):
        # 加载 SentencePiece 模型
        sp = spm.SentencePieceProcessor()
        sp.load(model_path)

        # 打开文件准备写入
        with open(output_file, 'w', encoding='utf-8') as f:
            # 遍历词汇表中的每一个词汇和其索引
            for piece_id in range(sp.get_piece_size()):
                piece = sp.id_to_piece(piece_id)
                score = sp.get_score(piece_id)
                # 将词汇和其索引（或分数）写入文件
                f.write(f'{piece_id}\t{piece}\t{score}\n')

    export_vocab_to_file('chinese_sp.model', 'vocab.txt')

  • vocab.txt文件格式

    1
    2
    3
    4
    5
    6
    7
    8
    9
    0	<unk>	0.0
    1	<s>	0.0
    2	</s>	0.0
    3	,	-2.814912796020508
    4	▁	-3.7144806385040283
    5	。	-3.715141534805298
    6	的	-3.7526743412017822
    7	、	-4.614748001098633
    ...

已有词表扩展

  • 代码如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    import sentencepiece as spm
    import os
    os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"]="python"
    from sentencepiece import sentencepiece_model_pb2 as sp_pb2_model

    chinese_sp_model_path = './chinese_sp.model'
    new_chinese_sp_model_path = './new_chinese_sp.model'
    chinese_sp_model = spm.SentencePieceProcessor()
    chinese_sp_model.Load(chinese_sp_model_path)

    print(len(chinese_sp_model))

    chinese_sp_model_mp = sp_pb2_model.ModelProto()
    chinese_sp_model_mp.ParseFromString(chinese_sp_model.serialized_model_proto())
    tokens_set=set(p.piece for p in chinese_sp_model_mp.pieces)

    ## 将特殊字符添加到词表中
    new_pieces = ['#_#_#', '$$%%##']
    for piece in new_pieces:
        if piece not in tokens_set:
            new_p = sp_pb2_model.ModelProto().SentencePiece()
            new_p.piece = piece
            # score越大，匹配的优先级越高，score为无穷小时约等于没有添加
            new_p.score = 0
            # 使用append添加词，token编码ID会累计+1
            chinese_sp_model_mp.pieces.append(new_p)
            # print(new_p)

    ## Save
    with open(new_chinese_sp_model_path, 'wb') as f:
        f.write(chinese_sp_model_mp.SerializeToString())

    ## load
    new_chinese_sp_model = spm.SentencePieceProcessor()
    new_chinese_sp_model.Load(new_chinese_sp_model_path)
    text='''#_#_#,$$%%##'''
    print("Test text:\n",text)
    print(f"Tokenized by Chinese-LLaMA tokenizer:{new_chinese_sp_model.tokenize(text)}")

sentencepiece 包使用总结

  • sentencepiece包中对分词模型的存储只需要一个文件xx.model
  • 分词模型存储文件中包含了词表信息，包括了piece,score等，同时按照顺序存储，也暗含了词表的编码信息
    • 排序按照score从大到小，score一般小于等于0
  • 想要在词表中增加编码，可以直接构建SentencePiece类对象，并添加到词表反序列化后的尾部即可
    • SentencePiece类对象包含两个关键信息piece,score

LightEval 工具

  • LightEval 是一个专注于 LLM 评估的开源工具库
  • 它提供了标准化的评估框架，支持模型性能测试、对比分析及结果可视化，帮助开发者更高效地衡量模型能力
  • 以下是 LightEval 主要能力：
    • 多维度评估指标 ：覆盖知识理解、推理能力、语言生成质量等多个维度
    • 丰富测试数据集 ：内置多个公开测试集，并支持自定义数据集扩展
    • 模型兼容性 ：支持主流LLM模型的直接接入与评估
    • 结果可视化 ：提供直观的图表展示，便于分析模型优势与不足

LightEval 代码示例

  • 一个简单的代码示例 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    from lighteval import Evaluator, load_dataset

    # 1. 加载评估数据集（内置或自定义）
    dataset = load_dataset("ceval")  # 加载C-Eval中文能力测试数据集

    # 2. 初始化评估器（支持多个评估维度）
    evaluator = Evaluator(metrics=["accuracy", "f1_score", "perplexity"])

    # 3. 准备待评估模型（示例使用Hugging Face模型）
    from transformers import AutoModelForCausalLM, AutoTokenizer
    model_name = "gpt2"  # 替换为实际模型
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    # 4. 定义预测函数（根据模型输入输出格式调整）
    def predict_fn(question):
        inputs = tokenizer(question, return_tensors="pt")
        outputs = model.generate(**inputs, max_length=100)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)

    # 5. 执行评估
    results = evaluator.evaluate(
        dataset=dataset,
        model_predict_fn=predict_fn,
        batch_size=8,  # 可调整参数
        verbose=True
    )

    # 6. 查看并可视化结果
    print("评估结果:", results)
    evaluator.plot_results()  # 生成评估报告图表

评估工具-OpenCompass

  • OpenCompass 是上海人工智能实验室开源的大模型评测平台，涵盖学科、语言、知识、理解、推理等五大评测维度，可全面评估大模型能力
  • OpenCompass 支持在线查看榜单，在线参与评测竞技
  • OpenCompass 还开源了大模型评测工具，使用很简便，GitHub 地址：OpenCompass 项目

本文主要介绍基于 Ollama 的 DeepSeek-R1 本地部署

安装 Ollama

  • 在 ollama官网 下载安装即可
  • ollama 是一个类 Docker 的大模型管理工具

安装 DeepSeek-R1 模型

  • 安装命令（以下命令会自动安装7B版本，即DeepSeek-R1-Distill-Qwen-7B）

    1
    ollama run deepseek-r1

  • 更多镜像可参考：ollama.com/library

基于 ChatBox 的可视化交互

  • 下载并安装：chatboxai.app
  • 点击设置选择指定本地模型即可启动

ollama API 调用

  • 参考链接：

  • API 调用 Demo

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    def http_api_demo():
        import requests
        import json

        # Ollama 服务器的地址
        OLLAMA_URL = "http://localhost:11434/api/generate"

        # 要调用的模型名称
        MODEL_NAME = "deepseek-r1:latest"

        # 要发送的提示文本
        prompt = "你好，DeepSeek！请求解方程x^3+x^2+x-3=0"

        # 请求的 payload
        payload = {
            "model": MODEL_NAME,
            "prompt": prompt,
            "stream": False  # 设置为 False 以获取完整的响应
        }

        # 发送 POST 请求
        response = requests.post(OLLAMA_URL, json=payload)

        # 检查响应状态码
        if response.status_code == 200:
            # 解析响应内容
            response_data = response.json()
            print("模型响应:", response_data.get("response"))
        else:
            print(f"请求失败，状态码: {response.status_code}")
            print("响应内容:", response.text)

    def ollama_api_stream_demo():
        import requests  # 使用 requests 库调用 Ollama 的 API
        import json

        # Ollama 的 API 地址
        url = "http://localhost:11434/api/chat"

        # 请求数据
        data = {
            "model": "deepseek-r1:latest",  # 使用的模型
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业人士，每次回答前请先说“解：”"
                },
                {
                    "role": "user",
                    "content": "9.9和9.11哪个更大?"
                }
            ],
            "stream": True  # 启用流式响应
        }

        # 发送 POST 请求
        response = requests.post(
            url,
            json=data,
            stream=True  # 启用流式接收
        )

        # 打印结果
        print("模型返回的内容：")
        for line in response.iter_lines():
            if line:  # 过滤掉空行
                # 解析 JSON 数据
                chunk = json.loads(line.decode('utf-8'))
                if "message" in chunk and "content" in chunk["message"]:
                    print(chunk["message"]["content"], end='', flush=True)  # 逐步打印内容

    def ollama_api_demo():
        import requests  # 使用 requests 库调用 Ollama 的 API
        import json

        # Ollama 的 API 地址
        url = "http://localhost:11434/api/chat"

        # 请求数据
        data = {
            "model": "deepseek-r1:latest",  # 使用的模型
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业人士，每次回答前请先说“解：”"
                },
                {
                    "role": "user",
                    "content": "9.9和9.11哪个更大?"
                }
            ],
            "stream": False  # 禁用流式响应
        }

        # 发送 POST 请求
        response = requests.post(
            url,
            json=data
        )

        # 检查响应状态
        if response.status_code == 200:
            # 解析 JSON 数据
            result = response.json()
            if "message" in result and "content" in result["message"]:
                print("模型返回的内容：")
                print(result["message"]["content"])  # 打印完整内容
        else:
            print(f"请求失败，状态码：{response.status_code}")
            print(response.text)  # 打印错误信息

    def openai_api_demo():
        from openai import OpenAI
        client = OpenAI(
            base_url='http://localhost:11434/v1/',
            # required but ignored
            api_key='ollama',
        )
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    'role': 'user',
                    'content': '9.9和9.11哪个更大?',
                },
                {
                    'role': 'system',
                    'content': '你是一个专业人士，每次回答前请先说“解：”',
                }
            ],
            model='deepseek-r1:latest',
            temperature=0.0,  # 可以根据需要调整温度值，决定生成的随机性程度
        )
        # 打印结果
        print("模型返回的内容：")
        print(chat_completion.choices[0].message.content)


    def openai_api_stream_demo():
        from openai import OpenAI
        client = OpenAI(
            base_url='http://localhost:11434/v1/',
            # required but ignored
            api_key='ollama',
        )

        # 启用流式响应
        stream = client.chat.completions.create(
            messages=[
                {
                    'role': 'user',
                    'content': '9.9和9.11哪个更大?',
                },
                {
                    'role': 'system',
                    'content': '你是一个专业人士，每次回答前请先说“解：”',
                }
            ],
            model='deepseek-r1:latest',
            temperature=0.0,  # 可以根据需要调整温度值，决定生成的随机性程度
            stream=True,  # 启用流式响应
        )

        # 打印结果
        print("模型返回的内容：")
        for chunk in stream:
            if chunk.choices[0].delta.content:  # 检查是否有内容
                print(chunk.choices[0].delta.content, end='', flush=True)  # 逐步打印内容


    # openai_api_demo()
    # openai_api_stream_demo()
    ollama_api_stream_demo()
    # ollama_api_demo()

  • 多线程调用 Demo

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    import requests
    import json
    from concurrent.futures import ThreadPoolExecutor, as_completed

    def ollama_api_demo(message):
        # Ollama 的 API 地址
        url = "http://localhost:11434/api/chat"

        # 请求数据
        data = {
            "model": "deepseek-r1:latest",  # 使用的模型
            # "model": "llama3.2:latest",
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业人士，每次回答前请先说“解：”"
                },
                {
                    "role": "user",
                    "content": message
                }
            ],
            "stream": False  # 禁用流式响应
        }

        # 发送 POST 请求
        response = requests.post(
            url,
            json=data
        )

        # 检查响应状态
        if response.status_code == 200:
            # 解析 JSON 数据
            result = response.json()
            if "message" in result and "content" in result["message"]:
                return result["message"]["content"]  # 返回完整内容
        else:
            return f"请求失败，状态码：{response.status_code}\n{response.text}"  # 返回错误信息

    def parallel_ollama_api_demo(messages):
        with ThreadPoolExecutor() as executor:
            # 提交任务到线程池
            futures = [executor.submit(ollama_api_demo, message) for message in messages]

            # 等待所有任务完成并获取结果
            results = []
            for future in as_completed(futures):
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    results.append(f"任务执行出错: {e}")

        return results

    if __name__ == "__main__":
        # 示例消息列表
        messages = [
            "9.9和9.11哪个更大?",
            # "Python 和 Java 哪个更适合初学者?",
            # "解释一下量子计算的基本概念。"
        ]

        import time
        x = time.time()
        # 并行调用 API
        results = parallel_ollama_api_demo(messages)

        # 打印结果
        for i, result in enumerate(results):
            print(f"结果 {i+1}:")
            print(result)
            print("-" * 40)
        y = time.time()
        print("time(s):", y-x)

  • 如果想要设定 temperature 等参数，可以在结构体中增加 option 参数，按照字典传入即可，比如，可以如下实现仅返回长度为 20 的 token

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    def ollama_api_demo():
            import requests  # 使用 requests 库调用 Ollama 的 API
            import json

            # Ollama 的 API 地址
            url = "http://localhost:11434/api/chat"

            # 请求数据
            data = {
                "model": "deepseek-r1:latest",  # 使用的模型
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个专业人士，每次回答前请先说“解：”"
                    },
                    {
                        "role": "user",
                        "content": "9.9和9.11哪个更大?"
                    }
                ],
                "options": {
                    "temperature": 0.7,
                    "top_p": 0.8,
                    "num_predict": 20 # 仅返回20个 token 长度
                },
                "stream": False  # 禁用流式响应
            }

            # 发送 POST 请求
            response = requests.post(
                url,
                json=data
            )

            # 检查响应状态
            if response.status_code == 200:
                # 解析 JSON 数据
                result = response.json()
                if "message" in result and "content" in result["message"]:
                    print("模型返回的内容：")
                    print(result["message"]["content"])  # 打印完整内容
            else:
                print(f"请求失败，状态码：{response.status_code}")
                print(response.text)  # 打印错误信息

整体说明

  • SGLang 全称是 Structured Generation Language，是由 LMSYS Org 发起的开源项目
  • SGLang 通过共同设计后端运行时和前端语言，使用户与模型的交互更快、更可控
  • 采用 RadixAttention 技术，通过基数树管理键值缓存（KV Cache），支持多轮对话中共享前缀的缓存复用，在多轮任务中可将缓存命中率提升 3-5 倍，显著降低延迟
  • SGLang 的前端采用编译器式设计，通过领域特定语言（DSL）简化复杂任务编程，后端运行时优化调度和资源分配，还可通过正则表达式和有限状态机（FSM）实现约束解码，直接生成 JSON 等结构化数据
  • SGLang 更适合处理复杂任务，如多轮对话、规划、工具调用（如调用 API 或数据库）等，以及需要生成 JSON、XML 等结构化数据的任务，如智能客服、数据分析等
    • 在 Llama-7B 多轮对话任务中，吞吐量比 vLLM 高 5 倍，延迟降低 30%-50%
  • 注：SGLang 接口经常变化，导致不同版本对应的接口不可复用，非常麻烦！

安装 SGLang

  • 通过 pip 安装 
    1
    2
    3
    4
    5
    pip install --upgrade pip
    pip install "sglang[all]"

    # Install FlashInfer CUDA kernels
    pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/

部署服务

  • 使用命令行启动 
    1
    python -m sglang.launch_server --model-path ~/llm/model/Qwen3-0.6B --port 30000

请求服务

  • 注：下面的命令暂未考虑 Qwen 模型的 Chat 模版

  • 使用命令行访问

    1
    2
    3
    4
    5
    6
    7
    8
    9
    curl http://localhost:30000/generate \
    -H "Content-Type: application/json" \
    -d '{
        "text": "白日依山尽,",
        "sampling_params": {
        "max_new_tokens": 16,
        "temperature": 0
        }
    }'

  • 使用 OpenAI 兼容的 API 访问

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    import openai
    client = openai.Client(
        base_url="http://127.0.0.1:30000/v1", api_key="EMPTY")

    # Text completion
    response = client.completions.create(
        model="default",
        prompt="The capital of France is",
        temperature=0,
        max_tokens=32,
    )
    print(response)

    # Chat completion
    response = client.chat.completions.create(
        model="default",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant"},
            {"role": "user", "content": "列出三个国家和他们的首都"},
        ],
        temperature=0,
        max_tokens=64,
    )
    print(response)

    # Text embedding，需要在服务启动命令中添加 --is-embedding 参数才能访问下面的接口
    response = client.embeddings.create(
        model="default",
        input="How are you today",
    )
    print(response)