配了三年 AI 数据管道,今天才知道分词速度一直是瓶颈——GigaToken 用 1000 倍速把这件事彻底翻了

配了三年 RAG 索引和微调数据准备,你大概已经习惯了这样的流程:丢进去一个十几 GB 的文本语料,然后去泡杯咖啡回来——等 HuggingFace Tokenizers 跑完。

分词(tokenization)这件事,长期被当作「反正够用」的角落。没人觉得它值得优化。Marcel Rød 觉得不行,搞出了 GigaToken,一个 Rust 写的 tokenizer,在 GPT-2 词表上跑出了 24.53 GB/s 的吞吐量——对比 HuggingFace Tokenizers 的 24.8 MB/s,这是 989 倍的差距。

同样的机器(双路 AMD EPYC 9565,144 核),HuggingFace Tokenizers 要花约 8 分钟处理 12GB OpenWebText;GigaToken 不到半秒。

它是怎么做到的

三个关键优化,没有一个是魔法:

SIMD 预分词(pretokenization)

传统方案把正则匹配丢给通用引擎,GigaToken 用了手写的 SIMD(AVX-512/AVX-2/NEON)来硬刚预分词这一步。这部分是 tokenization 里最重的 CPU 密集操作。

预分词缓存

对同一个词,只算一次,下次直接查缓存。整个语料里重复词越多,收益越大。

绕过 Python 开销

原生 API 让 Rust 直接读文件,跨线程零通信。不走 Python 的 GIL,不走 Python 对象序列化。

实际数字

词表 GigaToken HuggingFace 加速比
GPT-2 24.53 GB/s 24.8 MB/s 989×
Qwen 3 22.16 GB/s 34.2 MB/s 648×
DeepSeek V3/R1/V4 19.69 GB/s 26.2 MB/s 750×
Llama 3/3.1/3.2 22.15 GB/s 48.5 MB/s 457×
Phi-4 801×

在 Apple M4 Max(16 核)上,GPT-2 词表跑到 8.3 GB/s,是 HuggingFace 的 1353 倍

SentencePiece 风格的词表(Gemma、Llama 2、Mistral)提速幅度低一些,大约 7-22 倍——这部分还有优化空间。

怎么用

安装就一行:

pip install gigatoken

兼容模式(保留现有代码,换掉底层):

import gigatoken as gt

hf_tokenizer = ...  # 你现有的 HuggingFace tokenizer
tokenizer = gt.Tokenizer(hf_tokenizer).as_hf()

tokens = tokenizer.encode_batch(["你的文本", "再来一段"])

原生模式(榨干性能):

import gigatoken as gt

tokenizer = gt.Tokenizer("Qwen/Qwen3-8B")  # 直接接受 HuggingFace 模型名
file_source = gt.TextFileSource(["corpus.txt"], separator=b"")

tokens = tokenizer.encode_files(file_source)

什么时候用它

RAG 管道:文档分块+tokenize 是 embedding 前的瓶颈。1000 倍速意味着你可以用原来等一杯咖啡的时间把整个知识库重新索引一遍。

微调数据准备:大语料预处理从「过夜任务」变成「喝口水」。

Token 计数预算管理:跨百万文档统计 token 优化上下文窗口,慢 tokenizer 让迭代变成折磨。

局限

  • Windows 测试不充分,建议用 WSL
  • WordPiece 词表暂不支持
  • SentencePiece 优化不到位,速度只有 7-22 倍

下一步

跑一下 benchmark 看看你的词表实际数字:

uvx --with tokenizers gigatoken bench openai-community/gpt2 your_data.txt --validate

GitHub:marcelroed/gigatoken,MIT 许可证。

如果你的数据管道里有任何地方在跑 tokenization,换掉它。989 倍速不是 PPT 数字,是实测结果。

评论区

0 条评论

登录后可评论。

铁锈·Rust工具链 17 阅读