Split single model in multiple gpus

Ashwin_Raju: <p>I would like to train a model where it contains 2 sub-modules. I would like to train sub-model 1 in one gpu and sub-model 2 in another gpu. How would i do in pytorch? I tried specifying cuda device separately for each sub-module but it throws an error.<br>
Error: <code>RuntimeError: tensors are on different GPUs</code></p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5).cuda(device=1)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = x.cuda(device=1)
        conv2_in_gpu1 = self.conv2(x)
        x = x.cuda(device=0)
        x = F.relu(F.max_pool2d(self.conv2_drop(conv2_in_gpu1), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net()
if args.cuda:
    model.cuda(device=0)
</code></pre>
<p>This is just an example of what i was trying to achieve. I would like the <code>self.conv2</code> to be performed in gpu1 and rest in gpu0.</p>

ptrblck: <p>What kind of error do you get?</p>
<p>This should work:</p>
<pre><code class="lang-auto">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)
</code></pre>

Ashwin_Raju: <p>I have updated the post with error that I got. I tried your approach which was similar to mine except that I did not use P2P GPU transfer. I still get the same error. <code>RuntimeError: tensors are on different GPUs</code></p>

ptrblck: <p>Could you please post the code throwing this error?</p>

voxmenthe: <p>Is it absolutely necessary to specify? I have had excellent results with<br>
something like:</p>
<p><code>model = &lt;specify model here&gt;</code><br>
<code>model = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))</code></p>
<p>Which just runs your model as efficiently as possible over as many gpus as<br>
you have.</p>

ptrblck: <p>Your approach will use data parallelism, which splits the data/batch onto different GPUs.<br>
I think <a class="mention" href="/u/ashwin_raju">@Ashwin_Raju</a> would like to use model sharding. At least the error sounds like this to me. <img src="//discuss.pytorch.org/images/emoji/apple/wink.png?v=5" title=":wink:" class="emoji" alt=":wink:"></p>

Ashwin_Raju: <p>I have updated the post with the example that i was trying to achieve basically performing <code>self.conv2</code> in a separate gpu</p>

ptrblck: <p>It seems that in this line</p>
<pre><code>x = F.relu(F.max_pool2d(self.conv2_drop(conv2_in_gpu1), 2))
</code></pre>
<p><code>conv2_in_gpu1</code> is still on GPU1, while <code>self.conv2_drop</code> etc. are on GPU0. You only transferred <code>x</code> back to GPU0.</p>
<p>Btw, what is your use case?</p>

ptrblck: <p>Also, I think another error is in calling <code>model.cuda(device=0)</code>.<br>
This will move all model parameters and buffers to the GPU specified with the <code>id</code>.<br>
So your <code>self.conv2</code> layer will also be moved to GPU0, which might cause the same error again.</p>
<p>If you really need model sharding, I think you should define the splits in the <code>__init__</code> function of your model.</p>

Ashwin_Raju: <p>Since my entire module is computation intensive I would like to compute half of my work in 1 gpu and the other in another gpu. I will look into these corrections now</p>

ptrblck: <p>Sounds logical, but note that you won’t gain any performance gains when your model is built in a sequential manner like in your example. The GPUs will have to wait for each other, since they cannot start working without the preceding data.</p>
<p>In that case, you might follow <a class="mention" href="/u/voxmenthe">@voxmenthe</a> 's suggestion of using <code>DataParallel</code>.</p>

Ashwin_Raju: <p>Yeah. I agree . Your solution works <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Andy-jqa: <p>Can we do both DataParallel and splitting the model in different GPUs?</p>

ProGamerGov: <p>If I do something like this:</p>
<pre><code class="lang-auto"># GPU 0:
loss_c = a * b 

# GPU 1:
loss_f = d * e 
</code></pre>
<p>And then I add them together (converting one output to the other’s GPU):</p>
<pre><code class="lang-auto">total_loss = loss_c + loss_f.cuda(0)
</code></pre>
<p>Then when I run backward() on <code>total_loss</code>, will the backward pass split off onto both GPUs? Or just will it just take place on GPU 0?</p>

