仅需1.68GB GPU，低成本复制ChatGPT训练流程，现已开源！

当昨日我们还在讨论从大厂以及个人创业的角度来看，复制一家 OpenAI 和一款强大的 ChatGPT 可行性究竟有几成之际，苦于 OpenAI 并未将 ChatGPT 开源出来，所以这趟水究竟有多深，众人并不知。

不过，2 月 14 日情人节这一天，来自加州大学伯克利分校的教授 James Demmel 和新加坡国立大学计算机系的校长青年教授尤洋及其背后的研究团队悄悄提供了一些答案，其率先呈现了一个开源的低成本 ChatGPT 等效实现流程，瞬间吸引无数 AI 爱好者的目光。

话不多说，简单来看，只需实现如下图所示的三步走，便能实现：

详细来看，我们将通过官方发布的详细公告一探究竟。

当然，等不及的小伙伴现在可以直接通过 GitHub 地址了解详情：https: //github.com/hpcaitech/ColossalAI

ChatGPT 技术分析

ChatGPT 的爆火，引得无数英雄竞折腰，那么 ChatGPT 为什么这么神奇？复制的难点是什么？

根据尤洋创立的潞晨科技（HPC-AI TECH）研究团队解析，ChatGPT 取得惊人成绩的重要特点是在训练过程中引入了人类反馈强化学习（RLHF），由此可以让这款 AI 聊天机器人更好地捕捉到人类的偏好。

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

1. 从 Prompt 库中取样，收集人类的反应，并使用这些数据来微调预先训练好的大型语言模型。

2. 从 Prompt 库中取样，使用大型语言模型生成多个响应，手动对这些响应进行排序，并训练一个奖励模型（RM）以适应人类的偏好。

3. 基于第 1 阶段的监督微调模型和第 2 阶段的奖励模型，使用强化学习算法进一步训练大语言模型。

在第三阶段，也就是 RLHF 训练的核心部分，OpenAI 采用强化学习中的近端策略优化（Proximal Policy Optimization，PPO）算法，引入奖励信号，使语言模型生成更更符合人类偏好的内容。

这也是开头伊始显示的那张图：

RLHF 三个阶段

ChatGPT 模型的复杂性实际上就是来源于强化学习的引入，这样会带来更多的模型调用。例如，使用基于 Actor-Critic（AC）结构的 PPO 算法，我们需要在训练过程中对 Actor 和 Critical 模型进行正向推理和反向传播，并在训练过程中对监督微调模型和奖励模型进行多次正向推理。关于作为 ChatGPT 基础的 InstructGPT 的论文，Actor 和监督微调模型都使用了有 1750 亿个参数的 GPT-3 系列模型，而 critical 和 奖励模型则使用了有 60 亿个参数的 GPT-3 系列模型。

在如此庞大的模型参数下，要启动原始的 ChatGPT 训练过程需要数千 GB 的 GPU 内存，这显然远远超出了单个 GPU 的能力，普通的数据并行技术也是不够的。

然而，即使引入张量并行和流水线并行来划分参数，仍然需要至少 64 个 80GB 的 A100 GPU 作为硬件基础。更糟糕的是，流水线由于其复杂性，以及 bubble 和调度的效率，不适合 AIGC 的生成性任务。第三阶段涉及复杂的强化学习和四个模型的训练过程，进一步给 ChatGPT 的代码复制带来困难和挑战。

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

那该怎么解决？

在这里，潞晨科技研究团队用上了自家研发的面向大模型时代的通用深度学习系统 Colossal-AI 带来了一个开源地成本部的 ChatGPT 等效实现流程。

想必不少 AI 从业者对 Colossal-AI 这款开源模型也有过一定的了解，它是开源的（https://github.com/hpcaitech/ColossalAI），其运用了高效多维自动并行、异构内存管理、大规模优化库、自适应任务调度等多项技术，实现高效快速部署 AI 大模型训练和推理，降低 AI 大模型应用成本。

