训练自定义

TRL 的设计考虑了模块化，以便用户能够有效地根据自己的需求自定义训练循环。下面是一些关于如何应用和测试不同技术的示例。

在多个 GPU/节点上训练

TRL 中的训练器使用 🤗 Accelerate 来启用跨多个 GPU 或节点的分布式训练。为此，首先通过运行以下命令创建一个 🤗 Accelerate 配置文件：

accelerate config

并根据您的多 GPU/多节点设置回答问题。然后，您可以通过运行以下命令启动分布式训练：

accelerate launch your_script.py

我们还在示例文件夹中提供了可用作模板的配置文件。要使用这些模板，只需在启动作业时传递配置文件的路径，例如：

accelerate launch --config_file=examples/accelerate_configs/multi_gpu.yaml --num_processes {NUM_GPUS} path_to_script.py --all_arguments_of_the_script

有关更多详细信息，请参阅示例页面。

使用 DeepSpeed 进行分布式训练

TRL 中的所有训练器都可以与 DeepSpeed ZeRO-{1,2,3} 一起在多个 GPU 上运行，以有效地对优化器状态、梯度和模型权重进行分片。为此，请运行以下命令：

accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero{1,2,3}.yaml --num_processes {NUM_GPUS} path_to_your_script.py --all_arguments_of_the_script

请注意，对于 ZeRO-3，需要进行少量调整才能通过 zero3_init_context_manager() 上下文管理器在正确的设备上初始化您的奖励模型。特别是，这需要避免 DeepSpeed 在固定数量的训练步骤后挂起。以下是来自 sentiment_tuning 示例中涉及的内容片段

ds_plugin = ppo_trainer.accelerator.state.deepspeed_plugin
if ds_plugin is not None and ds_plugin.is_zero3_init_enabled():
    with ds_plugin.zero3_init_context_manager(enable=False):
        sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)
else:
    sentiment_pipe = pipeline("sentiment-analysis", model="lvwerra/distilbert-imdb", device=device)

有关 DeepSpeed 插件的更多信息，请参阅 🤗 Accelerate 文档。

使用不同的优化器

默认情况下，PPOTrainer 创建一个 torch.optim.Adam 优化器。您可以创建一个不同的优化器并将其传递给 PPOTrainer。

import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)


# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)

为了进行内存高效的微调，您还可以从 bitsandbytes 传递 Adam8bit 优化器。

import torch
import bitsandbytes as bnb

from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=config.learning_rate)

# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)

使用LION优化器

您也可以使用来自 Google 的新的 LION 优化器，首先获取优化器定义的源代码此处，并将其复制以便您可以导入优化器。请确保在初始化优化器时仅考虑可训练参数，以实现更内存高效的训练。

optimizer = Lion(filter(lambda p: p.requires_grad, self.model.parameters()), lr=self.config.learning_rate)

...
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer)

我们建议您使用您将用于 Adam 的学习率除以 3，如此处所述。我们观察到使用此优化器时与经典 Adam 相比有所改进（请参阅此处的完整日志）。

添加学习率调度器

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

import torch
from transformers import GPT2Tokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 2. define config
ppo_config = {'batch_size': 1, 'learning_rate':1e-5}
config = PPOConfig(**ppo_config)


# 2. Create optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=config.learning_rate)
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)

# 3. initialize trainer
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, optimizer=optimizer, lr_scheduler=lr_scheduler)

通过共享层进行内存高效的微调

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

import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m')
ref_model = create_reference_model(model, num_shared_layers=6)
tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')

# 2. initialize trainer
ppo_config = {'batch_size': 1}
config = PPOConfig(**ppo_config)
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)

传递 8 位参考模型

由于 trl 在使用 from_pretrained 从 transformers 加载模型时支持所有关键字参数，因此您还可以利用 transformers 中的 load_in_8bit 来进行更内存高效的微调。

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

# 0. imports
# pip install bitsandbytes
import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m')
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained('bigscience/bloom-560m', device_map="auto", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')

# 2. initialize trainer
ppo_config = {'batch_size': 1}
config = PPOConfig(**ppo_config)
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer)

使用CUDA缓存优化器

在训练大型模型时，您应该更好地处理CUDA缓存，方法是迭代地清除它。为此，只需将 optimize_cuda_cache=True 传递给 PPOConfig。

config = PPOConfig(..., optimize_cuda_cache=True)

使用分数缩放/标准化/裁剪

如大型语言模型中 RLHF 的秘密第一部分：PPO 所述，我们支持分数（又名奖励）缩放/标准化/裁剪，以通过 PPOConfig 提高训练稳定性。

from trl import PPOConfig

ppo_config = {
    use_score_scaling=True,
    use_score_norm=True,
    score_clip=0.5,
}
config = PPOConfig(**ppo_config)

要运行 ppo.py，您可以使用以下命令

python examples/scripts/ppo.py --log_with wandb --use_score_scaling --use_score_norm --score_clip 0.5

< > GitHub上的更新