Model parallelism in pytorch for large(r than 1 GPU) models?
Hi! I have a model that is too large to fit inside a single TITAN X (even with 1 batch size). I want to split it over several GPUs such that the memory cost is shared between GPUs. That is, place different parts of the same model on different GPUs and train it end-to-end. Questions:
ajdroid
<SCODE>class MyModel(nn.Module):
    def __init__(self, split_gpus):
        self.large_submodule1 = ...
        self.large_submodule2 = ...

        self.split_gpus = split_gpus
        if split_gpus:
            self.large_submodule1.cuda(0)
            self.large_submodule1.cuda(1)

    def forward(self, x):
        x = self.large_submodule1(x)
        if split_gpus:
            x = x.cuda(1) # P2P GPU transfer
        return self.large_submodule2(x)
<ECODE>
apaszke
This was so easy! I love your work with PyTorch. Minimum fuss! Cheers!
ajdroid
D-X-Y
<SCODE>   def forward(self, x):
         
        x1= self.resnet_1(x[:,0:3,:,:]).cuda(0)
        x2= self.resnet_2(x[:,3:6,:,:]).cuda(1)

        flat  = torch.cat([x1, x2],1).cuda(???)
        logit = self.fc(flat).cuda(???) 
        return logit

<ECODE>
Hengck
did you figure out how to do this?
isalirezag
dimtsi
Dear apaszke, I am trying to implement inter-GPU communiation by using pytorch+mpi+gpu. Following are the tested code, which is designed to make sure that process0 runs on GPU0 and process1 runs on GPU1. However, the code can not be run successfully. Do you know why? def run(rank, size): mpirun noticed that process rank 1 with PID 0 on node osherlab exited on signal 11 (Segmentation fault). Your help will be appreciated. Thank you.
LeiGuan
Would this also work for one 1 gpu with two sequential steps somehow? If my model is too large to fit on one gpu can I somehow do the forward/backward pass sequtially where I only have one part in gpu memory and somehow cache the other part for the backward pass later. Somehow like this: <SCODE>x = submodule1(x)
#somehow unload intermediate results of submodule1 from gpu here and cache for later backward pass 
#(and then load on gpu again when needed in backward pass of submodule1)
x = submodule2(x)
<ECODE> I could imagine how this works but then I don’t know how I would pass the gradients that come from submodule2 back to submodule1 and initiate the backward pass on submodule1.
marcel1991
Thx
RaayChen
shinx