Repository: horovod/horovod

Issue:
PID KeyError in broadcast_optimizer_state
**Environment:**
1. Framework: PyTorch
2. Framework version: `torch==1.4.0`, `torchvision==0.5.0`
3. Horovod version: `horovod==0.19.0`
4. MPI version: Open MPI 4.0.0
5. CUDA version: 10.0
6. NCCL version: 2.4.2-1+cuda10.0
7. Python version: 3.7.6
8. OS and version: Ubuntu 18.04
9. GCC version: 7.4.0

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?

**Bug report:**
Please describe erroneous behavior you're observing and steps to reproduce it.

When calling `broadcast_optimizer_state` we're seeing KeyErrors when trying to access the `state_dict` for apparently nonexistent pids on all processes, example stack trace attached:
```
Traceback (most recent call last):
  ...
  File ************
    HVD.broadcast_optimizer_state(optimizer, root_rank=0)
  File "/usr/local/lib/python3.7/dist-packages/horovod/torch/__init__.py", line 572, in broadcast_optimizer_state
    param_state = state_dict['state'][pid]
KeyError: 140137585983888
```
This is running on a single multi-GPU instance, and all processes are failing the same way (although with different pids). Any pointers appreciated, thanks!


Discussion:
This script seems to consistently reproduce the error in our environment:

```
import torch
import torchvision

import horovod.torch as HVD

HVD.init()
torch.cuda.set_device(HVD.local_rank())
torch.cuda.manual_seed(20)

MODEL = torchvision.models.detection.keypointrcnn_resnet50_fpn()

optimizer = torch.optim.SGD(
    MODEL.parameters(), lr=0.001, weight_decay=1e-6, momentum=0.9, nesterov=True
)

HVD.broadcast_optimizer_state(optimizer, root_rank=0)
optimizer = HVD.DistributedOptimizer(
    optimizer, named_parameters=MODEL.named_parameters(), backward_passes_per_step=1
)
```

Resolving patch:
diff --git a/horovod/torch/__init__.py b/horovod/torch/__init__.py
--- a/horovod/torch/__init__.py
+++ b/horovod/torch/__init__.py
@@ -569,6 +569,10 @@ def _from_tensor():
 
         # The params list here is ordered by the layers in the model
         for pid in group['params']:
+            if pid not in state_dict['state']:
+                # The param has not set requires_grad, so skip broadcast
+                continue
+
             param_state = state_dict['state'][pid]
             for name, p in param_state.items():
                 # Some parameter names may appear more than once, in which


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1144,14 +1144,17 @@ def forward(self, x):
 
         model = ModelNoGrad(a, b)
 
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.001, weight_decay=1e-6, momentum=0.9, nesterov=True)
         optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
 
         hvd.broadcast_parameters(model.state_dict(), root_rank=0)
         hvd.broadcast_optimizer_state(optimizer, root_rank=0)
 
+        grad = optimizer.param_groups[0]['params'][1].grad
+        bgrad = hvd.broadcast(grad, root_rank=0)
+
         assert optimizer.param_groups[0]['params'][0].grad is None
-        assert torch.all(torch.eq(optimizer.param_groups[0]['params'][1].grad, torch.zeros([4]))).item()
+        assert torch.all(torch.eq(grad, bgrad)).item()
 
     def test_broadcast_object(self):
         hvd.init()