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

大模型已成為 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<br><br><br>zero = dict(model_config=dict(shard_strategy=TensorShardStrategy(),<br>tensor_placement_policy="auto"),<br>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(), <br>shard_strategy=shard_strategy,<br>shard_param=True):<br>model = OPTForCausalLM.from_pretrained(<br>'facebook/opt-1.3b'<br>config=config<br>)

接着，只需要调用 colossalai.initialize，便可将配置文件里定义的异构内存功能统一注入到训练引擎中，即可启动相应功能。

engine, train_dataloader, eval_dataloader, lr_scheduler = colossalai.initialize(model=model,<br> optimizer=optimizer,<br> criterion=criterion,<br> train_dataloader=train_dataloader,<br> test_dataloader=eval_dataloader,<br> 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<br>parallel = dict(<br>pipeline=2,<br>tensor=dict(mode='2.5d', depth = 1, size=4)<br>)

具体来说，对于 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

以上是無縫支援Hugging Face社區，Colossal-AI低成本輕鬆加速大模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！