きっかけ

「松尾・岩澤研の大規模言語モデル2024参戦期」として書き残していきます。この講座に関連して自分で調べたり、なんだかんだしたものを書いていきます。

ちなみに、僕は普通のおじさんです。ちょっと他の人よりいろいろなことに興味があるというだけのおじさんです。
それが、日本の最高学府である東京大学、それも松尾研の講座を受け、最終課題のコンペで苦労したところを共有できればと思ったのがきっかけです。

最終課題が難問

最後の課題が、コンペだったんです。そう、大規模言語モデルをファインチューニングして精度を競うっていうコンペ。そのなかで「これは！」と思う内容をここに残しおかないともったいない！そうおもいました。

※コンペの建付け等はココでは話題にしませんのであしからず。あくまでコンペを通じて学んだことをここに残しておきたいというのが目的です。

コンペのざっくりした内容

参加規約上、おそらく詳しくは書けませんのでざっくりと。
「9月以降のテレビ情報を学習させ、うまくチャット形式で回答させる」というコンペです。いくつかの事前学習済みモデルが指定されていましたが、いずれも9月以降の情報では学習されていません。
今回僕が用いたモデルはllm-jp/llm-jp-3-13bです。2024年春に発表された事前学習済みです。つまり、まだ石破さんは総理大臣になっていません。新しく2024年9月、10月のニュース記事を継続事前学習することで、新しい知識を埋め込みます。
その後、チャット形式で回答できるようにSupervised fine-tuningをする予定ですが、ここでは継続事前学習を行います。

全体の流れは以下の通り

事前学習済みモデルの選択
継続事前学習（今回はここ！）
Supervised fine-tuning

ちなみに、llm-jpを選んだのはなんだかカッコいいから。笑

継続事前学習（CPT）とは

すばり、ここに書いてある！
というのは雑すぎるので、簡単に説明。
事前学習済みのモデルに追加して情報を入れるってことです。この学習後は単に文章を生成するだけの状態です。

SFTとの違いといっても、SFTについては今回は書かないのですが、学習する層が異なります。

ドキュメント、ドキュメントのサンプルコードを読んだ結果、肝は3点

UnslothTrainerを使うこと
target_modulesに「"lm_head", "embed_tokens"」を付け加えること（UnslothTrainerを使うことで追加可能）
embedding層の学習率を他の層よりも2~10倍小さく設定すること（embedding_learning_rateで設定）

継続事前学習のコードと解説

基本的にGoogleColaboratryでの実行を想定しています。

あと、松尾・岩澤研のPaper＆HacksのYoutubeにとってもいい動画があるので、マジでガン見しました。笑

環境構築

unslothのインストール

!pip install unsloth
!pip uninstall unsloth -y # ここまではエラー回避
!pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" # unslothのGitHubより

なぜか、一度目のインストールではエラーになるので、一度インストールしてアンインストール、そのうえでインストールします。（謎）
それ以外はドキュメントのまま。

事前Colabにインストールされているライブラリのアップデート

!pip install --upgrade torch
!pip install --upgrade xformers

flashアテンションの導入（GPUによるのでIF文入れてます）

import torch
if torch.cuda.get_device_capability()[0] >= 8:
    !pip install --no-deps packaging ninja einops "flash-attn>=2.6.3"

継続事前学習のメインとなるコード

ここからは（サンプルコードと似た（同じ？）ところもありますが）、ドキュメントやunslothのサンプルコードを見たりと、自分が調べた結果のコードになります。間違いの指摘、アドバイスなどいただけると嬉しい。

※講座のサンプルコードはおそらく規約上公開できないと思いますので公開しません。あしからず。（でも書き方が似ちゃったかも・・・💦）

ライブラリのインポートと環境変数の読み込み

import torch
from unsloth import FastLanguageModel
from peft import peft_model

# token類の呼び出し
from google.colab import userdata
HF_TOKEN = userdata.get('HF_TOKEN')
WANDB_KEY = userdata.get('wandb')

今回はwandbを使いました。とても便利で、学習の経過がわかるのはもちろん、学習中に外出してwandbをみて、学習のセッションを切ってGoogleColaboratryのコンピューティングユニットを無駄にしないようにも使ったりしました。

かなり、講座で家庭の時間を浪費したので、妻に感謝するとともに、妻との時間を大事にしながらコンピューティングユニットの消費を抑えるのにとても役立ちました。（目的違っ

もろもろの設定

max_seq_length = 1024 # unslothではRoPEをサポート
dtype = None
load_in_4bit = True

