Fetching a new iterator takes a long time with DistributedDataParallel but not if running normally on a single GPU

Whisky-Jack: <p>I am trying to adapt some existing training code to run on multiple GPUs on the same node using <code>DistributedDataParallel</code>.  I have a <code>main_worker_function</code> which sets up the dataset, model, etc then runs training and validation. I run either a single version of this to run on a single GPU, or call it using <code>mp.spawn</code> to run with DDP:</p>
<pre><code class="lang-auto">if is_distributed:
    print("Using multiprocessing...")
    devids = ["cuda:{0}".format(x) for x in config.gpu_list]
    print("Using DistributedDataParallel on these devices: {}".format(devids))
    mp.spawn(main_worker_function, nprocs=ngpus, args=(ngpus, is_distributed, config))
else:
    print("Only one gpu found, not using multiprocessing...")
    main_worker_function(0, ngpus, is_distributed, config)
</code></pre>
<p>Currently it seems like I obtain no increase in running speed when running on two GPUS with DDP, and upon closer inspection it seems to be because fetching a new validation iterator takes an enormous amount of time. To check this I replaced my training function with:</p>
<pre><code class="lang-auto">time0 = time()
print("Fetching iterator")
val_iter = iter(self.data_loaders["validation"])
time1 = time()
print("Done in time: ", time1 - time0)

