很多時(shí)候,內(nèi)存限制會(huì)阻礙 ViT 以及 LLM 的訓(xùn)練,這篇文章介紹了 9 種減少內(nèi)存消耗峰值的方法。難能可貴的是,這幾種方法可以同時(shí)使用,就好像降龍十八掌中最后一掌,正是將前幾張組合起來就能打出最強(qiáng)大的效果。
峰值內(nèi)存消耗是訓(xùn)練深度學(xué)習(xí)模型(如視覺 Transformer 和 LLM)時(shí)的常見瓶頸。本文提供了一系列可以在不犧牲建模性能和預(yù)測(cè)精度的情況下,將 PyTorch 中的內(nèi)存消耗降低到約 1/20 的技術(shù)。以 PyTorch 的 Torchvision 庫中的視覺 transformer 為基礎(chǔ),本文作者編寫了大約 100 行代碼的訓(xùn)練腳本,并且所有代碼示例都可以在 GitHub 上找到。- 微調(diào) vision transformer
- Reduced Batch Size 訓(xùn)練
- 在目標(biāo)設(shè)備上實(shí)例化模型
- 以上九種方法結(jié)合起來,就形成了一種可以用于 llm 的綜合方法,也可以稱之為第十種方法。
本文在實(shí)驗(yàn)中使用的 ViT 為 ViT-L-16 模型。在依次將上述方法添加后,研究者將訓(xùn)練 BigBird-Roberta LLM 來執(zhí)行文本分類任務(wù)。這些技術(shù)使得在消費(fèi)類硬件上訓(xùn)練這樣的模型成為可能。微調(diào) vision transformer為了簡(jiǎn)化實(shí)驗(yàn)中的 PyTorch 代碼,本文使用了開源庫 ——Fabric,十幾行代碼就能應(yīng)用各種先進(jìn)的 PyTorc 技術(shù)(自動(dòng)混合精度訓(xùn)練、多 GPU 訓(xùn)練、張量分片等)。原生 PyTorch 代碼和修改后的使用 Fabric 的代碼之間的區(qū)別很微妙,只有較小的修改,如下面的代碼所示:如上所述,改動(dòng)雖然不大,但是可以方便的使用 PyTorch 中的高級(jí)功能,而無需重新構(gòu)造任何現(xiàn)有代碼。總結(jié)上圖,將普通 PyTorch 代碼轉(zhuǎn)換為 PyTorch+Fabric 的主要 3 個(gè)步驟可以歸納如下:- 導(dǎo)入 Fabric 并實(shí)例化 Fabric 對(duì)象。
- 使用 Fabric 設(shè)置模型、優(yōu)化器和數(shù)據(jù)加載程序。
- 調(diào)用 fabric.backward () 構(gòu)造損失函數(shù),而不是通常使用的 loss.backward ()
使用普通 PyTorch 和 PyTorch with Fabric 的性能和內(nèi)存消耗幾乎完全相同:Plain PyTorch (01_pytorch-vit.py):Time elapsed 17.94 min
Memory used: 26.79 GB
Test accuracy 95.85%
PyTorch with Fabric (01-2_pytorch-fabric.py)
Time elapsed 17.88 min
Memory used: 26.84 GB
Test accuracy 96.06%
model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)
model = vit_l_16(weights=None)
替換后,將不再是微調(diào),而是從頭開始訓(xùn)練相同的 ViT 架構(gòu),預(yù)測(cè)準(zhǔn)確率會(huì)從 96% 以上下降到約 60%:上一節(jié)使用 Fabric 修改了 PyTorch 代碼,在此基礎(chǔ)上,使用混合精度和分布式訓(xùn)練,也只需更改一行代碼。應(yīng)用混合精度訓(xùn)練,只需一個(gè)小的修改,將下面這行代碼fabric = Fabric(accelerator='cuda', devices=1)
fabric = Fabric(accelerator='cuda', devices=1, precision='16-mixed')
之后,在不犧牲預(yù)測(cè)精度的情況下,內(nèi)存消耗從 26.84GB 減少到 18.21GB,如下所示:01-2_pytoch-fabric.py 和 02_mixed-precision.py 的結(jié)果對(duì)比此外,混合精確訓(xùn)練不僅減少了內(nèi)存使用,還將運(yùn)行時(shí)間減少了 6 倍(從 17.88 分鐘減少到 3.45 分鐘),這可以說是意外收獲。混合精度訓(xùn)練同時(shí)使用 16 位和 32 位精度,以確保不損失精度。16 位表示的梯度計(jì)算比 32 位格式快得多,并且節(jié)省了大量的內(nèi)存。這種策略是有益的,尤其是當(dāng)受到內(nèi)存或計(jì)算限制時(shí)。之所以被稱為「混合」而不是「低」精度訓(xùn)練的原因是,并不會(huì)將所有參數(shù)和操作都轉(zhuǎn)移成 16 位浮點(diǎn)數(shù)。實(shí)際上,在訓(xùn)練期間會(huì)在 32 位和 16 位運(yùn)算之間切換。如下圖所示,混合精度訓(xùn)練可以分解為:將權(quán)重轉(zhuǎn)換為較低精度(如 FP16)以實(shí)現(xiàn)更快的計(jì)算、計(jì)算梯度、將梯度轉(zhuǎn)換回較高精度(FP32)以實(shí)現(xiàn)數(shù)值穩(wěn)定性,以及用縮放的梯度更新原始權(quán)重等幾個(gè)步驟。這種方法在保證訓(xùn)練有效的前提下,還能保持神經(jīng)網(wǎng)絡(luò)的準(zhǔn)確性和穩(wěn)定性。感興趣的讀者還可以在本文作者的另一篇文章:《使用混合精度技術(shù)加速大型語言模型》中獲得更多底層概念。文章地址:https:///pages/community/tutorial/accelerating-large-language-models-with-mixed-precision-techniques/還可以更進(jìn)一步,嘗試以「完全」較低的 16 位精度運(yùn)行,而不是混合精度。fabric = Fabric(accelerator='cuda', precision='16-mixed')
fabric = Fabric(accelerator='cuda', precision='16-true')
但需要注意的是,這樣會(huì)在訓(xùn)練中產(chǎn)生 NaN 值:Epoch: 0001/0001 | Batch 0000/0703 | Loss: 2.4105
Epoch: 0001/0001 | Batch 0300/0703 | Loss: nan
Epoch: 0001/0001 | Batch 0600/0703 | Loss: nan
...
這是因?yàn)槌R?guī)的 16 位浮點(diǎn)只能表示 - 65504 和 65504 之間的數(shù)字:In [1]: import torch
In [2]: torch.finfo(torch.float16)
Out[2]: finfo(resolution=0.001, min=-65504, max=65504, eps=0.000976562, smallest_normal=6.10352e-05, tiny=6.10352e-05, dtype=float16)
因此,為了避免 NaN 問題,可以將參數(shù)修改為「bf16 true」:fabric = Fabric(accelerator='cuda', precision='bf16-true')
可以將內(nèi)存消耗進(jìn)一步降低到 13.82 GB(同樣,在不犧牲準(zhǔn)確性的情況下):將 03_bfloat16.py 與之前的代碼的結(jié)果進(jìn)行比較「bf16 mixed」中的「bf16」代表 Brain Floating Point(bfloat16)。谷歌為機(jī)器學(xué)習(xí)和深度學(xué)習(xí)應(yīng)用程序開發(fā)了這種格式,特別是在其張量處理單元(TPU)中。與傳統(tǒng) float16 格式相比,Bfloat16 以降低精度為代價(jià)擴(kuò)展了動(dòng)態(tài)范圍。擴(kuò)展的動(dòng)態(tài)范圍有助于 bfloat16 表示非常大和非常小的數(shù)字,使其更適合可能遇到廣泛值的深度學(xué)習(xí)應(yīng)用。然而,較低的精度可能會(huì)影響某些計(jì)算的準(zhǔn)確性,或在某些情況下導(dǎo)致舍入誤差。但在大多數(shù)深度學(xué)習(xí)應(yīng)用中,這種精度的降低對(duì)建模性能的影響微乎其微。雖然 bfloat16 最初是為 TPU 開發(fā)的,但這種格式從 A100 Tensor Core GPU 開始,也得到了其之后的 NVIDIA GPU 的支持。以下代碼可以檢查 GPU 是否支持 bfloat16:>>> import torch
>>> torch.cuda.is_bf16_supported()
True
減少批大小通常是減少內(nèi)存消耗的一個(gè)有效方法。然而,它有時(shí)會(huì)導(dǎo)致較差的預(yù)測(cè)性能,因?yàn)檫@樣要改變訓(xùn)練動(dòng)態(tài)。無論哪種方式,需要探討減少批量大小對(duì)結(jié)果有何影響。事實(shí)證明,可以在不犧牲性能的情況下將批大小降低到 16,從而將內(nèi)存消耗降至 5.69 GB:將 04_lower-batchsize.py 與以前的代碼進(jìn)行比較。梯度累積是一種在訓(xùn)練過程中虛擬增加批大小的方法,當(dāng)可用的 GPU 內(nèi)存不足以容納所需的批量大小時(shí),這是非常有用的。并且這種方法只會(huì)在運(yùn)行時(shí)產(chǎn)生影響,建模性能并不會(huì)受到影響。梯度累積中,每批計(jì)算的量較小,并在多次迭代中累積梯度(通常求和或求平均),而不是在每個(gè)批次之后立刻更新模型權(quán)重。一旦累積的梯度達(dá)到目標(biāo)「虛擬」批大小,模型權(quán)重就會(huì)用累積的梯度更新。為了實(shí)現(xiàn)梯度積累,只需要對(duì)向前和向后傳球進(jìn)行兩次小的修改:05_gradient-acum.py 中的代碼修改本文作者的另一篇文章《使用梯度累積在單個(gè) GPU 上微調(diào) LLM》,更詳細(xì)地介紹了梯度累積的細(xì)節(jié)。文章地址:https:///blog/gradient-accumulation/
有效批大小為 16,并且累積步數(shù)為 4,意味著實(shí)際批大小為 4(因?yàn)?16/4=4)。05_gradient-acum.py 的結(jié)果這種技術(shù)的缺點(diǎn)是運(yùn)行時(shí)間從 3.96 分鐘增加到 12.91 分鐘。值得注意的是,批大小最小可以減少到 1,進(jìn)一步減少 75% 的內(nèi)存消耗。時(shí)下流行的 Adam 優(yōu)化器其實(shí)附帶了額外的參數(shù),例如,Adam 為每個(gè)模型參數(shù)提供了 2 個(gè)額外的優(yōu)化器參數(shù)(平均值和方差)。因此,通過將 Adam 與 SGD 等無狀態(tài)優(yōu)化器進(jìn)行交換,可以將參數(shù)數(shù)量減少 2/3,這在使用 ViT 和 LLM 時(shí)非常重要。普通 SGD 的缺點(diǎn)是收斂性較差。因此,Adam 與 SGD 交換后,需要引入余弦衰減學(xué)習(xí)速率調(diào)度器來進(jìn)行補(bǔ)償。optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=num_steps)
通過這種變化,模型能夠在保持大約 97% 分類準(zhǔn)確率的同時(shí)減少峰值內(nèi)存消耗:06_sgd-with-scheduler.py 的結(jié)果在目標(biāo)設(shè)備上創(chuàng)建模型在 PyTorch 中實(shí)例化模型時(shí),通常是首先在 CPU 設(shè)備上創(chuàng)建它,然后將它轉(zhuǎn)移到目標(biāo)設(shè)備上,并將其轉(zhuǎn)換為所需的精度:model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)
model.cuda().float16()
但是 CPU 上生成完整精度的中間模型,是一種低效的方法。所以,可以使用 Fabric 中的 init_module 上下文在目標(biāo)設(shè)備(例如 GPU)上直接創(chuàng)建所需精度的模型:import lightning as L
fabric = Fabric(accelerator='cuda', devices=1, precision='16-true')
with fabric.init_module():
model = vit_l_16(weights=ViT_L_16_Weights.IMAGENET1K_V1)
在這種特定情況下(模型),前向通過期間的峰值內(nèi)存大于其全精度表示中的大小。對(duì)模型加載本身對(duì) fabric.init_module 方法進(jìn)行基準(zhǔn)測(cè)試,結(jié)果如下:- 沒有 init_module 的 GPU 峰值內(nèi)存:1.24 GB(07_01_init-module.py)
- GPU 帶 init_module 的峰值內(nèi)存:0.65 GB(07_03_init-module.py)
可以看到,在這種情況下,init_module 將模型加載的峰值內(nèi)存需求減少了 50%。有關(guān) init_module 的更多詳細(xì)信息,可以參閱這篇關(guān)于大型模型的高效初始化的的文章。文章地址:https:///pages/community/efficient-initialization-of-large-models/下一個(gè)修改是多 GPU 訓(xùn)練。多個(gè) GPU 可供使用是有效的,因?yàn)檫@樣做可以更快地訓(xùn)練模型。然而,本文探討的是內(nèi)存節(jié)省。因此,需要一種更先進(jìn)的分布式多 GPU 策略,稱為完全共享數(shù)據(jù)并行(FSDP),該策略利用數(shù)據(jù)并行性和張量并行性在多個(gè)設(shè)備上共享大權(quán)重矩陣。但是如果模型已經(jīng)很小了,例如將此技術(shù)添加到上面第 7 節(jié)的代碼中時(shí),是幾乎看不到任何效果的。因此,為了純粹地關(guān)注分片的效果,可以與第 1 節(jié)中的全精度基線進(jìn)行比較。fabric = Fabric(accelerator='cuda', devices=1)
auto_wrap_policy = partial(
transformer_auto_wrap_policy,
transformer_layer_cls={EncoderBlock})
strategy = FSDPStrategy(
auto_wrap_policy=auto_wrap_policy,
activation_checkpointing=EncoderBlock)
fabric = Fabric(accelerator='cuda', devices=4, strategy=strategy)
08_fsdp 與 - 01-2.py 的結(jié)果除了手動(dòng)定義,請(qǐng)也可以使用以下方法,自動(dòng)確定要分割哪些層:fabric = Fabric(accelerator='cuda', devices=4, strategy='fsdp')
在數(shù)據(jù)并行中,mini-batch 需要繼續(xù)被劃分,并且每個(gè) GPU 上都有一份模型副本。由于多個(gè) GPU 并行工作,能夠加快模型訓(xùn)練。- 然后,每個(gè) GPU 被饋送輸入數(shù)據(jù)的不同子集(不同的小批量)。
- 所有 GPU 獨(dú)立地執(zhí)行模型的前向和后向傳遞,計(jì)算各自的局部梯度。
- 然后,收集梯度并對(duì)所有 GPU 進(jìn)行平均。
這種方法的主要優(yōu)點(diǎn)是速度塊。由于每個(gè) GPU 都在與其他 GPU 同時(shí)處理一個(gè)獨(dú)特的小批量數(shù)據(jù),因此可以在更短的時(shí)間內(nèi)在更多數(shù)據(jù)上訓(xùn)練模型。這可以顯著減少訓(xùn)練模型所需的時(shí)間,尤其是在使用大型數(shù)據(jù)集時(shí)。然而,數(shù)據(jù)并行性有一些局限性。每個(gè) GPU 必須具有模型及其參數(shù)的完整副本。這限制了可訓(xùn)練模型的大小,因?yàn)槟P捅仨氝m合單個(gè) GPU 的內(nèi)存 —— 這對(duì)于現(xiàn)代 ViT 或 LLM 來說是不可行的。與數(shù)據(jù)并行不同,張量并行將模型本身劃分為 GPU。在數(shù)據(jù)并行中,每個(gè) GPU 都需要適應(yīng)整個(gè)模型,這在訓(xùn)練更大的模型時(shí)可能會(huì)成為一個(gè)限制。然而,張量并行性允許通過分解模型并將其分布在多個(gè)設(shè)備上來訓(xùn)練對(duì)于單個(gè) GPU 來說可能太大的模型。具體來說,其原理和矩陣乘法相似。按行或按列都可以對(duì)模型進(jìn)行拆解。簡(jiǎn)單起見,以按列拆解為例,可以將一個(gè)大型矩陣乘法運(yùn)算分解為單獨(dú)的計(jì)算,每個(gè)計(jì)算都可以在不同的 GPU 上執(zhí)行,如下圖所示。然后將結(jié)果連接起來以獲得原始結(jié)果,從而有效地分配了計(jì)算負(fù)載。除了上一節(jié)中解釋的 FSDP 策略之外,還可以將優(yōu)化器參數(shù)卸載到 CPU,可以通過將以下代碼strategy = FSDPStrategy(
auto_wrap_policy=auto_wrap_policy,
activation_checkpointing=EncoderBlock,
)
strategy = FSDPStrategy(
auto_wrap_policy=auto_wrap_policy,
activation_checkpointing=EncoderBlock,
cpu_offload=True
)
內(nèi)存消耗從 6.59 GB 減少到 6.03 GB:09_fsdp-cpu-offload-with-01-2.py 的結(jié)果。美中不足的小缺點(diǎn)是運(yùn)行時(shí)間從 5.5 分鐘增加到了 8.3 分鐘。將前面幾招連著打出,就成為了最強(qiáng)的降龍十八掌最后一掌!前幾節(jié)對(duì)優(yōu)化 ViT 進(jìn)行了大量介紹,其實(shí)這些技術(shù)也同樣適用于 LLM。作者在 Lit LLaMA 和 Lit GPT 存儲(chǔ)庫中使用了許多這些技巧,這些存儲(chǔ)庫支持 LLaMA、Falcon、Pythia 和其他流行的模型。盡管如此,為了創(chuàng)建一個(gè)更通用的例子,作者從流行的 HF transformers 庫中微調(diào) LLM,用于對(duì) IMDb 電影評(píng)論的情緒進(jìn)行分類。使用上述技術(shù),僅使用 1.15 Gb 內(nèi)存(bonus_DistilBERT-after.py)而不是 3.99 Gb(bonus_bigbird-before.py)就可以訓(xùn)練 DistilBERT 分類器。更令人印象深刻的是,通過將這些技術(shù)應(yīng)用于 transformers 庫中的 BigBird 模型,BigBird 僅消耗 4.03 GB(bonus_BigBird-after.py)。 strategy = FSDPStrategy(
cpu_offload=True
)
fabric = Fabric(
accelerator='cuda',
devices=4,
strategy=strategy,
precision='bf16-true'
)
with fabric.init_module():
model = AutoModelForSequenceClassification.from_pretrained(
'google/bigbird-roberta-base', num_labels=2)
本文展示了 9 種減少 PyTorch 模型內(nèi)存消耗的技術(shù)。當(dāng)將這些技術(shù)應(yīng)用于 ViT 時(shí),單個(gè) GPU 上減少了 20 倍的內(nèi)存消耗。可以看到,跨 GPU 的張量分片甚至可以降低內(nèi)存消耗。同樣的優(yōu)化還使 BigBird LLM 能夠僅使用 4GB 峰值 GPU RAM 進(jìn)行訓(xùn)練。這些技術(shù)都不是特定于模型的,可以與任何 PyTorch 訓(xùn)練腳本一起使用。使用開源 Fabric 庫,大多數(shù)優(yōu)化都可以通過一行代碼實(shí)現(xiàn)。參考鏈接:https:///pages/community/tutorial/pytorch-memory-vit-llm/