Rompiendo la Barrera de la Memoria: Escalado del Tamaño del Lote Casi Infinito para la Pérdida 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

Resumen

La pérdida contrastiva es un enfoque poderoso para el aprendizaje de representaciones, donde tamaños de lote más grandes mejoran el rendimiento al proporcionar más ejemplos negativos para distinguir mejor entre datos similares y disímiles. Sin embargo, escalar los tamaños de lote está limitado por el crecimiento cuadrático en el consumo de memoria de la GPU, principalmente debido a la instanciación completa de la matriz de similitud. Para abordar esto, proponemos una estrategia de cálculo basada en bloques que divide el cálculo de la pérdida contrastiva en bloques pequeños arbitrarios, evitando la materialización completa de la matriz de similitud. Además, introducimos una estrategia de división multinivel para aprovechar la estructura jerárquica de los sistemas distribuidos, utilizando comunicación basada en anillos a nivel de GPU para optimizar la sincronización y núcleos CUDA fusionados a nivel de núcleo para reducir la sobrecarga de E/S. Los resultados experimentales muestran que el método propuesto escala los tamaños de lote a niveles sin precedentes. Por ejemplo, permite el entrenamiento contrastivo de un modelo CLIP-ViT-L/14 con un tamaño de lote de 4M o 12M utilizando 8 o 32 A800 80GB sin sacrificar precisión alguna. En comparación con las soluciones eficientes en memoria del estado del arte, logra una reducción de dos órdenes de magnitud en memoria manteniendo una velocidad comparable. El código estará disponible públicamente.

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.