Is it ok to use methods other than 'forward' in DDP?

jwillette: <p>I am wondering if there is anything special about the forward method which DDP must use. Sometimes it is convenient to have a model call different methods during training for different tasks to avoid bundling everything up in the forward method with more complicated logic. For example, say I have the following module,</p>
<pre><code class="lang-auto">class Model(nn.Module):
  def forward_pretrain(self, x):
    ...

  def forward_finetune(self, x):
    ...
</code></pre>
<p>Is it ok to wrap this in <code>DistributedDataParallel</code> and call <code>model.module.forward_pretrain(x)</code> instead of bundling everything into the forward function? am I ruining the functionality of DDP in any way by calling the underlying module?</p>

ptrblck: <p>I would not use custom forward methods, as these would skip hooks (the same would happen if you manually call <code>model.forward(x)</code> instead of <code>model(x)</code>).<br>
E.g. <a href="https://github.com/pytorch/pytorch/blob/1f71ac785c3c690069300fa04e59da68ae598114/torch/nn/parallel/distributed.py#L834-L846">this code</a> uses <code>forward_pre_hook</code>s to copy to lower precision and could easily break.<br>
Also, <a href="https://github.com/pytorch/pytorch/blob/1f71ac785c3c690069300fa04e59da68ae598114/torch/nn/parallel/distributed.py#L1402-L1403">here</a> the <code>self.module</code> is called, which will internally call into <code>__call__</code> and then <code>forward</code>. Using your custom forward methods would break this again.</p>
<p>However, it should be possible to call your custom forward methods from the actual <code>forward</code> method in case this would keep your code cleaner.</p>

HELLORPG: <p>I think, if you have a DDP model <code>A</code>, and:</p>
<pre><code class="lang-python">class A(nn.Module):
    ...

model = A()
model = DDP(model)
</code></pre>
<p>Then you must use <code>output = model(input)</code> at least once, no matter what in this forward function. OR it may cause some problems.<br>
Like I have subsequent code:</p>
<pre><code class="lang-python"># A ResNet18 demo for MNIST:
class ResNet18(nn.Module):
    def __init__(self, num_classes: int):
        super(ResNet18, self).__init__()

        if is_main_process():
            self.model = torchvision.models.resnet18(weights=ResNet18_Weights.DEFAULT)
        else:
            self.model = torchvision.models.resnet18(weights=None)

        num_features = self.model.fc.in_features
        self.model.fc = nn.Linear(num_features, num_classes)

        return

    def forward(self, x):
        # return self.model(x)
        return x.detach()    # this forward function does not process any computing

    def fake_forward(self, x):
        return self.model(x)

def build(config: dict):
    return ResNet18(
        num_classes=10
    )
</code></pre>
<p>And in my training function, if I:</p>
<pre><code class="lang-python">outputs = model.module.fake_forward(images.to(device))
</code></pre>
<p>The parameters of my model will be inconsistent between different GPUs (nodes).</p>
<p>But If I:</p>
<pre><code class="lang-python">images = model(images)    # although do noting!
# or images = model.forward(images)
outputs = model.module.fake_forward(images.to(device))
</code></pre>
<p>The parameters will be consistent.</p>
<p>Thus, I think some code deep in the <code>forward()</code> function is the key to DDP parameter sync.</p>