在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

该论文介绍了一种名为 ReMax 的新算法,专为基于人类反馈的强化进修(RLHF)而设计。ReMax 在盘算效率(约减少 50% 的 GPU 内存和 2 倍的训练速度提升)和完成简易性(6 行代码)上超越了最常用的算法 PPO,且功能没有损失。论文链接::李子牛,许天,张雨舜,俞扬,孙若愚,罗智泉机构:香港中文大学(深圳),深圳市大数据研究院,南京大学,南栖仙策开源代码:,所有图片来自于论文。背景今年,以 ChatGPT 为首的大语言模型(Large Language Models, LLMs) 在各个方面大放光彩

该论文介绍了一种名为 ReMax 的新算法,专为基于人类反馈的强化进修(RLHF)而设计。ReMax 在盘算效率(约减少 50% 的 GPU 内存和 2 倍的训练速度提升)和完成简易性(6 行代码)上超越了最常用的算法 PPO,且功能没有损失。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

论文链接:https://arxiv.org/abs/2310.10505

作者:李子牛,许天,张雨舜,俞扬,孙若愚,罗智泉

机构:香港中文大学(深圳),深圳市大数据研究院,南京大学,南栖仙策

开源代码:https://github.com/liziniu/ReMax

如未额外说明,所有图片来自于论文。

背景

今年,以 ChatGPT 为首的大语言模型(Large Language Models, LLMs) 在各个方面大放光彩,由此引发了学术界和商业界对 GPU 等盘算资源的需求剧增。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                      左图来自 DALL・E3,右图来自 DALL・E3

比如监督训练地调优 (supervised fine-tuning, SFT) 一个 Llama2-7B 的模型,需要消耗 80GB 以上的内存。而这往往不够,为了和人类对齐(alignment),大语言模型还要经过 RLHF (reinforcement learning from human feedback) 的训练。RLHF 的 GPU 消耗往往是 SFT 的 2 倍以上,训练空儿更能达到 6 倍以上。

近日,美国政府宣布限制英伟达 GPU 产品 H100, H800等进入中国市场。这项条款无疑为中国发展大语言模型(LLMs) 和人工智能增添了很多阻力。减小 RLHF 的训练成本(GPU 消耗和训练空儿)对 LLMs 的发展非常重要。

动机

RLHF 包含三个阶段:

1. 监督式地调优(Supervised Fine-Tuning, SFT)。 

2. 从对比数据中进修嘉奖模型(reward model)。

3. 利用强化进修(RL)算法来最大化嘉奖。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                    图片来源自 InstructGPT 论文

我们发现 RLHF 的主要盘算开支来源于第三阶段(嘉奖最大化)。这一点可以从 DeepSpeed-Chat 的报告里看到,第三阶段的训练空儿是前两个阶段空儿总和的 4 倍以上。而且,根据我们的经验,第三阶段的 GPU 消耗是前两阶段的 2 倍以上。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                    图片来自 DeepSpeed-Chat 技术报告

目前 RLHF 第 3 阶段的主要盘算瓶颈是什么?

我们发现该阶段的盘算瓶颈主要来源用来目前使用的 RL 算法:PPO 算法。PPO 算法是用来解决普适 RL 问题的最流行的算法之一,有非常多成功的案例。我们在这里省略 PPO 的技术细节,着重介绍 PPO 的一个关键组件:价格模型  (The value model)。价格模型是一个需要被训练的神经网络,能够有效地预计给定计谋的预期长期回报。尽管价格模型为 PPO 带来了良好的功能,但它在 RLHF 义务中也引入了沉重的盘算开支。例如,为了更好地与人类偏好对齐,PPO 中的价格模型通常与 LLM 大小相似,这使存储需求翻了一番。此外,价格模型的训练需要存储其梯度、激活和优化器状态,这进一步增加了近 4 倍的 GPU 存储需求。总结来说,PPO 和它的价格模型(以及其训练相关部分)已成为 RLHF 嘉奖最大化阶段的主要盘算障碍。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                相比 PPO,ReMax 是轻量级算法

思路

是否有可能找到比 PPO 更适配 RLHF 的算法?

 我们得出的答案是肯定的。这是因为 PPO 和价格模型是为通用 RL 问题设计的,而不是针对像 RLHF 这样的特定问题(RLHF 只是 RL 问题中的一个子类)。有趣的是,我们发现 RLHF 具有三个在 PPO 中未使用的重要结构:

1. 快速模拟(fast simulation):  轨迹(即 LLM 中的整个响应)可以在很短的空儿内迅速执行(小于 1s),几乎没有空儿开支。

2. 确定性转移(deterministic transitions):上下文确定性依赖于过去的标记和当前生成的标记。

3. 轨迹级嘉奖(trajectory-level rewards):嘉奖模型只在响应完成时提供一个奖赏值。

通过这三个观察,我们不难发现 value model 在 RLHF 的问题中是 “冗余” 的。这是因为 value model 设计的初衷是为了随机环境下的样本效率和慢仿真环境的盘算效率。然而这在 RLHF 中是不需要的。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                        ReMax 是针对 RLHF 设计的算法,PPO 则是为通用 RL 设计的算法

方法

ReMax

ReMax 算法基于一个古老的计谋梯度算法 REINFORCE,REINFORCE 使用的计谋梯度预计器如下图所示:

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                         REINFORCE 梯度预计器

