Scalable toolkit for efficient model alignment
NVIDIA NeMo-Aligner
⚠️ 自 2025 年 5 月 15 日起,本仓库不再积极维护。我们建议迁移至 NeMo RL,这是一个可扩展的模块化训练后库,具有无缝的 Hugging Face 集成、Megatron Core 优化,并使用 Ray 作为调度主干。 ⚠️
最新动态
- 我们发布了 Nemotron-4-340B Base、Instruct、Reward 模型。其中 Instruct 和 Reward 变体是在 NeMo-Aligner 中训练的。有关奖励模型训练的更多详细信息,请参阅 Helpsteer2 论文。
- 我们很高兴地宣布,在我们的 RLHF pipeline 中支持使用 TensorRT-LLM 进行加速生成。有关更多信息,请参阅我们的 RLHF 文档。
- NeMo-Aligner 论文 现已在 arXiv 上发布!
简介
NeMo-Aligner 是一个用于高效模型对齐的可扩展工具包。该工具包支持最先进的模型对齐算法,如 SteerLM、DPO 和基于人类反馈的强化学习(RLHF)。这些算法使用户能够将语言模型调整得更安全、无害且更有帮助。用户可以对各种规模的模型执行端到端的模型对齐,并利用所有并行技术确保其模型对齐以高性能和资源高效的方式进行。有关更多技术细节,请参阅我们的 论文。
NeMo-Aligner 工具包基于 NeMo Framework 构建,支持通过张量并行、数据并行和管道并行在数千个 GPU 上进行可扩展训练,适用于所有对齐组件。此外,我们的检查点与 NeMo 生态系统交叉兼容,便于推理部署和进一步定制(https://github.com/NVIDIA/NeMo-Aligner)。
该工具包目前处于早期阶段。我们致力于改进此工具包,使开发人员能够更轻松地选择和使用不同的对齐算法,以构建安全、有帮助且可靠的模型。
核心特性
- SteerLM:属性条件下的监督微调(SFT),作为(用户可操控的)RLHF替代方案。
- Llama3-70B-SteerLM-Chat 通过 NeMo-Aligner 完成对齐。
- 相应的奖励模型 Llama3-70B-SteerLM-RM。
- 更多信息请参阅我们的 SteerLM 和 HelpSteer2 论文。
- 监督微调(Supervised Fine Tuning)
- 奖励模型训练(Reward Model Training)
- 基于人类反馈的强化学习(使用 PPO 算法)
- Llama3-70B-PPO-Chat 通过 NeMo-Aligner 结合 TRT-LLM 完成对齐。
- 基于人类反馈的强化学习(使用 REINFORCE 算法)
- Llama-3.1-Nemotron-70B-Instruct 通过 NeMo-Aligner 结合 TRT-LLM 完成对齐。
- 直接偏好优化(Direct Preference Optimization),如论文 Direct Preference Optimization: Your Language Model is Secretly a Reward Model 中所述
- Llama3-70B-DPO-Chat 通过 NeMo Aligner 完成对齐。
- 自博弈微调(Self-Play Fine-Tuning, SPIN),如论文 Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models 中所述
了解更多
最新发布
有关最新稳定版本,请参阅 发布页面。所有版本均随附预构建容器。每个版本的变更将记录在 CHANGELOG 中。
安装您自己的环境
要求
NeMo-Aligner 与 NeMo 工具包要求 相同,此外还需安装 PyTriton。
在 NeMo 容器内快速开始
NeMo 容器中已包含 NeMo Aligner。在安装了 NVIDIA GPU 和驱动程序的机器上,运行 NeMo 容器:
docker run --gpus all -it --rm --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/nemo:24.07
进入容器后,NeMo-Aligner 已完成安装,它与 NeMo 及其他工具一同位于 /opt/ 文件夹下。
安装 NeMo-Aligner
请按照 NeMo 工具包安装指南 中概述的相同步骤操作。安装 NeMo 后,执行以下附加命令:
pip install nemo-aligner
或者,如果你更倾向于安装最新的提交版本:
pip install .
Docker 容器
我们提供了官方的 NeMo-Aligner Dockerfile,该文件基于 NeMo、Megatron-LM 和 TransformerEngine 的稳定且经过测试的版本。此 Dockerfile 的主要目标是确保稳定性,尽管它可能并不总是反映这三个软件包的最新版本。您可以通过此处访问我们的 Dockerfile。
或者,您可以构建此处的 NeMo Dockerfile NeMo Dockerfile,并在末尾添加 RUN pip install nemo-aligner。
未来工作
- 我们将继续提高 PPO 学习阶段的稳定性。
- 改进 RLHF 的性能。
- 为拒绝采样添加 TRT-LLM 推理支持。
为 NeMo-Aligner 做贡献
我们欢迎社区贡献!请参考CONTRIBUTING.md了解相关指南。
在您的工作中引用 NeMo-Aligner
@misc{shen2024nemoaligner,
title={NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment},
author={Gerald Shen and Zhilin Wang and Olivier Delalleau and Jiaqi Zeng and Yi Dong and Daniel Egert and Shengyang Sun and Jimmy Zhang and Sahil Jain and Ali Taghibakhshi and Markel Sanz Ausin and Ashwath Aithal and Oleksii Kuchaiev},
year={2024},
eprint={2405.01481},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
许可协议
本工具包根据 Apache License, Version 2.0 许可。