[論文レビュー] 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

ひとこと要約

本論文は、1台のマシンに搭載された8枚のエントリーレベルのA5000 GPUを用いて、18時間未満で最先端の動画トランスフォーマーモデルをトレーニング可能な極めて効率的なパイプラインを提示している。FlashAttention、統合デコード・クロップ演算、チャンクベースの動画ローディングを活用してI/O、CPU、GPU計算を最適化することで、従来手法と比較してメモリ使用量を6.7倍、GPU時間（GPU-hours）を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.

研究の動機と目的

• 大規模な動画モデルのトレーニングにかかる計算コストとハードウェアコストの高さが、学術的アクセスを制限しているという問題に対処すること。
• 動画トレーニングにおける3つの主要なボトル neck（I/O、CPUプリプロセッシング、GPU計算）を同定し、スケーラブルに解決すること。
• 1日未満で、1台の8GPUワークステーション上で最先端の対照的動画言語モデルをトレーニング可能にする。
• 著しく少ない計算量とメモリ使用量で、従来モデルと同等またはそれ以上の精度を達成すること。
• ハードウェア要件とコストを低減することで、大規模な動画モデルトレーニングを学術的および小規模研究ラボにアクセス可能にする。

提案手法

• 自己注意機構のメモリフットプリントをO(N²)からO(N)に削減するFlashAttentionを採用することで、Nが動画トークン数である場合、より大きなバッチサイズを可能にする。
• ランダムリサイズドクロップを動画デコード段階に統合した統合デコード・クロップ演算を設計し、不要なデコードとCPUオーバーヘッドを最小限に抑える。
• 長時間の動画を固定長の圧縮セグメントに分割し、並列処理可能なチャンクベースの動画ローディングシステムを実装する。
• すべてのデータ拡張操作（例：クロッピング、フリップ）をGPUに移行し、並列処理を活用してCPUボトル neck を軽減する。
• 全パイプラインをエンドツーエンドの効率性を最大化するように最適化し、データ移動を最小限に抑え、各段階でのGPU利用率を最大化する。
• 32枚以上のA100 GPUを必要としていた従来手法と比較し、1台の8GPU A5000サーバーで合計バッチサイズ2,048までトレーニング可能にした。

(a) Decode-then-crop.

実験結果

リサーチクエスチョン

• RQ1 1台のマシンにエントリーレベルのGPUを搭載した環境で、24時間未満に最先端の動画モデルをトレーニング可能か？
• RQ2 動画モデルトレーニングにおける主なボトル neck は何か？スケーラブルに解決するにはどうすればよいか？
• RQ3 メモリ効率の良い注意機構と統合演算を用いることで、トレーニングコストとリソース要件をどの程度削減できるか？
• RQ4 同じパイプラインが、大幅に少ない計算量で動画言語および動画認識ベンチマークでSOTA性能を達成できるか？
• RQ5 従来手法と比較して、提案手法のメモリ使用量、トレーニング時間、ハードウェアコストはそれぞれどの程度異なるか？

主な発見

• 提案されたパイプラインは、1台の8GPU A5000マシンで400万件の動画・テキストペアを用いて対照的動画言語モデルを18時間でトレーニングし、Epic-Kitchens 100で最先端の性能を達成した。
• ViT-Baseの場合、同等のTimeSformer-Baseモデルと比較して、ゼロショット平均mAPが2.0%向上し、微調整後には1.3%高い精度を達成した。
• 32枚の40GB A100 GPUを必要としていた従来手法と比較し、本手法ではメモリ使用量を6.7倍、GPU時間（GPU-hours）を11.8倍、ハードウェアコストを15倍削減した。
• VideoMAEの事前学習において、データローディングのオーバーヘッドを3倍削減し、全体のトレーニング時間を35%短縮した。
• ViT-Baseのトレーニングでは、1台の8GPU A5000サーバーで合計バッチサイズ2,048を実現したが、32枚のA100でトレーニングされたモデルと同等の精度を維持した。
• 統合デコード・クロップ演算により、不要なデータ取得を最小限に抑え、デコードのオペレーション数を削減することで、デコード速度が向上し、CPU負荷が低減した。

(b) Fused DecodeCrop.