[GH-ISSUE #111] how to use multiple GPUs,the default is to use the first CUDA device

bsakaguchi commented

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

Owner

Originally created by @h13791565170 on GitHub (Jun 2, 2021).
Original GitHub issue: https://github.com/openai/CLIP/issues/111

the default is to use the first CUDA device if there is any, so,how to use Muti-GPU.I set up multiple GPUs, and it still only uses a single GPU

Originally created by @h13791565170 on GitHub (Jun 2, 2021). Original GitHub issue: https://github.com/openai/CLIP/issues/111 the default is to use the first CUDA device if there is any, so,how to use Muti-GPU.I set up multiple GPUs, and it still only uses a single GPU

bsakaguchi closed this issue

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

bsakaguchi commented

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

Author

Owner

@vinson2233 commented on GitHub (Jun 4, 2021):

Are you working on single machine or multiple machine?
Have you tried the DataParallel or DistributedDataParallel from pytorch?

@vinson2233 commented on GitHub (Jun 4, 2021): Are you working on single machine or multiple machine? Have you tried the DataParallel or DistributedDataParallel from pytorch?

bsakaguchi commented

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

Author

Owner

@vinson2233 commented on GitHub (Jun 4, 2021):

The simple way to utilize multiple GPU is using DataParallel on the CLIP model like this

model, preprocess = clip.load("ViT-B/32",device=device,jit=False) #Best model use ViT-B/32
model = torch.nn.DataParallel(model)

and remove .to(device) in your code since the DataParallel will handle how the data will be scattered across GPU.
But this approach is troublesome for the ground_truth since if we pass 100 images on 10 GPU's, then each GPU only receive 10 pairs, and each GPU only return 10x10 cosine similarity matrices, joining the result and you will get 100x10 cosine similarity matrix instead of 100x100 matrix.

My approach is separate the image part and the text part, create data parallel for both of them. So the code goes like this :

def create_logits(x1,x2,logit_scale):
    x1 = x1 / x1.norm(dim=-1, keepdim=True)
    x2 = x2 / x2.norm(dim=-1, keepdim=True)

    # cosine similarity as logits
    logits_per_x1 =  logit_scale*x1 @ x2.t()
    logits_per_x2 =  logit_scale*x2 @ x1.t()

    # shape = [global_batch_size, global_batch_size]
    return logits_per_x1, logits_per_x2

class TextCLIP(nn.Module):
    def __init__(self, model) :
        super(TextCLIP, self).__init__()
        self.model = model
        
    def forward(self,text):
        return self.model.encode_text(text)
    
class ImageCLIP(nn.Module):
    def __init__(self, model) :
        super(ImageCLIP, self).__init__()
        self.model = model
        
    def forward(self,image):
        return self.model.encode_image(image)

model, preprocess = clip.load("ViT-B/32",device=device,jit=False) #Best model use ViT-B/32
model_text = TextCLIP(model)
model_image = ImageCLIP(model)

model_text = torch.nn.DataParallel(model_text)
model_image = torch.nn.DataParallel(model_image)
<define loss function like usual>
optimizer = optim.Adam(model.parameters(), lr=5e-6,betas=(0.9,0.98),eps=1e-6,weight_decay=0.05) #We chain the optimizer to model, not to model_image or model_text, since the computaional graph is still connected. So updating model will also update the model_image and model_text

for epoch in range(EPOCH):
  for batch in dataloader :
    <unpacking, preprocessing and stuff>
    image_embedding = model_image(images)
    text_embedding = model_text(text)

    logit_scale = model.logit_scale.exp()
    logits_per_image, logits_per_text = create_logits(image_embedding,text_embedding,logit_scale)
    ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU  
   <calculate loss and update the model params>

@vinson2233 commented on GitHub (Jun 4, 2021): The simple way to utilize multiple GPU is using DataParallel on the CLIP model like this ```python model, preprocess = clip.load("ViT-B/32",device=device,jit=False) #Best model use ViT-B/32 model = torch.nn.DataParallel(model) ``` and remove `.to(device)` in your code since the DataParallel will handle how the data will be scattered across GPU. But this approach is troublesome for the ground_truth since if we pass 100 images on 10 GPU's, then each GPU only receive 10 pairs, and each GPU only return 10x10 cosine similarity matrices, joining the result and you will get 100x10 cosine similarity matrix instead of 100x100 matrix. My approach is separate the image part and the text part, create data parallel for both of them. So the code goes like this : ```python def create_logits(x1,x2,logit_scale): x1 = x1 / x1.norm(dim=-1, keepdim=True) x2 = x2 / x2.norm(dim=-1, keepdim=True) # cosine similarity as logits logits_per_x1 = logit_scale*x1 @ x2.t() logits_per_x2 = logit_scale*x2 @ x1.t() # shape = [global_batch_size, global_batch_size] return logits_per_x1, logits_per_x2 class TextCLIP(nn.Module): def __init__(self, model) : super(TextCLIP, self).__init__() self.model = model def forward(self,text): return self.model.encode_text(text) class ImageCLIP(nn.Module): def __init__(self, model) : super(ImageCLIP, self).__init__() self.model = model def forward(self,image): return self.model.encode_image(image) model, preprocess = clip.load("ViT-B/32",device=device,jit=False) #Best model use ViT-B/32 model_text = TextCLIP(model) model_image = ImageCLIP(model) model_text = torch.nn.DataParallel(model_text) model_image = torch.nn.DataParallel(model_image) <define loss function like usual> optimizer = optim.Adam(model.parameters(), lr=5e-6,betas=(0.9,0.98),eps=1e-6,weight_decay=0.05) #We chain the optimizer to model, not to model_image or model_text, since the computaional graph is still connected. So updating model will also update the model_image and model_text for epoch in range(EPOCH): for batch in dataloader : <unpacking, preprocessing and stuff> image_embedding = model_image(images) text_embedding = model_text(text) logit_scale = model.logit_scale.exp() logits_per_image, logits_per_text = create_logits(image_embedding,text_embedding,logit_scale) ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU <calculate loss and update the model params> ```

bsakaguchi commented

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

Author

Owner

@jongwook commented on GitHub (Jul 19, 2021):

The JIT model contains hard-coded CUDA device strings which needs to be manually patched by specifying the device option to clip.load(), but using a non-JIT model should be simpler. You can do that by specifying jit=False, which is now the default in clip.load().

Once the non-JIT model is loaded, the procedure shouldn't be any different from the standard PyTorch way. DataParallel as in @vinson2233's example (thanks again!) provides a simpler interface to be used in a single process, e.g. in Jupyter notebook. DistributedDataParallel can better utilize the GPUs by multiprocessing. Read this tutorial if interested.

@jongwook commented on GitHub (Jul 19, 2021): The JIT model contains hard-coded CUDA device strings which needs to be manually patched by specifying the `device` option to `clip.load()`, but using a non-JIT model should be simpler. You can do that by specifying `jit=False`, which is now the default in `clip.load()`. Once the non-JIT model is loaded, the procedure shouldn't be any different from the standard PyTorch way. `DataParallel` as in @vinson2233's example (thanks again!) provides a simpler interface to be used in a single process, e.g. in Jupyter notebook. `DistributedDataParallel` can better utilize the GPUs by multiprocessing. [Read this tutorial](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) if interested.

bsakaguchi commented

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

Author

Owner

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

@jongwook If I use DDP, one problem is that I have to use the logit_scale parameters outside the forward() function which occurs an error. Under DDP, I have to calculate the similarity outside the forward() function to gather all pairs on multi gpus first and then calculate similarity. Is there a solution? Thanks a lot !

@nbl97 commented on GitHub (Sep 30, 2021): @jongwook If I use DDP, one problem is that I have to use the logit_scale parameters outside the forward() function which occurs an error. Under DDP, I have to calculate the similarity outside the forward() function to gather all pairs on multi gpus first and then calculate similarity. Is there a solution? Thanks a lot !

bsakaguchi commented

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

Author

Owner

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

You can put the logit_scale parameter as a single-element nn.Parameter inside the module that DDP receives, and in each forward() you just need to calculate and return the losses corresponding to the local batch. This comment has an example.

@jongwook commented on GitHub (Sep 30, 2021): You can put the logit_scale parameter as a single-element `nn.Parameter` inside the module that DDP receives, and in each forward() you just need to calculate and return the losses corresponding to the local batch. [This comment](https://github.com/openai/CLIP/issues/132#issuecomment-908004353) has an example.

bsakaguchi commented

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

Author

Owner

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

You can put the logit_scale parameter as a single-element nn.Parameter inside the module that DDP receives, and in each forward() you just need to calculate and return the losses corresponding to the local batch. This comment has an example.

@jongwook Sorry, I understand the clever batch sharding but I still wonder does the gradient is correct bacause all_gather() does not support backward function. Do I need to write backward myself?
Besides, I still don't know what is "put the logit_scale parameter as a single-element nn.Parameter inside the module that DDP receives". Is the code snippet(the mentioned comment) inside or outside the forward? If it is outside, the "some parameter used outside the forward" error still exists.
Last, I have a solution but I don't know whether it is correct, could you pls help me have a look at it?

class CLIP(nn.Module):
    def __init__():
        <no change>
        self.logit_scale = nn.Parameter(...)
    def encode_image(self, image):
        <no change>
    def encode_text(self, text):
        <no change>
    def forward(self, image, text):
        image_features = self.encode_image(image)
        text_features = self.encode_text(text)
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        # scale the text_embedding
        logit_scale = self.logit_scale.exp()
        return image_features, logit_scale * text_features

clip_model = CLIP().cuda()
clip_model = torch.nn.parallel.DistributedDataParallel(clip_model)
# train code
for (images, texts) in data_loader:
    img_embed, scale_text_embed = clip_model(images, texts) # [local batch, C]   
    all_img = all_gather(img_embed) # [global batch, C]
    all_scale_text_embed = all_gather(scale_text_embed) # [global batch, C]

    logits_per_image = all_img @ all_scale_text_embed.t()
    logits_per_text = all_scale_text_embed @ all_img.t()

    loss = Loss(logits_per_image, logits_per_text, ground_truth)
    loss.backward()

In this way, the logit_scale is used inside the forward() and the logits_per_image/text are also correct.
The loss on all GPUs are the same, which is not so clever as batch sharding but still works. Is my solution right?
Thanks for your patience!

@nbl97 commented on GitHub (Sep 30, 2021): > You can put the logit_scale parameter as a single-element `nn.Parameter` inside the module that DDP receives, and in each forward() you just need to calculate and return the losses corresponding to the local batch. [This comment](https://github.com/openai/CLIP/issues/132#issuecomment-908004353) has an example. @jongwook Sorry, I understand the clever batch sharding but I still wonder does the gradient is correct bacause all_gather() does not support backward function. Do I need to write backward myself? Besides, I still don't know what is "put the logit_scale parameter as a single-element nn.Parameter inside the module that DDP receives". Is the code snippet([the mentioned comment](https://github.com/openai/CLIP/issues/132)) inside or outside the forward? If it is outside, the "some parameter used outside the forward" error still exists. Last, I have a solution but I don't know whether it is correct, could you pls help me have a look at it? ```python class CLIP(nn.Module): def __init__(): <no change> self.logit_scale = nn.Parameter(...) def encode_image(self, image): <no change> def encode_text(self, text): <no change> def forward(self, image, text): image_features = self.encode_image(image) text_features = self.encode_text(text) image_features = image_features / image_features.norm(dim=-1, keepdim=True) text_features = text_features / text_features.norm(dim=-1, keepdim=True) # scale the text_embedding logit_scale = self.logit_scale.exp() return image_features, logit_scale * text_features clip_model = CLIP().cuda() clip_model = torch.nn.parallel.DistributedDataParallel(clip_model) # train code for (images, texts) in data_loader: img_embed, scale_text_embed = clip_model(images, texts) # [local batch, C] all_img = all_gather(img_embed) # [global batch, C] all_scale_text_embed = all_gather(scale_text_embed) # [global batch, C] logits_per_image = all_img @ all_scale_text_embed.t() logits_per_text = all_scale_text_embed @ all_img.t() loss = Loss(logits_per_image, logits_per_text, ground_truth) loss.backward() ``` In this way, the logit_scale is used inside the forward() and the logits_per_image/text are also correct. The loss on all GPUs are the same, which is not so clever as batch sharding but still works. Is my solution right? Thanks for your patience!

bsakaguchi commented

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

Author

Owner

@jongwook commented on GitHub (Oct 1, 2021):

Your version will basically work! Backproping through all_gather will reduce the compute in the final matmul a little bit, but the majority of compute still happens in the encoders.

The following is a gist of what we were using for backproping through all_gather; you basically need to call reduce_scatter which will calculate the sum of the incoming gradients from all nodes and return the same sum to every node, as if it had calculated the losses for the full batch without sharding.

import torch
import torch.distributed as dist


class AllGatherFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor: torch.Tensor, reduce_dtype: torch.dtype = torch.float32):
        ctx.reduce_dtype = reduce_dtype

        output = list(torch.empty_like(tensor) for _ in range(dist.get_world_size()))
        dist.all_gather(output, tensor)
        output = torch.cat(output, dim=0)
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        grad_dtype = grad_output.dtype
        input_list = list(grad_output.to(ctx.reduce_dtype).chunk(dist.get_world_size()))
        grad_input = torch.empty_like(input_list[dist.get_rank()])
        dist.reduce_scatter(grad_input, input_list)
        return grad_input.to(grad_dtype)


def all_gather(tensor):
    return AllGatherFunction.apply(tensor)

@jongwook commented on GitHub (Oct 1, 2021): Your version will basically work! Backproping through all_gather will reduce the compute in the final matmul a little bit, but the majority of compute still happens in the encoders. The following is a gist of what we were using for backproping through `all_gather`; you basically need to call `reduce_scatter` which will calculate the sum of the incoming gradients from all nodes and return the same sum to every node, as if it had calculated the losses for the full batch without sharding. ```python import torch import torch.distributed as dist class AllGatherFunction(torch.autograd.Function): @staticmethod def forward(ctx, tensor: torch.Tensor, reduce_dtype: torch.dtype = torch.float32): ctx.reduce_dtype = reduce_dtype output = list(torch.empty_like(tensor) for _ in range(dist.get_world_size())) dist.all_gather(output, tensor) output = torch.cat(output, dim=0) return output @staticmethod def backward(ctx, grad_output: torch.Tensor): grad_dtype = grad_output.dtype input_list = list(grad_output.to(ctx.reduce_dtype).chunk(dist.get_world_size())) grad_input = torch.empty_like(input_list[dist.get_rank()]) dist.reduce_scatter(grad_input, input_list) return grad_input.to(grad_dtype) def all_gather(tensor): return AllGatherFunction.apply(tensor) ```

bsakaguchi commented

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

Author

Owner

@Wuziyi616 commented on GitHub (Dec 4, 2021):

@nbl97 @jongwook Sorry for disturbing you, I compared some versions of the DDP training code of CLIP. The code presented here calculates logits of shape [local_batch, global_batch] (also here), while some other code calculates logits of shape [global_batch, global_batch]. I'm wondering if there are any difference between these two implementations?

According to my understanding of DDP, they are the same, because PyTorch will propagate and average gradient across GPUs, so only calculating [local_batch, global_batch] is enough. Another question is regarding the ground-truth labels used here, should it still be torch.arange(B, dtype=torch.long, device=device), or we need to take the local_rank into consideration and use torch.arange(B, dtype=torch.long, device=device) + B * args.local_rank? I look into the all_gather doc, it seems that the gathered tensors look the same on all gpus? Therefore, the correspondence is different in each GPU.

Looking forward to your reply! Thanks in advance

EDIT: so I check the features and it seems that my understanding of GT labels is correct. And the loss seems good now.

@Wuziyi616 commented on GitHub (Dec 4, 2021): @nbl97 @jongwook Sorry for disturbing you, I compared some versions of the DDP training code of CLIP. The code presented here calculates logits of shape [local_batch, global_batch] (also [here](https://github.com/openai/CLIP/issues/132#issuecomment-908004353)), while some other [code](https://github.com/mlfoundations/open_clip/blob/bb665df657d9bf122890993b462a92e2db01bc12/src/training/train.py#L52) calculates logits of shape [global_batch, global_batch]. I'm wondering if there are any difference between these two implementations? According to my understanding of DDP, they are the same, because PyTorch will propagate and average gradient across GPUs, so only calculating [local_batch, global_batch] is enough. Another question is regarding the ground-truth labels used here, should it still be `torch.arange(B, dtype=torch.long, device=device)`, or we need to take the local_rank into consideration and use `torch.arange(B, dtype=torch.long, device=device) + B * args.local_rank`? I look into the `all_gather` [doc](https://pytorch.org/docs/stable/distributed.html), it seems that the gathered tensors look the same on all gpus? Therefore, the correspondence is different in each GPU. Looking forward to your reply! Thanks in advance EDIT: so I check the features and it seems that my understanding of GT labels is correct. And the loss seems good now.

bsakaguchi commented

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

Author

Owner

@durianer-D commented on GitHub (Mar 8, 2022):

I found it necessary to use nn.Module packages CLIP into class. I failed, when I try to just use jit to work on Muti-GPU. But I succeeded, when I use nn.Module packages CLIP into class without jit.

@durianer-D commented on GitHub (Mar 8, 2022): I found it necessary to use nn.Module packages CLIP into class. I failed, when I try to just use jit to work on Muti-GPU. But I succeeded, when I use nn.Module packages CLIP into class without jit.

bsakaguchi commented

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

Author

Owner

@ajain625 commented on GitHub (Jan 24, 2023):

Would I be correct in believing that
ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU
only works if all batches are of the same size? That is, the total number of training examples is a perfect multiple of BATCH_SIZE?
Wouldn't a more general implementation may be
ground_truth = torch.arange(images.shape[0]).to(device) # this one still need manually to put on GPU

@ajain625 commented on GitHub (Jan 24, 2023): Would I be correct in believing that `ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU` only works if all batches are of the same size? That is, the total number of training examples is a perfect multiple of BATCH_SIZE? Wouldn't a more general implementation may be `ground_truth = torch.arange(images.shape[0]).to(device) # this one still need manually to put on GPU`

bsakaguchi commented

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

Author

Owner

@noitanec commented on GitHub (Dec 8, 2023):

  ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU  
 <calculate loss and update the model params>

@vinson2233 , Will this remain as it is (#83 ) or we're doing this separately for model_text & model_image in the training loop?

        convert_models_to_fp32(model)
        optimizer.step()
        clip.model.convert_weights(model)

@noitanec commented on GitHub (Dec 8, 2023): > ``` logits_per_image, logits_per_text = create_logits(image_embedding,text_embedding,logit_scale) > ground_truth = torch.arange(BATCH_SIZE).to(device) # this one still need manually to put on GPU > <calculate loss and update the model params> >``` @vinson2233 , Will this remain as it is (#83 ) or we're doing this separately for `model_text` & `model_image` in the training loop? ``` convert_models_to_fp32(model) optimizer.step() clip.model.convert_weights(model) ```

bsakaguchi referenced this issue

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

[PR #100] [CLOSED] Add end to end steering + clip pipeline with styleGAN2 #455

Rows
Columns

[GH-ISSUE #111] how to use multiple GPUs,the default is to use the first CUDA device #100