詞元化器(tokenizer)在機器學習工作流程中,過去從未是效能瓶頸。就運算而言,詞元化相較於管線中其他繁重的模型運算,負擔輕微。然而,在某些情況下,它已迅速成為加速(或拖慢)機器學習工作的關鍵。

隨著模型變得更快且工作負載不斷擴展,這種平衡開始轉變。在大規模資料集上訓練、處理大量併發請求,或重複處理長輸入時,詞元化器可能會承受足夠壓力,導致模型缺乏資料。

這就是為什麼我們選擇在即將推出的 tokenizers v1 版本中,大幅專注於效能提升。詞元化應該輕量化,並能隨著您的工作流程擴展。您的 GPU 絕不應閒置等待 CPU 完成詞元化。

在本文中,我們將探討是什麼讓 v1 比 v0.23 快上數十倍。這項工作完全歸功於整個生態系。

詞元化是一個非常活躍的開源領域,諸如 gigatoken、tiktoken、kitoken、tokie、fastokens、wordchipper 和 ai-tokenizer 等函式庫,以及許多其他專案,都各自推進了快速詞元化器的可能性。我們參考了這些工作,其中一些想法之所以被採納,是因為其他專案證明了它們的價值。

在這次重構之前,tokenizers 的效能遠未達到其潛力,因此貢獻給它可能看起來不值得。透過這次重構,我們希望清楚表明,我們打算讓 tokenizers 成為一個值得貢獻的函式庫。

我們也感謝 IBM、NVIDIA 和 ExecuTorch 團隊貢獻補丁,並協助我們在廣泛的硬體上進行測試,以擴大平台支援。

### 成果

我們展示了 tokenizers v1 候選版本與其他廣泛使用的替代方案的比較結果。我們涵蓋了單執行緒、多執行緒、跨執行緒擴展、按模型比較、按語言比較、延遲、解碼吞吐量、記憶體堆積以及 crate 大小。

我們從 tokbench 儲存庫執行這些測試,並新增了一個指令,如果您願意,可以在您的硬體上重新執行基準測試。

### v1 是什麼

v1 將產生與 v0.23 相同的詞元 ID。目標是保留輸出、API、詞彙表和合併排名,並改進所有可以改進的地方。這包括廣度。

該函式庫保持對各種詞元化器家族的通用性,而非專注於 BPE,因此 v1 載入 v0.23 載入的所有內容。詞元化器將文字轉換為模型讀取的整數列表。

tokenizers 以四個階段執行該轉換。正規化(Normalization)將小寫轉換或 Unicode 正規化等操作應用於原始文字。預詞元化(Pre-tokenization)將文字分割成稱為預詞元(pre-tokens)的較小片段。

模型將每個預詞元轉換為詞元(tokens),並將它們映射到其詞彙表中的 ID。後處理(Post-processing)添加模型預期的任何特殊詞元。

模型階段是本文描述的大部分工作發生的地方。本文測量的十個模型家族中有八個使用位元組對編碼(Byte Pair Encoding,簡稱 BPE)。BPE 從預詞元的位元組開始,重複合併排名最高的相鄰對,直到沒有排名對為止。

排名是在詞元化器訓練時學習並隨之發布的,因此相同的文字總是產生相同的 ID。合併從不跨越預詞元邊界。另外兩個家族使用 WordPiece 和 Unigram,這是函式庫支援的另外兩種模型類型。

詞元化管線頁面記錄了這四個階段。詞元化演算法文件記錄了 BPE、WordPiece 和 Unigram。

每個階段都進行了改進。以下是重要的變更:

| 變更 | 作用 |

|---|---|

| 工作區拆分 | 一個 crate 變成了一個工作區:tk-encode 是必需的執行時,而 tk-serialize、tk-convert 和 tk-train 僅在應用程式需要時才連結。|

| 無記憶體分配模型 | 合併工作集存在於呼叫者擁有的暫存緩衝區中;迴圈從不觸及記憶體分配器。|

| bitcannon | 分割模式變為位元流上的布林運算,使用 SIMD 指令尋找分割點,而非正規表達式引擎。|

| 合併迴圈重寫 | 正在合併的片段在一個預先分配的緩衝區內形成一個侵入式雙向鏈結列表,因此合併只更新兩個索引,而不是移動資料。|

| 詞彙快取 | 一個執行緒本地的備忘錄,將預詞元位元組映射到完成的 ID,因此重複的詞彙只合併一次。|

