LLM実装におけるポイントまとめ

はじめに

こんにちは。今回はLLM実装について記事を書いてみました。
今年の秋からLLMを実装する機会があったのですが、その際に初め分からなかった箇所などを、いくつか備忘録の意も込めてまとめています。

コード実装

TrainingArgumentsにおける引数設定

TrainingArgumentsではoutput_dirなどトレーニングにおける引数を指定することができます。ここで特に引数設定をしなければ、デフォルト値が適用されます。
例えば、バッチサイズを1にして5エポックの学習を行わせ、1エポックごとにモデルを保存したい場合(設定変更)は以下のような設定になります。

training_args = TrainingArguments(
    save_steps=198,
    num_train_epochs=5,  # エポック数
    per_device_train_batch_size=1,  # GPUごとのバッチサイズ

ここで、上記[save_steps]の値は198となっていますが、この値は学習の際に使用するデータやバッチサイズによって変動します。今回は5エポックで990ステップを要する処理であったので、この値を198(990÷5=198)としています。

DeepSpeedにおけるCPUオフロード

DeepSpeedは、分散トレーニングや推論を効率的かつ効果的に行うための深層学習最適化ライブラリです。今回はDeepSpeedを使用して全パラメータ学習を行ってみました。ここでは、DeepSpeedにおけるCPUオフロードについて記載します。
CPUオフロードでは、オプティマイザのメモリや計算を GPU からホスト CPU にオフロードすることで、メモリの削減や高速化を目指します。これにより、単一のGPUでの処理など、限られたGPU環境の中でも処理を実行することができます。
今回は以下のような設定にしてみました（一部抜粋）。

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        }
        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

なお、ここで指定している[stage]では、オプティマイザーの状態や勾配、パラメーターの分割に対応する最適化ステージを指定します。
以下の図は標準的なデータ並列ベースラインと比較した3 ステージのメモリ節約量と通信量を表しています。

(引用：https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/)

LoRAConfigにおけるtarget_modulesの設定

LoRAを使用する場合はLoRAConfigにおいて引数を設定することができます。上記で記載したTraining Arguments同様、特に指定しなければデフォルト値が適用されます。
LoRAConfigでは、[r]や[lora_alpha]などの引数が指定できますが、ここでは[target_modules]の設定について記載します。
target_modulesでは、LoRAを適用させたい層(モジュール)を指定することができます。既存のモデルを使用する場合は、そのモデルごとにこの引数を変更する必要があります。
以下ElyzaとOpenCalmを使用した際のコードを記載しています。なお、Elyzaはllama2、OpenCalmはGPT-NeoXがベースとなっています。

#Elyzaの場合
peft_config = LoraConfig(
            r=8,
            lora_alpha=16,
            lora_dropout=0.1,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=[
            　"q_proj",
            　"o_proj",
        　　　"gate_proj",
            　"up_proj",
            　"down_proj",
            　"k_proj",
            　"v_proj",
            ]
        )

#OpenCalmの場合
peft_config = LoraConfig(
            r=8, 
            lora_alpha=16, 
            lora_dropout=0.1,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["query_key_value"])

さいごに

LLM実装をやり始めてからまだ少しの期間ですが、論文や公式ドキュメントなど見つつ行ってはいるもののまだまだ分からない事だらけです。場合にもよりますが、分からない点に対しては、基本的には公式ドキュメント(HuggingFaceやGitなど）を見に行ってみるのが、バージョン変更などの観点からも確実ではないかと考えています。
今回は3つのポイントを紹介しましたが、今後もこのようなナレッジはためていこうと思います！