arXiv:2410.17243cs.CV2024-10被引 7

突破内存限制,实现对比学习超大批次训练

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

  • 分块计算避免完整相似矩阵占用内存
  • 8卡A800下实现400万批次无精度损失
  • 适合大规模对比学习与分布式训练场景

对比损失是表征学习的有效方法,更大批次可提供更丰富的负样本以提升区分能力。然而,批大小受限于显存随批次平方增长,主要源于相似矩阵的完整构建。为此,我们提出基于分块的计算策略,将对比损失分解为任意小块,避免全量相似矩阵的显式存储。进一步引入多级分块机制,利用分布式系统的层次结构,在GPU层面采用环形通信优化同步,在CUDA核心层面使用融合核函数减少输入输出开销。实验表明,该方法实现了前所未有的批大小扩展:在8张A800 80GB GPU上,可对CLIP-ViT-L/14模型进行400万批次训练;在32张卡上达到1200万批次,且精度未下降。相比现有最先进内存优化方案,显存消耗降低两个数量级,同时保持相当的训练速度。代码将公开。

原文摘要 · Abstract (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.

对比学习显存优化大批次训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。