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

語言: CN / TW / HK
時間 2022-07-13 23:03:05
主題:
大模型已成為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， 僅需添加幾行代碼，便可實現大模型的低成本訓練和微調。

低成本加速大模型OPT

OPT模型

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模型，仍離不開高效的分佈式並行擴展。對於同時使用數據並行、流水並行、2.5維張量並行等 複雜並行策略，僅需簡單聲明，即可自動實現，Colossal-AI無需像其他系統和框架侵入代碼，手動處理複雜的底層邏輯。
  
parallel = dict(
    pipeline=2,
    tensor=dict(mode='2.5d', depth = 1, size=4)
)
在面對擴展至數十甚至數百張GPU的大規模並行場景時，相比英偉達Megatron-LM等現有系統，Colossal-AI在性能上仍顯示出顯著的加速性能和資源節約。這意味着對於 預訓練GPT-3等超大AI模型，可節省數百萬元的訓練費用。
 
 
Colossal-AI相關解決方案已成功在 自動駕駛、雲計算、零售、醫藥、芯片等行業知名廠商落地應用，廣受好評。
Colossal-AI注重開源社區建設，提供中文教程，開放用户社羣及論壇，對於用户反饋進行高效交流與迭代更新，不斷添加PaLM、AlphaFold等前沿應用。
自然開源以來，Colossal-AI已經 多次在GitHub及Papers With Code熱榜位列世界第一，與眾多已有數萬star的明星開源項目一起受到海內外關注！

傳送門

項目地址：

招聘

潞晨科技還在廣納英才，招聘全職/實習AI分佈式系統、架構、編譯器、網絡、CUDA、SaaS、k8s等核心系統研發人員，開源社區運營、銷售人員。
潞晨科技提供有競爭力的薪資回報，特別優秀的，還可以申請遠程工作。也歡迎各位向潞晨科技引薦優秀人才，如果您推薦優秀人才成功簽約潞晨科技，我們將為您提供數千元至數萬元的推薦費。
工作地點：北京，新加坡，美國 (可相互轉崗)。
簡歷投遞郵箱： [email protected]

參考文獻：

 
 
「」