GigaToken、Rust・SIMD・ファイル直接読み込みでBPEを高速化、GPT-2トークナイザーの独自テストで24.53 GB/sを達成

GigaToken 0.9.0は7月21日にPyPIで公開された。Rust製コアとPythonインターフェースを備え、Hugging Faceのモデル名を指定して直接読み込めるほか、既存のHugging Face Tokenizersまたはtiktokenオブジェクトをラップすることもできる。このプロジェクトが対象とするのはモデルのデコードではなく、学習データの準備、オフラインでのコーパス分析、大規模バッチのtoken数カウントといったCPU上のトークン化処理だ。

高いスループットは、複数の最適化を組み合わせることで実現されている。SIMDによりGPT系の正規表現を用いた事前トークン化を高速化し、繰り返し現れるpretokenに対応するtoken列をキャッシュしてBPE mergeの反復実行を回避する。また、tokenizerの規則に応じて専用の処理経路を採用し、Rust側では`TextFileSource`を通じてファイルを直接読み込むことで、Python文字列、オブジェクト割り当て、言語間境界をまたぐコストを削減している。互換モードは移行のハードルを下げるが、既存の出力と厳密に一致させるための変換コストにより、ネイティブAPIの最高性能には届かないことも公式に明記されている。

プロジェクトが独自に公表したベンチマークでは、2基のAMD EPYC 9565、合計144コアのホストで11.9 GBのOpenWebTextを処理した結果、GPT-2 tokenizerは24.53 GB/s、Hugging Face Tokenizersは24.8 MB/sだった。Phi-4とOLMo 2/3も23 GB/sを超えた。この種の結果は、GPU学習パイプラインがCPUによるtokenの準備を待っている場合、tokenizerが個別にプロファイリングする価値のあるボトルネックになり得ることを示している。

ただし、「989倍」は一般的なアプリケーションで期待できる高速化ではない。ベンチマークでは巨大なファイル、デュアルソケットサーバー、全コアによる並列処理が使用され、ファイル入力経路の違いも比較に含まれている。短いテキスト、ストリーミングリクエスト、Pythonのメモリ上にある入力、あるいは異なるNUMA構成では、効果が大幅に小さくなる可能性がある。エンジニアリングチームは次の段階として、tokenizer間でのビット単位の出力一致、Unicodeと特殊tokenの境界、ピークメモリ使用量、さらに自社環境のCPUと実際のバッチサイズにおけるエンドツーエンドのスループットを検証すべきだ。