注：本文包含 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

本文主要介绍 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，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

本文主要介绍 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 没有语义的设定

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)

Tulu3

  • 原始论文：(Tulu3)Tülu 3: Pushing Frontiers in Open Language Model Post-Training
  • Tulu3 数据集是艾伦人工智能研究所（Ai2）为训练Tulu3模型而创建的大规模多语言文本数据集
  • Tulu3 数据集包含 939,344 个样本，覆盖多种语言和任务，数据来源广泛，包括 Coconot、Flan v2、No Robots 等
  • Tulu3 数据集支持语言模型的训练和微调，特别是在多语言环境下，其结构包含标准的指令调整数据点，如 ID 等
  • Tulu3 数据集的意义在于为研究人员和开发人员提供了丰富的语言资源，以增强和优化多语言人工智能模型的性能，可用于教育和研究目的，但需遵循特定的许可协议

lmarena-ai/arena-human-preference-140k

  • HuggingFace：huggingface.co/datasets/lmarena-ai/arena-human-preference-140k
  • Blog：A Deep Dive into Recent Arena Data, 20250731
  • lmarena-ai/arena-human-preference-140k 数据集包含文本类别的用户投票（vote）数据，累计大约 14W 数据
    • 每行代表一次投票，记录用户在特定对话场景下对两个模型（model_a 和 model_b）的评判结果，同时包含完整对话历史及元数据
  • 核心字段说明如下：
    • id：每次投票/每行数据的唯一反馈ID
    • evaluation_session_id：每次评估会话的唯一 ID，一个会话可包含多次独立投票/评估
      • 经测试，同一个 evaluation_session_id 对应的模型可能不同，如何理解？
    • evaluation_order：当前投票的评估顺序（序号）
    • winner：对决结果，取值为
      • model_a（模型 A 获胜）
      • model_b（模型 B 获胜）
      • tie（平局）
      • both_bad（两者均差）
    • conversation_a/conversation_b：当前评估轮次中两个模型对应的完整对话内容
    • full_conversation：完整对话历史，包含上下文提示词及所有先前评估轮次的模型回复
      • 注意：每次投票后会重新采样模型，因此完整上下文中的响应模型会有所不同
      • 问题：发现多轮数据中，两个模型的上下文是分离的，所以这部分训练时相当于是多轮信号
    • conv_metadata：聚合元数据（含格式标记、令牌计数），用于风格控制
    • category_tag：标注标签，包括Math、创意写作（creative writing）、高难度提示词（hard prompts）、Instruction Following四类
    • is_code：对话是否涉及代码（布尔值）
  • 其他分析：
    • 大部分数据是英语，约 51.8%
    • 中文简体占比较少，仅 5%
    • 中文繁体占比，仅 0.65%
  • 注：这个数据集是用作 Reward Model 的优秀模型，也可以用于 SFT 和 DPO 等
    • 用作 SFT 时，多轮数据也可以使用；用作 Reward Model 和 DPO 时，多轮数据暂无法使用

