Understanding torch.distributed.DistributedDataParallel from a torchvision's reference example

f10w: <p>I would like to ask some questions regarding the DDP code used in the <code>torchvision</code>'s <a href="https://github.com/pytorch/vision/blob/master/references/classification/train.py" rel="nofollow noopener">reference example on classification</a>. An example of using this script is given as follows, on a machine with 8 GPUs:</p>
<pre><code class="lang-bash">python -m torch.distributed.launch --nproc_per_node=8 --use_env train.py --model resnext50_32x4d --epochs 100
</code></pre>
<p>My first question concerns the saving and loading of checkpoints.</p>
<p><a href="https://github.com/pytorch/vision/blob/d6ee8757eca7b74b98e5f0d434a565eb7b1c410b/references/classification/train.py#L211" rel="nofollow noopener">This</a> is how a checkpoint is saved in the script:</p>
<pre><code class="lang-auto">checkpoint = {
    'model': model_without_ddp.state_dict(),
    'optimizer': optimizer.state_dict(),
    'lr_scheduler': lr_scheduler.state_dict(),
    'epoch': epoch,
    'args': args}
utils.save_on_master(
    checkpoint,
    os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
utils.save_on_master(
    checkpoint,
    os.path.join(args.output_dir, 'checkpoint.pth'))
</code></pre>
<p>But in the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="nofollow noopener">DDP tutorial</a>, it seems necessary that <code>torch.distributed.barrier()</code> is called somewhere:</p>
<pre><code class="lang-auto"># Use a barrier() to make sure that process 1 loads the model after process 0 saves it.
dist.barrier()
...
# Use a barrier() to make sure that all processes have finished reading the checkpoint
dist.barrier()
</code></pre>
<p>Why is <code>dist.barrier()</code> not necessary in the above reference example?</p>
<p>My second question is about the validation stage.</p>
<p><a href="https://github.com/pytorch/vision/blob/d6ee8757eca7b74b98e5f0d434a565eb7b1c410b/references/classification/train.py#L204" rel="nofollow noopener">This</a> is how it’s done in the script:</p>
<pre><code class="lang-python">for epoch in range(args.start_epoch, args.epochs):
    if args.distributed:
        train_sampler.set_epoch(epoch)
    train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args.print_freq, args.apex)
    lr_scheduler.step()
    evaluate(model, criterion, data_loader_test, device=device)
</code></pre>
<p>Doesn’t this mean that the <code>evaluate()</code> function is called on <strong>all</strong> the processes (i.e. all the GPUs in this case)? Shouldn’t we rather do something like this:</p>
<pre><code class="lang-auto">for epoch in range(args.start_epoch, args.epochs):
    if args.distributed:
        train_sampler.set_epoch(epoch)
    train_one_epoch(model, criterion, optimizer, data_loader, device, epoch, args.print_freq, args.apex)
    lr_scheduler.step()
    if torch.distributed.get_rank() == 0: # master
        evaluate(model, criterion, data_loader_test, device=device)
        # save checkpoint here as well
</code></pre>
<p>But then, again, shouldn’t we wait, using <code>dist.barrier()</code>, for all the processes to finish the computations and for the master to gather the gradients, before evaluating the model?</p>
<p>Thank you very much in advance for your help!</p>

mrshenli: <blockquote>
<p>Why is  <code>dist.barrier()</code>  not necessary in the above reference example?</p>
</blockquote>
<p>IIUC, the torchvision example only saves the checkpoint to file in epoch but is not reading from it unless it is recovering from a crash? In that case it is not a hard requirement to perform the barrier, because</p>
<ol>
<li>it does not need to ensure non-master processes are reading stale checkpoints.</li>
<li>non-master processes will just block on DDP backward (AllReduce) and waiting for the master to join.</li>
</ol>
<blockquote>
<p>Doesn’t this mean that the  <code>evaluate()</code>  function is called on  <strong>all</strong>  the processes (i.e. all the GPUs in this case)?</p>
</blockquote>
<p>cc <a class="mention" href="/u/fmassa">@fmassa</a> for this implementation</p>
<blockquote>
<p>But then, again, shouldn’t we wait, using  <code>dist.barrier()</code> , for all the processes to finish the computations and for the master to gather the gradients, before evaluating the model?</p>
</blockquote>
<p>The gradients are synchronized in DDP backward using AllReduce operations. So, there is no need to add another barrier here to do that. As soon as <code>loss.backward()</code> returns, the local gradients should be representing the global average. However, it might need a <code>barrier</code> here for a different reason. If the <code>evaluate()</code> step takes too long, non-master processes could timeout on <code>AllReduce</code>. If that happens, <code>barrier</code> might help.</p>

