低成本复制 ChatGPT 训练流程仅需 168GB GPU 即可使用方法现已开

编译 屠敏

出品 CSDN(ID:CSDNnews)

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

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

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

在单台多 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 都可以满足。

底层优化

低成本的 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://finance.yahoo.com/news/chatgpt-on-track-to-surpass-100-million-users-faster-than-tiktok-or-instagram-ubs-214423357.html

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