編譯 | 屠敏

出品 | CSDN（ID：CSDNnews）

當(dāng)昨日我們還在討論從大廠以及個(gè)人創(chuàng)業(yè)的角度來看，復(fù)制一家 OpenAI 和一款強(qiáng)大的 ChatGPT 可行性究竟有幾成之際，苦于 OpenAI 并未將 ChatGPT 開源出來，所以這趟水究竟有多深，眾人并不知。

不過，2 月 14 日情人節(jié)這一天，來自加州大學(xué)伯克利分校的教授 James Demmel 和新加坡國立大學(xué)計(jì)算機(jī)系的校長青年教授尤洋及其背后的研究團(tuán)隊(duì)悄悄提供了一些答案，其率先呈現(xiàn)了一個(gè)開源的低成本 ChatGPT 等效實(shí)現(xiàn)流程，瞬間吸引無數(shù) AI 愛好者的目光。

話不多說，簡(jiǎn)單來看，只需實(shí)現(xiàn)如下圖所示的三步走，便能實(shí)現(xiàn)：

詳細(xì)來看，我們將通過官方發(fā)布的詳細(xì)公告一探究竟。

當(dāng)然，等不及的小伙伴現(xiàn)在可以直接通過 GitHub 地址了解詳情：https: //github.com/hpcaitech/ColossalAI

ChatGPT 技術(shù)分析

ChatGPT 的爆火，引得無數(shù)英雄競(jìng)折腰，那么 ChatGPT 為什么這么神奇？復(fù)制的難點(diǎn)是什么？

根據(jù)尤洋創(chuàng)立的潞晨科技（HPC-AI TECH）研究團(tuán)隊(duì)解析，ChatGPT 取得驚人成績的重要特點(diǎn)是在訓(xùn)練過程中引入了人類反饋強(qiáng)化學(xué)習(xí)（RLHF），由此可以讓這款 AI 聊天機(jī)器人更好地捕捉到人類的偏好。

ChatGPT 的訓(xùn)練過程主要分為三個(gè)階段：

1. 從 Prompt 庫中取樣，收集人類的反應(yīng)，并使用這些數(shù)據(jù)來微調(diào)預(yù)先訓(xùn)練好的大型語言模型。

2. 從 Prompt 庫中取樣，使用大型語言模型生成多個(gè)響應(yīng)，手動(dòng)對(duì)這些響應(yīng)進(jìn)行排序，并訓(xùn)練一個(gè)獎(jiǎng)勵(lì)模型（RM）以適應(yīng)人類的偏好。

3. 基于第 1 階段的監(jiān)督微調(diào)模型和第 2 階段的獎(jiǎng)勵(lì)模型，使用強(qiáng)化學(xué)習(xí)算法進(jìn)一步訓(xùn)練大語言模型。

在第三階段，也就是 RLHF 訓(xùn)練的核心部分，OpenAI 采用強(qiáng)化學(xué)習(xí)中的近端策略優(yōu)化（Proximal Policy Optimization，PPO）算法，引入獎(jiǎng)勵(lì)信號(hào)，使語言模型生成更更符合人類偏好的內(nèi)容。

這也是開頭伊始顯示的那張圖：

RLHF 三個(gè)階段

ChatGPT 模型的復(fù)雜性實(shí)際上就是來源于強(qiáng)化學(xué)習(xí)的引入，這樣會(huì)帶來更多的模型調(diào)用。例如，使用基于 Actor-Critic（AC）結(jié)構(gòu)的 PPO 算法，我們需要在訓(xùn)練過程中對(duì) Actor 和 Critical 模型進(jìn)行正向推理和反向傳播，并在訓(xùn)練過程中對(duì)監(jiān)督微調(diào)模型和獎(jiǎng)勵(lì)模型進(jìn)行多次正向推理。關(guān)于作為 ChatGPT 基礎(chǔ)的 InstructGPT 的論文，Actor 和監(jiān)督微調(diào)模型都使用了有 1750 億個(gè)參數(shù)的 GPT-3 系列模型，而 critical 和 獎(jiǎng)勵(lì)模型則使用了有 60 億個(gè)參數(shù)的 GPT-3 系列模型。

在如此龐大的模型參數(shù)下，要啟動(dòng)原始的 ChatGPT 訓(xùn)練過程需要數(shù)千 GB 的 GPU 內(nèi)存，這顯然遠(yuǎn)遠(yuǎn)超出了單個(gè) GPU 的能力，普通的數(shù)據(jù)并行技術(shù)也是不夠的。