Colossal-AI 背后的潞晨科技，其核心成员来自美国加州伯克利、斯坦福、清华、北大、新加坡国立、南洋理工等世界一流高校。这款模型于去年四月发布，通过一年不到的时间，截至目前，其获得了 8.9k 个 Star。

在实验过程中，研究人员以开源的方式复制了 ChatGPT 训练的基本过程，包括第一阶段的预训练，第二阶段的奖励模型训练，以及第三阶段的强化学习训练，这是流程中最复杂的阶段。

此外，Colossal-AI 通过使用 ZeRO、Gemini、LoRA、AutoChunk 内存管理等，大大降低了 ChatGPT 训练的 GPU 内存开销。它只需要一半的硬件资源就可以开始 1750 亿个参数的模型训练（从 64 张卡到 32 张卡），大大降低了 ChatGPT 应用的成本。

倘若在上述相同的硬件资源下，Colossal-AI 能够在更短的时间内进行训练，节省训练成本，加速产品迭代。

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

在单台多 GPU 服务器上，即使使用最高端的 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 训练过程。与 PyTorch 相比，在 14999 美元的 A100 80GB 上最多只能启动 7.8 亿个参数模型，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 训练代码。在此，以 GPT 为例，只需要一行代码就可以指定使用 Colossal-AI 作为系统策略来启动。

使用以下命令，开发者可以快速启动单 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

底层优化

低成本的 LoRA 微调

复制 ChatGPT 的实现过程通常依赖于 Colossal-AI。Colossal-AI 支持通过低秩矩阵微调（LoRA，Low-Rank Adaptation）方法进行高效微调。该方法假设大型语言模型是过度参数化的，微调过程中的参数变量是一个低秩矩阵，它可以分解为两个小矩阵的乘积：

由于大型语言模型的参数是固定的，在微调过程中只有调整矩阵的参数，从而减少训练参数的数量。在进行部署推理时，将矩阵的乘积加回到原始矩阵中，如，不影响推理延迟。

LoRA 结构，只训练 A、B

Zero+Gemini 来减少内存冗余

与传统的数据并行策略相比，Colossal-AI 使用零冗余优化器（ZeRO）来消除内存冗余，并在不影响计算粒度和通信效率的情况下，提高内存使用率。此外，为进一步提高了 ZeRO 的性能，Colossal-AI 还提出了基于 Chunk 的内存管理机制，它可以将连续的参数集按操作顺序存储在一个连续的、均匀分割的内存空间中，由此更能有效地利用网络带宽（PCI-e 和 GPU 之间），减少通信成本，并避免潜在的内存碎片。

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

写在最后

目前，该研究团队已经开源了完整的算法和软件设计来复制 ChatGPT 的实现过程：https://github.com/hpcaitech/ColossalAI。

不过，同样是基于成本考虑，他们表示，「对于这样一个巨型的人工智能模型，它需要更多的数据和计算资源来实际生效和部署。毕竟，用 1750 亿个参数训练一个 GPT-3 需要价值数百万美元的计算能力。因此，大型预训练模型长期以来只为少数大科技公司所拥有。」

因此，他们也希望能够以开源的方式，吸引更多的研究人员、机构共同参与进来，仅以上文中所复制 ChatGPT 训练流程的实践探索为起点，未来可以向大模型的时代做出努力。

更多详情内容可查阅官方公告：https://www.hpc-ai.tech/blog/colossal-ai-chatgpt

参考资料：

https://twitter.com/ArtificialAva/status/1623346998928723971

https://finance.yahoo.com/news/chatgpt-on-track-to-surpass-100-million-users-faster-than-tiktok-or-instagram-ubs-214423357.html

https://blogs.microsoft.com/blog/2023/02/07/reinventing-search-with-a-new-ai-powered-microsoft-bing-and-edge-your-copilot-for-the-web/

https://arxiv.org/abs/2106.09685

https://arxiv.org/pdf/2203.02155

https://openai.com/blog/chatgpt/

https://en.wikipedia.org/wiki/ChatGPT

https://www.benzinga.com/news/23/02/30850547/bill-gates-says-chatgpt-as-big-an-invention-as-the-internet-will-make-many-office-jobs