低成本復制 ChatGPT 訓練流程，僅需 1.68GB GPU 即可使用，方法現(xiàn)已開源！
編譯 | 屠敏
出品 | CSDN（ID：CSDNnews）
當昨日我們還在討論從大廠以及個人創(chuàng)業(yè)的角度來看，復制一家 OpenAI 和一款強大的 ChatGPT 可行性究竟有幾成之際，苦于 OpenAI 并未將 ChatGPT 開源出來，所以這趟水究竟有多深，眾人并不知。
不過，2 月 14 日情人節(jié)這一天，來自加州大學伯克利分校的教授 James Demmel 和新加坡國立大學計算機系的校長青年教授尤洋及其背后的研究團隊悄悄提供了一些答案，其率先呈現(xiàn)了一個開源的低成本 ChatGPT 等效實現(xiàn)流程，瞬間吸引無數(shù) AI 愛好者的目光。
話不多說，簡單來看，只需實現(xiàn)如下圖所示的三步走，便能實現(xiàn)：
詳細來看，我們將通過官方發(fā)布的詳細公告一探究竟。
當然，等不及的小伙伴現(xiàn)在可以直接通過 GitHub 地址了解詳情：https: //github.com/hpcaitech/ColossalAI
ChatGPT 技術(shù)分析
ChatGPT 的爆火，引得無數(shù)英雄競折腰，那么 ChatGPT 為什么這么神奇？復制的難點是什么？
根據(jù)尤洋創(chuàng)立的潞晨科技（HPC-AI TECH）研究團隊解析，ChatGPT 取得驚人成績的重要特點是在訓練過程中引入了人類反饋強化學習（RLHF），由此可以讓這款 AI 聊天機器人更好地捕捉到人類的偏好。
ChatGPT 的訓練過程主要分為三個階段：
-
從 Prompt 庫中取樣，收集人類的反應，并使用這些數(shù)據(jù)來微調(diào)預先訓練好的大型語言模型。
從 Prompt 庫中取樣，使用大型語言模型生成多個響應，手動對這些響應進行排序，并訓練一個獎勵模型（RM）以適應人類的偏好。
基于第 1 階段的監(jiān)督微調(diào)模型和第 2 階段的獎勵模型，使用強化學習算法進一步訓練大語言模型。
在第三階段，也就是 RLHF 訓練的核心部分，OpenAI 采用強化學習中的近端策略優(yōu)化（Proximal Policy Optimization，PPO）算法，引入獎勵信號，使語言模型生成更更符合人類偏好的內(nèi)容。
這也是開頭伊始顯示的那張圖：
RLHF 三個階段
ChatGPT 模型的復雜性實際上就是來源于強化學習的引入，這樣會帶來更多的模型調(diào)用。例如，使用基于 Actor-Critic（AC）結(jié)構(gòu)的 PPO 算法，我們需要在訓練過程中對 Actor 和 Critical 模型進行正向推理和反向傳播，并在訓練過程中對監(jiān)督微調(diào)模型和獎勵模型進行多次正向推理。關(guān)于作為 ChatGPT 基礎的 InstructGPT 的論文，Actor 和監(jiān)督微調(diào)模型都使用了有 1750 億個參數(shù)的 GPT-3 系列模型，而 critical 和 獎勵模型則使用了有 60 億個參數(shù)的 GPT-3 系列模型。
在如此龐大的模型參數(shù)下，要啟動原始的 ChatGPT 訓練過程需要數(shù)千 GB 的 GPU 內(nèi)存，這顯然遠遠超出了單個 GPU 的能力，普通的數(shù)據(jù)并行技術(shù)也是不夠的。
然而，即使引入張量并行和流水線并行來劃分參數(shù)，仍然需要至少 64 個 80GB 的 A100 GPU 作為硬件基礎。更糟糕的是，流水線由于其復雜性，以及 bubble 和調(diào)度的效率，不適合 AIGC 的生成性任務。第三階段涉及復雜的強化學習和四個模型的訓練過程，進一步給 ChatGPT 的代碼復制帶來困難和挑戰(zhàn)。
使用 Colossal-AI 低成本復制 ChatGPT 訓練過程
那該怎么解決？
在這里，潞晨科技研究團隊用上了自家研發(fā)的面向大模型時代的通用深度學習系統(tǒng) Colossal-AI 帶來了一個開源地成本部的 ChatGPT 等效實現(xiàn)流程。
想必不少 AI 從業(yè)者對 Colossal-AI 這款開源模型也有過一定的了解，它是開源的（https://github.com/hpcaitech/ColossalAI），其運用了高效多維自動并行、異構(gòu)內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應任務調(diào)度等多項技術(shù)，實現(xiàn)高效快速部署 AI 大模型訓練和推理，降低 AI 大模型應用成本。
Colossal-AI 背后的潞晨科技，其核心成員來自美國加州伯克利、斯坦福、清華、北大、新加坡國立、南洋理工等世界一流高校。這款模型于去年四月發(fā)布，通過一年不到的時間，截至目前，其獲得了 8.9k 個 Star。
在實驗過程中，研究人員以開源的方式復制了 ChatGPT 訓練的基本過程，包括第一階段的預訓練，第二階段的獎勵模型訓練，以及第三階段的強化學習訓練，這是流程中最復雜的階段。
此外，Colossal-AI 通過使用 ZeRO、Gemini、LoRA、AutoChunk 內(nèi)存管理等，大大降低了 ChatGPT 訓練的 GPU 內(nèi)存開銷。它只需要一半的硬件資源就可以開始 1750 億個參數(shù)的模型訓練（從 64 張卡到 32 張卡），大大降低了 ChatGPT 應用的成本。
倘若在上述相同的硬件資源下，Colossal-AI 能夠在更短的時間內(nèi)進行訓練，節(jié)省訓練成本，加速產(chǎn)品迭代。
為了讓更多的開發(fā)者跑完 ChatGPT 訓練過程，除了原有的 1750 億參數(shù)版本，Colossal-AI 還提供高效的單 GPU、獨立的 4/8GPU 的類 ChatGPT版本以減少硬件限制：
在單臺多 GPU 服務器上，即使使用最高端的 A100 80GB GPU，由于 ChatGPT 的復雜性和內(nèi)存碎片化，PyTorch 也只能啟動基于 GPT-L（774M）等小模型的 ChatGPT。因此，用 PyTorch 的 DistributedDataParallel（DDP） 將多 GPU 并行擴展到 4 或 8 個 GPU，結(jié)果性能提升有限。
根據(jù)研究團隊介紹，Colossal-AI 不僅在單 GPU 上有明顯的訓練和推理速度優(yōu)勢，而且可以隨著并行規(guī)模的擴大而進一步提高，單服務器訓練速度可達 7.73 倍，單 GPU 推理速度可達 1.42 倍，并且能夠繼續(xù)擴大到大規(guī)模的平行度，大大降低 ChatGPT 復制的成本。
為了最大限度地降低訓練成本和易用性，Colossal-AI 還提供了一個可以在單GPU 上試用的 ChatGPT 訓練過程。與 PyTorch 相比，在 14999 美元的 A100 80GB 上最多只能啟動 7.8 億個參數(shù)模型，Colossal-AI 將單個 GPU 的容量提高了 10.3 倍，達到 80 億參數(shù)。對于基于 1.2 億個參數(shù)的小模型的 ChatGPT 訓練，至少需要 1.62GB 的 GPU 內(nèi)存，任何單一消費級 GPU 都可以滿足。
此外，Colossal-AI 在致力于降低基于預訓練的大型模型的微調(diào)任務的成本。例如，關(guān)于 OPT 模型的 ChatGPT 的微調(diào)任務，Colossal-AI 能夠?qū)蝹€ GPU 上的微調(diào)模型的容量比 PyTorch 提高 3.7 倍，同時以足夠高的速度工作。
一行代碼即可啟動
一切準備就緒，只欠行動。根據(jù)研究人員介紹，Colossal-AI 提供了開箱即用的 ChatGPT 訓練代碼。在此，以 GPT 為例，只需要一行代碼就可以指定使用 Colossal-AI 作為系統(tǒng)策略來啟動。
使用以下命令，開發(fā)者可以快速啟動單 GPU 規(guī)模、單機多 GPU 規(guī)模、原始 1750 億參數(shù)規(guī)模版本的訓練，并對各種性能指標（包括最大GPU內(nèi)存使用率、吞吐量和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
底層優(yōu)化
低成本的 LoRA 微調(diào)
復制 ChatGPT 的實現(xiàn)過程通常依賴于 Colossal-AI。Colossal-AI 支持通過低秩矩陣微調(diào)（LoRA，Low-Rank Adaptation）方法進行高效微調(diào)。該方法假設大型語言模型是過度參數(shù)化的，微調(diào)過程中的參數(shù)變量是一個低秩矩陣，它可以分解為兩個小矩陣的乘積：
由于大型語言模型的參數(shù)是固定的，在微調(diào)過程中只有調(diào)整矩陣的參數(shù)，從而減少訓練參數(shù)的數(shù)量。在進行部署推理時，將矩陣的乘積加回到原始矩陣中，如，不影響推理延遲。
LoRA 結(jié)構(gòu)，只訓練 A、B
Zero Gemini 來減少內(nèi)存冗余
與傳統(tǒng)的數(shù)據(jù)并行策略相比，Colossal-AI 使用零冗余優(yōu)化器（ZeRO）來消除內(nèi)存冗余，并在不影響計算粒度和通信效率的情況下，提高內(nèi)存使用率。此外，為進一步提高了 ZeRO 的性能，Colossal-AI 還提出了基于 Chunk 的內(nèi)存管理機制，它可以將連續(xù)的參數(shù)集按操作順序存儲在一個連續(xù)的、均勻分割的內(nèi)存空間中，由此更能有效地利用網(wǎng)絡帶寬（PCI-e 和 GPU 之間），減少通信成本，并避免潛在的內(nèi)存碎片。
此外，Colossal-AI 的異構(gòu)內(nèi)存管理器 Gemini 通過將優(yōu)化器狀態(tài)卸載到 CPU ，以此減少 GPU 內(nèi)存占用，允許同時使用 GPU 內(nèi)存和 CPU 內(nèi)存（包括 CPU DRAM 或 NVMe SSD 內(nèi)存）來訓練超出單個 GPU 內(nèi)存限制的大規(guī)模模型。
寫在最后
目前，該研究團隊已經(jīng)開源了完整的算法和軟件設計來復制 ChatGPT 的實現(xiàn)過程：https://github.com/hpcaitech/ColossalAI。
不過，同樣是基于成本考慮，他們表示，「對于這樣一個巨型的人工智能模型，它需要更多的數(shù)據(jù)和計算資源來實際生效和部署。畢竟，用 1750 億個參數(shù)訓練一個 GPT-3 需要價值數(shù)百萬美元的計算能力。因此，大型預訓練模型長期以來只為少數(shù)大科技公司所擁有?！?/p>
因此，他們也希望能夠以開源的方式，吸引更多的研究人員、機構(gòu)共同參與進來，僅以上文中所復制 ChatGPT 訓練流程的實踐探索為起點，未來可以向大模型的時代做出努力。
更多詳情內(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
