ChatGPT 低成本再现流程开源！任意单张消费级显卡可体验，记忆体需求低至 1.62GB

AI与大数据 1小时前 0

▼

首个开源 ChatGPT 低成本再现流程来了！预训练、奖励模型训练、强化学习训练，一次性打通。

明敏发自凹非寺，量子位 | 公众号QbitAI

最小demo训练流程仅需1.62GB显卡记忆体，随便一张消费级显卡就能满足。单卡模型容量最多提升10.3倍，相比原生PyTorch，单机训练速度最高可提升7.73倍，单卡推理速度提升1.42倍，仅需一行代码。

微调任务可最多提升单卡微调模型容量3.7倍，同时保持高速运行，同样仅需一行代码。ChatGPT红是红，再现也真的难，毕竟ChatGPT不开源，市面至今没有开源预训练权重、完全开源的低成本训练流程，且千亿等级大模型训练本身就是难题。

但ChatGPT军备赛已愈演愈烈，为了抓住趋势，如Google等都在打造竞品。快速再现ChatGPT是应趋势所需，开源加速方案Colossal-AI正是为此而来，且提供开源完整再现流程同时，也降低成本。

降低显卡记忆体消耗是关键

ChatGPT效果好，主要是训练过程引入人类回馈强化学习（RLHF），但也直接导致ChatGPT再现训练难度飙升。流程分为三阶段：

1. 监督微调：从Prompt库采样，收集人工回答，利用这些数据微调预训练大语言模型。

2. 奖励模型：从Prompt库采样，使用大语言模型产生多个回答，人工排序回答后，训练奖励模型（RM）模拟人类价值判断。

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

▲ RLHF三个阶段。

对ChatGPT训练而言，第三阶段是核心部分。OpenAI采用强化学习近端策略最佳化演算法（PPO），引入奖励讯号，使语言模型产生内容更符合人类评判标准。但引入强化学习，代表有更多模型调用。如使用基于Actor-Critic（AC）结构的PPO演算法，需训练时进行Actor、Critic两个模型前向推理和反向传播，以及监督微调模型、奖励模型多次前向推理。

ChatGPT基础的InstructGPT论文，Actor和监督微调模型都使用1,750亿参数GPT-3系列模型，Critic和奖励模型使用60亿参数GPT-3系列模型。大规模模型参数，代表想启动原始ChatGPT训练流程，需要数千GB显卡记忆体，单张GPU显然无法容纳，常见数据并行技术也不能搞定。

即便引入张量并行、流水并行划分参数，也需要至少64张80GB的A100为硬件基础，且流水并行并不适合AIGC产生任务，bubble和调度复杂会导致效率受限。

单张消费级显卡都能体验

Colossal-AI基于ZeRO、Gemini、LoRA、Chunk-based记忆体管理等方法，提出一系列单卡、单机多卡、大规模并行解决方案。对基于GPT-3系列模型的ChatGPT，Colossal-AI能用原本一半硬件资源启动1,750亿参数模型训练，从64卡降低到32卡。如果续用64卡，则将训练时间压缩到更短，节省训练成本、加速产品更新。

为了让更多开发者体验再现ChatGPT，除了1750亿参数版本，Colossal-AI还提供单卡、单机4 / 8卡类ChatGPT版本，以降低硬件限制。

单机多卡服务器，即便显卡规格提升为A100 80GB，由于ChatGPT复杂性和内存碎片，PyTorch最大也只能启动基于GPT-L（774M）这类小模型ChatGPT。用PyTorch原生DistributedDataParallel（DDP）进行多卡并行扩展至4卡或8卡，性能提升有限。Colossal-AI最高可提升单机训练速度7.73倍，单卡推理速度1.42倍，还可继续扩大规模并行。

为了尽可能降低训练成本和上手门槛，Colossal-AI还提供单张GPU即可尝试的ChatGPT训练流程。相比PyTorch在约10万元的A100 80GB最大仅能启动7.8亿参数模型，Colossal-AI将单卡容量提升10.3倍至80亿参数。对于基于1.2亿参数小模型的ChatGPT训练，最低仅需1.62GB显卡记忆体，任意单张消费级GPU即可满足。

此外，Colossal-AI也致力降低基于预训练大模型的微调任务成本。以ChatGPT可选的开源基础模型OPT为例，相比PyTorch，Colossal-AI可提升单卡微调模型容量3.7倍（原始计算量显着增大），同时保持高速运行。

一行代码快速上手

操作部分，如再现流程多步骤，基于Colossal-AI开源方案，都能一行代码快速上手。先看模型部分。尽管ChatGPT背后大语言模型GPT-3.5不开源，但GPT、OPT、BLOOM等主流开源模型都可取代。

