Is it legitimate to optimize only the module part of DDP?

Fei_Liu: <p>Taking this code snippet as an example:</p>
<pre><code class="lang-python">model = Model()
ddp_model = DDP(model)

optimizer = optim.Adam(model.parameters())  # instead of ddp_model.parameters()

output = ddp_model(...)   # forward call
loss_fn(output).backward()  # backward

optimizer.step()  # This runs on `model.parameters()` (instead of `ddp_model.parameters()`). Is it problematic?
</code></pre>
<p>I tried this distributed training setting (meaning: only using DDP to do initialization, forward call while keeping the other parts only access <code>model</code>), and didn’t find any issue so far. However this is not following the routine as in the <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" rel="noopener nofollow ugc">official tutorial</a>.</p>
<p>I read from somewhere that <code>model.forward()</code> should never be called after the <code>ddp_model</code> is constructed (cannot find the reference), which makes sense as the DDP forward call takes care of certain ops needed for synchronization in the backward step. But how about the other parts? E.g. can I call the <code>optimizer.step()</code> above, as well as <code>model.train()</code> and <code>model.zero_grad()</code>? I know that <code>model.state_dict()</code> is still valid and <code>id(model) == id(ddp_model.module)</code> so there are probably many valid things to be called. Hence the question: where’s the boundary? Are there anything else “forbidden” like <code>model.forward()</code>?</p>
<p>Use case: I usually wrap the optimizer as well as many other helper/utility functions as class members of the Model itself for the ease of model version controls. So I’m wondering what are the restrictions of these helper members (e.g. init an optimizer member by <code>optim.Adam(self.parameters(), ...)</code> ) in the context of DDP.</p>

ptrblck: <p>Generally, I would of course stick to the example usage, as it would be used in other code bases and tested. I don’t have an example which would show a breaking behavior, but I also don’t know how all bells and whistles of <code>DDP</code> would interact with your approach. In particular, I don’t know if the “native mixed-precision”, fused optimizers, <code>ddp_comm_hook</code>s, etc. could break.</p>

Fei_Liu: <p>Got it. Thank you for your reply! Then is it fair to say that: other than the “mutable” operations, including:</p>
<ol>
<li><code>ddp_model.train()/.eval()</code></li>
<li>
<code>ddp_model.forward()</code> (along with the subsequent backward from the output of the forward call + some loss function)</li>
<li>
<code>optimizer.step()</code> where optimizer is constructed from <code>ddp_model.parameters()</code>
</li>
</ol>
<p>the other operations, esp. the ones that are essentially treating the original model as a const object, are still free to use without breaking DDP’s training framework? E.g.:</p>
<ol>
<li>
<code>model.state_dict()</code> (or <code>ddp_model.module.state_dict()</code>, this should be the most common one)</li>
<li>some “const” member function from the Model class e.g. some custom evaluation method called after <code>ddp_model.eval()</code> and with <code>torch.no_grad()</code> context?</li>
</ol>
<p>These statements appears true per my understanding of the DDP framework, but I think we need more clarity, in case that something is broken unintentionally. Custom Model class methods are quite common and I need to make myself clear of what can/cannot be called with the DDP framework.</p>

ptrblck: <p>There are still some caveats:</p>
<ul>
<li>You should never call the <code>.forward</code> method of any model manually, but instead the model itself (via <code>model(input)</code>) which calls the <code>__call__</code> method internally. This will make sure all hooks are properly fired before the <code>forward</code> is executed. Calling <code>model.forward(input)</code> could easily break it (standalone and DDP).</li>
<li>The <code>state_dict</code> of the DDP vs. standalone model will differ as the former will contain the <code>.module</code> keys. It’s generally a good practice to save the <code>ddp_model.module.state_dict()</code>.</li>
<li>
<code>model.eval()</code>, <code>torch.no_grad()</code> etc. will work.</li>
</ul>

Fei_Liu: <p>Understood. Thank you!</p>