CPU bound DataLoader, can't increase throughput with multiprocessing (num_workers > 1)?

sam_l: <p>I have a compute-bound data loading step, and was hoping to improve things by scaling up num_workers. Why can’t I increase throughput by parallelizing compute with multiprocessing (via num_workers &gt; 0 in DataLoader)?</p>
<p><strong>Setup</strong><br>
I have a script (below) to demonstrate the issue. High level setup:</p>
<ul>
<li>
<code>IterableDataset</code> that generates a batch for each iteration. I’m parallelizing with <code>num_workers</code> in <code>DataLoader</code>. I’m manually sharding within the dataset based on worker id.</li>
<li>Batch is randomly generated to rule out any file or network IO</li>
<li>Extra (useless) computation is done when generating a batch to simulate compute bound data preparation and to ensure computation dominates any overhead due to passing data back from workers to main process</li>
<li>I’m doing almost nothing with the batch in the main process, to simulate the extreme situation where data prep completely dominates actual training</li>
<li>Running on a machine w/ 24 CPUs</li>
<li>Pytorch 2.0.0</li>
</ul>
<details>
<summary>
Full script</summary>
<pre><code class="lang-auto"># data_test.py

import sys

import numpy as np
import torch
import torch.utils.data as data
from tqdm.auto import tqdm


class MyIterDataset(data.IterableDataset):
    def __init__(self, num_batches, extra_work):
        self.num_batches = num_batches
        self.extra_work = extra_work
        
    def __iter__(self):
        worker_info = data.get_worker_info()
        
        for ix in range(self.num_batches):
            #manual sharding
            if worker_info is None or ((ix + worker_info.id) % worker_info.num_workers == 0):
                yield self.gen_batch(ix)
        
        
    def gen_batch(self, ix):
        #generate some arbitrary data

        x = np.random.randn(1024) + ix
        for _ in range(self.extra_work):
            #computing things here rather than time.sleep
            # to ensure this really is compute bound
            np.dot(x,x)
    
        x = torch.from_numpy(x)
        return x
        
        
if __name__ == "__main__":
    if len(sys.argv) &lt; 2:
        print(f"""Usage: python {__file__} &lt;num_workers&gt; &lt;num_batches:50&gt; &lt;extra_work:100000&gt;
num_workers arg to DataLoader
num_batches - number of bathes to produce 
extra_work - how much extra computation to do to slow things down""")
        exit(0)
    
    num_workers = int(sys.argv[1])
    num_batches = 50 if len(sys.argv) &lt; 3 else int(sys.argv[2])
    extra_work = 100000 if len(sys.argv) &lt; 4 else int(sys.argv[3])
    
    ds = MyIterDataset(num_batches, extra_work)
    dl = data.DataLoader(ds, batch_size = None, num_workers=num_workers, collate_fn=None)

    n=0
    #minimal compute in main process
    for df in tqdm(dl):
        n += len(df)

    print(f"got {n} samples, expected {1024 * num_batches}")
</code></pre>
</details>
<p><strong>Observations</strong></p>
<p><strong>Total time</strong> to generate all batches increases as I add workers. I would expect 1 worker to be slightly slower than no multiprocessing because of comms overhead, but why no advantage to parallelizing w/ more than 1 worker?</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>Workers</th>
<th>Time</th>
</tr>
</thead>
<tbody>
<tr>
<td>No multiprocessing</td>
<td>0m8.989s</td>
</tr>
<tr>
<td>1 worker</td>
<td>0m10.037s</td>
</tr>
<tr>
<td>2 workers</td>
<td>0m10.596s</td>
</tr>
<tr>
<td>4 workers</td>
<td>0m11.036s</td>
</tr>
<tr>
<td>8 workers</td>
<td>0m12.752s</td>
</tr>
</tbody>
</table>
</div><p>The machine I’m running on has 24 cores, but from looking at <code>top</code>, it seems that <strong>total compute</strong> for all processes (main + workers) only sums up to about 200%, regardless of how many workers are used. <em>Why is this?</em></p>
<p>I would expect to be able to scale up to about 24 workers for this test script fairly efficiently. This makes me suspect I’m either doing something wrong or missing something obvious. It would be appreciated if someone could point my mistake out to me.</p>
<p><strong>Other thoughts</strong></p>
<p>I’ve tried playing with <code>prefetch_factor</code> but it did not seem to help.</p>
<p>In the current script, the dataset is returning batches of torch Tensors and <code>collate_fn</code> is set to None. I’ve also tried returning numpy arrays with and without a collate function. I’ve also tried calling <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.share_memory_.html#torch.Tensor.share_memory_" rel="noopener nofollow ugc"><code>share_memory_</code></a> before returning a batch from the dataset. None of theses seem to affect things.</p>
<p>To test whether communication/serialization of batch data is the culprit, I’ve tried increasing “extra work” by orders of magnitude so that amount of data being passed is tiny relative to computation. Total compute across all processes always is capped at about 200%, regardless of num_workers. This seems to be the smoking gun, but I don’t know what to make of it.</p>

sam_l: <p>Update: I’ve tried out different values for DataLoader’s <code>multiprocessing_context</code> argument. The machine I’m running on is Debian based. The default is <code>fork</code> on linux. Using either <code>spawn</code> or <code>forkserver</code>, I do get a speedup roughly proportional to num_workers, as expected.</p>
<p>I don’t fully understand the difference between these options, but my understanding is fork is generally lighter weight than spawn, and should be preferred. However, CPU seems limited when using <code>fork</code>, but not with <code>spawn</code>. Why might that be?</p>

Haydnspass: <p>Hey,</p>
<p>this is very interesting. In my code I also observer utilization up to 200% on the server machine, whereas my local machine does the job properly. So I figure it is something related with the OS or some environment?</p>
<p>See my question here <a href="https://discuss.pytorch.org/t/dataloader-cpu-utilization-and-slow-training/183062" class="inline-onebox">DataLoader CPU utilization and slow training</a></p>

Haydnspass: <p>Please have a look here: <a href="https://github.com/pytorch/builder/pull/1389" class="inline-onebox" rel="noopener nofollow ugc">Fix processor affinity for fork child by jjyyxx · Pull Request #1389 · pytorch/builder · GitHub</a><br>
It seems to fix it.</p>

sam_l: <p><a class="mention" href="/u/haydnspass">@Haydnspass</a> thanks for those links. This looks like possibly the same issue <a href="https://github.com/pytorch/pytorch/issues/101850" class="inline-onebox" rel="noopener nofollow ugc">Unexpected modification to CPU affinity of Dataloader workers · Issue #101850 · pytorch/pytorch · GitHub</a></p>

sam_l: <p>Yep installing <code>llvm-openmp&lt;16</code> fixes it for me as well. And using the default <code>fork</code> mp start method with dataloader is significantly faster than either <code>spawn</code> or <code>forkserver</code> using the toy example I created above.</p>
<p>For posterity, if you’re seeing only 200% CPU utilization with multiprocessing and running in a conda environment, it might be due to a bug in <code>llvm-openmp 16</code>. Try installing an older version with</p>
<blockquote>
<p><code>[conda|mamba|micromamba] install 'llvm-openmp&lt;16' -c conda-forge</code></p>
</blockquote>
<p>For details follow the links the other replies</p>