ptrblck: <p>This is a bit tricky, but is possible.<br>
I’ve created a small code example, which uses model sharing and <code>DataParallel</code>.<br>
It’s using 4 GPUs, where each submodule is split on 2 GPUs as a <code>DataParallel</code> module:</p>
<pre><code class="lang-auto">class SubModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(SubModule, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
        
    def forward(self, x):
        print('SubModule, device: {}, shape: {}\n'.format(x.device, x.shape))
        x = self.conv1(x)
        return x


class MyModel(nn.Module):
    def __init__(self, split_gpus, parallel):
        super(MyModel, self).__init__()
        self.module1 = SubModule(3, 6)
        self.module2 = SubModule(6, 1)
        
        self.split_gpus = split_gpus
        self.parallel = parallel
        if self.split_gpus and self.parallel:
            self.module1 = nn.DataParallel(self.module1, device_ids=[0, 1]).to('cuda:0')
            self.module2 = nn.DataParallel(self.module2, device_ids=[2, 3]).to('cuda:2')
        
    def forward(self, x):
        print('Input: device {}, shape {}\n'.format(x.device, x.shape))
        x = self.module1(x)
        print('After module1: device {}, shape {}\n'.format(x.device, x.shape))
        x = self.module2(x)
        print('After module2: device {}, shape {}\n'.format(x.device, x.shape))
        return x


model = MyModel(split_gpus=True, parallel=True)
x = torch.randn(16, 3, 24, 24).to('cuda:0')
output = model(x)
</code></pre>
<p>The script will output:</p>
<pre><code class="lang-auto">Input: device cuda:0, shape torch.Size([16, 3, 24, 24])

SubModule, device: cuda:0, shape: torch.Size([8, 3, 24, 24])

SubModule, device: cuda:1, shape: torch.Size([8, 3, 24, 24])

After module1: device cuda:0, shape torch.Size([16, 6, 24, 24])

SubModule, device: cuda:2, shape: torch.Size([8, 6, 24, 24])
SubModule, device: cuda:3, shape: torch.Size([8, 6, 24, 24])


After module2: device cuda:2, shape torch.Size([16, 1, 24, 24])
</code></pre>
<p>EDIT: As you can see, I just implemented this one use case. So the conditions on <code>self.split_gpu</code> and <code>self.parallel</code> are a bit useless. However, this should give you a starter for your code.</p>

yolay: <p>Thanks to <a class="mention" href="/u/ptrblck">@ptrblck</a> for your great solution answer.</p>
<p>I just want to remind that it also requires you to move your predicted results tensor and the ground-truth target tensor into the same GPU because of the loss function.</p>
<p>When the loss function calculates the difference between target and prediction tensor, it needs to have both of them on the same GPU.</p>

PeterXiaoGuo: <p>Hi voxmenthe, I think your solution is the answer that most people are looking for!</p>
<p>thank you very much! I solve my problem use this code simply.</p>
<p>Best,<br>
Peter</p>

Rao_Shivansh: <p>hi thanks for your solution, but when i do this I get error in my loss function :<br>
buffer[torch.eq(target, -1.)] = 0<br>
RuntimeError: invalid argument 2: sizes do not match at /opt/conda/conda-bld/pytorch_1512946747676/work/torch/lib/THC/generated/…/generic/THCTensorMasked.cu:13<br>
this is not an error in my code but an error popping up after using parallelism of data ( as i tried to run my less intensive code both with and without data parallelism and it throws up same error while using it with data parallelsim)</p>
<p>my model is intensive and I have 2 gpu’s 12206MiB each. I just need to split my model to use both gpu’s while training as well as testing.</p>
<p>thanks</p>
<p>btw my model is a fcn and its batch size is 1</p>

PeterXiaoGuo: <p>Hi Rao,</p>
<p>I guess the problem results from your model is too huge to hold for 1 GPU only.</p>
<p>You need to split model directly on different GPU when writing it from scratch, split training data cannot help and check the below thread.</p>
<p>sorry I don’t have experience to write a multi-GPU training model.</p>
<p>Best,<br>
Xiao</p>

isalirezag: <aside class="quote no-group quote-modified" data-username="ptrblck" data-post="15" data-topic="13239">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>self.module1 = nn.DataParallel(self.module1, device_ids=[0, 1]).to(‘cuda:0’)</p>
</blockquote>
</aside>
<p>can you explain what does this means, so you are using gpus 0 and 1 to be used for datapralel and then what <code>.to('cuda:0')</code> means/does?</p>