GigaToken:實現 1000 倍速的語言模型 Tokenization

GigaToken:實現 1000 倍速的語言模型 Tokenization

GigaToken 是一個高性能的 tokenization 函式庫,旨在作為 HuggingFace 和 Tiktoken 的直接替代方案,在高端 CPU 硬體上可達到高達 24 GB/s 的吞吐量。它顯著減少了處理海量數據集所需的時間,在雙插槽 AMD EPYC 系統上,有可能在不到 6.5 小時內完成整個 Common Crawl 數據集(約 130 兆個 tokens)的 tokenization。

技術實現與效能提升

// GigaToken 通過專注於 tokenization 過程中計算最密集的部分(特別是 pretokenization 和 caching)來實現其速度。

SIMD 與 Pretokenization

大多數 tokenizer 將 pretokenization 外包給 Regex 引擎。GigaToken 使用經過高度優化的實現來取代此機制,利用 SIMD (Single Instruction, Multiple Data) 並最小化分支 (branching)。這種方法允許該函式庫在 x86 和 ARM 架構上都能以極高的效率在位元組層級處理文本數據。

優化的快取層級 (Optimized Cache Hierarchies)

快取 pretoken 映射是關鍵的效能瓶頸,因為 pretoken 分佈呈長尾效應,且快取會迅速增長。GigaToken 實現了專門的快取層級,以高效查找已出現單詞的編碼 tokens,從而減少冗餘計算。

最小化執行時開銷 (Minimizing Runtime Overhead)

為了最大化吞吐量,GigaToken 最小化了與 Python runtime 的交互,並避免了執行緒之間的通信。通過在 Rust 中實現核心邏輯並允許 Rust 實現直接從文件讀取數據,它繞過了與傳遞 Python 數據結構相關的開銷。

基準測試與硬體效能

當使用其原生 API 時,GigaToken 的效能提升最為顯著,該 API 允許最大程度的並行處理和直接文件讀取。

高端伺服器硬體 (AMD EPYC 9565)

在雙插槽 144 核心 AMD EPYC 9565 處理器上,GigaToken 對於常見 tokenizer 實現了以下吞吐量:

Tokenizer GigaToken Throughput vs HuggingFace
GPT-2 24.53 GB/s 989x
Phi-4 24.00 GB/s 801x
Llama 3 / 3.1 / 3.2 22.15 GB/s 457x
DeepSeek V3 / R1 / V4 19.69 GB/s 750x
Qwen 2 / 2.5 19.12 GB/s 19.12 GB/s

消費級硬體 (Apple M4 Max)

在 16 核心 Apple M4 Max 上,GigaToken 展示了類似數量級的加速:

Tokenizer GigaToken Throughput vs HuggingFace
GPT-2 8.79 GB/s 1,268x
OLMo 2 / 3 7.56 GB/s 1,299x
Llama 3 / 3.1 / 3.2 7.60 GB/s 676x
Phi-4 7.76 GB/s 1,012x

桌上型硬體 (AMD Ryzen 7 9800X3D)

在 16 核心 AMD Ryzen 7 9800X3D 上,吞吐量範圍從 6.27 GB/s (GPT-2) 到 1.12 GB/s (Gemma 3)。

使用與相容性

GigaToken 提供兩種主要方式來整合到現有工作流中:

相容模式 (Compatibility Mode)

此模式允許 GigaToken 作為 HuggingFace 或 Tiktoken 的直接替代方案。雖然它能確保輸出完全匹配,但由於維護相容性所產生的開銷,會帶來效能成本。用戶可以期待顯著的加速,但無法達到原生 API 所見的完整 1000 倍增益。

原生 GigaToken API

為了獲得最大效能,應使用原生 API。它接受 HuggingFace 模型名稱,並利用 TextFileSource 在 Rust 中直接讀取數據,跳過 Python 開銷。

import gigatoken as gt

tokenizer = gt.Tokenizer("Qwen/Qwen3-8B")
file_source = gt.TextFileSource(["owt_train.txt"], separator=b"<|endoftext|>")
tokens = tokenizer.encode_files(file_source)

限制與已知問題

雖然針對 BPE tokenizer 進行了高度優化,但 GigaToken 仍有一些已知的限制:

  • SentencePiece: 基於 SentencePiece 的模型(如 Google 模型或 BERT 風格的模型)的 tokenization 優化程度較低,導致吞吐量較慢。
  • WordPiece: 目前尚不支持 WordPiece tokenization。
  • OS 支援: Windows 支援尚未經過測試;在 Windows 上建議使用 WSL。
  • Python ABI: 目前的實現使用 ABI3,作者打算為每個 Python 版本進行專門化,以潛在地將受開銷限制的情況下的速度提高一倍。

社群討論

業界從業者指出,雖然 tokenization 通常僅佔總推理時間的一小部分(通常 <0.1%),但對於需要對訓練數據進行大規模 pre-tokenization 的相同應用來說,它是至關重要的。

"這太棒了,但 tokenization 通常僅佔總推理時間的 <0.1%。不過,猜測應該有很多應用僅僅需要進行 tokenization,而這對那些應用來說會非常棒!"

Sources