無縫支援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]

參考文獻：

 
 
「」