LMFlowによる日本語LISAトレーニング#1

LISA(Layerwise Importance Sampled AdamW)は、イリノイ大学と香港科技大学の共同研究によって提唱されたLoRA(Low Rank Adaptation)とフルファインチューニングに代替するファインチューニング手法です。

メモリ消費がLoRAと同等に低く、なおかつパフォーマンスはフルパラメータのファインチューニングに匹敵もしくは上回る効果を持つと言われています。

Rui Pan et al. 2024

FreeAIでは、LISAが実際に使用できるか、日本語コーパスで実験することにしました。

LISAがいち早く実装されたフレームワークはLMFlowというものです。LMFlowの論文はこちらです。

ただし、日本語で使おうとすると色々と問題があることがわかりました。この記事はLISAとLMFlowによる日本語ファインチューニングの方法をまとめたものです。

まず、日本語Wikipediaを元に作られた東京大学和泉研究室のデータセットをLMFlowで扱える形に変換します。

その変換のためのコードは以下のとおりです(Clause-3によって開発されました)。

from datasets import load_dataset
import json
from tqdm import tqdm

# データセットの読み込み
dataset = load_dataset("izumi-lab/wikipedia-ja-20230720")

# JSONファイルに書き込むデータの作成
output_data = {
    "type": "text_only",
    "instances": []
}

# データセットの各サンプルに対して処理を行う
for sample in tqdm(dataset["train"], desc="Processing samples"):
    # textフィールドを抽出
    text = sample["text"]
    
    # インスタンスをJSONファイルに追加
    output_data["instances"].append({"text": text})

# JSONファイルに書き込む
with open("output.json", "w") as f:
    json.dump(output_data, f, indent=2)

こうして得られたJSONファイルをdata/wikija/trainの下に置きます。

次にトークナイザーの拡張が必要です。

幸い、LMFlowは中国語と英語のバイリンガル運用ができるようになっていたので、中国語のトークナイザー拡張と同じ方法でトークナイザーを拡張できます。

トークナイザーの拡張に関してはLMFlow/scripts/vocab_extensionに必要なコードがあります。

このコードを一部変更しながら使います。

まず、jsonファイルをテキストファイルに変換します

python utils/convert_json_to_txt.py --dataset_path ./data/wikija/train \
        --output_path ./data/wikija/converted_data.txt \
        --overwrite True

次に、トークナイザーを学習させます。

mkdir -p ./output_models/merged_tokenizer
python utils/train_tokenizer.py --dataset_path ./data/wikija/converted_data.txt \
        --model_type bpe \
        --output_dir ./output_models/new_tokenizer \
        --user_defined_symbols 0,1,2,3,4,5,6,7,8,9,% \
        --vocab_size 60000 \
        --max_sentencepiece_length 5

元のトークナイザーの設定はvocab_sizeが20000でmax_sentencepiece_lengthが4でしたが、日本語の性質を考えてvocab_sizeを60000、max_sentencepiece_lengthを5にしてみました。これは「岩崎弥太郎」のように五文字の人名が日本人には少なくないからです。

トークナイザーの学習が終わったら、元のモデルとトークナイザーをマージします。 意外と、トークナイザーをマージするというのがこの手法の重要なところではないかと個人的には感じました。例えばaxolotlなどのやり方ではトークナイザーを丸ごと変えることができるのですが、丸ごと変えてしまうと、そもそも事前学習と食い違ってくるからです。

wikipediaの全文をトークナイザーに学習させようとするとプログラムから「デカすぎる」とか結構文句を言われますが無視します。弊社の社長であるAIスーパーコンピュータ継之助では数時間程度でトークナイザーの学習が終了しました。

トークナイザーのマージは以下のように行います。

mkdir -p ./output_models/new_tokenizer
python utils/merge_tokenizer.py --tokenizer_dir openlm-research/open_llama_3b \
        --chinese_sp_model_file ./output_models/new_tokenizer/example.model \
        --output_dir ./output_models/merged_tokenizer 

