Преодоление барьера памяти: масштабирование размера пакета практически до бесконечности для функции потерь контраста.

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

October 22, 2024

Авторы: Zesen Cheng, Hang Zhang, Kehan Li, Sicong Leng, Zhiqiang Hu, Fei Wu, Deli Zhao, Xin Li, Lidong Bing

cs.AI

Аннотация

Контрастная потеря - мощный подход для обучения представлений, где увеличение размера пакета улучшает производительность за счет предоставления большего количества отрицательных примеров для лучшего различения между похожими и различными данными. Однако увеличение размеров пакетов ограничено квадратичным ростом потребления памяти GPU, в основном из-за полной инстанциации матрицы сходства. Для решения этой проблемы мы предлагаем стратегию вычислений на основе плиток, которая разбивает расчет контрастной потери на произвольные маленькие блоки, избегая полной материализации матрицы сходства. Более того, мы предлагаем многоуровневую стратегию разделения для использования иерархической структуры распределенных систем, применяя кольцевую коммуникацию на уровне GPU для оптимизации синхронизации и объединенные ядра на уровне ядер CUDA для снижения накладных расходов на ввод-вывод. Экспериментальные результаты показывают, что предложенный метод масштабирует размеры пакетов до рекордных уровней. Например, он позволяет проводить контрастное обучение модели CLIP-ViT-L/14 с размером пакета 4M или 12M с использованием 8 или 32 A800 80GB без ущерба для точности. По сравнению с передовыми решениями по экономии памяти, он обеспечивает сокращение памяти на два порядка величины при сохранении сопоставимой скорости. Код будет доступен публично.

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.