Unsloth 与 Hugging Face TRL 结合，让 LLM 微调最高提速 2.7 倍

因为 LLM 微调耗时过长而抓狂？在这篇文章中，我们介绍一个由社区开发的轻量级工具，让 LLM 微调变得超级快！

在深入了解 Unsloth 之前，阅读我们的 QLoRA 博客文章，或者熟悉使用 🤗 PEFT 库进行 LLM 微调，可能会有所帮助。

Unsloth - 速度提升 2 倍，内存占用减少 40%，精度零损失

Unsloth 是一个轻量级库，可实现更快的 LLM 微调，并且与 Hugging Face 生态系统（Hub、transformers、PEFT、TRL）完全兼容。该库由 Unsloth 团队（Daniel 和 Michael）以及开源社区积极开发。该库支持大多数 NVIDIA GPU——从 GTX 1070 一直到 H100——并且可与 TRL 库中的整套训练器套件（SFTTrainer、DPOTrainer、PPOTrainer）一起使用。在撰写本文时，Unsloth 支持 Llama（CodeLlama、Yi 等）和 Mistral 架构。

Unsloth 的工作原理是用优化后的操作覆盖部分建模代码。通过手动推导反向传播步骤，并将所有 Pytorch 模块重写为 Triton 内核，Unsloth 既能减少内存占用，又能加快微调速度。至关重要的是，与普通 QLoRA 相比，精度损失为 0%，因为优化后的代码中没有进行任何近似。

基准测试

Unsloth 在 Tesla T4 和 A100 Google Colab 实例上使用 4 个数据集进行了 59 次运行的基准测试。QLoRA 应用于所有线性层（注意力和 MLP），秩为 16，并启用了梯度检查点。通过与最新的 Transformers 版本 (4.36) 进行测试（如果你有 Pytorch 2.1.1，该版本原生集成了 SDPA），Unsloth 的速度最高可提升 2.7 倍，内存占用最高可减少 74%。我们还在免费的 Google Colab 实例（低 RAM、1 块 T4 GPU、Pytorch 2.1.0 CUDA 12.1）上测试了 Unsloth。所有 59 个 notebook 均已提供，以确保完全可复现，更多细节见 Unsloth 的基准测试详情此处

如何使用 Unsloth？

只需使用 FastLanguageModel.from_pretrained 加载你的模型！目前，Unsloth 支持 Llama 和 Mistral 类型架构（Yi、Deepseek、TinyLlama、Llamafied Qwen）。如果你希望支持其他架构，请提交 Github issue！此外，在最新的 Transformers main 分支上，你现在可以直接加载预量化的 4bit 模型！这使得模型下载速度提升 4 倍，并减少约 500MB 的内存碎片，从而让你能够容纳更大的批次！为方便起见，我们提供了一些预量化模型，包括 unsloth/llama-2-7b-bnb-4bit、unsloth/llama-2-13b-bnb-4bit、unsloth/mistral-7b-bnb-4bit 和 unsloth/codellama-34b-bnb-4bit。

你需要向 from_pretrained 提供你期望的最大序列长度。Unsloth 内部会执行 RoPE Scaling，因此会自动支持更大的最大序列长度。除此之外，API 与 transformers 的 from_pretrained 几乎相同，只是 FastLanguageModel.from_pretrained 还会为了方便而返回模型的 tokenizer。

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/mistral-7b-bnb-4bit", # Supports Llama, Mistral - replace this!
    max_seq_length = 2048, # Supports RoPE Scaling internally, so choose any!
    load_in_4bit = True,
)

模型加载完成后，使用 FastLanguageModel.get_peft_model 附加适配器，以便进行 QLoRA 微调。

# Do model patching and add fast LoRA weights
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    use_gradient_checkpointing = True,
)

附加适配器后，你可以在 HF 生态系统的任何类中直接使用该模型，例如 TRL 中的 SFTTrainer！

Unsloth + TRL 集成

要将 Unsloth 与 TRL 库一起使用，只需将 Unsloth 模型传入 SFTTrainer 或 DPOTrainer！训练后的模型与 Hugging Face 生态系统完全兼容，因此你可以将最终模型推送到 Hub，并开箱即用地使用 transformers 进行推理！

import torch

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

from unsloth import FastLanguageModel

max_seq_length = 2048 # Supports RoPE Scaling interally, so choose any!
# Get dataset
dataset = load_dataset("imdb", split="train")

# Load Llama model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/mistral-7b-bnb-4bit", # Supports Llama, Mistral - replace this!
    max_seq_length = max_seq_length,
    dtype = None,
    load_in_4bit = True,
)

# Do model patching and add fast LoRA weights
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 16,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    use_gradient_checkpointing = True,
    random_state = 3407,
    max_seq_length = max_seq_length,
)

trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    tokenizer = tokenizer,
    args = TrainingArguments(
      per_device_train_batch_size = 2,
      gradient_accumulation_steps = 4,
      warmup_steps = 10,
      max_steps = 60,
      fp16 = not torch.cuda.is_bf16_supported(),
      bf16 = torch.cuda.is_bf16_supported(),
      logging_steps = 1,
      output_dir = "outputs",
      optim = "adamw_8bit",
      seed = 3407,
  ),
)
trainer.train()

可复现的 notebook

我们在下面分享完全可复现的 notebook，供任何想在免费层级的 Google Colab 实例上使用 SFTTrainer 试用 Unsloth 的人使用。

Llama 7b 免费 Tesla T4 colab 示例见此处

Mistral 7b 免费 Tesla T4 colab 示例见此处

CodeLlama 34b A100 colab 示例见此处

Zephyr DPO 复现 T4 colab 示例见此处