DDP with Gradient accumulation and clip grad norm

shivammehta007: <p>Hello,</p>
<p>I am trying to do gradient accumulation</p>
<pre><code class="lang-auto">model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, labels)       # Compute loss function
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # Backward pass
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()                           # Reset gradients tensors
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()   
</code></pre>
<p>as mentioned <a href="https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255" rel="noopener nofollow ugc">here</a></p>
<p>But my model is on multiple GPUs</p>
<p>So when my loss is output from the model how do I scale it? also I want to clip grad norms while training it. How should my training loop modify?</p>
<p>My code is this, does this looks correct?</p>
<pre><code class="lang-auto">for i, batch in enumerate(train_loader_with_distributed_sampler):

    for param_group in optimizer.param_groups:
        param_group['lr'] = learning_rate

    x, y = batch
    y_hat = model(x)
    loss = criterion(y_hat, y).mean()

    loss = loss / hparams.gradient_accumulation_step

    reduced_loss = reduce_tensor(loss.data, n_gpus).item()

    loss.backward()

    current_accumulation_run = iteration % hparams.gradient_accumulation_step + 1

    # my grad clip thres is 1.0 so it will be multiplied with 1, 2, 3, 4, 5 based on my gradient accumulation step size
    # Or maybe I don't need to manage this? 
    grad_norm = torch.nn.utils.clip_grad_norm_(
        model.parameters(), hparams.grad_clip_thresh * current_accumulation_run)

    grad_norm = grad_norm

    if (i + 1) % hparams.gradient_accumulation_step == 0:
        optimizer.step()
        model.zero_grad()

    if rank == 0:
        print("Optimizing Step")
        print("Train loss {} {:.6f} Grad Norm {:.6f}".format(
            i, reduced_loss, grad_norm))

</code></pre>

mrshenli: <aside class="quote no-group" data-username="shivammehta007" data-post="1" data-topic="115672">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/shivammehta007/48/31230_2.png" class="avatar"> shivammehta007:</div>
<blockquote>
<p>But my model is on multiple GPUs</p>
<p>So when my loss is output from the model how do I scale it?</p>
</blockquote>
</aside>
<p>You mean your gradients are on multiple GPUs, so that you cannot directly pass it to the same operator? <a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/alband">@albanD</a> does this mean the gradients need to be move to the same device first, and then calculate the scaling ratio? Or is the recommended way to do per device/model-shard scaling?</p>

mrshenli: <p>I think I misunderstood in the above comment. Would I be correct if I assume each model replica is on one GPU, and then you have DDP on top for gradient averaging?</p>

mrshenli: <p>If you have single-GPU model replica + DDP, will it be acceptable to let DDP first do gradient averaging, and then do gradient scaling/clipping independently on every process before calling <code>optimizer.step()</code>. Since DDP will make sure that all model replicas have the same gradient, their should reach the same scaling/clipping result.</p>
<p>Another thing is that, to accumulate gradients from multiple iterations, you can try using the <a href="https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync" rel="noopener nofollow ugc"><code>ddp.no_sync()</code></a>, which can help avoid unnecessary communication overheads.</p>

shivammehta007: <aside class="quote no-group" data-username="mrshenli" data-post="3" data-topic="115672" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>I think I misunderstood in the above comment. Would I be correct if I assume each model replica is on one GPU, and then you have DDP on top for gradient averaging?</p>
</blockquote>
</aside>
<p>I think this is the default behaviour of DDP no? Since it is initialised with distributed.init, and triggered by the different script where I have set the world size and everything as mentioned in the DDP basic tutorial. So let’s assume I have 4 GPUs my each model is replicated on each GPU and distributed sampler is providing each GPU with equal number of batches.</p>

mrshenli: <aside class="quote no-group" data-username="shivammehta007" data-post="5" data-topic="115672">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/shivammehta007/48/31230_2.png" class="avatar"> shivammehta007:</div>
<blockquote>
<p>I think this is the default behaviour of DDP no?</p>
</blockquote>
</aside>
<p>Yep, it’s the default behavior. Just wanna make sure this is how DDP is used in your application so that my comment can be relevant.</p>
<p>Could you please elaborate more on your requirements for gradient clipping? Do you need to do that before DDP gradient averaging comm or after?</p>
<p>If you need to before DDP comm, then you probably can use the <a href="https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.register_comm_hook" rel="noopener nofollow ugc">DDP communication hook</a>. More specifically, you can wrap the gradient bucket clipping with the allreduce communication in the hook.</p>
<p>If it is OK to do clipping after DDP comm, then you can run the clipping ops after DDP <code>backward()</code> and before optimizer step.</p>

shivammehta007: <p>I have no specific requirements actually, I just want to prevent gradient explosion in my implementation, so I was clipping the norm, but then I switched to DDP module and was accumulating the gradients as well and was not sure where to put the clipping now.</p>