DDP error due to loss only depending on forward hook's output

Jiaji_Huang: <p>I’m training a model whose loss is a function of some very intermediate output. As I import the model from another package, I don’t want to change anything in the package, e.g., modify its forward method to output the intermediate tensor I want. That can result in many changes to the package, and is not scalable.</p>
<p>So instead, I register a forward hook to the model’s submodule that produces the intermediate output. Code is lengthy, so I put here a snippet that should be sufficient to show the idea.</p>
<pre><code class="lang-auto">from the_package import MODEL
model = MODEL() # suppose I have created the model by calling that package
optimizer = Adam(model.parameters(), lr=1e-3)
# register hook, and store the intermediate output
intermediate = None
def hook_function(module, input, output):
    intermediate = output
my_submodule = locate_my_submodule(model) # some function that locates the submodule
my_submodule.register_forward_hook(hook_function)
model_output = model.forward(input_batch) # In fact, model_output is not used
loss = loss_function(intermediate, labels_or_targets)
loss.backward()
optimizer.step()
</code></pre>
<p>The code is runnable on single GPU. But if I wrap the model and optimizer with DDP on 8xV100, it fails after the 1st iteration. And I get the following error:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/8119e8111e83eabd3d2524f00d0953facf9ec623.png" data-download-href="https://discuss.pytorch.org/uploads/default/8119e8111e83eabd3d2524f00d0953facf9ec623" title="Screen Shot 2022-04-19 at 10.35.10 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8119e8111e83eabd3d2524f00d0953facf9ec623_2_690x46.png" alt="Screen Shot 2022-04-19 at 10.35.10 AM" data-base62-sha1="iq565dEN3n7Lp4dKPpo6Eikendh" width="690" height="46" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8119e8111e83eabd3d2524f00d0953facf9ec623_2_690x46.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8119e8111e83eabd3d2524f00d0953facf9ec623_2_1035x69.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/8119e8111e83eabd3d2524f00d0953facf9ec623_2_1380x92.png 2x" data-dominant-color="1D1D1D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2022-04-19 at 10.35.10 AM</span><span class="informations">1915×128 36.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Note I have set <code>find_unused_parameters=True</code> by referring to posts like <a href="https://github.com/pytorch/pytorch/issues/43259" rel="noopener nofollow ugc">here</a> , and <a href="https://github.com/pytorch/pytorch/issues/55582" rel="noopener nofollow ugc">here</a></p>
<p>One get-around might be building a “dummy” loss on the actual <code>model_output</code>. Then add this “dummy” loss to total_loss but with a weight multiplier of 0. But that can cause unwanted overhead. Is there a better solution?</p>