然而，即使引入張量并行和流水線并行來劃分參數(shù)，仍然需要至少 64 個(gè) 80GB 的 A100 GPU 作為硬件基礎(chǔ)。更糟糕的是，流水線由于其復(fù)雜性，以及 bubble 和調(diào)度的效率，不適合 AIGC 的生成性任務(wù)。第三階段涉及復(fù)雜的強(qiáng)化學(xué)習(xí)和四個(gè)模型的訓(xùn)練過程，進(jìn)一步給 ChatGPT 的代碼復(fù)制帶來困難和挑戰(zhàn)。

使用 Colossal-AI 低成本復(fù)制 ChatGPT 訓(xùn)練過程

那該怎么解決？

在這里，潞晨科技研究團(tuán)隊(duì)用上了自家研發(fā)的面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng) Colossal-AI 帶來了一個(gè)開源地成本部的 ChatGPT 等效實(shí)現(xiàn)流程。

想必不少 AI 從業(yè)者對(duì) Colossal-AI 這款開源模型也有過一定的了解，它是開源的（https://github.com/hpcaitech/ColossalAI），其運(yùn)用了高效多維自動(dòng)并行、異構(gòu)內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應(yīng)任務(wù)調(diào)度等多項(xiàng)技術(shù)，實(shí)現(xiàn)高效快速部署 AI 大模型訓(xùn)練和推理，降低 AI 大模型應(yīng)用成本。

Colossal-AI 背后的潞晨科技，其核心成員來自美國加州伯克利、斯坦福、清華、北大、新加坡國立、南洋理工等世界一流高校。這款模型于去年四月發(fā)布，通過一年不到的時(shí)間，截至目前，其獲得了 8.9k 個(gè) Star。

在實(shí)驗(yàn)過程中，研究人員以開源的方式復(fù)制了 ChatGPT 訓(xùn)練的基本過程，包括第一階段的預(yù)訓(xùn)練，第二階段的獎(jiǎng)勵(lì)模型訓(xùn)練，以及第三階段的強(qiáng)化學(xué)習(xí)訓(xùn)練，這是流程中最復(fù)雜的階段。

此外，Colossal-AI 通過使用 ZeRO、Gemini、LoRA、AutoChunk 內(nèi)存管理等，大大降低了 ChatGPT 訓(xùn)練的 GPU 內(nèi)存開銷。它只需要一半的硬件資源就可以開始 1750 億個(gè)參數(shù)的模型訓(xùn)練（從 64 張卡到 32 張卡），大大降低了 ChatGPT 應(yīng)用的成本。

倘若在上述相同的硬件資源下，Colossal-AI 能夠在更短的時(shí)間內(nèi)進(jìn)行訓(xùn)練，節(jié)省訓(xùn)練成本，加速產(chǎn)品迭代。

為了讓更多的開發(fā)者跑完 ChatGPT 訓(xùn)練過程，除了原有的 1750 億參數(shù)版本，Colossal-AI 還提供高效的單 GPU、獨(dú)立的 4/8GPU 的類 ChatGPT版本以減少硬件限制：

在單臺(tái)多 GPU 服務(wù)器上，即使使用最高端的 A100 80GB GPU，由于 ChatGPT 的復(fù)雜性和內(nèi)存碎片化，PyTorch 也只能啟動(dòng)基于 GPT-L（774M）等小模型的 ChatGPT。因此，用 PyTorch 的 DistributedDataParallel（DDP） 將多 GPU 并行擴(kuò)展到 4 或 8 個(gè) GPU，結(jié)果性能提升有限。

根據(jù)研究團(tuán)隊(duì)介紹，Colossal-AI 不僅在單 GPU 上有明顯的訓(xùn)練和推理速度優(yōu)勢(shì)，而且可以隨著并行規(guī)模的擴(kuò)大而進(jìn)一步提高，單服務(wù)器訓(xùn)練速度可達(dá) 7.73 倍，單 GPU 推理速度可達(dá) 1.42 倍，并且能夠繼續(xù)擴(kuò)大到大規(guī)模的平行度，大大降低 ChatGPT 復(fù)制的成本。

為了最大限度地降低訓(xùn)練成本和易用性，Colossal-AI 還提供了一個(gè)可以在單GPU 上試用的 ChatGPT 訓(xùn)練過程。與 PyTorch 相比，在 14999 美元的 A100 80GB 上最多只能啟動(dòng) 7.8 億個(gè)參數(shù)模型，Colossal-AI 將單個(gè) GPU 的容量提高了 10.3 倍，達(dá)到 80 億參數(shù)。對(duì)于基于 1.2 億個(gè)參數(shù)的小模型的 ChatGPT 訓(xùn)練，至少需要 1.62GB 的 GPU 內(nèi)存，任何單一消費(fèi)級(jí) GPU 都可以滿足。

