LLMを安価に自分専用にする「LoRAファインチューニング」実践ガイド

# 仮想環境の作成（必須だ、システムを汚すな）
python -m venv venv
source venv/bin/activate

# 必要なライブラリのインストール
pip install -U pip
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers datasets accelerate peft bitsandbytes bitsandbytes

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# 1. モデルを4-bitでロード（メモリ節約の要）
model_id = "meta-llama/Llama-2-7b-hf"  # または任意のモデル
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# 2. LoRA設定
# r（ランク）は8〜16がコスパ最強。128とかにしても精度は大して変わらず重くなるだけだ。
config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"], # モデルによって変える必要がある
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# 3. LoRAモデルの適用
model = get_peft_model(model, config)
model.print_trainable_parameters()

from trl import SFTTrainer
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./lora-output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4, # LoRAならこの程度が妥当
    num_train_epochs=3,
    logging_steps=10,
    fp16=True, # 16-bit混合精度
    save_strategy="steps",
    save_steps=100,
)

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset, # 事前にロードしたdatasetsオブジェクト
    dataset_text_field="text",
    max_seq_length=512,
    args=training_args,
)

trainer.train()

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate ...