無縫支持Hugging Face社區，Colossal-AI低成本輕松加速大模型

作者：機器之心編輯部 2022-07-13 16:45:34

?大模型已成為 AI 圈的一種潮流，不僅橫掃各大性能榜單，更產生了諸多有趣應用。例如，微軟和 OpenAI 開發的自動代碼建議補全神器 Copilot，化身程序員最佳助手，提升工作效率。

OpenAI 剛剛發布能以假亂真的文本生成圖像模型 DALL-E 2，Google 便緊接著發布了 Imagen，在大模型上，大公司也是相當的卷，絲毫不比 CV 刷榜差。

文本到圖像生成樣例“一個被貓絆倒的希臘人雕像”（左側兩列為 Imagen，右側兩列為 DALL·E 2）模型增大帶來的神奇表現，使得近幾年預訓練模型規模呈現爆炸式增長。然而，訓練甚至微調大模型都需要非常高的硬件成本，動輒幾十、上百張 GPU。此外，PyTorch、TensorFlow 等現有深度學習框架也難以有效處理超大模型，通常需要專業的 AI 系統工程師做針對具體模型做適配和優化。更重要的是，不是每一個實驗室以及研發團隊都具備 “鈔” 能力，能夠隨時調用大規模 GPU 集群來使用大模型，更不用提僅有一張顯卡的個人開發者。因此，盡管大模型已經吸引了大量關注，高昂的上手門檻卻令大眾“望塵莫及”。

導致大模型使用成本增高的核心原因是顯存限制。GPU 計算雖快，但顯存容量有限，無法容納大模型。Colossal-AI 針對這一痛點，通過異構內存系統，高效地同時使用 GPU 顯存以及價格低廉的 CPU 內存，在僅有一塊 GPU 的個人 PC 上便能訓練高達 180 億參數 GPT，可提升模型容量十余倍，大幅度降低了 AI 大模型微調和推理等下游任務和應用部署的門檻，還能便捷擴展至大規模分布式。Hugging Face 為深度學習社區提供了超過 5 萬個 AI 模型的實現，最其中也不乏像 GPT, OPT 這樣的大模型，現已成為最流行的 AI 庫之一。

Colossal-AI 無縫支持 Hugging Face 社區模型，讓大模型對每一位開發者都變得觸手可及。接下來，我們將以 Meta 發布的大模型 OPT 為例，展現如何使用 Colossal-AI，僅需添加幾行代碼，便可實現大模型的低成本訓練和微調。

開源地址：https://github.com/hpcaitech/ColossalAI低成本加速大模型 OPTOPT 模型OPT 的全稱為 Open Pretrained Transformer，是 Meta(Facebook) AI 實驗室發布的對標 GPT-3 的大規模 Transformer 模型。與 OpenAI 尚未公開模型權重的 GPT-3 相比，Meta AI 慷慨地開源了所有的代碼以及模型權重，極大推動了 AI 大模型的民主化，每一位開發者都能以此為基礎開發個性化的下游任務。接下來，我們將用 Hugging Face 提供的 OPT 模型的預訓練權重進行 Casual Language Modelling 的微調。添加配置文件想要使用 Colossal-AI 中各個強大功能，用戶無需更改代碼訓練邏輯，只用添加一個簡單的配置文件，即可賦予模型所期望的功能，比如混合精度、梯度累積、多維并行訓練、冗余內存優化等。在一張 GPU 上，以異構訓練為例，我們只需在配置文件里加上相關配置項。其中 tensor_placement_policy 決定了我們異構訓練的策略，這個參數可以為 cuda、cpu 以及 auto。各個策略有不同的優點：

cuda: 將全部模型參數都放置于 GPU 上，適合不 offload 時仍然能進行訓練的傳統場景；
cpu 則會將模型參數都放置在 CPU 內存中，僅在 GPU 顯存中保留當前參與計算的權重，適合超大模型的訓練；
auto 則會根據實時的內存信息，自動決定保留在 GPU 顯存中的參數量，這樣能最大化利用 GPU 顯存，同時減少 CPU-GPU 之間的數據傳輸。

對于一般用戶而言，僅需選擇 auto 策略，由 Colossal-AI 自動化地實時動態選擇最佳異構策略，最大化計算效率。

from colossalai.zero.shard_utils import TensorShardStrategy


zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(),
                              tensor_placement_policy="auto"),
            optimizer_config=dict(gpu_margin_mem_ratio=0.8)

運行啟動在配置文件準備好之后，我們只需插入幾行代碼即可啟動聲明的新功能。首先，通過一行代碼，使用配置文件啟動 Colossal-AI，Colossal-AI 會自動初始化分布式環境，并讀取相關配置，之后將配置里的功能自動注入到模型以及優化器等組件中。

colossalai.launch_from_torch(config='./configs/colossalai_zero.py')

接下來，用戶可以照常定義數據集、模型、優化器、損失函數等，例如直接使用原生 PyTorch 代碼。在定義模型時，只需將模型放置于 ZeroInitContext 下初始化即可。在例子里，我們使用 Hugging Face 提供的 OPTForCausalLM 模型以及預訓練權重，在 Wikitext 數據集上進行微調。

with ZeroInitContext(target_device=torch.cuda.current_device(), 
                    shard_strategy=shard_strategy,
                    shard_param=True):
    model = OPTForCausalLM.from_pretrained(
                'facebook/opt-1.3b'
                config=config
            )

接著，只需要調用 colossalai.initialize，便可將配置文件里定義的異構內存功能統一注入到訓練引擎中，即可啟動相應功能。

engine, train_dataloader, eval_dataloader, lr_scheduler = colossalai.initialize(model=model,
                                                                               optimizer=optimizer,
                                                                               criterion=criterion,
                                                                               train_dataloader=train_dataloader,
                                                                               test_dataloader=eval_dataloader,
                                                                               lr_scheduler=lr_scheduler）

優勢顯著在單張 GPU，與微軟 DeepSpeed 相比，Colossal-AI 的使用自動化的 auto 策略，在不同的模型規模上相比 DeepSpeed 的 ZeRO Offloading 策略，均體現出顯著優勢，最快可實現 40% 的加速。而 PyTorch 等傳統深度學習框架，在單張 GPU 上已經無法運行如此大的模型。

對于使用 8 張 GPU 的并行訓練，Colossal-AI 僅需在啟動命令中添加 - nprocs 8 即可實現！背后秘訣如此顯著的提升來自于 Colossal-AI 的高效異構內存管理子系統 Gemini。簡單的來說，在模型訓練時，Gemini 在前面的幾個 step 進行預熱，收集 PyTorch 動態計算圖中的內存消耗信息；在預熱結束后，計算一個算子前，利用收集的內存使用記錄，Gemini 將預留出這個算子在計算設備上所需的峰值內存，并同時從 GPU 顯存里移動一些模型張量到 CPU 內存。

Gemini 內置的內存管理器給每個張量都標記一個狀態信息，包括 HOLD、COMPUTE、FREE 等。然后，根據動態查詢到的內存使用情況，不斷動態轉換張量狀態、調整張量位置，相比起 DeepSpeed 的 ZeRO Offload 的靜態劃分，Colossal-AI Gemini 能更高效利用 GPU 顯存和 CPU 內存，實現在硬件極其有限的情況下，最大化模型容量和平衡訓練速度。

對于大模型的代表 GPT，使用 Colossal-AI 在搭載 RTX 2060 6GB 的普通游戲筆記本上，也足以訓練高達 15 億參數模型；對于搭載 RTX3090 24GB 的個人電腦，更是可以直接訓練 180 億參數的模型；對于 Tesla V100 等專業計算卡，Colossal-AI 也能顯示出顯著改善。更進一步：便捷高效并行擴展并行分布式技術是進一步加速模型訓練的重要手段，想要以最短時間訓練當今世界最大最前沿的 AI 模型，仍離不開高效的分布式并行擴展。針對現有方案并行維度有限、效率不高、通用性差、部署困難、缺乏維護等痛點，Colossal-AI 通過高效多維并行和異構并行等技術，讓用戶僅需極少量修改，即可高效快速部署 AI 大模型訓練。例如，對于同時使用數據并行、流水并行、2.5 維張量并行等復雜并行策略，僅需簡單聲明，即可自動實現，Colossal-AI 無需像其他系統和框架侵入代碼，手動處理復雜的底層邏輯。

Python
parallel = dict(
    pipeline=2,
    tensor=dict(mode='2.5d', depth = 1, size=4)
)

具體來說，對于 GPT-3 這樣的超大 AI 模型，相比英偉達方案，Colossal-AI 僅需一半的計算資源，即可啟動訓練；若使用相同計算資源，則能提速 11%，可降低 GPT-3 訓練成本超百萬美元。Colossal-AI 相關解決方案已成功自動駕駛、云計算、零售、醫藥、芯片等行業知名廠商落地應用，廣受好評。

例如，對于蛋白質結構預測應用 AlphaFold2，基于 Colossal-AI 的加速方案的 FastFold，成功超越谷歌和哥倫比亞大學的方案，將 AlphaFold2 訓練時間從 11 天減少到 67 小時，且總成本更低，在長序列推理中也實現 9.3～11.6 倍的速度提升。

Colossal-AI 注重開源社區建設，提供中文教程，開放用戶社群及論壇，對于用戶反饋進行高效交流與迭代更新，不斷添加 PaLM、AlphaFold 等前沿應用。自然開源以來，Colossal-AI 已經多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數萬 star 的明星開源項目一起受到海內外關注！

傳送門項目地址：https://github.com/hpcaitech/ColossalAI?

責任編輯：未麗燕來源：機器之心

?大模型 AI 微軟