本文主要介绍LLM预训练相关的笔记

  • 参考链接
    • [如何从零开始训练大模型（minicpm分享&讨论）](https : //zhuanlan.zhihu.com/p/686664720)：minicpm模型训练分享
    • [LLM训练-Pretrain](https : //zhuanlan.zhihu.com/p/718354385)：非常详尽的一些实操笔记和思考，论文的许多笔记是参考了这篇博客
    • [花费千万试出来的LLM预训练经验](https : //mp.weixin.qq.com/s/MF6MuZC_fvKu25S1MzXkCQ)

预训练的重要性

  • 预训练阶段是知识注入的关键阶段
  • 已有研究基本都证明了，预训练阶段是知识注入阶段，后续的 post-training 阶段主要是激发模型的能力

预训练一个大模型有哪些需要关注的点？

数据爬取和解析

  • 数据需要上 T，甚至几十 T 才够用
  • 爬虫不好做，容易被封 IP
  • 很多论文是 PDF，使用 Python 解析可能不准确，使用 GPT4 等解析成本较高
  • 最好多找开源的数据

数据清洗、去重、配比

  • 清洗很重要
  • 需要自己训练打分器，最好是使用 BERT 系列模型，而不是 Decoder-Only 模型训练打分器，分类器的 label 可以使用 GPT4 标注，或者按照优质源与劣质源作为正负label标记

    需要注意的是，基本上大家都认同：同等 size 下，BERT 结构的模型的表征能力是强于 transformer-decoder 模型的，因此打分模型最好还是从 BERT 家族中选一个来训，效果好、速度还快

  • 除了打分器，规则也很重要
  • 需要做好数据脱敏，避免人名等隐私被侵犯出来
  • 数据去重很重要，网页互相引用会造成爬取到的数据重复性太高，做不了 sentence 粒度的，就做 document 粒度的
  • 可以训练一个文档分类器来对文档进行分类，不同类别使用不同的相似度阈值去重
  • 数据配比一般是 中 : 英 : code = 4 : 4 : 2

数据顺序

  • 同样大小的数据，按照不同顺序输入模型，结果也会不同，基本上可以类似课程学习的思想，想简单后复杂，先优质再劣质等
  • 不同文档的句子之间应该互相看见吗？llama3.1 强调了不能看见，但是作者理解看见其实也还好

流水化数据线

  • 数据需要提前 tokenization
  • 训练和数据处理应该是两个进程
  • 训练过程中多存 save_checkpoint，为了训练时决定是否即时保存，可以使用一些规则，比如“遇到存在 save 文件的文件夹就保存一下”

数据实验

  • 训练前可以先用小模型验证一下，如何设置数据配比，顺序更合适

模型结构

  • RoPE + GQA + RMS_Norm + SwiGLU
  • 预训练不建议创新，试错成本太高，一般来说照抄开源的优秀的技术报告即可

模型大小

  • 按照合适来选择，可以通过时间和数据来评估自己需要多少参数量的模型
  • layer_num 和 hidden_size 同步增加或者减少更好些
  • seq_len 不要上来就很大，RoPE 的 NTK 外推方法已经是各大厂标配的方案：4K/8K + RoPE 小 Base + 90% 数据量 –> 32K/64K + RoPE 大 Base + 10% 数据量

训练框架

  • Megatron 和 DeepSpeed之间选：

    Megatron 和 DeepSpeed 该怎么选？直接说结论：从零开始的 Pretrain 必须选 Megatron，Continue-Pretrain 可以考虑使用 DeepSpeed， 换句话说，T 级别的 token 训练量必须是 Megatron，B 的级别 token 训练量无所谓

  • DeepSpeed 加载速度慢
  • 其他：

    无论是用哪个训练框架，都要记得把 attention 的默认方式换成 flash_attention

训练时评估

  • loss 分不同种类观察，即 channel_loss，比如英文、中文、代码分开观察

训练后评估

  • 可以使用 ACC 来衡量评估结果

概率探针评估

  • 可以用于查看某个目标条件概率是否在增大

一个特别干货的分享

  • 参考链接：[如何从零开始训练大模型（minicpm分享&讨论）](https : //zhuanlan.zhihu.com/p/686664720)， 这个博客分享了包含了一些地数据集处理、清洗等
  • 文中很有意思的一段话

    但我们回到2020年，当大部分人都在基于 bert 做各种魔改的时候
    OpenAI 发现了这么一个规律。数据，训练，参数一直增长下去，好像 loss 的确是在不断的下降哎？
    于是，他们拿着这个 paper 去问微软的 CTO，你想不想看看这个 loss 下降到一定程度会发生什么？
    会发生什么？
    chatgpt 就出来了

LLM 预训练实战经验

  • 以下内容参考自：[花费千万试出来的LLM预训练经验](https : //mp.weixin.qq.com/s/MF6MuZC_fvKu25S1MzXkCQ)

随机初始化的痛点

  • 训练成本高 ：早期使用 Qwen1.5-0.5B 结构，训练 8T+token 通用数据仅达勉强效果，相比 1-2T 训练量提升有限
  • 核心问题 ：能否复用已有模型？直接续训开源模型存在两大问题：
    • 参数量不匹配（如业务需要 3B 模型，开源仅有 7B/4B/1.5B，则还是需要重头训练）
    • 版权风险（如 Qwen2.5-4B 无商用许可），注：Qwen 的同一系列不同模型的开源许可也可能不同

模型初始化：大模型 to 小模型

  • 适用场景 ：当拥有可商用大模型（如14B），需初始化小模型（如3B）
  • 主流方法对比 ：
    方法 核心原理 数据要求 初始化效果 限制条件
    Sheared LLaMA 训练mask筛选重要参数 几B-几十B高质量数据（如代码/数学） 初始loss更低，几十步收敛至2.x 模块类型不可改（激活函数/attention类型固定）
    Weight Subcloning 计算神经元重要度裁剪 需输入数据获取激活值 后期loss略低，评测效果持平 同上
  • 实践效果 ：14B -> 3B 裁剪可行，结合蒸馏后 500B token 可达原 8T 效果
  • 关键结论 ：0.5B-72B 范围内，已有模型初始化效果优于随机初始化

模型初始化：小模型 to 大模型

  • 代表方法 ：Bert2BERT、Llama Pro
  • 适用场景 ：封闭域对话能力提升，训练数据充足
  • 效果预期 ：仅能带来1-2个点的提升，不适合大规模预训练

模型初始化：Dense to Sparse（MoE模型）

  • 初始化挑战 ：随机初始化易出现 loss spike，需 50k+step 恢复
    • 理解：“loss spike”（损失激增）是指模型的损失函数值在训练过程中出现突然且显著的上升现象
  • Sparse Upcycling 方案 ：
    • 1）用同结构 Dense 模型通过 FFN 复制生成 MoE
    • 2）关键操作：
      • 裁剪前打乱 neuron 顺序打破专家对称性
      • 保留 50%-70% 原 FFN 参数，其余随机初始化
    • 3）优化技巧：全单精度训练可提升稳定性，但 GPU 利用率 < 50%

通用预训练数据准备

基础清洗
  • 清洗原则 ：数据量充足时可严格筛选
  • 具体规则 ：
    • 过滤 ppl 异常、格式混乱（多分行/短词）、长度过短数据
    • 排除 url、安全词、重复内容（防 LLM 复读）
    • 结合 fasttext 二分类模型筛除低质量数据
  • 辅助工具 ：
    • 人工打标 + FastText 二分类模型筛选低质量数据
    • 参考 Llama3.1 和 Qwen2 技术报告的清洗流程
去重
  • 技术方案 ：minhash 实现文章/段落粒度去重
  • 执行原则 ：宁杀错不放过，重复数据严重影响模型能力（重复数据会导致模型 “记忆偏差”）
    • 理解：目前数据量已经很大了，可以多删除一些，对于数据量不大的场景，还是要小心错杀现象
分类与配比
  • 数据类型 ：代码、数学、高教育性数据（如 chinese-fineweb-edu ）、通用数据（体育/音乐/时政等）
  • 经验配比 ：提高 education score、代码、数学数据比例，无需过度纠结具体数值（如 30% vs 40%）
  • 关键逻辑 ：通用数据先夯实语言基础，再切入专业领域数据

长文本预训练数据被准

  • 数据要求 ：几B-几十B 即可
  • 质量排序 ：
    • 1）天然长文本（大学课本、GitHub 项目）【效果最好】
    • 2）相关文档拼接（如带 reference 的论文）【效果次之】
    • 3）无关文档拼接（仅训练位置编码）【这个方法效果不好，不建议使用】
  • 注意事项 ：避免领域分布集中（避免数据 Bias），构造阅读理解类长文本效果有限

退火阶段数据准备

合成数据核心原则

  • 多样性优先 ：

    • 方法1：调整解码参数（温度/top k/top p）
    • 方法2：Prompt 注入随机变量（如腾讯 Persona Hub 的 10亿 人物描述）
    • 示例：生成数学题时，在 Prompt 中加入不同人物背景描述 
      1
      2
      3
      4
      5
      6
      7
      8
      9
      # 数学题生成示例
      prompt  =  """
      根据以下人物描述生成生活场景数学题：{persona}
      """
      persona_list  =  [
      "23岁卡车司机，身高178cm，单身，爱吃海鲜...",
      "北京化学家小A，本科就读于...",
      # 更多个性化描述
      ]

  • 质量保障 ：

    • 模型规模：越大越好（专用模型 > 通用模型）
    • 质量评估：LLM-as-judge 打分、代码可执行性验证（如 OpenAI 工具）
数据类型与来源
  • 推理能力 ：数学/代码数据（大量合成）
  • 知识储备 ：educational data（学科试题/教科书）
  • 语言能力 ：高质量网页数据（严格阈值筛选）

训练优化策略-蒸馏策略

  • 适用场景 ：大模型初始化的小模型训练
  • 实施要点 ：
    • teacher 模型与初始化大模型同系列
    • 温度策略：先高后低（前期广覆盖，后期精收敛）
  • 效率优势 ：训练 token 数可减少至直接训练的 10% 以下

超参迁移优化-MuP参数化

  • MuP，即 Maximal Update Parameterization
  • 核心价值 ：
    • 设计模型超参实现大小模型 LR/BS 通用
    • 小模型调参与大模型训练超参一致性强
  • 实践应用 ：
    • 小模型（如 10B 数据）网格搜索调参，直接迁移至大模型
    • 搜参过程中关注 loss 下降速度而非绝对数值

学习率策略-WSD vs Cosine

  • Cosine 衰减（余弦衰减） ：
    • 需预先设定总步数，收敛效果更稳定
  • WSD（Warmup-Stable-Decay）学习率调度 ：
    • 分为预热阶段、稳定阶段和衰减阶段
    • 优势：灵活调整阶段，适合实验场景，无需提前设定总步数
      • 虽然平时也常将预热阶段设置为总步数的一定百分比，但也可以设定为固定值，且稳定阶段的步数一般没有特定要求
    • 劣势：同 token 数下效果不如 Cosine（需足够长 Decay 阶段 + 低最终 LR）

多阶段退火训练（提升效果的核心模块）

  • Step1：学科能力提升

    • 数据比例：教材 30%+合成选择题 50%+通用选择题 10%+通用数据 10%
    • 关键参数：LR = 3e-4，训练量 120B+ token（最佳 checkpoint 在几十 B 时达成）
    • 效果：MMLU/Ceval 超越 Qwen 官方模型

  • Step2：加入数学数据

    • 数据比例：学科 40% : 数学 60%
    • 关键发现：
      • LR = 8e-5 时，效果优于 1e-4
      • 数学数据超过 60% 会导致学科能力下降
    • 效果：Ceval比基线高6分，学科能力稳定

  • Step3：加入代码数据

    • 数据比例：学科 20% : 数学 40% : 代码 40%
    • 关键参数：LR = 5e-5，训练量几十B token
    • 代码和数学关联效应：数学与代码能力正相关，无需过高代码比例

  • Step4：综合能力提升

    • 数据比例：学科15% : 数学 25% : 代码 25% : 通用 10% : SFT 30%
    • 关键参数：LR = 3e-5 衰减至0 ，训练量几十B token
    • 训练目标：
      • 保持语言能力 + 引入 SFT 数据（如 Tool Use 的 function call）
    • 数据特性：加入 Tool Use 等 SFT 数据

关键原则

  • 数据不重复使用，预留充足退火数据
  • 小阶段训练顺序：学科 > 数学 > 代码 > 通用（按最佳学习率排序）
  • 单一数据训练时需优化内部配比（如代码数据中 Python 占比、GitHub 项目比例等）

Batch 内数据分配

  • 实时跟踪各类数据 token 数，严格按预定比例混合
  • 避免某类数据在 batch 中占比波动过大

评测数据集

  • 通用能力：opencompass（代码/数学/语言/知识）
  • 下游任务：业务相关数据集（如 Tool Use 评测集）

评测结果分析经验

  • loss 与效果关系 ：loss 低不一定效果好，但 loss 高效果一定差
  • 通用与下游平衡 ：两者指标并非总是正相关，需针对性优化

本文主要介绍 TRL 库的使用

整体说明

  • TRL（Transformer Reinforcement Learning）是 huggingface 中的一个完整的库，用于微调和对齐大型语言模型，可用于优化 Transformer 语言和扩散模型
  • 这个库支持 SFT、PPO、DPO 等模型微调、对齐流程
  • TRL 库目前被很多开源框架依赖，是 LLM 领域的标准基础框架
  • TRL 支持了很多开源的微调方法，而且还在持续更新，详情见：huggingface.co/docs/trl 的 API 部分
  • TRL 集成了很多底层框架
    • 待补充

安装 TRL 库

  • 通过 pip 安装 TRL 库：

    1
    pip install trl

  • 也可以通过 Git 克隆并直接从源代码安装：

    1
    2
    3
    git clone https://github.com/huggingface/trl.git
    cd trl
    pip install .
    • 一些未发布功能和修复在最新版本里面，此时需要通过上述源码方式安装
    • 安装方式说明：
      • pip install . 会安装并复制文件到默认目录（标准稳定版 pip 包安装）
      • pip install -e . 则会创建链接到当前目录，同时对当前的目录会立即生效到包上（--editable）

SFT 示例

  • 参考链接：huggingface.co/docs/trl/sft_trainer
  • SFTTrainer用于在自定义数据集上进行监督微调。示例代码如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    from datasets import load_dataset
    from trl import SFTConfig, SFTTrainer
    from transformers import AutoModelForCausalLM, AutoTokenizer

    # 定义基础模型
    model = AutoModelForCausalLM.from_pretrained("模型名称")
    tokenizer = AutoTokenizer.from_pretrained("模型名称")

    # 加载训练数据集
    dataset = load_dataset("trl-lib/Capybara", split="train")
    # 配置训练参数
    training_args = SFTConfig(output_dir="Qwen/Qwen2.5-0.5B-SFT")
    # 初始化SFTTrainer
    trainer = SFTTrainer(
    	model=model,  # 基础模型
    	tokenizer=tokenizer,  # 对应的tokenizer
    	args=training_args,
    	train_dataset=dataset,
    )
    # 开始训练
    trainer.train()

奖励模型（Reward Model）训练示例

  • RewardTrainer用于训练奖励模型，该模型可以评估文本生成的质量。示例代码如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    from trl import RewardConfig, RewardTrainer
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    from datasets import load_dataset

    # 加载预训练的模型和分词器
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
    model = AutoModelForSequenceClassification.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", num_labels=1)
    # 加载适合于奖励模型的数据集
    dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
    # 配置训练参数
    training_args = RewardConfig(output_dir="Qwen2.5-0.5B-Reward", per_device_train_batch_size=2)
    # 初始化RewardTrainer
    trainer = RewardTrainer(
    	model=model,
    	tokenizer=tokenizer,
    	train_dataset=dataset,
    	args=training_args
    )
    # 开始训练
    trainer.train()

PPO 训练示例

  • PPOTrainer用于基于近端策略优化算法对语言模型进行微调。示例代码如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    import torch
    from transformers import AutoTokenizer
    from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model
    from trl.core import respond_to_batch

    # 加载预训练的模型和分词器
    model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
    model_ref = create_reference_model(model)
    tokenizer = AutoTokenizer.from_pretrained('gpt2')

    # 初始化PPO训练器
    ppo_config = PPOConfig(batch_size=1)

    # 编码一个查询
    query_txt = "This morning I went to the "
    query_tensor = tokenizer.encode(query_txt, return_tensors="pt")

    # 获取模型响应
    response_tensor = respond_to_batch(model, query_tensor)

    # 创建PPO训练器
    ppo_trainer = PPOTrainer(ppo_config, model, model_ref, tokenizer)

DPO 训练示例

  • DPOTrainer用于根据人类偏好直接优化语言模型。示例代码如下： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import DPOTrainer
    import datasets

    # 加载模型和分词器
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
    model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
    # 加载数据集
    dataset = datasets.load_dataset("trl-lib/ultrafeedback_binarized", split="train")

    # 初始化DPOTrainer
    trainer = DPOTrainer(
    	model=model,
    	tokenizer=tokenizer,
    	train_dataset=dataset,
    )

    # 开始训练
    trainer.train()

GRPO 训练示例