开源解决方案复制 ChatGPT 训练过程!仅需 1.6GB GPU 内存即可使用,训练速度提高 7.73 倍!

939 阅读11分钟

开源解决方案复制 ChatGPT 训练过程!仅需 1.6GB GPU 内存即可使用,训练速度提高 7.73 倍!

ChatGPT 在大西洋杂志 2022 年的“年度突破”中被 Dereck Thompson 认为是“生成式 AI 爆发”的一部分,它增加了该领域的显着兴趣和竞争。伴随着这一趋势,全球许多科技巨头如谷歌、微软都推出了 ChatGPT 风格的服务,加入与 ChatGPT 的竞争。

但是,由于OpenAI一直没有发布ChatGPT的代码,如何有效复制ChatGPT就成了摆在大家面前的一个巨大问题,急需一个开源的ChatGPT等价物。

好消息是,现在Colossal-AI作为最火的大型AI模型开源方案之一,率先呈现了一个开源的低成本ChatGPT等效实现流程,亮点包括:

  • 一个开源完整的基于PyTorch的ChatGPT等效实现流程,涵盖所有3个阶段,可以帮助您构建基于预训练模型的ChatGPT式服务。
  • 我们提供了一个迷你演示训练过程供用户试玩,它只需要 1.62GB 的 GPU 显存,并且可能在单个消费级 GPU 上实现, 在一个 GPU 上模型容量增长 高达10.3 倍。
  • 与原始 PyTorch 相比,单机训练过程可以快 7.73 倍,单 GPU 推理可以快 1.42 倍,这可以在 一行代码 中实现。
  • 在微调任务上,一行代码,就可以在保持足够 高的运行速度的情况下,将微调模型的容量 在单GPU上提升3.7倍。****
  • 我们提供多个版本 的单 GPU 规模、单节点多 GPU 规模和原始 1750 亿参数规模。我们还支持从 Hugging Face 导入 OPT、GPT-3、BLOOM 和许多其他预训练的大型模型到您的训练过程中。
  • 收敛性验证正在进行中,欢迎所有用户共同建设社区
  • GitHub 回购: https: //github.com/hpcaitech/ColossalAI

ChatGPT——AIGC引领的工业革命

新年伊始,你可能会想,人工智能领域最热门的技术是什么?那应该是 ChatGPT。

它就像一个六角战士,可以聊天、编写代码、修复错误、创建表单、发表论文、做作业、翻译,甚至是谷歌搜索引擎的有力竞争者。

ChatGPT自发布以来,风靡各行各业,仅5天就吸引了数百万用户的关注,上线2个月月活跃用户突破1亿。它已经成为历史上增长最快的应用程序,远远超过当今其他知名应用程序,例如推特用了 5 年达到 1 亿用户,Meta(Facebook)用了 4 年半,TikTok 用了 9 个月甚至 16 年。移动电话达到 1 亿用户。

比尔·盖茨称赞ChatGPT的意义不亚于互联网的发明,而微软CEO萨蒂亚·纳德拉则更直言不讳地说,“人工智能将从根本上改变每一个软件类别”。作为 OpenAI 百亿美元投资的主要投资方,微软已经迅速将 ChatGPT 集成到自家搜索引擎 Bing 和 Edge 浏览器中,并计划加入 Teams、Office 等办公套件。它的股票一夜之间飙升超过 800 亿美元。

两次 AI 演示后的谷歌股票与微软股票

由于担心在搜索方面被 ChatGPT 打败,谷歌迅速发布了竞品 Bard。然而,由于演示首次亮相的失败,它的股票立即蒸发了 1000 亿美元。

一夜之间,全球科技巨头正试图率先推出自己的类似 ChatGPT 的服务,以主导该领域的市场。

虽然 ChatGPT 已经发布几个月了,但既没有开源的预训练权重,也没有完整的低成本开源训练过程向公众开放。事实上,基于各种千亿参数模型的ChatGPT-style服务的全流程高效复制是很难实现的。最近临时上线了一堆ChatGPT替代品。但很难区分,因为它们是封闭源代码。

为什么 ChatGPT 如此神奇?复制的难点是什么?

ChatGPT技术分析

ChatGPT 取得惊人成绩的一个重要特点是在训练过程中引入了人类反馈强化学习(RLHF),以更好地捕捉人类的偏好。

