Checkpointing DDP.module instead of DDP itself

username: mazabou3
cooked: <p>I am using DDP to distribute training across multiple gpu.</p>
<pre><code class="lang-auto">model = Net(...)
ddp_model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 
ddp_model = DDP(ddp_model, device_ids=[gpu], find_unused_parameters=True)
</code></pre>
<p>When checkpointing, is it ok to save <code>ddp_model.module</code> instead of <code>ddp_model</code>? I need to be able to use the checkpoint for 1. evaluation using a single gpu 2. resume training with multiple gpus</p>
<pre><code class="lang-auto">def save_ckpt(ddp_model, path):
    state = {'model': ddp_model.module.state_dict(),
             'optimizer': optimizer.state_dict(),
            }
    torch.save(state, path)

def load_ckpt(path, distributed):
    checkpoint = torch.load(path, map_location=map_location)
    model = Net(...)
    optimizer = ...
    model.load_state_dict(checkpoint['model'], strict=False)
    optimizer.load_state_dict(checkpoint['optimizer'])
    if distributed:
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 
        model = DDP(model, device_ids=[gpu], find_unused_parameters=True)
    return model
</code></pre>
<p>I am not sure how this would behave, does manipulating ddp_model.module in general break things?<br>
Similarly what if the number of gpus changes, would that impact  the optimizer for example? does it need to be reinitialized? Thank you!</p>

username: mrshenli
cooked: <aside class="quote no-group" data-username="mazabou3" data-post="1" data-topic="115714">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/a5b964/48.png" class="avatar"> mazabou3:</div>
<blockquote>
<p>When checkpointing, is it ok to save <code>ddp_model.module</code> instead of <code>ddp_model</code>? I need to be able to use the checkpoint for 1. evaluation using a single gpu 2. resume training with multiple gpus</p>
</blockquote>
</aside>
<p>Yep, this is actually the recommended way:</p>
<ol>
<li>On save, use one rank to save <code>ddp_model.module</code> to checkpoint.</li>
<li>On load, first use the checkpoint to load a local model, and then wrap the local model instance with DDP on all processes (i.e., sth like <code>DistributedDataParallel(local_model, device_ids=[rank])</code>).</li>
</ol>

username: Brando_Miranda
cooked: <p>I’d like to double-check because I see multiple question and <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">the official doc</a> is confusing to me. I want to do:</p>
<ol>
<li>resume from a checkpoint to continue training on multiple gpus</li>
<li>save checkpoint correctly during training with multiple gpus</li>
</ol>
<p>For that my guess is the following:</p>
<ol>
<li>to do 1 we have all the processes load the checkpoint from the file, then call <code>DDP(mdl)</code> for each process. I assume the checkpoint saved a <code>ddp_mdl.module.state_dict()</code>.</li>
<li>to do 2 simply check who is rank = 0 and have that one do the torch.save({‘model’: ddp_mdl.module.state_dict()})</li>
</ol>
<p>Is this correct?</p>
<p>I am not sure why there are so many <a href="https://discuss.pytorch.org/t/ddp-and-gradient-checkpointing/132244">other posts</a> (<a href="https://github.com/pytorch/pytorch/issues/23138" rel="noopener nofollow ugc">also this</a>) or what <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">the official doc</a> is talking about. I don’t get why I’d want to optimize writing (or reading) once given that I will train my model from either 100 epochs to even 600K iterations…that one write seems useless to me.</p>
<hr>
<p>I think the correct code is this one:</p>
<pre><code class="lang-auto">def save_ckpt(rank, ddp_model, path):
    if rank == 0:
        state = {'model': ddp_model.module.state_dict(),
             'optimizer': optimizer.state_dict(),
            }
        torch.save(state, path)

def load_ckpt(path, distributed, map_location=map_location=torch.device('cpu')):
    # loads to
    checkpoint = torch.load(path, map_location=map_location)
    model = Net(...)
    optimizer = ...
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    if distributed:
        model = DDP(model, device_ids=[gpu], find_unused_parameters=True)
    return model
</code></pre>
<p>also, I added more details here: <a href="https://stackoverflow.com/questions/70386800/what-is-the-proper-way-to-checkpoint-during-training-when-using-distributed-data" class="inline-onebox" rel="noopener nofollow ugc">python - What is the proper way to checkpoint during training when using distributed data parallel (DDP) in PyTorch? - Stack Overflow</a></p>