大規模言語モデルを自作しよう！2 (C++コーパスクリーニング+Patch-Level Training）
本記事は、LLM・LLM活用 Advent Calendar 2024 22日目の記事です。
はじめに
本記事は、「大規模言語モデルを自作しよう！」シリーズの第2段です。
第1段では、300Mサイズのmistralモデルをベースに、wikipediaデータセットとcc100データセットを用いた事前学習と、databricks-dolly-15k-jaを用いたファインチューニングを実現しました。しかし、事前学習モデルの出力がそもそも文章になっていない場合や、引用符、記事フッターやコピーライトのようなノイズが含まれる場合があり、性能が低いという課題がありました。
本記事では、生成テキストの品質向上を目的に、事前学習用コーパスのクリーニング処理を拡充し、各データに対して11種類のクリーニング処理を行います。また、さらなる性能向上を目的に、シングルRTX4090で2Bモデルの事前学習を行います。学習には40Bサイズのデータセットを用いつつ、Patch-Level Trainingを導入することで、学習コストを約60%削減します。
学習用のスクリプト、学習済みモデル等は、それぞれ以下の場所にアップロードしています。
| 項目 | リンク | 備考 |
|---|---|---|
| 学習用ツール | ce-lery/mistral-2b-recipe | |
| コーパスクリーニングツール | ce-lery/corpus-cleaner | mistral-2b-baseのexternalsフォルダにsubmoduleとして同梱 |
| 事前学習モデル | mistral-2b-base |
なお、事前学習モデルは以下から直ぐに試せます。是非、遊んでみてください。
- 下記の「Open In Colab」ボタンをクリックし、Google Colaboratoryを開く
- ページ上部のタブから「ランタイム」>「すべてのセルを実行」をクリック
検証環境
第1段の記事とほぼ同一です。差分は太字で表示します。
| 項目 | バージョン | 備考 |
|---|---|---|
| OS | Ubuntu 22.04.3 LTS | |
| CPU | AMD® Ryzen 5 3600x 6-core processorx12 | |
| RAM | DDR4 80GB | |
| GPU | RTX4090 VRAM24GB | |
| python | 3.11.6 | Dockerfile参照 |
| CUDA toolkit | 12.1 | Dockerfile参照 |
| cudnn | 8.8.0.121 | Dockerfile参照 |
| NVIDIA Driver | 550.107.02 | |
| pythonライブラリ | transformers==4.46.2 torch==2.1.1+cu121 |
Dockerfile参照 |
| SSD | 4TB | |
| その他ハードディスク | HDD 12TB SSD 4TB |
学習手順
それぞれのモデルの学習手順は次の通りです。
コーパスクリーニングと事前学習については、次章にて説明します。
その他の詳細は各スクリプト、および第1段の記事をご参照ください。
| step | 詳細 | 使用スクリプト | 処理時間 |
|---|---|---|---|
| 環境構築 | Dockerを使用してpython実行環境を構築します。 その他追加で必要なインストールをsetup.shで実行します。 |
README setup.sh |
0.2h |
| コーパス構築 | wikipedia、oscar等のデータをダウンロードします。 | dataset.sh | 12h |
| コーパスクリーニング | ダウンロードしたデータセットをクリーニングします。 クリーニング後のデータはマージし、train.jsonlとします。 |
dataset.sh | 36h |
| トークナイザー学習 | コーパスクリーニング済みのwikipediaデータ、cc100データを用いて、SentencePieceトークナイザーを学習します。 学習したSentencePieceトークナイザーは、huggingface llamatokenizer形式に変換します。 |
tokenizer.sh | 1h |
| 事前学習 | train.jsonlを用いて、mistral-2bモデルを事前学習します。 train.jsonlの前半80%(train_head.jsonl)をpatch-level trainingにて、残り20%(train_tail.jsonl)をtoken-level trainingにて使用します。 |
pretrain_patch.sh | 600h (patch-level 314h, token-level 286h) |
| 推論 | 事前学習済みモデルの出力を確認します。 | inference.sh | 10s |
コーパスのクリーニング
下表の11種類のコーパスのクリーニングを行いました。
高速化のため、処理はC++で実装しました。ソースコードはこちらです。
| 特徴 | 詳細 | 備考 |
|---|---|---|
| 制御文字除去 | 文字コードが0x1F以下(0x0A(LF)と0x0D(CR)は除く)の文字、および0x7Fと\u2028(LS)と\u2029(PS)の文字を除去します。 | OSCAR2019の生データにはNULL文字が含まれており、データをjsonl化しようとした際にエラーとなったため、この処理を行いました。 |
| 正規化 |
mecab-ipadic-NEologdの作成時に使用されたこちらの正規化処理を実行します。 正規化内容は、例えば「半角カタカナを全角に置き換える」や「全角スペースを半角スペースに置き換える」「NFKCによるUnicode正規化」などです。 |
すべての処理内容はこちらを参照。 |
| URL除去 | 正規表現「"(https?|ftp)(:\/\/[-_\.!~*\'()a-zA-Z0-9;\/?:\@&=\+\$,%#]+)"」に一致する URL を削除します。 |
|
| 特殊文字除去 | ☀、♡、☆のような特定の絵文字を削除します。 特定の Unicode 範囲内の特殊文字を削除します。 |
こちらの URL を参照してください。 |
| 絵文字除去 | 🤗、 🐉、 📊のような特定の絵文字を削除します。 文字コードは、 \U0001F300(🌀)から\U0001F9FF(🧿)を対象に、文中で完全一致するものを削除します。 |
|
| 引用符除去 | [1], {245}のような引用符を削除します。 正規表現「 "(\[([0-9]+)\]|\{([0-9]+)\})"」にマッチする文字列を削除しました。 |
|
| 長さフィルター | 1行が5文字以下、5000文字以上のデータを削除します。 | |
| 言語フィルター | fastTextを使用して文章の構成言語と品質を分類し、日本語のみを抽出します。 | fastTextについては、こちらを参照。 |
| 名詞比率フィルター | 先行文献に基づき、名詞が80%以上含まれるデータを削除します。 | |
| Minhash重複排除 | minhash を使用して文の重複を除去します。 今回は不使用です。 |
|
| ゼロ句読点フィルター |
先行文献に基づき、句読点（'、', '､', '。', '｡', '．', '.', '？', '?', '！', '!'）のない文を削除します。 |
|
| Perplexityフィルター | KenLMにより文章のPerplexityを算出し、Perplexityが高いものを低品質とし削除します。 文をトークン化するときにSentencePieceが使用されます。 |
最終的なコーパス構成は次のとおりです。
総クリーニング時間は約7日間、クリーニング後の総トークン数は41.01Bでした。
| データセット名 | 詳細 | 最終トークン数 |
|---|---|---|
| Wikipedia |
dumps.wikimedia.orgよりダウンロード。 bert-japaneseのスクリプトを用いてデータ整形。 |
0.61B |
| Wikibooks |
dumps.wikimedia.orgよりダウンロード。 bert-japaneseのスクリプトを用いてデータ整形。 |
0.01B |
| Wikiversity |
dumps.wikimedia.orgよりダウンロード。 bert-japaneseのスクリプトを用いてデータ整形。 |
0.00B |
| CC-100 |
data.statmt.orgよりダウンロード。 bert-japaneseのスクリプトを用いてデータ整形。 |
8.89B |
| OSCAR 2019 | oscar-corpus/oscarよりダウンロード。 | 22.61B |
| mC4 | allenai/c4より、先頭の150GB分をダウンロード。 | 8.89B |
事前学習
学習の高速化
第1段の記事では、学習速度高速化のため、「bf16混合精度学習」「torch.compile」「flash_attention2」の3つを実行しました。
今回は「bf16混合精度学習」「flash_attention2」に加え、Patch-Level Trainingによる学習速度の向上を行います。
Patch-Level Trainingとは、複数のトークンの分散表現を平均化したパッチという単位でLLMを学習することで、学習時間を高速化する手法です。イメージ図は次の通りです。
出典： https://arxiv.org/html/2407.12665, Figure 3.
Patch-Level Trainingを用いる際は、次のように、Patch-LevelとToken-Levelの2段階で学習を行います。
- Patch-Level Training: 言語モデルにパッチ単位で入力し、次のパッチを予測するようにトレーニングを行う
- Token-Level Training: Patch-Level Trainingのモデルから重みを引き継ぎ、残りのデータでトークン単位のトレーニング（通常の学習）を行う
この2段階の学習により、学習時間が
論文では、Patch-Level Trainingを用いることで、学習済みモデルの性能はそのままに、学習コストの削減が可能であることが示されています。
出典： https://arxiv.org/html/2407.12665, Figure 1.
今回は
Patch-Level Trainingは、Transformers v4.46.3のmodeling_mistral.pyに対して実装しました。
実装差分は以下の通りです。
変更後のmodeling_mistral.pyは、run_clm.pyから呼び出して使用します。
次のスクリプトでrun_clmを実行することで、2段階の学習を実現しています。
# ------------------------------------------
# 1st: patch-level training
# ------------------------------------------
python ../../source/train/run_clm_patch.py \
--model_type "mistral" \
--config_name ../../source/model/config/config_mistral_2b.json \
--tokenizer_name ./results/tokenizer/llamatokenizer \
--train_file ./results/dataset/train_head.jsonl \
--patch_size 4 \# patch sizeを4（=K）に設定
--output_dir ./results/pretrain/$DIR_NAME/trial1 \
# 省略
# ------------------------------------------
# 2nd: token-level training
# ------------------------------------------
python ../../source/train/run_clm_patch.py \
--model_type "mistral" \
--model_name_or_path ./results/pretrain/$DIR_NAME/trial1\
--train_file ./results/dataset/train_tail.jsonl \
--patch_size 1 \ # patch sizeを1に設定
--output_dir ./results/pretrain/$DIR_NAME/trial2 \
--cache_dir ./results/pretrain/cache/patch_train \
--do_train \
--prediction_loss_only \
# 省略
その他の変更点
その他、第1段の記事からの変更点は次の通りです。
- tokenizerをT5tokenizerからllamatokenizerに変更
- tokenizerのvocabulary sizeを50000から32000へ変更
- データセットのtoken化時、各データに対してbosとeosを付与するように変更(rinna TODOの構成を参照した)
- 学習率スケジューラとしてcosine_with_min_lrを採用
- 最終学習率（min_lr）として、最大学習率の10%の値を採用
- 実行時間削減のため、学習中の評価を不実施に（非推奨。過学習有無の確認のため、数回でいいので実施しておいた方が良かったのではと思います。）
出力例
作成したモデルの入出力例は以下の通りです。
一部誤った情報もありますが、比較的適切な情報が多く、流暢な日本で出力できている印象です。
| プロンプト | 出力 |
|---|---|
| <s>香川県の名物は、 | 讃岐うどんだけではありません。香川県には、瀬戸内海に面する香川県が誇る「小豆島」があります。小豆島の名物といえば、やはり「オリーブソフトクリーム」でしょう。</s> |
| <s>香川県の県庁所在地は、 | 香川県高松市です。</s> |
| <s>兵庫県の県庁所在地は、 | 姫路市です。</s> |
| <s>栃木県の県庁所在地は、 | 宇都宮市です。</s> |
| <s>日本の首都は、 | 東京です。</s> |
| <s>日本で一番高い山は、 | 富士山です。</s> |
| <s>日本で二番目に高い山は、 | 富士山です。富士山のふもとには、富士宮やきそばの店が立ち並び、観光客でにぎわっています。</s> |
| <s>日本で一番大きな湖は、 | 琵琶湖。琵琶湖の南岸に琵琶湖大橋が架かっている。</s> |
| <s>世界で一番高い山は、 | エベレスト(8848m)です。</s> |
| <s>世界で一番大きな湖は、 | 南極大陸の南端にある氷河湖。その面積は1,300平方キロメートルで、世界第2位の大きさを誇ります。</s> |
| <s>赤信号 | 、みんなで渡れば怖くない</s> |
| <s>ジョークとは、 | 冗談や冗談めかした冗談。</s> |
| <s>自然言語処理とは、 | 人間の言語を機械的に処理する技術です。</s> |
| <s>自動車を運転する際、青信号は進む、 | 赤信号で止まる、右折する、左折するなどのルールを守らないと罰金が科せられる。</s> |
| <s>人工知能 | (AI)が人間の仕事を奪う日も近い?</s> |
おわりに
本記事では、2BサイズのLLMの事前学習を実施しました。学習用コーパスは生成テキストの品質向上を目的に、11種類のクリーニング処理を実施しました。また、事前学習では40Bのデータセットを用いつつ、Patch-Level Trainingを導入することで、学習コストを約60%削減しました。
今後の展望として、今回作成した事前学習済みモデルを元にして、指示チューニング及び選好チューニングを行い、llm-jp-evalにて評価を行いたいと思います。
Discussion