[论文解读] Training a Large Video Model on a Single Machine in a Day

Yue Zhao, Philipp Krähenbühl|arXiv (Cornell University)|Sep 28, 2023

Human Pose and Action Recognition被引用 5

一句话总结

该论文提出了一种高度高效的训练流水线，仅使用一台配备八张消费级A5000 GPU的单机设备，即可在18小时内训练出当前最先进（SOTA）的视频Transformer模型。通过优化I/O、CPU和GPU计算，采用FlashAttention、融合的解码-裁剪操作以及基于数据块的视频加载方式，该方法相比之前的工作，将内存使用量减少6.7倍，GPU小时数减少11.8倍，硬件成本降低15倍。

ABSTRACT

Videos are big, complex to pre-process, and slow to train on. State-of-the-art large-scale video models are trained on clusters of 32 or more GPUs for several days. As a consequence, academia largely ceded the training of large video models to industry. In this paper, we show how to still train a state-of-the-art video model on a single machine with eight consumer-grade GPUs in a day. We identify three bottlenecks, IO, CPU, and GPU computation, and optimize each. The result is a highly efficient video training pipeline. For comparable architectures, our pipeline achieves higher accuracies with $\frac{1}{8}$ of the computation compared to prior work. Code is available at https://github.com/zhaoyue-zephyrus/AVION.

研究动机与目标

• 解决大规模视频模型训练带来的高计算与硬件成本问题，从而提升学术界对这类训练的可及性。
• 识别并解决视频训练中的三大主要瓶颈：I/O、CPU预处理和GPU计算。
• 实现在单台8-GPU工作站上，于一天内完成当前最先进对比性视频-语言模型的训练。
• 在显著减少计算与内存使用量的前提下，实现与之前模型相当或更优的准确率。
• 通过降低硬件需求与成本，使大规模视频模型训练对学术界和小型研究实验室更加可及。

提出的方法

• 采用FlashAttention将自注意力机制的内存占用从O(N²)降低至O(N)，其中N为视频token的数量，从而支持更大的批量大小。
• 设计一种融合的解码-裁剪操作，将RandomResizedCrop集成至视频解码阶段，最大限度减少重复解码与CPU开销。
• 实现基于数据块的视频加载系统，将长视频分割为固定长度的压缩片段，并行处理。
• 将所有数据增强操作（如裁剪、翻转）移至GPU执行，以利用其并行能力，缓解CPU瓶颈。
• 对整个流水线进行端到端优化，最小化数据移动，最大化各阶段的GPU利用率。
• 使用单台8-GPU A5000服务器，将模型总批量大小提升至2,048，而此前工作需32块以上40GB的A100 GPU。

(a) Decode-then-crop.

实验结果

研究问题

• RQ1 能否在单台机器上，使用消费级GPU，在24小时内完成当前最先进视频模型的训练？
• RQ2 视频模型训练中的主要瓶颈是什么？能否在大规模场景下系统性地加以解决？
• RQ3 内存高效的注意力机制与融合操作在多大程度上可降低训练成本与资源需求？
• RQ4 该流水线是否能在显著减少计算量的前提下，实现视频-语言与视频识别基准上的SOTA性能？
• RQ5 与之前工作相比，该流水线在内存使用、训练时间与硬件成本方面表现如何？

主要发现

• 所提出的流水线在单台8-GPU A5000机器上，18小时内完成对400万对视频-文本数据的训练，实现了Epic-Kitchens 100数据集上的SOTA性能。
• 对于ViT-Base模型，该方法相比同等规模的TimeSformer-Base模型，在零样本平均mAP上提升2.0%，微调后准确率提升1.3%。
• 相比此前需要32张40GB A100 GPU的工作，该方法将内存消耗降低6.7倍，GPU·小时数减少11.8倍，硬件成本降低15倍。
• 在VideoMAE预训练中，该流水线将数据加载开销降低3倍，整体训练时间缩短35%。
• 该系统在单台8-GPU A5000服务器上实现了ViT-Base的总批量大小2,048，同时保持与使用32张A100训练的模型相当的准确率。
• 融合的解码-裁剪操作通过减少冗余数据检索，显著降低了解码操作次数，提升了解码速度并减轻了CPU负载。

(b) Fused DecodeCrop.