Ü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
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.