如何在GPU資源受限的情況下訓(xùn)練transformers庫(kù)上面的大模型

前言

自BERT出現(xiàn)以來(lái)，nlp領(lǐng)域已經(jīng)進(jìn)入了大模型的時(shí)代，大模型雖然效果好，但是畢竟不是人人都有著豐富的GPU資源，在訓(xùn)練時(shí)往往就捉襟見(jiàn)肘，出現(xiàn)顯存out of memory的問(wèn)題，或者訓(xùn)練時(shí)間非常非常的久，因此，這篇文章主要解決的問(wèn)題就是如何在GPU資源受限的情況下訓(xùn)練transformers庫(kù)上面的大模型。

這篇文章源自Vadim Irtlach大佬在kaggle的開(kāi)源notebook，感謝原作者的分享，本nlp小白覺(jué)得受益良多，因此搬運(yùn)到知乎分享給大家，已取得作者授權(quán)，大部分內(nèi)容是照搬翻譯過(guò)來(lái)的，小部分內(nèi)容結(jié)合自己的理解進(jìn)行了補(bǔ)充和修改，不對(duì)的地方請(qǐng)大家批評(píng)指正，正文開(kāi)始！

盡管Huggingface開(kāi)源的Transformers在自然語(yǔ)言處理（NLP）任務(wù)中取得了驚人的成功，但由于里面的模型參數(shù)數(shù)量龐大，即使是使用GPU進(jìn)行訓(xùn)練或者部署，也仍具有非常大的挑戰(zhàn)性，因?yàn)橛萌绱舜蟮哪Ｐ瓦M(jìn)行訓(xùn)練或推理，會(huì)很容易發(fā)生顯存不足（OOM）以及訓(xùn)練時(shí)間過(guò)長(zhǎng)的問(wèn)題。（這里想吐槽一句的是，kaggle上面的nlp比賽現(xiàn)在動(dòng)不動(dòng)就用五折debert-large-v3，沒(méi)幾塊V100根本玩不起這種比賽，所以這篇文章對(duì)我這種只能用colab的p100來(lái)跑實(shí)驗(yàn)的窮學(xué)生來(lái)說(shuō)真的是福音?。。?/p>

然而，有很多方法可以避免顯存不足以及訓(xùn)練時(shí)間過(guò)長(zhǎng)的方法，這篇文章的主要貢獻(xiàn)就是介紹了這些方法的原理以及如何實(shí)現(xiàn)，具體包括以下幾種方法：

梯度累積（Gradient Accumulation）

凍結(jié)（Freezing）

自動(dòng)混合精度（Automatic Mixed Precision）

8位優(yōu)化器（8-bit Optimizers）

梯度檢查點(diǎn)（Gradient Checkpointing）

快速分詞器（Fast Tokenizers）

動(dòng)態(tài)填充（Dynamic Padding）

均勻動(dòng)態(tài)填充（Uniform Dynamic Padding）

其中1-5是神經(jīng)網(wǎng)絡(luò)通用的方法，可以用在任何網(wǎng)絡(luò)的性能優(yōu)化上，6-8是針對(duì)nlp領(lǐng)域的性能優(yōu)化方法。

梯度累積

梯度累積背后的想法非常簡(jiǎn)單，就是為了模擬更大的批量（batch）。有時(shí)，為了更好地收斂或提高性能，需要使用大批量進(jìn)行訓(xùn)練，但是，這通常需要更大的顯存。這個(gè)問(wèn)題的一種可能的解決方案是使用較小的批量，但是，一方面，小批量訓(xùn)練會(huì)增加訓(xùn)練和推理時(shí)間，另一方面，梯度下降算法對(duì)批量大小的選擇非常敏感，小批量可能會(huì)導(dǎo)致不穩(wěn)定的收斂和性能降低。所以，我們可以先執(zhí)行幾次前向傳播和反向傳播，使得梯度進(jìn)行累積，當(dāng)我們有足夠的計(jì)算梯度時(shí)，再對(duì)參數(shù)進(jìn)行優(yōu)化，從而利用小顯存，模擬大批量的效果，并且訓(xùn)練時(shí)間也不會(huì)大幅增加。