Colossal-AI为Hugging Face社群这些模型提供开箱即用的ChatGPT再现代码，涵盖三阶段训练。以GPT为例，加一行代码指定使用Colossal-AI为系统策略即可快速启用。

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）

使用下列指令即可快速启动单卡、单机多卡、1750亿版训练，并测试各种性能指标（包括最大记忆体占用、吞吐率和TFLOPS等）：

# 使用单机单卡训练GPT2-S，使用最小的batch size，Colossal-AI Gemini CPU策略
torchrun –standalone –nproc_pero_node 1 benchmark_gpt_dummy.py –model s –strategy colossalai_gemini_cpu –experience_batch_size 1 –train_batch_size 1
# 使用单机4卡训练GPT2-XL，使用Colossal-AI Zero2策略
torchrun –standalone –nproc_per_node 4 benchmark_gpt_dummy.py –model xl –strategy colossalai_zero2
# 使用4机32卡训练GPT-3，使用Colossal-AI Gemini CPU策略
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_

背后原理如何？

核心方案还是Colossal-AI。毕竟诞生起就是大模型应用专用，可基于PyTorch高效快速部署AI大模型训练和推理，是这领域的明星，GitHub Star超过8千，并成功入选SC、AAAI、PPoPP、CVPR等国际AI与HPC顶级会议的官方教程。

目前Colossal-AI已成功帮某世界500强企业开发具增强线上搜寻引擎能力的类ChatGPT聊天机器人模型，之前也为Stable Diffusion、OPT、AlphaFold等先进模型提供多样高效大规模多维并行分布式解决方案。主创者为加州柏克莱大学教授James Demmel和新加坡国立大学校长尤洋。

▲ Colossal-AI与主要开源专案同期开源数据对比。

细节原理方面，LoRA、ZeRO＋Gemini是关键。

低成本微调的LoRA

微调部分，Colossal-AI支援使用低秩矩阵微调（LoRA）法。LoRA法认为大语言模型过参数化，微调参数改变量是低秩矩阵，可分解为两个更小矩阵的乘积：

微调时固定大模型参数，只调整低秩矩阵参数，减少训练参数量。微调后推理部署前，只需将参数加回原有矩阵即可：

不增加模型的推理延迟：

▲ LoRA示意图，仅需训练A、B。

减少记忆体冗余的ZeRO＋Gemini

Colossal-AI支援使用无冗余优化器（ZeRO）最佳化记忆体使用，可有效减少记忆体冗余，并比传统数据并行策略更不牺牲计算粒度和通讯效率，同时大幅提高记忆体使用效率。

为了提升ZeRO性能，Colossal-AI引入自动Chunk机制。透过将运算顺序连续一组参数存入同Chunk（Chunk是一段连续记忆体空间），确保每个Chunk大小相同，提高记忆体使用效率。使用Chunk法组织记忆体可确保PCI-e和GPU-GPU之间网络频宽有效利用，减少通讯次数，同时避免潜在记忆体碎片。

▲ Chunk机制。

此外，Colossal-AI异构记忆体空间管理器Gemini支援将优化器状态从GPU卸除到CPU ，以节省GPU记忆体。同时利用GPU、CPU记忆体（由CPU DRAM或NVMe SSD记忆体组成）突破单GPU记忆体墙限制，扩展可训练模型规模。

▲ 透过ZeRO＋Gemini提升硬件模型容量。

One More Thing

尽管这次开源包含再现ChatGPT的完整演算法流程和必要软件系统，但想要实际应用，至少还需数据、演算力等支援。参考开源大模型BLOOM、开源AI画画工具Stable Diffusion的经验，背后都需个人开发者、演算力、数据模型等可能合作方支援共建。

以前超过千名科学家联合发起、耗时一年多炼成的号称和GPT-3一样强大的语言模型BLOOM，还有AI绘图头号明星Stable Diffusion，也是由Stability AI、EleutherAI和LAION联合完成。再现ChatGPT也是如此，Colossal-AI正在发起活动。如果读者对这工作感兴趣或有好想法，可透过以下方式与他们联络：

GitHub发issue或提交PR。
加入Colossal-AI用户微信或Slack群。
点击“阅读原文”填写合作提案。
发送合作提案到信箱contact@hpcaitech.com。

Open source solution replicates ChatGPT training process! Ready to go with only 1.6GB GPU memory and gives you 7.73 times faster training!

（作者：明敏；本文由量子位 QbitAI授权转载；首图来源：shutterstock）

▼

特别声明本页内容仅供参考，版权终归原著者所有，若有侵权，请联系我们删除。