機器之心釋出
機器之心編輯部
「百模大戰」 正風起雲湧，AIGC 相關企業融資和并購金額也屢創新高，全球科技企業争相入局。然而，AI 大模型風光無限的背後是成本極其高昂，單次預訓練成本或高達上千萬元。基于 LLaMA 等現有開源大模型的微調，也難以滿足企業打造核心競争力和多樣化商業使用等需求。是以，如何低成本量身打造預訓練基礎大模型，已成為 AI 大模型浪潮的關鍵瓶頸。
Colossal-AI 作為全球最大、最活躍的大模型開發工具與社群，以目前被最廣泛使用的 LLaMA 為例，提供開箱即用的 650 億參數預訓練方案，可提升訓練速度 38%，為大模型企業節省大量成本。
開源位址：https://github.com/hpcaitech/ColossalAI
LLaMA 點燃開源熱情
Meta 開源的 7B~65B LLaMA 大模型進一步激發了打造類 ChatGPT 的熱情，并由此衍生出 Alpaca, Vicuna，ColossalChat 等微調項目。
但 LLaMA 隻開源了模型權重且限制商業使用，微調能夠提升和注入的知識與能力也相對有限。對于真正投身大模型浪潮的企業來說，仍必須預訓練自己的核心大模型。開源社群為此做了諸多努力：
- RedPajama：開源可商用類 LLaMA 資料集，無訓練代碼和模型
- OpenLLaMA：開源可商用類 LLaMA 7B,13B 模型，使用 EasyLM 基于 JAX 和 TPU 訓練
- Falcon：開源可商用類 LLaMA 7B,40B 模型，無訓練代碼
但對于最主流的 PyTorch + GPU 生态，仍缺乏高效、可靠、易用的類 LLaMA 基礎大模型預訓練方案。
最佳大模型預訓練方案提速 38%
針對上述空白與需求，Colossal-AI 首個開源了 650 億參數 LLaMA 低成本預訓練方案，相比業界其他主流選擇，可提升預訓練速度 38%，僅需 32 張 A100/A800 即可使用，并且不限制商業使用。
而像原生 PyTorch、FSDP 等，則因顯存溢出無法運作該任務。Hugging Face accelerate、DeepSpeed、Megatron-LM 也未對 LLaMA 預訓練進行官方支援。
開箱即用
1. 安裝 Colossal-AI
git clone -b example/llama https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI
# install and enable CUDA kernel fusion
CUDA_EXT=1 pip install . 2. 安裝其他依賴
cd examples/language/llama
# install other dependencies
pip install -r requirements.txt
# use flash attention
pip install xformers 3. 資料集
預設資料集 togethercomputer/RedPajama-Data-1T-Sample 将在首次運作時自動下載下傳，也可通過 -d 或 --dataset 指定自定義資料集。
4. 運作指令
已提供 7B 和 65B 的測速腳本，僅需根據實際硬體環境設定所用多節點的 host name 即可運作性能測試。
cd benchmark_65B/gemini_auto
bash batch12_seq2048_flash_attn.sh 對于實際的預訓練任務，使用與速度測試一緻，啟動相應指令即可，如使用 4 節點 * 8 卡訓練 65B 的模型。
colossalai run --nproc_per_node 8 --hostfile YOUR_HOST_FILE --
master_addr YOUR_MASTER_ADDR pretrain.py -c '65b' --plugin
"gemini" -l 2048 -g -b 8 -a 例如，使用 Colossal-AI gemini_auto 并行政策，可便捷實作多機多卡并行訓練，降低顯存消耗的同時保持高速訓練。還可根據硬體環境或實際需求，選擇流水并行 + 張量并行 + ZeRO1 等複雜并行政策組合。
其中，通過 Colossal-AI 的 Booster Plugins，使用者可以便捷自定義并行訓練，如選擇 Low Level ZeRO，Gemini，DDP 等并行政策。Gradient checkpointing 通過在反向傳播時重新計算模型的 activation 來減少記憶體使用。通過引入 Flash attention 機制加速計算并節省顯存。使用者可以通過指令行參數便捷控制數十個類似的自定義參數，在保持高性能的同時為自定義開發保持了靈活性。
Colossal-AI 最新的 ShardFormer 極大降低了使用多元并行訓練 LLM 的上手成本。現已支援包括 LLaMA 的多種等主流模型，且原生支援 Huggingface/transformers 模型庫。無需改造模型，即可支援多元并行 (流水、張量、ZeRO、DDP 等) 的各種配置組合，能夠在各種硬體配置上都發揮卓越的性能。
AI 大模型系統基礎設施 Colossal-AI
Colossal-AI 為該方案提供了核心系統優化與加速能力支援，它由加州伯克利大學傑出教授 James Demmel 和新加坡國立大學校長青年教授尤洋上司開發。Colossal-AI 基于 PyTorch，可通過高效多元并行、異構記憶體等，降低 AI 大模型訓練 / 微調 / 推理的開發與應用成本，降低 GPU 需求等。
Colossal-AI上述解決方案已在某世界500強落地應用，在千卡叢集性能優異，僅需數周即可完成千億參數私有大模型預訓練。上海AI Lab與商湯等新近釋出的InternLM也基于Colossal-AI在千卡實作高效預訓練。
自開源以來，Colossal-AI 已經多次在 GitHub 熱榜位列世界第一，獲得 GitHub Star 超 3 萬顆，并成功入選 SC、AAAI、PPoPP、CVPR、ISC 等國際 AI 與 HPC 頂級會議的官方教程，已有上百家企業參與共建 Colossal-AI 生态。其背後的潞晨科技，近期獲得數億元 A 輪融資，已在成立 18 個月内已迅速連續完成三輪融資。
開源位址：
https://github.com/hpcaitech/ColossalAI
參考連結：
https://www.hpc-ai.tech/blog/large-model-pretraining