Skip to main content

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

If our project helps you, please give us a star ⭐ on GitHub to support us. 🙏🙏

arXiv hf_paper PyPI
License Hits GitHub issues GitHub closed issues
zhihu Twitter

💡 Some other multimodal foundation model projects from our team may interest you ✨.

VCD: Mitigating Object Hallucinations in Large Vision-Language Models through Visual Contrastive Decoding
Sicong Leng, Hang Zhang, Guanzheng Chen, Xin Li, Shijian Lu, Chunyan Miao, Lidong Bing
github github arXiv

VideoLLaMA 2: Advancing Spatial-Temporal Modeling and Audio Understanding in Video-LLMs
Zesen Cheng, Sicong Leng, Hang Zhang, Yifei Xin, Xin Li, Guanzheng Chen, Yongxin Zhu, Wenqi Zhang, Ziyang Luo, Deli Zhao, Lidong Bing
github github arXiv

The Curse of Multi-Modalities: Evaluating Hallucinations of Large Multimodal Models across Language, Visual, and Audio
Sicong Leng, Yun Xing, Zesen Cheng, Yang Zhou, Hang Zhang, Xin Li, Deli Zhao, Shijian Lu, Chunyan Miao, Lidong Bing
github github arXiv

📰 News

  • [2024.10.18] Release training and evaluation codes of Inf-CLIP.

🛠️ Requirements and Installation

Basic Dependencies:

  • Python >= 3.8
  • Pytorch >= 2.0.0
  • CUDA Version >= 11.8

[Remote] Install Inf-CL:

# remote installing
pip install inf_cl -i https://pypi.org/simple

[Local] Install Inf-CL:

pip install -e .

Install required packages:

git clone https://github.com/DAMO-NLP-SG/Inf-CLIP
cd Inf-CLIP
pip install -r requirements.txt

⭐ Features

inf_cl is the triton implementation of Inf-CL loss:

inf_clip is the CLIP training codebase with Inf-CL loss and other training features:

🔑 Usage

A simple example about how to adopt our Inf-CL loss for contrastive learning. Using such command for attempting:

torchrun --nproc_per_node 2 tests/example.py
import torch
import torch.nn.functional as F
import torch.distributed as dist
import numpy as np

from inf_cl import cal_inf_loss


def create_cl_tensors(rank, world_size):
    # Parameters
    dtype = torch.float32
    num_heads = 3        # Number of attention heads
    seq_length_q = 32768 # Sequence length
    seq_length_k = 32768
    d_model = 256        # Dimension of each head (must be 16, 32, 64, or 128)

    # Randomly initialize inputs
    q = torch.rand((seq_length_q // world_size, num_heads * d_model), dtype=dtype, device=f"cuda:{rank}")
    k = torch.rand((seq_length_k // world_size, num_heads * d_model), dtype=dtype, device=f"cuda:{rank}")
    l = torch.ones([], dtype=dtype, device=f"cuda:{rank}") * np.log(1 / 0.07)

    q = F.normalize(q, p=2, dim=-1).requires_grad_() # Query
    k = F.normalize(k, p=2, dim=-1).requires_grad_() # Key
    l = l.requires_grad_() # Logit scale

    return q, k, l


if __name__ == "__main__":
    # Assume that the distributed environment has been initialized
    dist.init_process_group("nccl")

    rank = dist.get_rank()
    world_size = dist.get_world_size()

    torch.cuda.set_device(rank)

    # Exampled by Image-Text Contrastive Learning, q is the global image features, 
    # k is the text features, and l is the logit scale.
    q, k, l = create_cl_tensors(rank, world_size)

    # labels are diagonal elements by default. 
    # labels = torch.arange(q.shape[0])
    loss = cal_inf_loss(q, k, scale=l.exp())

    print(loss)

🚀 Main Results

Memory Cost

* denotes adopting "data offload" strategy.

Max Supported Batch Size

Speed

Batch Size Scaling

Training with larger data scale needs larger batch size.

🗝️ Training & Evaluation

Quick Start

To facilitate further development on top of our codebase, we provide a quick-start guide on how to use Inf-CLIP to train a customized CLIP and evaluate the trained model on the mainstream clip benchmarks.

  1. Training Data Structure:
Inf-CLIP
├── datasets
│   ├── cc3m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/cc3m.md
|   |   ├── 0000.tar
|   |   ├── 0001.tar
|   |   ├── ...
|   |   └── 0301.tar
│   ├── cc12m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/cc12m.md
|   |   ├── 0000.tar
|   |   ├── 0001.tar
|   |   ├── ...
|   |   └── 1044.tar
│   ├── laion400m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/laion400m.md
|   |   ├── 00000.tar
|   |   ├── 00001.tar
|   |   ├── ...
|   |   └── 41407.tar
  1. Command:
bash scripts/cc3m/lit_vit-b-32_bs16k.sh
bash scripts/cc12m/lit_vit-b-32_bs32k.sh
bash scripts/laion400m/lit_vit-b-32_bs256k.sh
  1. Evaluation Data Structure:
Inf-CLIP
├── datasets
│   ├── imagenet-1k/ # download val_images.tar.gz of imagenet
|   |   └── val/
|   |   |   ├── n01440764
|   |   |   ├── n01443537
|   |   |   ├── ...
|   |   |   └── n15075141
│   ├── clip-benchmark/ # bash datasets/benchmarks_download.sh
|   |   ├── wds_mscoco_captions
|   |   ├── wds_flickr8k
|   |   ├── wds_flickr30k
|   |   ├── wds_imagenet1k
|   |   ├── wds_imagenetv2
|   |   ├── wds_imagenet_sketch
|   |   ├── wds_imagenet-a
|   |   ├── wds_imagenet-r
|   |   ├── wds_imagenet-o
|   |   └── wds_objectnet
  1. Command:
# imagenet evaluation
bash scripts/imagenet_eval.sh
# overall evaluation
bash scripts/benchmarks_eval.sh

📑 Citation

If you find Inf-CLIP useful for your research and applications, please cite using this BibTeX:

@article{damovl2024infcl,
  title={Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss},
  author={Zesen Cheng, Hang Zhang, Kehan Li, Sicong Leng, Zhiqiang Hu, Fei Wu, Deli Zhao, Xin Li, Lidong Bing},
  journal={arXiv preprint arXiv:2410.17243},
  year={2024},
  url={https://arxiv.org/abs/2410.12787}
}

👍 Acknowledgement

The codebase of Inf-CLIP is adapted from OpenCLIP. We are also grateful for the following projects our Inf-CL arose from:

🔒 License

This project is released under the Apache 2.0 license as found in the LICENSE file. The service is a research preview intended for non-commercial use ONLY, subject to the model Licenses of CLIP, Terms of Use of the data generated by OpenAI, and Laion. Please get in touch with us if you find any potential violations.

Metadata

Release files for inf-cl 1.2.1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for inf-cl 1.2.1
File Size Uploaded
inf_cl-1.2.1.tar.gz 17.1 kB Details

Built distribution (wheel)

Table of built distributions (wheels) for inf-cl 1.2.1
File Interpreter ABI Platform
inf_cl-1.2.1-py3-none-any.whl Python 3 none any Details

Total release size: 32.7 kB

Release files / inf_cl-1.2.1.tar.gz

Download URL inf_cl-1.2.1.tar.gz
Size 17.1 kB
Tags Source
SHA-256 checksum
How to use checksums
0447e409666f921e9f265fff17668604effdaf109127d38c104ec2704cf28190
BLAKE2b-256 checksum
How to use checksums
908bdbfc791cb19105c5fee114a7b902933563499318d6c465f7755222893b3e
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/5.1.1 CPython/3.8.10

Release files / inf_cl-1.2.1-py3-none-any.whl

Download URL inf_cl-1.2.1-py3-none-any.whl
Size 15.6 kB
Tags Python 3
SHA-256 checksum
How to use checksums
c9657b6ddd592a08ce41906eb81f865f4d7ddafda2287215702be1e0e5bd8acf
BLAKE2b-256 checksum
How to use checksums
6fddc8a92556a3b761b15fe61fd9e77f7d21b7256994e61648ea89072d16bd98
Upload date
Uploaded using Trusted Publishing?
What is trusted publishing?
No
Uploaded via twine/5.1.1 CPython/3.8.10

Release history Release notifications | RSS feed

This release

1.2.1 This release

2 release files

1.2

2 release files

1.1

2 release files

1.0

2 release files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page