I have designed a big model, its architecture looks like this:
class BigModel(nn.Module):
    def __init__(self, encoder: nn.Module, component1: nn.Module, component2: nn.Module, component3: nn.Module):
        super(BigModel, self).__init__()
        self.encoder = nn.DataParallel(encoder, device_ids=["cuda:0", "cuda:1","cuda:2", "cuda:3"])
       self.component1 = component1
       self.component2 = component2
       self.component3 = component3
    def deploy(self):
       self.component1 = self.component1.to("cuda:4")
       self.component2 = self.component2.to("cuda:4")
       self.component3 = self.component3.to("cuda:5")

I have a single machine with 6 Tesla V100 GPU (32 GB)
The encoder is very big (BERT-like model) so I want to use 4 GPUs to perform encoding process and then use the rest GPUs for other works of this model.
It can work but some times the so of the outputs of the encoder are  lost.
For example, I input a batch of sequences of size (16, 256) to the encoder, data parallel should split it into 4 tensor of size (4, 256) and encode them in parallel. Then gather those output and merge into a tensor of size (16, 256, 1024) .
But, sometimes I only got the output of size  (12, 256, 1024),  that means one split of data is lost. I can not figure out the reason of this problem…
So is there anyone can explain this problem or give a way to combine model parallel and data parallel ?