Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2024
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866914984270757888 |
|---|---|
| author | Cheng, Zesen Zhang, Hang Li, Kehan Leng, Sicong Hu, Zhiqiang Wu, Fei Zhao, Deli Li, Xin Bing, Lidong |
| author_facet | Cheng, Zesen Zhang, Hang Li, Kehan Leng, Sicong Hu, Zhiqiang Wu, Fei Zhao, Deli Li, Xin Bing, Lidong |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_17243 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss Cheng, Zesen Zhang, Hang Li, Kehan Leng, Sicong Hu, Zhiqiang Wu, Fei Zhao, Deli Li, Xin Bing, Lidong Computer Vision and Pattern Recognition 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. |
| title | Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss |
| topic | Computer Vision and Pattern Recognition |
| url | https://arxiv.org/abs/2410.17243 |