Multi-GPU operation for general tensor operation

yzhao062: <p>Hi there,</p>
<p>Have a question regarding how to leverage torch for general tensor operations (e.g., matmul, cdist, etc.) other than deep learning. So I do not have a training process but a simple calculation.</p>
<p>For instance, I would like to calculate the pairwise distance of two large matrices (100,000 samples, 128 dimensions) with four GPUs (cuda:0,1,2,3). A single GPU does not have enough memory for doing so.</p>
<blockquote>
<p>A = torch.randn(100000, 128).cuda()</p>
<p>B = torch.randn(100000, 128).cuda()</p>
<p>pdist = torch.nn.PairwiseDistance(p=2)</p>
<p>pairwise_distance = pdist(A, B)</p>
</blockquote>
<p>My questions are:</p>
<ul>
<li>how to easily split the task to multiple GPUs (just like joblib with native Python)?</li>
<li>how to do the calculation in parallel (I could retrieve the gpu id and split the matrix into multiple fold with a for loop)</li>
<li>does pytorch multiprocessing also handle data split with multiple GPU? I am afraid that is not the case.</li>
</ul>
<p>Thanks for the help and happy new year!</p>

mrshenli: <p>Hey <a class="mention" href="/u/yzhao062">@yzhao062</a></p>
<blockquote>
<p>how to easily split the task to multiple GPUs (just like joblib with native Python)?</p>
</blockquote>
<p>Not sure if this is helpful. The code below is how <code>DataParallel</code> parallelizes work on multiple GPUs using multi-threading. But <code>matmul</code> is more complicated than data parallel. You will need to implement your own parallel version if you would like to parallelize one <code>matmul</code> operation across multiple GPUs.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/963f7629b591dc9750476faf1513bc7f1fb4d6de/torch/nn/parallel/parallel_apply.py#L23-L88">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/963f7629b591dc9750476faf1513bc7f1fb4d6de/torch/nn/parallel/parallel_apply.py#L23-L88" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/963f7629b591dc9750476faf1513bc7f1fb4d6de/torch/nn/parallel/parallel_apply.py#L23-L88" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/963f7629b591dc9750476faf1513bc7f1fb4d6de/torch/nn/parallel/parallel_apply.py#L23-L88</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="23" style="counter-reset: li-counter 22 ;">
          <li>def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):</li>
          <li>    r"""Applies each `module` in :attr:`modules` in parallel on arguments</li>
          <li>    contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)</li>
          <li>    on each of :attr:`devices`.</li>
          <li></li>
          <li>    Args:</li>
          <li>        modules (Module): modules to be parallelized</li>
          <li>        inputs (tensor): inputs to the modules</li>
          <li>        devices (list of int or torch.device): CUDA devices</li>
          <li></li>
          <li>    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and</li>
          <li>    :attr:`devices` (if given) should all have same length. Moreover, each</li>
          <li>    element of :attr:`inputs` can either be a single object as the only argument</li>
          <li>    to a module, or a collection of positional arguments.</li>
          <li>    """</li>
          <li>    assert len(modules) == len(inputs)</li>
          <li>    if kwargs_tup is not None:</li>
          <li>        assert len(modules) == len(kwargs_tup)</li>
          <li>    else:</li>
          <li>        kwargs_tup = ({},) * len(modules)</li>
      </ol>
    </code></pre>


  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/963f7629b591dc9750476faf1513bc7f1fb4d6de/torch/nn/parallel/parallel_apply.py#L23-L88" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

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

<p>You can also use <code>torch.multiprocessing</code> and use <a href="https://pytorch.org/docs/stable/notes/multiprocessing.html" rel="noopener nofollow ugc"><code>shared_memory</code></a> to share tensors.</p>
<blockquote>
<p>how to do the calculation in parallel (I could retrieve the gpu id and split the matrix into multiple fold with a for loop).</p>
</blockquote>
<p>It depends on which operator you need. Some operators like <code>add</code> and <code>abs</code> can be easily parallelized, and you can use the same strategy as used by <code>DataParallel</code> to achieve that. Other operators will be harder, especially if they require cross-device communication/data-sharing during the computation.</p>
<blockquote>
<p>does pytorch multiprocessing also handle data split with multiple GPU? I am afraid that is not the case.</p>
</blockquote>
<p>PyTorch supports splitting a tensor in one process, and then share each split with a different process with <code>torch.multiprocessing</code> and <code>shared_memory</code> tensors, but the difficult part will be how to implement multi-device computation.</p>