大規模言語モデルを自作しよう！(Transformers+DeepSpeed+torch.compile+flash_attn2）
本記事は、LLM Advent Calendar 2023 13日目の記事です。
はじめに
🤗 Transformersは、自然言語処理、マルチモーダル、音声処理、コンピュータビジョン分野の事前学習済モデルを簡単にダウンロードしトレーニングすることが可能なpythonライブラリです。このライブラリを使用し、大規模言語モデル（LLM）の事前学習済モデルをローカルPC上にダウンロードし、それを使用した言語生成や、要約・翻訳・質問応答などの個別のタスクへのファインチューニング、チャットAIへの組み込みなどが盛んに行われています。
LLMの事前学習方法に関する情報としては、GPT-NeoXやMegatron-LM、TinyLlama、lit-llamaなど、他のpythonライブラリを使用したものが増えてきています。一方で、Transformersライブラリを使用したLLMの事前学習に関する情報は未だ少ない現状にあります。
そこで本記事では、300M規模のMistralモデルを題材とし、Transformersを使用してLLMの事前学習・ファインチューニングを実施する方法を紹介します。本記事で作成できるbaseモデルはこちらの「japanese-mistral-300m-base」、instructionモデルはこちらの「japanese-mistral-300m-instruction」に公開しています。
実装のためのソースコードは、japanese-mistral-300m-recipeにてv0.1.0として公開しており、本記事の位置づけはその解説です。以下のコマンドを実行することで、環境構築、事前学習、ファインチューニング、推論のすべてが実施可能です。
git clone japanese-mistral-300m-recipe
cd japanese-mistral-300m-recipe
git checkout v0.1.0
docker build -t cuda12.1-cudnn8-python3.11.6 ./
docker run -v ./:/home/japanese-gpt2/ -it --gpus all cuda12.1-cudnn8-python3.11.6
bash run_all.sh
この記事の特徴は、以下の通りです。
- SentencePieceトークナイザーでのbyte fallback使用によるunknown_token生成抑制と、huggingface Tokenizers形式への変換
- torch.compileを利用した学習高速化（2倍程度）
- flash attention2を使用した学習高速化（1.2倍程度）
- DeepSpeed ZEROを用いたRAMオフロードにより、小規模VRAMでの学習にも対応
- Mistral 300Mの利用
Quick Start with Google Colaboratory
事前学習・ファインチューニングを手っ取り早く試したい方のために、Google Colaboratory上で本記事の簡易的な内容を試せるipynbスクリプトを用意しました。
無料枠のT4 GPUを用いて実行可能です。T4でのデータセット構築〜ファインチューニング完了・推論までの時間は6時間程度ですので、コンピューティングユニットの消費にご注意ください。
以下の手順で実行してください。
- 上記の「Open In Colab」ボタンをクリックし、Google Colaboratoryを開く
- ページ上部のタブから「ランタイム」>「すべてのセルを実行」をクリック
以降の章では、japanese-mistral-300m-recipeのv0.1.0を使用してローカルPCにて事前学習&ファインチューニングを実施する方法を説明します。
検証環境
| 項目 | バージョン | 備考 |
|---|---|---|
| OS | Ubuntu 22.04.3 LTS | |
| CPU | AMD® Ryzen 5 3600x 6-core processor × 12 | |
| RAM | DDR4 80GB | DDR5の方が学習時間が高速化されるかもしれません（未検証） |
| GPU | RTX4090 VRAM24GB | |
| python | 3.11.6 | pyenv localにて設定 |
| CUDA toolkit | 12.1 | Dockerfile参照 |
| cudnn | 8.8.0.121 | Dockerfile参照 |
| pythonライブラリ | transformers==4.35.2 torch==2.1.1+cu121 deepspeed==0.12.3 flash_attn==2.3.4 |
その他はrequirements.txt参照 |
| SSD | 1.5TB | 中間生成ファイル等が作られますが、このサイズがあれば十分です |
| その他ハードディスク | HDD 12TB SSD 4TB |
ワークフロー
本記事の実施内容と所要時間は以下の通りです。LLMは、基本的に以下の流れで作成されます。
所要時間はハードウェア性能に起因して前後します。
| No. | ステップ | 所要時間 |
|---|---|---|
| 1 | python仮想環境構築 | 10min |
| 2 | データセット構築 | 10h |
| 3 | トークナイザー学習 | 30min |
| 4 | 事前学習 | 120h |
| 5 | 評価 | 1.5h |
| 6 | 推論 | 1min |
| 7 | ファインチューニング | 10min |
| 8 | 評価 | 20min |
| 9 | 推論2 | 1min |
python仮想環境構築
Dockerを使用した環境構築を推奨します。
以下のコマンドを実行することで、CUDAT toolkitとcudnn、その他のツールをインストールしたコンテナが作成されます。
docker build -t cuda12.1-cudnn8-python3.11.6 ./
次に、以下のコマンドでコンテナを起動し、コンテナ上での開発に移ります。
docker run -v ./:/home/japanese-gpt2/ -it --gpus all cuda12.1-cudnn8-python3.11.6
最後に、以下のコマンドでpython仮想環境を構築します。
bash setup.sh
setup.sh内では、pyenvによるpythonバージョン指定、venv環境の構築、pythonライブラリのインストールを実施します。
データセット構築
事前学習用のデータセットは、wikipediaデータセットと、cc100データセットをマージして構築します。
以下のコマンドを実行することで、データセットの構築が可能です。
cd pretrain
bash dataset/dataset.sh
データセットのダウンロード
以下のコードを実行し、wikipediaデータと、cc100データのダウンロード、マージを実施します。
クリーニング
データのマージ前に、データセットの最低限のクリーニングとして、文章の正規化を実施します。
今回はこちらで紹介されているneologdnを使用し、データセットの正規化を実施しました。neologdnの正規化内容は、例えば「半角カタカナは全角に置き換え」「全角スペースは半角スペースに置き換える」などがあります。詳細は以下をご覧ください。
データセットのtrain-test分割
wikipediaデータとcc100データのそれぞれを、ファイル行数でtrain:test=95:5で分割し、train、testごとにマージしてデータセットとします。
また、トークナイザー学習用に、wikipediaデータとcc100データの全てをマージしたmerge_dataset_for_spm.txtを作成します。
トークナイザー学習
TransformersのTransformerモデルは、入力として文字列を受けとることができません。代わりに、文字列をトークンという単位に分割し、トークンごとに割り当てられた数値（ID）を入力値として使用します。トークナイザーは、文字列をトークンに分割し、トークンごとにID変換するものです。
トークナイザーのトークン分割単位を決定づけるアルゴリズムには、バイトペアエンコーディング（BPE）、WordPiece、Unigramどがあります。各アルゴリズムの詳細はこちらに譲るとして、今回はSentensePieceライブラリにてUnigramアルゴリズムを使用して、トークン分割単位の決定とトークンごとのID割り当てを行います。本資料では、この処理を「トークナイザーの学習」と呼称します。
トークナイザーの学習は、以下のスクリプトで実行可能です。
SentencePiece学習
SentensePieceライブラリを使用し、トークナイザーの学習を実施します。トークナイザーは、wiki+cc100の全データを用いてUnigramアルゴリズムで学習します。
SentencePieceトークナイザーは、15行目の「vocab_size=50000」で、トークナイザーのリストに登録する語彙数を設定しています。逆に、この数以上の語彙は登録されません。トークナイザーは、リストに登録されていない文字列は処理することができず、未知語（unknown_token）として処理します。すなわち、リストに登録されていない文字列は、学習時・推論時に[unk]となってしまいます。
これを防ぐための機能がbyte-fallbackです。13行目のようにこれを有効にすることで、SentencePieceトークナイザーは渡されたリスト未登録文字をバイト単位でIDにエンコードすると共に、バイト単位のIDをUTF-8形式でデコードすることが可能になります。つまり、トークナイザーのリストに登録されていない文字列も処理することができます。
Tokenizers T5Tokenizer形式への変換
Transformersライブラリは、SentencePieceライブラリで生成されたトークナイザーをそのまま使用することができません。Transformersライブラリでも使用できるように、huggingface TokenizersライブラリのT5Tokenizersクラス形式に変換します。
データセットのトークン化処理
先ほど作成したトークナイザーを使用して、データセット全体をトークン化します。
この処理は、後述のrun_clm.pyに実装されており、run_clm.py実行時に実施されます。
事前学習
事前学習およびファインチューニングには、Transformersのexampleスクリプトの1つであるrun_clm.pyを使用します。run_clm.pyの一部を、以下のように修正しています。
@@ -56,7 +56,7 @@
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
- check_min_version("4.34.0")
+ check_min_version("4.35.0")
require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
@@ -248,12 +248,12 @@ def main():
# We now keep distinct sets of args, for a cleaner separation of concerns.
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
- if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
- # If we pass only one argument to the script and it's the path to a json file,
- # let's parse it to get our arguments.
- model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
- else:
- model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+ # if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+ # # If we pass only one argument to the script and it's the path to a json file,
+ # # let's parse it to get our arguments.
+ model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+ # else:
+ # model_args, data_args, training_args = parser.parse_args_into_dataclasses()
if model_args.use_auth_token is not None:
warnings.warn(
@@ -437,7 +437,8 @@ def main():
if model_args.torch_dtype in ["auto", None]
else getattr(torch, model_args.torch_dtype)
)
- model = AutoModelForCausalLM.from_pretrained(
+ from transformers import MistralForCausalLM, MistralConfig
+ model = MistralForCausalLM.from_pretrained(
model_args.model_name_or_path,
from_tf=bool(".ckpt" in model_args.model_name_or_path),
config=config,
@@ -447,9 +448,31 @@ def main():
trust_remote_code=model_args.trust_remote_code,
torch_dtype=torch_dtype,
low_cpu_mem_usage=model_args.low_cpu_mem_usage,
+ use_flash_attention_2=True
)
else:
- model = AutoModelForCausalLM.from_config(config, trust_remote_code=model_args.trust_remote_code)
+ from transformers import MistralForCausalLM, MistralConfig
+ import json
+
+ def load_config_from_json(config_file):
+ with open(config_file, 'r') as f:
+ config = json.load(f)
+ config = MistralConfig.from_dict(config)
+ return config
+
+ # model = AutoModelForCausalLM.from_config(config, trust_remote_code=model_args.trust_remote_code)
+ config = load_config_from_json(config_file = os.path.join(os.path.dirname(__file__),"mistral-300m/config.json"))
+ # model = MistralForCausalLM(config)
+ #refer:https://github.com/huggingface/transformers/issues/21610
+ from collections import OrderedDict
+
+ model = MistralForCausalLM.from_pretrained(pretrained_model_name_or_path=None,
+ config=config,
+ state_dict=OrderedDict(),
+ use_flash_attention_2=True)
+
+ print("mistral config:",config)
+ print("mistral model architecture:",model)
n_params = sum({p.data_ptr(): p.numel() for p in model.parameters()}.values())
logger.info(f"Training new model from scratch - Total size={n_params/2**20:.2f}M params")
@@ -513,6 +536,7 @@ def tokenize_function(examples):
f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
)
block_size = min(data_args.block_size, tokenizer.model_max_length)
+ print("block_size:",block_size)
# Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
def group_texts(examples):
@@ -659,4 +683,4 @@ def _mp_fn(index):
if __name__ == "__main__":
- main()
+ main()
このrun_clm.pyに対し、以下のhf_config.jsonを渡すことで、TrainingArguments()等のパラメータを設定しています。
Mistral 300Mモデルの設定
今回は、mistralai/Mistral-7B-v0.1のconfig.jsonを改良し、以下のようにモデルサイズが338Mになるように変更します。 各 モデルのパラメータは、japanese-gpt2-mediumのconfig.json、GPT2Config、MistralConfigを参照し、設定しています。
上記のconfig.jsonを読み込み、以下のようにrun_clm.py内でモデル初期化&学習を実施しています。
#refer:https://huggingface.co/learn/nlp-course/ja/chapter7/6?fw=pt
from transformers import AutoModelForCausalLM, MistralForCausalLM, MistralConfig
import json
def load_config_from_json(config_file):
with open(config_file, 'r') as f:
config = json.load(f)
config = MistralConfig.from_dict(config)
return config
config = load_config_from_json(config_file = "mistral-300m/config.json")
#refer:https://github.com/huggingface/transformers/issues/21610
from collections import OrderedDict
model = MistralForCausalLM.from_pretrained(pretrained_model_name_or_path=None,
config=config,
state_dict=OrderedDict(),
use_flash_attention_2=True)
# ~Training Argumentsなどの設定説明は省略~
trainer = Trainer(
model=model,
tokenizer=tokenizer,
args=args,
data_collator=data_collator,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["valid"],
)
モデルサイズは、以下の処理を実行することで確認できます。
model_size = sum(t.numel() for t in model.parameters())
print(f"Mistral size: {model_size/1000**2:.1f}M parameters")
学習速度高速化
学習時に課題となる要素の1つとして、「学習速度」が挙げられます。例えば学習時間が3600時間（150日）である場合、実質学習できないことと同義です。また、学習時間を短縮することで、GPUの電気料金やGoogle Colabo等の課金料金を削減することができるため、可能な限り学習速度を高速化することが望まれます。
学習速度の高速化方法は、こちらが参考になります。これらのうち、以下の設定を実施します。
| 設定項目 | 内容 | 備考 |
|---|---|---|
| fp16の混合精度トレーニング | モデルパラメータの一部をfp16精度で表現し、残りはfp32精度とすることで、計算を高速化する | こちらの通り、モデルが16ビットと32ビットの両方の精度 (GPU 上の元のモデルの1.5倍)でGPU上に存在するため、多くのGPUメモリが使用される可能性がある |
| torch.compile | PyTorch コードを最適化されたカーネルにJITコンパイルすることで、Pytorch処理時の高速化を実現する | こちらとこちらも参照 |
| flash attention2 | attention機構の計算を、並列化とWork Partitioningにより効率的に実施するflash attentionに置き換える | 現状は、Ampere、Ada、またはHopper GPU (A100、RTX3090、RTX4090、H100 など)のみ対応。T4では使用できない。 |
これらの有効化方法は、以下の通りです。
# flash_attention_2の有効化
model = MistralForCausalLM.from_pretrained(pretrained_model_name_or_path=None,
config=config,
state_dict=OrderedDict(),
use_flash_attention_2=True)
# fp16混合精度使用とtorch_compile使用の有効化
training_args = TrainingArguments(...,
fp16=True,
torch_compile=True)
trainer = Trainer(...)
trainer.train()
比較として、torch.comileとflash_attn2を共に無効にした場合の学習時間は236時間、共に有効にした場合は98時間でした。これらの機能を有効にすることで、学習速度が2.4倍になることがわかります。
VRAM確保
学習時に課題となるもう1つの要素として、「GPU VRAM容量」が挙げられます。VRAM容量が不足する場合、そもそも学習ができません。
VRAM容量の不足に対しては、DeepSpeed ZeROが有効です。これは、複数のGPUおよびCPUに対してさまざまなモデルトレーニング状態 (重み、勾配、オプティマイザー)を分割することでGPU VRAMの消費を削減する機能です。
以下のコマンドでrun_clm.pyを実行することで、学習に必要なVRAM容量を削減しつつ、高速に学習することを可能としています。
deepspeed --no_local_rank run_clm.py hf_config.json --deepspeed_config ds_config_zero.json
学習時の進捗状況確認
Transformersでは、以下の手順を踏むことで、学習時のtrain_loss、eval_lossなどをダッシュボード表示する機能が存在します。
- 以下のコマンドで、python仮想環境にtensorboardライブラリをインストールする
pip install tensorboard - TrianingArgumentsのlogging_dirにログ出力ディレクトリ（例：checkpoints-mistral-300M-FA2/logs）を設定し、学習を開始する
- 別のterminalを開き、python仮想環境をactivateした状態で以下のコマンドを実行する
tensorboard --logdir ./ - ブラウザで「http://localhost:6006/ 」を開く
以下は、epoch 0.82時の分析結果の様子です。train_loss、eval_lossともに順調に下がっている様子が確認できます。また、learning_rateのWarmupとDecayが効いていることが確認できます。
評価
run_clm.pyは以下のように、学習終了後に自動でperplexityを算出します。
perplexityは35.1016でした。
***** eval metrics *****
epoch = 1.0
eval_loss = 3.5582
eval_runtime = 1:44:34.83
eval_samples = 551057
eval_samples_per_second = 87.82
eval_steps_per_second = 21.955
perplexity = 35.1016
ここまでの手順で、以下の「japanese-mistral-300m-base」が完成します。
推論
推論用ソースコードは以下の通りです。
入力プロンプトを「大規模言語モデルとは、」とした時の、出力結果の一例は以下の通りです。今回は256文字の制限（max_new_tokens=256）を加えていますが、1024 tokenまで出力可能です。
大規模言語モデルとは、言語仕様の異なる複数の言語の集合である。 言語は、各言語ごとに異なる言語体系を持つが、その言語がどの言語で記述されているかは必ずしも明確ではない(例:C言語、C++、Perl、Python、Ruby、PHP、Java、Swift、Objective-C、D言語など)。ただし、これらの言語では、必ずしも言語を記述できるとは限らない。また、すべての言語には、特定の言語特有の言語特性がある。例えば、日本語、英語、フランス語、ドイツ語、イタリア語、スペイン語、ポルトガル語、トルコ語、アラビア語、ヘブライ語など、他の言語とは異なる言語によって記述されている言語もある。たとえば、ラテン文字とラテン文字を区別する言語として、アルメニア語や、スロバキア語などがある。しかし、これらは、いずれも言語の特徴を特徴づける言語ではない。そのため、多くの言語に共通した言語構造を持つ言語は存在しない。言語学の分野は多岐にわたるため、それぞれの言語の特性を理解した上で言語設計を行う必要がある。したがって、一つの言語(言語)で言語の構造を記述する言語としては、いくつかの言語が挙げられる。さらに、ある言語において、それぞれ独自の言語が存在する。それらは、文法、語法、語彙、構文、文理、構造、意味、規則、および規則などを含む言語
大規模言語モデルとは、どのようなものなのか。 また、その言語がどのように進化してきたのか、そしてどのように進化していったのかについて、さまざまな視点から解説する。... [Read more...] »「言語学入門」第3回「なぜ、言語が進化したのか?」[第2回] (2017年12月3日) (C) 2018 SQUARE ENIX CO., LTD. [JP/JP];東京都港区南青山5丁目3番3号2-1, Shinjuku University of Tokyo, LLC. ALL RIGHTS RESERVED.〒150-0003東京都渋谷区神宮前6丁目5番1号(GoogleMapsで見る) (Google Mapsで開く) [P], [M]、[F]の順にクリックし、図3に示すように[B]と[C]にカーソルを合わせ、【C】を押下する(B)をドラッグして[D]キーを押す(D)と,[E]キーを押すと,図4のようになる([R]は[H]).[V]を押す(A),または[S]を押しながら[G]を押して[L]を押すと
ファインチューニング
事前学習で作成されたcheckpoints-mistral-300M-FA2モデルをベースに、databricks-dolly-15k-jaを用いてinstructionチューニングを実施します。
データセットの形式は、以下のようにalpacaと同様のものとしています。
<s>
以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。
[SEP]
指示:
あなたは何でも正確に答えられるAIです。
[SEP]
入力:
User:日本で一番高い山は?
[SEP]
応答:
富士山
</s>
fine-tuning/dataset/dataset.shを実行することで、以下の手順で学習データセットの整形を実施します。
- 「databricks-dolly-15k-ja」のdatabricks-dolly-15k-ja.jsonをダウンロードする
- alpaca_preprocess.pyを実行し、databricks-dolly-15k-jaを前述のalpacaフォーマット形式のtxtデータ（databricks-dolly-15k-ja.txt）に変換する
また、fine-tuning/train/train.shを実行することで、databricks-dolly-15k-ja.txtを元にしたファインチューニングを実施します。前述の事前学習と同様にrun_clm.pyを使用して、ファインチューニングも実施しています。
run_clm.pyには、以下のhf_config_ft.jsonを渡しています。事前学習時のhf_config.jsonとの大きな違いは、model_name_or_pathに対して事前学習済みモデルを設定している点と、validation_split_percentageを10%に設定することで、train_fileから10%をeval用データ・セットとして使用している点です。
ここまでの手順で、以下の「japanese-mistral-300m-instruction」が完成します。
推論2
ファインチューニングで生成されたモデルを使用し、推論を実行します。
推論用のコードは以下の通りです。
推論結果は、以下の通りです。
Assistant:エベレスト山(標高3,929m)は、インドで最も高い山である。標高2,530mの山で、世界で最も高い山であり、世界第3位である[1][2][3][4]。 また、インドの最高峰である[2][4][5]。
Assistant:中国は、世界で最も人口の多い国の一つです。中国の人口は、世界第2位の人口を誇り、世界の人口の約3分の1を占めています。また、人口密度は世界最高で、世界で6番目に人口が多い国でもあります。中国には、中国最大の都市があり、
Assistant:世界で最も高い山の1つで、世界第2位の山である。標高は2,620mで世界最高地点であり、世界で2番目に高い。また、地球上で最も低い山であり[1]、世界で最も標高の高い山でもある[2][3][4][5][6][7]
Assistant:地球上で最も広い大陸はインドで、2番目に大きな国です
eedenは、インドで最も人口の多い国の一つです。インド最大の都市であり、世界で最も人口密度の高い都市の一つでもあります。また、インドの人口の約4分の1が住んでおり、世界第
Assistant:私は、AIが人間を判別できるAIであると考えています。AIは、人間の感情や行動のパターンを学習し、人間とAIを区別する能力を持っています。また、感情分析、言語理解、記憶、学習など、さまざまな能力も持っています。しかし、これらの
おわりに
本記事では、japanese-mistral-300m-recipeを使用し、LLMの事前学習とファインチューニングの方法について説明しました。
学習自体は回ったものの、最終perplexityが低く、出力結果にURLやwebページのフッターのような文字が混じっているなど、いくつか課題が見られました。今後の展望は以下の通りです。
- T5TokenizerではなくLlamaTokenizerを使用して学習する
- データセットのクリーニング（重複排除、品質フィルタリング）
- lrのスケジューラをconstantに変更&epoch増やして学習
- 学習時間のさらなる高速化(numbaなどを使用したJITコンパイル、Mixture of Expertsの利用)
- llm-evaluation-harnessを使用した評価
本記事に関するご意見・改善点等がありましたら、是非コメント欄へ記載をお願いいたします。特に、学習時間の高速化についての情報を必要としております。japanese-mistral-300m-recipeへのissue、pull requestも歓迎します。
また、よろしければ本記事へのいいねをお願いいたします。著者の励みになります。
最後に、本記事内でリンクしている情報をご提供いただきました皆様に、心より感謝申し上げます。
Discussion
素晴らしい記事をありがとうございます。
LLMのpre-trainingについて、実装方法が分からず悩んでいたのでとても参考になりました。
一つ質問なのですが、こちらはhttps://huggingface.co/rinna/japanese-gpt2-mediumの継続事前学習ではなく、configを使用して新しいものを1から作成しているという認識であっているでしょうか。
コメント頂きありがとうございます！
はい、その通りです。継続事前学習ではなく、事前学習です。新しいモデルを1から学習させています。
返信が遅くなり申し訳ございません。
ありがとうございます。こちらの記事のおかげで理解が深まりました！