Superamento della barriera della memoria: scalabilità della dimensione del batch quasi infinita per la perdita contrastiva

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

October 22, 2024

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

cs.AI

Abstract

La perdita contrastiva è un approccio potente per l'apprendimento della rappresentazione, dove dimensioni di batch più grandi migliorano le prestazioni fornendo più campioni negativi per distinguere meglio tra dati simili e dissimili. Tuttavia, l'incremento delle dimensioni di batch è limitato dalla crescita quadratica del consumo di memoria GPU, principalmente a causa dell'istanziazione completa della matrice di similarità. Per affrontare questo problema, proponiamo una strategia di calcolo basata su blocchi che suddivide il calcolo della perdita contrastiva in blocchi arbitrariamente piccoli, evitando la completa materializzazione della matrice di similarità. Inoltre, introduciamo una strategia di suddivisione a più livelli per sfruttare la struttura gerarchica dei sistemi distribuiti, utilizzando comunicazioni basate su anelli a livello di GPU per ottimizzare la sincronizzazione e kernel fusi a livello di core CUDA per ridurre l'overhead di I/O. I risultati sperimentali mostrano che il metodo proposto scala le dimensioni di batch a livelli senza precedenti. Ad esempio, consente l'addestramento contrastivo di un modello CLIP-ViT-L/14 con una dimensione di batch di 4M o 12M utilizzando 8 o 32 A800 80GB senza sacrificare alcuna precisione. Rispetto alle soluzioni efficienti in termini di memoria SOTA, ottiene una riduzione di due ordini di grandezza della memoria pur mantenendo una velocità comparabile. Il codice sarà reso pubblicamente disponibile.

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.