[논문 리뷰] Training a Large Video Model on a Single Machine in a Day

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

Human Pose and Action Recognition인용 수 5

한 줄 요약

이 논문은 소비자용 A5000 그래픽 카드 8개를 탑재한 단일 머신에서 상태최고수준의 비디오 트랜스포머 모델을 18시간 이내에 훈련시킬 수 있는 매우 효율적인 훈련 파이프라인을 제시한다. FlashAttention, 융합된 디코딩-자르기 연산, 청크 기반 비디오 로딩을 통해 I/O, CPU, GPU 계산을 최적화함으로써 이전 작업 대비 메모리 사용량을 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 계산을 규명하고 해결한다.
• 단일 8GPU 워크스테이션에서 상태최고수준의 대비 비디오-언어 모델을 하루 이내에 훈련시킬 수 있도록 한다.
• 이전 모델과 비교해 유사하거나 더 높은 정확도를 달성하면서도 훨씬 적은 계산과 메모리 자원을 사용한다.
• 하드웨어 요구사항과 비용을 줄여 대규모 비디오 모델 훈련을 학술계 및 소규모 연구 랩에 접근 가능하게 한다.

제안 방법

• 자기주의성 주의(self-attention)의 메모리 사용량을 O(N²)에서 O(N)으로 줄이는 FlashAttention를 도입하여, 더 큰 배치 크기를 가능하게 한다. 여기서 N은 비디오 토큰의 수이다.
• 무작위 크기 조정 및 자르기(ReducedResizedCrop)를 비디오 디코딩 단계에 통합한 융합된 디코딩-자르기 연산을 설계하여, 중복 디코딩과 CPU 오버헤드를 최소화한다.
• 긴 비디오를 고정 길이의 압축 세그먼트로 나누어 병렬로 처리하는 청크 기반 비디오 로딩 시스템을 구현한다.
• 모든 데이터 증강 연산(예: 자르기, 뒤집기)을 GPU로 이관하여 병렬 처리 능력을 활용하고 CPU 병목 현상을 줄인다.
• 모든 단계에서 데이터 이동을 최소화하고 GPU 활용도를 극대화하는 방식으로 전체 파이프라인을 종단 간 최적화한다.
• 기존 작업에서 32개 이상의 A100 GPU가 필요로 했던 것에 비해, 단일 8GPU A5000 서버를 사용해 총 배치 크기가 최대 2,048에 이르는 모델을 훈련시킨다.

(a) Decode-then-crop.

실험 결과

연구 질문

• RQ1 소비자용 GPU를 탑재한 단일 머신에서 24시간 이내에 상태최고수준의 비디오 모델을 훈련시킬 수 있는가?
• RQ2 비디오 모델 훈련에서의 주요 병목 현상은 무엇이며, 스케일링을 고려해 체계적으로 해결할 수 있는가?
• RQ3 메모리 효율적인 주의 메커니즘과 융합된 연산이 훈련 비용과 자원 요구량을 얼마나 줄일 수 있는가?
• RQ4 동일한 파이프라인이 유의미하게 감소된 계산량으로 비디오-언어 및 비디오 인식 벤치마크에서 상태최고수준의 성능을 달성할 수 있는가?
• RQ5 기존 작업 대비 메모리 사용량, 훈련 시간, 하드웨어 비용 측면에서 제안된 파이프라인은 어떻게 비교되는가?

주요 결과

• 제안된 파이프라인이 단일 8GPU A5000 머신에서 400만 개의 비디오-텍스트 쌍을 사용해 대비 비디오-언어 모델을 18시간 내에 훈련시켜 Epic-Kitchens 100 벤치마크에서 상태최고수준의 성능을 달성했다.
• ViT-Base 모델의 경우, 유사한 TimeSformer-Base 모델 대비 0.02% 높은 제로샷 평균 mAP와 테스트 후 미세조정 후 0.013% 높은 정확도를 기록했다.
• 기존 작업에서 32× 40GB A100 GPU가 필요로 했던 것에 비해, 파이프라인이 메모리 소비량을 6.7배 줄이고, GPU 시간을 11.8배 감소시키며 하드웨어 비용을 15배 절감했다.
• VideoMAE 사전학습에서, 파이프라인이 데이터 로딩 오버헤드를 3배 줄이고 전체 훈련 시간을 35% 단축시켰다.
• ViT-Base 모델의 경우, 단일 8GPU A5000 서버에서 총 배치 크기가 2,048에 이르게 되었으며, 32× A100에서 훈련된 모델과 유사한 정확도를 유지했다.
• 융합된 디코딩-자르기 연산은 중복된 데이터 검색을 최소화해 디코딩 연산을 줄여 디코딩 속도를 향상시키고 CPU 부담을 감소시켰다.

(b) Fused DecodeCrop.