Franchir la barrière de la mémoire : Mise à l'échelle de la taille du lot quasi infinie pour la perte contrastiveBreaking the Memory Barrier: Near Infinite Batch Size Scaling for
Contrastive Loss
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.