PaLM + RLHF - Pytorch (进行中)
基于PaLM架构实现RLHF(基于人类反馈的强化学习)。也许我会添加类似RETRO的检索功能。
如果你有兴趣在公开领域中复制类似ChatGPT的项目,请考虑加入Laion
潜在的继任者:直接偏好优化 - 此仓库中的所有代码变为~二元交叉熵损失,少于5行代码。关于奖励模型和PPO的内容就到此为止了。
常见问题
- 这里包含用于推理的模型吗?
没有训练好的模型。这只是整个框架和总体地图。我们仍然需要数百万美元的计算资源和数据来航行到高维参数空间的正确位置。即便如此,你还需要像Stable Diffusion的Robin Rombach这样的专业水手来实际引导船只度过动荡的时刻,到达那个点。
社区
CarperAI 在ChatGPT发布前的几个月中,一直在为大语言模型开发一个RLHF框架。
Yannic Kilcher 也在为一个开源实现进行工作。
AI Coffeebreak w/ Letitia | Code Emporium | Code Emporium 第二部分
感谢
-
感谢Stability.ai的慷慨赞助,使得我们能够进行前沿的人工智能研究。
-
感谢🤗 Hugging Face和CarperAI撰写博客文章Illustrating Reinforcement Learning from Human Feedback (RLHF),以及前者提供的accelerate库。
-
感谢@kisseternity和@taynoel84进行代码审查和发现bug。
-
感谢Enrico将Pytorch 2.0中的Flash Attention整合进来。
安装
$ pip install palm-rlhf-pytorch
使用方法
首先训练 PaLM
,像任何其他自回归Transformer一样
import torch
from palm_rlhf_pytorch import PaLM
palm = PaLM(
num_tokens = 20000,
dim = 512,
depth = 12,
flash_attn = True # https://arxiv.org/abs/2205.14135
).cuda()
seq = torch.randint(0, 20000, (1, 2048)).cuda()
loss = palm(seq, return_loss = True)
loss.backward()
# 经过大量训练后,你现在可以生成序列了
generated = palm.generate(2048) # (1, 2048)
然后使用经过筛选的人类反馈训练你的奖励模型。在原始论文中,他们无法在不发生过拟合的情况下从预训练Transformer微调奖励模型,但我仍然提供了使用 LoRA
微调的选项,因为这仍然是一个开放的研究领域。
import torch
from palm_rlhf_pytorch import PaLM, RewardModel
palm = PaLM(
num_tokens = 20000,
dim = 512,
depth = 12,
causal = False
)
reward_model = RewardModel(
palm,
num_binned_output = 5 # 例如从1到5的评分
).cuda()
# 模拟数据
seq = torch.randint(0, 20000, (1, 1024)).cuda()
prompt_mask = torch.zeros(1, 1024).bool().cuda() # 序列中哪些部分是提示,哪些部分是响应
labels = torch.randint(0, 5, (1,)).cuda()
# 训练
loss = reward_model(seq, prompt_mask = prompt_mask, labels = labels)
loss.backward()
# 经过大量训练后
reward = reward_model(seq, prompt_mask = prompt_mask)
然后你将Transformer和奖励模型传递给 RLHFTrainer
import torch
from palm_rlhf_pytorch import PaLM, RewardModel, RLHFTrainer
# 加载预训练的palm
palm = PaLM(
num_tokens = 20000,
dim = 512,
depth = 12
).cuda()
palm.load('./path/to/pretrained/palm.pt')
# 加载预训练的奖励模型
reward_model = RewardModel(
palm,
num_binned_output = 5
).cuda()
reward_model.load('./path/to/pretrained/reward_model.pt')
# 准备好你的提示列表以进行强化学习
prompts = torch.randint(0, 256, (50000, 512)).cuda() # 5万个提示
# 将它们全部传递给训练器并开始训练
trainer = RLHFTrainer(
palm = palm,
reward_model = reward_model,
prompt_token_ids = prompts
)
trainer.train(num_episodes = 50000)
# 然后,如果成功的话...
# 生成例如10个样本,并使用奖励模型返回最佳样本
answer = trainer.generate(2048, prompt = prompts[0], num_samples = 10) # (<= 2048,)
待办事项
-
使用不同的lora克隆基础Transformer以便用于评论员
-
允许非LoRA基础的微调
-
重新做归一化以便能够实现掩码版本,不确定是否有人会使用每个token的奖励/值,但作为一种好的实践来实现
-
配备最好的注意力机制
-
添加Hugging Face的accelerate并测试wandb的工具
-
搜索文献以确定PPO的最新SOTA,假设RL领域仍在进步
-
使用预训练的情感网络作为奖励模型来测试系统
-
将PPO的内存写入memmapped numpy文件
-
让带有可变长度提示的采样工作,即使不需要,因为瓶颈是人类反馈
-
允许仅在预训练时微调最后N层,无论是在演员还是评论员中
-
从Sparrow中汲取一些学习点,参考Letitia的视频
-
使用django+htmx创建一个简单的Web界面用于收集人类反馈
-
考虑RLAIF
引用
@article{Stiennon2020LearningTS,
title = {Learning to summarize from human feedback},
author = {Nisan Stiennon and Long Ouyang and Jeff Wu and Daniel M. Ziegler and Ryan J. Lowe and Chelsea Voss and Alec Radford and Dario Amodei and Paul Christiano},
journal = {ArXiv},
year = {2020},
volume = {abs/2009.01325}
}
@inproceedings{Chowdhery2022PaLMSL,
title = {PaLM: Scaling Language Modeling with Pathways},
author = {Aakanksha Chowdhery and Sharan Narang and Jacob Devlin and Maarten Bosma and Gaurav Mishra and Adam Roberts and Paul Barham and Hyung Won Chung and Charles Sutton and Sebastian Gehrmann