Transformer强化学习X

trlX是一个从头开始设计的分布式训练框架，专注于使用强化学习对大型语言模型进行微调，既可以利用提供的奖励函数，也可以使用带有奖励标签的数据集。

对于🤗 Hugging Face模型的训练支持由Accelerate后端的训练器提供，允许用户微调高达200亿参数的因果和T5类语言模型，例如facebook/opt-6.7b、EleutherAI/gpt-neox-20b和google/flan-t5-xxl。对于超过200亿参数的模型，trlX提供了基于NVIDIA NeMo的训练器，这些训练器利用高效的并行化技术来实现有效扩展。

目前实现了以下RL算法：

算法	Accelerate训练器	NeMo训练器
近端策略优化(PPO)	✅	✅
隐式语言Q学习(ILQL)	✅	✅

📖 文档

🧀 CHEESE 使用我们的“人机交互”数据收集库为您的RL应用收集人类标注。

安装

git clone https://github.com/CarperAI/trlx.git
cd trlx
pip install torch --extra-index-url https://download.pytorch.org/whl/cu118
pip install -e .

示例

更多用法请参阅示例。您还可以尝试下面的Colab笔记本：

描述	链接
模拟物(GPT2, ILQL)
情感(GPT2, ILQL)

示例的最新运行结果可在我们的Weights & Biases上查看。

如何训练

您可以使用奖励函数或带有奖励标签的数据集来训练模型。

使用奖励函数

trainer = trlx.train('gpt2', reward_fn=lambda samples, **kwargs: [sample.count('cats') for sample in samples])

有关奖励模型训练，请参阅我们的autocrit库。

使用带有奖励标签的数据集

trainer = trlx.train('EleutherAI/gpt-j-6B', samples=['dolphins', 'geese'], rewards=[1.0, 100.0])

使用提示-完成数据集

trainer = trlx.train('gpt2', samples=[['问题：1 + 2 答案：', '3'], ['问题：解这个方程：∀n>0, s=2, sum(n ** -s). 答案：', '(pi ** 2)/ 6']])

训练器为其底层模型提供封装

trainer.generate(**tokenizer('问：谁统治世界？答：', return_tensors='pt'), do_sample=True)

配置超参数

from trlx.data.default_configs import default_ppo_config

config = default_ppo_config()
config.model.model_path = 'EleutherAI/gpt-neox-20b'
config.tokenizer.tokenizer_path = 'EleutherAI/gpt-neox-20b'
config.train.seq_length = 2048

trainer = trlx.train(config=config, reward_fn=lambda samples, **kwargs: [len(sample) for sample in samples])

为了减少内存使用（如果您遇到CUDA Out of Memory错误），请先尝试将以下超参数设置为最低值，然后逐步提高：

# 每个GPU的微批次大小
config.train.batch_size = 1
# 冻结所有Transformer层
config.model.num_layers_unfrozen = 0
# 最大样本长度，超过该长度的提示或样本将会被截断
config.train.seq_length = 128

# 采样的微批次大小（仅适用于PPO）
config.method.chunk_size = 1
# 使用额外的Q-head（仅适用于ILQL）
config.method.two_qs = False

将训练好的模型保存为Hugging Face预训练语言模型。（已准备好上传到Hub！）

trainer.save_pretrained('/path/to/output/folder/')

使用🤗 Accelerate启动分布式训练

accelerate config # 选择DeepSpeed选项
accelerate launch examples/simulacra.py

使用NeMo-Megatron启动分布式训练

请按照NeMo README中的设置说明操作。

python examples/nemo_ilql_sentiments.py

更多用法请参阅NeMo README。

使用Ray Tune启动超参数搜索

ray start --head --port=6379
python -m trlx.sweep --config configs/sweeps/ppo_sweep.yml --accelerate_config configs/accelerate/ddp.yaml --num_gpus 4 examples/ppo_sentiments.py

将您的trlX分支与trlX的main分支进行基准测试

python -m trlx.reference octocat/trlx-fork:fix-branch

日志记录

trlX 使用标准的 Python logging 库将训练信息记录到控制台。默认的日志记录器级别设置为 INFO，这意味着 INFO、WARNING、ERROR 和 CRITICAL 级别的消息都会打印到标准输出。

要直接更改日志级别，可以使用 verbosity 设置器。例如，要将日志级别设置为 WARNING，可以使用以下代码：

import trlx

trlx.logging.set_verbosity(trlx.logging.WARNING)

这将抑制 INFO 级别的消息，但仍会打印 WARNING、ERROR 和 CRITICAL 级别的消息。

你还可以通过设置 TRLX_VERBOSITY 环境变量来控制日志的详细程度，该变量可以取标准日志记录 级别名称 中的一个：

• CRITICAL (trlx.logging.CRITICAL)
• ERROR (trlx.logging.ERROR)
• WARNING (trlx.logging.WARNING)
• INFO (trlx.logging.INFO)
• DEBUG (trlx.logging.DEBUG)

export TRLX_VERBOSITY=WARNING

默认情况下，使用 tqdm 进度条来显示训练进度。你可以通过调用 trlx.logging.disable_progress_bar() 来禁用它们，或者调用 trlx.logging.enable_progress_bar() 来启用。

