Franchir la barrière de la mémoire : Mise à l'échelle de la taille du lot quasi infinie pour la perte contrastive

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

Résumé

La perte contrastive est une approche puissante pour l'apprentissage de représentations, où des tailles de lots plus grandes améliorent les performances en fournissant plus d'échantillons négatifs pour mieux distinguer les données similaires et différentes. Cependant, l'augmentation des tailles de lots est limitée par la croissance quadratique de la consommation de mémoire GPU, principalement due à l'instanciation complète de la matrice de similarité. Pour remédier à cela, nous proposons une stratégie de calcul basée sur des tuiles qui partitionne le calcul de la perte contrastive en petits blocs arbitraires, évitant la matérialisation complète de la matrice de similarité. De plus, nous introduisons une stratégie de tuilage multi-niveaux pour exploiter la structure hiérarchique des systèmes distribués, en utilisant une communication basée sur des anneaux au niveau du GPU pour optimiser la synchronisation et des noyaux fusionnés au niveau du cœur CUDA pour réduire les surcharges d'E/S. Les résultats expérimentaux montrent que la méthode proposée permet de faire évoluer les tailles de lots à des niveaux sans précédent. Par exemple, elle permet l'entraînement contrastif d'un modèle CLIP-ViT-L/14 avec une taille de lot de 4M ou 12M en utilisant 8 ou 32 A800 80GB sans sacrifier la précision. Comparé aux solutions efficaces en mémoire de l'état de l'art, elle permet une réduction de deux ordres de grandeur de la mémoire tout en maintenant une vitesse comparable. Le code sera rendu publiquement disponible.

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.