メモリバリアの突破：対比損失のためのほぼ無限のバッチサイズスケーリング

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モデルのコントラスティブトレーニングを可能にしますが、精度を犠牲にすることなく。SOTAのメモリ効率の高いソリューションと比較して、同等の速度を維持しながらメモリの削減率が2桁向上しています。コードは公開されます。

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.