| 原生平行化 | 一個共享的詞元化器可以同時從多個執行緒進行編碼;每個執行緒從自己的子池中提取其暫存緩衝區和詞彙快取,因此執行緒不再在單一鎖上排隊(#2365)。|

### 分割:位元流取代正規表達式

BPE 模型使用正規表達式將輸入文字分割成更小、更容易處理的區塊,稱為預詞元。合併發生在預詞元內部,從不跨越兩個預詞元之間的邊界,因此這個分割決定了管線其餘部分所看到的內容。

該正規表達式是模型的固定參數。它隨詞元化器一起發布,在執行時從不改變,因此無需通用正規表達式引擎在每次編碼時解釋它。可以為特定模型實際使用的模式,手動編寫一次等效的分割函數。

手寫函數隨後可以使用現代 CPU 的 SIMD 指令(單指令多資料),這些指令可以同時對多個位元組執行一個操作,非常適合 UTF-8 文字。bitcannon 將輸入的位元組視為平行的位元流,因此邊界是透過對整個暫存器進行布林運算而產生,而非一次推進一個字元的掃描。

它每個暫存器操作處理 64 個位元組。相同的想法也驅動了用於文字處理的 Parabix 和用於 JSON 的 simdjson。這取決於識別模式。少數語法涵蓋了大多數位元組級 BPE 模型,而模式不在此列的詞元化器將保留正規表達式路徑,因此無法獲得此加速。

### 詞彙快取

實際文字中包含許多重複的詞彙。由於 BPE 總是為給定的預詞元產生相同的詞元 ID,v1 可以在處理一次後儲存結果。一個執行緒本地快取將每個預詞元的位元組映射到其詞元 ID,允許後續出現的詞彙跳過合併過程。

自然地,隨著輸入的增長,唯一詞彙的數量可能比總詞彙量增長得慢。重複詞彙在輸入中所佔的比例會越來越高。新詞彙仍然會出現,這解釋了下方動畫中偶爾的快取未命中。

當輸入包含重複的預詞元時,快取效果最佳。包含少量重複預詞元的輸入可能會付出查詢成本,卻沒有獲得太多快取命中。

### 合併迴圈

下一個主要成本來自 BPE 合併迴圈。對於每個預詞元,迴圈重複尋找優先級最高的相鄰對並將其合併。先前實作在每次呼叫時都分配新記憶體,並為每個預詞元建立一個新的優先級佇列。

v1 重複使用呼叫者擁有的暫存緩衝區,消除了這些重複的記憶體分配。它將符號儲存在一個平面陣列中,並透過它們在陣列中的位置連結相鄰符號,這使得合併期間的更新成本更低。

它還在單次模型呼叫中處理一批預詞元。每個候選對也被打包成一個單一的 64 位元值,合併排名位於高位元。比較兩個候選對就只是比較兩個整數,而「此處無合併」是可能的最大值,因此迴圈無需分支即可找到下一個合併。

### 方法

基準測試設計中的微小差異可能會導致詞元化器效能的巨大差異。我們使用以下規則來保持不同引擎之間的比較一致性。

| 規則 | 原因 |

|---|---|

| 單一計時迴圈 | 每個引擎都執行相同的迴圈;沒有針對特定引擎的快速路徑。|

| 排除載入時間 | 詞彙表載入時間單獨計時,從不包含在編碼時間內。|

| ID 雜湊驗證 | 輸出 ID 的 FNV-1a 雜湊必須與基準線完全匹配。|

| 僅限共同單元格 | 中位數是所有引擎都執行並驗證過的單元格。|

| 每個程序完整掃描 | 每次重複都在一個新程序中開始並保留每個單元格。|

| 實體核心綁定 | 工作者綁定到八個不同的實體核心,從不綁定到同級 SMT 執行緒。|

| 獨立作業 | 獨立作業測量主機間的差異。|

重複編碼一個文件可能比在相同建構上編碼一系列不同文件更快。第一種方法測量當整個文件已在快取中表示時的效能。第二種方法測量新輸入的效能,同時允許先前見過的預詞元保持快取。

這兩種情況有時都被描述為「暖機」,儘管它們測量的是不同的工作負載。我們的主要結果使用不同的文件,且完整的語料庫太大無法完全放入快取。詞元化器基準測試應說明其使用的工作負載,因為選擇會主導結果。

### 這些改進的總和

在 v1 編碼路徑涵蓋的十個模型家族中,它在 Apple M4 Max 上使用單執行緒編碼文字的速度比 v0.23 快 3 到 30 倍。低端是 t5-base,高端是 gpt2。在八個工作者上,其擴展效率達到線性的 76%。

透過這些改變,v1 產生與已發布函式庫完全相同的詞元 ID。整體改進來自於多項變革的協同作用:手寫分割器取代正規表達式引擎、快取能回應重複詞彙而無需再次合併、從不觸及記憶體分配器的合併迴圈,以及每批預詞元一次模型呼叫,而非每個預詞元一次。每項改進都減少了管線中不同環節的工作量。

下一個優先事項是支援更多模型家族。我們將在 1.0.0 版本之前將更多模型遷移到新的合併迴圈。一旦候選版本穩定,下一步將是在依賴 tokenizers 函式庫的 transformers 函式庫和整個生態系中實現這些改進。

這篇文章是根據 tokbench 結果生成的,並將隨著支援範圍的擴大而更新。

### 如何取得

v1 的候選版本已在 crates.io 上發布。您呼叫的 API 與您目前使用的相同,因此唯一改變的是您安裝的版本。這是普通的安裝方式:`cargo add tokenizers --pre`。

訓練功能預設啟用,並會引入一個 C++ 相依性。如果您只需要編碼,可以將其關閉以排除訓練實作:`cargo add tokenizers --pre --no-default-features --features http`。

編碼功能不變:相同的呼叫,相同的 ID。

```rust

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` 是跨核心擴展的關鍵。它是上述擴展視圖測量的呼叫。

```rust

let encodings = tokenizer.encode_batch(documents, false)?;

```

本文中的每個圖表都是根據這個 crate 測量的。Python 綁定包裝了相同的程式碼,並從 bindings/python 建構,但它們增加了每次呼叫的開銷,這些測量中並未包含。

### v1 的進度

本文中的基準測試涵蓋了首先列出的已完成的候選版本工作。其餘部分展示了 1.0.0 版本仍需完成的工作以及我們之後計劃探索的內容。

### 候選版本:已實作

這項工作已在 crates.io 上的 Rust 預發布版本中:`cargo add tokenizers --pre`。

工作區拆分:將單一 crate 拆分為 tk-encode、tk-seri。