Überwindung der Speicherbarriere: Skalierung der Batch-Größe nahezu unendlich für Kontrastverlust.

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

October 22, 2024

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

cs.AI

Zusammenfassung

Der kontrastive Verlust ist ein leistungsstarker Ansatz für die Repräsentationslernen, bei dem größere Batch-Größen die Leistung verbessern, indem sie mehr negative Beispiele bereitstellen, um ähnliche und unähnliche Daten besser zu unterscheiden. Die Skalierung der Batch-Größen wird jedoch durch das quadratische Wachstum des GPU-Speicherverbrauchs eingeschränkt, hauptsächlich aufgrund der vollständigen Instantiierung der Ähnlichkeitsmatrix. Um dies zu lösen, schlagen wir eine kachelbasierte Berechnungsstrategie vor, die die Kontrastverlustberechnung in beliebig kleine Blöcke unterteilt, um die vollständige Materialisierung der Ähnlichkeitsmatrix zu vermeiden. Darüber hinaus führen wir eine mehrstufige Kachelstrategie ein, um die hierarchische Struktur von verteilten Systemen zu nutzen, wobei ringbasierte Kommunikation auf GPU-Ebene zur Optimierung der Synchronisation und fusionierte Kerne auf der CUDA-Core-Ebene zur Reduzierung des I/O-Overheads verwendet werden. Experimentelle Ergebnisse zeigen, dass die vorgeschlagene Methode Batch-Größen auf beispiellose Niveaus skalieren kann. Beispielsweise ermöglicht sie das kontrastive Training eines CLIP-ViT-L/14-Modells mit einer Batch-Größe von 4M oder 12M unter Verwendung von 8 oder 32 A800 80GB, ohne dabei Genauigkeit zu opfern. Im Vergleich zu SOTA-speichereffizienten Lösungen erreicht sie eine Reduzierung des Speichers um zwei Größenordnungen, während sie eine vergleichbare Geschwindigkeit beibehält. Der Code wird öffentlich zugänglich gemacht.

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.