Weird behavior when dealing with uneven inputs using the join context manager

username: Coderlemon
cooked: <h1>
<a name="background-1" class="anchor" href="#background-1"></a>Background</h1>
<p>Hi, I was trying to reproduce the tutorial given by <a href="https://tutorials.pytorch.kr/advanced/generic_join.html#what-is-join" rel="noopener nofollow ugc">https://tutorials.pytorch.kr/advanced/generic_join.html#what-is-join</a>, when I notice some weird behavior.<br>
The code is:</p>
<pre><code class="lang-python">import os
from statistics import mode
from bleach import clean
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

BACKEND = "nccl"
WORLD_SIZE = 2
NUM_INPUTS = 5

def worker(rank):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(BACKEND, rank=rank, world_size=WORLD_SIZE)

    model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank])
    # Rank 1 gets one more input than rank 0
    inputs = [torch.tensor([1]).float() for _ in range(NUM_INPUTS + rank)]

    num_inputs = 0

    # Case 1:
    with model.join():
        for input in inputs:
            num_inputs += 1
            loss = model(input)
            loss.backward()

    # Case 2:
    for input in inputs:
        num_inputs += 1
        loss = model(input)
        loss.backward()

    # Case 3:
    for input in inputs:
        num_inputs += 1
        loss = model(input)
        loss.backward()
        dist.barrier()
    
    # Case 4:
    with model.join():
        for input in inputs:
            num_inputs += 1
            loss = model(input)
            loss.backward()
            dist.barrier()
    
    # Case 5:
    for input in inputs:
        num_inputs += 1
        loss = model(input)
        loss.backward()
    dist.barrier()

    # Case 6:
    with model.join():
        for input in inputs:
            num_inputs += 1
            loss = model(input).sum()
            loss.backward()
    dist.barrier()

    print(f"Rank {rank} has exhausted all {num_inputs} of its inputs!")

def main():
    mp.spawn(worker, nprocs=WORLD_SIZE, join=True)

if __name__ == "__main__":
    main()
</code></pre>
<p>The basic idea of the code is that I spawn two processes: Rank0 and Rank1 in one Node with two GPUs, which have <code>5</code> and <code>6</code> batches separately.<br>
And I want to see how the <em>distributed</em> works when meeting uneven data.</p>
<p>When I run those cases separately, those are the results I got:</p>
<ul>
<li>Case 1: Rank0 and Rank1 are both finished.</li>
<li>Case 2: Rank0 and Rank1 are both finished.</li>
<li>Case 3: Only Rank0 finished and Rank1 is stuck.</li>
<li>Case 4: Only Rank1 finished and Rank0 is stuck.</li>
<li>Case 5: Only Rank0 finished and Rank1 is stuck.</li>
<li>Case 6: Rank0 and Rank1 are both finished.</li>
</ul>
<p>And some results are <strong>counterintuitive</strong> for me:</p>
<ul>
<li>
<strong>Case 2</strong>: I assume Rank1 should get stuck since without the <code>model.join</code>, it should not be able to deal with the extra uneven data, it should just hang there.</li>
<li>
<strong>Case 3</strong>: As I add a <code>barrier</code> compared to Case2, this time it works as I expected. But why do I have to add this <code>barrier</code>? I thought PyTorch will have a synchronization at the <code>loss.backward()</code> line, so the barrier should be redundant.</li>
<li>
<strong>Case 4</strong>: Same as my questions in Case3, I thought adding the <code>dist.barrier()</code> right after <code>loss.backward()</code> won’t have any effect.</li>
<li>
<strong>Case 5</strong>: Given that in Case2, both Rank0 and Rank1 can finish properly, I assume this barrier won’t have any effect again, but it does have.</li>
</ul>
<h1>
<a name="my-environment-2" class="anchor" href="#my-environment-2"></a>My Environment</h1>
<p>I run this script on <code>Ubuntu 7.5.0-3ubuntu1~18.04</code> with <code>Pytorch 1.9.1</code> and <code>Python 3.7.11</code>.</p>