ここではサンプルの通りopen_llama_3bのトークナイザーをベースにしましたが、本当は学習するベースモデルで使うトークナイザーに結合した方がいいような気がします。

こうして作ったトークナイザーを使うようにするためには、examplesのシェルスクリプトに直接手を入れる必要があります。この辺り、LMFlowも多言語対応をするならもう少し丁寧に作って欲しいところです。

examples/run_finetune_with_lisa.shを参考に改造を加えます

# finetune_lisa.sh

# Parses arguments
model_name_or_path=meta-llama/Llama-2-7b-hf
dataset_path=data/wikija/train
output_dir=output_models/finetune_lisa
lisa_activated_layers=1
lisa_interval_steps=20
deepspeed_args="--master_port=11000"

while [[ $# -ge 1 ]]; do
  key="$1"
  case ${key} in
    -m|--model_name_or_path)
      model_name_or_path="$2"
      shift
      ;;
    -d|--dataset_path)
      dataset_path="$2"
      shift
      ;;
    -o|--output_model_path)
      output_dir="$2"
      shift
      ;;
    --deepspeed_args)
      deepspeed_args="$2"
      shift
      ;;
    --lisa_activated_layers)
      lisa_activated_layers="$2"
      shift
      ;;
    --lisa_interval_steps)
      lisa_interval_steps="$2"
      shift
      ;;
    *)
      echo "error: unknown option \"${key}\"" 1>&2
      exit 1
  esac
  shift
done

# Finetune
exp_id=finetune
project_dir=$(cd "$(dirname $0)"/..; pwd)
log_dir=${project_dir}/log/${exp_id}
mkdir -p ${output_dir} ${log_dir}

deepspeed ${deepspeed_args} \
  examples/finetune.py \
    --model_name_or_path ${model_name_or_path} \
    --dataset_path ${dataset_path} \
    --output_dir ${output_dir} --overwrite_output_dir \
    --num_train_epochs 1 --tokenizer_name ./output_models/merged_tokenizer/merged_tokenizer_hf/ \
    --learning_rate 2e-5 \
    --block_size 512 \
    --per_device_train_batch_size 1 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    --fp16 \
    --run_name finetune \
    --validation_split_percentage 0 \
    --logging_steps 20 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 5000 \
    --dataloader_num_workers 1 \
    --gradient_checkpointing True \
    --use_lisa 1 \
    --lisa_activated_layers ${lisa_activated_layers} \
    --lisa_interval_steps ${lisa_interval_steps} \
    | tee ${log_dir}/train.log \
    2> ${log_dir}/train.err

重要なポイントは、deepspeedの呼び出し時にトークナイザーをtokenizer_nameとして明確に指定することです。この時、トークナイザ名の先頭に「./」を入れないとhuggingfaceにないとエラーが出ます。

あとはwandbにログインした状態で起動します。

学習が勝手に止まらないためにtmuxを起動してから学習を走らせます。

$ tmux
$ cd git/LMFlow
$ wandb login
$ chmod +x finetune_lisa.sh
$ ./finetune_lisa.sh   --model_name_or_path meta-llama/Llama-2-7b-hf   --dataset_path data/wikija/train   --output_model_path output_models/finetuned_llama   --lisa_activated_layers 1

wandbがなくても動きます。

起動すると綺麗にlossが下がって行きます。

手始めにllama_7bをやってみましたが、継之助に8つ搭載されているA100 80GBのGPUが平等に使われているのがわかります。

メモリ消費量は各GPUともに25%程度なので30B以上のモデルの学習もできる可能性があります。これは他の手法では難しかったのでもしこれで学習がうまくいくようなら非常に有望な方法と言えます。

Llama2-7Bモデルに日本語版WIkipediaを学習させる場合、継之助を使った学習にかかる時間はおよそ12時間と推定されています。

LMFlowにはマルチモーダルモデルの学習も含まれているようです。

こちらももし学習が上手くいけば、今までより簡単に独自のマルチモーダルモデルを作れるようになるはずです。

実際に学習ができたかどうか、また、もっと大きいモデルは学習可能かといった続報は今後もこのブログで配信していきます。