A100 買不到了,只有小顯示卡怎麼訓大模型

語言: CN / TW / HK

為了達到更好的訓練效果,通常煉丹師們會使用更大的模型和更大的 Batch size,但因此帶來的大視訊記憶體佔用,卻成為不可避免的硬傷。

尤其是如今 GPU 越來越貴,甚至還可能買不到......

MegEngine v1.5 及以上版本,支援動態圖和靜態圖的視訊記憶體優化,視訊記憶體佔用可降至 1/4。

先上對比效果

背後的邏輯很簡單:計算換視訊記憶體 - MegEngine 使用重計算策略,實現了用時間換空間。

使用方法

動態圖下開啟:

在訓練程式碼之前只需新增一行程式碼:

megengine.dtr.enable()

靜態圖下開啟:

在編譯靜態圖時使用 DTRConfig 設定 trace 的引數 dtr_config

from megengine.jit import trace, DTRConfig
 
config = DTRConfig(eviction_threshold=8*1024**3)
 
@trace(symbolic=True, dtr_config=config)
def train_func(data, label, * , net, optimizer, gm):
    ...

更多使用技巧,見官方文件

使用 DTR 進行視訊記憶體優化 - MegEngine 1.11 文件

延展閱讀:

To 新朋友的特別提示:

DTR 作為 MegEngine 原生特性,無法脫離 MegEngie 獨立使用。

好在 MegEngine 頂層 API 基於 Python,採取了類似於 PyTorch 的風格,已被眾多使用者反饋易上手,入門簡單。

還有豐富的使用者指南文件:

所以,不用顧慮太多,可以勇敢嘗試哦~

開源地址:

(歡迎 star~ :D