How to handle scheduler during DDP training?

username: mariosconsta
cooked: <p>I have this function here for a training epoch:</p>
<pre><code class="lang-auto">def _run_epoch(self, epoch):
        self.train_data.sampler.set_epoch(epoch)
        self.model.train()

        assert self.model.training == True, "Model is in Eval mode while training"

        for idx, (_, img, fidt_map, point_map) in enumerate(self.train_data):
            img = img.to(self.gpu_id)
            fidt_map = fidt_map.type(torch.FloatTensor).unsqueeze(1).to(self.gpu_id)
            point_map = point_map.type(torch.FloatTensor).unsqueeze(1).to(self.gpu_id)

            output = self.model(img)

            loss = self.binary_fl(output, point_map)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
</code></pre>
<p>And this training loop:</p>
<pre><code class="lang-auto">    def train(self, max_epochs: int):
        for epoch in range(self.start_epoch, max_epochs):
            self._run_epoch(epoch=epoch)
            self.scheduler.step()
            self.writer.add_scalar(
                "Learning Rate", self.optimizer.param_groups[0]["lr"], epoch
            )
</code></pre>
<p>This is the scheduler:</p>
<pre><code class="lang-auto">    scheduler = CosineAnnealingLR(
        optimizer, T_max=args["epochs"], eta_min=args["eta_min"], verbose=False
    )
</code></pre>
<p>The reason I set T_max equal to epochs is because I want this curve:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/9/09b6457d37587b4e9a3b473dc0e046b56e4f7cee.png" data-download-href="https://discuss.pytorch.org/uploads/default/09b6457d37587b4e9a3b473dc0e046b56e4f7cee" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/09b6457d37587b4e9a3b473dc0e046b56e4f7cee_2_690x213.png" alt="image" data-base62-sha1="1nUOljGws7JOXh19mZx3Cr4H4Cy" width="690" height="213" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/09b6457d37587b4e9a3b473dc0e046b56e4f7cee_2_690x213.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/09b6457d37587b4e9a3b473dc0e046b56e4f7cee_2_1035x319.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/09b6457d37587b4e9a3b473dc0e046b56e4f7cee_2_1380x426.png 2x" data-dominant-color="323232"><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">image</span><span class="informations">1534×475 18 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><br>
The above LR curve is from using DP training. When I switch to DDP this happens:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/8153ec2aaa28278afef5bffd1fc1df8cd5e91371.png" data-download-href="https://discuss.pytorch.org/uploads/default/8153ec2aaa28278afef5bffd1fc1df8cd5e91371" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8153ec2aaa28278afef5bffd1fc1df8cd5e91371_2_690x202.png" alt="image" data-base62-sha1="is5oBzogZaVHb4tvN1YAlM7nMbf" width="690" height="202" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8153ec2aaa28278afef5bffd1fc1df8cd5e91371_2_690x202.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8153ec2aaa28278afef5bffd1fc1df8cd5e91371_2_1035x303.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8153ec2aaa28278afef5bffd1fc1df8cd5e91371_2_1380x404.png 2x" data-dominant-color="313233"><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">image</span><span class="informations">1535×451 21.5 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><br>
Which is really weird.</p>
<p>Check these values here:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/7/4/74cfc0047594d2fc42c95e39a7db7824c5cf99d0.png" alt="image" data-base62-sha1="gFmujRWA5Y5ltVIJn6iHmpNrylq" width="368" height="87"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/9/69271dbf9593e7f823329a15d9ab0e00e3453357.png" alt="image" data-base62-sha1="f0dWjcttePiZOn1jTWetChXB9wb" width="383" height="109"><br>
The learning rate at epoch 660 goes crazy…</p>
<p>Does anyone know why this is happening?</p>