Multiple forward functions in DP and DDP

seungjun: <p><strong>I updated the below script to reflect the interaction between different models.</strong></p>
<p>Hi,</p>
<p>I am implementing a model with multiple types of <code>forward</code> functions.<br>
An example should look like:</p>
<pre><code class="lang-python">import torch.nn as nn
from torch.nn.parallel import DataParallel, DistributedDataParallel

class Model1(nn.Module):
    def __init__(self):
        ...
    def forward(self, x):
        ...
    def forward2(self, x, y):
        ...
    def forward3(self, z, w):
        ...

# model1 = DP(Model1(), ...)
model1 = DDP(Model1(), ...)
model2 = DDP(Model2(), ...)
...

out1 = model1(x)                       # parallized
out2 = model1.module.forward2(x, y)    # not parallelized in DP and

z, w = model2(y)                       # model2 is also used somewhere else
out3 = model1.module.forward3(z, w)    # no communication in DDP

</code></pre>
<p>These forward functions are there for different purposes and are all necessary for training and inference.<br>
However, DP or DDP-wrapped models do not directly parallelize those functions other than the default <code>forward</code> that could be called without explicitly naming it.<br>
How can I parallelize other <code>forward2</code> and <code>forward3</code>?</p>
<p>Here are several candidates that I could think of for now.</p>
<ol>
<li>define a big forward function with an option flag so that it could call other functions in it.</li>
</ol>
<p>I think this would work but would cause a lot of if else statements in the functions.<br>
Also, the input argument parsing part would be cluttered.</p>
<ol start="2">
<li>register a function to DDP such that it would recognize other functions</li>
</ol>
<p>I looked into some hooks but didn’t find a way for my case.</p>
<ol start="3">
<li>create a custom <code>myDDP</code> class that inherits <code>DDP</code> and implement other functions similar to forward.</li>
</ol>
<p>This might be possible but I would need to update myDDP every time I define new functions.</p>
<p>If 2 is possible, that would be the most elegant solution and 1 could be a safe and dirty way.<br>
Are there any suggestions?</p>
<p>p.s. I checked <a href="https://discuss.pytorch.org/t/custom-methods-in-distributeddataparallel/77114/7">this thread</a> but it does not apply to me.<br>
My actual code is more complex and a more fundamental solution is necessary.</p>

cbalioglu: <p>How do you train your model locally? Do you call your forward functions in a specific order? Do they train different parts of your model?</p>

seungjun: <p>Hi <a class="mention" href="/u/cbalioglu">@cbalioglu</a></p>
<p>Thank you for looking into this issue:)</p>
<ol>
<li>training</li>
</ol>
<p>I am training an invertible network with forward/inverse path and optionally, jointly with other models.<br>
The invertible network has several other functions, too.</p>
<ol start="2">
<li>function call order</li>
</ol>
<p>I (usually) call the functions in order but not always:</p>
<pre><code class="lang-python">model1.func1
model1.func2

model2.func1    # sometimes omitted for different purposes

model1.func3    # also called outside the training loop
model1.func4
</code></pre>
<p>I am using the intermediate outputs to compute loss functions.<br>
I sometimes call func3 independently from the above flow.</p>
<ol start="3">
<li>training parts per function</li>
</ol>
<p>The training weights are partially shared across different functions while some don’t require weights.</p>

mrshenli: <p>Hey <a class="mention" href="/u/seungjun">@seungjun</a>, yep, I confirm that your 1st option should work as DDP only calls the default <code>forward()</code> function.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/a33d3d84dfcaef93c17d78cf7983d62ee01d28e3/torch/nn/parallel/distributed.py#L915-L919">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/a33d3d84dfcaef93c17d78cf7983d62ee01d28e3/torch/nn/parallel/distributed.py#L915-L919" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/a33d3d84dfcaef93c17d78cf7983d62ee01d28e3/torch/nn/parallel/distributed.py#L915-L919" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/a33d3d84dfcaef93c17d78cf7983d62ee01d28e3/torch/nn/parallel/distributed.py#L915-L919</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="915" style="counter-reset: li-counter 914 ;">
          <li>if self.device_ids:</li>
          <li>    inputs, kwargs = self.to_kwargs(inputs, kwargs, self.device_ids[0])</li>
          <li>    output = self.module(*inputs[0], **kwargs[0])</li>
          <li>else:</li>
          <li>    output = self.module(*inputs, **kwargs)</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>However, there is a caveat. DDP has internal states that requires alternating forward and backward. So if you call things like <code>foward</code>, <code>forward</code>, <code>backward</code>, DDP is likely to hang or crash.</p>

seungjun: <p>Hi <a class="mention" href="/u/mrshenli">@mrshenli</a>,</p>
<p>Thank you for confirming the 1st option and pointing to the related part of the DDP source code.<br>
I checked the DDP implementation and it seems that option 1 is the only possible way for now.</p>
<p><code>forward</code> is the only function that DDP supports safe parallelization and going for option 3 would be an adventure.</p>
<p>By the way, I’m not sure if I could avoid the function call patterns like <code>forward</code>, <code>forward</code>, <code>backward</code> you mentioned.<br>
Thank you very much and I will post here when I come up with a nice solution.</p>
<p>Best,<br>
Seungjun</p>