Gathering dictionaries of DistributedDataParallel

username: jteuwen
cooked: <p>I am training a model to segment 3D images in a slice by slice-fashion. To distribute over multiple GPUs I am using <code>DistributedDataParallel</code> and I use <code>DistributedSampler</code> to split the dataset across the GPUs.</p>
<p>During prediction of new cases, I use a similar <code>Dataset</code> and <code>DataLoader</code> setup and I basically can gather a dictionary like: <code>{'filename': [(slice_no_1, slice_1_pred), (slice_no_2, slice_2_pred)], ...}</code> which I can subsequently sort on the first index to get an output. However, when I use <code>DistributedSampler</code> the <em>slices</em> are distributed along two GPUs, and I therefore end up with two dictionaries which most likely are both incomplete (one containing slices of the other and vice versa).</p>
<p>How do I gather these two dictionaries? As I preferably cast the predictions to a numpy array, it might be most convenient to gather these in the CPU memory.</p>

username: pietern
cooked: <p>Hi <a class="mention" href="/u/jteuwen">@jteuwen</a>,</p>
<p>I’m not sure I understand your issue, but I’ll give it a shot. You’re doing validation of your model and you’re using a distributed sampler on the validation set. This means you have a partial result on each process, and you’re looking to combine them into a single result, for the final accuracy numbers?</p>
<p>Or… perhaps the sampler splits the array of slices for a single filename across GPUs and you want to keep the predictions for a single filename on a single process?</p>

username: jteuwen
cooked: <p>Hi <a class="mention" href="/u/pietern">@pietern</a></p>
<p>No, training and validation is done on a slice-by-slice basis while the data are 3D MRI images. My Dataset outputs a dictionary with the data including a key which says to which file the slice belongs, and what the index of the slice is. I use the same setup for predicting new cases. However, for that I would like to recombine data again into a 3D volume. When you have one GPU that is fine: when processing a complete dataset you can combine based on the dictionary key denoting the filename and the slice number.</p>
<p>When doing this with a DistributedSampler, you have multiple python processing having a part of the dataset. In this case, even when the sampling is sequential, it can be that part of the slices end up in one process, and the others in another process. To recombine them I would need to have access in the process with rank 0 to <em>all</em> dictionaries of the other processes containing the slices.</p>
<p>Solutions I have come up with now:</p>
<ul>
<li>Dump each slice to disk and when done, combine them in process rank 0</li>
<li>Use a memorymap to do the same thing (but can do with pickled dictionaries)</li>
<li>Use something such as Redis to store the results in. Extra bonus is that it would be easier to distribute as we already use Redis.</li>
</ul>
<p>However, that seems to be quite convoluted for a reasonably simple problem. I could change the Dataset classes and the sampler specifically for this purpose, but that has the disadvantage that (1) if I change something in the dataset / dataloader I would need to change in two places, and a source for bugs (2) also tricky to implement a multiGPU model which scales well across a cluster.</p>

username: pietern
cooked: <p>I understand, thanks for the clarification.</p>
<p>You can use existing <code>torch.distributed</code> primitives <code>gather</code> or <code>all_gather</code> to get all results to a single or all processes, respectively. You say you’re outputting dictionaries, so you can’t do it with functions in core yet, and would need to serialize the dictionaries yourself. Coincidentally, <a class="mention" href="/u/mrshenli">@mrshenli</a> and I were talking about adding this to core yesterday, and he created an issue to track it: <a href="https://github.com/pytorch/pytorch/issues/23232" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/23232</a>. This doesn’t solve your issue right now though.</p>
<p>To solve this today, I would indeed write everything to a shared filesystem if you have one (with <code>torch.save</code>), probably named after the rank of the process, run a barrier to ensure all writes are done, and then <code>torch.load</code> all of them on the process where you want to collect the results.</p>

username: jteuwen
cooked: <p>Thanks for your reply - that does seem like a good addition to the code base. By the way: if I would use the <code>torch.distributed</code> primitives, then since <code>nccl</code> is the backend, wouldn’t that transfer through the GPU memory (since <code>nccl</code> does not support CPU ipc)? That might be unconvenient, also for my use case.</p>

