Transformer Reinforcement Learning X

trlX 是一个从头构建的分布式训练框架，专注于使用强化学习对大型语言模型进行微调，支持通过提供的奖励函数或带有奖励标签的数据集进行训练。

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

目前已实现以下 RL 算法：

📖 文档

🧀 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 .

示例

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

描述	链接
Simulacra (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=[['Question: 1 + 2 Answer:', '3'], ['Question: Solve this equation: ∀n>0, s=2, sum(n ** -s). Answer:', '(pi ** 2)/ 6']])

训练器提供对底层模型的封装

trainer.generate(**tokenizer('Q: Who rules the world? A:', 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 内存不足错误），可先尝试将以下超参数设置为最低值，再逐步增加：

# 每 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 头（特别适用于 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 fork 与 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] Message...

💡 提示：为了减少日志输出量，您可能会发现更改 trlX 使用的第三方库的日志级别很有帮助。例如，尝试在 trlX 脚本顶部添加 transformers.logging.set_verbosity_error() 以静默 transformers 库的详细消息（更多细节请参阅他们的日志文档）。

贡献

有关开发指南，请查看这些准则，并阅读我们的文档。

引用 trlX

@inproceedings{havrilla-etal-2023-trlx,
    title = "trl{X}: A Framework for Large Scale Reinforcement Learning from Human Feedback",
    author = "Havrilla, Alexander  and
      Zhuravinskyi, Maksym  and
      Phung, Duy  and
      Tiwari, Aman  and
      Tow, Jonathan  and
      Biderman, Stella  and
      Anthony, Quentin  and
      Castricato, Louis",
    booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing",
    month = dec,
    year = "2023",
    address = "Singapore",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2023.emnlp-main.530",
    doi = "10.18653/v1/2023.emnlp-main.530",
    pages = "8578--8595",
}

致谢

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