Google Colab で Llama 3 のファインチューニングを試す

# パッケージのインストール
!pip install -U transformers accelerate bitsandbytes
!pip install trl peft wandb
!git clone https://github.com/huggingface/trl
%cd trl

# HuggingFaceへのログイン
!huggingface-cli login

    raw_datasets = load_dataset(args.dataset_name)
    train_dataset = raw_datasets["train"]
    eval_dataset = raw_datasets["test"]

    # データセットの読み込み
    dataset = load_dataset("bbz662bbz/databricks-dolly-15k-ja-gozarinnemon", split="train")
    dataset = dataset.filter(lambda example: example["category"] == "open_qa")

    # プロンプトの生成
    def generate_prompt(example):
        messages = [
            {
                'role': "system",
                'content': "あなたは日本語で回答するAIアシスタントです。"
            },
            {
                'role': "user",
                'content': example["instruction"]
            },
            {
                'role': "assistant",
                'content': example["output"]
            }
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False)

    # textカラムの追加
    def add_text(example):
        example["text"] = generate_prompt(example)
        return example

    dataset = dataset.map(add_text)
    dataset = dataset.remove_columns(["input", "category", "output", "index", "instruction"])

    # データセットの分割
    train_test_split = dataset.train_test_split(test_size=0.1)
    train_dataset = train_test_split["train"]
    eval_dataset = train_test_split["test"]

# 学習
!python examples/scripts/sft.py \
    --model_name meta-llama/Meta-Llama-3-8B-Instruct \
    --dataset_name bbz662bbz/databricks-dolly-15k-ja-gozaru \
    --dataset_text_field text \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 1 \
    --learning_rate 2e-4 \
    --optim adamw_torch \
    --save_steps 50 \
    --logging_steps 50 \
    --max_steps 500 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16 \
    --lora_dropout 0.1 \
    --load_in_4bit \
    --report_to wandb \
    --output_dir Llama-3-Gozaru-8B-Instruct

{'loss': 1.9094, 'grad_norm': 1.3531537055969238, 'learning_rate': 0.00018, 'epoch': 0.03}
{'loss': 1.354, 'grad_norm': 0.7856459617614746, 'learning_rate': 0.00016, 'epoch': 0.06}
{'loss': 1.3859, 'grad_norm': 0.6068201661109924, 'learning_rate': 0.00014, 'epoch': 0.09}
{'loss': 1.3009, 'grad_norm': 1.1571404933929443, 'learning_rate': 0.00012, 'epoch': 0.12}
{'loss': 1.2169, 'grad_norm': 0.7323430776596069, 'learning_rate': 0.0001, 'epoch': 0.15}
{'loss': 1.287, 'grad_norm': 0.7343457341194153, 'learning_rate': 8e-05, 'epoch': 0.18}
{'loss': 1.31, 'grad_norm': 0.8455875515937805, 'learning_rate': 6e-05, 'epoch': 0.22}
{'loss': 1.2687, 'grad_norm': 0.8726828694343567, 'learning_rate': 4e-05, 'epoch': 0.25}
{'loss': 1.3166, 'grad_norm': 0.543333888053894, 'learning_rate': 2e-05, 'epoch': 0.28}
{'loss': 1.3012, 'grad_norm': 0.7505847811698914, 'learning_rate': 0.0, 'epoch': 0.31}
{'train_runtime': 542.2354, 'train_samples_per_second': 1.844, 'train_steps_per_second': 0.922, 'train_loss': 1.3650647735595702, 'epoch': 0.31}
100% 500/500 [08:02<00:00,  1.04it/s]

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct"
)
model = AutoModelForCausalLM.from_pretrained(
    "./Llama-3-Gozaru-8B-Instruct",
    device_map="auto",
    torch_dtype="auto",
)

import torch

# プロンプトの準備
chat = [
    { "role": "system", "content": "あなたは日本語で回答するAIアシスタントです。" },
    { "role": "user", "content": "まどか☆マギカでは誰が一番かわいい?" },
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

# 推論の実行
token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
with torch.no_grad():
    output_ids = model.generate(
        token_ids.to(model.device),
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
        max_new_tokens=256,
        eos_token_id=[
            tokenizer.eos_token_id,
            tokenizer.convert_tokens_to_ids("<|eot_id|>")
        ],
    )
output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1) :], skip_special_tokens=True)
print(output)

import gc
from peft import PeftModel

# メモリ解放
del model
gc.collect()
gc.collect()
torch.cuda.empty_cache()

# モデル名
base_model = "meta-llama/Meta-Llama-3-8B-Instruct"
new_model = "Llama-3-Gozaru-8B-Instruct"

# トークナイザの準備
tokenizer = AutoTokenizer.from_pretrained(base_model)

# モデルの準備
fp16_model = AutoModelForCausalLM.from_pretrained(
    base_model,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="auto",
)

# アダプタをベースモデルにマージ
model = PeftModel.from_pretrained(fp16_model, new_model)
model = model.merge_and_unload()

# HuggingFace Hubへのアップロード
new_model = "Llama-3-Gozaru-8B-Instruct"
model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)