username: pietern
cooked: <p>That’s correct. It would require serialization on the CPU side, copy to GPU memory to perform the collective, copy back to CPU memory, and then deserialization. Alternatively, you can create a separate Gloo group and use that, with <a href="https://pytorch.org/docs/stable/distributed.html#torch.distributed.new_group" rel="nofollow noopener"><code>torch.distributed.new_group</code></a>.</p>

username: solarflarefx
cooked: <p><a class="mention" href="/u/jteuwen">@jteuwen</a> I have a question for you.  Are you loading a single 3D file that then gets sliced and passed into training a model?  If this is the case one question I have is that when doing multiGPU training via DDP, could you not run into a situation where multiple processes would get different slices that originate from the same 3D file and would therefore want to access the file at the same time?</p>

username: jteuwen
cooked: <p><a class="mention" href="/u/solarflarefx">@solarflarefx</a> Apologies that I missed your question.</p>
<p>It might be possible that in this way multiple processes try to access the same volume. What I do is make a BatchSampler, which takes this into account and tries to split the volumes over the different processes instead (using the rank).</p>

username: mrs2023
cooked: <p>Hi Jteuwen,</p>
<p>I am wondering if you could resolve your issue. I have a similar problem you have described here. In my experience, different parts of my output are handled by different GPUs, and each GPU returns back a dictionary. I expected to get all the results from GPUs so, I could recombine them using the keys of dictionaries. However, the GPUs return incomplete results likely because some GPUs (processes) have finished their job sooner than others. I have used dist.all_gather_object(), but still I cannot resolve my problem.</p>

username: ptrblck
cooked: <p>You might want to add a <code>barrier</code> in case you need to synchronize all processes.</p>

username: Zuzanna_Skorniewska
cooked: <p>I am coming late to this party, however maybe this will be useful to some…</p>
<pre><code class="lang-auto">def reduce_dict(input_dict, average=True):
    """
    Args:
        input_dict (dict): all the values will be reduced
        average (bool): whether to do average or sum
    Reduce the values in the dictionary from all processes so that all processes
    have the averaged results. Returns a dict with the same fields as
    input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size &lt; 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.all_reduce(values)
        if average:
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict
</code></pre>

username: 5o1
cooked: <p>The response might be a bit late, but here are some suggestions for anyone encountering this issue in the future. There are several solutions to address this problem; please consider your specific needs:</p>
<ol>
<li>Use <code>all_gather()</code> for Tensor objects or <code>all_gather_object()</code> for dict objects in PyTorch. As most distributed frameworks use the <code>NCCL</code> backend by default, mentioning in the PyTorch official documentation, tensors must on the GPU.</li>
<li>The <code>TorchDict</code> toolkit supports using file identifiers as intermediaries to enable communication between GPUs. Additionally, you could create file identifiers in a <code>tmpfs</code>-mounted path to avoid disk I/O. This is equivalent to implementing a custom wrapper for the <code>torch.save</code>, <code>torch.load</code>.</li>
<li>Use the <code>Torch RPC</code> to globally share and ensure the uniqueness of <code>dicts</code> across ranks. However, <code>Torch RPC</code> is rarely maintained and ags behind updates for a long time in the past. Moreover, it only supports tensors on the CPU (which is the opposite of 1). This may be because <code>TorchRPC</code> uses the <code>Gloo</code> as backend. Since <code>Torch RPC</code> and <code>Torch Distribute</code> are independent of each other, they do not need to use the same communication channel. Here’s my implementation, which may not work perfectly: <a href="https://github.com/5o1/TorchUnique" class="inline-onebox" rel="noopener nofollow ugc">GitHub - 5o1/TorchUnique: Convert any serializable object into a shared object across ranks. This is useful for aggregating results from different ranks during distributed training, implemented based on Torch RPC.</a></li>
</ol>