Superamento della barriera della memoria: scalabilità della dimensione del batch quasi infinita per la perdita contrastivaBreaking the Memory Barrier: Near Infinite Batch Size Scaling for
Contrastive Loss
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.