代碼實(shí)現(xiàn)

steps=len(loader)

##performvalidationloopeach`validation_steps`trainingsteps!
validation_steps=int(validation_steps*gradient_accumulation_steps)

forstep,batchinenumerate(loader,1):

#prepareinputsandtargetsforthemodelandlossfunctionrespectively.

#forwardpass
outputs=model(inputs)

#computingloss
loss=loss_fn(outputs,targets)

#accumulatinggradientsoversteps
ifgradient_accumulation_steps>1:
loss=loss/gradient_accumulation_steps

#backwardpass
loss.backward()

#performoptimizationstepaftercertainnumberofaccumulatingstepsandattheendofepoch
ifstep%gradient_accumulation_steps==0orstep==steps:
torch.nn.utils.clip_grad_norm_(model.parameters(),max_norm)
optimizer.step()
model.zero_grad()

#performvalidationloop
ifstep%validation_steps==0:
validation_loop()

凍結(jié)

凍結(jié)是一種非常有效的方法，通過(guò)取消計(jì)算模型某些層中的梯度計(jì)算（如embedding層，bert的前幾層），可以大大加快訓(xùn)練速度并且降低了顯存占用，而且?guī)缀醪粫?huì)損失模型的性能。

深度學(xué)習(xí)中的一個(gè)眾所周知的事實(shí)是，網(wǎng)絡(luò)的底層學(xué)習(xí)輸入數(shù)據(jù)的通用特征，而網(wǎng)絡(luò)頂層學(xué)習(xí)目標(biāo)任務(wù)特定的高級(jí)特征，所以在對(duì)預(yù)訓(xùn)練模型進(jìn)行微調(diào)時(shí)，一般網(wǎng)絡(luò)底層的參數(shù)都不怎么需要變，這些都是通用的知識(shí)，需要學(xué)習(xí)的是頂層的那些參數(shù)，當(dāng)使用某種優(yōu)化算法（如SGD、AdamW或RMSprop）執(zhí)行優(yōu)化步驟時(shí)，網(wǎng)絡(luò)的底層的梯度就都很小，因此參數(shù)幾乎保持不變，這也被稱為梯度消失，因此，與其花費(fèi)大量的時(shí)間和算力來(lái)計(jì)算底層這些“無(wú)用”梯度，并對(duì)此類梯度很小的參數(shù)進(jìn)行優(yōu)化，不如直接凍結(jié)它們，直接不計(jì)算梯度也不進(jìn)行優(yōu)化。

PyTorch為關(guān)閉梯度計(jì)算提供了一個(gè)舒適的API，可以通過(guò)torch.Tensor的屬性requires_grad設(shè)置。

代碼實(shí)現(xiàn)

deffreeze(module):
"""
Freezesmodule'sparameters.
"""
forparameterinmodule.parameters():
parameter.requires_grad=False

defget_freezed_parameters(module):
"""
Returnsnamesoffreezedparametersofthegivenmodule.
"""

freezed_parameters=[]
forname,parameterinmodule.named_parameters():
ifnotparameter.requires_grad:
freezed_parameters.append(name)

returnfreezed_parameters

importtorch
fromtransformersimportAutoConfig,AutoModel


##initializingmodel
model_path="microsoft/deberta-v3-base"
config=AutoConfig.from_pretrained(model_path)
model=AutoModel.from_pretrained(model_path,config=config)


##freezingembeddingsandfirst2layersofencoder
freeze(model.embeddings)
freeze(model.encoder.layer[:2])

freezed_parameters=get_freezed_parameters(model)
print(f"Freezedparameters:{freezed_parameters}")

##selectingparameters,whichrequiresgradientsandinitializingoptimizer
model_parameters=filter(lambdaparameter:parameter.requires_grad,model.parameters())
optimizer=torch.optim.AdamW(params=model_parameters,lr=2e-5,weight_decay=0.0)

