UnslothでLLMを超高速ファインチューニング:VRAMを90%削減し、Llama 3を「普及型GPU」で動かす

Artificial Intelligence tutorial - IT technology blog
Artificial Intelligence tutorial - IT technology blog

LLMのファインチューニングにおけるVRAMの悩み

「Out of Memory」(OOM)という悪夢は、Llama 3やMistralなどの大規模言語モデルを再学習(ファインチューニング)しようとする際、常に最大の障壁となります。正直なところ、以前はA100やH100のような専用のグラフィックボードがなければ、8Bや70Bクラスのモデルに手を出すのはかなり贅沢なことでした。

私は以前、従来のHuggingFaceスクリプトを使用して24GBのGPUで Llama-3-8B を学習させようとしたことがあります。結果として、VRAMは開始からわずか数秒で22GBを使い果たし、速度は気が遠くなるほど遅かったです。しかし、Unslothを知ってから、状況は一変しました。今では、ミドルレンジのカードや、Google Colabの無料版でさえ、非常に印象的なパフォーマンスでスムーズに実行できています。

現在主流のファインチューニング手法

Unslothを深掘りする前に、トレーニング手法の全体像を振り返ってみましょう。これにより、なぜこのツールがDevOpsやAIエンジニアの新たな「武器」になりつつあるのかが理解できるはずです。

フルファインチューニング (Full Fine-tuning)

これは最も「コストがかかる」方法です。モデルのすべてのパラメータを更新する必要があり、数十億のレイヤーの勾配を保存するために膨大なハードウェアリソースを必要とします。予算が無限にない限り、実務においては一旦忘れていい方法でしょう。

PEFTとLoRA/QLoRA

PEFT(Parameter-Efficient Fine-Tuning)は救世主として登場しました。すべてを更新する代わりに、いくつかの小さなレイヤー(アダプター)を追加し、そこだけを学習させます。LoRAやQLoRA(Quantized LoRA)は、元のモデルを4ビット形式で凍結することでVRAMを節約します。Unslothが登場するまでは、これがゴールデンスタンダードでした。

なぜUnslothはこれほど違うのか?

多くの人はUnslothをHuggingFaceライブラリの単なるラッパーだと思い込んでいます。実際には、開発チームはOpenAIのTriton言語を使用して、逆伝播(Backpropagation)プロセスの計算カーネルを完全に書き直しました。

このアプローチによりPyTorchの冗長な計算が排除され、メモリが低レベルで最適化されます。その結果、通常のLoRAと比較して、トレーニング速度は2〜3倍になり、VRAM消費量は最大90%削減されるという、非常に実用的な成果が得られています。

なぜ実戦プロジェクトでUnslothを選ぶのか?

最近のベトナム語技術文書処理プロジェクトで、本番環境にUnslothを適用し、高い安定性を得ることができました。最大の利点は、モデルをGGUF形式に極めて高速にエクスポートして、OllamaやvLLMで直接実行できることです。古いQLoRAスクリプトを使用している場合でも、UnslothはHuggingFaceのエコシステムと完全に互換性があるため、移行は数分で完了します。

UnslothによるLlama-3ファインチューニングの実装ガイド

始めるにあたって、GPUドライバの競合を避けるために、クリーンなLinux環境またはGoogle Colabを使用することをお勧めします。

ステップ1:環境構築

ライブラリのインストールは最も重要なステップです。Unslothはハードウェアの性能を最大限に引き出すために、特定のTorchバージョンを必要とします。

# Unslothと最適化された依存関係をインストール
pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes

ステップ2:モデルの読み込みとデータ準備

Unslothは最適化済みのモデル(Pre-quantized 4-bit)を提供しています。これにより、モデルのダウンロードと起動が瞬時に完了します。

from unsloth import FastLanguageModel
import torch

max_seq_length = 2048 # コンテキスト長
dtype = None # 自動検出 (Float16 または Bfloat16)
load_in_4bit = True # VRAMを最小限に抑える

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-bnb-4bit", 
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

データは Instruction - Input - Output の構造にフォーマットする必要があります。チャットボットを開発する場合、データの量よりも質を優先するようにしてください。

ステップ3:LoRAの設定とトレーニング開始

Unslothを使用すれば、OOMエラーを心配することなく自信を持って batch_size を増やすことができます。パラメータは次のように設定します。

model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # Rank: ほとんどのタスクには8または16で十分です
    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", # VRAMを節約するための鍵
    random_state = 3407,
)

次に、SFTTrainer を使用してトレーニングを実行します。12GBのGPUでは、per_device_train_batch_size = 2 を安全に設定できます。

from trl import SFTTrainer
from transformers import TrainingArguments

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        optim = "adamw_8bit",
        output_dir = "outputs",
    ),
)

ステップ4:モデルのエクスポート

完了後、モデルをアダプターとして保存するか、直接マージして使用できます。個人的には、ローカルで軽量に動かすためにGGUF形式でエクスポートすることが多いです。

# アダプターを保存
model.save_pretrained("lora_model")
# Ollamaで使用するためにGGUF形式でエクスポート
model.save_pretrained_gguf("model_output", tokenizer, quantization_method = "q4_k_m")

実装における「痛い目を見て学んだ」経験

何度もの実証実験を経て、時間を無駄にしないための4つの重要な注意点をまとめました。

  • CUDAバージョン: Unslothは環境に対して非常に繊細です。Tritonカーネルのパワーを最大限に活用するために、必ずCUDA 12.1以上を使用してください。
  • Rank (r) を使いすぎない: r=64 や r=128 にすればモデルが賢くなると思わないでください。実際には、通常のタスクには r=16 で十分であり、学習を速め、過学習を避けることができます。
  • 言語の問題: 特定の言語を学習させる際は、元のモデルのトークナイザーを確認してください。Llama 3は多言語対応が比較的良好で、ファインチューニング後の「言語の乱れ」を防ぐのに役立ちます。
  • 実際の数値: RTX 3060 12GB では、Llama-3-8B モデルに対して Unsloth は約 6.5GB の VRAM しか消費しません。処理速度は HuggingFace のデフォルトスクリプトより2倍速いです。

終わりに

LLMのファインチューニング分野に足を踏み入れたいなら、Unslothは試すべき必須のツールです。GPUのレンタルコストを節約できるだけでなく、アイデアの試行錯誤にかかる時間を大幅に短縮できます。

一般的なGPUやGoogle Colabで今すぐ始めてみましょう。もし設定やドライバのエラーで困ったことがあれば、コメントを残してください。私とDevOpsチームが一緒にお答えします。

Share: