Elyza-7BをずんだもんデータセットでQLoRAファインチューニング

# Googleドライブのマウント
from google.colab import drive 
drive.mount("/content/drive")

# 作業フォルダへの移動
!mkdir -p "/content/drive/My Drive/work/"
%cd "/content/drive/My Drive/work/"

# パッケージのインストール
!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7 sentencepiece
!pip install -q wandb

# パッケージのインポート
import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

# データセットの読み込み
dataset = load_dataset("takaaki-inada/databricks-dolly-15k-ja-zundamon", split="train")

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    if data_point["input"]:
        result = f"[INST] {data_point['instruction']}\n\n{data_point['input']} [/INST] {data_point['output']}"
    else:
        result = f"[INST] {data_point['instruction']} [/INST] {data_point['output']}"
    return result

# text列の追加
def add_text(example):
    example["text"] = generate_prompt(example)
    for key in ["index", "category", "instruction", "input", "output"]:
        del example[key]
    return example
dataset = dataset.map(add_text)
print(dataset)
print(dataset[0]["text"])

Dataset({
    features: ['text'],
    num_rows: 15015
})
[INST] ヴァージン・オーストラリア航空はいつから運航を開始したのですか？

ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社なのだ。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長しました。 [/INST] ヴァージン・オーストラリア航空は、2000年8月31日にヴァージン・ブルー航空として、2機の航空機で単一路線の運航を開始したのだ。

# 量子化パラメータ
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,  # 量子化の有効化
    bnb_4bit_quant_type="nf4",  # 量子化種別 (fp4 or nf4)
    bnb_4bit_compute_dtype=torch.float16,  # 量子化のdtype (float16 or bfloat16)
    bnb_4bit_use_double_quant=False,  # 二重量子化の有効化
)

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    "elyza/ELYZA-japanese-Llama-2-7b-instruct",  # モデル名
    quantization_config=bnb_config,  # 量子化パラメータ
    device_map="auto" # モデル全体をGPU0にロード
)
model.config.use_cache = False  # キャッシュ (学習時はFalse)
model.config.pretraining_tp = 1  # 事前学習で使用したテンソル並列ランク

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "elyza/ELYZA-japanese-Llama-2-7b-instruct",  # モデル名
    use_fast=False,  # Fastトークナイザーの有効化
    add_eos_token=True,  # データへのEOSの追加を指示
    trust_remote_code=True
)
tokenizer.pad_token = tokenizer.unk_token
tokenizer.padding_side = "right" # fp16でのオーバーフロー問題対策

# LoRAパラメータ
peft_config = LoraConfig(
    r=8,  # LoRAアテンションの次元
    lora_alpha=16,  # LoRAスケーリングのAlphaパラメータ
    lora_dropout=0.05,  # LoRA レイヤーのドロップアウト確率
    bias="none",  # LoRAのバイアス種別 ("none","all", "lora_only")
    task_type="CAUSAL_LM",  # タスク種別
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "embed_tokens", "lm_head"],
)

# 学習パラメータ
training_arguments = TrainingArguments(
    output_dir="./train_logs",  # 出力ディレクトリ
    fp16=True,  # fp16学習の有効化
    bf16=False,  # bf16学習の有効化
    max_steps=300,  # 学習ステップ数
    per_device_train_batch_size=4,  # 学習用のGPUあたりのバッチサイズ
    gradient_accumulation_steps=1,  # 勾配を蓄積するための更新ステップの数
    optim="paged_adamw_32bit",  # オプティマイザ
    learning_rate=1e-4,  # 初期学習率
    lr_scheduler_type="cosine",  # 学習率スケジュール
    max_grad_norm=0.3,  # 最大法線勾配 (勾配クリッピング)
    warmup_ratio=0.03,  # 線形ウォームアップのステップ比率 (0から学習率まで)
    weight_decay=0.001,  # bias/LayerNormウェイトを除く全レイヤーに適用するウェイト減衰
    save_steps=25,  # 何ステップ毎にチェックポイントを保存するか
    logging_steps=25,  # 何ステップ毎にログを記録するか
    group_by_length=True,  # シーケンスを同じ長さのバッチにグループ化 (メモリ節約)
    report_to="wandb"  # レポート
)

# SFTパラメータ
trainer = SFTTrainer(
    model=model,  # モデル
    tokenizer=tokenizer,  # トークナイザー
    train_dataset=dataset,  # データセット
    dataset_text_field="text",  # データセットのtext列
    peft_config=peft_config,  # PEFTパラメータ
    args=training_arguments,  # 学習パラメータ
    max_seq_length=None,  # 使用する最大シーケンス長
    packing=False,  # 同じ入力シーケンスに複数サンプルをパッキング(効率を高める)
)

# モデルの学習
trainer.train()
trainer.model.save_pretrained("./lora_model")

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
import torch

# モデルの読み込み
model = AutoPeftModelForCausalLM.from_pretrained(
    "./lora_model",
    torch_dtype=torch.float16,
    load_in_4bit=True,  # 4bit量子化
    device_map="auto",
)

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "elyza/ELYZA-japanese-Llama-2-7b-instruct",  # モデル名
    trust_remote_code=True
)
tokenizer.pad_token = tokenizer.unk_token
tokenizer.padding_side = "right"

# プロンプトの準備
prompt = "[INST] 富士山の高さは？ [/INST] "

# 推論の実行
input_ids = tokenizer(prompt, add_special_tokens=False, return_tensors='pt')
output_ids = model.generate(
    **input_ids.to(model.device),
    max_new_tokens=100,
    do_sample=True,
    temperature=0.3,
)
output = tokenizer.decode(output_ids.tolist()[0])
print(output)

[INST] 富士山の高さは？ [/INST] 富士山の高さは3,776メートルなのだ。</s>