DisCo-CLIP：一种用于内存高效 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 以 32K 或 196K 的批大小对 ViT-B/32 模型进行对训练，而原始 CLIP 方案需要 128 块 A100 40GB GPU 才能以 32K 批大小训练 ViT-B/32 模型。代码将于 https://github.com/IDEA-Research/DisCo-CLIP 发布。