Custom methods in DistributedDataParallel

hij: <p>I am trying to do multi gpu training with DistributedDataParallel. I wrap it around my model. However my model has a custom function that now i call by doing model.module.function(x). I was wondering if this is ok and if something bad will happen. Thanks</p>

mrshenli: <p>What does this custom function do? and when do you call this custom function? If it does not modify parameters and the autograd graph built during the forward pass, it should be OK.</p>

hij: <p>The pseudo code is something like this</p>
<pre><code class="lang-auto">output = model(input)
output2 = model(input2)
final_output = model.module.function(output, output2)
loss = loss_function(final_output)
optimizer.zero_grad()
loss.backward()
optimizer.step()
</code></pre>
<p>Would this be fine? The custom function is just a MLP to classify something. It does not change anything, but I want it to get updated when I call my optimizer.step()</p>

mrshenli: <p>If <code>model</code> is a <code>DistributedDataParallel</code> (DDP) instance, this won’t work. Because setup some internal states at the end of the forward pass, and does not work if you call <code>forward</code> twice without a backward in between.</p>
<p>However, this can be easily solve by wrapping the two forward and the <code>function</code> invocation into a wrapper model, and then pass that wrapper model to DDP, sth like:</p>
<pre><code class="lang-auto">class WrapperModel(nn.Module):
  def __init__(self, model) :
    super(WrapperModel, self).__init__()
    self.model = model

  def forward(input, input2): 
    output = model(input)
    output2 = model(input2)
    final_output = model.module.function(output, output2)
    return final_output

ddp = DistributedDataParallel(WrapperModel(model).to(device), device_ids=[device])

final_output = ddp.forward(input, input2)
loss = loss_function(final_output)
optimizer.zero_grad()
loss.backward()
optimizer.step()
</code></pre>

hij: <p>I called broadcast_buffers=False so I didnt have an issue calling forward twice. In that case, is it fine if i call my custom function the way I did and will the gradients be correct?</p>

mrshenli: <p>If the <code>model.module.function</code> is not using the parameters in the model, it should work.</p>

hij: <p>A little more details on my method. Pseudo code is</p>
<pre><code class="lang-auto">class model(nn.Module):
  def __init__(self) :
    super(model, self).__init__()
    self.encoder = Encoder()
    self.decoder = Decoder()
    self.mlp = MLP()
  def encode(self, x):
    return self.encoder(x)
  def decode(self, x): 
    return self.decoder(x)
  def classify(self, a, b)
    return self.mlp(a, b)
  def forward(self, x):
    enc = self.encode(x)
    out = self.decode(enc)
    return enc, out
# this is my main training script
enc, out = model(x)
enc2 = enc + d #d is some random perturbations
out2 = model.module.decode(enc2)
pred = model.module.classify(enc, enc2)
</code></pre>
<p>There are a bunch of other stuff, but in this scenario, my decode function is using the parameters in model? Would this be an issue? There are no errors when running.</p>

mrshenli: <p>how do yo compute the final loss (the one where <code>backward</code> is launched from)? I assume both <code>end</code> and <code>out</code> contribute to that loss? If so, this looks OK to me.</p>
<p>This should be an issue for your current use case,  but I want to mention that this probably won’t work correctly with <code>find_unused_parameters=True</code> mode. Because the <code>mlp</code> is used outside of <code>forward</code>, and DDP will find unused parameters using <code>forward</code> output. So in that mode, DDP would treat parameters in <code>mlp</code> as unused parameters although they are actually part of the autograd graph.</p>

hij: <p>my loss functions is something like</p>
<pre><code class="lang-auto">loss1 = adv_loss(out) #make output image look realistic
loss2 = adv_loss(out2)
loss3 = adv_loss(enc) #make encoding normal distributed
loss4 = adv_loss(enc2)
loss5 = l1_loss(out, x) # reconstruction loss
loss6 = l1_loss(out2, x)
loss7 = cross_entropy_loss(pred, GT)
</code></pre>
<p>I dont have find_unuse_parameters=True and have no error. If i understand what you are saying, the gradients are fine?</p>

mrshenli: <p>Yes, I think the gradients should be fine.</p>