650億參數(shù)大模型的預(yù)訓(xùn)練方案，發(fā)布即開(kāi)源。

訓(xùn)練速度較傳統(tǒng)方案提升38%。

這就是由Colossal-AI最新發(fā)布的類(lèi)LLaMA基礎(chǔ)大模型預(yù)訓(xùn)練方案。

要知道，在“百模大戰(zhàn)”背景下，誰(shuí)擁有自家大模型，往往被視為核心競(jìng)爭(zhēng)力。

在這個(gè)節(jié)點(diǎn)下，愿意開(kāi)源大模型的公司少之又少。

但自己從頭訓(xùn)練一個(gè)大模型，對(duì)技術(shù)、資金都有很高要求。

由此，Colossal-AI最新的開(kāi)源動(dòng)作，可以說(shuō)是應(yīng)時(shí)勢(shì)所需了。

并且它還不限制商業(yè)使用，開(kāi)箱即用僅需4步。

具體項(xiàng)目有哪些內(nèi)容？一起往下看~

開(kāi)源地址：https://github.com/hpcaitech/ColossalAI

32張A100/A800即可使用

實(shí)際上，自從Meta開(kāi)源LLaMA后，掀起了一波微調(diào)項(xiàng)目熱潮，如Alpaca、Vicuna、ColossalChat等都是在其基礎(chǔ)上打造的。

但是LLaMA只開(kāi)源了模型權(quán)重且限制商業(yè)使用，微調(diào)能夠提升和注入的知識(shí)與能力也相對(duì)有限。

對(duì)于真正想要投身大模型浪潮的企業(yè)來(lái)說(shuō)，訓(xùn)練自己的核心大模型非常重要。

開(kāi)源社區(qū)也此前已獻(xiàn)了一系列工作：

• RedPajama：開(kāi)源可商用類(lèi)LLaMA數(shù)據(jù)集（無(wú)訓(xùn)練代碼和模型）
• OpenLLaMA：開(kāi)源可商用類(lèi)LLaMA 7B/13B模型，使用EasyLM基于JAX和TPU訓(xùn)練
• Falcon：開(kāi)源可商用類(lèi)LLaMA 7B/40B模型（無(wú)訓(xùn)練代碼）

但這些都還不夠，因?yàn)閷?duì)于最主流的PyTorch+GPU生態(tài)，仍缺乏高效、可靠、易用的類(lèi)LLaMA基礎(chǔ)大模型預(yù)訓(xùn)練方案。

所以Colossal-AI交出了最新的開(kāi)源答卷。

僅需32張A100/A800，即可搞定650億參數(shù)類(lèi)LLaMA大模型預(yù)訓(xùn)練，訓(xùn)練速度提升38%。

而像原生PyTorch、FSDP等，則因顯存溢出無(wú)法運(yùn)行該任務(wù)。

Hugging Face accelerate、DeepSpeed、Megatron-LM也未對(duì)LLaMA預(yù)訓(xùn)練進(jìn)行官方支持。

開(kāi)箱即用、4步搞定

而這一項(xiàng)目真正上手起來(lái)也很簡(jiǎn)易。共有四步：

• 1、安裝Colossal-AI
• 2、安裝其他依賴(lài)項(xiàng)
• 3、數(shù)據(jù)集
• 4、運(yùn)行命令

具體代碼如下：

第一步、安裝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 .

第二步、安裝其他依賴(lài)項(xiàng)。

cd examples/language/llama
# install other dependencies
pip install -r requirements.txt
# use flash attention
pip install xformers

第三步、數(shù)據(jù)集。

默認(rèn)數(shù)據(jù)集togethercomputer/RedPajama-Data-1T-Sample將在首次運(yùn)行時(shí)自動(dòng)下載，也可通過(guò)-d或—dataset指定自定義數(shù)據(jù)集。

第四步、運(yùn)行命令。

已提供7B和65B的測(cè)速腳本，僅需根據(jù)實(shí)際硬件環(huán)境設(shè)置所用多節(jié)點(diǎn)的host name即可運(yùn)行性能測(cè)試。

cd benchmark_65B/gemini_auto
bash batch12_seq2048_flash_attn.sh

對(duì)于實(shí)際的預(yù)訓(xùn)練任務(wù)，使用與速度測(cè)試一致，啟動(dòng)相應(yīng)命令即可，如使用4節(jié)點(diǎn)*8卡訓(xùn)練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并行策略，可便捷實(shí)現(xiàn)多機(jī)多卡并行訓(xùn)練，降低顯存消耗的同時(shí)保持高速訓(xùn)練。

還可根據(jù)硬件環(huán)境或?qū)嶋H需求，選擇流水并行+張量并行+ZeRO1等復(fù)雜并行策略組合。

其中，通過(guò)Colossal-AI的Booster Plugins，用戶(hù)可以便捷自定義并行訓(xùn)練，如選擇Low Level ZeRO、Gemini、DDP等并行策略。

Gradient checkpointing通過(guò)在反向傳播時(shí)重新計(jì)算模型的activation來(lái)減少內(nèi)存使用。

通過(guò)引入Flash attention機(jī)制加速計(jì)算并節(jié)省顯存。用戶(hù)可以通過(guò)命令行參數(shù)便捷控制數(shù)十個(gè)類(lèi)似的自定義參數(shù)，在保持高性能的同時(shí)為自定義開(kāi)發(fā)保持了靈活性。

Colossal-AI最新的ShardFormer極大降低了使用多維并行訓(xùn)練LLM的上手成本。

現(xiàn)已支持包括LLaMA的多種等主流模型，且原生支持Huggingface/transformers模型庫(kù)。

無(wú)需改造模型，即可支持多維并行(流水、張量、ZeRO、DDP等)的各種配置組合，能夠在各種硬件配置上都發(fā)揮卓越的性能。

Colossal-AI：大模型系統(tǒng)基礎(chǔ)設(shè)施

帶來(lái)如上新工作的Colossal-AI，如今已是大模型趨勢(shì)下的明星開(kāi)發(fā)工具和社區(qū)了。

Colossal-AI上述解決方案已在某世界500強(qiáng)落地應(yīng)用，在千卡集群性能優(yōu)異，僅需數(shù)周即可完成千億參數(shù)私有大模型預(yù)訓(xùn)練。

上海AI Lab與商湯等新近發(fā)布的InternLM也基于Colossal-AI在千卡實(shí)現(xiàn)高效預(yù)訓(xùn)練。

自開(kāi)源以來(lái)，Colossal-AI多次在GitHub熱榜位列世界第一，獲得 GitHub Star超3萬(wàn)顆，并成功入選SC、AAAI、PPoPP、CVPR、ISC等國(guó)際 AI 與HPC頂級(jí)會(huì)議的官方教程，已有上百家企業(yè)參與共建Colossal-AI生態(tài)。

它由加州伯克利大學(xué)杰出教授 James Demmel 和新加坡國(guó)立大學(xué)校長(zhǎng)青年教授尤洋領(lǐng)導(dǎo)開(kāi)發(fā)。

Colossal-AI基于PyTorch，可通過(guò)高效多維并行、異構(gòu)內(nèi)存等，主打?yàn)锳I大模型訓(xùn)練/微調(diào)/推理的開(kāi)發(fā)與應(yīng)用成本，降低GPU需求等。

其背后公司潞晨科技，近期獲得數(shù)億元A輪融資，已在成立18個(gè)月內(nèi)已迅速連續(xù)完成三輪融資。

開(kāi)源地址：https://github.com/hpcaitech/ColossalAI。

參考鏈接：https://www.hpc-ai.tech/blog/large-model-pretraining。