title: Breaking the Memory Barrier: Near Infinite Batch Size Scaling for
    Contrastive Loss
summary: 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.
paper_authors._id: 671b005026d87414f5a61579
paper_authors.name: Zesen Cheng
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-25T09:30:02.067Z
paper_authors._id: 65b2529285b6c21448a10d65
paper_authors.avatarUrl: /avatars/1b09e2742aecce1bbdc57f0c4504cf38.svg
paper_authors.fullname: Zesen Cheng
paper_authors.type: user
paper_authors.user: ClownRat
paper_authors._id: 671b005026d87414f5a6157a
paper_authors.name: Hang Zhang
paper_authors._id: 671b005026d87414f5a6157b
paper_authors.name: Kehan Li
paper_authors._id: 671b005026d87414f5a6157c
paper_authors.name: Sicong Leng
paper_authors._id: 671b005026d87414f5a6157d
paper_authors.name: Zhiqiang Hu
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-25T09:30:00.203Z
paper_authors._id: 637f228152229c63921119c3
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/637f228152229c63921119c3/acwXorra1r9_7i3KlBFjS.jpeg
paper_authors.fullname: Zhiqiang Hu
paper_authors.type: user
paper_authors.user: Zhiqiang007
paper_authors._id: 671b005026d87414f5a6157e
paper_authors.name: Fei Wu
paper_authors._id: 671b005026d87414f5a6157f
paper_authors.name: Deli Zhao
paper_authors._id: 671b005026d87414f5a61580
paper_authors.name: Xin Li
paper_authors._id: 671b005026d87414f5a61581
paper_authors.name: Lidong Bing
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-25T09:29:58.206Z
paper_authors._id: 6454685a548f22be598414c4
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/eMjMWKJ-AouF7eY1-RzGF.jpeg
paper_authors.fullname: Lidong Bing
paper_authors.type: user
paper_authors.user: LidongBing
paper_title: Breaking the Memory Barrier: Near Infinite Batch Size Scaling for
    Contrastive Loss
paper_summary: 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.
paper_ai_keywords: contrastive loss representation learning negative samples similarity matrix tile-based computation strategy multi-level tiling strategy hierarchical structure distributed systems ring-based communication fused kernels CLIP-ViT-L/14 contrastive training SOTA memory-efficient techniques