Interaction between DataLoader's num_workers parameter and multi-GPU training

SnoopKilla98: <p>I am training a transformer with an encoder architecture using PyTorch and Lightning. In particular, I am using a machine with 8 GPUs, each one processing batches of 10 samples. I have defined my custom PyTorch’s DataSet class, and in the <code>__getitem__</code> method I have inserted the command <code>time.sleep(10)</code>​ because I want to understand the effect of increasing the <code>num_workers</code> parameter of the DataLoader.</p>
<p>My reasoning is the following: 1. Since I have 8 GPUs, if I set <code>num_workers=32</code> then I will have 4 workers per GPU to load the data. 2. Since I inserted the command <code>time.sleep(10)</code>, data loading will be for sure the bottleneck of computations. 3. To complete one step, every model replica in each of the 8 GPUs needs to process one batch. 4. Since I have 4 workers per GPU, I expect that the steps should increase “quickly” from 0 to 4, then stop progressing for a while, then increase again by 4 and so on.</p>
<p>However, what I observe is that steps increase 32 by 32 instead of 4 by 4, as if every GPU has 32 workers available for data loading.</p>
<p>What am I getting wrong?</p>

SnoopKilla98: <p>I further investigated the matter by reducing the number of workers to 2 and adding:</p>
<pre><code class="lang-auto">print(f"Processing index {idx} from process number {torch.utils.data.get_worker_info().id}")
</code></pre>
<p>to the <code>__getitem__</code> procedure of the DataSet.</p>
<p>This is the output I am getting:</p>
<pre><code class="lang-auto">Processing index 328984 from process number 0
Processing index 756612 from process number 1
Processing index 633928 from process number 0
Processing index 254308 from process number 1
Processing index 1269380 from process number 0
Processing index 1156174 from process number 1
Processing index 1023439 from process number 0
Processing index 328514 from process number 1
Processing index 659258 from process number 0
Processing index 528768 from process number 1
Processing index 1373253 from process number 0
Processing index 610593 from process number 1
Processing index 115510 from process number 0
Processing index 890381 from process number 1
Processing index 671759 from process number 0
Processing index 183167 from process number 1
Epoch 0:   0%|                                                         | 0/195720 [00:00&lt;?, ?it/s]
Processing index 704502 from process number 1
Processing index 1442105 from process number 0
Processing index 16286 from process number 1
Processing index 1344068 from process number 0
Processing index 807109 from process number 1
Processing index 1277123 from process number 0
Processing index 1270663 from process number 0
Processing index 8758 from process number 1
Processing index 1346151 from process number 1
Processing index 1218119 from process number 0
Processing index 320834 from process number 1
Processing index 980333 from process number 0
Processing index 16215 from process number 1
Processing index 712106 from process number 1
Processing index 439272 from process number 0
Processing index 646609 from process number 0
Epoch 0:   0%|                                                         | 2/195720 [00:11&lt;306:38:53,  0.18it/s, v_num=2, train_loss_step=8.340]
Processing index 266000 from process number 1
Processing index 204153 from process number 0
Processing index 1162778 from process number 1
Processing index 312658 from process number 1
Processing index 1529420 from process number 0
Processing index 1565556 from process number 0
Processing index 868215 from process number 0
Processing index 160078 from process number 1
Processing index 1411870 from process number 1
Processing index 909244 from process number 0
Processing index 1193179 from process number 0
Processing index 413422 from process number 1
Processing index 749390 from process number 1
Processing index 11874 from process number 1
Processing index 929870 from process number 0
Processing index 722520 from process number 0
Epoch 0:   0%|                                                         | 4/195720 [00:20&lt;278:57:18,  0.19it/s, v_num=2, train_loss_step=8.170]
</code></pre>
<p>It looks like every worker starts loading 8 batches of data in parallel <em>within the same process</em>, which is not the expected behaviour according to my understanding.<br>
This is very interesting and I would like to learn more about how this is working.</p>

Zacharias: <p>Hi <a class="mention" href="/u/snoopkilla98">@SnoopKilla98</a>,</p>
<p>I think a minimal example of a dataset and dataloader would be helpful.<br>
The behavior of the multiprocessing dataloading can be investigated without any nn or gpu in this instance, I believe, just by iterating through the dataloader itself.</p>
<p>Also note that the dataloader workers are independent of the GPUs, not actually assigned 4 per GPU or anything like that and that there is a buffer that the workers fill, which defaults to a size of <code>2 * num_workers</code> (and the dataloader retrieves round-robin from the workers afaik), but you can change the buffer size via <code>torch.utils.data.DataLoader(prefetch_factor=N, ...)</code> to simplify further.</p>
<p>Cheers</p>

SnoopKilla98: <p>Hi <a class="mention" href="/u/zacharias">@Zacharias</a>,</p>
<p>I think in my example the presence of multiple GPUs is relevant as I am experimenting with the interaction between the <code>num_workers</code> parameter of the DataLoader and the multiple processes dedicated to each GPU.</p>
<p>As for my DataSet, it is a mapped dataset that loads one batch of data from the hard drive (basically with <code>pandas.read_pickle(data_{idx})</code>, whereas the DataLoader is just the standard PyTorch one where I disable automatic batching.</p>
<p>My question is about the fact that, even though I have only 2 workers in the DataLoader, each one of them starts loading 8 batches of data at the same time (1 for each GPU), whereas I would expect them to start loading one batch, then wait for around 10 seconds, then start loading the second one and so on.</p>
<p>Is it possible that each worker performs one loading for each GPU in parallel?</p>

Zacharias: <blockquote>
<p>Is it possible that each worker performs one loading for each GPU in parallel?</p>
</blockquote>
<p>I think this may depends on your parallization strategy.</p>
<p>For example in Data Parallel (DP) this should not happen: What happens on linux for example is that the N=2 worker processes each collect data and write that into shared memory, from where the main thread can then read it and send it to any of the GPU devices.</p>
<p>For DDP, could you have a look at the docs again, in particular at DistributedSampler (<a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.Sampler" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.3 documentation</a>), it seems possible that there is a replication for each GPU going.</p>
<p>Hope that helps, otherwise kindly post a minimal example that reproduces the behavior please.</p>