TRL 文档

训练自定义

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

训练自定义

TRL 的设计考虑了模块化,以便用户能够根据自己的需求高效地自定义训练循环。以下是一些关于如何应用和测试不同技术的示例。注意:尽管这些示例使用 DPOTrainer,但自定义方法适用于大多数(如果不是全部)训练器。

使用不同的优化器和调度器

默认情况下,DPOTrainer 创建一个 torch.optim.AdamW 优化器。你可以创建并定义一个不同的优化器,并将其传递给 DPOTrainer,如下所示:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch import optim
from trl import DPOConfig, DPOTrainer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
training_args = DPOConfig(output_dir="Qwen2.5-0.5B-DPO")

optimizer = optim.SGD(model.parameters(), lr=training_args.learning_rate)

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    optimizers=(optimizer, None),
)
trainer.train()

添加学习率调度器

你还可以通过添加学习率调度器来调整你的训练。

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch import optim
from trl import DPOConfig, DPOTrainer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
training_args = DPOConfig(output_dir="Qwen2.5-0.5B-DPO")

optimizer = optim.AdamW(model.parameters(), lr=training_args.learning_rate)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    optimizers=(optimizer, lr_scheduler),
)
trainer.train()

通过共享层实现内存高效的微调

另一个可以用来实现更高效内存微调的工具是在参考模型和你想要训练的模型之间共享层。

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import create_reference_model, DPOConfig, DPOTrainer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
ref_model = create_reference_model(model, num_shared_layers=6)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:1%]")
training_args = DPOConfig(output_dir="Qwen2.5-0.5B-DPO")

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()

传递 8 位参考模型

由于 `trl` 在使用 `from_pretrained` 从 `transformers` 加载模型时支持所有关键字参数,你也可以利用 `transformers` 的 `load_in_8bit` 功能来进行更高效的内存微调。

在此处阅读更多关于在 `transformers` 中加载 8 位模型的信息。

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import DPOConfig, DPOTrainer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config= quantization_config)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
training_args = DPOConfig(output_dir="Qwen2.5-0.5B-DPO")

trainer = DPOTrainer(
    model=model,
    ref_model=ref_model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)
trainer.train()

使用加速器缓存优化器

在训练大型模型时,你应该通过迭代地清除加速器缓存来更好地管理它。要做到这一点,只需将 `optimize_device_cache=True` 传递给 `DPOConfig`

training_args = DPOConfig(..., optimize_device_cache=True)
< > 在 GitHub 上更新