此外，Colossal-AI 在致力于降低基于預(yù)訓(xùn)練的大型模型的微調(diào)任務(wù)的成本。例如，關(guān)于 OPT 模型的 ChatGPT 的微調(diào)任務(wù)，Colossal-AI 能夠?qū)蝹€(gè) GPU 上的微調(diào)模型的容量比 PyTorch 提高 3.7 倍，同時(shí)以足夠高的速度工作。

一行代碼即可啟動(dòng)

一切準(zhǔn)備就緒，只欠行動(dòng)。根據(jù)研究人員介紹，Colossal-AI 提供了開箱即用的 ChatGPT 訓(xùn)練代碼。在此，以 GPT 為例，只需要一行代碼就可以指定使用 Colossal-AI 作為系統(tǒng)策略來啟動(dòng)。

使用以下命令，開發(fā)者可以快速啟動(dòng)單 GPU 規(guī)模、單機(jī)多 GPU 規(guī)模、原始 1750 億參數(shù)規(guī)模版本的訓(xùn)練，并對(duì)各種性能指標(biāo)（包括最大GPU內(nèi)存使用率、吞吐量和TFLOPS）進(jìn)行評(píng)估。

# 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


底層優(yōu)化

低成本的 LoRA 微調(diào)

復(fù)制 ChatGPT 的實(shí)現(xiàn)過程通常依賴于 Colossal-AI。Colossal-AI 支持通過低秩矩陣微調(diào)（LoRA，Low-Rank Adaptation）方法進(jìn)行高效微調(diào)。該方法假設(shè)大型語言模型是過度參數(shù)化的，微調(diào)過程中的參數(shù)變量是一個(gè)低秩矩陣，它可以分解為兩個(gè)小矩陣的乘積：

由于大型語言模型的參數(shù)是固定的，在微調(diào)過程中只有調(diào)整矩陣的參數(shù)，從而減少訓(xùn)練參數(shù)的數(shù)量。在進(jìn)行部署推理時(shí)，將矩陣的乘積加回到原始矩陣中，如，不影響推理延遲。

LoRA 結(jié)構(gòu)，只訓(xùn)練 A、B

Zero Gemini 來減少內(nèi)存冗余

與傳統(tǒng)的數(shù)據(jù)并行策略相比，Colossal-AI 使用零冗余優(yōu)化器（ZeRO）來消除內(nèi)存冗余，并在不影響計(jì)算粒度和通信效率的情況下，提高內(nèi)存使用率。此外，為進(jìn)一步提高了 ZeRO 的性能，Colossal-AI 還提出了基于 Chunk 的內(nèi)存管理機(jī)制，它可以將連續(xù)的參數(shù)集按操作順序存儲(chǔ)在一個(gè)連續(xù)的、均勻分割的內(nèi)存空間中，由此更能有效地利用網(wǎng)絡(luò)帶寬（PCI-e 和 GPU 之間），減少通信成本，并避免潛在的內(nèi)存碎片。

此外，Colossal-AI 的異構(gòu)內(nèi)存管理器 Gemini 通過將優(yōu)化器狀態(tài)卸載到 CPU ，以此減少 GPU 內(nèi)存占用，允許同時(shí)使用 GPU 內(nèi)存和 CPU 內(nèi)存（包括 CPU DRAM 或 NVMe SSD 內(nèi)存）來訓(xùn)練超出單個(gè) GPU 內(nèi)存限制的大規(guī)模模型。

寫在最后

目前，該研究團(tuán)隊(duì)已經(jīng)開源了完整的算法和軟件設(shè)計(jì)來復(fù)制 ChatGPT 的實(shí)現(xiàn)過程：https://github.com/hpcaitech/ColossalAI。

不過，同樣是基于成本考慮，他們表示，「對(duì)于這樣一個(gè)巨型的人工智能模型，它需要更多的數(shù)據(jù)和計(jì)算資源來實(shí)際生效和部署。畢竟，用 1750 億個(gè)參數(shù)訓(xùn)練一個(gè) GPT-3 需要價(jià)值數(shù)百萬美元的計(jì)算能力。因此，大型預(yù)訓(xùn)練模型長期以來只為少數(shù)大科技公司所擁有。」

因此，他們也希望能夠以開源的方式，吸引更多的研究人員、機(jī)構(gòu)共同參與進(jìn)來，僅以上文中所復(fù)制 ChatGPT 訓(xùn)練流程的實(shí)踐探索為起點(diǎn)，未來可以向大模型的時(shí)代做出努力。

更多詳情內(nèi)容可查閱官方公告：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