定义与性质

TRLX是一个用于通过人类反馈强化学习（Reinforcement Learning via Human Feedback，RLHF）来分布式训练语言模型的开源框架，由CarperAI开发。旨在为研究人员和开发人员提供便利，使其能够利用RLHF技术对大型语言模型进行微调，以更好地使模型与人类偏好对齐。

技术特点

  • 支持多种训练方式：支持使用提供的奖励函数或带奖励标签的数据集对模型进行训练，还支持近端策略优化（Proximal Policy Optimization，PPO）和隐式语言Q学习（Implicit Language Q-learning，ILQL）等强化学习算法。

  • 模型支持广泛：对于至多200亿参数的模型，提供基于🤗Hugging Face的Accelerate训练器，可微调如facebook/opt-6.7b、EleutherAI/gpt-neox-20b和google/flan-t5-xxl等基于因果和T5的语言模型。对于超过200亿参数的模型，提供NVIDIA NeMo支持的训练器，利用高效的并行技术进行扩展。

  • 分布式训练支持：支持多种分布式训练方式，包括分布式数据并行、模型分片、张量并行、顺序并行和管道并行等，能够充分利用计算资源，实现高效的大规模训练。

应用场景

  • 学术研究：为研究人员提供了一个便捷的平台，用于探索和实验新的强化学习算法与语言模型微调策略，推动自然语言处理领域的学术研究进展，例如研究如何更好地利用人类反馈来优化语言模型的生成效果等。

  • 工业应用：在智能客服、智能写作助手、内容推荐等实际应用场景中，可通过TRLX对语言模型进行微调，使其更好地满足特定业务需求，提高用户体验和业务效率。比如在智能客服中，让模型能够更准确、更符合人类偏好地回答用户问题。

安装与使用

  • 安装：可通过从GitHub克隆仓库，进入项目目录后使用pip install torch安装PyTorch，再使用pip install -e.安装TRLX及其依赖。

  • 使用：可以通过定义奖励函数、准备奖励标签数据集或提示-完成数据集等方式，结合配置好的超参数，调用trlx.train函数进行模型训练，最后可使用trainer.save_pretrained保存训练好的模型。