为视觉语言多模态模型进行偏好优化
训练模型使得它能够理解并预测人类偏好是一项比较复杂的任务。诸如 SFT (Supervised finetuning) 的传统的方法一般都需要耗费较大成本,因为这些算法需要对数据打上特定的标签。而偏好优化 (Preference Optimization) 作为一种替代选项,通常可以简化这一过程,并产出更准确的结果。通过对候选回答的对比和排序,而不是赋予固定的标签,偏好优化使得模型能更高效地捕捉人类偏好中的细微差别。
偏好优化已经在大语言模型中广泛使用了,但现在,它也可以用在视觉语言模型 (VLM) 上。得益于 TRL 的开发,现在我们可以 使用 TRL 对 VLM 进行直接偏好优化 (Direct Preference Optimization)。本文将会介绍使用 TRL 和 DPO 对视觉语言模型进行训练的全过程。
偏好数据集
进行偏好优化,首先我们需要有一个能体现用户偏好的数据集。在双项选择的设定下,相应的数据一般包含一个提示词 (Prompt) 和两个候选回答,两个回答中一个被记为选中 (chosen),另一个被记为淘汰 (rejected)。模型将要去学习着给出选中的回答,而不是被淘汰的那个。下图就是一个例子:
图片来自 openbmb/RLAIF-V-Dataset 数据集
❔ Question: How many families?
- ❌ Rejected: The image does not provide any information about families.
- ✅ Chosen: The image shows a Union Organization table setup with 18,000 families.
需要注意的是,尽管选中的回答也不是完全正确的 (回答 18000 个家庭还是不对,应该是 18000000),但它也好于那个被淘汰的回答。
本文将使用 openbmb/RLAIF-V-Dataset 作为示例数据集,它包含了超过 83000 条标注的数据。可以通过下面代码查看一下数据集:
>>> from datasets import load_dataset
>>> dataset = load_dataset("openbmb/RLAIF-V-Dataset", split="train[:1%]")
>>> sample = dataset[1]
>>> sample["image"].show()
>>> sample["question"]
'how many families?'
>>> sample["rejected"]
'The image does not provide any information about families.'
>>> sample["chosen"]
'The image shows a Union Organization table setup with 18,000 families.'
我们将要训练的 VLM 模型需要文本和图像同时作为输入,所以这里的第一步还是要对数据集格式进行改造。一条数据应该被结构化成能模拟人机对话的形式。用户提供一个提示语,其中包含一张图片和一个问题,然后模型需要能够给出一个回答。我们用以下代码实现格式转换:
from datasets import features
from transformers import AutoProcessor
processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b", do_image_splitting=False)
def format(example):
# Prepare the input for the chat template
prompt = [
{
"role": "user",
"content": [{"type": "image"}, {"type": "text", "text": example["question"]}],
},
]
chosen = [
{
"role": "assistant",
"content": [{"type": "text", "text": example["chosen"]}],
},
]
rejected = [
{
"role": "assistant",
"content": [{"type": "text", "text": example["rejected"]}],
},
]
# Apply the chat template
prompt = processor.apply_chat_template(prompt, tokenize=False)
chosen = processor.apply_chat_template(chosen, tokenize=False)
rejected = processor.apply_chat_template(rejected, tokenize=False)
# Resize the image to ensure it fits within the maximum allowable
# size of the processor to prevent OOM errors.
max_size = processor.image_processor.size["longest_edge"]
example["image"].thumbnail((max_size, max_size))
return {"images": [example["image"]], "prompt": prompt, "chosen": chosen, "rejected": rejected}
# Apply the formatting function to the dataset,
# remove columns to end up with only "images", "prompt", "chosen", "rejected" columns
dataset = dataset.map(format, remove_columns=dataset.column_names)
# Make sure that the images are decoded, it prevents from storing bytes.
# More info here https://github.com/huggingface/blog/pull/2148#discussion_r1667400478
f = dataset.features
f["images"] = features.Sequence(features.Image(decode=True)) # to avoid bytes
dataset = dataset.cast(f)
完成了格式转换,我们来看看第一条数据:
>>> dataset[1]
{'images': [<PIL.JpegImagePlugin.JpegImageFile image mode=L size=980x812 at 0x154505570>],
'prompt': 'User:<image>how many families?<end_of_utterance>\n',
'rejected': 'Assistant: The image does not provide any information about families.<end_of_utterance>\n',
'chosen': 'Assistant: The image shows a Union Organization table setup with 18,000 families.<end_of_utterance>\n'}
OK!接下来准备好 GPU,训练马上开始。
训练
我们将使用 Idefics2-8b 作为我们的示例模型,但 TRL 里的 DPO 也是能用在像 Llava 1.5 和 PaliGemma 这样的模型上的 (可参考这篇文章: Finetuning Llava 1.5, PaliGemma and others)。不过训练之前,我们先检查一下我们的 GPU 显存是否够用:
训练需要多大的 GPU 显存?
一个 80GB VRAM 的 GPU 足够用来对 Idefics2-8b 进行 DPO 训练吗?我们可以先计算一下:
我们用 $ N $ 表示参数的数量,用 $ P $ 表示训练使用的精度。训练过程中,下列部分需要共同放入显存中:
- 要训练的模型: $ N \times P $
- 用以防止模型产生偏离的参考模型: 和要训练的模型一样大,所以也是 $ N \times P $
- 梯度: 我们对所有参数都进行训练,所以每个参数都有梯度: $ N \times P $
- 优化器的状态量: 我们使用 AdamW,一个参数会保存两个状态量,所以需要: $ 2 \times N \times P $
Idefics2-8b 有 80 亿 (8B) 参数,我们使用 float32
精度,每个参数占 4 个字节。所以总的显存需求是:
参数来源 | 计算公式 | 显存需求 |
---|---|---|
要训练的模型 | $ 8 \times 10^9 \times 4 $ | 32 GB |
参考模型 | $ 8 \times 10^9 \times 4 $ | 32 GB |
梯度 | $ 8 \times 10^9 \times 4 $ | 32 GB |
优化器状态量 | $ 2 \times 8 \times 10^9 \times 4 $ | 64 GB |
合计 | 160 GB |
这远超我们前面说的 80GB 显存了!幸运的是,我们可以使用量化、LoRA 等技术来大幅度地减少显存需求,让训练可以进行。接下来我们将介绍这些技术。
量化
量化会降低模型权重和激活值的精度,但也同时显著减少内存需求。将精度从 float32
改为 bfloat16
,会让每个参数需要的比特数从 4 比特减少到 2 比特。这一策略不仅能减少内存使用,还会显著加速训练,确保以最小代价保证足够高的性能。具体做法如下:
import torch
from transformers import AutoModelForVision2Seq
model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/idefics2-8b", torch_dtype=torch.bfloat16)
通过如下 bf16=True
的设置, bfloat16
也可以被用在优化器上:
from transformers import TrainingArguments
training_args = TrainingArguments(..., bf16=True)
LoRA
LoRA 对参数矩阵进行低秩分解; 在训练时,固定住原参数矩阵,仅训练分解出的两个矩阵。是一种大规模减少 LLM 训练参数的方法。LoRA 已被集成在了 PEFT 库里,使用非常方便:
from transformers import AutoModelForVision2Seq
+ from peft import get_peft_model, LoraConfig
model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/idefics2-8b")
+ peft_config = LoraConfig(target_modules="all-linear")
+ model = get_peft_model(model, peft_config)
PEFT 像是给原模型进行了一次封装 (代码中称为 adapter )。训练时,实际上是这个 adapter 在被训练,而原有的模型保持不动。我们现在算算 LoRA 帮我们减少了多少要训练的参数:
>>> model.print_trainable_parameters()
trainable params: 55,348,736 || all params: 8,458,116,848 || trainable%: 0.6543860411799315
它帮我们把要训练的参数从八十亿降到了五千五百万!差距真大!这将显著减少显存需求。
使用 bfloat16 和 LoRA 后的显存需求
现在我们来算算新的显存需求:
参数来源 | 计算公式 | 显存需求 |
---|---|---|
要训练的模型 | $ 8 \mathrm{G} \times 2 $ | 16 GB |
参考模型 | $ 8 \mathrm{G} \times 2 $ | 16 GB |
梯度 | $ 55 \mathrm{M} \times 2 $ | 0.1 GB |
优化器状态量 | $ 2 \times 55 \mathrm{M} \times 2 $ | 0.2 GB |
合计 | 32.3 GB |
现在我们仅需 32GB 的显存就可以训练我们的 Idefics2-8b 模型了。这合理多了,用 80GB 显存的 GPU 就可以训练了。
PEFT 文档 和 谷歌这篇关于 LoRA 和 QLoRA 文章 也提供了很多关于显存优化的帮助指南,读者感兴趣可以阅读。
训练时 batch size 怎么设定?
上述关于显存占用的计算还不算准确,因为实际训练时,激活值也需要占用显存。激活值是神经网络各层的输出。作为中间产物,它们的显存占用量取决于模型结构和训练时的 batch size。准确计算这些显存需求还是很困难的,我们一般依赖实验观察。
若想找到一个合适的 batch size ( per_device_train_batch_size
),你可以先随便选取一个你认为合适的数值 (比如 64) 然后试着开始训练。当然这大多数情况下会爆显存 (OOM)。如果这样,你可以减半 batch size,同时将 gradient_accumulation_steps
翻倍,以获得和原先 batch size 设定相同的效果。反复重复这一过程,最终当 OOM 不再出现时,你就可以训练了。我们的实验参数是: per_device_train_batch_size
设为 2, gradient_accumulation_steps
设为 32。
你还可以使用 gradient_checkpointing
来减少激活值所需的内存。这一技术在计算梯度时,会重新计算一遍前向过程,而不是在前向过程中保存用于计算梯度的中间结果。需要使用时,设置 gradient_checkpointing=True
即可。
完整训练代码
一切就绪,我们可以开始训练了。下面是我们的完整训练代码。除了上面提到的部分外,我们还设置了 dataset_num_proc
和 dataloader_num_workers
等参数,用于加速数据预处理。
# dpo_idefics2-8b.py
from datasets import features, load_dataset
from transformers import AutoModelForVision2Seq, AutoProcessor
import torch
from trl import DPOConfig, DPOTrainer
from peft import LoraConfig
def main():
# Load the model and processor
model = AutoModelForVision2Seq.from_pretrained("HuggingFaceM4/idefics2-8b", torch_dtype=torch.bfloat16)
processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b", do_image_splitting=False)
# Load the dataset
dataset = load_dataset("openbmb/RLAIF-V-Dataset", split="train")
def format(example):
# Prepare the input for the chat template
prompt = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": example["question"]}]}]
chosen = [{"role": "assistant", "content": [{"type": "text", "text": example["chosen"]}]}]
rejected = [{"role": "assistant", "content": [{"type": "text", "text": example["rejected"]}]}]
# Apply the chat template
prompt = processor.apply_chat_template(prompt, tokenize=False)
chosen = processor.apply_chat_template(chosen, tokenize=False)
rejected = processor.apply_chat_template(rejected, tokenize=False)
# Resize the image to ensure it fits within the maximum allowable
# size of the processor to prevent OOM errors.
max_size = processor.image_processor.size["longest_edge"]// 2
example["image"].thumbnail((max_size, max_size))
return {"images": [example["image"]], "prompt": prompt, "chosen": chosen, "rejected": rejected}
# Apply the formatting function to the dataset
dataset = dataset.map(format, remove_columns=dataset.column_names, num_proc=32)
# Make sure that the images are decoded, it prevents from storing bytes.
# More info here https://github.com/huggingface/blog/pull/2148#discussion_r1667400478
f = dataset.features
f["images"] = features.Sequence(features.Image(decode=True))
dataset = dataset.cast(f)
# Train the model
training_args = DPOConfig(
output_dir="idefics2-8b-dpo",
bf16=True,
gradient_checkpointing=True,
per_device_train_batch_size=2,
gradient_accumulation_steps=32,
num_train_epochs=1,
dataset_num_proc=32, # tokenization will use 32 processes
dataloader_num_workers=32, # data loading will use 32 workers
logging_steps=10,
)
trainer = DPOTrainer(
model,
ref_model=None, # not needed when using peft
args=training_args,
train_dataset=dataset,
tokenizer=processor,
peft_config=LoraConfig(target_modules="all-linear"),
)
trainer.train()
if __name__ == "__main__":
main()
启动脚本开始训练,接下来就等待结果吧
标签:模态,显存,训练,image,dataset,偏好,视觉,times,size From: https://www.cnblogs.com/huggingface/p/18306310