TRL训练自定义

TRL训练自定义

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

在多个 GPU/节点上进行训练

TRL 中的训练器使用 🤗 Accelerate 来实现跨多个 GPU 或节点的分布式训练。为此，首先通过运行创建一个 🤗 加速配置文件

accelerate config

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

accelerate launch your_script.py

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

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，需要做一个小调整，通过ZeRO-3 init context manager()上下文管理器在正确的设备上初始化奖励模型。特别是，这需要避免DeepSpeed在固定数量的训练步骤后挂起。下面是情绪调优示例中所涉及的内容的一个片段

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')
model_ref = 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, model_ref, tokenizer, optimizer=optimizer)

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

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')
model_ref = 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, model_ref, 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, model_ref, tokenizer, optimizer=optimizer)

添加学习率调度程序

您还可以通过添加学习率调度程序来进行训练！

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

# 1. load a pretrained model
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
model_ref = 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, model_ref, 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')
model_ref = 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, model_ref, tokenizer)

通过 8 位参考模型

由于使用 usingtrl加载模型时支持所有关键字参数，因此您还可以利用from进行更有效的内存微调。transformersfrom_pretrainedload_in_8bittransformers

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')
model_ref = 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, model_ref, tokenizer)

使用 CUDA 缓存优化器

在训练大型模型时，您应该通过迭代清除来更好地处理 CUDA 缓存。这样做，只需传递optimize_cuda_cache=True到PPOConfig：

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

使用分数缩放/归一化/剪辑

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

在这里插入代码片

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

from trl import PPOConfig

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

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

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