如何顯著提升Vision Transformer的訓(xùn)練效率

近期MetaAI發(fā)布了一篇博客，關(guān)于如何顯著提升Vision Transformer的訓(xùn)練效率。

原文：[Significantly faster Vision Transformer training]

鏈接：https://ai.facebook.com/blog/significantly-faster-vision-transformer-training

What the research is

Vision Transformer模型幾乎火遍計(jì)算機(jī)視覺(jué)各個(gè)領(lǐng)域，其性能隨著參數(shù)增加和更久的訓(xùn)練過(guò)程而得到提升。隨著模型越來(lái)越大，超過(guò)了萬(wàn)億次浮點(diǎn)運(yùn)算的規(guī)模，該領(lǐng)域達(dá)到了瓶頸：訓(xùn)練一個(gè)模型往往要一個(gè)月，需要幾百上千個(gè)GPU，導(dǎo)致大部分人無(wú)法接觸到大規(guī)模ViT模型，并進(jìn)而增加了對(duì)加速器的需求。

為了降低門(mén)檻，讓更多人能夠應(yīng)用ViT，我們開(kāi)發(fā)一系列方法來(lái)加速整個(gè)訓(xùn)練。我們基于MetaAI的圖像分類(lèi)模型庫(kù)PyCls實(shí)現(xiàn)了一系列優(yōu)化，這些優(yōu)化極大的提升了模型訓(xùn)練過(guò)程的吞吐量：

How it works ？

我們首先對(duì)代碼庫(kù)進(jìn)行分析，以定位訓(xùn)練效率低下的原因，最后關(guān)注點(diǎn)落在計(jì)算類(lèi)型上：大部分模型都是用FP32進(jìn)行訓(xùn)練，如果使用FP16訓(xùn)練的話，可以降低顯存占用，并提高模型訓(xùn)練速度，但這一做法經(jīng)常會(huì)導(dǎo)致準(zhǔn)確率下降

所以我們選了一個(gè)折中的方法：自動(dòng)混合精度。在該方法下，我們用half類(lèi)型進(jìn)行計(jì)算，以加快訓(xùn)練，減少顯存使用。并以fp32類(lèi)型存儲(chǔ)參數(shù)，以保證模型準(zhǔn)確率。其中我們沒(méi)有手動(dòng)將網(wǎng)絡(luò)各部分轉(zhuǎn)換成half類(lèi)型，而是應(yīng)用AMP各種模式（如O1, O2, O3)，以尋找一個(gè)既能提升速度又不影響精度的平衡點(diǎn)。

FSDP

為了讓訓(xùn)練更加高效，我們應(yīng)用了FSDP訓(xùn)練策略，他能夠?qū)?shù)，梯度，優(yōu)化器狀態(tài)分片到各GPU上。在FSDP的幫助下，我們可以用更少的GPU資源構(gòu)建更大的模型。

FSDP策略可以參考 [數(shù)據(jù)并行Deep-dive: 從DP 到 Fully Sharded Data Parallel （FSDP）完全分片數(shù)據(jù)并行] 鏈接：https://zhuanlan.zhihu.com/p/485208899

MTA Optimizer

前向計(jì)算完畢后，優(yōu)化器需要對(duì)各個(gè)參數(shù)進(jìn)行修改。而當(dāng)參數(shù)比較多的情況下，對(duì)應(yīng)啟動(dòng)的Optimizer Kernel就會(huì)變得很多，通常這些Kernel都比較小，計(jì)算負(fù)擔(dān)不大，啟動(dòng)Kernel的開(kāi)銷(xiāo)反而占了大頭。

在ContiguousParams中，它將模型參數(shù)放置到一塊連續(xù)的顯存中進(jìn)行計(jì)算，這樣就能減少優(yōu)化器這部分的時(shí)間。下圖是Resnet50+SGD是否應(yīng)用ContiguousParams的比較，可以看到OptimizerStep這部分時(shí)間顯著減少了。

而NVIDIA的Apex庫(kù)的做法則是在底層重新實(shí)現(xiàn)了一系列MultiTensorOptimizer，如Adam, Adagrad等等。

Apex這種方法比較硬核，普通用戶(hù)如果想要自己自定義優(yōu)化器并應(yīng)用Multi Tensor的優(yōu)化，就必須改動(dòng)底層CUDA代碼。而最近PyTorch也在計(jì)劃提供了一系列foreach接口[Replace optimizers in torch.optim with the ones from torch.optim._multi_tensor] 鏈接：https://github.com/pytorch/pytorch/pull/49039，讓用戶(hù)只需要在Python層即可享受到優(yōu)化，對(duì)應(yīng)的MultiTensor版Momentum優(yōu)化器代碼如下所示：

torch._foreach_mul_(bufs,momentum)
torch._foreach_add_(bufs,grads,alpha=1-dampening)

Pooled Classifier

原版的ViT是額外加了一個(gè)分類(lèi)token，來(lái)輸出最后的分類(lèi)結(jié)果。而這里采用平均池化 如：https://github.com/facebookresearch/pycls/blob/main/pycls/core/config.py#L205 處理最后的分類(lèi)

Batch Second Input Tensor Layout

這里的數(shù)據(jù)格式與以往不同，將batch維度放在第二維，并在調(diào)用nn.MultiheadAttention的時(shí)候，設(shè)置batch_first=False，以減少不必要的轉(zhuǎn)置

ifself.batch_firstandis_batched:
returnattn_output.transpose(1,0),attn_output_weights
else:
returnattn_output,attn_output_weights

總感覺(jué)這個(gè)實(shí)現(xiàn)怪怪的

其他優(yōu)化

我們?cè)诓扇?60大小的batchsize下，達(dá)到了1.51倍的加速比，進(jìn)一步的我們將batchsize設(shè)置為384，并將圖片大小增大到256，達(dá)到了1.86倍加速比。在全FP16運(yùn)算下，能夠達(dá)到2.18倍加速比，盡管這偶爾會(huì)降低準(zhǔn)確率（在實(shí)驗(yàn)中，準(zhǔn)確率降低不到10%）。

使用上述優(yōu)化，我們將Imagenet1K數(shù)據(jù)集每epoch訓(xùn)練時(shí)間從0.65小時(shí)降低到0.43小時(shí)

我們還研究了不同GPU配置對(duì)訓(xùn)練速度的影響，在不同配置下我們都實(shí)現(xiàn)了比DDP baseline更高的吞吐量。隨著GPU增加，吞吐量會(huì)因?yàn)樵O(shè)備之間的通信開(kāi)銷(xiāo)略微下降。然而即使在64塊GPU下，我們?nèi)匀槐菵DP基線快1.83倍

文中鏈接

PyCls ：https://github.com/facebookresearch/pycls

ContiguousParams：https://github.com/PhilJd/contiguous_pytorch_params

Adam：https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_adam.cu

審核編輯 ：李倩