突破大模型“内存边界” | Inf-CL重塑{图文对齐新基线}，相同BN下内存最低降低281倍！
突破大模型“内存边界” | Inf-CL重塑{图文对齐新基线}，相同BN下内存最低降低281倍！
原创 AI产品汇 2024年10月28日 07:20 广东
“ 可能很多朋友们只听说过大模型、多模态大模型等，却很少有人知道CLIP底层基准模型。预测哪个字幕对应哪个图像的简单预训练任务是一种高效且可扩展的方法，可以在从互联网收集的4亿对（图像、文本）数据集中从头开始学习SOTA图像表示。CLIP模型可以轻松地转移到大多数任务中，并且通常与完全监督的基线竞争，而不需要任何特定于数据集的训练。它的出现很好的解决了文本与图像特征的对齐问题，使得大模型得到了快速发展，很多大模型里面都可以看到它的身影。对比学习在多种应用场景中都有着广泛的应用，如多模态检索、自监督表征学习和密集文本检索等。本文介绍了Inf-CL，这是一种减轻对比学习中二次记忆成本的新方法，这种成本是由对数求和（LSE）计算的相似性矩阵的完全实例化引起的。Inf-CL没有存储整个矩阵，而是利用LSE的累积特性将LSE计算划分为更小的、顺序计算的图块。这将内存使用限制在图块大小和并行图块的数量上，允许在内存和计算效率之间进行权衡。使用8×A800，CLIP和OpenCLIP的内存消耗呈二次方增长，而Inf-CL实现了线性增长，在256k的批量大小下将内存成本降低了78倍。在1024k的批量大小下，即使使用128个GPU，以前的方法也会超过内存限制，而Inf-CL将内存需求减少了281倍！”
代码链接-https://github.com/DAMO-NLP-SG/Inf-CLIP
论文链接-https://arxiv.org/pdf/2410.17243
01-CLIP系列算法发展历程
01.01-CLIP算法简介
最先进的计算机视觉系统经过训练，可以预测一组固定的预定对象类别。这种受限的监督形式限制了它们的通用性和可用性，因为需要额外的标记数据来指定任何其他视觉概念。直接从原始文本中学习图像是一种有前景的替代方案，它利用了更广泛的监督来源。
大量实验证明，预测哪个字幕对应哪个图像的简单预训练任务是一种高效且可扩展的方法，可以在从互联网收集的4亿对（图像、文本）数据集中从头开始学习SOTA图像表示。在预先训练之后，使用自然语言来参考学习到的视觉概念（或描述新概念），从而使模型能够零样本转移到下游任务。作者通过对30多个不同的现有计算机视觉数据集进行基准测试来研究该方法的性能，这些数据集涵盖了OCR、视频中的动作识别、地理定位和许多类型的细粒度对象分类等任务。
该模型可以轻松地转移到大多数任务中，并且通常与完全监督的基线竞争，而不需要任何特定于数据集的训练。例如，作者在ImageNet零样本上匹配原始ResNet-50的精度，而无需使用它所训练的128万个训练示例中的任何一个。
01.02-OpenCLIP算法简介
扩大神经网络的规模已经在广泛的任务中取得了显著的性能。此外，性能通常遵循可靠的缩放规律，作为训练集大小、模型大小和计算的函数，这为大规模实验变得越来越昂贵提供了有价值的指导。然而，之前关于缩放定律的研究主要使用私人数据和模型，或者专注于单峰语言或视觉学习。
为了解决这些局限性，作者使用公共LAION数据集和开源OpenCLIP存储库研究了对比语言图像预训练（CLIP）的缩放规律。该大规模实验涉及在多达20亿个图像-文本对上训练的模型，并识别多个下游任务的幂律缩放，包括零样本分类、检索、线性探测和端到端微调。
作者发现，训练分布在缩放定律中起着关键作用，虽然OpenAI和OpenCLIP模型具有相同的模型架构和相似的训练配方，但是它们表现出不同的缩放行为。作者开源了其评估工作流程和所有模型，包括最大的公共CLIP模型，以确保可重复性，并使比例定律研究更容易获得。
01.03-DisCo-CL算法简介
本文提出了DisCo CLIP，这是一种分布式内存高效的CLIP训练方法，可以在训练对比学习模型时减少对比损失的内存消耗。该方法将对比损失及其梯度计算分解为两部分，一部分计算GPU内梯度，另一部分计算CPU间梯度。根据作者的分解，只有GPU内梯度在当前GPU上计算，而GPU间梯度是通过其他GPU的all-reduce收集的，而不是在每个GPU上重复计算。
通过这种方式，作者可以将对比损失计算的GPU内存消耗从O（B2）减少到O（B2/N），其中B和N是批量大小和用于训练的GPU数量。这种分布式解决方案在数学上等价于原始的非分布式对比损失计算，而不会牺牲任何计算精度。它对于大批量CLIP训练特别有效。例如，DisCoCLIP可以使用8或64个A100 40GB GPU对批量大小为32K或196K的ViT-B/32模型进行对比训练，而原始CLIP解决方案需要128个A100 40GB GPU来训练批量大小为32K的ViT-B/32模型。
02-Inf-CLIP背景简介
对比学习在多种应用场景中都有着广泛的应用，如多模态检索、自监督表征学习和密集文本检索等。它学习了一个嵌入空间表示，其中相似的数据对保持紧密，而不同的数据对相距甚远。Batch Size的大小对对比学习的成功至关重要，因为它们依赖于块内否定。具体来说，较大的批次能够提供一组不同的负样本，从而增强了模型学习判别表示的能力。
尽管有上述好处，但对比学习中批量大小的缩放受到GPU内存的严重限制。计算和存储图像-文本相似性矩阵所需的内存随批处理大小呈二次增长，即使使用高级硬件，进一步扩展也不切实际，并限制了潜在的性能提升。学者们已经相继提出了多种方法来减轻对比学习中缩放批量大小时的内存限制。梯度缓存通过解耦模型和损失计算来减少内存使用，但损失的内存成本仍然是一个重大的瓶颈。OpenCLIP和DisCo CLIP通过在n个GPU上分布对比损失计算来提高效率，将内存消耗减少了n倍。尽管内存高效技术取得了进步，但大多数研究仅限于128k的批量大小，限制了对比学习的潜力以及现代模型和数据集的缩放需求。
上图展示了Inf-CL和以前的方法（CLIP、OpenCLIP）之间的GPU内存使用比较。虚线标记了常见的GPU内存限制。通过曲线拟合估算超过80G A800瓶颈的内存成本。左图展示了使用8×A800，CLIP和OpenCLIP的内存消耗呈二次方增长，而Inf-CL实现了线性增长，在256k的批量大小下将内存成本降低了78倍。右图展示了在1024k的批量大小下，即使使用128个GPU，以前的方法也会超过内存限制，而Inf-CL将内存需求减少了281倍！
03-Inf-CLIP算法简介
本文介绍了Inf-CL，这是一种减轻对比学习中二次记忆成本的新方法，这种成本是由对数求和（LSE）计算的相似性矩阵的完全实例化引起的。Inf-CL没有存储整个矩阵，而是利用LSE的累积特性将LSE计算划分为更小的、顺序计算的图块。这将内存使用限制在图块大小和并行图块的数量上，允许在内存和计算效率之间进行权衡。
为了提高实际运行效率，作者提出了一种多层次的平铺策略。在粗粒度级别，图像和文本批处理分布在多个GPU上，每个GPU对多行执行串行LSE计算。随着计算的进行，异步列数据交换将通信开销降至最低，如图b所示。在细粒度级别，行计算在每个GPU内的CUDA内核之间并行化，将迭代合并到单个内核中以减少I/O开销。理论上，Inf-CL可以使用较小的图块大小计算具有几乎无限批大小的对比损失，尽管速度有所降低。多级平铺策略对于实现实际的可扩展性和效率、平衡内存减少和计算速度至关重要。
作者在图像-文本对比学习任务中评估了Inf-CL。Inf-CL将空间复杂度从二次型（例如，CLIP为O（b 2），OpenCLIP为O（b2/n））降低到线性型（Inf-CL为O（b/n2）），其中b和n是批量大小和GPU数量。内存使用的大幅减少允许使用大批量进行高效训练。例如，在32个A800 GPU（每个80 GB）上训练批量大小超过10M的ViT-L/14 CLIP模型，每个GPU只需要1.44 GB的内存，比以前的方法提高了30倍以上。
此外，Inf-CL保持了与现有方法一致的精度。在计算时间方面，Inf-CL与先前方法的性能相匹配，在8个A800 GPU上处理64k批量大约需要59个小时。时间成本与批量大小几乎呈线性关系，批量大小从64k增加到256k，导致训练时间大约增长4倍（220.3/49.4≈4）。
04-Inf-CLIP算法整体流程
批量大小的缩放通常伴随着GPU数量的缩放。为了充分利用多个GPU之间的并行性，同时利用单个GPU上的部分串行计算来降低内存成本，作者提出了一种多级平铺方法，将上述LSE计算分配给粗粒度的跨GPU平铺和细粒度的GPU平铺。
上图展示了多层平铺策略。顶部展示了对于跨GPU平铺，每个GPU都分配了多行。计算和列通信是异步执行的，从而降低成本。底部展示了对于GPU内拼接，每个GPU中的计算被进一步划分为块，行计算被分配到多个CUDA内核。每行的累积操作被合并到一个内核中，从而减少SRAM和HBM之间的I/O时间。
05-Inf-CLIP算法实现细节
05.01-数据预处理细节
作者评估了Inf-CL在Laion400M数据集上面的有效性，由于剩余样本中没有图像，作者使用了280M（400M中的280M）样本进行训练。使用RandomResizedDrop对图像进行预处理，裁剪比为[0.75，1.33]，比例为[0.08，1.0]。
05.02-训练超参数
根据ViT-g中的设置，作者采用改进的AdaFactor优化器进行训练。优化器的学习率为1×10−3，权重衰减为1×10-4，系数β1=0.9和β2=0.95。训练跨越8个周期，使用余弦学习率计划，在前0.5个周期内进行线性预热。
05.03-模型实现细节
对于分布式训练而言，作者采用数据并行和自动混合精度（float16）。为了支持更大的批量大小，作者采用了梯度缓存，该缓存将对比损失计算与模型的前向和后向传递解耦。因此，每次迭代的峰值内存成本Mpeak计算如下：Mpeak≈Mdata+max（Mloss，Mbackbone），其中Mdata是数据的内存，Mloss用于损失计算，Mbackbone用于模型的前后操作。
作者将该方法与两个基线进行了比较：CLIP的vanilla损失和OpenCLIP/DisCo-CLIP的局部损失。vanilla损失通过从所有GPU收集行和列特征来计算b×b相似性矩阵，而局部损失只需要列特征即可计算b/n×b相似度矩阵，其中b和n是批量大小和GPU数量。
06-Inf-CLIP算法环境搭建与运行
06.01-搭建运行环境
# 克隆代码到本地06.02-运行样例代码
利用“torchrun --nproc_per_node 2 tests/example.py”运行下面的样例代码。
import torch07-Inf-CLIP算法性能评估
上表展示了跨不同硬件和批量大小训练需要的内存成本。该实验利用具有自动混合精度的数据并行性进行高效的分布式训练。基线包括香草损失（CLIP）和局部损失（OpenCLIP）。
为了说明Inf-CL的内存效率，作者将其与使用相同批大小的先前方法进行了比较。如上表显示，对于损失计算，Inf-CL需要的内存比其前身少得多。具体来说，在8×A800上的批处理大小为128k时，Inf-CL仅消耗0.72 GB，而OpenCLIP需要33.64 GB。然而，尽管使用Inf-CL进行损失计算的内存成本很小，但由于数据内存的增长，峰值内存使用量仍会随着批处理大小的增加而迅速增加，如“最大支持批处理大小”中所述。通过将Inf-CL与“数据卸载”集成，它们可以减轻这种内存增加，使其能够在8×A800上训练一个批处理大小为1024k的ViT-L/14模型。
如上表所示，作者将Inf-CL的最大批量大小与各种模型架构（ViT-B/16或ViT-L/14）和训练预算（8×A800或32×A800）下的SOTA方法进行了比较。Inf-CL的性能明显优于之前的SOTA方法，ViT-B/16在8×A800上的性能提高了4.65倍，在使用32×A800时进一步提高到9.60倍。值得注意的是，随着我们扩大模型尺寸，改进程度会降低；例如当从ViT-B/16变为ViT-L/14时从4.65变为2.94。
为了了解这一趋势，作者分析了其峰值内存使用情况。由于Inf-CL的内存需求可以忽略不计，因此峰值内存主要由Mbackbone+Mdata驱动。Mbackbone是恒定的，这意味着峰值记忆的快速增长主要是由于Mdata的增加。由于ViT-L/14具有较大的Mbackbone，因此剩余的内存只能容纳较小的Mdata批大小。
为了解决这个问题，作者实现了“数据卸载”，这允许我们在每个累积步骤中只将一小部分数据加载到GPU上，从而有效地稳定了数据内存的使用。因此，通过将数据卸载与Inf-CL相结合，其可以在32×A800上将批处理大小扩展到10M以上。
如上图所示，作者将Inf-CL的训练速度与之前的方法进行了比较。使用Inf-CL在8×A800上训练ViT-L/14的速度与之前的方法几乎相同。即使当批量大小增加到超出以前方法的限制时，Inf-CL也会保持迭代时间的线性增加，一个迭代周期持续需要大约59小时。将训练速度结果与内存成本结果相结合表明，nf-CL具有更高的内存效率，同时只引入了一点额外的时间成本。
如上表所示，作者评估了用不同损失实现训练的CLIP模型。在相同的批量下，Inf-CL的性能与之前的方法相似，性能差异在误差范围内，这证实了其设计在损失计算中没有精度损失。此外，结果表明，在一定范围内增加批量大小可以提高性能，从而强调了我们的方法在帮助扩大批量大小方面的重要性。然而，在作者的实验条件下，我们目前观察到，过大的批量生产会导致次优性能。这可能是由于未优化的超参数、训练迭代不足或与数据大小相关的约束等因素造成的。由于该工作主要集中在如何实现大批量训练上，这些因素值得在未来的工作中进一步研究。
如上表所示，作者在表中进行了多级平铺实验，并表明其设计在损耗计算中没有精度损失。这允许任意组合实现几乎相同的零样本分类精度（对于64k批量大小，在ImageNet上约为74.8%），同时显著降低内存成本。根据方程式10，它们的Mpeak由Mbackbone+Mdata决定，而不是像以前的方法那样由Mlass+Mdata确定。
对于Loss复杂性分析，跨GPU平铺为O（b 2/n2），导致内存成本为OpenCLIP的1/n。在此基础上，引入In-GPU平铺可以进一步降低内存成本，使内存成本的增长呈线性，即O（b 2/n2）→O（b/n2）。
