We propose DisCo-CLIP, a distributed memory-efficient CLIP training approach, to reduce the memory consumption of contrastive loss when training contrastive learning models. Our approach decomposes the contrastive loss and its gradient computation into two parts, one to calculate the intra-GPU gradients and the other to compute the inter-GPU gradients. According to our decomposition, only the intra-GPU gradients are computed on the current GPU, while the inter-GPU gradients are collected via all_reduce from other GPUs instead of being repeatedly computed on every GPU. In this way, we can reduce the GPU memory consumption of contrastive loss computation from $\bigO(B^2)$ to $\bigO(\frac{B^2}{N})$, where $B$ and $N$ are the batch size and the number of GPUs used for training. Such a distributed solution is mathematically equivalent to the original non-distributed contrastive loss computation, without sacrificing any computation accuracy. It is particularly efficient for large-batch CLIP training. For instance, DisCo-CLIP can enable contrastive training of a ViT-B/32 model with a batch size of 32K or 196K using 8 or 64 A100 40GB GPUs, compared with the original CLIP solution which requires 128 A100 40GB GPUs to train a ViT-B/32 model with a batch size of 32K. The code will be released at https://github.com/IDEA-Research/DisCo-CLIP


翻译:我们提出了 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 上发布。

0
下载
关闭预览

相关内容

专知会员服务
19+阅读 · 2021年9月16日
专知会员服务
88+阅读 · 2021年6月29日
近期必读的七篇NeurIPS 2020【对比学习】相关论文和代码
专知会员服务
65+阅读 · 2020年10月20日
【Google】平滑对抗训练,Smooth Adversarial Training
专知会员服务
48+阅读 · 2020年7月4日
GNN 新基准!Long Range Graph Benchmark
图与推荐
0+阅读 · 2022年10月18日
浅聊对比学习(Contrastive Learning)
极市平台
2+阅读 · 2022年7月26日
浅聊对比学习(Contrastive Learning)第一弹
PaperWeekly
0+阅读 · 2022年6月10日
深度自进化聚类:Deep Self-Evolution Clustering
我爱读PAMI
15+阅读 · 2019年4月13日
强化学习的Unsupervised Meta-Learning
CreateAMind
17+阅读 · 2019年1月7日
无监督元学习表示学习
CreateAMind
27+阅读 · 2019年1月4日
Unsupervised Learning via Meta-Learning
CreateAMind
42+阅读 · 2019年1月3日
ResNet, AlexNet, VGG, Inception:各种卷积网络架构的理解
全球人工智能
19+阅读 · 2017年12月17日
【推荐】图像分类必读开创性论文汇总
机器学习研究会
14+阅读 · 2017年8月15日
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2010年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
国家自然科学基金
0+阅读 · 2008年12月31日
Arxiv
0+阅读 · 2023年6月2日
Arxiv
16+阅读 · 2022年11月1日
VIP会员
相关资讯
GNN 新基准!Long Range Graph Benchmark
图与推荐
0+阅读 · 2022年10月18日
浅聊对比学习(Contrastive Learning)
极市平台
2+阅读 · 2022年7月26日
浅聊对比学习(Contrastive Learning)第一弹
PaperWeekly
0+阅读 · 2022年6月10日
深度自进化聚类:Deep Self-Evolution Clustering
我爱读PAMI
15+阅读 · 2019年4月13日
强化学习的Unsupervised Meta-Learning
CreateAMind
17+阅读 · 2019年1月7日
无监督元学习表示学习
CreateAMind
27+阅读 · 2019年1月4日
Unsupervised Learning via Meta-Learning
CreateAMind
42+阅读 · 2019年1月3日
ResNet, AlexNet, VGG, Inception:各种卷积网络架构的理解
全球人工智能
19+阅读 · 2017年12月17日
【推荐】图像分类必读开创性论文汇总
机器学习研究会
14+阅读 · 2017年8月15日
相关基金
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2010年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
国家自然科学基金
0+阅读 · 2008年12月31日
Top
微信扫码咨询专知VIP会员