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