REINFORCE可以在盘算层面利用好RLHF义务的三个性质,因为REINFORCE直接利用一个响应的嘉奖来进行优化,不需要像一般的RL算法一样需要知道中间步骤的嘉奖和值函数。然而,由于计谋的随机性, REINFORCE梯度预计器存在高方差问题(在Richard Sutton的RL书里有指出),这一问题会影响模型训练的有效性,因此REINFORCE在RLHF义务中的效果较差,见下面两张图片。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                          REINFORCE 的盘算代价小,但功能差

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                 REINFORCE 的(随机)梯度值远远大于 ReMax

为解决这一问题,ReMax 使用贪婪生成的回答(greedy response)的嘉奖作为基准值(baseline value)来构建梯度预计器,具体公式如下:

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                             ReMax 梯度预计器

注意到,贪婪回复的嘉奖在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了可以看作为期望嘉奖在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了的好的近似。在理想情形下(在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了),对于随机变量在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了,因此我们能够期望预计器在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了具有更小的方差。

下图展示了 ReMax 的算法流程,红色方框中的是核心算法改变。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                               ReMax 算法流程

理论保证 

我们证明了 ReMax 使用的梯度预计器仍然是真实计谋梯度的一个无偏预计器。

详细理论介绍见论文。

算法优点

ReMax 的核心部分可以用 6 行代码来完成。相比之下,PPO 要额外引入重要性采样(importance sampling),广义优势预计(generalized advantage estimation,GAE),价格模型进修等额外模块。

ReMax 的超参数很少。相比之下,PPO 有额外的超参数,例如重要性采样剪切阈值(importance sampling clipping ratio)、GAE 系数、价格模型进修率,离计谋训练轮次(off-policy training epoch)等,这些超参数都需要花大量空儿去调优。

ReMax 能理论上节俭约 50% 内存。相比于 PPO,ReMax 成功移除了所有和价格模型相关的部件,大大减小了内存开支。通过盘算,我们发现相比于 PPO,ReMax 能节俭约 50% 内存。

效果

有效性

ReMax 可以像 PPO 一样有效地最大化嘉奖

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                        在 OPT-1.3B 上,ReMax 可以有效地最大化嘉奖

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                     在 OPT-1.3B 上,ReMax 的训练非常稳定

在 GPT-4 评估下(LIMA Test Questions),ReMax 得到的计谋比 SFT 和 PPO 会更好

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                    GPT4 打分显示 ReMax 得到的模型会更好

高效性

ReMax 能节俭近 50% 的 GPU 内存。ReMax 移除掉了价格模型和它的训练部分(梯度,优化器,激活值),从而极大节俭了 GPU 内存需求。考虑 Llama2-7B,PPO 无法在 8xA100-40GB 的机器上跑起来,但是 ReMax 可以。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                   在 Llama2-7B 上,ReMax 可以节俭近 50% 的 GPU 内存

ReMax 能加快 2 倍的训练速度。在每一轮中,ReMax 调用 2 次生成(generation),1 次反向传播(backpropagation);而 PPO 使用 1 次生成,2 次反向传播。对于大模型而言,生成会比反向传播的空儿小,从而 ReMax 可以完成理论上接近 2 倍的训练减速。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

通用性

除了 RLHF 义务,作为一个 RL 算法,ReMax 对于经典的 NLP 义务也适用。本文考虑了在 GPT-2 上进行一个电影评论续写的义务,这里嘉奖模型不是从对比数据进修的。实验观测到,ReMax 可以完成 2.2 倍的训练减速和 60% 的 GPU 内存节俭。

在RTX 4090被限制的时代下,让大模型使用RLHF更高效的方法来了

                                          在经典的 NLP 义务(文本续写)上,ReMax 相比 PPO 完成了 2.2 倍减速

总结

最后,我们从实验中简要总结了 ReMax 相对于 PPO 的主要优势。

更简单的完成:  ReMax 的核心部分 6 行代码即可完成。这与 PPO 中的众多复杂的代码构建块形成鲜明对比。

更少的内存开支:由于移除了价格模型及其全部训练组件,相比 PPO,ReMax 节俭了大约 50% 的 GPU 内存。

更少的超参数:   ReMax 成功移除了所有和价格模型训练相关的超参数,其中包括:GAE 系数、价格模型进修率、重要性采样时期、小批量(mini-batch)大小。这些超参数往往对问题敏感且难以调整。我们相信 ReMax 对 RLHF 研究者更加友好。

更快的训练速度:在 GPT2(137M)的实验中,我们观察到 ReMax 在真实运行空儿方面相比于 PPO 有 2.2 倍的减速。减速来自 ReMax 每次迭代中较少的盘算开支。通过我们的盘算,该减速优势在更大的模型上也能维持(假设在足够大的内存下 PPO 可以被成功部署)。

优异的功能:如前所示,ReMax在中等规模实验中与PPO完成了相当的功能,并且有时甚至超越它(可能是由于 ReMax 更容易找到合适的超参数)。我们推测这种良好的功能可以拓展到更大规模的模型中。

给TA打赏
共{{data.count}}人
人已打赏
工程

阿里达摩院颁布遥感AI大模型,让AI进一步下沉到田间地头

2023-10-20 15:13:00

工程

OpenAI终于Open一回:DALL-E 3论文公布、上线ChatGPT,作者一半是华人

2023-10-20 15:29:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
有新私信 私信列表
搜索