Tanuki-8B, 8x8B - Supervised Fine-Tuning (SFT) 実行(11/24日勉強会公開用)

GENIAC 松尾研LLM開発プロジェクトメンバーのknishimaeです。
Team「たぬき」Phase 2では、日本語対話能力に優れた大規模言語モデルTanuki-8BとTanuki-8x8Bを開発しました。本記事では、このプロジェクトにおいて、モデルの性能を向上させるために実施したSupervised Fine-Tuning (以降SFT) で活用したソースコードとライブラリについて説明させて頂きます。

はじめに

この記事の内容は、以下の勉強会で報告したものになります。

開催日時：2024年11月24日（日）17:00 - 19:00

テーマ：【勉強会 #3】小型LlamaモデルのTRLライブラリを用いた事後学習

この勉強会、TRLライブラリを用いたSTFについて報告を行っています。
本記事と合わせてご活用ください。

また、以下の記事についても参考にしてください。
https://zenn.dev/matsuolab/articles/62c75674190a41

ソースコードは、以下のURLを活用して解説しています。
https://github.com/matsuolab/nedo_project_code/tree/team_hatakeyama_phase2/team_hatakeyama_phase2/sftlab

1. TRLライブラリとは

HuggingFaceにアップされている事後学習用ライブラリ

https://huggingface.co/docs/trl/index

from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("lucasmccabe-lmi/CodeAlpaca-20k", split="train")

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = " ### Answer:"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(output_dir="/tmp"),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

trainer.train()

2. sftlabについて

sftの実験コードおよびテンプレートを作成しました。

(Phase1で活用したコードを誰でも活用できるようにカスタマイズ)

実験管理のため、一定のディレクトリ構造や命名規則に従って構成されています。

3. ディレクトリ構成

ディレクトリ構成は、以下の通りとなっています。

sftlab
├── LICENSE
├── README.md
├── base_config  各実験で共通する設定を記載するconfigを格納する
├── experiments　実験用コードを格納する
├── llm-jp-sft    学習コードの元となったコードが格納されている
├── playground    experimentsに含めない実験コードを格納する(wandbのprojectは分離され、コードもgitignoreされるので自由に使って良い)
├── scripts　　　　学習データ作成用コードが格納されている(学習コードに統合されているので削除予定)
├── template　　　実験用コードのテンプレートが格納されている
└── test　　　　　開発時に使用したコードが格納されている

実験用コードは基本的に以下の構成に従って作ります。

experiments                  
└── your_project_name                         検証テーマごとに作成する(e.g. ハルシネーションの効果を検証する)
    └── your_exp_name                         学習コード(train.pyやrun.py)ごとに作成する
        ├── accelerate_config
        │   └── your_accelerate_config.yaml   accelerateを使用する場合の設定を記載したファイル
        ├── exp_config
        │   └── your_exp_config.yaml          実験設定(モデル、データ、学習パラメータなど)を記載したファイル
        ├── run.py　　　　　　　　　　　　　　　 train.pyを実行するコード
        └── train.py                          sftを行う学習コード

projectとexpが1:n、expとrun.py/train.pyが1:1、expとexp_config.yaml/accelerate_config.yamlが1:nで紐づくようになっています。

学習コードが異なるような実験はexpを分け、設定の違いはconfig.yamlで切り替えるようにしています。

4. config.yaml

4.1 config.yamlサンプル

template:
  system: 以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。
  instruction: \n\n### 指示:\n
  input: \n\n### 入力:\n
  output: \n\n### 応答:\n

data:
  - name: kanhatakeyama/ramdom-to-fixed-multiturn-Calm3
    preprocess:
      - name: preprocess_openai_messages
        args:
          messages: messages
          role: role
          content: content
          add_eos: true
      - name: add_bos
    split:
      train: 20240806filtered[1000:]
      eval: 20240806filtered[:1000]

model:
  name: /storage5/shared/Nishijima/Llama-3-8b-MoE/11th_tonyu_iter_0000220
  
tokenizer:
  name: team-hatakeyama-phase2/tanuki-tokenizer
  
exp_params:
  num_train_epochs: 1
  per_device_train_batch_size: 2
  per_device_eval_batch_size: 2
  gradient_accumulation_steps: 4
  save_strategy: epoch
  save_steps: 0.5
  logging_steps: 1
  learning_rate: 5e-7
  warmup_ratio: 0.1
  lr_scheduler_type: cosine
  dtype: bf16
  use_fast: true
  gradient_checkpointing: true
  max_seq_length: 4096
  use_peft: false
  peft_target_model: tanuki-8x8
  use_flash_attention_2: true
  peft_lora_r: 128
  peft_lora_alpha: 256
  peft_lora_dropout: 0.05
  neftune_noise_alpha: null

  do_eval: true
  eval_strategy: steps
  eval_steps: 40

4.2各パラーメータの説明

1.template：プロンプトのカスタマイズするために設定

テンプレートプロンプトの例

yaml設定
template:
  system: 以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。
  instruction: \n\n### 指示:\n
  input: \n\n### 入力:\n
  output: \n\n### 応答:\n

出力されるプロンプト
system: 以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。
instruction: 

### 指示:
日本の四季について説明してください。

input: 

### 入力:
なし

output: 

### 応答:
日本には四季があり、春、夏、秋、冬それぞれに独特の特徴があります。春は桜が咲き、夏は暑く湿度が高いです。秋は紅葉が美しく、冬は寒く雪が降る地域もあります。

2.data:data:STFで使うデータを設定する
　dataの各種パラメータ
　name:データセットのパスやHuggingFaceのデータソースの名前を設定します。　　　
　preprocess:データの前処理に関する設定をリスト形式で記述します。
　　name:train.pyでコーディングしている各種データ変換関数を設定します。
　　　　　またサンプルの前処理結果にbosトークンを追加する。(add_bos)を追加します。
　　args：データパラメータを設定します。
　　　　　①apply_template
　　　　　　sinstruction: サンプル内の指示部分のキーを設定します。
　　　　　　sinput: サンプル内の入力部分のキーを設定します。デフォルトはNone。
　　　　　　soutput: サンプル内の出力部分のキーを設定します。
　　　　　　
　　　　　②apply_chat_template
　　　　　　messages: サンプル内のメッセージのキーを設定します。
　　　　　　add_system_message: システムメッセージを有無を設定します。
　　　　　　　　　　　　　　　　　　　デフォルトはFalse。
　　　　　　
　　　　　③preprocess_openai_messages
　　　　　　messages: サンプル内のメッセージのキーを設定します。
　　　　　　role: メッセージ内の役割のキーを設定します。デフォルトは'role'。
　　　　　　content: メッセージ内のコンテンツのキーを設定します。
　　　　　　　　　　　デフォルトは'content'。
　　　　　　add_system_message: システムメッセージの有無を設定する。
　　　　　　　　　　　　　　　　　　デフォルトはTrue。
　　　　　　ignore_original_system_message: 元のシステムメッセージの有無を
　　　　　　　　　　　　　　　　　　　　　　　　　設定する。デフォルトはTrue。
　　　　　　add_eos: eosトークンを各ターンの終わりの有無を設定する。
　　　　　　　　　　　　　デフォルトはFalse。
　　　　　
　　　　　④apply_preprocessing　
　　　　　　sample: 処理するデータを含むサンプルを設定する。
　　　　　　steps: 前処理関数の名前と引数を含む辞書のリストを設定する。
　
　split：データセットをトレーニングセットと評価セットに分割する設定。
　　train:トレーニング部分のデータ設定
　　　　例: 20240806filtered[1000:]
　　　　インデックス1000以降をトレーニングデータとして使用します。
　　　eval:評価部分のデータ設定
　　　　　例：20240806filtered[:1000]
　　　　インデックス0〜999を評価データとして使用します。

3.model：モデルのパスやHuggingFaceのモデルパス名前を設定します。

4.tokenizer：tokenizerのパスやHuggingFaceのtokenizerパス名前を設定します。

5.exp_params：モデルの引数を設定します。
　トレーニング関連
　　num_train_epochs: SFTのエポック数を設定します。
　　per_device_train_batch_size: 1デバイスあたりのトレーニング用バッチサイズ
　　　　　　　　　　　　　　　　　　を設定します。
　　per_device_eval_batch_size: 1デバイスあたりの評価用バッチサイズを設定します。
　　gradient_accumulation_steps: 勾配を更新する前に蓄積するステップ数を設定します。
　　save_strategy: モデル保存する所を設定します。
　　save_steps: モデル保存の頻度で設定します。
　　logging_steps: ログ出力の頻度（1ステップごと）。
　　gradient_checkpointing: 勾配チェックポイントを使用する設定します。

学習率とスケジュールリング
　　learning_rate: 初期学習率を設定します。
　　warmup_ratio: ウォームアップステップの比率を設定します。
　　lr_scheduler_type: 学習率スケジューラのタイプ設定します。

モデル構成
　　dtype: モデルのデータ型を設定します。
　　use_fast: 高速トークナイザーの有無を設定します。
　　max_seq_length: 最大シーケンス長を設定します。
　　
　PEFT (Parameter-Efficient Fine-Tuning) 関連
　　use_peft: Parameter-Efficient Fine-Tuning（PEFT）を設定します
　　peft_target_model: PEFT用のターゲットモデル名を設定します。
　　use_flash_attention_2: Flash Attention v2の有無を設定します
　　peft_lora_r: LoRA（Low-Rank Adaptation）の次元を設定します。
　　peft_lora_alpha: LoRAのスケーリング係数を設定します。
　　peft_lora_dropout: LoRAのドロップアウト率を設定します。
　
　その他
　　neftune_noise_alpha: パラメータノイズ用のアルファ値を設定します。
　　do_eval: 評価の有無を設定します。
　　eval_strategy: 評価の頻度戦略を設定します。
　　eval_steps: 評価の頻度を設定します。

5. sftlab構築手順

①condaで仮想環境を作成します。
※condaがインストールされている前提です。

# Python仮想環境を作成
conda create -n sft python=3.11 -y

# 作成したPython仮想環境を有効化
conda activate sft

# Python仮想環境を有効化した後は (python3コマンドだけでなく) pythonコマンドも使えることを確認
which python && echo "====" && python --version

②リポジトリをgit cloneし、必要なライブラリをインストールします。

# リポジトリをgit clone
git clone https://github.com/team-hatakeyama-phase2/sftlab.git
cd sftlab

# 必要なライブラリのインストール
pip install -r requirements.in 

# この後のflash-attnでエラーになったので以下を実行
export LD_LIBRARY_PATH={path/to/your/miniconda3}/envs/sft/lib/python3.11/site-packages/nvidia/nvjitlink/lib:$LD_LIBRARY_PATH

pip install flash-attn --no-build-isolation
pip install --upgrade accelerate
pip install datasets

③「base_config_template.yaml」を「base_config.yaml 」にコピーする

# base_config_template.yamlをコピー
cp base_config/base_config_template.yaml base_config/base_config.yaml 

④「base_config.yaml 」をReadme（Setup）に従って環境構築をします（base_configの設定まで）

output_dir: /storage5以下のお好きなディレクトリ
hf_cache_dir: /storage5/shared/huggingface_cache # 全体の方針でここにcacheを貯めています
wandb:
  entity: weblab-geniac1 # チームのwandbです

6. sftlabの実行

{project_name}：プロジェクト名
{exp_name}：検証名
{config_name}：ファイル名
{your_accelerate_config}:accelerateの名称
accelerate:deepspeedのstage種類

①experimets以下に新しくプロジェクトディレクトリを作成します。

mkdir experiments/{project_name}

②実験ディレクトリのテンプレートをコピーします

cp -r template/basic_tanuki8b_v3.2 experiments/{project_name}/{exp_name}

③configのテンプレートをコピーします

cd experiments/{project_name}/{exp_name}/exp_config

# 8B_loraの場合の例
cp tanuki8b_lora_001.yaml {config_name}.yaml

# 8B_fullの場合の例
cp tanuki8b_full_001.yaml {config_name}.yaml　

④configを書き換えます

• exp_configのdata:name:, model:name: は、モデルとデータを設定します

• 他のパラメータは基本的にtemplete（v4.1）の値を使用してください

  • lora（1GPU）のexp_config（v4.1）

    exp_params:
      num_train_epochs: 1
      per_device_train_batch_size: 4
      per_device_eval_batch_size: 4
      gradient_accumulation_steps: 64
      save_strategy: steps
      save_steps: 1000
      logging_steps: 1
      learning_rate: 1e-4
      warmup_ratio: 0.1
      lr_scheduler_type: cosine
      dtype: bf16
      use_fast: true
      instruction_template: \n\n### 指示:\n
      response_template: \n\n### 応答:\n
      gradient_checkpointing: true
      max_seq_length: 4096
      use_peft: true
      peft_target_model: llama-all
      use_flash_attention_2: true
      peft_lora_r: 128
      peft_lora_alpha: 256
      peft_lora_dropout: 0.05
      neftune_noise_alpha: null
    
      eval_strategy: steps
      eval_steps: 50
    

  • full（8GPU）のexp_config（v4.1）

    exp_params:
      num_train_epochs: 1
      per_device_train_batch_size: 4
      per_device_eval_batch_size: 4
      gradient_accumulation_steps: 8
      save_strategy: steps
      save_steps: 1000
      logging_steps: 1
      learning_rate: 5e-5
      warmup_ratio: 0.1
      lr_scheduler_type: cosine
      dtype: bf16
      use_fast: true
      instruction_template: \n\n### 指示:\n
      response_template: \n\n### 応答:\n
      gradient_checkpointing: true
      max_seq_length: 4096
      use_peft: false
      peft_target_model: llama-all
      use_flash_attention_2: true
      peft_lora_r: 128
      peft_lora_alpha: 256
      peft_lora_dropout: 0.05
      neftune_noise_alpha: null
    
      eval_strategy: steps
      eval_steps: 50
    

• loraでaccelerateを使用してGPUを増やす場合は、batch_sizeが一定になるようにしてください

  • batch_size = per_device_train_batch_size * gradient_accumulation_steps * GPU数

• マルチGPUを使用する場合はaccelerate_configを設定します

  • フルパラ（8GPU）の場合はデフォルトのzero1.yamlを使用できます

  • GPUを1→2にしたらexp_configのgradient_accumulation_stepsを半分にします

    # lora 1GPUの場合
      per_device_train_batch_size: 4
      per_device_eval_batch_size: 4
      gradient_accumulation_steps: 64
    

    # lora 2GPUの場合
      per_device_train_batch_size: 4
      per_device_eval_batch_size: 4
      gradient_accumulation_steps: 32
    

  • 例えば、lora（2GPU）の場合はaccelerate_configでnum_processes: 2 のようにGPU数を設定します

⑤SFTを実行します

# run.py ファイルがあるディレクトリに戻る
cd ..

# シングルGPUの場合
python run.py {config_name}.yaml

# マルチGPUを使用する場合
python run.py {config_name}.yaml --accelerate_config your_accelerate_config.yaml

• 動作確認のときは--debug付けてもらって、動作確認取れたら外して実行してもらうと、wandbのログ（project）が分離されて見やすくなります
  python run.py your_config.yaml --debug
• 長時間学習を回すときはsbatchを使った方が良さそうです

  • （参考）バッチファイル

    #!/bin/bash
    
    #SBATCH --nodelist=slurm0-a3-ghpc-0
    #SBATCH --time=12:00:00
    #SBATCH --gpus-per-node=1
    #SBATCH --cpus-per-task=16
    #SBATCH --mem=160GB
    #SBATCH --job-name=sft
    #SBATCH --output=%x_%j.log
    
    source $EXP_HOME/miniconda3/etc/profile.d/conda.sh
    conda activate sft
    
    LD_LIBRARY_PATH="$EXP_HOME/miniconda3/envs/sft/lib" python run.py tanuki8b_lora_001.yaml
    

上記でSFTが可能となります。

おわりに

TRLライブラリとsftlabの使い方について説明することが出来ました。

この概要等が今後のSFTのナレッジ活用となることを期待しております。

最後に、このナレッジにご協力いただいた関係者の皆様に、深く感謝申し上げます。皆様のご協力のおかげで、このようなナレッジを作成することが出来ました。

---

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。