无限批扩展可能么？达摩院Inf-CL打破对比学习显存瓶颈，提效100倍！

文章预览

达摩院研究员提出了一种对比损失（Contrastive Loss）的高效实现方式（ Inf-CL ），通过分块计算策略，在单台 A800 机器上就能把 batch size 扩展到 400 万。该方案突破了领域内 “Contrastive loss 由于显存限制不能放大 batch size” 的“共识”，实现了 对比损失的 batch size 近乎无限的扩展 。 论文标题： Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss 论文链接： https://arxiv.org/pdf/2410.17243 代码链接： https://github.com/DAMO-NLP-SG/Inf-CLIP 先看显著结果： ▲ 图1：Inf-CL 与现有方法（CLIP 和 OpenCLIP）的 GPU 显存使用对比。 图中标出了常见的 GPU 显存限制。对于超过 80GB A800 显存瓶颈 的情况，通过曲线拟合估算显存消耗。 1. 左图：在 8×A800 GPU 配置下，CLIP 和 OpenCLIP 的显存消耗呈 二次增长，而 Inf-CL 实现了 线性增长。在 256k batch size 下，Inf-CL 将显存消耗降低 ………………………………