Breaking the Memory Barrier: Near Infinite Batch Size Scaling for
  Contrastive Loss
Authors: name: Zesen Cheng; name: Hang Zhang; name: Kehan Li; name: Sicong Leng; name: Zhiqiang Hu; name: Fei Wu; name: Deli Zhao; name: Xin Li; name: Lidong Bing
Published: 2024-10-22 17:59:30

Abstract: 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.