Using pin_memory with DistributedDataParallel

username: yiftach
cooked: <p>Hi,<br>
To fully reap the benefits of using <code>pin_memory=True</code> in the <code>DataLoader</code>, it is advised to modify the CPU to GPU transfers to be <code>non_blocking=True</code> (as advised <a href="https://pytorch.org/docs/stable/notes/cuda.html#use-pinned-memory-buffers" rel="noopener nofollow ugc">here</a>). However, when using <code>DataDistributedParallel</code>, I am no longer the one manually calling the <code>.to(device)</code> - the model’s forward is changed to move the inputs to the same device.</p>
<p>When I use them both, does DDP automatically know I am using <code>pin_memory=True</code> and changes to <code>non_blocking=True</code>? Do I need to modify anything else myself?<br>
Also, should using <code>pin_memory</code> without <code>non_blocking=True</code> also give me an improvement, or is it just when I combine both?</p>
<p>Thanks,<br>
Yiftach</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="yiftach" data-post="1" data-topic="183730">
<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/yiftach/48/58678_2.png" class="avatar"> yiftach:</div>
<blockquote>
<p>However, when using <code>DataDistributedParallel</code>, I am no longer the one manually calling the <code>.to(device)</code> - the model’s forward is changed to move the inputs to the same device.</p>
</blockquote>
</aside>
<p>That’s not the common approach and you can still call the <code>to()</code> operation on the data outside the model.</p>

username: yiftach
cooked: <p>Is it not the default when I supply <code>device_ids</code>? I see a call to <code>to_kwargs</code> from <code>pre_forward</code> called by <code>forward</code> (in <a href="https://github.com/pytorch/pytorch/blob/91dcc3b2726aa33ca9a30dfb8d165d8ec7b7bba3/torch/nn/parallel/distributed.py#L1417" rel="noopener nofollow ugc">here</a>). How do I disable/enable it?</p>

username: ptrblck
cooked: <p>You can refer to e.g. the <a href="https://github.com/pytorch/examples/blob/7f7c222b355abd19ba03a7d4ba90f1092973cdbc/imagenet/main.py">ImageNet example</a> or most of other examples which do not change the model architecture or the <code>forward</code> pass and just wrap the model in the <code>DDP</code> wrapper.</p>

username: yiftach
cooked: <p>I am still confused. In the example code you shared, the model is still wrapped with DDP in the <a href="https://github.com/pytorch/examples/blob/7f7c222b355abd19ba03a7d4ba90f1092973cdbc/imagenet/main.py#L162C24-L162C24" rel="noopener nofollow ugc">line</a>:</p>
<pre><code class="lang-auto">model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
</code></pre>
<p>As far as I understand, this means that when we call forward, we go to:</p>
<pre><code class="lang-auto">    def forward(self, *inputs, **kwargs):
        with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
            inputs, kwargs = self._pre_forward(*inputs, **kwargs)
            output = (
                self.module.forward(*inputs, **kwargs)
                if self._delay_all_reduce_all_params
                else self._run_ddp_forward(*inputs, **kwargs)
            )
            return self._post_forward(output)
</code></pre>
<p>i.e. before <em>our</em> forward is called (<code>self.module.forward</code>), <code>pre_forward</code> will be called, moving inputs to devices.<br>
What am I missing?</p>

username: ptrblck
cooked: <p>If you are sticking to the standard approach of moving the data inside the <code>DataLoader</code> loop to the <code>rank</code> where should the data be moved to again inside the <code>forward</code> method? By then it’s already on the desired rank isn’t it?</p>

username: yiftach
cooked: <p>So if I call <code>.to(device, non_blocking=True)</code> from my own training loop, and then DDP calls <code>.to(device)</code> (without <code>non_blocking</code>), does the second call not make my code wait for the device?<br>
I thought (but I may be wrong here) that the gain from <code>non_blocking=True</code> is meaningful when there is overhead in my code that happens between moving the data to the device and starting the calculations on that data, e.g. inside the <code>forward</code> call. But maybe the assumption here is that if any such overhead exists, its before the call to <code>forward</code>, and in my case, it may be inside my <code>forward</code> - so I want to make sure that I do not miss the optimization if I use DDP.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="yiftach" data-post="7" data-topic="183730">
<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/yiftach/48/58678_2.png" class="avatar"> yiftach:</div>
<blockquote>
<p>does the second call not make my code wait for the device?</p>
</blockquote>
</aside>
<p>No, since your data is already on the <code>device</code> and thus calling <code>to(device)</code> is a no-op.</p>
<aside class="quote no-group" data-username="yiftach" data-post="7" data-topic="183730">
<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/yiftach/48/58678_2.png" class="avatar"> yiftach:</div>
<blockquote>
<p>But maybe the assumption here is that if any such overhead exists, its before the call to <code>forward</code>,</p>
</blockquote>
</aside>
<p>The earlier to move the data to the device asynchronously w.r.t the host the more work it can perform and overlap before the first device operation.<br>
I would assume moving the data inside the <code>forward</code> with <code>non_blocking=True</code> might not help at all as most likely the next call is a CUDA operation depending on the data or which operations do you want to overlap in the <code>forward</code>?</p>