Unsloth 和 Hugging Face TRL 集成以实现更快的 LLM 微调
TL;DR
Unsloth 是一个轻量级库，旨在在减少 VRAM 使用的同时加速大型语言模型（LLM）的微调。它完全与 Hugging Face 生态系统集成，在不牺牲模型精度的情况下提供显著的加速和内存效率。
技术实现：Triton 内核和手动反向传播
Unsloth 通过用优化操作替换标准建模代码来实现性能提升。该库手动推导反向传播步骤，并将 PyTorch 模块重写为 Triton 内核。
由于这些优化不使用近似方法，与标准 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% |
免费 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.7x 的更快训练速度和最高 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 Scaling 以支持更大的最大序列长度。
为了进一步优化性能，该库提供了预量化的 4-bit 模型（例如 unsloth/llama-2-7b-bnb-4bit），这可以将内存碎片减少约 500MB，并使模型下载速度提高四倍。
对于 QLoRA 微调，用户可以通过 FastLanguageModel.get_peft_model 附加适配器。当 lora_dropout 设置为 0 且 bias 设置为 'none' 时，库的优化效果最佳。