기억 장벽 극복: 대조 손실을 위한 거의 무한한 배치 크기 확장

Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss

October 22, 2024

저자: Zesen Cheng, Hang Zhang, Kehan Li, Sicong Leng, Zhiqiang Hu, Fei Wu, Deli Zhao, Xin Li, Lidong Bing

cs.AI

초록

대조 손실은 표현 학습에 대한 강력한 방법으로, 더 큰 배치 크기는 유사하고 비유사한 데이터를 더 잘 구별하기 위해 더 많은 부정적 샘플을 제공하여 성능을 향상시킵니다. 그러나 배치 크기를 확장하는 것은 GPU 메모리 소비의 이차적 증가로 제한되며, 주로 유사성 행렬의 완전한 인스턴스화 때문입니다. 이를 해결하기 위해 우리는 대조 손실 계산을 임의의 작은 블록으로 분할하여 유사성 행렬의 완전한 자료화를 피하는 타일 기반 계산 전략을 제안합니다. 더불어, 분산 시스템의 계층 구조를 활용하기 위해 다중 수준의 타일링 전략을 소개하여 GPU 수준에서 링 기반 통신을 활용하여 동기화를 최적화하고 CUDA 코어 수준에서 통합 커널을 사용하여 I/O 오버헤드를 줄입니다. 실험 결과는 제안된 방법이 배치 크기를 전례 없는 수준으로 확장시킨다는 것을 보여줍니다. 예를 들어, 8개 또는 32개의 A800 80GB를 사용하여 4M 또는 12M의 배치 크기로 CLIP-ViT-L/14 모델의 대조 훈련을 가능하게 하면서 정확도를 희생시키지 않습니다. 최신 기술에 비해 메모리 효율적인 솔루션을 비교했을 때, 메모리를 두 단계 줄이면서 비슷한 속도를 유지합니다. 코드는 공개적으로 제공될 예정입니다.

English

Contrastive loss is a powerful approach for representation learning, where larger batch sizes enhance performance by providing more negative samples to better distinguish between similar and dissimilar data. However, scaling batch sizes is constrained by the quadratic growth in GPU memory consumption, primarily due to the full instantiation of the similarity matrix. To address this, we propose a tile-based computation strategy that partitions the contrastive loss calculation into arbitrary small blocks, avoiding full materialization of the similarity matrix. Furthermore, we introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems, employing ring-based communication at the GPU level to optimize synchronization and fused kernels at the CUDA core level to reduce I/O overhead. Experimental results show that the proposed method scales batch sizes to unprecedented levels. For instance, it enables contrastive training of a CLIP-ViT-L/14 model with a batch size of 4M or 12M using 8 or 32 A800 80GB without sacrificing any accuracy. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed. The code will be made publicly available.