Het doorbreken van de geheugenbarrière: Schalen van de batchgrootte naar bijna oneindig voor Contrastive Loss.

Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss

October 22, 2024

Auteurs: Zesen Cheng, Hang Zhang, Kehan Li, Sicong Leng, Zhiqiang Hu, Fei Wu, Deli Zhao, Xin Li, Lidong Bing

cs.AI

Samenvatting

Contrastieve verlies is een krachtige benadering voor representatie-leren, waarbij grotere batchgroottes de prestaties verbeteren door meer negatieve voorbeelden te bieden om beter onderscheid te maken tussen vergelijkbare en niet-vergelijkbare gegevens. Het schalen van batchgroottes wordt echter beperkt door de kwadratische groei in GPU-geheugenverbruik, voornamelijk als gevolg van de volledige instantiëring van de gelijkheidsmatrix. Om dit aan te pakken, stellen we een strategie voor voor berekening van contrastieve verlies op basis van tegels voor, waarbij de berekening wordt opgedeeld in willekeurig kleine blokken om volledige materialisatie van de gelijkheidsmatrix te vermijden. Bovendien introduceren we een meerlaagse tegelstrategie om gebruik te maken van de hiërarchische structuur van gedistribueerde systemen, waarbij ringgebaseerde communicatie op GPU-niveau wordt toegepast om synchronisatie te optimaliseren en gefuseerde kernels op CUDA-kernniveau om I/O-overhead te verminderen. Experimentele resultaten tonen aan dat de voorgestelde methode batchgroottes schaalt naar ongekende niveaus. Zo maakt het contrastieve training van een CLIP-ViT-L/14-model met een batchgrootte van 4M of 12M mogelijk met 8 of 32 A800 80GB zonder enig verlies aan nauwkeurigheid. Vergeleken met toonaangevende geheugenefficiënte oplossingen, bereikt het een vermindering van twee ordes van grootte in geheugen terwijl het een vergelijkbare snelheid behoudt. De code zal openbaar beschikbaar worden gesteld.

English

Contrastive loss is a powerful approach for representation learning, where larger batch sizes enhance performance by providing more negative samples to better distinguish between similar and dissimilar data. However, scaling batch sizes is constrained by the quadratic growth in GPU memory consumption, primarily due to the full instantiation of the similarity matrix. To address this, we propose a tile-based computation strategy that partitions the contrastive loss calculation into arbitrary small blocks, avoiding full materialization of the similarity matrix. Furthermore, we introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems, employing ring-based communication at the GPU level to optimize synchronization and fused kernels at the CUDA core level to reduce I/O overhead. Experimental results show that the proposed method scales batch sizes to unprecedented levels. For instance, it enables contrastive training of a CLIP-ViT-L/14 model with a batch size of 4M or 12M using 8 or 32 A800 80GB without sacrificing any accuracy. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed. The code will be made publicly available.