用 Unsloth 和 🤗 TRL 让 LLM 微调提速 2 倍

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

在深入了解 Unsloth 之前，不妨先阅读我们的 QLoRA 博客文章，或者熟悉使用 🤗 PEFT 库进行 LLM 微调。

Unsloth——速度提升 2 倍，内存占用降低 40%，精度零损失

Unsloth 是一个用于更快进行 LLM 微调的轻量级库，与 HF Mirror 生态（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！训练后的模型与 HF Mirror 生态系统完全兼容，因此你可以将最终模型推送到 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()

可复现的 notebooks

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

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

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

CodeLlama 34b A100 colab 示例见此处

Zephyr DPO 复现 T4 colab 示例见此处