自動(dòng)混合精度

自動(dòng)混合精度（AMP）是另一種在不損失最終質(zhì)量的情況下減少顯存消耗和訓(xùn)練時(shí)間的方法，該方法由NVIDIA和百度研究人員在2017年的Mixed Precision Training論文中提出。該方法背后的關(guān)鍵思想是使用較低的精度將模型的梯度和參數(shù)保留在內(nèi)存中，即不使用全精度（float32），而是使用半精度（例如float16）將張量保存在內(nèi)存中。然而，當(dāng)以較低精度計(jì)算梯度時(shí)，某些值可能太小，以至于被視為零，這種現(xiàn)象被稱為“溢出”。為了防止“溢出”，原始論文的作者提出了一種梯度縮放方法。

PyTorch從1.6的版本開(kāi)始提供了一個(gè)包：torch.cuda.amp，具有使用自動(dòng)混合精度所需的功能（從降低精度到梯度縮放），自動(dòng)混合精度作為上下文管理器實(shí)現(xiàn)，因此可以隨時(shí)隨地的插入到訓(xùn)練和推理腳本中。

代碼實(shí)現(xiàn)

fromtorch.cuda.ampimportautocast,GradScaler


scaler=GradScaler()

forstep,batchinenumerate(loader,1):

#prepareinputsandtargetsforthemodelandlossfunctionrespectively.

#forwardpasswith`autocast`contextmanager
withautocast(enabled=True):
outputs=model(inputs)

#computingloss
loss=loss_fn(outputs,targets)

#scalegradintandperformbackwardpass
scaler.scale(loss).backward()

#beforegradientclippingtheoptimizerparametersmustbeunscaled.
scaler.unscale_(optimizer)

#performoptimizationstep
torch.nn.utils.clip_grad_norm_(model.parameters(),max_norm)

scaler.step(optimizer)
scaler.update()

8位優(yōu)化器

8-bit Optimizers的思想類似于自動(dòng)混合精度（模型的參數(shù)和梯度使用較低的精度保存），但8-bit Optimizers還讓優(yōu)化器的狀態(tài)使用低精度保存。作者（Meta Research）在最初的論文8-bit Optimizers via Block-wise Quantization中詳細(xì)介紹了8-bit Optimizers，表明8-bit Optimizers顯著降低了顯存占用，略微加快了訓(xùn)練速度。此外，作者研究了不同超參數(shù)設(shè)置的影響，表明8-bit Optimizers對(duì)不同的學(xué)習(xí)率、beta和權(quán)重衰減參數(shù)的效果是穩(wěn)定的，不會(huì)降低性能或影響收斂性。因此，作者為8位優(yōu)化器提供了一個(gè)高級(jí)庫(kù)，叫做bitsandbytes。

代碼實(shí)現(xiàn)

!pipinstall-qbitsandbytes-cuda110

defset_embedding_parameters_bits(embeddings_path,optim_bits=32):
"""
https://github.com/huggingface/transformers/issues/14819#issuecomment-1003427930
"""

embedding_types=("word","position","token_type")
forembedding_typeinembedding_types:
attr_name=f"{embedding_type}_embeddings"

ifhasattr(embeddings_path,attr_name):
bnb.optim.GlobalOptimManager.get_instance().register_module_override(
getattr(embeddings_path,attr_name),'weight',{'optim_bits':optim_bits}
)

importbitsandbytesasbnb


##selectingparameters,whichrequiresgradients
model_parameters=filter(lambdaparameter:parameter.requires_grad,model.parameters())

##initializingoptimizer
bnb_optimizer=bnb.optim.AdamW(params=model_parameters,lr=2e-5,weight_decay=0.0,optim_bits=8)
##bnb_optimizer=bnb.optim.AdamW8bit(params=model_parameters,lr=2e-5,weight_decay=0.0)#equivalenttotheaboveline

##settingembeddingsparameters
set_embedding_parameters_bits(embeddings_path=model.embeddings)

