Unslothで指示チューニング

はじめに

UnslothでInstructionTuningをやってみます。
以下の記事で継続事前学習を行ったモデルに対して、指示チューニングを行います。

継続事前学習では、新しい知識を埋め込むことができたとは言えない状況でしたが、そのまま指示チューニングまで続けました。

実装

準備

ライブラリ群のインポートとパラメータの設定。モデルのロードまで実行します。

import os
from unsloth import FastLanguageModel
from unsloth import is_bfloat16_supported
from unsloth import UnslothTrainingArguments
from trl import SFTTrainer
from datasets import load_dataset
import wandb
os.environ["WANDB_PROJECT"] = "llama3.2-sft" 

# 基本設定
max_seq_length = 2048
dtype = None
load_in_4bit = True
num_proc = 4
sft_output = 'sft_output'
random_seed = 3407

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = 'model',
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

データセットの準備

データセットにはdatabricks-dolly-15k-jaを使用します。

kunishou/databricks-dolly-15k-ja · Datasets at Hugging Face

コンテキストがある場合とない場合のデータが存在するので、2つの場合に対応できるようにテンプレートを2つ用意します。

# データダウンロード
dataset = load_dataset("llm-jp/databricks-dolly-15k-ja")
print(dataset['train'][3])

# 学習時のプロンプトフォーマットの定義
prompt1 = """
<|start_header_id|>system<|end_header_id|>
あなたは以下のツールを使用できる指示に忠実で優秀なAIアシスタントです。
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
{}
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
{}
"""

prompt2 = """
<|start_header_id|>system<|end_header_id|>
あなたは以下のツールを使用できる指示に忠実で優秀なAIアシスタントです。
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
{}

## 指示
{}
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
{}
"""

EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
def formatting_prompts_func(examples):
    input = examples["instruction"] # 入力データ
    output = examples["response"] # 出力データ
    context = examples['context']
    if context == '':
        text = prompt1.format(input, output) + EOS_TOKEN
    else:
        text = prompt2.format(context, input, output) + EOS_TOKEN
    return { "text" : text, } # 新しいフィールド "formatted_text" を返す

# 各データにフォーマットを適用
dataset = dataset.map(
    formatting_prompts_func,
    num_proc= num_proc, # 並列処理数を指定
)
dataset

学習設定

指示チューニングは、継続事前学習で追加したターゲットlm_headとembed_tokensを除外します。これは、学習対象のパラメータを削減する目的と指示チューニングでは知識に変更を加えず、回答時のアライメントのみを変更することを目的としてチューニングを行うためです。

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    use_gradient_checkpointing=False,
    random_state = random_seed,
    max_seq_length = max_seq_length,
    use_rslora=True,
    # use_dora=True,
)

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset['train'],
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    tokenizer = tokenizer,
    neftune_noise_alpha=5,
    packing=False,

    args = UnslothTrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=8,
        num_train_epochs=3,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        logging_steps = 10,
        save_total_limit = 2,
        output_dir = sft_output,
    #   optim = "schedule_free_adamw",
        learning_rate = 5e-5, 
        optim="adamw_8bit",
        lr_scheduler_type="constant", 
        seed = random_seed,
        report_to = "wandb",
  ),
)

学習の実行

trainer_stats = trainer.train()
wandb.finish()

学習の実行はこれだけ。wandb.finish()を実行しないとwandbが正しく終了しない。

モデルの保存

指示チューニングの結果をAdapterとして、LoRAで学習したパラメータのみを保存することにします。

model.save_pretrained_merged("instruct_lora_model", tokenizer, save_method = "lora",)
del model

結果の確認

結果の確認には以下のコードを使います。

import os
from unsloth import FastLanguageModel
from peft import PeftModel

# 基本設定
max_seq_length = 2048
dtype = None
load_in_4bit = False
num_proc = 4

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = './model',
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)
model = PeftModel.from_pretrained(model, "./instruct_lora_model")

prompt_template = """
<|start_header_id|>system<|end_header_id|>
あなたは以下のツールを使用できる指示に忠実で優秀なAIアシスタントです。
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
2020年以降の流行語大賞を振り返ると、日本社会の重要な出来事や話題が反映されています：
2020年の「3密」は、新型コロナウイルス感染症対策として小池百合子東京都知事が提唱した「密閉・密集・密接」を避けるという概念を表し、パンデミック初期の日本社会を象徴する言葉となりました。
2021年は、メジャーリーグで投打の「二刀流」として歴史的な活躍を見せた大谷翔平選手の「リアル二刀流/ショータイム」が選ばれ、日本人選手の世界での躍進を表しました。
2022年の「村神様」は、プロ野球で史上初となる高校生ドラフト1位から4年連続本塁打王を達成した村上宗隆選手のニックネームで、若き天才打者の台頭を示しています。
2023年の「アレ（A.R.E）」は、阪神タイガースの岡田彰布監督が頻繁に使用した言葉で、44年ぶりのリーグ優勝を果たした阪神の快進撃を象徴しました。
2024年は、TBSドラマ『不適切にもほどがある!』から生まれた「ふてほど」が選ばれ、厳しい社会規範の中で生きる現代人の共感を呼んだことを示しています。
この5年間の流行語は、感染症対策から、スポーツ界の活躍、そして時代を反映したドラマまで、その時々の日本社会の関心事を端的に表現しています。

### 指示
2024年の流行語大賞は何ですか？
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
"""
inputs = tokenizer([prompt_template]*1, return_tensors = "pt").to("cuda")

FastLanguageModel.for_inference(model) # Enable native 2x faster inference
text_streamer = TextStreamer(tokenizer)
_ = model.generate(**inputs, streamer = text_streamer,
    max_new_tokens =1024,
    use_cache=True,
    do_sample=False,
    repetition_penalty=1.2
)

出力

「ふてほど」というキーワードを抽出することに成功したものの、文章としてはおかしな状態になっています。

<|begin_of_text|>
<|start_header_id|>system<|end_header_id|>
あなたは以下のツールを使用できる指示に忠実で優秀なAIアシスタントです。
<|eot_id|>
<|start_header_id|>user<|end_header_id|>
2020年以降の流行語大賞を振り返ると、日本社会の重要な出来事や話題が反映されています：
2020年の「3密」は、新型コロナウイルス感染症対策として小池百合子東京都知事が提唱した「密閉・密集・密接」を避けるという概念を表し、パンデミック初期の日本社会を象徴する言葉となりました。
2021年は、メジャーリーグで投打の「二刀流」として歴史的な活躍を見せた大谷翔平選手の「リアル二刀流/ショータイム」が選ばれ、日本人選手の世界での躍進を表しました。
2022年の「村神様」は、プロ野球で史上初となる高校生ドラフト1位から4年連続本塁打王を達成した村上宗隆選手のニックネームで、若き天才打者の台頭を示しています。
2023年の「アレ（A.R.E）」は、阪神タイガースの岡田彰布監督が頻繁に使用した言葉で、44年ぶりのリーグ優勝を果たした阪神の快進撃を象徴しました。
2024年は、TBSドラマ『不適切にもほどがある!』から生まれた「ふてほど」が選ばれ、厳しい社会規範の中で生きる現代人の共感を呼んだことを示しています。
この5年間の流行語は、感染症対策から、スポーツ界の活躍、そして時代を反映したドラマまで、その時々の日本社会の関心事を端的に表現しています。

### 指示
2024年の流行語大賞は何ですか？
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>
2024年に発表されたのは、「ふてほど」のタイトルのもの。  この名前は、テレビ番組『不適切にもほどがない！』（2007～2018）の出演者たちによって作られたとされている。このシリーズでは、さまざまな役柄で登場する人物もいくつかいるため、この名前にちなんだことだろう。しかし、これらはすべて別個のキャラクターである可能性があります。その中には、『不適切にもほどはない！』に出場していた元テニスのエースでもあることがあります。また、同じような性格を持った他の俳優も含まれていて、どのような意味になるのかわかりません。
<|end_of_text|>

抽出タスクを行うだけなら、元の状態のLlama3.2の方がよい結果を出しますね。