Async data loading has huge GPU bubble

HSL: <p>Hi folks,<br>
I wanted to experiment with asynchronous data loading to overlap the data transfer time with GPU compute workload. I used multiple data loader workers with a prefetch factor of 4 and I used non-blocking data transfer to copy data from cpu to gpu. Here is my code</p>
<pre><code class="lang-auto">def async_data_loading():
    """copying data to gpu asynchronously"""
    size = 1 &lt;&lt; 10
    batch_size = 1 &lt;&lt; 4
    dim = 1 &lt;&lt; 10
    input_list = [torch.randn(dim, dim) for _ in range(size)]

    # mp.set_start_method("spawn")
    dl = torch.utils.data.DataLoader(
        input_list,
        batch_size=batch_size,
        collate_fn=collate_fn,
        pin_memory=True,
        num_workers=4,
        prefetch_factor=4,
        persistent_workers=True,
    )
    f = nn.Linear(dim, dim).cuda()

    for i, batch in enumerate(dl, 1):
        if i &gt; 5:  # warmup
            torch.cuda.cudart().cudaProfilerStart()
        
        batch = batch.cuda(non_blocking=True)
        range_push("forward")
        f(batch)
        range_pop()
    
    torch.cuda.cudart().cudaProfilerStop()
</code></pre>
<p>After I ran the code with nsys, I see huge bubbles from the profiler<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/5/d5d891f89cdf9061f49f7b02f4bea6945a47bc86.png" data-download-href="https://discuss.pytorch.org/uploads/default/d5d891f89cdf9061f49f7b02f4bea6945a47bc86" title="Screenshot 2024-04-05 at 2.24.20 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5d891f89cdf9061f49f7b02f4bea6945a47bc86_2_690x122.png" alt="Screenshot 2024-04-05 at 2.24.20 PM" data-base62-sha1="uvLHryTss12cr1ZsiJ2nbbaYlPU" width="690" height="122" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5d891f89cdf9061f49f7b02f4bea6945a47bc86_2_690x122.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/5/d5d891f89cdf9061f49f7b02f4bea6945a47bc86_2_1035x183.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/5/d5d891f89cdf9061f49f7b02f4bea6945a47bc86.png 2x" data-dominant-color="DFE7E7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-04-05 at 2.24.20 PM</span><span class="informations">1185×211 17.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’d like to understand what if I did anything wrong or my expectation of seeing zero bubble is incorrect.</p>

ptrblck: <p>Based on your code snippet it seems you are directly using the <code>batch</code> in the next line of code (i.e. in the forward pass of your model) after moving it to the GPU. In this case there won’t be any operation to overlap the data transfer since the data would directly be needed. You might also want to profile the <code>DataLoader</code> and check if it’s the current bottleneck in your code creating these idle times.</p>

HSL: <p>The collate function for the data loader is</p>
<pre><code class="lang-auto">def collate_fn(batch):
    # range_push("data_loading")
    batch = torch.stack(batch)
    # range_pop()
    return batch
</code></pre>
<p>do you think it is causing the white space between memory + kernel regions? My understanding is that by setting a high prefetch factor, you can overlap the collating with computation. Also, I tried to record it with nvtx as you can see from the commented out part of the function, but the <code>data_loading</code> section didn’t show up in the profile</p>

ptrblck: <p>I’m unsure which part of your data loading pipeline could be the bottleneck.</p>
<aside class="quote no-group" data-username="HSL" data-post="3" data-topic="200326">
<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/hsl/48/64447_2.png" class="avatar"> HSL:</div>
<blockquote>
<p>My understanding is that by setting a high prefetch factor, you can overlap the collating with computation.</p>
</blockquote>
</aside>
<p>Not necessarily, since the <code>prefetch_factor</code> defined the size of the queue new batches will be pushed into. The workers (used via <code>num_workers&gt;=1</code>) are already preloading the data in the background. However, if they are already too slow and cannot keep up with the actual training loop, increasing the <code>prefetch_factor</code> won’t help since the queue will be empty most of the time.</p>
<aside class="quote no-group" data-username="HSL" data-post="3" data-topic="200326">
<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/hsl/48/64447_2.png" class="avatar"> HSL:</div>
<blockquote>
<p>but the <code>data_loading</code> section didn’t show up in the profile</p>
</blockquote>
</aside>
<p>Was the marker entirely missing in the profile?</p>

