在消费级 GPU 上使用 TRL 和直接偏好优化 (DPO) 微调 SmolVLM

在本食谱中，我们将指导您使用 Transformer Reinforcement Learning (TRL) 库通过直接偏好优化 (DPO) 来微调一个 smol 🤏 视觉语言模型 (VLM)，以演示即使在使用消费级 GPU 的情况下，您也可以如何定制 VLM 以满足您的特定需求。

我们将使用偏好数据集微调 SmolVLM，以帮助模型与期望的输出对齐。SmolVLM 是一款高性能且内存高效的模型，使其成为此任务的理想选择。如果您是语言或视觉语言模型的偏好优化新手，请查看此博客以获得深入的介绍。

我们将使用的数据集是 HuggingFaceH4/rlaif-v_formatted，其中包含 prompt + image 对，以及每对的 chosen 和 rejected 答案。此微调过程的目标是使模型始终偏好数据集中的chosen 答案，从而减少幻觉。

此 notebook 已使用 NVIDIA L4 GPU 进行测试。

1. 安装依赖

让我们从安装微调所需的基本库开始！🚀

!pip install  -U -q transformers trl datasets bitsandbytes peft accelerate
# Tested with transformers==4.46.3, trl==0.12.2, datasets==3.2.0, bitsandbytes==0.45.0, peft==0.14.0, accelerate==1.2.0

!pip install -q flash-attn --no-build-isolation

使用您的 Hugging Face 帐户进行身份验证，以直接从此 notebook 保存和共享您的模型 🗝️。

from huggingface_hub import notebook_login

notebook_login()

2. 加载数据集 📁

我们将使用 HuggingFaceH4/rlaif-v_formatted 数据集，该数据集提供 prompt + image 对，以及每对的 chosen 和 rejected 答案。这种结构化格式非常适合使用直接偏好优化 (DPO) 训练模型。

该数据集已为此任务预先格式化。如果您正在使用自定义数据集，则需要将其预处理为相同的格式。

在此示例中，我们将使用数据集的子集来演示该过程。但是，在实际场景中，您应该使用完整的数据集以获得更好的性能。

from datasets import load_dataset

dataset_id = "HuggingFaceH4/rlaif-v_formatted"
train_dataset, test_dataset = load_dataset(dataset_id, split=["train[:6%]", "test[:1%]"])

我们将确保所有图像均为 RGB 格式

from PIL import Image


def ensure_rgb(example):
    # Convert the image to RGB if it's not already
    image = example["images"][0]
    if isinstance(image, Image.Image):
        if image.mode != "RGB":
            image = image.convert("RGB")
        example["images"] = [image]
    return example


# Apply the transformation to the dataset
train_dataset = train_dataset.map(ensure_rgb, num_proc=32)
test_dataset = test_dataset.map(ensure_rgb, num_proc=32)

让我们探索数据集中的一个示例，以更好地了解其结构和我们正在处理的数据类型。

train_dataset[20]

>>> train_dataset[20]["images"][0]

3. 使用 TRL 微调模型

3.1 加载用于训练的量化模型 ⚙️

让我们首先使用 bitsandbytes 加载 SmolVLM-Instruct 模型的量化版本，并且让我们也加载处理器。我们将使用 SmolVLM-Instruct。

import torch
from transformers import Idefics3ForConditionalGeneration, AutoProcessor

model_id = "HuggingFaceTB/SmolVLM-Instruct"

from transformers import BitsAndBytesConfig

# BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
)

# Load model and tokenizer
model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config,
    _attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(model_id)

3.2 设置 QLoRA 和 DPOConfig 🚀

在此步骤中，我们将为我们的训练设置配置 QLoRA。QLoRA 是一种强大的微调技术，旨在减少内存占用，从而可以在有限的硬件上高效地微调大型模型。

QLoRA 通过为适配器权重引入量化，建立在传统的 LoRA（低秩自适应）之上。这种增强功能可显著降低内存使用量并加快训练速度，使其成为资源受限环境的理想选择。

>>> from peft import LoraConfig, get_peft_model

>>> # Configure LoRA
>>> peft_config = LoraConfig(
...     r=8,
...     lora_alpha=8,
...     lora_dropout=0.1,
...     target_modules=["down_proj", "o_proj", "k_proj", "q_proj", "gate_proj", "up_proj", "v_proj"],
...     use_dora=True,
...     init_lora_weights="gaussian",
... )

>>> # Apply PEFT model adaptation
>>> peft_model = get_peft_model(model, peft_config)

>>> # Print trainable parameters
>>> peft_model.print_trainable_parameters()

trainable params: 11,269,248 || all params: 2,257,542,128 || trainable%: 0.4992

接下来，我们将使用 DPOConfig 配置训练选项。

from trl import DPOConfig

training_args = DPOConfig(
    output_dir="smolvlm-instruct-trl-dpo-rlaif-v",
    bf16=True,
    gradient_checkpointing=True,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=32,
    num_train_epochs=5,
    dataset_num_proc=8,  # tokenization will use 8 processes
    dataloader_num_workers=8,  # data loading will use 8 workers
    logging_steps=10,
    report_to="tensorboard",
    push_to_hub=True,
    save_strategy="steps",
    save_steps=10,
    save_total_limit=1,
    eval_steps=10,  # Steps interval for evaluation
    eval_strategy="steps",
)

我们将使用 TRL 库中的 DPOTrainer 类，为直接偏好优化 (DPO) 定义训练参数。

DPO 使用标记的偏好数据来引导模型生成与偏好对齐的响应。TRL 的 DPOTrainer 将在训练之前对数据集进行标记化并将其保存到磁盘。此过程可能会占用大量磁盘空间，具体取决于用于训练的数据量。请据此计划，以避免存储空间不足。

此步骤可能需要一段时间，因此请随意放松并享受这个过程！ 😄

from trl import DPOTrainer

trainer = DPOTrainer(
    model=model,
    ref_model=None,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    peft_config=peft_config,
    tokenizer=processor,
)

是时候训练模型了！ 🎉

trainer.train()

让我们保存结果 💾

trainer.save_model(training_args.output_dir)

4. 测试微调后的模型 🔍

我们的视觉语言模型 (VLM) 已经过微调，现在是时候评估其性能了！在本节中，我们将使用 HuggingFaceH4/rlaif-v_formatted 数据集中的示例来测试该模型。让我们深入研究结果，并评估模型与首选响应的对齐程度！ 🚀

在开始之前，让我们清理 GPU 内存，以确保平稳和最佳的性能。 🧹

>>> import gc
>>> import time


>>> def clear_memory():
...     # Delete variables if they exist in the current global scope
...     if "inputs" in globals():
...         del globals()["inputs"]
...     if "model" in globals():
...         del globals()["model"]
...     if "processor" in globals():
...         del globals()["processor"]
...     if "trainer" in globals():
...         del globals()["trainer"]
...     if "peft_model" in globals():
...         del globals()["peft_model"]
...     if "bnb_config" in globals():
...         del globals()["bnb_config"]
...     time.sleep(2)

...     # Garbage collection and clearing CUDA memory
...     gc.collect()
...     time.sleep(2)
...     torch.cuda.empty_cache()
...     torch.cuda.synchronize()
...     time.sleep(2)
...     gc.collect()
...     time.sleep(2)

...     print(f"GPU allocated memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
...     print(f"GPU reserved memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")


>>> clear_memory()

GPU allocated memory: 1.64 GB
GPU reserved memory: 2.01 GB

我们将使用与之前相同的管道重新加载基础模型。

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

我们会将训练好的适配器附加到预训练模型。此适配器包含训练期间进行的微调调整，使基础模型能够利用新知识，同时保持其核心参数完整。通过集成适配器，我们增强了模型的功能，而无需更改其原始结构。

adapter_path = "sergiopaniego/smolvlm-instruct-trl-dpo-rlaif-v"
model.load_adapter(adapter_path)

让我们在未见过的样本上评估模型。

test_dataset[20]

>>> test_dataset[20]["images"][0]

让我们创建一个通用函数，我们可以使用不同的样本调用该函数，以简化测试过程。此函数将使我们能够有效地评估模型在多个示例上的性能，而无需为每个示例重写代码。通过使用此可重用函数，我们可以快速评估模型在各种输入下的性能。

def generate_text_from_sample(model, processor, sample, max_new_tokens=1024, device="cuda"):
    # Prepare the text input by applying the chat template
    text_input = processor.apply_chat_template(sample["prompt"], add_generation_prompt=True)

    image_inputs = []
    image = sample["images"][0]
    if image.mode != "RGB":
        image = image.convert("RGB")
    image_inputs.append([image])

    # Prepare the inputs for the model
    model_inputs = processor(
        text=text_input,
        images=image_inputs,
        return_tensors="pt",
    ).to(
        device
    )  # Move inputs to the specified device

    # Generate text with the model
    generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens)

    # Trim the generated ids to remove the input ids
    trimmed_generated_ids = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)]

    # Decode the output text
    output_text = processor.batch_decode(
        trimmed_generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    return output_text[0]  # Return the first decoded output text

现在，我们准备好调用该函数并评估模型了！ 🚀

output = generate_text_from_sample(model, processor, test_dataset[20])
output

该模型现在能够根据提供的图像和提示生成响应。对于此类任务，将模型的性能与基准进行比较以查看其改进程度以及与其他选项的比较情况非常有用。有关此比较的更多信息和详细信息，请查看此帖子。

💻 我开发了一个示例应用程序来测试该模型，您可以在此处找到它。

由于这里我们只使用数据集的子集运行示例训练，因此对于 Space，我使用了官方的 Hugging Face DPO 微调模型。您可以轻松地将其与另一个以预训练模型为特色的 Space 进行比较，该模型可在此处获得。

from IPython.display import IFrame

IFrame(src="https://sergiopaniego-smolvlm-trl-dpo-rlaif-v.hf.space", width=1000, height=800)

5. 继续学习之旅 🧑‍🎓️

通过以下资源扩展您对视觉语言模型和相关工具的知识

食谱中的多模态食谱： 探索多模态模型的实用食谱，包括检索增强生成 (RAG) 管道和微调。我们已经发布了使用 SFT 通过 TRL 微调 smol VLM 的食谱，该食谱与本指南完美互补——查看以获取更多详细信息。
TRL 社区教程： 探索丰富的教程集合，这些教程深入探讨了 TRL 及其在现实世界中的应用的复杂性。

您还可以回顾使用 Hugging Face 生态系统 (TRL) 微调视觉语言模型 (Qwen2-VL-7B) 中的继续学习之旅部分。

这些资源将帮助您加深在多模态学习方面的知识和专业知识。

< > 更新在 GitHub 上

开源 AI 食谱