Repository: pytorch
Issue: [rfc] Build a debug tool to help users find which parameters are unused during DDP training
Labels: oncall: distributed,triaged

## 🚀 Feature
One of the most common errors in DDP is as follows:

```
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 (1) passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`; (2) making sure all `forward` function outputs participate in calculating loss. If you already have done the above two steps, 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).
```

This error usually comes from 2 scenarios:
1) Parameter unused in forward pass. This can be resolved with passing `find_unused_parameters=True` into DDP
2) Parameter is used in forward pass, but result is not used in loss computation, so param does not get gradient. DDP currently does not support this, but we aim to land support for this soon for static graph training.

However, in some use cases users want to debug this further and understand which parameters aren't getting gradients. There are a couple of options to already do this:
1. Let DDP crash and enable training with TORCH_DISTRIBUTED_DEBUG=INFO, this will print out unused parameter names. However, this is somewhat cumbersome as it means the user must let the application crash to get this debug information. Iterating with this sort of setup is not a great developer experience since crash logs are more cumbersome to go through.

2. Run local-only model, and manually inspect parameter grads to see which parameter didn't have their grad touched. Although, this is also a bit cumbersome as it involves writing additional code to run local model, run backwards, and look at gradients and compare them to the previous iteration.

To help resolve this issue, we are proposing an API `torch.distributed.utils.get_unused_params(model: nn.Module, loss: torch.Tensor) -> Dict[str, torch.nn.parameter])`  which returns the * fully qualified name and parameter * that did not get gradient. The implementation will be based on searching for gradient accumulators that don't show up in the backwards graph of `loss`. 

Thus, to debug these type of issues, we can just drop this line of code into a DDP loop:

```
model = DDP(....)
for i in range(x):
    out = model(inp)
    loss = out.sum()
    print(torch.distributed.utils.get_unused_params(model, loss)) # <--- one-liner drop in
    loss.backward()
    opt.step()
```


## Motivation

In many situations developers (either PT distributed developers or model authors themselves) try to debug which parameters went unused/were not used in computing loss, since they may want to fix it in their model. Internal examples of this exist and an OSS example is https://discuss.pytorch.org/t/need-help-runtimeerror-expected-to-have-finished-reduction-in-the-prior-iteration-before-starting-a-new-one/119247, and there are several other issues/forum posts reporting this issue. Providing this API will help users quickly detect the unused parameters.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23