Multiprocess data loading speed becomes slow after importing torch

PlusLight: <p>My actual problem:<br>
I am training a tiny mlp network (~1M parameters) with lots of data (~5TB). The data is 2D matrices saved in hdf5 format with blosc compression. Each matrix is saved to a separate file and is around 25MB on disk (50MB after decompression). The matrices are intended to be passed to the network one by one, and no batching is needed (just shuffling for each epoch).<br>
The hardware setup is 4090x4 ddp on a single node, and I find the training speed is severely bottlenecked by data loading speed. The gpus are running around 20% utilization each, and total disk read speed is around 4GB/s (the upper limit of my disk setup should be able to reach 20GB/s read speed, I reach &gt; 10GB/s by directly <code>cp</code> the dataset).<br>
I tried increasing <code>num_workers</code> for <code>Dataloader</code>, but the total speed does not change from <code>num_workers=1</code> to <code>num_workers=4</code> or <code>num_workers=16</code>. I find in <code>htop</code> that there seems to be a lock between processes: If I use <code>num_workers=1</code>, the worker process can run at 100% cpu ultilization; if <code>num_workers=4</code>, each worker can only run at 25%, and <code>num_workers=16</code> each worker at 7%, rendering increasing workers useless.</p>
<p>Minimal example:<br>
After some messing around, I managed to create a minimal example that a single <code>import torch</code> could create such a bottleneck (even if <code>torch</code> is never used!).</p>
<pre><code class="lang-auto">import os
import numpy as np
import multiprocessing
# import torch  # comment or uncomment this line

class CSDataset():  # simulate torch interface, this is not a subclass
    def __init__(self, dir):
        self.files = [f"{dir}/{file}" for file in os.listdir(file)]

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        with open(self.file[idx], "rb") as f:
            return f.read()

dset = CSDataset(data_dir)
indices = np.arange(len(dset))
ids = np.array_split(indices, 8)

def work(id):  # fake workload to read the data
    return len([dset[i] for i in id])

with multiprocessing.Pool(8) as p:
    print(p.map(work, ids))
</code></pre>
<p>If <code>import torch</code> is commented out: runtime 16s, disk read speed &gt;10GB/s, each process cpu at ~80%;<br>
If <code>import torch</code> is kept: runtime 70s, disk read speed ~4GB/s, each process cpu at ~11%;</p>
<p>So what is <code>torch</code> doing and results in poor disk reading performance and inter-process lock?</p>
<p>torch version: 2.0.1+cu118</p>

nivek: <p>I’m curious what OS are you running on?</p>

PlusLight: <p>I run on Oracle Linux 7, you can think of it as centos 7. The torch environment is installed through conda.</p>

nivek: <p>Can you try adding <code>multiprocessing.set_start_method('spawn')</code> before <code>with multiprocessing.Pool(8) as p:</code>?</p>

PlusLight: <p>It would report:</p>
<pre><code class="lang-auto">RuntimeError: context has already been set
</code></pre>

nivek: <p>You might need to do something like this:</p>
<pre><code class="lang-python">ctx = multiprocessing.get_context('spawn')
pool = multiprocessing.pool.Pool(context=ctx)
</code></pre>

PlusLight: <p>Changing multiprocessing to <code>spawn</code> or <code>forkserver</code> fixed the minimal example.<br>
However still no luck with my actual problem even if I pass <code>multiprocessing_context</code> to <code>Dataloader</code>. The cpu utilization up a bit, but the training it/s does not improve much. I think probably it is hitting the bottleneck of python’s ipc and <code>Queue</code>? Maybe I have to implement my custom <code>Dataloader</code> in C++?</p>

nivek: <p>Can you provide a code snippet of your DataLoader and Dataset set up?</p>

PlusLight: <p>Its like this (almost the same as the example, just read actual hdf5 data):</p>
<pre><code class="lang-auto">class CSDataset(Dataset):
    def __init__(self, dir):
        self.files = [f"{dir}/{file}" for file in os.listdir(file)]

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        with h5py.File(self.file[idx]) as f:
            sample = {"X": f["X"][:], "y":f["y"][:].reshape(-1, 1)}
            return sample

ctx = multiprocessing.get_context("forkserver")
dset = CSDataset(path_to_data)
dataloader = DataLoader(dset, batch_size=None, shuffle=True, pin_memeory=True, num_workers=4, multiprocessing_context=ctx, persistent_workers=True)
</code></pre>