🔥

Unslothを使った高速なLlama 3.2ファインチューニング入門 (📒ノートブック付)

2024/10/28に公開

高速化

はじめに

Unslothは、LLM（大規模言語モデル）のファインチューニングを大幅に高速化するライブラリーです。従来の手法と比較して約2倍の速度向上を実現し、メモリ使用量も削減できます。この記事では、Llama 3.2モデルを例に、Unslothを使用したファインチューニングの手順を解説します。

主な特徴

従来比2倍以上の学習速度
メモリ使用量の大幅な削減
多様なモデルのサポート（Llama、Mistral、Phi-3、Gemma、Yi、DeepSeekなど）
16bit LoRAまたは4bit QLoRA対応
柔軟なシーケンス長設定

環境セットアップ

まず、Unslothをインストールします。Google Colabの無料GPUインスタンスでも実行可能です。

%%capture
!pip install unsloth
# 最新版の取得
!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

モデルの準備

Llama 3.2モデルを読み込み、基本設定を行います。

from unsloth import FastLanguageModel
import torch

# 基本設定
max_seq_length = 2048  # RoPEスケーリングで自動調整
dtype = None  # GPU種別に応じて自動選択
load_in_4bit = True  # メモリ効率化のための4bit量子化

# モデルとトークナイザーの読み込み
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Llama-3.2-3B-Instruct",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

LoRAアダプターの追加

パラメータの効率的な更新のため、LoRAアダプターを追加します。

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,  # ランク（推奨値: 8, 16, 32, 64, 128）
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                     "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
)

データの準備

FineTome-100kデータセットを使用し、会話形式のデータを準備します。

from unsloth.chat_templates import get_chat_template
from datasets import load_dataset

# チャットテンプレートの設定
tokenizer = get_chat_template(
    tokenizer,
    chat_template = "llama-3.1",
)

# データセットの読み込みと整形
dataset = load_dataset("mlabonne/FineTome-100k", split = "train")

def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
    return { "text" : texts, }

dataset = dataset.map(formatting_prompts_func, batched = True)

モデルの学習

TrainingArgumentsを設定し、SFTTrainerを使用して学習を実行します。

from trl import SFTTrainer
from transformers import TrainingArguments, DataCollatorForSeq2Seq
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        output_dir = "outputs",
    ),
)

# 学習の実行
trainer_stats = trainer.train()

モデルの推論

学習したモデルを使って実際に推論を行います。

FastLanguageModel.for_inference(model)  # 高速推論の有効化

messages = [
    {"role": "user", "content": "フィボナッチ数列を続けてください: 1, 1, 2, 3, 5, 8,"},
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize = True,
    add_generation_prompt = True,
    return_tensors = "pt",
).to("cuda")

from transformers import TextStreamer
text_streamer = TextStreamer(tokenizer, skip_prompt = True)
_ = model.generate(
    input_ids = inputs, 
    streamer = text_streamer, 
    max_new_tokens = 128,
    temperature = 1.5, 
    min_p = 0.1
)

モデルの保存とHugging Faceへのアップロード

学習したモデルは様々な形式で保存およびHugging Faceにアップロードできます。

Hugging Faceへのログイン

まず、Hugging Faceにログインします。トークンは設定ページで取得できます。

from huggingface_hub import login
login()  # あるいは login(token="your_token_here")

LoRAアダプターとしての保存

# ローカル保存
model.save_pretrained("lora_model")
tokenizer.save_pretrained("lora_model")

# Hugging Faceへのアップロード
model.push_to_hub("your-username/model-name")
tokenizer.push_to_hub("your-username/model-name")

GGUF形式での保存とアップロード

llama.cppなどで使用するためのGGUF形式での保存とアップロードができます。

# 複数の量子化形式でアップロード
model.push_to_hub_gguf(
    "your-username/model-name-gguf",
    tokenizer,
    quantization_method = ["q4_k_m", "q8_0", "q5_k_m"],
)

16bit形式での保存とアップロード

VLLMなどで使用する場合は16bit形式が便利です。

# 16bit形式でのアップロード
model.push_to_hub_merged(
    "your-username/model-name-16bit", 
    tokenizer, 
    save_method = "merged_16bit"
)

4bit形式での保存とアップロード

メモリ効率を重視する場合は4bit形式が有用です。

# 4bit形式でのアップロード
model.push_to_hub_merged(
    "your-username/model-name-4bit", 
    tokenizer, 
    save_method = "merged_4bit"
)

まとめ

Unslothを使用することで、以下のメリットが得られました：

学習速度の2倍以上の向上
メモリ使用量の削減
柔軟なモデル保存オプション
簡単なHugging Faceへのモデルアップロード

学習したモデルは、ローカルでの保存だけでなく、Hugging Faceへの直接アップロードも可能で、様々な形式（LoRA、GGUF、16bit、4bit）に対応しています。これにより、異なるユースケースや環境での再利用が容易になります。

より詳細な情報や最新のアップデートについては、Unslothの公式ドキュメントを参照してください。

📒 ノートブック

リポジトリ

原点

https://github.com/unslothai/unsloth