val_iter = iter(self.data_loaders["validation"])
time2 = time()
print("Done in time: ", time2 - time1)
</code></pre>
<p>Running the above code on a single GPU and with <code>num_workers = 4</code> yields outputs:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  0.40811872482299805
Done in time:  0.7533595561981201
</code></pre>
<p>Running the same code on a single GPU and <code>num_workers=4</code>, but spawned as single subprocess with  <code>mp.spawn</code> yields:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  49.32914686203003
Done in time:  50.69797325134277
</code></pre>
<p>As far as I can tell it has something to do with the spawning of the worker functions, because spawning a subprocess with <code>num_workers=2</code> yields:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  25.56034803390503
Done in time:  28.712769746780396
</code></pre>
<p>But I’m not sure what behaviour is changing between single and multiprocessing mode to cause this increase in time, and I’m not quite sure how to check (I’ve been consulting <a href="https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading</a> but I’m not sure if it specifies any difference between default and multiprocessing worker spawning). Note that the initialization of the dataset and sampler should currently be performed the same (I was previously using a <code>DistributedSampler</code> in DDP mode, but this problem persists regardless of or not this sampler is used).</p>

pritamdamania87: <p><a class="mention" href="/u/whisky-jack">@Whisky-Jack</a> Can you share the code you’re using to initialize the DataLoader on each process?</p>
<p>cc <a class="mention" href="/u/vitalyfedyunin">@VitalyFedyunin</a></p>

Whisky-Jack: <p>Note that I’ve made some important progress on this since last week, and so I included that first, with the dataloader information below (you can jump to that part if that’s more important).</p>
<p>I have done more investigation on this since last week, and I have a bit of a better handle on where the problem is occurring. The problem seems to involve differences in spawning the worker functions with different multiprocessing start methods. This is described a little bit in the platform specific behaviours section found in here <a href="https://pytorch.org/docs/stable/data.html#single-and-multi-process-data-loading" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.1 documentation</a>, although it doesn’t really discuss differences between running in a single process vs processes launched with <code>spawn</code>.</p>
<p>Calling <code>print(mp.get_context())</code> in the <code>main_worker_function</code> in single GPU mode yields:</p>
<pre><code class="lang-auto">&lt;multiprocessing.context.ForkContext object at 0x7fc14dd64da0&gt;
</code></pre>
<p>While in the spawned process</p>
<pre><code class="lang-auto">&lt;multiprocessing.context.SpawnContext object at 0x7f8e02fd0ef0&gt;
</code></pre>
<p>Consequently, I think the workers processes are being spawned using <code>fork()</code> in the single process case and <code>spawn()</code> in the multiprocessing case.</p>
<p>The documentation states:</p>
<blockquote>
<ul>
<li>On Unix, <code>fork()</code> is the default <a href="https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing" rel="noopener nofollow ugc"><code>multiprocessing</code></a> start method. Using <code>fork()</code>, child workers typically can access the <code>dataset</code> and Python argument functions directly through the cloned address space.</li>
<li>On Windows, <code>spawn()</code> is the default <a href="https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing" rel="noopener nofollow ugc"><code>multiprocessing</code></a> start method. Using <code>spawn()</code>, another interpreter is launched which runs your main script, followed by the internal worker function that receives the <code>dataset</code>, <code>collate_fn</code> and other arguments through <a href="https://docs.python.org/3/library/pickle.html#module-pickle" rel="noopener nofollow ugc"><code>pickle</code></a> serialization.</li>
</ul>
</blockquote>
<p>Since I’m running on linux, I’m guessing that with <code>fork</code>, the cloned address space is working correctly with low overhead, while in the case of the <code>spawn</code> something about pickling the dataset is causing a big overhead when trying to pass it to the workers.</p>
<p>To check that this difference is causing the issue, I tried running in single GPU mode (so NOT spawning any subprocesses using <code>mp.spawn</code>, or wrapping with <code>DistributedDataParallel</code>), but changing the default start method using <code>mp.set_start_method('spawn')</code> in my main function. This results in the same context:</p>
<pre><code class="lang-auto">&lt;multiprocessing.context.SpawnContext object at 0x7f5e3b4a29e8&gt;
</code></pre>
<p>And produces the same lag as in the multiprocessing case:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  52.480613708496094
Fetching iterator
Done in time:  52.47240161895752
Fetching iterator
Done in time:  53.32112169265747
</code></pre>
<p>I then took a look at my dataset to see what might be causing problems. Our dataset is stored in h5 form, and is loaded into numpy arrays when the dataset is initialized.</p>
<p>The dataset consists of a number of types of data, some of which are too large to load into memory, and so are stored in memmaps (these memmaps are not initialized until the first call to <code>get_item</code>):</p>
<pre><code class="lang-auto">with h5py.File(self.h5_path, 'r') as h5_file:
    self.dataset_length = h5_file["labels"].shape[0]

    hdf5_hit_pmt    = h5_file["hit_pmt"]
    hdf5_hit_time   = h5_file["hit_time"]
    hdf5_hit_charge = h5_file["hit_charge"]

    # initialize memmap param dict
    self.pmt_dict    = {'shape':hdf5_hit_pmt.shape,    'offset':hdf5_hit_pmt.id.get_offset(),   'dtype':hdf5_hit_pmt.dtype}
    self.time_dict   = {'shape':hdf5_hit_time.shape,   'offset':hdf5_hit_time.id.get_offset(),  'dtype':hdf5_hit_time.dtype}
    self.charge_dict = {'shape':hdf5_hit_charge.shape, 'offset':hdf5_hit_charge.id.get_offset(),'dtype':hdf5_hit_charge.dtype}
</code></pre>
<p>The other data is small enough to be loaded into memory, and is stored as attributes of the dataset:</p>
<pre><code class="lang-auto">self.labels    = np.array(h5_file["labels"])
self.energies  = np.array(h5_file["energies"])
self.positions = np.array(h5_file["positions"])
self.angles    = np.array(h5_file["angles"])
</code></pre>
<p>This latter data seems to be what is causing the issue. If none of this data is stored as part of the dataset, very little lag is observed:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  1.6751770973205566
Fetching iterator
Done in time:  2.464519500732422
Fetching iterator
Done in time:  2.2533977031707764
</code></pre>
<p>Whereas if only the labels is stored, the lag increases:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  9.610902309417725
Fetching iterator
Done in time:  10.206432342529297
Fetching iterator
Done in time:  9.656260013580322
</code></pre>
<p>With more parameters producing more lag.</p>
<p>I then tried moving this initialization step to the same initialization function in <code>get_item</code> that creates the memmaps (so this information is initialized in the worker functions after they have spawned), and this seems to eliminate most of the lag:</p>
<pre><code class="lang-auto">Fetching iterator
Done in time:  1.7485809326171875
Fetching iterator
Done in time:  8.036261796951294
Fetching iterator
Done in time:  3.001110076904297
</code></pre>
<p>It does add overhead to the first call to <code>get_item</code>, but not the 50s lag observed previously. To see this I ran:</p>
<pre><code class="lang-auto">print("Fetching iterator...")
val_iter = iter(self.data_loaders["validation"])
time1 = time()
print("Done in time: ", time1 - time0)

print("Fetching data...")
data = next(val_iter)
time0 = time()
print("Done in time: ", time0 - time1)

print("Fetching data...")
data = next(val_iter)
time1 = time()
print("Done in time: ", time1 - time0)
</code></pre>
<p>Yielding:</p>
<pre><code class="lang-auto">Fetching iterator...
Done in time:  1.7992744445800781
Fetching data...
Done in time:  6.465835094451904
Fetching data...
Done in time:  0.10752415657043457
</code></pre>
<p>Which is not great, but probably tolerable. It would be nice to know what is going wrong with storing this data in the<code> __init__</code>/passing it to the workers if anyone has an idea why this lag is being produced.</p>
<p><strong>DataLoader</strong></p>
<p>Regarding the instantiation of the dataloader, the process is a little convoluted because we’re running using hydra configs and it is made in a series of calls, but I have tried to pare it down into a reasonable example.</p>
<p>We have an engine class which handles training and so has the dataloaders as attributes. In our <code>main_worker_function</code>, we set this up using the config information using:</p>
<pre><code class="lang-auto"># Configure data loaders
for task, task_config in config.tasks.items():
    if 'data_loaders' in task_config:
        engine.configure_data_loaders(config.data, task_config.data_loaders, is_distributed, config.seed)
</code></pre>
<pre><code class="lang-auto">def configure_data_loaders(self, data_config, loaders_config, is_distributed, seed):
        """
        Set up data loaders from loaders config
        """
        for name, loader_config in loaders_config.items():
            self.data_loaders[name] = get_data_loader(**data_config, **loader_config, is_distributed=is_distributed, seed=seed)
</code></pre>
<p>We use the standard DataLoader class, passing a SubsetRandomSampler on a set of indices and our own dataset.</p>
<pre><code class="lang-auto">from hydra.utils import instantiate
from torch.utils.data import DataLoader
from torch.utils.data import SubsetRandomSampler

def get_data_loader(dataset, batch_size, sampler, num_workers, is_distributed, seed, split_path=None, split_key=None, transforms=None):
    split_indices = np.load(split_path, allow_pickle=True)[split_key]
    
    sampler = SubsetRandomSampler(split_indices)
    dataset = instantiate(dataset, transforms=transforms, is_distributed=is_distributed)
    return DataLoader(dataset, sampler=sampler, batch_size=batch_size, num_workers=num_workers)
</code></pre>
<p>Hydra instantiate simply returns the object specified in the dataset (which comes from the config). I’ve checked and the dataset at this point in the code is just our dataset object.</p>
<p>For details on the dataset, refer to the code in the section above this one. I can clarify details further if this is too complicated.</p>