Unused parameters error post upgrade.

[vikigenius]

So I am using two modules like this

        self._query_embedder = BasicTextFieldEmbedder({
            'tokens': PretrainedTransformerEmbedder(transformer_model, last_layer_only=False),
        })
        self._document_embedder = BasicTextFieldEmbedder({
            'tokens': PretrainedTransformerEmbedder(transformer_model, last_layer_only=False),
        })

This worked fine in allennlp 2.3. But on a recent upgrade I get the following error:

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`, and by
making sure all `forward` function outputs participate in calculating loss.
If you already have done the above, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).
Parameters which did not receive grad for rank 2: _document_embedder.token_embedder_tokens.transformer_model.pooler.dense.bias, _document_embedder.token_embedder_tokens.transformer_model.pooler.dense.weight, _query_embedder.token_embedder_tokens.transformer_model.pooler.dense.weight, _query_embedder.token_embedder_tokens.transformer_model.pooler.dense.bias
Parameter indices which did not receive grad for rank 2: 101 102 211 212

The only relevant issue that I found is #5228

So I added the following to my trainer

  "distributed": {
    "cuda_devices": if num_gpus > 1 then std.range(0, num_gpus - 1) else 0,
    "ddp_accelerator": {
      "type": "torch",
      "find_unused_parameters": true
    } 
  },

And now it no longer shows that error. I am a bit confused as to what is happening here and why my solution worked. Any explanation would be appreciated.

[epwalsh]

Based on the error message, it sounds like the pooler layer within your _query_embedder and _document_embedder aren't part of the graph that leads to the loss, which is not surprising. Depending on which transformer model you're using you may be able to disable the pooler layer, which should avoid this error without having to set find_unused_parameters to True. With BERT, for example, see here: https://github.com/huggingface/transformers/blob/ca0b82bbd7eadd4f91e48ebedbf647d12b654f38/src/transformers/models/bert/modeling_bert.py#L864