訓練大模型時,顯存都哪去了? 原創
GPT-2(XL)有15億個參數,使用16位精度,一個參數占用2個字節的內存,因此這些參數大約占用3GB的內存。
按照如下超參數設置:
- 優化器 → Adam
- 批量大小 → 32
- 變換層數量 → 48
- 序列長度 → 1000
要想在單個GPU上訓練GPT-2,所需的最小內存大概是多少?
答案可能會嚇到你。
在一個擁有32GB內存的單個GPU上,幾乎無法訓練一個3GB的GPT-2模型。
但這怎么可能呢?內存都去哪了?讓我們來了解一下。
模型在訓練過程中有很多方面會持續占用內存。
#1)優化器狀態,梯度,模型參數
混合精度訓練廣泛用于加速模型訓練。
顧名思義,這個方法的思想是在訓練過程中同時使用float16低精度(在卷積和矩陣乘法等操作中)和高精度(如32位浮點數,float32)。
這就是“混合精度”名稱的由來。
前向傳播和反向傳播都使用16位浮點數表示權重和梯度。
因此,如果模型有Φ個參數,那么:
● 權重將占用2 * Φ字節的內存。
● 梯度將占用2 * Φ字節的內存。
這里的“2”表示每個參數占用2個字節的內存(16位)。
Adam 是最受歡迎的模型訓練優化器之一。
雖然許多實踐者僅僅因為它流行而使用它,但他們沒有意識到,在訓練過程中,Adam 會存儲兩種優化器狀態來計算更新——梯度的動量和方差。
因此,如果模型有Φ個參數,那么這兩個優化器狀態將消耗:
● 4 * Φ 字節用于動量。
● 另需 4 * Φ 字節用于方差。
這里的“4”表示每個參數占用 4 個字節的內存(32 位)。
此外,反向傳播結束時的更新仍然在32位精度下進行,以確保有效的計算。這導致:
● 另需 4 * Φ 字節用于模型參數。
讓我們把它們加起來:
這就是 16 * Φ,或者 24GB 的內存,遠遠高于 16 位參數所使用的 3GB 內存。
而且我們還沒有考慮到所有的因素。
2#)激活值
對于像大型深度學習模型(如大語言模型,LLMs)來說,激活值在訓練過程中占用了大量內存。
更確切地說,在GPT-2的一個Transformer塊中計算的激活值總數是:
因此,在所有的Transformer塊中,總計就是:
這是 GPT-2-XL 的配置:
總共大約是 300 億個激活值。由于每個激活值使用 16 位表示,所有激活值總共占用 60GB 的內存。
通過使用像梯度檢查點(在上一章討論過的)這樣的技術,可以將內存消耗降低到大約 8-9GB,但這也會帶來額外 25-30% 的計算開銷。
除了可以計算的內存開銷外,還有一些額外的開銷,例如內存碎片化。
內存碎片化是指在分配的內存塊之間存在小的未使用間隙,導致可用內存的低效使用。
內存分配請求失敗是因為沒有足夠的連續內存塊可用。
在上述討論中,我們考慮了一個相對較小的模型——GPT-2(XL),它有 15 億個參數,與如今訓練的模型規模相比非常小。
然而,這個討論可能幫助你反思構建大規模語言模型(LLMs)時的固有挑戰。很多人常說,GPT 模型只是簡單地堆疊更多的層并使網絡變得更大。
如果真是那么簡單,大家都會在做了。從這個討論中,你可能已經理解到,這并不像僅僅添加更多層那么簡單。
即便是增加一層,也可能導致額外數 GB 的內存需求。多 GPU 訓練是這些模型的核心技術,我們將在另一篇文章中討論。
本文轉載自公眾號人工智能大講堂
原文鏈接:??https://mp.weixin.qq.com/s/PFQZnqJJ-tjFcv6oSjaV-A??