model_id = "llm-jp/llm-jp-3-13b"
peft_model_id = "ikedachin/llm-jp-3-13b-october-news-e2" # Huggingfaceにアップロードするときのリポジトリ名

ベースモデルを量子化してダウンロード

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_id,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
    trust_remote_code = True,
)

そのままダウンロードするとGoogleColaboratryのメモリーが足りずに学習できないので、量子化してダウンロードします。
モデルのダウンロードが終わったらprint(model)をしてどんな層があるのかしっかり確認しておきましょう。
そして、transformerの理解と、Llama構造の違いもざっくり理解しておくと後で役立ちます。

いったんここでデータセットを用意

CPT用データセットの用意

from datasets import load_dataset
dataset = load_dataset("ikedachin/october-news-dataset") # このデータセットはダミーです

ここで継続事前学習のためのデータセットをHuggingFaceから探してきます。
本来はライセンスを気にしながらモデルを選択したり、データセットを選択する必要があります。講義のコンペでは皆さんの議論から抽出して選択しましたが、よくわかりません。終了後のアンケートにはライセンスの講義もしてほしい旨のお願いをさせていただきました。やってくれると嬉しいなぁ・・・。お願いします！！

PEFTモデルを作る

ここからPEFTモデルを作っていきます。
PEFTモデルってどういうことかというと、ベースモデルとは別に二つの小さな行列を使ってパラメータを作り、その小さな行列を学習させて、ベースモデルと合わせて推論する方法です。パラメ数がぐっと減るんです。こんな感じ。

https://arxiv.org/abs/2106.09685 より
この右側の赤いところがPEFTの小さな二つの行列
パラメータの元の行列がd×dの行列だとした場合、d×rの行列とr×dの行列の内積をとったものと同じサイズになりますよね。これを使った仕組みです。（これ以上詳しく言えない。💦）
このd×rとr×dの行列を学習していきます。

参考にしたのはunsloth.aiのこのページ。

これはマジでしっかり読みました。サンプルコードも。
大事なのはどの層にPEFTモデルを適用するかという点と、それらのパラメータの学習率をどうするかという点。

まずはココで、どの層にPEFTモデルを適用するかということをしっかりと設定します。
ここで、上に書いた「print(model)をしてどんな層があるのかしっかり確認しておきましょう。」と言うのが生きます。なのでモデル構造を絶対に確認したほうが理解の進みが良いです。

["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head",]を指定しました。
ちなみに、SFT（出力調整のための学習、Supervised fine-tuning）では"embed_tokens", "lm_head"は指定しないようです。

学習するPEFTモデルを作成

model = FastLanguageModel.get_peft_model(
    model,
    r = 32, # LoRAのランク
    target_modules =  ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head",],
    lora_alpha = 16,
    lora_dropout = 0.2,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
    max_seq_length = max_seq_length,
)

こちらのドキュメントも非常に有用でした！（1/15追記）

学習のための準備

ライブラリの読み込み

from transformers import TrainingArguments
from unsloth import is_bfloat16_supported
from unsloth import UnslothTrainer, UnslothTrainingArguments

TrainingArgsの設定

training_args = UnslothTrainingArguments(
    run_name = "continue_learning_test", # wandbのプロジェクトのRUN名
    per_device_train_batch_size = 2, # デバイスごとのトレーニングバッチサイズ
    gradient_accumulation_steps = 2, # 勾配を更新する前にステップを積み重ねる回数
    optim = "adamw_8bit", # オプティマイザの設定
    # weight_decay = 0.01, # 指定する場合
    num_train_epochs = 1, # エポック数
    logging_steps = 10, # ログを出力するステップ間隔
    warmup_steps = 10, # 学習率のウォームアップステップ数
    save_strategy = 'steps',
    save_steps = 40, # モデルを保存するステップ間隔
    save_total_limit = 2, # 保存しておくcheckpointの数
    max_steps = -1, # トレーニングの最大ステップ数
    learning_rate = 3e-5, # 学習率
    embedding_learning_rate = 3e-6, # embed層の学習率（他の2-10x小さい値）
    fp16 = not is_bfloat16_supported(), # GPUによってはbf16が使えないときはこっちがTrue
    bf16 = is_bfloat16_supported(), # bf16が使えるときはTrue
    group_by_length=True,
    seed = 3407,
    output_dir = 'drive/MyDrive/model', # モデルの保存場所（ここではGoogleDriveに接続したとき）
    report_to = "wandb", # ログの送信先 ("wandb"/"tensorboard"など)
    auto_find_batch_size = True, # バッチサイズの自動設定
)

