Redian新闻
>
StackLLaMA: 用 RLHF 训练 LLaMA 的手把手教程

StackLLaMA: 用 RLHF 训练 LLaMA 的手把手教程

公众号新闻
来自:Hugging Face
进NLP群—>加入NLP交流群

如 ChatGPT,GPT-4,Claude 这样的语言模型之所以强大,是因为它们采用了 基于人类反馈的强化学习 (Reinforcement Learning from Human Feedback, RLHF) 来使之更符合我们的使用场景。

本博客旨在展示用 RLHF 训练一个 LLaMA 模型,以回答 Stack Exchange 上的问题。具体而言,包含以下几个方面:

  • 有监督的微调 (Supervised Fine-tuning,SFT)。
  • 奖励 / 偏好建模 (Reward / preference modeling,RM)。
  • 基于人类反馈的强化学习 (RLHF)。

摘自 InstructGPT 论文,Ouyang, Long, et al. “Training language models to follow instructions with human feedback.” arXiv preprint arXiv:2203.02155 (2022).

结合了上述方法,我们发布了 StackLLaMA 模型,该模型在 🤗 Hub 上开源 (访问链接查看 Meta 的原始 LLaMA ),整个 训练的流程 已经集成到了 Hugging Face TRL 库中 。你可以通过下面的 demo 来尝试该模型。

LLaMA 模型

在实践 RLHF 时,选取一个合适的模型很重要: RLHF 只是一个让模型满足我们交互形式的需求的微调过程 。所以我们选取了最近上线的 LLaMA 模型。LLaMA 模型是 Mata AI 最近推出的大语言模型。其参数量大小涵盖 7B 到 65B,以及训练在 1T 和 1.4T 的 token 上,这让其很实用。我们这里采用 7B 的模型。(请填写 Meta AI 的这份 表单 来下载模型)。

Stack Exchange 数据集

收集人类的反馈数据集是很复杂且昂贵的劳动。为了做到这个,并且还能保证模型的有效性,我们使用 StackExchange 数据集。该数据集涵盖了 StackExchange 平台上的问题和答案 (包含 StackOverflow 的编程等话题下的)。这很适合我们的实践,因为其包含了每个答案的赞和踩的数量。

我们按照 Askell et al. 2021 中的方法,给每个答案赋分:

score = log2 (1 + upvotes) rounded to the nearest integer, plus 1 if the questioner accepted the answer (we assign a score of −1 if the number of upvotes is negative).