username: agu
cooked: <p>Thanks for trying out the join context manager!</p>
<p><strong>TL;DR:</strong> The DDP join context manager uses its own collective communications, so adding <code>dist.barrier()</code> in the training loop can misalign those communications.</p>
<p><strong>Some Details</strong></p>
<ul>
<li>NCCL barrier() is implemented as an <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp#L1987-L1999" rel="noopener nofollow ugc">all-reduce of <code>Tensor([X])</code>s</a>, where X represents uninitialized data since it uses <a href="https://pytorch.org/docs/stable/generated/torch.empty.html" rel="noopener nofollow ugc"><code>empty</code></a>. <code>X</code> may be <code>0</code> or some garbage value — it is runtime dependent.</li>
</ul>
<p>Not under <code>join()</code>:</p>
<ul>
<li>DDP <code>backward()</code> uses one all-reduce per gradient bucket to synchronize gradients. For the <code>Linear(1, 1)</code> model, there is a single gradient bucket.</li>
</ul>
<p>Under <code>join()</code> (assuming not under <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync" rel="noopener nofollow ugc"><code>no_sync()</code></a>):</p>
<ul>
<li>DDP <code>forward()</code> schedules an all-reduce of <code>Tensor([1])</code> to indicate that the calling process has not yet joined. Label this <code>forward() 1</code>.</li>
<li>The <code>join()</code> context manager schedules a matching all-reduce of <code>Tensor([0])</code> to count the number of processes that have not yet joined. This is via <a href="https://github.com/pytorch/pytorch/blob/1f29b3130af218847a043e58fdc64511bbe072fe/torch/distributed/algorithms/join.py#L283" rel="noopener nofollow ugc"><code>_get_num_nonjoined_procs()</code></a>.</li>
<li>DDP <code>forward()</code> schedules an additional all-reduce of <code>Tensor([1])</code> to indicate that backward synchronization is required on this iteration. Label this <code>forward() 2</code>.</li>
<li>DDP <code>join()</code> schedules a matching all-reduce of <code>Tensor([0])</code>.</li>
<li>DDP <code>backward()</code> uses one all-reduce per gradient bucket to synchronize gradients.</li>
<li>DDP <code>join()</code> shadows every gradient all-reduce with an all-reduce of zero-valued tensors. This is via the <code>join()</code> hook (specifically, the <a href="https://github.com/pytorch/pytorch/blob/1f29b3130af218847a043e58fdc64511bbe072fe/torch/nn/parallel/distributed.py#L189" rel="noopener nofollow ugc">main hook</a>).</li>
<li>The last joining process broadcasts the final model parameters to the joined processes. However, since there may be multiple last joining processes, an authoritative rank is found via <a href="https://github.com/pytorch/pytorch/blob/1f29b3130af218847a043e58fdc64511bbe072fe/torch/nn/parallel/distributed.py#L1574" rel="noopener nofollow ugc"><code>_find_common_rank()</code></a>, which takes the max rank over the candidates using an all-reduce.</li>
</ul>
<p>Denote rank 0 and rank 1 by R0 and R1, respectively. R0 has <code>N</code> inputs, and R1 has <code>N + 1</code> inputs. For simplicity, generally assume <code>N = 2</code>.</p>
<p><strong>Case 1</strong></p>
<ul>
<li>R0 and R1 finish.</li>
</ul>
<details>
<summary>
Trace</summary>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (DDP forward() 1)</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP forward() 2)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>all-reduce (DDP forward() 1)</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP forward() 2)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td>inputs[2]</td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() check bwd grad sync)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() hook)</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (_find_common_rank())</td>
<td>all-reduce (_find_common_rank())</td>
</tr>
<tr>
<td></td>
<td>broadcast (sync final model params)</td>
<td>broadcast (sync final model params)</td>
</tr>
<tr>
<td></td>
<td>finish</td>
<td>finish</td>
</tr>
</tbody>
</table>
</div></details>
<p><strong>Case 2</strong></p>
<ul>
<li>R0 and R1 finish.</li>
</ul>
<details>
<summary>
Trace</summary>
<p>When I tried to reproduce the behavior, this is what I saw: For <code>N = 1</code>, R0 hangs, and R1 finishes. For <code>N &gt; 1</code>, both finish, but the program errors with a <code>NCCL error (Process Group destroyed on rank 1)</code>.<br>
<code>N = 1</code>:</p>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (backward)</td>
<td>all-reduce (backward)</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>finish</td>
<td>all-reduce (backward)</td>
</tr>
<tr>
<td></td>
<td></td>
<td>hang</td>
</tr>
</tbody>
</table>
</div>
<p><code>N = 2</code>:</p>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (backward)</td>
<td>all-reduce (backward)</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>all-reduce (backward)</td>
<td>all-reduce (backward)</td>
</tr>
<tr>
<td>inputs[2]</td>
<td>finish</td>
<td>all-reduce (backward)</td>
</tr>
<tr>
<td></td>
<td></td>
<td>finish with error (NCCL error: Process Group destroyed on rank 1)</td>
</tr>
</tbody>
</table>
</div>
<p>This <code>N &gt; 1</code> case is somewhat surprising. I am not entirely sure why R1 still prints, but I imagine that it can be explained with further digging.</p>
</details>
<p><strong>Case 3</strong></p>
<ul>
<li>R0 finishes. R1 hangs.</li>
</ul>
<details>
<summary>
Trace</summary>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (barrier)</td>
<td>all-reduce (barrier)</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (barrier)</td>
<td>all-reduce (barrier)</td>
</tr>
<tr>
<td>inputs[2]</td>
<td>finish</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td></td>
<td>hang</td>
</tr>
</tbody>
</table>
</div></details>
<p><strong>Case 4</strong></p>
<ul>
<li>R1 finishes. R0 hangs.</li>
</ul>
<details>
<summary>
Trace</summary>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (DDP forward() 1)</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP forward() 2)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (barrier)</td>
<td>all-reduce (barrier)</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>all-reduce (DDP forward() 1)</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP forward() 2)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (barrier)</td>
<td>all-reduce (barrier)</td>
</tr>
<tr>
<td>inputs[2]</td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
<td>all-reduce (DDP forward() 1)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() check bwd grad sync)</td>
<td>all-reduce (DDP forward() 2)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() hook)</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
<td>all-reduce (barrier)</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() check bwd grad sync)</td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
</tr>
<tr>
<td></td>
<td>all-reduce (join() _get_num_nonjoined_procs())</td>
<td>all-reduce (_find_common_rank())</td>
</tr>
<tr>
<td></td>
<td></td>
<td>error (resulting from garbage authoratitative_rank value)</td>
</tr>
</tbody>
</table>
</div>
<p>Consider the step where R0 corresponds to “check bwd grad sync” and R1 corresponds to <code>_get_num_nonjoined_procs()</code>. (Third row from the bottom)</p>
<ul>
<li>R0 contributes a 0-valued tensor to check if it needs a backward gradient synchronization.</li>
<li>R1 contributes a 0-valued tensor to check if there are any non-joined processes remaining.</li>
<li>R0 interprets the 0-valued result to mean that no backward gradient synchronization is needed.</li>
<li>R1 interprets the 0-valued result to mean that there are no remaining non-joined processes, so it is the last joining process.</li>
<li>In the next step, R0 calls <code>_get_num_nonjoined_procs()</code>. However, under the current implementation, <code>_get_num_nonjoined_procs() </code>uses the default dtype, <code>float32</code>, while <code>_find_common_rank()</code> uses <code>int64</code>. This results in a cast to floating point in the all-reduce result, and yields a garbage value as the authoritative rank for R1.</li>
</ul>
</details>
<p><strong>Case 5</strong></p>
<ul>
<li>R0 finishes. R1 hangs.</li>
</ul>
<details>
<summary>
Trace</summary>
<div class="md-table">
<table>
<thead>
<tr>
<th></th>
<th>R0</th>
<th>R1</th>
</tr>
</thead>
<tbody>
<tr>
<td>inputs[0]</td>
<td>all-reduce (DDP backward())</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td>inputs[1]</td>
<td>barrier</td>
<td>all-reduce (DDP backward())</td>
</tr>
<tr>
<td></td>
<td>finish</td>
<td>barrier</td>
</tr>
<tr>
<td></td>
<td></td>
<td>hang</td>
</tr>
</tbody>
</table>
</div></details>
<p><strong>Case 6</strong></p>
<ul>
<li>R0 and R1 finish.</li>
</ul>
<details>
<summary>
Trace</summary>
<p>This is the same as Case 1 only with an “all-reduce (barrier)” at the end for both R0 and R1.</p>
</details>
<p>As you can see, the join context manager introduces several more synchronization points (using all-reduces), so adding extraneous collective communications can lead to strange behavior. Let me know if you have any more questions.</p>

username: Coderlemon
cooked: <p>Thank you soooo much for your detailed reply and illustration, I found them very informative!</p>