不堆參數、不靠時長，Meta加速ViT訓練流程，吞吐量4倍提升

現階段，視覺transformer（ViT）模型已經在影像分類、目標偵測與分割等各樣各樣的電腦視覺任務中得到了廣泛應用，並可以在視覺表徵與識別中實現SOTA 結果。由於電腦視覺模型的表現往往與參數量和訓練時間呈正相關，AI 社群已經實驗了越來越大規模的 ViT 模型。

但應該看到，隨著模型開始超出萬億次浮點運算的規模，該領域已經遇到了一些主要的瓶頸。訓練單一模型可能耗費數月，需要數以千計的 GPU，進而增加了加速器需求並導致大規模 ViT 模型將許多從業者「排除在外」。

為了擴展 ViT 模型的使用範圍，Meta AI 的研究者已經開發了更有效率的訓練方法。非常重要的一點是對訓練進行最佳化以實現最佳的加速器利用。但是，這一過程耗時費力且需要大量的專業知識。為了設定有序的實驗，研究者必須從無數可能的最佳化方案中進行選擇：一次訓練過程中執行的百萬次運算中的任何一個都有可能受到低效率的影響和阻礙。

Meta AI 發現，透過將一系列最佳化應用到其影像分類程式碼庫 PyCls 中的 ViT 實現，可以提升運算和儲存效率。 對於使用 PyCIs 訓練的 ViT 模型，Meta AI 的方法可以提升訓練速度和每加速器吞吐量（TFLOPS）。

下圖展示了使用最佳化程式碼庫PyCIs 後每晶片（per chip）加速器吞吐量相較於V100 基準的相對增加，而A100 優化的加速器吞吐量是V100 基準的4.05 倍。

運行原理

#Meta AI 首先對PyCIs 程式碼庫進行分析以確認低訓練效率的潛在來源，最終將注意力放在了數字格式的選擇上。在預設情況下，大多數應用使用 32-bit 單精度浮點格式來表徵神經網路值。轉換至 16-bit 半精度格式（FP16）可以減少模型的記憶體佔用和執行時間，但往往也會降低準確率。

研究者採取了折中方案，即混合精確度。利用它，系統透過單精度格式執行計算以加速訓練並減少記憶體使用，同時透過單精度儲存結果以保持準確率。他們沒有手動地將部分網路轉換至半精度，而是實驗了不同模式的自動混合精度訓練，這樣可以在數位格式之間自動切換。更高級模式的自動混合精度主要依賴半精度運算和模型權重。研究者採用的平衡設定既能大幅加速訓練，同時也不犧牲準確率。

為了讓流程更有效率，研究者充分利用了FairScale 庫中的完全分片資料並行（Fully Sharder Data Parallel, FSDP）訓練演算法，它在GPU 上對參數、梯度和優化器狀態進行分片。透過 FSDP 演算法，研究者可以使用更少的 GPU 來建立更大量級的模型。此外，研究者還使用了 MTA 優化器、一個池化的 ViT 分類器和一個 batch-second 輸入張量佈局來跳過冗餘轉置運算。

下圖 X 軸為可能的最佳化，Y 軸為採用 ViT-H/16 訓練時加速器吞吐量相較於分散式資料並行（DDP）基準的相對增加。

研究者在總patch 大小為560 時實現了1.51 倍的加速器吞吐量提升，以每個加速器晶片上每秒執行的浮點運算數量衡量。透過將影像大小從 224 像素增加至 256 像素，他們可以將吞吐量提升至 1.86 倍。但是，改變影像大小意味著超參數的變化，這會對模型的準確率造成影響。在完全 FP16 模式下訓練時，相對吞吐量增加至 2.18 倍。儘管有時會降低準確率，但在實驗中準確率降低少於 10%。

下圖 Y 軸為 epoch 時間，在整個 ImageNet-1K 資料集上一次訓練的持續時間。這裡專注於現有配置的實際訓練時間，這些配置通常使用 224 像素的影像大小。

#

Meta AI 的研究者使用最佳化方案，將 epoch 時間（在整個 ImageNet-1K 資料集上一次訓練的持續時間）從 0.65 小時減少到 0.43 小時。

下圖 X 軸表示特定配置中 A100 GPU 加速器晶片的數量，Y 軸表示每晶片 TFLOPS 的絕對吞吐量。

該研究也討論了不同 GPU 配置的影響。在每種情況下，系統都實現了比分散式資料並行（DDP）基線水平更高的吞吐量。隨著晶片數量的增加，由於設備間通訊的開銷，我們可以觀察到吞吐量略有下降。然而，即使用 64 塊 GPU，Meta 的系統也比 DDP 基準快 1.83 倍。

新研究的意義

將ViT 訓練中可實現的吞吐量翻倍可以有效讓訓練集群規模翻倍，提高加速器利用率直接減少了AI 模型的碳排放。由於最近大模型的發展帶來了更大模型和更長訓練時間的趨勢，這種優化有望幫助研究領域進一步推動最先進的技術，縮短週轉時間並提高生產力。

以上是不堆參數、不靠時長，Meta加速ViT訓練流程，吞吐量4倍提升的詳細內容。更多資訊請關注PHP中文網其他相關文章！