加速视觉-语言模型:BridgeTower 在 Habana Gaudi2 上

BridgeTower 模型概述

BridgeTower 引入多个桥接层,将单模态编码器的顶层连接到跨模态编码器的每一层,从而在不同语义层次上实现有效的自底向上跨模态对齐和融合。仅在 Conceptual Captions、SBU Captions、MSCOCO Captions 和 Visual Genome 的 400 万张图像上进行预训练,它在视觉-语言任务上实现了最先进的性能,例如在 VQAv2 test‑std 上达到 78.73% 的准确率,相比 METER 提高了 1.09%,且额外参数和计算开销可忽略不计。

硬件平台

Habana Gaudi2 提供八个 HPU,每个具有 96 GB 内存,并得到 Optimum Habana 的支持,以便轻松移植 Transformers 脚本。Nvidia H100 是最新的 GPU 一代,配备了用于 fp8 混合精度的专用 Transformer Engine,内存为 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 样本/秒 (workers=0),747.4 样本/秒 (workers=1),768.7 样本/秒 (workers=2)。
  • H100 GPU:336.5 样本/秒 (workers=0),580.1 样本/秒 (workers=1),602.1 样本/秒 (workers=2)。
  • A100 GPU:227.5 样本/秒 (workers=0),339.7 样本/秒 (workers=1),345.4 样本/秒 (workers=2)。 这些结果表明,在使用两个工作进程时,Gaudi2 比 H100 快 x1.28,比 A100 快 x2.23;增加工作进程数可在 Gaudi2 上获得 x1.28 的加速,在 H100 上获得 x1.79 的加速,在 A100 上获得 x1.52 的加速。

使用 Optimum Habana 进行硬件加速数据加载的影响

通过将图像解码和增强转移到加速器(使用仅在 Gaudi2 上可用的 --mediapipe_dataloader 标志)可以进一步提高吞吐量。在使用两个 dataloader 工作进程的相同实验中得到:

  • Gaudi2 HPU:847.7 样本/秒 (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 的 TrainerTrainingArguments,移除 Gaudi 配置引用,并从 Transformers 导入 set_seed

结论

在训练视觉-语言模型时,两种简单技术——增加 dataloader 工作进程数以及将图像解码/增强卸载到加速器——可以带来显著的加速。在 Habana Gaudi2 上使用 Optimum Habana 进行 BridgeTower 微调时,大约可以达到 Nvidia H100 速度的 1.4×、Nvidia A100 80GB 速度的 2.5×,只需添加少量额外的训练参数。

Sources