前言

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

這篇文章源自Vadim Irtlach大佬在kaggle的開源notebook，感謝原作者的分享，本nlp小白覺得受益良多，因此搬運到知乎分享給大家，已取得作者授權，大部分內容是照搬翻譯過來的，小部分內容結合自己的理解進行了補充和修改，不對的地方請大家批評指正，正文開始！

盡管Huggingface開源的Transformers在自然語言處理（NLP）任務中取得了驚人的成功，但由于里面的模型參數數量龐大，即使是使用GPU進行訓練或者部署，也仍具有非常大的挑戰性，因為用如此大的模型進行訓練或推理，會很容易發生顯存不足（OOM）以及訓練時間過長的問題。（這里想吐槽一句的是，kaggle上面的nlp比賽現在動不動就用五折debert-large-v3，沒幾塊V100根本玩不起這種比賽，所以這篇文章對我這種只能用colab的p100來跑實驗的窮學生來說真的是福音啊！）

然而，有很多方法可以避免顯存不足以及訓練時間過長的方法，這篇文章的主要貢獻就是介紹了這些方法的原理以及如何實現，具體包括以下幾種方法：

梯度累積（Gradient Accumulation）

凍結（Freezing）

自動混合精度（Automatic Mixed Precision）

8位優化器（8-bit Optimizers）

梯度檢查點（Gradient Checkpointing）

快速分詞器（Fast Tokenizers）

動態填充（Dynamic Padding）

均勻動態填充（Uniform Dynamic Padding）

其中1-5是神經網絡通用的方法，可以用在任何網絡的性能優化上，6-8是針對nlp領域的性能優化方法。

梯度累積

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

代碼實現

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()

凍結

凍結是一種非常有效的方法，通過取消計算模型某些層中的梯度計算（如embedding層，bert的前幾層），可以大大加快訓練速度并且降低了顯存占用，而且幾乎不會損失模型的性能。

深度學習中的一個眾所周知的事實是，網絡的底層學習輸入數據的通用特征，而網絡頂層學習目標任務特定的高級特征，所以在對預訓練模型進行微調時，一般網絡底層的參數都不怎么需要變，這些都是通用的知識，需要學習的是頂層的那些參數，當使用某種優化算法（如SGD、AdamW或RMSprop）執行優化步驟時，網絡的底層的梯度就都很小，因此參數幾乎保持不變，這也被稱為梯度消失，因此，與其花費大量的時間和算力來計算底層這些“無用”梯度，并對此類梯度很小的參數進行優化，不如直接凍結它們，直接不計算梯度也不進行優化。

PyTorch為關閉梯度計算提供了一個舒適的API，可以通過torch.Tensor的屬性requires_grad設置。

代碼實現

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)

自動混合精度

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

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

代碼實現

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位優化器

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

代碼實現

!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}")

梯度檢查點

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

演示梯度檢查點如何在正向和反向傳播過程中工作

PyTorch框架里也有梯度檢查點的實現，通過這兩個函數：torch.utils.checkpoint.checkpoint和torch.utils.checkpoint.checkpoint_sequential

這邊引用一段torch官網對梯度檢查點的介紹：

梯度檢查點通過用計算換取內存來工作。檢查點部分不是存儲整個計算圖的所有中間激活以進行反向計算，而是不保存中間激活，而是在反向過程中重新計算它們。它可以應用于模型的任何部分。

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

此外，HuggingFace Transformers也支持梯度檢查點。梯度檢查點可以通過PreTrainedModel實例的gradient_checkpointing_enable方法執行，一行代碼直接搞定！

代碼實現

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提供兩種類型的分詞器：基本分詞器和快速分詞器。它們之間的主要區別在于，快速分詞器是在Rust上編寫的，因為Python在循環中非常慢，但在分詞的時候又要用到循環。快速分詞器是一種非常簡單的方法，允許我們在分詞的時候獲得額外的加速。要使用快速分詞器也很簡單，只要把transformers.AutoTokenizer里面的from_pretrained方法的use_fast的值修改為True就可以了。

分詞器是如何工作的

代碼實現

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}")

動態填充

通常來說，模型是用批量數據輸入訓練的，批中的每個輸入必須具有固定大小，即一批量的數據必須是矩陣的表示，所有批量數據的尺寸都一樣。固定尺寸通常是根據數據集中的長度分布、特征數量和其他因素來選擇的。在NLP任務中，輸入大小稱為文本長度，或者最大長度（max length）。然而，不同的文本具有不同的長度，為了處理這種情況，研究人員提出了填充標記和截斷。當最大長度小于輸入文本的長度時，會使用截斷，因此會刪除一些標記。當輸入文本的長度小于最大長度時，會將填充標記，比如[PAD]，添加到輸入文本的末尾，值得注意的是，填充標記不應包含在某些任務的損失計算中（例如掩蔽語言建模或命名實體識別）

固定長度填充

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

動態填充

均勻動態填充

還有一種基于動態填充的方法，叫做均勻動態填充。其思想是在分batch時，先按文本的長度對文本進行排序，這樣同一個batch里面的文本長度就都差不多。這種方法非常有效，在訓練或推理期間的計算量都比動態填充要來的少。但是，不建議在訓練期間使用均勻動態填充，因為訓練時數據最好是要shuffer的，但是推理時如果一次性要推理很多文本的話可以考慮這么做

均勻動態填充

總結

即使在現代GPU上，優化內存和時間也是開發模型的必要步驟，因此，本文介紹了加速訓練和減少transformers等大型模型內存消耗的最強大、最流行的方法。

審核編輯：劉清