GPUまたはTPU上のJAXでLLaMA、OpenLLaMA、XGenを微調整する方法

LLaMA、OpenLLaM、XGen

LLaMAモデルは2023年2月にMeta社からリリースされた。この生成AIモデルは、いくつかのサイズで提案されているオープンソースのモデルである：7Bパラメータ、13Bパラメータ、33Bパラメータ、65Bパラメータである。

2021年6月、GPT-Jがリリースされ、オープンソースの生成AIモデルがOpenAI GPT-3に本気で対抗できることを世界が認識し始めました。そして今、LLaMAによってそのハードルは再び明らかに上がり、このモデルはOpenAIのChatGPTやGPT-4に代わる非常に優れたオープンソースとなったようです。

LLaMAのライセンスはビジネスには不向きだ。しかし、他のモデルが存在することは朗報だ。

2023年6月にリリースされたOpenLLaMAは、バークレーAI研究チームによって開発されたLLaMAの代替バージョンで、非常に優れた結果をもたらし、ビジネスにも使用できる。本稿執筆時点で2つのバージョンが利用可能：7Bパラメータと13Bパラメータである。

Salesforceが2023年6月にリリースしたXGenも、商用アプリケーションで使用できる非常に強力な基盤モデルです。本稿執筆時点では、7Bパラメータ・バージョンのみが利用可能である。LLaMAとOpenLLaMAが2kトークンのコンテンツしかサポートしていないのに対し、このモデルは8kトークンのコンテキストをサポートしていることは注目に値する。

なぜ独自のモデルを微調整するのか？

上記のモデルは基礎的なモデルであり、大規模なテキスト・コーパスを教師なしで学習したものである。

このような基礎的なAIモデルは、通常は良い基礎となるが、あなたが望むものを適切に理解し、良い結果を返すようにチューニングする必要がある。これを達成する最も簡単な方法は、数発学習（「プロンプト・エンジニアリング」とも呼ばれる）を使うことだ。 こちらで、私たちの数発の専用ラーニング・ガイドを自由にお読みください。

数ショット学習は、生成AIモデルの新しいバージョンを作成することなく、その場で実行できるので便利だが、それだけでは不十分な場合もある。

最先端の結果を得るためには、あなた自身のユースケースに合わせてAIモデルを微調整する必要があります。微調整とは、あなた自身のデータに基づいてモデルのパラメーターをいくつか変更し、あなた自身のバージョンのモデルを得ることを意味します。

ファインチューニングは、生成AIモデルをゼロからトレーニングするよりもはるかに安価だが、それでも計算能力が必要なため、独自のモデルをファインチューニングするには高度なハードウェアが必要になる。最近のファインチューニングの代替技術には、より少ない計算パワーで済むものもあります（p-チューニング、プロンプトチューニング、ソフトチューニング、パラメータ効率的ファインチューニング、アダプター、LoRA、QLoRA...を参照）が、今のところ、これらの技術で同じレベルの品質を得ることはできていませんので、このチュートリアルでは触れません。

JAXとEasyLMによるTPU上のLLaMAの微調整

このチュートリアルでは、Berkeley AI ResearchチームがリリースしたEasyLMライブラリを使ったLLaMAの微調整に焦点を当てます： https://github.com/young-geng/EasyLM. This library is based on JAX which makes the fine-tuning process fast and compatible with both GPUs and Google TPUs.

同じ手法でOpenLLaMAやXGenを微調整することもできる。

ここではGoogle TPU V3-8でLLaMA 7Bを微調整していますが、A100 GPUでも完璧に同じことができます（EasyLMドキュメントの "Installation "の部分を注意深く読んでください。）もちろん、より大きなバージョンのLLaMA（13B、33B、65B...）を微調整することもできますが、TPU V3-8や単一のA100 GPUよりもはるかに多くのものが必要になります。

さあ、始めよう！

まず、ユースケース用のテキスト生成データセットを、JSONL 形式で、各例のキーとして "text" を使用して作成します。ここに簡単なセンチメント分析のデータセットがあります：

{"text":"[Content]: I love NLP Cloud, this company is awesome!\n[Sentiment]: Positive"}
{"text":"[Content]: Training LLMs is a complex but rewarding process.\n[Sentiment]: Neutral"}
{"text":"[Content]: My fine-tuning keeps crashing because of an OOM error! It just does not work at all!\n[Sentiment]: Negative"}

いくつかの重要なことに注意してほしい。第一に、このデータセットには簡単のために3つの例しか含まれていないが、実際の生活ではもっと多くの例が必要になるだろう。通常は300例が良いスタートです。次に、微調整したモデルを推論に使うときは、"[Content]: "と"[Sentiment]: "を接頭辞として使い、同じ書式に厳密に従う必要がある。最後に、"</s> "トークンは、ここでモデルの生成を停止することを意味するので、重要である。その他のデータセット例はNLPクラウドのドキュメントにあります： 詳しくはこちらをご覧ください。

Google Cloud上にV2 AlphaソフトウェアバージョンでTPU V3-8 VMを作成する：

VM に SSH して EasyLM をインストールします：

git clone https://github.com/young-geng/EasyLM
cd EasyLM
bash ./scripts/tpu_vm_setup.sh

LLaMAウエイトをダウンロードして変換することができます。最初の選択肢は、Metaに公式ウエイトを求めることです： https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. そして、このスクリプトで重みをEasyLMに変換します： https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. 2つ目の選択肢は、ハギングフェイスのLLaMAウェイトを使うことだ： https://huggingface.co/decapoda-research/llama-7b-hf. そして、このスクリプトで重みをEasyLMに変換します： https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

