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

要約

【タイトル】メモリ効率の高いCLIPトレーニングのための分散対比損失「DisCo-CLIP」

【要約】
・DisCo-CLIPは、対比損失をトレーニングする際のメモリ消費を削減するための分散メモリ効率の高いCLIPトレーニング手法である。
・対比損失とその勾配計算を2つの部分に分解し、GPU内グラデーションとGPU間グラデーションの計算に使用する。
・GPU内グラデーションのみを現在のGPUで計算し、GPU間グラデーションは他のGPUからall_reduceで収集するため、各GPUで反復して計算する必要がなく、対比損失計算のGPUメモリ消費を$\bigO(B^2)$から$\bigO(\frac{B^2}{N})$に削減することができた。
・分散ソリューションは、あらゆる計算精度を犠牲にすることなく、元の非分散対比損失計算と数学的に同等である。
・大規模バッチCLIPトレーニングに特に適している。
・DisCo-CLIPは、8個または64個のA100 40GB GPUを使用して、バッチサイズが32Kまたは196KのViT-B / 32モデルの対比トレーニングを可能にするが、元のCLIPソリューションは、バッチサイズが32KのViT-B/32モデルをトレーニングするには128個のA10040BGPUが必要である。
・DisCo-CLIPのコードは、https：//github.com/IDEA-Research/DisCo-CLIPで公開される予定。

要約(オリジナル)

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

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI