DisCo-CLIP: A Distributed Contrastive Loss for Memory Efficient CLIP Training

解决问题： 本篇论文旨在解决在训练对比学习模型时，对比损失函数所需的内存消耗过大的问题。通过提出分布式内存高效的DisCo-CLIP训练方法，将对比损失函数分解为两部分，一部分计算内部GPU梯度，另一部分通过all_reduce从其他GPU收集外部GPU梯度，从而将GPU内存消耗从$\bigO(B^2)$降至$\bigO(\frac{B^2}{N})$，其中$B$和$N$分别为批量大小和用于训练的GPU数量。

关键思路： DisCo-CLIP的关键思路是将对比损失函数分解为内部GPU梯度和外部GPU梯度，通过all_reduce从其他GPU收集外部GPU梯度，降低GPU内存消耗。相比于当前领域的研究，DisCo-CLIP的思路在解决对比学习模型训练时的内存消耗问题上有新意。

其他亮点： DisCo-CLIP可以使ViT-B/32模型的对比训练批量大小达到32K或196K，使用8或64个A100 40GB GPU进行训练，而原始的CLIP解决方案需要128个A100 40GB GPU才能训练一个批量大小为32K的ViT-B/32模型。DisCo-CLIP的代码将在https://github.com/IDEA-Research/DisCo-CLIP上发布。

关于作者： 本文的主要作者为Yihao Chen、Xianbiao Qi、Jianan Wang和Lei Zhang。他们来自IDEA研究组和香港科技大学。在此之前，他们的代表作包括：

  • Chen等人的“Cross-Domain Self-Supervised Learning for Speech Recognition”（2021）
  • Qi等人的“Learning to Compose with Professional Photographs on the Web”（2020）
  • Wang等人的“Deep High-Resolution Representation Learning for Human Pose Estimation”（2019）
  • Zhang等人的“Residual Attention Network for Image Classification”（2017）

相关研究： 与本文相关的其他研究包括：

  • “Big Self-Supervised Models are Strong Semi-Supervised Learners”（Chen等人，2021，Google Research）
  • “Distributed Deep Learning with Mixed-Precision: Quantizing Gradient and Accumulation Tensors”（Liu等人，2021，Tsinghua University）
  • “A Survey of Distributed Deep Learning”（Wang等人，2020，University of California, Los Angeles）

论文摘要：本文提出了DisCo-CLIP，一种分布式的、节省内存的CLIP训练方法，旨在减少对对比损失的内存消耗。该方法将对比损失及其梯度计算分解为两个部分，一个用于计算内部GPU梯度，另一个用于计算跨GPU梯度。根据分解，只有内部GPU梯度在当前GPU上计算，而跨GPU梯度通过all_reduce从其他GPU收集，而不是在每个GPU上重复计算。这样，我们可以将对比损失计算的GPU内存消耗从$\bigO(B^2)$降低到$\bigO(\frac{B^2}{N})$，其中$B$和$N$分别是批量大小和用于训练的GPU数量。这种分布式解决方案在数学上等价于原始的非分布式对比损失计算，不会损失任何计算精度。它特别适用于大批量CLIP训练。例如，DisCo-CLIP可以使用8或64个A100 40GB GPU对ViT-B/32模型进行对比训练，批量大小为32K或196K，而原始CLIP解决方案需要128个A100 40GB GPU才能对批量大小为32K的ViT-B/32模型进行训练。该代码将在https://github.com/IDEA-Research/DisCo-CLIP上发布。

内容中包含的图片若涉及版权问题，请及时与我们联系删除