現階段，視覺 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 模型的碳排放。由于最近大模型的發展帶來了更大模型和更長訓練時間的趨勢，這種優化有望幫助研究領域進一步推動最先進的技術，縮短周轉時間并提高生產力。