对奖励模型,我们将看到每个问题总是需要两个答案对比。有些问题有很多答案,可以产生很多对,我们只取十个以限制每个问题的数据量。最后,我们把格式从 HTML 转化到 Markdown 以提高输出的可读性。你可以看到数据集和处理过程的 [笔记本]。(https://huggingface.co/datasets/lvwerra/stack-exchange-paired。)

高效训练策略

即使是最小 LLaMA 模型的训练,都需要大量内存。估算一下: 以 bf16 半精度,每个参数用 2 个字节 (以 fp32 精度四字节的标准),训练时需要 8 个字节 (例如 Adam 优化器,参见 Tramsformers 的 性能文档)。可见 7B 参数量的模型将用 (2+8)* 7B = 70 GB 的内存,并且还可能需要更多用于计算诸如注意力分数的中间值。所以很难在一张 80GB 显存的 A100 上训练。或许你可以使用一些技巧,比如用更高效的半精度训练的优化器来压缩内存,但溢出是迟早的。

另外的可能是 参数高效的微调(Parameter-Efficient Fine-Tuning, PEFT) 技术,比如 peft 库,它可以对使用 8-bit 加载的模型做 低秩优化(Low-Rank Adaptation,LoRA)。

线性层的低秩优化: 额外参数 (橙色) 被加在 Frozen 层 (蓝色),编码后的隐藏状态与 Frozen 层的隐藏状态叠加在一起。

以 8bit 加载模型会大幅降低内存占用,因为每个参数只要一字节 (比如 7B LLaMA 是 7GB 内存)。与直接训练原始模型不同,LoRA 在特定层 (一般是注意力层) 添加少量新参数,大幅降低了需要训练的参数。

此情此景,一个衡量标准是 1B 的参数在整个微调过程中占 ~1.2-1.4GB (和具体 batch size 及序列长度有关)。在参考的博客中具体讨论了,这使得低成本下微调较大参数规模的模型成为可能 (比如在一张 A100 上微调 50-60B 的参数)。

这些技术能让微调大模型的任务,在消费级设备和 Google Colab 上执行。这里提供一些值得关注的演示 demo: facebook/opt-6.7b (在 float16 精度下 13GB) 和 openai/whisper-large跑在 Google Colab (15GB 显存) 上。欲了解 peft 的使用,请参见 github 仓库 或者之前的 博客介绍: 在客户端训练 20B 参数量的模型。

现在我们能在一张 GPU 上微调很大的模型了,但训练还是会很慢。此时最简单的策略便是并行化: 把一个训练同时放到不同的 GPU 上,各 GPU 接受不同的 batch。这样我们可以并行执行前向传播和后向传播,通过增加 GPU 的数量实现并行能力提升。

我们可以选用 trainsformers.Traineraccelerate,因为它们都支持无代码变更进行数据并行化。只需注意调用 torchrun 或者 accelerate launch 脚本时的参数即可实现。比如以下就是在一个 8 显卡的机器上分别用 accelerate launchtorchrun的方法:

accelerate launch --multi_gpu --num_machines 1  --num_processes 8 my_accelerate_script.py
torchrun --nnodes 1  --nproc_per_node 8 my_torch_script.py

有监督的微调

在训练奖励模型和用 RL 之前,模型若是已经在我们感兴趣的方面表现好将会很有帮助。在我们的示例中,我们想要其能回答问题,而其他时候,我们可能它能听指令 (这时对指令执行的微调是理想的)。实现这个最简单的方法便是面向该语言任务,用该任务和领域的文本,继续训练。StackExchange 数据集 含 10M 的指令量,所以我们能用其子集很容易地训练。

在用 RLHF 之前的模型微调没有特别的,就是一般的面向语言任务的预训练模型微调。为了高效利用数据,我们采用了称之为 打包 的技术: 与 batch 中的每个样本均由单一文本组成,最后基于最长的文本来 padding (填充),我们把很多文本拼接起来,用 EOS token 来隔开,然后分割成一些 chunk (切块) 来做成 batch,避免 padding。

该方法大大提高了效率,因为模型输入的所有 token 都对 loss 有所训练,而非 padding 作为掩码被丢弃了。如果你没有足够数据,并且担心随意地分开 token 会失去上下文语义,你也可以用传统的数据加载器ConstantLengthDataset 解决了 打包技术,并且我们能在用 peft 加载模型后用 Trainer。首先,我们用 int8 加载模型,准备训练,然后加入 LoRA 微调器。

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_int8_training(model)

# add LoRA to model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, config)

我们根据相应的语言任务,对模型训练几千个 step (步),并保存模型。由于我们将会有其他微调模型的目的,我们将 LoRA 的微调器权重合并到原模型中。

声明: 因为 LLaMA 的许可证规定,我们只能发布微调器的权重,你需要填 Meta AI 的 表格 来获取模型,然后用这个 脚本 来转成 🤗 Transformers 格式。注意 🤗 Transformers 应该从源码安装,或者 v4.28 版。

现在我们已经微调好了模型,可以训练奖励模型了。

奖励模型和人类偏好

原则上,我们可以直接用人类标注来对模型做 RLHF 微调。然而,这将需要我们给人类发送一些样本,在每轮优化后计分。这是贵且慢的,因为收敛需要的训练样本量大,而人类阅读和标注的速度有限。

一个比直接反馈更好的策略是,在进入 RL 循环之前用人类标注集来训练一个奖励模型。奖励模型的目的是模拟人类对文本的打分。构建奖励模型有许多能用的策略: 最直接的便是预测标注 (比如根据好与坏,输出比分或者布尔值)。最佳实践是,预测结果的排序,即对每个 prompt (输入文本) 对应的两个结果 ,模型预测人类标注的比分哪个更高。

或者表示为 loss (损失) 函数:

其中 是模型对可能的标注 的预测分数。

在 StackExchange 数据集上,我们能得到两个答案的受欢迎程度。有了这个信息和上面的损失函数,我们就能自定义 loss 来改 transformers.Trainer 了。


class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"], attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

我们用数据集中的 100000 对,并在 50000 对上评估。在比较小的 batch size,为 4 下,我们用 LoRA 的  peft 微调器来训练 LLaMA 模型,在 BF16 精度下用 Adam 优化器。我们的 LoRA 设置是:

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

训练用 Weights & Biases 来记日志,并在 🤗 训练集群上,用 8 卡 A-100,要数小时,最后准确率为 **67%**。尽管看上去可能低了,但想想这个任务的难度。

如下文要细说的,训练结果将作为固定参数,以供下游使用。

基于人类反馈的强化学习

现在我们手头有了微调的语言模型和奖励模型,可以开始执行 RL 循环了: 这个过程大致分为三步

  1. 生成对 prompt (输入文本) 的反馈。
  2. 用奖励模型来对反馈评分。
  3. 对评分,进行一轮策略优化的强化学习。

在被 token 化并输入奖励模型前,提问和回答的 prompt 模版如下:

Question: <Query>
Answer: <Response>

在有监督训练 (SFT),奖励模型训练 (RM) 和 RLHF 的阶段都用此模版。

用 RL 训练语言模型出现的常见问题是,模型可能学会胡说八道以糊弄奖励模型,后者可能给高分。为了权衡,我们对奖励增加惩罚: 留一份没有训练的模型,如何比较两者输出的 KL 散度

其中 是奖励模型的结果, 是当前模型和对比模型的 KL 散度差。

再提一遍,我们用 peft 来实现内存高效的训练,其对 RLHF 阶段提供了优势。这里参考的模型和训练的模型用同一个基底,也就是有监督训练 (SFT) 的结果,它是用 8-bit 来加载,并且自始自终是固定的。我们仅用 PPO 方法优化最终模型的 LoRA 权重,同时全部共享一个基底模型。

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]
        
    # sample from the policy and generate responses
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    # Log stats to WandB
    ppo_trainer.log_stats(stats, batch, rewards)

