Posted inAI

Linux環境でUnslothを使った高速LLMファインチューニング：コンシューマー向けGPUでVRAM使用量を80%削減

AI tutorial - IT technology blog
AI tutorial - IT technology blog

背景と導入理由：コンシューマー向けGPUでのファインチューニングの現実

コンシューマー向けグラフィックカードで標準的なHugging FaceのPEFT/LoRAを実行しようとすると、誰もが一度は「CUDA out-of-memory（メモリ不足）」という馴染み深い壁に突き当たります。RTX 3060（12GB）やRTX 4070（12GB）でLlama 3やQwen 2.5といった8Bクラスのモデルをファインチューニングする場合、積極的なCPUオフロードや極小のバッチサイズを設定するか、クラウドインスタンスをレンタルせざるを得ません。処理速度は大幅に低下し、環境構築のストレスも溜まる一方です。

Unslothはこの問題を根本から解決します。高レベルなAPIレイヤーでPyTorchにパッチを当てるのではなく、主要な計算カーネルをOpenAIのTriton言語でゼロから書き直しています。開発チームはRoPE埋め込み、クロスエントロピー損失、MLPプロジェクション向けに手動でバックワードパス（逆伝播）を構築しました。これにより、フォワードパス中に肥大化した中間テンソルを保持し続ける必要がなくなります。

これが実践において何を意味するかは以下の通りです：

  • 最大80%のメモリ削減： 8Bパラメータのモデルでも、わずか6GB〜8GBのVRAMで快適に学習可能です。低速なシステムRAMへスワップすることもありません。
  • 2倍〜5倍の学習高速化： PCIeバスを通過するデータ量が減ることで、メモリ帯域幅によるGPUコアのボトルネックが解消されます。
  • 厳密な数学的精度： 情報落ちを伴う枝刈りや大雑把な近似は一切ありません。勾配は標準的なFP16のPyTorch学習とビット単位で完全に一致します。

特定のドメインに特化したモデルを構築する場合、この変化は開発ワークフローを一変させます。高額なマルチA100インスタンスの時間課金を気にすることなく、ローカルのデスクトップ環境で迅速に試行錯誤を繰り返すことができます。

インストール：クリーンなLinux環境のセットアップ

依存関係をクリーンに保つことで、CUDA関連のトラブルの9割は防げます。インストールを始める前に、Linuxマシン上に専用のCondaまたはvenv環境を作成しましょう。

1. NVIDIAドライバとCUDAの確認

UnslothはカスタムTritonカーネルを多用します。これらは、Turing（RTX 2000）、Ampere（RTX 3000）、Ada Lovelace（RTX 4000）などの最新NVIDIAアーキテクチャで真価を発揮します。まずはドライバの状態を確認します：

nvidia-smi

インストールされているドライバがCUDA 12.1または12.4をサポートしていること（ドライババージョン535.xx以降）を確認してください。

2. 独立した環境の作成

# クリーンなconda環境を作成して有効化
conda create -n unsloth_env python=3.10 -y
conda activate unsloth_env

3. PyTorchとUnslothのインストール

CUDA 12.1向けにビルドされたPyTorchバイナリを取得し、続いてUnslothとその主要な依存パッケージをインストールします：

# CUDA 12.1対応のPyTorchをインストール
pip install --upgrade --no-cache-dir torch torchvision torchaudio \
    --index-url https://download.pytorch.org/whl/cu121

# Unslothおよび主要な高速化ライブラリをインストール
pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
pip install --no-deps trl peft accelerate bitsandbytes datasets xformers

設定：Llama 3およびQwenモデルのチューニング

UnslothはHugging FaceのTRL（Transformer Reinforcement Learning）エコシステムにそのまま組み込むことができます。以下のスクリプトは、4-bit NormalFloat（NF4）量子化とLoRAアダプタを組み合わせて、Llama 3（8B）またはQwen 2.5（7B）を学習させる構成です。

1. 4-bit量子化ベースモデルの読み込み

以下の設定をtrain_lora.pyとして保存します：