ChatGPT 的训练过程分为三个主要阶段。

  1. 从 Prompt 库中抽样,收集其人类反应,并使用这些数据微调预训练的大型语言模型。
  2. 从 Prompt 库中采样,使用大型语言模型生成多个响应,手动对这些响应进行排序,并训练奖励模型 (RM) 以适应人类偏好。
  3. 基于第一阶段的监督微调模型和第二阶段的奖励模型,使用强化学习算法进一步训练大型语言模型。

在RLHF训练的核心部分stage 3,OpenAI采用强化​​学习中的近端策略优化(Proximal Policy Optimization,PPO)算法引入奖励信号,使语言模型生成更符合人类偏好的内容。

RLHF的三个阶段

ChatGPT 模型的复杂性源于强化学习的引入,会带来更多的模型调用。例如,使用基于Actor-Critic(AC)结构的PPO算法,我们需要在训练过程中对Actor和Critical模型进行前向推理和反向传播,以及有监督微调模型和奖励模型的多重前向推理。关于InstructGPT的论文,它是ChatGPT的基础,Actor和supervised fine-tuning模型都使用1750亿参数的GPT-3系列模型,而critical和reward模型使用60亿参数的GPT-3系列模型.

如此庞大的模型参数,要启动原始的ChatGPT训练过程, 需要数千GB的GPU显存,这显然远远超出了单个GPU的能力,普通的数据并行技术也不够用。然而,即使在分区参数上引入了张量并行和流水线并行,仍然至少 需要64个80GB A100 GPU作为硬件基础。更糟糕的是,流水线不适合AIGC的生成任务,因为它的泡沫带来的复杂性和效率和调度。Stage 3涉及四种模型复杂的强化学习和训练过程,进一步给ChatGPT的代码复现带来困难和挑战。

使用 Colossal-AI 低成本复制 ChatGPT 训练过程

Colossal-AI开源复制了ChatGPT训练的基本流程,包括stage 1预训练,stage 2奖励模型训练,stage 3强化学习训练,这是流程中最复杂的阶段。

此外,Colossal-AI通过使用ZeRO、Gemini、LoRA、AutoChunk内存管理等大大降低了ChatGPT训练的GPU内存开销,只需要一半的硬件资源就可以启动1750亿参数模型训练(从64张卡到32张卡) ,显着降低了 ChatGPT 样式应用程序的成本。使用上述相同的硬件资源,Colossal-AI 能够以更短的时间进行训练,从而节省训练成本并加速产品迭代。

为了让更多开发者跑通ChatGPT训练流程,除了原有的1750亿参数版本,Colossal-AI还提供了高效的单GPU、独立4/8-GPUs ChatGPT-like版本, 以减少硬件限制。

在单个多 GPUs 服务器上,即使使用最高端的 A100 80GB GPU,由于 ChatGPT 的复杂性和内存碎片,PyTorch 也只能基于 GPT-L(774M)等小型模型启动 ChatGPT。因此,使用 PyTorch 的 DistributedDataParallel (DDP) 将多 GPU 并行扩展到 4 或 8 个 GPU 会导致性能提升有限。

Colossal-AI 不仅在单 GPU 上的加速方面具有显着的训练和推理优势,而且可以随着并行性的扩大而进一步提高,单服务器训练速度最高可达 7.73 倍,单 GPU 推理速度最高可达 1.42 倍,并且 能够继续扩展到大规模并行,显着降低 ChatGPT 复制的成本

为了最大限度地降低培训成本和易用性,Colossal-AI 还提供了可以在单个 GPU 上试用的 ChatGPT 培训流程。与在 14999 美元的 A100 80GB 上最多只能启动 7.8 亿个参数模型的 PyTorch 相比,Colossal-AI 将单个 GPU 的容量提升了 10.3 倍,达到80 亿个参数。对于基于 1.2 亿参数的小模型的 ChatGPT 训练,至少需要 1.62GB 的 GPU 内存,任何单个消费级 GPU 都可以满足。

此外,Colossal-AI 一直在努力降低基于预训练大型模型的微调任务的成本。例如,在 OPT 模型上关于 ChatGPT 的微调任务,Colossal-AI 能够在单 GPU 上将微调模型的容量提高到PyTorch 的3.7 倍 ,同时工作速度足够高

