UnslothとHugging Face TRLの統合によるLLMファインチューニングの高速化

TL;DR

Unslothは、VRAM使用量を削減しながらLarge Language Model（LLM）のファインチューニングを加速するように設計された軽量ライブラリです。Hugging Faceエコシステムと完全に統合され、モデルの精度を損なうことなく大幅なスピードアップとメモリ効率を提供します。

技術的実装: Tritonカーネルと手動バックプロパゲーション

Unslothは、標準のモデリングコードを最適化された操作に置き換えることで性能向上を達成します。ライブラリはバックプロパゲーションの手順を手動で導出し、PyTorchモジュールをTritonカーネルに書き換えます。

これらの最適化は近似を使用しないため、標準のQLoRAと比較して0%の精度劣化が発生しません。Unslothは、GTX 1070からH100までのほとんどのNVIDIA GPUと互換性があります。

パフォーマンスベンチマーク

Tesla T4とA100のGoogle Colabインスタンスで4つのデータセットを使用して59回の実行を行ったベンチマークでは、Unslothは標準のHugging Face実装およびFlash Attention 2を大幅に上回ることが示されました。

A100 40GB パフォーマンス

モデル	データセット	Hugging Face	HF + Flash Attention 2	Unsloth	VRAM削減
Code Llama 34b	Slim Orca	1x	1.01x	1.94x	-22.7%
Llama-2 7b	Slim Orca	1x	0.96x	1.87x	-39.3%
Mistral 7b	Slim Orca	1x	1.88x	-65.9%
Tiny Llama 1.1b	Alpaca	1x	1.55x	2.74x	-57.8%
DPO with Zephyr	Ultra Chat	1x	1.24x	1.88x	-11.6%

無料Colab T4 パフォーマンス

モデル	データセット	Hugging Face	HF + Pytorch 2.1.1	Unsloth	VRAM削減
Llama-2 7b	OASST	1x	1.19x	1.95x	-43.3%
Mistral 7b	Alpaca	1x	1.07x	1.56x	-13.3%
Tiny Llama 1.1b	Alpaca	1x	2.06x	3.87x	-73.8%
DPO with Zephyr	Ultra Chat	1x	1.09x	1.55x	-18.6%

Transformersバージョン4.36と比較して、Unslothは最大2.7倍の高速なトレーニングと最大74%少ないメモリ使用量を提供します。

Hugging Faceエコシステムとの統合

Unslothは、Hugging Face Hub、transformers、PEFT、TRLライブラリと完全に互換性があります。

モデルサポート

Unslothは現在、以下のLlamaタイプのアーキテクチャをサポートしています：

Llama (and CodeLlama)
Mistral
Yi
Deepseek
TinyLlama
Llamafied Qwen

TRL統合

Unslothのモデルは、TRLのSFTTrainer、DPOTrainer、PPOTrainerに直接渡すことができます。これにより、ユーザーは標準のTRLトレーニングワークフローを維持しながらUnslothの速度を活用できます。その結果得られたトレーニング済みモデルは、Hugging Faceエコシステムと完全に互換性があり、transformersライブラリを使用してHubにプッシュして推論を行うことができます。

使用方法と最適化

ユーザーはFastLanguageModel.from_pretrainedを使用してモデルをロードできます。これは内部でRoPEスケーリングを実行し、より長い最大シーケンス長をサポートします。

パフォーマンスをさらに最適化するため、ライブラリは事前に量子化された4ビットモデル（例: unsloth/llama-2-7b-bnb-4bit）を提供します。これにより、メモリフラグメンテーションが約500MB削減され、モデルのダウンロードが4倍高速になります。

QLoRAファインチューニングでは、ユーザーはFastLanguageModel.get_peft_modelを介してアダプターを取り付けます。ライブラリはlora_dropoutが0に設定され、biasが"none"に設定されているときに最も最適化されます。