import torch
from unsloth import FastLanguageModel

max_seq_length = 2048 # データセットの要件に応じて調整
dtype = None          # None = 自動検出（旧世代カードはFloat16、Ampere以降はBfloat16）
load_in_4bit = True   # ベースモデルのメモリ使用量を6GB未満に維持

# Qwenを使用する場合は unsloth/Qwen2.5-7B-Instruct-bnb-4bit に変更
model_name = "unsloth/Meta-Llama-3-8B-Instruct-bnb-4bit"

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

2. 高速LoRAアダプタのアタッチ

次に、線形射影層にパッチを適用します。UnslothはPythonのオーバーヘッドを発生させることなく、AttentionおよびMLPの射影層に直接介入します：

model = FastLanguageModel.get_peft_model(
    model,
    r=16, # Rankの次元数（16または32で十分な推論能力を発揮）
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    lora_dropout=0,       # Tritonカーネルによる高速化のため0%に最適化
    bias="none",
    use_gradient_checkpointing="unsloth", # アクティベーションメモリを約30%削減
    random_state=3407,
)

3. データセットのフォーマットとTrainerの実行

対話データをモデル固有のテンプレートにマッピングします。学習ループの実行にはHugging FaceのSFTTrainerを使用します：

from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

# データをモデル固有のチャットテンプレートにマッピング
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 = load_dataset("philschmid/dolly-15k-oai-style", split="train")
dataset = dataset.map(formatting_prompts_func, batched=True)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    dataset_num_proc=2,
    packing=False, # 短いサンプルを結合して学習を高速化する場合はTrueに設定
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=60,
        learning_rate=2e-4,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=5,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
    ),
)

trainer_stats = trainer.train()

検証とモニタリング：VRAMの確認と重みのエクスポート

1. ハードウェアのリアルタイムモニタリング

スクリプトを実行し、別のターミナルを開いてハードウェアのリソース消費状況をリアルタイムで監視します：

# GPUメトリクスを毎秒取得
watch -n 1 nvidia-smi

メモリ消費の推移を確認してください。RTX 3090環境において、2,048トークンのウィンドウで8Bモデルを学習させた場合、VRAM消費量は約6.8GB〜7.4GBでピークを迎えます。最適化されていない標準的なLoRAでは、まったく同じパラメータでも22GBを容易に超えてしまいます。この十分な空き容量があれば、OOMによるクラッシュを心配することなく、バッチサイズを2倍に増やしたりコンテキスト長を拡張したりすることが可能です。

2. ローカルでの推論精度の検証

学習が完了したら、スクリプト内で直接推論テストを行います。新しいアダプタをインメモリでマージし、モデルがプロンプトへ正確に従うかを確認します：

# ネイティブの2倍高速な推論エンジンを有効化
FastLanguageModel.for_inference(model)

messages = [
    {"role": "user", "content": "Unslothが学習時のVRAM使用量を削減する仕組みをわかりやすく説明してください。"}
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt"
).to("cuda")

outputs = model.generate(input_ids=inputs, max_new_tokens=128, use_cache=True)
print(tokenizer.batch_decode(outputs))

3. GGUFまたは16-bit Hugging Face形式へのエクスポート

デプロイの準備が整ったら、別個のPython変換スクリプトを用意することなく、任意のランタイム向けにアダプタを直接エクスポートできます：

# vLLMやTGI向けに16-bitでマージしたHugging Face形式でエクスポート
model.save_pretrained_merged("llama3_custom_model", tokenizer, save_method="merged_16bit")

# またはOllamaやllama.cpp向けに量子化済みGGUF形式へ直接書き出し
model.save_pretrained_gguf("llama3_gguf", tokenizer, quantization_method="q4_k_m")

Tritonレベルのカーネル最適化と4-bit量子化を活用すれば、最新の大規模言語モデルのファインチューニングに何千ドルものクラウドインフラ費用を投じる必要はありません。手元のデスクトップGPUで十分にやり遂げることができます。

Share:
Tags: