Splitting batch at desired proportions in DataParallel

luchinoprince: <p>Hello to everyone.</p>
<p>I just added a second GPU to my machine and I would like to leverage the functionalities of <code>DataParallel</code> <a href="https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html" rel="noopener nofollow ugc">https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html</a> to train my model on both. The problem is that the two GPUs have different memories, namely one is 24GBs, while the other is 12GBs. Hence to fully leverage them it would be logical not to split the bath into two equal halves, but into 2/3 and 1/3.</p>
<p>To my knowledge there is no function that allows you to specify such a split, does anybody now if there is one or if there are other libraries that allow such a split?</p>

mrshenli: <blockquote>
<p>The problem is that the two GPUs have different memories, namely one is 24GBs, while the other is 12GBs. Hence to fully leverage them it would be logical not to split the bath into two equal halves, but into 2/3 and 1/3.</p>
</blockquote>
<p>Hey <a class="mention" href="/u/luchinoprince">@luchinoprince</a>,  DataParallel does not offer this API. Technically you can do that with DistributedDataParallel, but it could hang, as running two concurrent processes using NCCL on the same GPU can lead to undefined behavior.</p>
<p>Also splitting the 24GB GPU into 2 12GB virtual ones might not give you the desired speedup, unless it is also 2X faster than the 12GB GPU.</p>
<p>A work around could be copy+modify or monkey-patch this part of code, and send 2X batch size to the 24GB GPU.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/8b55b86dbdc500c441c2312c79a8ac8b747de17f/torch/nn/parallel/data_parallel.py#L150-L172">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/8b55b86dbdc500c441c2312c79a8ac8b747de17f/torch/nn/parallel/data_parallel.py#L150-L172" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/8b55b86dbdc500c441c2312c79a8ac8b747de17f/torch/nn/parallel/data_parallel.py#L150-L172" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/8b55b86dbdc500c441c2312c79a8ac8b747de17f/torch/nn/parallel/data_parallel.py#L150-L172</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="150" style="counter-reset: li-counter 149 ;">
          <li>def forward(self, *inputs, **kwargs):</li>
          <li>    with torch.autograd.profiler.record_function("DataParallel.forward"):</li>
          <li>        if not self.device_ids:</li>
          <li>            return self.module(*inputs, **kwargs)</li>
          <li></li>
          <li>        for t in chain(self.module.parameters(), self.module.buffers()):</li>
          <li>            if t.device != self.src_device_obj:</li>
          <li>                raise RuntimeError("module must have its parameters and buffers "</li>
          <li>                                   "on device {} (device_ids[0]) but found one of "</li>
          <li>                                   "them on device: {}".format(self.src_device_obj, t.device))</li>
          <li></li>
          <li>        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)</li>
          <li>        # for forward function without any inputs, empty list and dict will be created</li>
          <li>        # so the module can be executed on one device which is the first one in device_ids</li>
          <li>        if not inputs and not kwargs:</li>
          <li>            inputs = ((),)</li>
          <li>            kwargs = ({},)</li>
          <li></li>
          <li>        if len(self.device_ids) == 1:</li>
          <li>            return self.module(*inputs[0], **kwargs[0])</li>
      </ol>
    </code></pre>


  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/8b55b86dbdc500c441c2312c79a8ac8b747de17f/torch/nn/parallel/data_parallel.py#L150-L172" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


luchinoprince: <p>Hey <a class="mention" href="/u/mrshenli">@mrshenli</a> , thanks for the reply and for the information.</p>
<p>Unfortunately the 24GBs GPU is also faster than the latter, I will definetly try to change the code and see if I can leverage both the GPUs. I guess it will take some time to understand how to properly change it, especially where.</p>
<p>Thanks again for the information,<br>
Luca</p>