f10w: <p><a class="mention" href="/u/mrshenli">@mrshenli</a> Thanks for your prompt reply!</p>
<blockquote>
<p>IIUC, the torchvision example only saves the checkpoint to file in epoch but is not reading from it unless it is recovering from a crash?</p>
</blockquote>
<p>Well, the script has a resume code as follows:</p>
<pre><code class="lang-auto">if args.resume:
    checkpoint = torch.load(args.resume, map_location='cpu')
    model_without_ddp.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
    args.start_epoch = checkpoint['epoch'] + 1
</code></pre>
<p>so I guess it does what you’ve described, which is a usual scenario. But then are you saying that the <code>demo_checkpoint()</code> example given in the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">tutorial</a> handles a different scenario (other than resuming a training)?</p>
<blockquote>
<p>The gradients are synchronized in DDP backward using AllReduce operations. So, there is no need to add another barrier here to do that. As soon as  <code>loss.backward()</code>  returns, the local gradients should be representing the global average.</p>
</blockquote>
<p>Thanks! This has cleared up a lot of things for me. I guess the gradients are averaged? In this case, it seems that the learning rate should be scaled up by the number of GPUs.</p>
<p>Besides, the tutorial also notes that</p>
<blockquote>
<p>if training starts from random parameters, you might want to make sure that all DDP processes use the same initial values. Otherwise, global gradient synchronizes will not make sense.</p>
</blockquote>
<p>but I don’t see this being taken into account anywhere in the reference script.</p>

mrshenli: <blockquote>
<p>Well, the script has a resume code as follows:</p>
</blockquote>
<p>Not 100% sure about how this example would be used. <a class="mention" href="/u/fmassa">@fmassa</a> would know more.  Given the code, it looks the resume mode is designed for starting from pre-trained models or resume from crash. In these cases, the checkpoint file is ready before launching the script, so it should be fine.</p>
<blockquote>
<p>you saying that the  <code>demo_checkpoint()</code>  example given in the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">tutorial</a> handles a different scenario</p>
</blockquote>
<p>It is not targeting any specific use case, just wanted to make sure the example code can run as is. The main information it tries to convey is that applications need to make sure checkpoints are ready before loading them. We previous saw users running into weird errors caused by reading too soon.</p>
<blockquote>
<p>if training starts from random parameters, you might want to make sure that all DDP processes use the same initial values. Otherwise, global gradient synchronizes will not make sense.</p>
</blockquote>
<p>DDP handles this by broadcasting model weights from rank 0 to others at construction time. However, if the application modified model weights after constructing DDP and if that resulted in inconsistent weights across processes, DDP won’t be able to recover, as the broadcast only happens once in ctor.</p>

f10w: <p><a class="mention" href="/u/mrshenli">@mrshenli</a> Great. Thank you so much for the explanations!<br>
I hope <a class="mention" href="/u/fmassa">@fmassa</a> could join the discussion and clarify the points for which you mentioned him earlier, especially the one related to <code>evaluate()</code>. I tested the code and it seems that this function is called only once across the processes.</p>
<p>The training loop looks like this:</p>
<pre><code class="lang-auto">for epoch in range(args.start_epoch, args.epochs):
    if args.distributed:
        train_sampler.set_epoch(epoch)
    train_one_epoch(...)
    lr_scheduler.step()
    evaluate(model, criterion, data_loader_test, device=device)
</code></pre>
<p>where the evaluation function looks like:</p>
<pre><code class="lang-auto">def evaluate(model, criterion, data_loader, device, print_freq=100):
    model.eval()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'
    with torch.no_grad():
        for image, target in metric_logger.log_every(data_loader, print_freq, header):
            ...
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()

    print(' * Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'
          .format(top1=metric_logger.acc1, top5=metric_logger.acc5))
    return metric_logger.acc1.global_avg
</code></pre>
<p>This function has some <code>print()</code> to display the accuracy. In my experiment, that string is only displayed once, which means the function is called only once. Why? There is no <code>is_main_process()</code> check. Why isn’t this function called on all processes? I’m confused…</p>

f10w: <p>I’ve just realized that we shouldn’t wrap the evaluation phase inside <code>if torch.distributed.get_rank() == 0:</code>, because <code>data_loader_test</code> also splits the data across all the processes.</p>
<p>And for the <code>print()</code> part, <a href="https://github.com/pytorch/vision/blob/d6ee8757eca7b74b98e5f0d434a565eb7b1c410b/references/classification/utils.py#L194" rel="nofollow noopener">this code</a> explains why the message is displayed only once.</p>
<p>So now I understand what <a class="mention" href="/u/fmassa">@fmassa</a> did. Thanks.</p>