tokenizers v1:编码、解码与扩展性能实测
长期以来,分词器(tokenizer)并非机器学习工作流中的瓶颈。从计算开销来看,相较于流水线中其他重型建模环节,分词过程十分轻量。然而在某些场景下,它却迅速成为加速(或拖慢)机器学习任务的关键因素。
随着模型运行速度提升和任务规模扩大,这种平衡开始发生倾斜。在海量数据集上训练、处理高并发请求,或反复处理长文本输入时,分词器面临的压力足以让模型因等待数据而陷入“饥饿”状态。
正因如此,我们在即将发布的 tokenizers 1.0 版本中将性能作为核心重点。分词应当保持轻量,并能随工作流规模灵活扩展。GPU 不应因等待 CPU 完成分词而空转。
本文将介绍 v1 相比 v0.23 更快(通常快数十倍)的原因。
这一成果的取得离不开整个生态系统的支持。分词是开源社区非常活跃的领域,gigatoken、tiktoken、kitoken、tokie、fastokens、wordchipper 和 ai-tokenizer 等库,以及其他众多项目,都在不断拓宽快速分词器的边界。我们研读了这些工作,下文中提到的多项创意,正是因为其他项目证明了其价值才引入到本方案中。
在本次重构之前,tokenizers 的性能远未达到其潜力,因此参与贡献可能显得意义不大。通过此次重构,我们希望明确表达:tokenizers 是一个值得投入贡献的库。
同时,我们也感谢 IBM、NVIDIA 和 ExecuTorch 团队提交补丁,并协助我们在广泛的硬件平台上进行测试,从而扩大了平台兼容性。
结果展示
我们将 tokenizers v1 发布候选版的性能与其他常用替代方案进行了对比,涵盖单线程、多线程、线程扩展性、各模型对比、各语言对比、延迟、解码吞吐量、内存堆占用以及 crate 大小等多个维度。
我们在 tokbench 仓库中运行这些测试,并提供了相应命令,方便你在自己的硬件上复现基准测试。
V1 是什么
v1 的输出 token ID 与 v0.23 完全一致。我们的目标是保持输出、API、词表和合并优先级不变,同时改进一切可以改进的地方——包括适用范围。这个库仍然通用于各类 tokenizer 家族,而不只是专注于 BPE,因此 v1 能加载 v0.23 支持的所有格式。
tokenizer 的作用是把文本转换为模型读取的整数列表。tokenizers 的转换分四个阶段:Normalization 对原始文本执行小写化或 Unicode 规范化等操作;Pre-tokenization 把文本切分成更小的片段,即 pre-token;Model 阶段将每个 pre-token 转成 token,并映射为词表中的 ID;Post-processing 则添加模型所需的特殊 token。
本文描述的大部分优化工作都发生在 Model 阶段。文中测量的十个模型家族里有八个使用 byte pair encoding(BPE)。BPE 从 pre-token 的字节出发,反复合并优先级最高的相邻字节对,直到不存在可合并的字节对为止。优先级在训练 tokenizer 时学得并随模型一同分发,因此同样的文本总能产生相同的 ID。合并永远不会跨越 pre-token 边界。另外两个家族使用 WordPiece 和 Unigram,这是库支持的其他两种模型类型。
四阶段流程详见 tokenization pipeline 页面,BPE、WordPiece 和 Unigram 则在 Tokenization algorithms 中有说明。
每个阶段都做了优化,其中比较关键的变化有:
| 变更 | 作用 |
|---|---|
| workspace 拆分 | 原来的单个 crate 拆成了一个 workspace:tk-encode 是必需的运行时,tk-serialize、tk-convert 和 tk-train 只在应用需要时才链接 |
| 无内存分配的 model | 合并操作的工作数据存放在调用方提供的临时缓冲区中,循环过程完全不触碰内存分配器 |
| bitcannon | 拆分模式转变为位流的布尔运算,使用 SIMD 指令替代正则引擎来寻找拆分位置 |
| 合并循环重写 | 待合并的片段在一个预分配缓冲区内部形成一个侵入式双向链表,因此合并操作只需更新两个索引,无需移动数据 |
| 单词缓存 | 线程局部的备忘录,映射从预分词字节到最终 ID,使重复出现的单词只合并一次 |
| 原生并行 | 一个共享的分词器可同时从多个线程进行编码;每个线程从自己的子池获取临时缓冲区和单词缓存,从而线程不再在单个锁上排队等待(#2365) |
拆分:用位流替代正则表达式
BPE 模型使用正则表达式将输入文本拆分为更小、更易于处理的块,称为预分词(pre-tokens)。合并操作发生在单个预分词内部,绝不会跨越两个预分词之间的边界,因此拆分方式决定了后续流程看到的内容。
该正则表达式是模型的固定参数,随分词器一起分发,运行时永不更改,因此无需在每次编码时都用通用正则引擎进行解释。针对特定模型实际使用的模式,可以手写一个等效的拆分函数,只需编写一次。
手写的函数可以利用现代 CPU 的 SIMD(单指令多数据)指令,它一次对多个字节执行相同操作,非常适合处理 UTF-8 文本。bitcannon 将输入字节视为并行的位流,使得边界通过整个寄存器的布尔运算直接得出,而非逐字符扫描。每个寄存器操作即可判定 64 个字节。同样理念也应用于文本处理的 Parabix 和 JSON 处理的 simdjson。
这依赖于识别特定模式。大多数字节级 BPE 模型由少数几种语法涵盖,若分词器的模式不在其中,则保留正则路径,无法获得上述加速效果。这正是上述收益差异巨大的原因。
单词缓存
实际文本中包含大量重复词语。由于 BPE 对特定预分词始终生成相同的 token ID,v1 在处理完一次后即可保存结果。线程本地缓存将每个预分词的字节映射到其对应的 token ID,使得后续出现时可以直接跳过合并过程。 随着输入数据量的增长,唯一单词的数量增长速度通常慢于总单词数,因此重复词语在输入中所占比例会越来越高。虽然新词仍会出现,但这解释了下述动画中偶尔出现的未命中情况。tokbench measure prefix-sharing \
--engine pipeline \
--engine hf-tokenizers \
--compare-to pipeline-no-cache \
--corpus agentic_swe
缓存机制在处理包含大量重复预分词的输入时效果最佳。如果输入中重复预分词较少,查找成本可能超过收益,导致命中率低下。
合并循环
主要的性能开销来自 BPE 合并循环。对于每个预分词,该循环需反复查找优先级最高的相邻对并进行合并。此前的实现在每次调用时都分配新内存,并为每个预分词构建新的优先队列。 v1 复用调用方拥有的临时缓冲区,消除了这些重复分配。它使用扁平数组存储符号,并通过数组中的位置关联相邻符号,从而降低合并过程中的更新成本。此外,它还能在单次模型调用中处理一批预分词。 每个候选对也被打包为一个 64 位值,其中合并排名占用高位。这样比较两个候选项就等同于比较两个整数,而“此处不合并”对应可能的最大值,因此循环无需分支即可找到下一个合并操作。方法
基准测试设计的细微差异可能导致分词器性能出现巨大差别。为了保持各引擎之间对比的一致性,我们采用了以下规则。| 规则 | 原因 |
|---|---|
| 单一计时循环 | 所有引擎运行完全相同的循环;不为特定引擎设置快速路径 |
| 加载不计入 | 词表加载单独计时,绝不包含在 encode 时间内 |
| ID 哈希校验 | 输出 id 的 FNV-1a 哈希值必须与基线完全一致 |
| 仅通用单元格 | 中位数仅针对所有引擎均执行并验证过的单元格计算 |
| 每进程完整扫描 | 每次重复都在新进程中启动,并保留每个数据格 |
| 物理核心绑定 | worker 绑定到八个不同的物理核心,绝不使用同核 SMT 兄弟线程 |
| 独立 Jobs | 用独立的 Job 测量主机间的差异 |
反复编码同一篇文档,可能比在同一构建上编码多篇不同文档更快。前者测的是文档已完全缓存时的性能,后者测的是新输入的性能,同时允许之前见过的 pre-token 继续留在缓存中。
这两种情况有时都被称作“热缓存”,但它们测的是不同的工作负载。我们的核心结果采用不同文档的方式,且整个语料库大到根本放不进缓存。Tokenizer 基准测试应当说明自己用的是哪种负载,因为这个选择可能直接主导测试结果。
这些优化的整体成效
在 v1 的 encode 路径覆盖的十个模型家族上,单线程下 Apple M4 Max 的编码速度比 v0.23 快 3 到 30 倍。低端是 t5-base,高端是 gpt2。八个 worker 下的扩展效率达到线性的 76%。在这些改动之后,v1 生成的 token ID 与已发布版本完全一致。
整体性能提升来自多项改动的协同:用手写的 splitter 取代 regex 引擎、用缓存直接返回重复单词而无需再次合并、合并循环完全不触碰内存分配器,以及每批 pre-token 只调用一次模型而非每个 pre-token 调用一次。每一项都在流水线的不同环节减少了工作量。
下一个优先事项是支持更多模型家族。我们会在 1.0.0 之前把更多模型迁移到新的合并循环上。等候选版本稳定后,下一步是将这些改进带入 transformers 库以及生态中其他依赖 tokenizers 库的项目。
本文基于 tokbench 的结果生成,并将随着支持的扩展持续更新。
如何获取
v1 的候选版本已发布在 crates.io。API 完全不变,唯一的区别只是安装哪个构建版本。
安装方式和平时一样:
cargo add tokenizers --pre
训练功能默认开启,且依赖 C++ 库。若只需编码功能,可禁用该特性以排除训练实现:
cargo add tokenizers --pre --no-default-features --features http
编码逻辑保持不变:调用方式相同,返回的 ID 也一致。
use tokenizers::tokenizer::{Result, Tokenizer};
fn main() -> Result<()> {
let tokenizer = Tokenizer::from_pretrained("deepseek-ai/DeepSeek-V4-Flash", None)?;
let encoding = tokenizer.encode("The tokenizer is no longer the bottleneck.", false)?;
println!("{:?}", encoding.get_ids()); // [671, 17840, 9160, 344, 1119, 5827, 270, 111127, 16]
println!("{:?}", encoding.get_tokens()); // ["The", "Ġtoken", "izer", "Ġis", "Ġno", "Ġlonger", "Ġthe", "Ġbottleneck", "."]
Ok(())
}
对于批量处理,encode_batch 利用多核进行扩展,也是前文性能测试的目标函数。
let encodings = tokenizer.encode_batch(documents, false)?;
文中所有数据均基于该 Rust 库实测得出。Python 绑定封装了相同的底层代码,构建于 bindings/python 目录,但其额外的每次调用开销未包含在上述测量中。
迈向 1.0 版本的进展
本文基准测试覆盖了已完成的发布候选版工作,下文将列出 1.0.0 正式版仍需完成的事项及后续探索计划。
发布候选版:已实现功能
以下工作已发布于 crates.io 的 Rust 预发布版本:
cargo add tokenizers --pre
- 工作区拆分:将单一 crate 拆分为
tk-encode、tk-serialize、tk-convert和tk-train,使应用仅链接实际使用的模块 - Bitcannon:在编码路径中用位流操作替代正则表达式分割,支持 GPT-2、cl100k、o200k、Tekken 和 DeepSeek。此变更取代了早期在#2201和#2317中发布的有限状态机
- WordCache:复用先前处理过的预词元 ID#2262,提交
af5a3e3
STAGE_POST 流水线阶段 #2182role_to_token 支持 #23431.0.0
- 统一的编码实现:在训练验证期间使用
tk-encode,确保训练和推理阶段不会产生不同的分词结果 - 可选的 offsets 和 masks:仅在请求时计算这些元数据,使其不占用仅处理 token-ID 的路径
- 重构 normalizers
- 基于 atomnorm 的 bitnorm 支持 #2209
- spm 预编译
- 更简洁的 Python 绑定:在保留子类化、序列化、自定义解码器、变异行为以及无锁 CPython 支持的前提下,减少锁、包装类型和手写分发代码
- 面向 ExecuTorch 和 llama.cpp 的纯推理 C 和 C++ 绑定,未来可能跟进 JVM、Swift 和 Go 绑定
1.0.0 之后
- tok-devices:探索 GPU 编码和批量解码,同时保持文本和 token ID 驻留在设备上。解码器将一次性上传词汇表,并行计算输出位置,并在 GPU 上收集对应的字节。这将是一个面向大批量场景的可选组件,具体方案仍需进一步原型验证和测量。