print(f"8-bitOptimizer:

{bnb_optimizer}")

梯度檢查點(diǎn)

有時(shí)候，即使用了上面的幾種方法，顯存可能還是不夠，尤其是在模型足夠大的情況下。那么梯度檢查點(diǎn)（Gradient Checkpointing）就是壓箱底的招數(shù)了，這個(gè)方法第一次在 Training Deep Nets With Sublinear Memory Cost ，作者表明梯度檢查點(diǎn)可以顯著降低顯存利用率，從降低到，其中n是模型的層數(shù)。這種方法允許在單個(gè)GPU上訓(xùn)練大型模型，或者提供更多內(nèi)存以增加批量大小，從而更好更快地收斂。梯度檢查點(diǎn)背后的思想是在小數(shù)據(jù)塊中計(jì)算梯度，同時(shí)在正向和反向傳播過(guò)程中從內(nèi)存中移除不必要的梯度，從而降低內(nèi)存利用率，但是這種方法需要更多的計(jì)算步驟來(lái)再現(xiàn)整個(gè)反向傳播圖，其實(shí)就是一種用時(shí)間來(lái)?yè)Q空間的方法。

演示梯度檢查點(diǎn)如何在正向和反向傳播過(guò)程中工作

PyTorch框架里也有梯度檢查點(diǎn)的實(shí)現(xiàn)，通過(guò)這兩個(gè)函數(shù)：torch.utils.checkpoint.checkpoint和torch.utils.checkpoint.checkpoint_sequential

這邊引用一段torch官網(wǎng)對(duì)梯度檢查點(diǎn)的介紹：

梯度檢查點(diǎn)通過(guò)用計(jì)算換取內(nèi)存來(lái)工作。檢查點(diǎn)部分不是存儲(chǔ)整個(gè)計(jì)算圖的所有中間激活以進(jìn)行反向計(jì)算，而是不保存中間激活，而是在反向過(guò)程中重新計(jì)算它們。它可以應(yīng)用于模型的任何部分。

具體而言，在前向傳播中，該函數(shù)將以torch.no_grad()的方式運(yùn)行，即不存儲(chǔ)中間激活。然而，前向傳播保存了輸入元組和函數(shù)參數(shù)。在反向傳播時(shí)，檢索保存的輸入和函數(shù)，然后再次對(duì)函數(shù)進(jìn)行前向傳播，現(xiàn)在跟蹤中間激活，然后使用這些激活值計(jì)算梯度。

此外，HuggingFace Transformers也支持梯度檢查點(diǎn)。梯度檢查點(diǎn)可以通過(guò)PreTrainedModel實(shí)例的gradient_checkpointing_enable方法執(zhí)行，一行代碼直接搞定！

代碼實(shí)現(xiàn)

fromtransformersimportAutoConfig,AutoModel
##https://github.com/huggingface/transformers/issues/9919
fromtorch.utils.checkpointimportcheckpoint

##initializingmodel
model_path="microsoft/deberta-v3-base"
config=AutoConfig.from_pretrained(model_path)
model=AutoModel.from_pretrained(model_path,config=config)

##gradientcheckpointing
model.gradient_checkpointing_enable()
print(f"GradientCheckpointing:{model.is_gradient_checkpointing}")

快速分詞器

HuggingFace Transformers提供兩種類型的分詞器：基本分詞器和快速分詞器。它們之間的主要區(qū)別在于，快速分詞器是在Rust上編寫(xiě)的，因?yàn)?a href="http://ttokpm.com/tags/python/" target="_blank">Python在循環(huán)中非常慢，但在分詞的時(shí)候又要用到循環(huán)?？焖俜衷~器是一種非常簡(jiǎn)單的方法，允許我們?cè)诜衷~的時(shí)候獲得額外的加速。要使用快速分詞器也很簡(jiǎn)單，只要把transformers.AutoTokenizer里面的from_pretrained方法的use_fast的值修改為True就可以了。

