Quebrando a Barreira de Memória: Escalonamento de Tamanho de Lote Quase Infinito para Perda Contrastiva

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

October 22, 2024

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

cs.AI

Resumo

A perda contrastiva é uma abordagem poderosa para aprendizado de representações, onde tamanhos de lote maiores aprimoram o desempenho ao fornecer mais amostras negativas para distinguir melhor entre dados semelhantes e dissimilares. No entanto, o dimensionamento dos tamanhos de lote é limitado pelo crescimento quadrático no consumo de memória da GPU, principalmente devido à instanciação completa da matriz de similaridade. Para lidar com isso, propomos uma estratégia de cálculo baseada em blocos que divide o cálculo da perda contrastiva em blocos arbitrariamente pequenos, evitando a materialização completa da matriz de similaridade. Além disso, introduzimos uma estratégia de divisão em vários níveis para aproveitar a estrutura hierárquica de sistemas distribuídos, utilizando comunicação baseada em anel no nível da GPU para otimizar a sincronização e kernels fundidos no nível do núcleo CUDA para reduzir a sobrecarga de E/S. Resultados experimentais mostram que o método proposto dimensiona os tamanhos de lote para níveis sem precedentes. Por exemplo, ele permite o treinamento contrastivo de um modelo CLIP-ViT-L/14 com um tamanho de lote de 4M ou 12M usando 8 ou 32 A800 80GB sem sacrificar precisão alguma. Comparado às soluções de eficiência de memória do estado da arte, ele alcança uma redução de dois ordens de magnitude na memória mantendo uma velocidade comparável. O código será disponibilizado publicamente.

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.