HSL: <blockquote>
<p>However, if they are already too slow and cannot keep up with the actual training loop, increasing the <code>prefetch_factor</code> won’t help since the queue will be empty most of the time.</p>
</blockquote>
<p>That makes sense. And I suppose increasing <code>num_workers</code> wont help. I tried different values for num_workers and didn’t see the improvement. I’d like to understand more on how workers are being used to create a batch.</p>
<blockquote>
<p>Was the marker entirely missing in the profile?</p>
</blockquote>
<p>Yeah, that’s right. In the nvtx section, I only see <code>forward</code> being recorded.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/9/b9fa37d64a445c8e78a3db57045b461800380b2a.png" data-download-href="https://discuss.pytorch.org/uploads/default/b9fa37d64a445c8e78a3db57045b461800380b2a" title="Screenshot 2024-04-06 at 7.52.30 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b9fa37d64a445c8e78a3db57045b461800380b2a_2_690x205.png" alt="Screenshot 2024-04-06 at 7.52.30 AM" data-base62-sha1="qxeq9FbnHMOYDA5LQ8rivuJHkXo" width="690" height="205" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b9fa37d64a445c8e78a3db57045b461800380b2a_2_690x205.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b9fa37d64a445c8e78a3db57045b461800380b2a_2_1035x307.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/9/b9fa37d64a445c8e78a3db57045b461800380b2a_2_1380x410.png 2x" data-dominant-color="E2E7E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-04-06 at 7.52.30 AM</span><span class="informations">3362×1002 182 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>I don’t know why it’s not working for you as I cannot reproduce the issue using this simple example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.utils.data import TensorDataset, DataLoader

# simple use case
def my_collate(batch):
    torch.cuda.nvtx.range_push("collate_fn")
    x = torch.stack([a for a, b in batch])
    y = torch.stack([b for a, b in batch])
    torch.cuda.nvtx.range_pop()
    return x, y

device = "cuda"
dataset = TensorDataset(torch.arange(10).to(device), torch.arange(10, 20).to(device))
loader = DataLoader(dataset, batch_size=5, collate_fn=my_collate)

torch.cuda.nvtx.range_push("loader loop")
for data, target in loader:
    print(data, target)
torch.cuda.nvtx.range_pop()
</code></pre>
<p>The profile shows the <code>loader loop</code> as well as the <code>collate_fn</code> ranges:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/5/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0.png" data-download-href="https://discuss.pytorch.org/uploads/default/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0_2_690x156.png" alt="image" data-base62-sha1="7AMYsIWKOT8wGalciSgto3JUGoU" width="690" height="156" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0_2_690x156.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0_2_1035x234.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/5/3537daf06d94b2f8b35a3d9c6d536c857f9f3ce0_2_1380x312.png 2x" data-dominant-color="E9E6E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2224×506 48.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

HSL: <p>interesting. what was your command to run nsys? perhaps I was not using it correctly. here’s mine</p>
<pre><code class="lang-auto">nsys profile --output &lt;profile name&gt; --force-overwrite true python &lt;my program&gt;
</code></pre>

ptrblck: <p>My profile was created using just <code>nsys profile python tmp.py</code> and I’m using Nsight Systems 2024.1.1 with a current PyTorch build from source. Which setup are you using?</p>

HSL: <p>I’m using nsys 2023.3.1 and pytorch 2.1.0. They are both in a docker image and I believe it is based on a recent Nvidia PyTorch image. Sorry that I cannot be more precise, there are too many layers of teams between me and the image I am using</p>