[GH-ISSUE #132] Batch Sharding Details

bsakaguchi commented

2026-10-31 21:11:23 +00:00

Owner

Originally created by @Zasder3 on GitHub (Jul 24, 2021).
Original GitHub issue: https://github.com/openai/CLIP/issues/132

Within the paper you write:

The calculation of embedding similarities was also sharded with individual GPUs comput- ing only the subset of the pairwise similarities necessary for their local batch of embeddings.

I've seen many different interpretations of this line that vary between repo. How did the sharding function within the training code? Did each device receive the 32768x32768 matrix containing the inner-products? Did each device only compute a local NxN matrix s.t. N << 32768? Did each device recieve only 2 32768xEmbed_dim matrices?

Originally created by @Zasder3 on GitHub (Jul 24, 2021). Original GitHub issue: https://github.com/openai/CLIP/issues/132 Within the paper you write: > The calculation of embedding similarities was also sharded with individual GPUs comput- ing only the subset of the pairwise similarities necessary for their local batch of embeddings. I've seen many different interpretations of this line that vary between repo. How did the sharding function within the training code? Did each device receive the `32768x32768` matrix containing the inner-products? Did each device only compute a local `NxN` matrix s.t. `N << 32768`? Did each device recieve only 2 `32768xEmbed_dim` matrices?

bsakaguchi closed this issue

2026-10-31 21:11:23 +00:00

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@jongwook commented on GitHub (Aug 30, 2021):

It was briefly explained in this comment. Each device received the entire embeddings of size 32768 x embed_dim to be matmul'ed with the local embedding vectors, like:

# image_features.shape = [local_batch_size, embed_dim]
# text_features.shape = [local_batch_size, embed_dim]

all_image_features = all_gather(image_features)  # shape = [global_batch_size, embed_dim]
all_text_features = all_gather(text_features)    # shape = [global_batch_size, embed_dim]

logits_per_image = logit_scale * image_features @ all_text_features.t()  # shape = [local_batch_size, global_batch_size]
logits_per_text = logit_scale * text_features @ all_image_features.t()   # shape = [local_batch_size, global_batch_size]

where all_gather is a function gathering and concatenating the tensors across GPUs. The logits are then used as inputs to the cross entropy loss for 32768-way classification, resulting in a loss value corresponding to local_batch_size inputs in each GPU. DistributedDataParallel takes care of averaging these across GPUs, which becomes equivalent to calculating the loss over 32768 x 32768 similarities.

@jongwook commented on GitHub (Aug 30, 2021): It was briefly explained in [this comment](https://github.com/openai/CLIP/issues/29#issuecomment-770135793). Each device received the entire embeddings of size `32768 x embed_dim` to be matmul'ed with the local embedding vectors, like: ```python # image_features.shape = [local_batch_size, embed_dim] # text_features.shape = [local_batch_size, embed_dim] all_image_features = all_gather(image_features) # shape = [global_batch_size, embed_dim] all_text_features = all_gather(text_features) # shape = [global_batch_size, embed_dim] logits_per_image = logit_scale * image_features @ all_text_features.t() # shape = [local_batch_size, global_batch_size] logits_per_text = logit_scale * text_features @ all_image_features.t() # shape = [local_batch_size, global_batch_size] ``` where `all_gather` is a function gathering and concatenating the tensors across GPUs. The logits are then used as inputs to the cross entropy loss for 32768-way classification, resulting in a loss value corresponding to `local_batch_size` inputs in each GPU. `DistributedDataParallel` takes care of averaging these across GPUs, which becomes equivalent to calculating the loss over `32768 x 32768` similarities.

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@mutasem-mattar commented on GitHub (Sep 4, 2021):

I have a question about all_gather used for training, I know there is two versions for all_gather that can be implemented, one which records the gradients from all GPUs and the other one does a detach. I am wondering which one was used for training.

@mutasem-mattar commented on GitHub (Sep 4, 2021): I have a question about `all_gather` used for training, I know there is two versions for `all_gather` that can be implemented, one which records the gradients from all `GPUs` and the other one does a `detach`. I am wondering which one was used for training.

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@jongwook commented on GitHub (Sep 5, 2021):

@mutasem-mattar the all_gather function needs a full, undetached backward operation (calling reduce_scatter or similar operations) for training to work most efficiently.

@jongwook commented on GitHub (Sep 5, 2021): @mutasem-mattar the `all_gather` function needs a full, undetached backward operation (calling `reduce_scatter` or similar operations) for training to work most efficiently.

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@mutasem-mattar commented on GitHub (Sep 26, 2021):

I have one follow up question. How did you manage to all_gather the image embeddings and text embeddings and do dot product without getting OOM problem. specially your batch size is around 32768.

@mutasem-mattar commented on GitHub (Sep 26, 2021): I have one follow up question. How did you manage to `all_gather` the image embeddings and text embeddings and do dot product without getting OOM problem. specially your batch size is around 32768.

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@jongwook commented on GitHub (Sep 30, 2021):

The memory that all-gathered features take is:

 32768 (global batch size) x 1024 (largest CLIP embedding dim) x 2 (fp16) ≈ 64 MB

and the logits after the dot product would take:

32768 (global batch size) x 32768 / n_gpu (local batch size) x 2 (fp16) ≈ 4096 MB / n_gpu

so it's manageable.

@jongwook commented on GitHub (Sep 30, 2021): The memory that all-gathered features take is: 32768 (global batch size) x 1024 (largest CLIP embedding dim) x 2 (fp16) ≈ 64 MB and the logits after the dot product would take: 32768 (global batch size) x 32768 / n_gpu (local batch size) x 2 (fp16) ≈ 4096 MB / n_gpu so it's manageable.

bsakaguchi commented

2026-10-31 21:11:24 +00:00

Author

Owner

@zsnoob commented on GitHub (Nov 24, 2023):

It was briefly explained in this comment. Each device received the entire embeddings of size 32768 x embed_dim to be matmul'ed with the local embedding vectors, like:
# image_features.shape = [local_batch_size, embed_dim]
# text_features.shape = [local_batch_size, embed_dim]

all_image_features = all_gather(image_features)  # shape = [global_batch_size, embed_dim]
all_text_features = all_gather(text_features)    # shape = [global_batch_size, embed_dim]

logits_per_image = logit_scale * image_features @ all_text_features.t()  # shape = [local_batch_size, global_batch_size]
logits_per_text = logit_scale * text_features @ all_image_features.t()   # shape = [local_batch_size, global_batch_size]
where all_gather is a function gathering and concatenating the tensors across GPUs. The logits are then used as inputs to the cross entropy loss for 32768-way classification, resulting in a loss value corresponding to local_batch_size inputs in each GPU. DistributedDataParallel takes care of averaging these across GPUs, which becomes equivalent to calculating the loss over 32768 x 32768 similarities.

I'm confused about trainning implementation. When we do contrastive learning, the diagonal elements are positive samples. But if we calculate the loss locally with shape (local_batch_size, global_batch_size), how do we know the positive samples each GPU?

I have an unconfirmed idea that we may get the local_rank and know the distributed local batch "index" to get the corresponding position of the positive samples. But i don't know whether the local rank is exactly the "local batch index".

@zsnoob commented on GitHub (Nov 24, 2023): > It was briefly explained in [this comment](https://github.com/openai/CLIP/issues/29#issuecomment-770135793). Each device received the entire embeddings of size `32768 x embed_dim` to be matmul'ed with the local embedding vectors, like: > > ```python > # image_features.shape = [local_batch_size, embed_dim] > # text_features.shape = [local_batch_size, embed_dim] > > all_image_features = all_gather(image_features) # shape = [global_batch_size, embed_dim] > all_text_features = all_gather(text_features) # shape = [global_batch_size, embed_dim] > > logits_per_image = logit_scale * image_features @ all_text_features.t() # shape = [local_batch_size, global_batch_size] > logits_per_text = logit_scale * text_features @ all_image_features.t() # shape = [local_batch_size, global_batch_size] > ``` > > where `all_gather` is a function gathering and concatenating the tensors across GPUs. The logits are then used as inputs to the cross entropy loss for 32768-way classification, resulting in a loss value corresponding to `local_batch_size` inputs in each GPU. `DistributedDataParallel` takes care of averaging these across GPUs, which becomes equivalent to calculating the loss over `32768 x 32768` similarities. I'm confused about trainning implementation. When we do contrastive learning, the diagonal elements are positive samples. But if we calculate the loss locally with shape (local_batch_size, global_batch_size), how do we know the positive samples each GPU? I have an unconfirmed idea that we may get the local_rank and know the distributed local batch "index" to get the corresponding position of the positive samples. But i don't know whether the local rank is exactly the "local batch index".

bsakaguchi referenced this issue

2026-10-31 21:15:12 +00:00

[PR #120] [CLOSED] Update model.py #459

Rows
Columns

[GH-ISSUE #132] Batch Sharding Details #120