分詞器是如何工作的

代碼實(shí)現(xiàn)

fromtransformersimportAutoTokenizer

##initializingBaseversionofTokenizer
model_path="microsoft/deberta-v3-base"
tokenizer=AutoTokenizer.from_pretrained(model_path,use_fast=False)
print(f"BaseversionTokenizer:

{tokenizer}",end="
"*3)

##initializingFastversionofTokenizer
fast_tokenizer=AutoTokenizer.from_pretrained(model_path,use_fast=True)
print(f"FastversionTokenizer:

{fast_tokenizer}")

動(dòng)態(tài)填充

通常來(lái)說(shuō)，模型是用批量數(shù)據(jù)輸入訓(xùn)練的，批中的每個(gè)輸入必須具有固定大小，即一批量的數(shù)據(jù)必須是矩陣的表示，所有批量數(shù)據(jù)的尺寸都一樣。固定尺寸通常是根據(jù)數(shù)據(jù)集中的長(zhǎng)度分布、特征數(shù)量和其他因素來(lái)選擇的。在NLP任務(wù)中，輸入大小稱為文本長(zhǎng)度，或者最大長(zhǎng)度（max length）。然而，不同的文本具有不同的長(zhǎng)度，為了處理這種情況，研究人員提出了填充標(biāo)記和截?cái)?。?dāng)最大長(zhǎng)度小于輸入文本的長(zhǎng)度時(shí)，會(huì)使用截?cái)?，因此?huì)刪除一些標(biāo)記。當(dāng)輸入文本的長(zhǎng)度小于最大長(zhǎng)度時(shí)，會(huì)將填充標(biāo)記，比如[PAD]，添加到輸入文本的末尾，值得注意的是，填充標(biāo)記不應(yīng)包含在某些任務(wù)的損失計(jì)算中（例如掩蔽語(yǔ)言建?；蛎麑?shí)體識(shí)別）

固定長(zhǎng)度填充

然而，填充標(biāo)記有明顯的缺點(diǎn)。比如在輸入文本相對(duì)于選定的最大長(zhǎng)度非常短的情況下，效率就很低，需要更多的額外內(nèi)存，比如我有一條文本長(zhǎng)度512，然后其他文本長(zhǎng)度都在10左右，那么如果將max seq設(shè)置為512，就會(huì)導(dǎo)致很多無(wú)效計(jì)算。為了防止額外的計(jì)算操作，研究人員提出了一種非常有效的方法，就是將批量的輸入填充到這一批量的最大輸入長(zhǎng)度，如下圖所示，這種方法可以將訓(xùn)練速度提高35%甚至50%，當(dāng)然這種方法加速的效果取決于批量的大小以及文本長(zhǎng)度的分布，批量越小，加速效果越明顯，文本長(zhǎng)度分布越不均，加速效果也越好。

動(dòng)態(tài)填充

均勻動(dòng)態(tài)填充

還有一種基于動(dòng)態(tài)填充的方法，叫做均勻動(dòng)態(tài)填充。其思想是在分batch時(shí)，先按文本的長(zhǎng)度對(duì)文本進(jìn)行排序，這樣同一個(gè)batch里面的文本長(zhǎng)度就都差不多。這種方法非常有效，在訓(xùn)練或推理期間的計(jì)算量都比動(dòng)態(tài)填充要來(lái)的少。但是，不建議在訓(xùn)練期間使用均勻動(dòng)態(tài)填充，因?yàn)橛?xùn)練時(shí)數(shù)據(jù)最好是要shuffer的，但是推理時(shí)如果一次性要推理很多文本的話可以考慮這么做

均勻動(dòng)態(tài)填充

總結(jié)

即使在現(xiàn)代GPU上，優(yōu)化內(nèi)存和時(shí)間也是開(kāi)發(fā)模型的必要步驟，因此，本文介紹了加速訓(xùn)練和減少transformers等大型模型內(nèi)存消耗的最強(qiáng)大、最流行的方法。

審核編輯：劉清