预处理才是 Tokenization 的性能死角
在双路 144 核 AMD EPYC 9565 服务器上,处理 11.9 GB 的 OpenWebText 数据集,HuggingFace tokenizers 需要耗费接近 8 分钟,Python 社区常用的 tiktoken 需要 5.5 分钟。采用 Rust 纯手写的 Gigatoken 仅需 0.48 秒,单节点编解码吞吐达到 24.53 GB/s。
Gigatoken 在该硬件上实现了 HuggingFace 的 989 倍、tiktoken 的 681 倍吞吐;在 Apple M4 Max 芯片上同样录得 8.79 GB/s,比 HuggingFace 快 1268 倍。过去业界普遍将 Tokenization 视作微不足道的预处理开销,忽视了数据管线在 TB 级预训练中造成的 CPU 阻塞瓶颈。
长期以来,开发者习惯将 BPE 算法中的树形词表合并与二分查找当作性能优化的主要战场。真实数据剖析显示,阻塞文本吞吐的元凶是被普遍视作理所当然的正则表达式预切分阶段。基于 fancy-regex 的基准切分吞吐仅有 47 MiB/s,直接锁定了后续并发处理的上限。
用 SWAR 替代正则引擎的无分支解析
Gigatoken 放弃了通用正则表达式引擎状态机的复杂回溯,改用 SWAR(SIMD Within A Register)技术配合查找表进行字符类扫描。通过在 64 位寄存器内部并行执行位运算与掩码比较,预处理阶段的解析吞吐从 462 MiB/s 提升至 830 MiB/s,获得了 1.8 倍的纯计算收益。在特定规则的文本扫描场景中,硬编码的寄存器级并行效率能够碾压通用正则表达式状态机。
通过 Rust 的 get_unchecked 切断冗余的数组边界检查,吞吐从 830 MiB/s 提升至 840 MiB/s。紧接着将字符指针推进与数量计数逻辑彻底分离,使数据流避开了中间状态的频繁存储,吞吐进一步递增至 848 MiB/s。这些细节优化消除了编译后汇编代码中冗余的比较指令,为硬件流水线的饱和运行扫清了障碍。
图:Gigatoken 在 GPT-2 编码任务中的吞吐量表现。来源:marcelroed/gigatoken GitHub
榨干 CPU 流水线:多指针 ILP 压榨与硬件适配
为了跨越 1 GiB/s 的单核吞吐门槛,Gigatoken 引入了双指针交叉迭代(Dual-cursor ILP)机制。这种设计让单个 CPU 核心在处理数据流时能够同时维护两个独立的解析游标,直接将预处理吞吐从 840 MiB/s 拉升至 1,049 MiB/s,单核性能增幅达 25%。现代 CPU 超标量流水线的多发射端口在单游标循环中往往处于半饥饿状态,指令级并行度的显式压榨能直接填满流水线深处的执行单元。
这种极致的指令优化不仅体现在 AMD EPYC 节点,在 Intel Xeon、Apple M4 Max 以及桌面级的 AMD Ryzen 9800X3D 上均表现出极高的一致性。项目保持了对标准词表的全面支持,涵盖 GPT-2、Llama 3/4、Qwen 2/3、DeepSeek V3/V4、GLM 4/5、Gemma 4 等主流模型。开发者只需通过 pip install gigatoken 即可完成替换,完全兼容 HuggingFace 和 tiktoken 的原有 API 接口。
从数据处理到集群吞吐的工程重构
在集群工程层面,Gigatoken 展示了极强的多核扩展能力。在 144 核双路 EPYC 服务器上,它可以在约 6.5 小时内完成包含 130T tokens 的整个 Common Crawl 数据集的 Tokenization 处理。预训练数据准备从以往按天计算的大规模集群调度任务,缩减到了单节点一班倒的工时范畴,降低了数据流水线运维成本。
整个项目包含 357 次代码提交,底层算法完全由作者手动完成 Rust 架构设计与汇编级调试,AI 仅用于 API 包装与最终的端口迁移。这种回归裸金属性能极致的工程实践,展示了基础基础设施在 AI 时代依然蕴藏着的巨大优化空间。
文本吞吐管线的范式转向
Gigatoken 的意义在于打破了 LLM 工具链性能优化的思维定势。当绝大部分团队仍在关注 BPE 字典查找的树结构时,Gigatoken 证明了预处理层面的正则切分才是拖慢全局管线的真实瓶颈。通过 SWAR 指令级并行与流水线填充分析,它将数据打包效率推向了硬件极限,也为海量文本清洗与前处理基建树立了新的性能标杆。
参考链接:
- marcelroed/gigatoken GitHub 项目仓库
- OpenWebText 基准测试集