Weird behavior when saving checkpoint in DDP

username: danbochman
cooked: <p>Hi,</p>
<p>I am training on 4 A100s using <code>accelerate</code> in bf16 precision and torch.compile.<br>
I am not sure if this is an <code>accelerate</code> issue or the underlying problem is some logic I don’t understand in <code>torch</code>.<br>
Training is overall stable but sometimes collapses (but doesn’t diverge) when saving a checkpoint. The collapse manifests as very low gradient norms and the optimizer skipping steps. Sometimes the training recovers after the next checkpoint save (every 5k steps in the charts). Something weird happens whenever I save a checkpoint (or whenever I call <code>accelerator.wait_for_everyone()</code>), but I have no idea what…</p>
<p>This is how I save:</p>
<pre data-code-wrap="python"><code class="lang-python">    def save(self, path, overwrite=True):
        if not self.can_checkpoint:
            return

        fs = self.fs

        assert not (fs.exists(path) and not overwrite)

        save_obj = dict(
            model=self.unwrapped_model.state_dict(),
            steps=self.steps.cpu(),
            true_steps=self.true_steps.cpu(),
            optimizer=self.optimizer.state_dict(),
        )
        if exists(self.scheduler):
            save_obj["scheduler"] = self.scheduler.state_dict()

        if self.use_ema:
            save_obj["ema"] = self.unwrapped_ema_model.state_dict()

        # save to path
        with fs.open(path, "wb") as f:
            self.accelerator.save(save_obj, f)

        self.logger.info(f"checkpoint saved to {path}")
</code></pre>
<p>I would greatly appreciate any help… and many TIA</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60dde2f28ba093d62015516bfc980a12eafc94fe.png" data-download-href="https://discuss.pytorch.org/uploads/default/60dde2f28ba093d62015516bfc980a12eafc94fe" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60dde2f28ba093d62015516bfc980a12eafc94fe_2_690x284.png" alt="image" data-base62-sha1="dOVe3UULDlTkEsxWnEwk892n2Me" width="690" height="284" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60dde2f28ba093d62015516bfc980a12eafc94fe_2_690x284.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60dde2f28ba093d62015516bfc980a12eafc94fe_2_1035x426.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60dde2f28ba093d62015516bfc980a12eafc94fe_2_1380x568.png 2x" data-dominant-color="F5F8F8"><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">1511×623 92.1 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>

username: fduwjj
cooked: <p>Looks like this is some training code from HuggingFace, maybe you want to report this to them? I don’t remember we have accelerator in our DDP.</p>

username: danbochman
cooked: <p>Hi indeed I am using the <code>accelerate</code> package. I also tried asking there.<br>
Anyway I was wondering if this kind of effect seems similar to you because at the end of the day HuggingFace/Accelerate is just a wrapper (no offense) here to DDP.</p>

username: fduwjj
cooked: <p>Do you see the same behavior when using DDP directly for training?</p>