我们用 🤗 集群,在 3x8 A100-80GB 的机器上训练了 20h,但一个差不多的结果很快 (大概,在 8 A100-80GB 上训练 20h)。所有的训练过程都在 Weight & Biases 上找到。

每个 batch 的奖励,对每步的训练,在  ~1000 步时模型的效果最好。

所以模型训好了能干啥嘞 ? 我们拭目以待 !

尽管我们不该太相信其结果,至少目前。但结果已经很好了,甚至附上了 Google 链接。我们来看看训练时的挑战。

挑战,不稳定和突破口

用 RL 训练 LLM (Large Language Models,大语言模型) 不总是一帆风顺的,你看到的本文也是经历无数实验,无数失败和无数调参的。即便如此,该模型也不能说变现完美。这儿,我们分享一些遇到的观察和问题。

奖励更高代表更好表现?

天呐,这个实验肯定表现很好 ! 看奖励的曲线多甜啊 !

在 RL 中,一般而言,奖励越高越好。在 RLHF 中,我们用了一个奖励模型,它不完美,所以留给了 PPO 算法捡漏的机会。这能导致奖励突然上升,然而当检查文本结果时,却充斥了字符 “```”,因为奖励模型对含有代码 stack exchange 的答案更信任。幸运的是,该问题碰到的很少,应该是采取的 KL 散度的惩罚项起到了作用。

KL 散度总是正的?

如我们前面所提到的,一个 KL 惩罚项被用来保证训练后的分布和原始分布接近。一般地 , KL 散度来度量两个分布的相似程度,并且总是正的。然而,在 trl 我们用了一个 KL 的近似,期望值和真的 KL 散度相同。

显然,当训练中一个 token 比原始模型概率低,这会导致 KL 散度为负,合适的取样和平均总能得到正的。但是一些采样的生成策略导致了不匀称的采样。比如,当生成被 padding 的序列 batch 时和当设置 EOS token 被压缩的最小长度是,模型会有很大/很小的概率到负 KL 散度的 token。同时 PPO 算法是面向奖励优化的,模型就会追逐负的惩罚,导致训练不稳定。

对生成和采样,你需要特别小心。我们建议一开始用最简单的方式,如何在逐渐复杂。

任然存在的问题

任然有很多问题我们不懂,比如下面,loss 间断地跳跃,导致之后的不稳定

一旦我们解决了这些问题,我们就会上传变化到 trl 上,以保证社区受益。

总结

在本博客,我们走过了 RLHF 训练的整个流程,从准备人类标注的数据集开始,调整语言模型到特定领域,训练奖励模型,并最终用 RL 训练一个模型。

通过使用 peft,任何人都能在一张 GPU 上跑我们的实验 ! 如果训练慢了,可以用数据并行化的方法,不需要改任何代码,或者用多张 GPU 并行提高训练速度。

对实际应用,这仅仅是第一步 ! 一旦你有了模型,你就要和其他模型比较优劣。这个可以用一个面向不同模型的排名生成做到,和我们训练奖励数据集类似。

一旦你加入了评估的步骤,好玩的就开始了: 你可以在原数据集上反复炼丹,也可以增加数据集或者对原数据集提纯。另外,你可以对奖励模型和生成试不同大小和结构的模型,这需要时间。

我们在积极提高 TRL 以保证 RLHF 的每一步都可见,并且十分激动能看到人们用它来构建的东西。如果你想有所贡献,欢迎看我们的 Github Issue。

引用

@misc {beeching2023stackllama,
    author = { Edward Beeching and
                     Younes Belkada and
                     Kashif Rasul and
                     Lewis Tunstall and
                     Leandro von Werra and
                     Nazneen Rajani and
                     Nathan Lambert
                   },
    title = { StackLLaMA: An RL Fine-tuned LLaMA Model for Stack Exchange Question and Answering },
    year = 2023,
    url = { https://huggingface.co/blog/stackllama },
    doi = { 10.57967/hf/0513 },
    publisher = { Hugging Face Blog }
}

感谢

我们感谢 Philipp Schmid 分享了他对文本生成绝妙的 demo, 我们的 demo 也是基于他的。我们也感谢 Omar Sanseviero 和 Louis Castricato 对我们博客的草稿提供宝贵详尽的反馈。

进NLP群—>加入NLP交流群

NLP、CV、推荐等方向需要有偿解决问题的快到这边来~




英文原文: https://hf.co/blog/stackllama

作者: Edward Beeching, Kashif Rasul, Younes Belkada, Lewis Tunstall, Leandro von Werra Nazneen Rajani, Nathan Lambert

译者: Vermillion-Qi(张奇), zhongdongy (阿东)

微信扫码关注该文公众号作者

戳这里提交新闻线索和高质量文章给我们。
相关阅读
最新RLHF拯救语言模型「胡说八道」!微调效果比ChatGPT更好,两名华人共同一作5分钟搞定春日编发,手把手教会!手把手教儿子写十年作文,儿子考上北大中文系后,他出了一本作文秘籍!24小时内、200美元复制RLHF过程,斯坦福开源「羊驼农场」模拟器女人是否会忘记第一次睡的男人?三个女人回答出奇一致 (不花一毛钱,手把手教你专柜千元换头术!就是这么牛!【英伦看法】手把手教你把商品进口至英国!10行代码媲美RLHF!谷歌DeepMind用游戏数据让大模型更像人类多GPU通信效率提升4倍,RLHF生成提升2.25倍!DeepSpeed ZeRO++重磅升级RLHF中的「RL」是必需的吗?有人用二进制交叉熵直接微调LLM,效果更好油烟机都不用开!肉这样蒸着吃比炒更香,手把手教你5种家常做法吃旧瓜建议收藏 | 手把手教你查询移民签证状态!手把手教你搭建属于自己的服务器Stability AI连扔两个王炸!首个开源RLHF模型登基,DeepFloyd IF像素级出图保姆级祛痘攻略!手把手教你终结痘痘Meta AI 重磅推出LIMA!媲美GPT-4、无需RLHF就能对齐!闲说跨年在墨西哥城 (10) -科约阿坎 (野狼谷)ChatGPT是男?是女?零门槛复现ChatGPT:预训练模型数据集直接用,包含完整RLHF流程,在线可体验爆!IKEA赔款2425万美金,手把手教你理赔;美国十大生活最贵城市,加州占三席!RLHF中的「RL」是必需的吗?斯坦福提出用二进制交叉熵直接微调LLMRLHF 实践中的框架使用与一些坑 (TRL, LMFlow)2023ESC-HFA | ARNI再添HFpEF新证!沙库巴曲缬沙坦PARAGLIDE-HF研究结果重磅公布放弃RLHF吧!无需手动训练模型价值观,达特茅斯学院华人领衔发布全新对齐算法:「AI社会」是最好的老师赏画、评画放弃RLHF吧!无需手动训练模型价值观,达特茅斯学院发布全新对齐算法没有更新疫苗记录收到停学通知!手把手教上传安省学龄儿童的黄色疫苗卡lāo dao?láo dao!ChatGPT破圈的「秘密武器」:详解RLHF如何影响人类社会!玩不起RLHF?港科大开源高效对齐算法RAFT「木筏」,GPT扩散模型都能用人手一个ChatGPT!微软DeepSpeed Chat震撼发布,一键RLHF训练千亿级大模型【Zoom活动】EducationUSA 线上讲座:互动游戏,手把手教你如何选择美国大学,4月8日上午10:00-11:00RLHF魔法的衍生研究方向无需RLHF就能对齐人类,性能比肩ChatGPT!华人团队提出袋熊Wombat模型
logo
联系我们隐私协议©2024 redian.news
Redian新闻
Redian.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Redian.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。