Cuda synchronize between loss.backward()/clip_grad_norm_

username: YashasSamaga
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/0/2049955819e8c2ae7458e2917403f2009873b250.png" data-download-href="https://discuss.pytorch.org/uploads/default/2049955819e8c2ae7458e2917403f2009873b250" title="Screenshot from 2022-05-30 13-13-49"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/2049955819e8c2ae7458e2917403f2009873b250_2_690x131.png" alt="Screenshot from 2022-05-30 13-13-49" data-base62-sha1="4BCVHE7qPzZImPNOowE9QffGITe" width="690" height="131" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/2049955819e8c2ae7458e2917403f2009873b250_2_690x131.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/0/2049955819e8c2ae7458e2917403f2009873b250_2_1035x196.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/0/2049955819e8c2ae7458e2917403f2009873b250.png 2x" data-dominant-color="E3E7E2"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2022-05-30 13-13-49</span><span class="informations">1330×253 28.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">for _, (product, target_reactants) in enumerate(tqdm(self.train_loader, desc="Epoch progress", unit=" batches", leave=False)):               
                product, target_reactants = product.to(self.device, non_blocking=True), target_reactants.to(self.device, non_blocking=True)

                predicted_reactants = self(product, target_reactants, 0.5)
                target_reactants = target_reactants[1:]

                vocab_size = predicted_reactants.shape[-1]

                predicted_reactants = predicted_reactants.view(-1, vocab_size)
                target_reactants = target_reactants.view(-1)

                torch.cuda.nvtx.range_push("criterion")
                loss = self.criterion(predicted_reactants, target_reactants)
                torch.cuda.nvtx.range_pop()

                torch.cuda.nvtx.range_push("loss.detach()")
                accumlated_train_loss += loss.detach()
                torch.cuda.nvtx.range_pop()

                torch.cuda.nvtx.range_push("zero_grad")
                self.optimizer.zero_grad()
                torch.cuda.nvtx.range_pop()

                torch.cuda.nvtx.range_push("loss_backward")
                loss.backward()
                torch.cuda.nvtx.range_pop()

                torch.cuda.nvtx.range_push("clip_grad_norm")
                torch.nn.utils.clip_grad_norm_(self.parameters(), 0.1)
                torch.cuda.nvtx.range_pop()

                torch.cuda.nvtx.range_push("optimizer_step")
                self.optimizer.step()
                torch.cuda.nvtx.range_pop()
</code></pre>
<p>Why is there a synchronization in the <code>loss_backward</code> range in the profiler output? This shows up only once every 3-4 iterations.</p>

username: ptrblck
cooked: <p>Activate the backtraces in Nsight Systems which should then show where the stream synchronization in called from and needed.</p>