データセットをVMにアップロードし、HF LLaMAトークナイザーを使ってトークンの数を数える：

pip install -U transformers
python -c "from transformers import LlamaTokenizer; tokenizer = LlamaTokenizer.from_pretrained('decapoda-research/llama-7b-hf'); f = open('/path/to/your/dataset', 'r'); print(len(tokenizer.encode(f.read())))"

1024トークンのコンテキストでモデルをトレーニングする場合、返されるトークンの数を1024で割る必要があります。

2048トークンのコンテキストでモデルをトレーニングする場合、返されるトークンの数を2048で割る必要があります。

この数値が1エポックあたりのステップ数となる。そのため、例えば5エポック（これは通常良い設定です）の学習を行いたい場合は、この数値を5倍し、その結果を以下の-total_stepsに入力する必要があります。

データセットに100,000トークンが含まれ、1024トークンのコンテキストで5エポックが必要な場合、総ステップ数は(100,000/1024)*5 = 488となる。

コンテキストの長さに応じて、-train_dataset.json_dataset.seq_lengthを1024または2048に設定する。2048トークンのコンテキストでモデルを微調整するには、より多くのメモリを必要とすることに注意してください。

これで微調整プロセスを開始できる：

nohup python -u EasyLM/EasyLM/models/llama/llama_train.py \
--total_steps=your number of steps \
--save_model_freq=usually same as your number of steps \
--optimizer.adamw_optimizer.lr_warmup_steps=usually 10% of total steps \
--train_dataset.json_dataset.path='/path/to/your/dataset' \
--train_dataset.json_dataset.seq_length=1024 or 2048 \
--load_checkpoint='params::/path/to/converted/model' \
--tokenizer.vocab_file='/path/to/tokenizer' \
--logger.output_dir=/path/to/output  \
--mesh_dim='1,4,2' \
--load_llama_config='7b' \
--train_dataset.type='json' \
--train_dataset.text_processor.fields='text' \
--optimizer.type='adamw' \
--optimizer.accumulate_gradient_steps=1 \
--optimizer.adamw_optimizer.lr=0.002 \
--optimizer.adamw_optimizer.end_lr=0.002 \
--optimizer.adamw_optimizer.lr_decay_steps=100000000 \
--optimizer.adamw_optimizer.weight_decay=0.001 \
--optimizer.adamw_optimizer.multiply_by_parameter_scale=True \
--optimizer.adamw_optimizer.bf16_momentum=True &

いくつかの説明がある：

-save_model_freq：処理中にモデルを保存する頻度。小さなデータセットで微調整を行うだけであれば、プロセスの最後だけ保存することができ、その場合この値は --total_steps と等しくなる。

--optimizer.adamw_optimizer.lr_warmup_steps：通常、総ステップの10％が良い値です。

-tokenizer.vocab_file：tokenizer.modelファイルへのパス。例えば、HuggingFaceのdecapodaリポジトリを使う場合、以下はトークナイザーへのリンクです: https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

-logger.output_dir：最終モデルとログへのパス

他のパラメーターはそのままでいい。

微調整が終わると、-logger.output_dirで指定したパスからモデルを取り出すことができます。

微調整されたモデルを推論に使う

自分なりに微調整したモデルを手に入れたのだから、もちろんそれを使いたいだろう！

最初の戦略は、推論にEasyLMライブラリを使うことです。その場合、次のように推論サーバを起動することができます：

python EasyLM/EasyLM/models/llama/llama_serve.py  \
--mesh_dim='1,1,-1' \
--load_llama_config='7b' \
--load_checkpoint='params::/path/to/your/model' \
--tokenizer.vocab_file='/path/to/tokenizer'

そして、次のようにcURLでリクエストを送るだけでいい：

curl "http://0.0.0.0:5007/generate" \
-H "Content-Type: application/json" \
-X POST -d '{"prefix_text":["[Content]: EasyLM works really well!\n[Sentiment]:"]}'

2番目の戦略は、他のフレームワークで推論を行うために、モデルをHuggingFace形式にエクスポートすることです。以下は、それをエクスポートする方法です：

python EasyLM/EasyLM/models/llama/convert_easylm_to_hf.py \
--load_checkpoint='params::/path/to/output/model/streaming_params' \
--tokenizer_path='/path/to/tokenizer' \
--model_size='7b' \
--output_dir='/path/to/converted/model'

結論

2023年は、オープンソースの生成AIモデルにとって大きな節目となった。この記事を書いている時点では、LLaMA、OpenLLaMA、XGen、Orca、Falcon...のような素晴らしいモデルを誰でも使うことができる。

これらのモデルを微調整することは、ChatGPT（GPT-3.5）、GPT-4、Claude...のような最高の独自AIモデルを大幅に上回ることができる、あなた自身のユースケースに合わせた最先端の結果を得るための最良の方法です。

このガイドでは、LLaMA、OpenLLaMA、XGenを微調整する方法を紹介した。もし質問があれば、遠慮なく私に連絡してほしい。また、技術的に複雑なことなく、高度な生成AIモデルを簡単に微調整し、導入したいのであれば、遠慮なく私に連絡してほしい、 have a look at the NLP Cloud dedicated documentation!

Mark
NLPクラウドの機械学習エンジニア