一行代码启动

Colossal-AI 提供开箱即用的 ChatGPT 训练代码。用户会发现使用 Colossal-AI 以类似 ChatGPT 的方法训练来自 Hugging Face 社区的 GPT、OPT 和 BLOOM 等主流预训练模型很容易。例如,在 GPT 的情况下,只需要一行代码 来指定使用 Colossal-AI 作为系统策略来 kickstart。

from chatgpt.nn import GPTActor, GPTCritic, RewardModel
from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import ColossalAIStrategy

strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')

with strategy.model_init_context():
    actor = GPTActor().cuda()
    critic = GPTCritic().cuda()
    initial_model = deepcopy(actor).cuda()
    reward_model = RewardModel(deepcopy(critic.model)).cuda()

trainer = PPOTrainer(strategy, actor, critic, reward_model, initial_model, ...)
trainer.fit(prompts)

 

使用以下命令,可以快速启动单GPU规模、单机多GPU规模、原始1750亿参数规模版本的训练,并对各项性能指标(包括最大GPU内存占用、吞吐量和 TFLOPS):

# Training GPT2-S using a single card, a minimum batch size, Colossal-AI Gemini CPU strategy
torchrun --standalone --nproc_pero_node 1 benchmark_gpt_dummy.py --model s --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1
# Training GPT2-XL with a 4-GPU machine, Colossal-AI Zero2 strategy
torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2
# Training GPT-3 with 4 8-GPU servers, Colossal-AI Gemini CPU strategy
torchrun --nnodes 4 --nproc_per_node 8 \
 --rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR \
 benchmark_gpt_dummy.py --model 175b --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1

底层优化

核心系统:Colossal-AI

复制ChatGPT实现过程通常依赖大AI模型时代的深度学习系统Colossal-AI,支持基于PyTorch的大AI模型高效快速训练和推理,降低大AI模型部署成本。自开源以来,Colossal-AI 多次登上 GitHub Trending 榜首,获得超过 8000 颗 Star。被SC、AAAI、PPoPP、CVPR等顶级AI和HPC会议选为官方教程。除了上述优化,Colossal-AI还为大AI模型的大规模并行分布式训练提供了多种高效解决方案,并在Stable Diffusion、OPT、AlphaFold等前沿模型中展现了优越性。

Colossal-AI与目前同期主要开源项目对比

Colossal-AI 由美国加州大学伯克利分校特聘教授 James Demmel 先生和新加坡国立大学校长青年教授 You Yang 先生领导。其解决方案已在自动驾驶、云计算、零售、医药、芯片等领域得到多家知名科技巨头的成功应用和好评。例如,Colossal-AI 已成功帮助一家世界 500 强企业开发了具有增强的在线搜索能力的类 ChatGPT 聊天机器人模型。

LoRA的低成本微调

Colossal-AI 支持通过低秩适应 (LoRA) 方法进行高效微调。该方法假设大型语言模型是过度参数化的,微调过程中参数的变化可以表示为一个低秩矩阵,可以分解为两个较小矩阵的乘积

在大语言模型参数固定的情况下,fine-tuning时只调整分解矩阵的参数;因此,这种方法显着减少了训练参数的数量。为了进行推理部署,将矩阵乘积添加回原始矩阵,如,而不影响推理延迟。

LoRA结构,我们只训练A和B。

Zero+Gemini 减少内存冗余

与传统的数据并行性相比,Colossal-AI 使用零冗余优化器 (ZeRO) 来消除内存冗余并提高内存使用率,而不会影响计算粒度和通信速度。Colossal-AI还提出了基于chunk的内存管理,进一步提升了ZeRO的性能。基于块的内存管理将连续的参数集按操作顺序存储在称为块的连续、均匀分区的内存空间中,以更有效地使用网络带宽(在 PCI-e 和 GPU 之间)、降低通信成本并避免潜在的内存碎片化。

块策略

此外,Colossal-AI 的异构内存管理器 Gemini 通过将优化器状态卸载到 CPU 来减少 GPU 内存占用,允许同时使用 GPU 内存和 CPU 内存(包括 CPU DRAM 或 NVMe SSD 内存)来训练内存以外的更大规模模型单个 GPU 的限制。

ZeRO + Gemini 使用相同硬件增加模型容量