Hi all,
I have a model based on Bert (by using HuggingFace’s implementation) and MLP. I am trying to train it by using 3 gpus I have.
Unfortunately, my code uses 10 Gb of available 11 GB gpu memory in the first gpu and only 500 megabytes in the second and third GPUs. Here is the screenshot of it:

Here is the model and the code I use to initialize and train the model:
class MODEL_BERT_SUM_MLP(nn.Module):
    def __init__(self, input_size, hidden_size, dropout, output_size=2):
        super(MODEL_BERT_SUM_MLP, self).__init__()
        self.input_size   = input_size
        self.hidden_size  = hidden_size
        self.output_size  = output_size
        self.bert         = BertModel.from_pretrained('bert-base-cased')
        self.outlin       = nn.Linear(300,self.output_size)
        self.outlin1       = nn.Linear(self.input_size,300)
        self.activation   = nn.ReLU()
        self.dropout      =nn.Dropout(p=dropout)
        self.loss_func    = nn.CrossEntropyLoss()
        

    def forward(self,input_ids=None, attention_mask=None, token_type_ids=None, labels=None):
        a,pooled_output = self.bert(input_ids,
                                    attention_mask=attention_mask,
                                    token_type_ids=token_type_ids)
        # Here I am summing over the batch dimension. This is not a mistake, I need to do this,
        # is it safe to do that on Multi-gpu setting ? 
        out = torch.sum(pooled_output,dim=0,keepdim=True)
        out = self.dropout(self.activation(self.outlin1(out)))
        out = self.outlin(out)
        lss = self.loss_func(out,labels)
        return (out,lss)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MODEL_BERT_SUM_MLP(args.embed_size,args.hidden_size,args.dropout).to(device)
model = torch.nn.DataParallel(model)
optimizer = AdamW(model.parameters(),lr=args.lr, eps=1e-8)
scheduler = get_linear_schedule_with_warmup(optimizer, 
                                            num_warmup_steps = 0, 
                                            num_training_steps = total_steps)
for step,batch in enumerate(corpus.iter_batches(batch_size=args.bs,split_name='train',shuffle=True)):
                
     batch_input_ids=[torch.tensor(d_input_ids).to(device) for d_input_ids in batch[0]]
     batch_segment_ids = [torch.tensor(d_segment_ids).to(device) for d_segment_ids in batch[1]]
     batch_attention_masks = [torch.tensor(d_attn_masks).to(device) for d_attn_masks in batch[2]]
     batch_labels = torch.tensor(batch[3]).to(device)
                
     logits,loss = model(input_ids=batch_input_ids,
                         token_type_ids=batch_segment_ids,
                         attention_mask=batch_attention_masks,
                         labels = batch_labels)
     loss = loss.mean()

How can I solve this problem ? Is there anyone to help me on this ? Thanks.
EDIT
Note that, Unlike here, I am not using a DataLoader inhereted from (Dataset) class of PyTorch. Would it be a problem ?
EDIT2
Also note that, in my forward function, I am summing the elements in a batch:
out = torch.sum(pooled_output,dim=0,keepdim=True)
I did it because each row in my batch is a different section of a document. I get a representation for each section by using bert independently and then I sum them all to get a single representation for the document. Is this why PyTorch could not divide the batch into gpus?
EDIT3
When I modify the constructor of the model as below ( and remove the line  model = torch.nn.DataParallel(model) ), I observe a better memory balance across GPUs. However, I am now even not sure if these two codes are identical or not.
class MODEL_BERT_SUM_MLP(nn.Module):
    def __init__(self, input_size, hidden_size, dropout, output_size=2):
        super(MODEL_BERT_SUM_MLP, self).__init__()
        self.input_size   = input_size
        self.hidden_size  = hidden_size
        self.output_size  = output_size
        self.bert         = BertModel.from_pretrained('bert-base-cased')
        # THIS LINE IS NEW, now I am only making bert component multi-gpu
        self.bert          = torch.nn.DataParallel(self.bert) 
        self.outlin       = nn.Linear(300,self.output_size)
        self.outlin1       = nn.Linear(self.input_size,300)
        self.activation   = nn.ReLU()
        self.dropout      =nn.Dropout(p=dropout)
        self.loss_func    = nn.CrossEntropyLoss()