LLM 日本語データセット(コーパス)構築向け: C++ Exact dedup at scale(Suffix Array 構築編)
Exact dedup とは
LLM 向けデータセット(コーパス)構築での exact dedup とは,
- 文章の全探索を行い, 重複する部分文字列(センテンス)を除去する
- (英語では)探索する部分文字列は 50 tokens 以上のセンテンス
- 日本語だと 20 tokens くらいがよいかしらん?
The RefinedWeb Dataset for Falcon LLM のメモ
Deduplicating Training Data Makes Language Models Better
あたりを参照ください.
中身としては, Suffix Array を構築して部分文字列の matching を行います.
実装
コードはこちら(現状 Suffix Array 構築まで)
Suffix Array 概要
Suffix Array はいろいろ解説されているので, 検索したり ChatGPT に聞いてみてください.
SA-IS で O(n)時間でリニア時間で計算できることが知られています.
幸いなことに SA-IS を実装している libsais でぺろっとリニア時間で計算できます.
ので今回は libsais を使います.
ストレージ問題
N を入力のサイズ(UTF-8 バイト数, token 数)とすると, Suffix Array は 8 N(64bit index の場合)を必要とします. N が 2 GB 以下であれば, 32bit index にすることで, 4 N に抑えることができます.
ただし, libsais 利用の場合はインデックスが signed(負数 or sign bit を sentinel かなにかに使っている模様)のため扱えるサイズは 1/2 になります.
基本的には Suffix Array は全体をオンメモリで保持する必要があるので, ストレージと同じだけのメモリ量が必要になります.
(mmap とかで頑張れば out-of-core 処理できるかもであるが...)
libsais は 16bit の入力データにも対応していますので, 文字列を tokenize して 0~65535 で表現できればさらにストレージの削減ができます.
方針
4 TB(~1 T tokens)規模の日本語文章に対して Suffix Array を構築できるようにしたい.
そのままでは 4 * 8 = 32 TB と, 2024/02 時点の PC ではメモリ量, ストレージ量ともに現実的ではありません.
そこで, 入力文章は数 100 MB ~ 2 GB くらいに分割してき, 32bit で Suffix Array を構築する(libsais の場合, 32bit では 最大 2 GB まで)ことにします.
ファイル(Suffix Array の個数)が増えてしまい, 探索するときにはすべてのファイルに対して行うことになってしまいますが, 探索速度よりメモリ量とストレージスペースの節約を優先することにします.
tokenize することでさらにストレージコストを減らします. 後述しますが, tokenize で, UTF-8 バイト列に対しておよそ 1/3.5 程度に減らすことができました.
探索時(dedup 時)は, 複数のセンテンスをまとめて Exact match することで, ファイルが複数あることによる無駄な探索はある程度 compensate されると想定できます.
あとはストレージ速度がボトルネックになるでしょうから, これもある程度無駄な探索が発生しても許容されると思料します.
したがって, tokenize + 32bit により, 4 TB 日本語テキストは,
1.15(4/3.5) * 4 = 4.6 TB
程度に Suffix Array のストレージを抑えることできます. 追加で zstd 圧縮(圧縮 level 9)で 5~10% 程度さらに削減できます.
これなら 2024/02 時点の汎用的 PC でも NVMe 4 TB x 1 + 1 TB x 1 くらいで処理できます!
tokenize
tokenize することで, UTF-8 バイト列に対して Suffix Array のストレージの節約ができるのが期待できます.
tokenize はそれなりに速いことが求められます.
既存の tokenizer で速い(とされている)tiktoken では, 日本語はほとんどが UTF-8 バイト列になってしまい, 効果がありません.
(日本語データセットで vocab を train すればよいであるが... めんどい)
今回は最長一致で高速な tokenize が期待できる RWKV world を使いました.
RWKV world tokenizer の情報と C++ での実装メモ
vocab(語彙)を 65535 に抑え, 入力文章を uint16 の整数配列にすることで, さらにスペースを節約して Suffix Array を計算できます.
日本語に特化させるために, ソースコード関連や英語の vocab を削り, 絵文字と顔文字(ASCII 絵文字)を追加しました.
(絵文字は UTF-8 4byte になる. 顔文字も web 日本語文章ではそこそこあるので, これを一つの id で表現できるとコンパクトになる)
UTF-8 byte fallback も対応して, 未知語の場合は UTF-8 文字列にしてシッカリ変換できるようにしました.
最長一致のトークナイズの実装には cedar を使いました.
cedar で最長一致でトークナイズするメモ
処理時間とファイルサイズ
テストしたテキストは, cc100ja の一部(cc100ja を 65535 文章ごとに分けたもの)です.
まずクリーニング前. 圧縮したのは73 MB ほどです.
-rw-r--r-- 1 syoyo syoyo 73291560 Feb 25 02:05 cc100-ja.00000.jsonl.zstd
ファイル読み込み, Suffix Array 構築, ファイル書き出しの一式で, Ryzen 3900x + WSL2 環境で, 1 コアで
- char ベース(UTF-8 バイト列): 13 秒
- tokenize : 7 秒
- unicode codepoint 形式利用だと tokenize 処理(最長一致探索)が少し早くなって 6 秒
でした. tokenize のほうが速いのは, データサイズが小さくなったため, Suffix Array 計算と ファイル書き出し時間が削減されたためです.
char ベースでは, 900 MB ほどの Suffix Array バイナリが生成され, 圧縮すると zstd で 826 MB でした.
(lz4 はあまり圧縮が効かなかった)
-rw-r--r-- 1 syoyo syoyo 864M Feb 25 02:34 output-sa.lz4
-rw-r--r-- 1 syoyo syoyo 826M Feb 25 02:34 output-sa.zstd
tokenize したものは, 圧縮前 240 MB が, 圧縮後 214 MB になりました! バイト列に対して 1/4 程度となり, かなりの節約です! 🎉
-rw-r--r-- 1 syoyo syoyo 214M Feb 25 02:32 output-sa-tokenized.zstd
一部クリーニング後のデータに対しても試しました.
-rw-rw-r-- 1 syoyo syoyo 42659193 9月 3 2023 ../../data/02_clean_step1/cc100ja/cc100-ja.00000.jsonl.zstd
元データは zstd 圧縮時 42 MB ほど.
Suffix array サイズは...
バイト列: zstd compress: 562789052 -> 528458701
tokenize: zstd compress: 163064216 -> 146216852
概ね tokenize では 1/3.5 となりました(圧縮時で 1/3.6).
したがって, tokenize での実務上のサイズ削減はバイト列に対して 1/3.5 くらいですかね.
さらなる高みへ
シリアライズ
のちのち Exact dedup を Python で処理したりを考えて, Suffix Array データのシリアライズには safetensors 形式を使うとよいでしょう.
今回は C++ で safetensors 形式で書き出す safetensors-cpp を使いました.
マルチノード処理
並列処理は容易です.
富岳などでは 256 nodes 利用くらいで容易に 800 TB(100 T tokens) くらいまでスケールできるでしょう.
Discussion