Module.state_dict() is wrong when using DataParallel

username: LuChengTHU
cooked: <p>I have a module like this:</p>
<pre><code class="lang-auto">class Block(nn.Module):
    def __init__(self, net):
        super(Block, self).__init__()
        self.net = net
        self.net_copy = copy.deepcopy(net)

    def forward(self, x):
        self.net_copy.load_state_dict(self.net.state_dict())
        return self.net(x)
</code></pre>
<p>The <code>net</code> is an nn.Sequential() module. When I use Pytorch&gt;=1.5 and use nn.DataParallel in multi-GPUs, It shows that net_copy.state_dict().keys() is different with net.state_dict().keys(). However, when I use Pytorch==1.4 or single-GPU, this problem doesn’t appear. How can I make sure that <code>net</code> and <code>net_copy</code> is exactly the same?</p>

username: mrshenli
cooked: <p>This is probably due to this PR: <a href="https://github.com/pytorch/pytorch/pull/33907" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/33907</a></p>
<p>In v1.5, parameters on replicated models are no longer considered as leaves, as they shouldn’t be. If you really need to access those replicated parameters, you probably can get them from <code>_former_parameters</code> and manually add them into the <code>stat_dict</code>?</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/c93e96fbd9903e576c6c1aa2fe12d8d548ae2d5b/torch/nn/parallel/replicate.py#L148" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/c93e96fbd9903e576c6c1aa2fe12d8d548ae2d5b/torch/nn/parallel/replicate.py#L148" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/c93e96fbd9903e576c6c1aa2fe12d8d548ae2d5b/torch/nn/parallel/replicate.py#L148</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="138" style="counter-reset: li-counter 137 ;">
<li>            replica._parameters[key] = None</li>
<li>    else:</li>
<li>        param_idx = param_indices[param]</li>
<li>        for j in range(num_replicas):</li>
<li>            replica = module_copies[j][i]</li>
<li>            param = param_copies[j][param_idx]</li>
<li>            # parameters in replicas are no longer leaves,</li>
<li>            # so setattr them as non-parameter attributes</li>
<li>            setattr(replica, key, param)</li>
<li>            # expose the parameter for DDP</li>
<li class="selected">            replica._former_parameters[key] = param</li>
<li>for key, buf in module._buffers.items():</li>
<li>    if buf is None:</li>
<li>        for j in range(num_replicas):</li>
<li>            replica = module_copies[j][i]</li>
<li>            replica._buffers[key] = None</li>
<li>    else:</li>
<li>        if buf.requires_grad and not detach:</li>
<li>            buffer_copies = buffer_copies_rg</li>
<li>            buffer_idx = buffer_indices_rg[buf]</li>
<li>        else:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>cc <a class="mention" href="/u/ngimel">@ngimel</a> please correct me if I am wrong. And any thoughts on whether we should make <code>state_dict()</code> consistent between v1.4 vs v1.5?</p>

username: aashaka
cooked: <p>In order to access the <code>_former_parameters</code>, we would need to access replica, right? Can you help me figure out how to access <code>_former_parameters</code> in OP’s example?<br>
Or how to recreate state dict in some other manner?</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/aashaka">@aashaka</a></p>
<p>Below is the implementation of the <code>DataParallel.forward</code> method. It basically calls <code>replicas[i].forward(inputs[i], ...)</code>. So during execution, the <code>self</code> variable in the <code>forward</code> function is the replica. Hence, you can use <code>self._former_parameters</code> to access the field in <code>forward</code> function.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/c3466dabaae9328b207804afb043b7b519f64825/torch/nn/parallel/data_parallel.py#L147-L162" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/c3466dabaae9328b207804afb043b7b519f64825/torch/nn/parallel/data_parallel.py#L147-L162" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/c3466dabaae9328b207804afb043b7b519f64825/torch/nn/parallel/data_parallel.py#L147-L162</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="147" style="counter-reset: li-counter 146 ;">
<li>def forward(self, *inputs, **kwargs):</li>
<li>    if not self.device_ids:</li>
<li>        return self.module(*inputs, **kwargs)</li>
<li>
</li><li>    for t in chain(self.module.parameters(), self.module.buffers()):</li>
<li>        if t.device != self.src_device_obj:</li>
<li>            raise RuntimeError("module must have its parameters and buffers "</li>
<li>                               "on device {} (device_ids[0]) but found one of "</li>
<li>                               "them on device: {}".format(self.src_device_obj, t.device))</li>
<li>
</li><li>    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)</li>
<li>    if len(self.device_ids) == 1:</li>
<li>        return self.module(*inputs[0], **kwargs[0])</li>
<li>    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])</li>
<li>    outputs = self.parallel_apply(replicas, inputs, kwargs)</li>
<li>    return self.gather(outputs, self.output_device)</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: aashaka
cooked: <p>I managed to recreate the state_dict using code similar to <a href="https://github.com/pytorch/pytorch/blob/f11f9a8c1f200c0798dd5043a958c563dfb00777/torch/nn/modules/module.py#L904" rel="noopener nofollow ugc">state_dict</a>. Thanks for your help.</p>
<p>I noticed that _former_parameters exists in 1.5.1 but not in 1.5.0. It seems tricky to get the parameters in 1.5.0 if we do not know the names of the parameters in advance (but still possible since we are setting attr). Any suggestions for this?</p>

username: mrshenli
cooked: <p>Hey <a class="mention" href="/u/aashaka">@aashaka</a>, yep, we added <code>_former_parameters</code> after v1.5 to fix the regression caused on <a href="https://github.com/pytorch/pytorch/pull/33907" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/pull/33907</a>.</p>
<p>If this has become very inconvenient for you, I would suggest switch to DistributedDataParallel. There are more discussions here: <a href="https://github.com/pytorch/pytorch/issues/36268" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/36268</a></p>

username: aashaka
cooked: <p>Thanks a lot. I have one last question. Like the OP, I need to recreate the state dict every time in the forward pass. I see about 8x increase in training time when compared to original PyTorch DataParallel. Any ideas why this might be the case?</p>
<pre><code class="lang-auto">def create_state_dict_new(main_module):
    state_dict_data = OrderedDict()

    def state_dict_recursion(this_module, state_dict_data, prefix=''):
        if hasattr(this_module,"_former_parameters"):
            for name, param in this_module._former_parameters.items():
                if param is not None:
                    state_dict_data[prefix + name] = param
        for name, buf in this_module._buffers.items():
            if buf is not None:
                state_dict_data[prefix + name] = buf
        for name, module in this_module._modules.items():
            if module is not None:
                state_dict_recursion(module, state_dict_data, prefix + name + '.')
    state_dict_recursion(main_module._modules['model'], state_dict_data)
    return state_dict_data

class ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super(ModelWrapper, self).__init__()
        self.model = model
    def forward(self, x):
        state_list = create_state_dict_new(self)
        return model(x)

model = torch.nn.DataParallel(ModelWrapper(model))
</code></pre>

username: mrshenli
cooked: <p>Could you please measure the time spent on the <code>create_state_dict_new</code>?</p>
<p>The forward function will be launched in each thread. If you have 4 GPUs, it means that there will be 4 threads executing that <code>create_state_dict_new</code> independently. However, due to Python GIL, the 4 threads cannot run the function concurrently, which would further exacerbate the delay.</p>