可以通过设置 trlx.logging.enable_explicit_format() 来以更详细的方式格式化日志消息。这会在每条日志中注入调用站点的信息，有助于调试。

[2023-01-01 05:00:00,000] [INFO] [ppo_orchestrator.py:63:make_experience] [RANK 0] 消息...

💡 小贴士：为了减少日志输出量，你可能会发现调整 trlX 使用的第三方库的日志级别很有帮助。例如，尝试在 trlX 脚本的顶部添加 transformers.logging.set_verbosity_error()，以静默来自 transformers 库的冗长消息（更多详情请参阅其 日志文档）。

贡献

如需参与开发，请查看这些 指南，并阅读我们的 文档。

引用 trlX

@inproceedings{havrilla-etal-2023-trlx,
    title = "trl{X}: 用于大规模人类反馈强化学习的框架",
    author = "Havrilla, Alexander 与  Zhuravinskyi, Maksym 与  Phung, Duy 与  Tiwari, Aman 与  Tow, Jonathan 与  Biderman, Stella 与  Anthony, Quentin 与  Castricato, Louis",
    booktitle = "2023年自然语言处理经验方法会议论文集",
    month = dec,
    year = "2023",
    address = "新加坡",
    publisher = "计算语言学协会",
    url = "https://aclanthology.org/2023.emnlp-main.530",
    doi = "10.18653/v1/2023.emnlp-main.530",
    pages = "8578--8595",
}

致谢

非常感谢 Leandro von Werra 对 trl 的贡献，该库最初启发了这个仓库。

trlx 快速上手指南

trlX 是一个专为大规模语言模型强化学习（RL）微调设计的分布式训练框架。它支持使用自定义奖励函数或带奖励标签的数据集，对高达 20B 参数的 Hugging Face 模型（如 OPT、GPT-NeoX、Flan-T5）进行微调，并支持通过 NVIDIA NeMo 扩展至更大规模模型。

环境准备

• 系统要求：Linux 操作系统，配备 NVIDIA GPU（推荐 CUDA 11.8+）。
• 前置依赖：
  • Python 3.8+
  • PyTorch (需匹配 CUDA 版本)
  • Git
• 网络建议：由于模型权重和依赖包较大，建议配置稳定的网络环境或使用国内镜像源加速 pip 下载。

安装步骤

1. 克隆仓库

   git clone https://github.com/CarperAI/trlx.git
   cd trlx
   

2. 安装 PyTorch 请根据您的 CUDA 版本选择对应的安装命令（以下为 CUDA 11.8 示例）：

   pip install torch --extra-index-url https://download.pytorch.org/whl/cu118
   

   (注：国内用户可使用清华或阿里镜像源加速 PyTorch 安装)

3. 安装 trlx

   pip install -e .
   

基本使用

trlX 的核心是通过 trlx.train 接口启动训练。您可以直接使用奖励函数或提供带奖励的数据集。

1. 使用奖励函数训练 (最简单示例)

以下示例使用 GPT-2 模型，定义一个简单的奖励函数（统计样本中 "cats" 出现的次数），并启动训练：

import trlx

trainer = trlx.train('gpt2', reward_fn=lambda samples, **kwargs: [sample.count('cats') for sample in samples])

2. 使用带奖励标签的数据集训练

如果您已经准备好了样本及其对应的奖励分数：

import trlx

trainer = trlx.train('EleutherAI/gpt-j-6B', samples=['dolphins', 'geese'], rewards=[1.0, 100.0])

3. 使用提示 - 完成 (Prompt-Completion) 数据集

适用于问答或特定任务微调：

import trlx

trainer = trlx.train('gpt2', samples=[['Question: 1 + 2 Answer:', '3'], ['Question: Solve this equation: ∀n>0, s=2, sum(n ** -s). Answer:', '(pi ** 2)/ 6']])

4. 生成文本与保存模型

训练完成后，可以使用 trainer 直接生成文本或将模型保存为 Hugging Face 格式：

# 生成文本
output = trainer.generate(**tokenizer('Q: Who rules the world? A:', return_tensors='pt'), do_sample=True)

# 保存模型到本地 (可直接上传至 Hugging Face Hub)
trainer.save_pretrained('/path/to/output/folder/')

5. 配置超参数 (可选)

若遇到显存不足 (OOM) 错误，可调整配置文件中的批大小和序列长度：

from trlx.data.default_configs import default_ppo_config

config = default_ppo_config()
config.model.model_path = 'EleutherAI/gpt-neox-20b'
config.tokenizer.tokenizer_path = 'EleutherAI/gpt-neox-20b'
config.train.seq_length = 2048

# 降低显存占用的关键设置
config.train.batch_size = 1          # 每个 GPU 的微批次大小
config.model.num_layers_unfrozen = 0 # 冻结所有 Transformer 层
config.train.seq_length = 128        # 缩短最大序列长度

trainer = trlx.train(config=config, reward_fn=lambda samples, **kwargs: [len(sample) for sample in samples])

6. 分布式训练

使用 🤗 Accelerate 启动多卡训练：

accelerate config # 首次运行需配置，建议选择 DeepSpeed 选项
accelerate launch examples/simulacra.py