加速視覺語言模型：BridgeTower 在 Habana Gaudi2 上

BridgeTower 模型概覽

BridgeTower 引入多個橋接層，將單模態編碼器的頂層連接到跨模態編碼器的每一層，從而在不同語義層級實現有效的自底向上跨模態對齊與融合。僅使用來自 Conceptual Captions、SBU Captions、MSCOCO Captions 和 Visual Genome 的 4M 張圖像進行預訓練，它在視覺語言任務上達到最先進的性能，例如在 VQAv2 test‑std 上達到 78.73% 的準確率，相較於 METER 高出 1.09%，且額外參數與計算開銷可忽略不計。

硬體平台

Habana Gaudi2 提供八個 HPU，每個擁有 96 GB 記憶體，並透過 Optimum Habana 支援，以便輕鬆移植 Transformers 腳本。Nvidia H100 是最新的 GPU 世代，配備專用的 Transformer Engine 以支援 fp8 混合精度，並擁有 80 GB 記憶體。Nvidia A100（80 GB 變體）使用第三代 Tensor Core 技術，仍是雲端供應商中廣泛可用的最快 GPU。

基準設置

我們在 New Yorker Caption Contest 資料集上微調了 BridgeTower Large 檢查點（866 M 參數），在所有加速器上保持超參數相同，並以每秒樣本數衡量訓練吞吐量。

dataloader_num_workers 的影響

增加用於資料載入的子行程數量可提升所有設備的吞吐量。在每個設備的批次大小為 48 時，我們觀察到：

  • Gaudi2 HPU: 601.5 samples/s (workers=0), 747.4 samples/s (workers=1), 768.7 samples/s (workers=2)。
  • H100 GPU: 336.5 samples/s (workers=0), 580.1 samples/s (workers=1), 602.1 samples/s (workers=2)。
  • A100 GPU: 227.5 samples/s (workers=0), 339.7 samples/s (workers=1), 345.4 samples/s (workers=2)。 這些結果顯示，在使用兩個工作程序時，Gaudi2 比 H100 快 x1.28，比 A100 快 x2.23；而增加工作程序數可在 Gaudi2 上帶來 x1.28 的加速，在 H100 上帶來 x1.79 的加速，在 A100 上帶來 x1.52 的加速。

使用 Optimum Habana 進行硬體加速資料載入的影響

透過將圖像解碼與增強操作移至加速器，使用 --mediapipe_dataloader 標誌（僅在 Gaudi2 上可用）可進一步提升吞吐量。在使用兩個資料載入工作程序的相同實驗中得到：

  • Gaudi2 HPU: 847.7 samples/s (workers=2 + mediapipe_dataloader)。
  • H100 與 A100：不適用（該標誌在此些 GPU 上無法運作）。 相比基準運行（workers=0），這在 Gaudi2 上代表 x1.41 的加速，使得在使用兩個工作程序及媒體管線時，Gaudi2 比 H100 快 x1.41，比 A100 快 x2.45。

重現此基準

要重現結果，請透過 Intel Developer Cloud 存取 Gaudi2，安裝最新的 Optimum Habana，克隆儲存庫，安裝依賴項，並使用適當的參數執行提供的 run_bridgetower.py 腳本。在 Gaudi2 上，基礎命令包含 --use_habana --use_lazy_mode --use_hpu_graphs_for_inference --gaudi_config_name Habana/clip。加入 --dataloader_num_workers N 和 --mediapipe_dataloader 可測試其他配置。對於 A100 與 H100，請將 Habana 特有的類別替換為標準 Transformers 的 Trainer 和 TrainingArguments，移除 Gaudi 設定參考，並從 Transformers 匯入 set_seed。

結論

在訓練視覺語言模型時，兩種簡單技巧——增加更多資料載入工作程序以及將圖像解碼/增強卸載到加速器——能帶來顯著的加速效果。使用 Optimum Habana 在 Habana Gaudi2 上微調 BridgeTower 時，其速度約為 Nvidia H100 的 1.4 倍、Nvidia A100 80GB 的 2.5 倍，僅需少數額外的訓練參數。

Sources