本地微调AI大模型神器！

计算机视觉深度学习和自动驾驶 · 公众号 · · 2024-06-22 00:28

正文

作为AI研究者和开发者，我们经常面临的一个问题就是模型训练的速度慢，内存消耗大，尤其是在进行大型语言模型（LLMs）的微调时。

今天，我要给大家介绍一个神奇的工具——Unsloth.ai，它能让你的Llama、Mistral和Gemma等LLMs的训练速度提升2-5倍，同时内存消耗减少80%！

Unsloth.ai是什么？

Unsloth.ai ^[1] 是一个开源项目，它使用OpenAI的Triton语言编写，提供了一套性能卓越的内核和手动反向传播引擎。这个工具不仅能够在不损失精度的情况下提升训练速度，还能够在现有的硬件上运行，支持自2018年以来的NVIDIA GPU！

亮点功能

快速微调 ：Unsloth.ai提供了一系列的Colab笔记本，只需添加你的数据集并点击"Run All"，就可以得到一个微调速度提升2倍的模型。
内存使用降低 ：相比传统方法，Unsloth.ai能够显著减少内存使用，这意味着你可以在资源有限的环境下也能训练大型模型。
易于安装 ：支持Conda和Pip安装，安装过程简单明了。

性能基准测试

Unsloth.ai在多个性能基准测试中展现出了卓越的性能。例如，在使用A100 40GB GPU时，与Hugging Face的Flash Attention相比，Unsloth能够提供接近两倍的速度提升，并且内存节省超过47%。支持的模型在Colab上的性能提升和内存使用减少如下：

Llama-3 8b：2倍速度提升，60%内存减少
Gemma 7b：2.4倍速度提升，71%内存减少
Mistral 7b：2.2倍速度提升，73%内存减少
TinyLlama：3.9倍速度提升，82%内存减少
CodeLlama 34b：1.9倍速度提升，49%内存减少

安装教程

Conda安装

conda create --name unsloth_env python=3.10
conda activate unsloth_env
conda install pytorch-cuda=<12.1/11.8> pytorch cudatoolkit xformers -c pytorch -c nvidia -c xformers
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
pip install --no-deps trl peft accelerate bitsandbytes

Pip安装

pip install --upgrade --force-reinstall --no-cache-dir torch==2.1.0 triton \
  --index-url https://download.pytorch.org/whl/cu121
pip install "unsloth[cu118] @ git+https://github.com/unslothai/unsloth.git"

使用示例

以下是如何使用Unsloth.ai进行模型微调的一个简单示例：

from unsloth import FastLanguageModel
import torch
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset

# 加载数据集
url = "https://huggingface.co/datasets/laion/OIG/resolve/main/unified_chip2.jsonl"
dataset = load_dataset("json", data_files={"train" : url}, split="train")

# 微调模型
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/mistral-7b-bnb-4bit",
    max_seq_length = 2048,
    dtype = None,
    load_in_4bit = True,
)
trainer = SFTTrainer(
    model = model,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = 2048,
    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()