ざくざくとコメントを書いておきました。
この中で大事なのがlearning_rateとembedding_learning_rateの二つ
ドキュメントに書いてあったのが、「embedding_learning_rateは``の1/2~1/10に設定しなさい！」ということ。これで破綻なく、継続事前学習が実行できるみたいです。

Trainerクラスを使わずに、UnslothのUnslothTrainerクラスを使います。
使い方は多分ほとんど同じ。

UnslothTrainerの設定

trainer = UnslothTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset["train"],
    max_seq_length = max_seq_length,
    dataset_text_field = "text", # datasetのテキストが入ったkeyを指定
    packing = False,
    args = training_args,
    neftune_noise_alpha = 0.2
)

ここで、データセットの中のkeyに学習データが入っているかを指定する必要があります。

学習の実行

学習

trainer_stats = trainer.train()

もうこれだけです。
待つだけ。
しかも、パラメ数のわりに速い！量子化とunslothのおかげだと思います。

推論

出力

FastLanguageModel.for_inference(model) # 推論モードへの切り替え

text = '石破'
input = tokenizer([text], return_tensors = "pt").to(model.device)
output = model.generate(**input, max_new_tokens = max_seq_length, use_cache = True, do_sample=False, repetition_penalty=1.2)
print(tokenizer.decode(output[0], skip_special_tokens=True))

# 石破茂首相は30日、衆院選の結果を受けて記者会見し、自民党総裁を続投する意向を表明した。「国民に信任された」と強調。<以下省略>

それでは推論してみましょう。
unslothでは推論モードと学習モードがあるので、推論モードに切り替えてから出力します。
入力に「石破」（呼び捨てごめんなさい）と入れると石破さんが自民党総裁を続投すると意思表明してますね。自民党総裁であるという知識は入っているようです。

モデルをHuggingfaceへアップロード

Huggingfaceにアップロードする方法です。

LoRAモデルのみのアップロード

まずはLoRAモデルをアップロードする方法。
この方法はHuggingfaceの自分の容量をあまり消費することがないです。無料枠だと110GBまでなので注意が必要です。

Huggingfaceへのアップロード（LoRAモデルのみのアップロード）

model.push_to_hub_merged(
    new_model_id, # 保存するレポジトリ名
    tokenizer=tokenizer, # tokenizerも保存する場合（基本入れますわなぁ・・・）
    save_method="lora", # LoRAモデルだけを保存するとき
    token=HF_TOKEN, # 保存するアカウントのアクセストークン
    private=True # private設定の時はTrue、publicの時はFalse
)

これで保存できます。

ベースモデルとマージしてアップロード

実はもう一つの方法があります。
これを知るのに時間がかかりました。

Huggingfaceへのアップロード（ベースモデルとマージしてアップロード）

model = model.to(dtype=torch.bfloat16) # 念の為、bf16に変換する（必要かどうかは未検証）
model.push_to_hub_merged(
    new_model_id, # 保存するレポジトリ名
    tokenizer=tokenizer, # tokenizerも保存する場合（基本入れますわなぁ・・・）
    save_method="merged_16bit", # LoRAモデルとベースモデルをマージして保存するとき
    token=HF_TOKEN, # 保存するアカウントのアクセストークン
    private=True # private設定の時はTrue、publicの時はFalse
)

レポジトリ名にはtrlではアカウント名を除いたレポジトリ名のみを指定する必要がありますが、unslothではアカウント名が入っていても自動で切り替えてくれました。いつもやらかしているのがバレますね。笑

もっとも注意しなければならないのは、使ったGPUで使えるfloatタイプが違うということかな。

実はアップロード時にマージする方法を知らなかった！

上記のマージモデルを保存する方法を知る前は以下のように後でマージすることをやっていました。

しかしながら、さすがにマージして保存する方法が用意されてましたね。さすが、unsloth。
試行錯誤や悩み、困りごとを諦めていなかったので、無事に解決できたので嬉しいもんです。

終わりに

ということで、継続事前学習ができました。
これで、新しい知識をモデルに埋め込むことができました。
あとはHuggingFaceにアップロードして、model_cardを記載。公開したり、しなかったり。

次に、PEFTモデルとベースモデルをマージする方法について書きたいと思います。

この投稿は・・・

松尾・岩澤研の「大規模言語モデル2024」を受講して、自分で調べたことを書いています。
この講座、めっちゃ楽しかったし、有用なのでぜひ25年度には参加してみてください。

合言葉は「仕事している場合じゃない！」です。（なにそれ？ｗ）