Saving and loading optimizers in Distributed Data Parallel situations

username: shubhvachher
cooked: <p>As is given <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-torch-nn-dataparallel-models" rel="noopener nofollow ugc">here</a>:</p>
<blockquote>
<p><code>torch.nn.DataParallel</code> is a model wrapper that enables parallel GPU utilization. To save a <code>DataParallel</code> model generically, save the <code>model.module.state_dict()</code> . This way, you have the flexibility to load the model any way you want to any device you want.</p>
</blockquote>
<p>Considering the discussion just above this, of saving GPU models and loading on CPU etc. , I’m guessing this line refers to the data distributed models being on any one of the available GPUs and the <code>model.module</code> being the underlying module part that somehow will be device agnostic. (Please correct me here)</p>
<p>That being said, what happens to the optimizer internal state variables (mentioned <a href="https://discuss.pytorch.org/t/clarification-on-re-initializing-optimizer-in-every-epoch/48997/4">here</a>). They too would be on whichever GPU rank they were saved from. Is this issue just solved by using the <code>map_location</code> argument to <code>torch.load()</code>? If so, why the special treatment for the DistributedDataModel (<code>model.module.state_dict()</code> used while saving)?</p>

username: pietern
cooked: <aside class="quote no-group" data-username="shubhvachher" data-post="1" data-topic="49036">
<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/shubhvachher/48/29925_2.png" class="avatar"> shubhvachher:</div>
<blockquote>
<p>Considering the discussion just above this, of saving GPU models and loading on CPU etc. , I’m guessing this line refers to the data distributed models being on any one of the available GPUs and the <code>model.module</code> being the underlying module part that somehow will be device agnostic. (Please correct me here)</p>
</blockquote>
</aside>
<p>That’s correct. If you try to serialize the <code>nn.DataParallel</code> module itself then it contains the list of devices you parallelize for, the dimension to split the input batch on, etc. When you serialize the inner module then none of that is included and is up to you to do again (or not) after you load it.</p>
<aside class="quote no-group" data-username="shubhvachher" data-post="1" data-topic="49036">
<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/shubhvachher/48/29925_2.png" class="avatar"> shubhvachher:</div>
<blockquote>
<p>That being said, what happens to the optimizer internal state variables (mentioned <a href="https://discuss.pytorch.org/t/clarification-on-re-initializing-optimizer-in-every-epoch/48997/4">here </a>). They too would be on whichever GPU rank they were saved from. Is this issue just solved by using the <code>map_location</code> argument to <code>torch.load()</code> ? If so, why the special treatment for the DistributedDataModel ( <code>model.module.state_dict()</code> used while saving)?</p>
</blockquote>
</aside>
<p>Any optimizer state will likely include the devices that those state variables live on. You’re correct to say you can use <code>map_location</code> to remap at load time. Alternatively, you can copy the optimizer state to CPU first, then serialize, and then not worry about it at load time. What special treatment exactly are you talking about?</p>

username: shubhvachher
cooked: <aside class="quote no-group" data-username="pietern" data-post="2" data-topic="49036">
<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/pietern/48/4981_2.png" class="avatar"> pietern:</div>
<blockquote>
<p>the <code>nn.DataParallel</code> module itself then it contains the list of devices you parallelize for, the dimension to split the input batch on, etc. When you serialize the inner module then none of that is included and is up to you to do again (or not) after you load it.</p>
</blockquote>
</aside>
<p>Nothing more! My model is training but the error doesn’t seem to be coming down… I was just exploring the possibilities… This clears it up! Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: pietern
cooked: <p>Glad to help. Good luck figuring it out.</p>