Unsloth 與 Hugging Face TRL 整合,實現更快的 LLM 微調
TL;DR
Unsloth 是一個輕量級函式庫,旨在加速大型語言模型 (LLM) 的微調,同時減少 VRAM 使用量。它與 Hugging Face 生態系統完全整合,在不損害模型準確度的情況下,提供顯著的加速效果與記憶體效率。
技術實作:Triton Kernels 與手動反向傳播
Unsloth 透過將標準建模代碼替換為優化後的運算來實現性能提升。該函式庫手動推導反向傳播步驟,並將 PyTorch 模組重寫為 Triton kernels。
由於這些優化並未使用近似值,因此與標準 QLoRA 相比,準確度下降為 0%。Unsloth 相容於大多數 NVIDIA GPU,範圍從 GTX 1070 到 H100。
性能基準測試
在 Tesla T4 和 A100 Google Colab 實例上使用四個數據集進行 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% |
Free 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 (以及 CodeLlama)
- Mistral
- Yi
- Deepseek
- TinyLlama
- Llamafied Qwen
TRL 整合
Unsloth 模型可以直接傳遞給 TRL 的 SFTTrainer、DPOTrainer 和 PPOTrainer。這使用戶能夠在保持標準 TRL 訓練工作流的同時,利用 Unsloth 的速度。訓練後的模型與 Hugging Face 生態系統完全相容,並可以推送到 Hub 以使用 transformers 函式庫進行推論。
使用與優化
用戶可以使用 FastLanguageModel.from_pretrained 來載入模型,該函式庫內部執行 RoPE Scaling 以支援更大的最大序列長度。
為了進一步優化性能,該函式庫提供了預先量化後的 4-bit 模型(例如 unsloth/llama-2-7b-bnb-4bit),這可以減少約 500MB 的記憶體碎片,並使模型下載速度提高四倍。
對於 QLoRA 微調,用戶可以透過 FastLanguageModel.get_peft_model 附加適配器 (adapters)。當 lora_dropout 設置為 0 且 bias 設置為 "none" 時,該函式庫的優化效果最佳。