Tanuki-8B, 8x8B - Direct Preference Optimization (DPO)実行(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

1. TRLライブラリとは

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

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

from datasets import load_dataset
from trl import DPOConfig, DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

training_args = DPOConfig(output_dir="Qwen2-0.5B-DPO", logging_steps=10)
trainer = DPOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

2. polabについて

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

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

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

3. ディレクトリ構成

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

polab
├── LICENSE
├── README.md
├── base_config   各実験で共通する設定を記載するconfigを格納する
├── experiments　　実験用コードを格納する
├── playground    experimentsに含めない実験コードを格納する(wandbのprojectは分離され、コードもgitignoreされるので自由に使って良い)
├── src           実験用コードに使えるモジュールなどが格納されている　　　　　　　　　　　　　　
├── template　　　 実験用コードのテンプレートが格納されている
└── test　　　　　　開発時に使用したコードが格納されている

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

experiments                  
└── your_project_name                         検証テーマごとに作成する(projectごとにwandbのprojectが作成される)
    └── 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                          dpoなどを行う学習コード

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: team-hatakeyama-phase2/aya-ja-evol-instruct-calm3-dpo
    preprocess:
      - name: preprocess_openai_messages
        args:
          messages: prompt
          chosen: chosen
          rejected: rejected
          add_eos: true
      - name: add_bos
    split:
      train: train[:10000]
      eval: train[-1000:]

model:
  name: /storage5/someya/outputs/sftlab-experiments/8x8B/someya-sft_021-zero3_multi_node_no_offload

tokenizer:
  # name: null
  # name: team-hatakeyama-phase2/tanuki-tokenizer
  name: team-hatakeyama-phase2/tanuki-tokenizer-bos

exp_params:
  num_train_epochs: 2
  per_device_train_batch_size: 4
  per_device_eval_batch_size: 4
  gradient_accumulation_steps: 8
  #dataloader_num_workers: 24
  save_strategy: epoch
  save_steps: 50
  logging_steps: 1
  learning_rate: 2e-6
  warmup_ratio: 0.1
  lr_scheduler_type: cosine
  weight_decay: 0.0
  optim: adamw_torch
  dtype: bf16
  use_fast: true
  gradient_checkpointing: true
  max_length: 1150
  max_prompt_length: 925
  loss_type: sigmoid # loss for DPO
  beta: 0.1
  use_peft: true
  peft_target_model: tanuki-8x8-mini
  use_flash_attention_2: true
  peft_lora_r: 128
  peft_lora_alpha: 256
  peft_lora_dropout: 0.05
  ropo_alpha: 0.2 #weight of ropo loss for ROPO.
  ropo_gamma: 0.1 #weight of dpo loss for ROPO.
  rpo_alpha: null # weight of sft loss for rpo.
  do_eval: true
  eval_strategy: steps
  eval_steps: 25

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:DPOで使うデータを設定する
　dataの各種パラメータ
　name:データセットのパスやHuggingFaceのデータソースの名前を設定します。　　
　preprocess:データの前処理に関する設定をリスト形式で記述します。
　　name:train.pyでコーディングしている各種データ変換関数を設定します。
　　　　　またサンプルの前処理結果にbosトークンを追加する。(add_bos)を追加します。
　　args：データパラメータを設定します。
　　　　　①preprocess_trl
　　　　　　prompt: サンプル内の指示部分のキーを設定します。デフォルトは'prompt'。
　　　　　　chonsen: サンプル内の好ましい回答のキーを設定します。デフォルトは'prompt'。
　　　　　　rejected: サンプル内の好ましくない回答のキーを設定します。
　　　　　　　　　　　デフォルトは'rejected'。
　　　　　②preprocess_trl_with_score
　　　　　　prompt: サンプル内の指示部分のキーを設定します。デフォルトは'prompt'。
　　　　　　response1: サンプル内の指示部分のキーを設定します。
　　　　　　　　　　　　デフォルトは'response1'。
　　　　　　score1: サンプル内の指示部分のキーを設定します。デフォルトは'score1'。
　　　　　　response2: サンプル内の指示部分のキーを設定します。
　　　　　　　　　　　　デフォルトは'response2'。
　　　　　　score2: サンプル内の指示部分のキーを設定します。デフォルトは'score2'
　　　　　
　　　　　③preprocess_openai_messages
　　　　　　prompt: サンプル内の指示部分のキーを設定します。
　　　　　　role: サンプル内の指示部分のキーを設定します。デフォルトは'role'
　　　　　　content: サンプル内の指示部分のキーを設定します。デフォルトは'content'
　　　　　　chonsen: サンプル内の指示部分のキーを設定します。デフォルトは'prompt'
　　　　　　rejected: サンプル内の指示部分のキーを設定します。デフォルトは'rejected'
　　　　　　add_system_message: システムメッセージを追加するかどうかを設定します。
　　　　　　　　　　　　　　　　　　デフォルトは'True'
　　　　　　ignore_original_system_message: 元のシステムメッセージを無視するか
　　　　　　　　　　　　　　　　　　　　　　　　どうか。デフォルトはTrue。
　　　　　　add_eos: eosトークンを各ターンの終わりに追加するかどうか。
　　　　　　　　　　　デフォルトはFalse。

split：データセットをトレーニングセットと評価セットに分割する設定。
　　train:トレーニング部分のデータ設定
　　　　例: 20240806filtered[1000:]
　　　　インデックス1000以降をトレーニングデータとして使用します。
　　　eval:評価部分のデータ設定
　　　　　例：20240806filtered[:1000]
　　　　インデックス0〜999を評価データとして使用します。

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

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

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

学習率とスケジュールリング
　　learning_rate: 初期学習率を設定します。
　　warmup_ratio: ウォームアップステップの比率を設定します。
　　lr_scheduler_type: 学習率スケジューラのタイプ設定します。
　　weight_decay: ウェイト減衰率を設定します。
　　optim: 最適化アルゴリズムを設定します。

モデル構成
　　dtype: モデルのデータ型を設定します。
　　use_fast: 高速トークナイザーの有無を設定します。
　　radient_checkpointing: 勾配チェックポイントの有効化を設定します。
　　max_seq_length: 最大シーケンス長を設定します。
　　max_prompt_length: プロンプト部分の最大トークン長を設定します。
　　loss_type: DPO用にシグモイド損失を設定します。
　　beta: DPOの正則化ハイパーパラメータを設定します。
　　
　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のドロップアウト率を設定します。

ROPO（Ranked Ordered Preference Optimization）
　　ropo_alpha: ROPO損失の重みを設定します。
　　ropo_gamma: DPO損失の重み（ROPO内で使用）を設定します。
　　rpo_alpha: SFT（Supervised Fine-Tuning）損失の重みを設定します。

その他
　　do_eval: 評価の有無を設定します。
　　eval_strategy: 評価の頻度戦略を設定します。
　　eval_steps: 評価の頻度を設定します。

5. polab構築手順

①condaで仮想環境を作成します。

※condaがインストールされている前提です。

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

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

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

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

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

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

# この後のflash-attnでエラーになったので以下を実行
export LD_LIBRARY_PATH={path/to/your/miniconda3}/envs/poenv/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

pip install -e .

③「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. polabの実行

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

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

mkdir experiments/{project_name}

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

cp -r template/basic_v1.0 experiments/{project_name}/{exp_name}

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

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

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

# 8B_fullの場合の例
cp tanuki8b_dpo_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: 2
      per_device_train_batch_size: 2
      per_device_eval_batch_size: 2
      gradient_accumulation_steps: 64
      #dataloader_num_workers: 24
      save_strategy: steps
      save_steps: 1000
      logging_steps: 1
      learning_rate: 5e-7
      warmup_ratio: 0.1
      lr_scheduler_type: cosine
      weight_decay: 0.0
      optim: adamw_torch
      dtype: bf16
      use_fast: true
      gradient_checkpointing: false
      max_length: 1150
      max_prompt_length: 925
      loss_type: sigmoid # loss for DPO
      beta: 0.1
      use_peft: true
      peft_target_model: llama-mini
      use_flash_attention_2: true
      peft_lora_r: 128
      peft_lora_alpha: 256
      peft_lora_dropout: 0.05
      ropo_alpha: 0.2 #weight of ropo loss for ROPO.
      ropo_gamma: 0.1 #weight of dpo loss for ROPO.
      
      do_eval: true
      eval_strategy: steps
      eval_steps: 50
    

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

    exp_params:
      num_train_epochs: 2
      per_device_train_batch_size: 1
      per_device_eval_batch_size: 2
      gradient_accumulation_steps: 32
      #dataloader_num_workers: 24
      save_strategy: epoch
      save_steps: 1000
      logging_steps: 1
      learning_rate: 5e-7
      warmup_ratio: 0.1
      lr_scheduler_type: cosine
      weight_decay: 0.0
      optim: adamw_torch
      dtype: bf16
      use_fast: true
      gradient_checkpointing: false
      max_length: 1150
      max_prompt_length: 925
      loss_type: sigmoid # loss for DPO
      beta: 0.1
      use_peft: false
      peft_target_model: llama-mini
      use_flash_attention_2: true
      peft_lora_r: 128
      peft_lora_alpha: 256
      peft_lora_dropout: 0.05
      ropo_alpha: 0.2 #weight of ropo loss for ROPO.
      ropo_gamma: 0.1 #weight of dpo loss for ROPO.
      
      do_eval: true
      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: 2
      per_device_eval_batch_size: 2
      gradient_accumulation_steps: 64
    

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

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

⑤DPOを実行します

# 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
    

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

おわりに

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

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

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

---

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