Repository: horovod/horovod

Issue:
pytorch + horovod 0.15.1 distributed optimizer not working anymore
I just upgraded horovod 0.15.0 -> 0.15.1 on a ubuntu image `4.4.0-137-generic #163-Ubuntu SMP Mon Sep 24 13:14:43 UTC 2018 x86_64 x86_64 x86_64 GNU/Linux`. When using the DistributedOptimizer from horovod.torch I now encounter the error

```
Traceback (most recent call last):
  File "train.py", line 641, in <module>
    train_images(hps)
  File "train.py", line 444, in train_images
    train_step(batch, batch_idx, epoch, hps, model, opt, train_logger)
  File "train.py", line 457, in train_step
    opt.step()
  File "/opt/conda/lib/python3.6/site-packages/horovod/torch/__init__.py", line 97, in step
    return super(self.__class__, self).step(closure)
  File "/opt/conda/lib/python3.6/site-packages/torch/optim/adamax.py", line 75, in step
    exp_avg.mul_(beta1).add_(1 - beta1, grad)
TypeError: mul_() received an invalid combination of arguments - got (numpy.float32), but expected one of:
 * (Tensor other)
      didn't match because some of the arguments have invalid types: (numpy.float32)
 * (float other)
      didn't match because some of the arguments have invalid types: (numpy.float32)
```

Downgrading to 0.15.0 fixes the issue. The behavior is independent of CPU, GPU or MultipleGPU training.


Discussion:
Hey @jotterbach, what version of PyTorch are you using?
Sorry, I forgot to specify that. It's 0.4.0.
That's curious.  I haven't been able to repro yet, but it looks like the value of `beta1` is being set to a numpy float32 when a pure Python float is expected.  Do you have a snippet of code that initializes the Adamax optimizer?
Sure thing @tgaddair . The code is fairly straight forward:

```
if hps.cuda:
    model = model.cuda()
opt_s = torch.optim.Adamax(model.parameters(), hps.lr)
opt = hvd.DistributedOptimizer(opt_s,
                               named_parameters=model.named_parameters())

hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(opt, root_rank=0)
```

The relevant code in the training loop looks like:

```
model.train()
opt.zero_grad()
if hps.cuda:
    batch = batch.cuda()
loss = model(batch)

loss.backward()
opt.synchronize()
opt.step()
```

The optimizer is called nowhere else.

Ah, okay, I see what's going on here.  This is the result of a change we made to `hvd.broadcast_optimizer_state()` to broadcast the scalar options.  Looks like `beta1` is getting coerced into a `numpy.float32`, which is causing your issue.  

I'll start working on a fix for this.  Thanks for reporting!

Resolving patch:
diff --git a/horovod/torch/__init__.py b/horovod/torch/__init__.py
--- a/horovod/torch/__init__.py
+++ b/horovod/torch/__init__.py
@@ -227,6 +227,22 @@ def broadcast_optimizer_state(optimizer, root_rank):
     callbacks = {}
     occurrences = collections.defaultdict(int)
 
+    # Returns the full type structure of the possibly nested objects for recursive casting back
+    def _get_types(x):
+        if isinstance(x, collections.Iterable):
+            return type(x), [_get_types(xi) for xi in x]
+        else:
+            return type(x)
+
+    # Casts an object encoded in a tensor back into its original type and subtypes
+    def _recursive_cast(x, dtype):
+        if isinstance(dtype, tuple):
+            t, dtypes = dtype
+            x = t(x)
+            return t([_recursive_cast(x[i], dtypes[i]) for i in range(len(x))])
+        else:
+            return dtype(x)
+
     # Some optimizer parameters may be represented as scalars instead of
     # tensors.  In such cases, we need to wrap the scalar in a tensor, then
     # broadcast, then update the appropriate value in the state_dict with the
@@ -236,9 +252,9 @@ def _from_tensor():
             state_dict['state'][pid][name] = t(p.numpy()[0])
         return _from_tensor
 
-    def _create_option_callback(index, option_key, option_tensor, dtype):
+    def _create_option_callback(index, option_key, option_tensor, dtypes):
         def _from_tensor():
-            optimizer.param_groups[index][option_key] = dtype(option_tensor.numpy()[0])
+            optimizer.param_groups[index][option_key] = _recursive_cast(option_tensor.numpy()[0], dtypes)
         return _from_tensor
 
     # Param groups are an ordered list, normally there is only one per model,
@@ -252,9 +268,9 @@ def _from_tensor():
 
             # Options like the learning rate are scalar, and need to be wrapped in tensors
             key = '%s.%d' % (option_key, index)
-            dtype = type(option_value)
+            dtypes = _get_types(option_value)
             option_tensor = torch.Tensor([option_value])
-            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtype)
+            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtypes)
             params.append((key, option_tensor))
 
         # The params list here is ordered by the layers in the model


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -873,9 +873,9 @@ def test_broadcast_state_options(self):
         y = torch.randn(N, D_out).requires_grad_()
 
         params_0 = dict(lr=0.1, momentum=0.8, weight_decay=0.2, nesterov=True,
-                        etas=(0.8, 2.4), step_sizes=(1e-5, 100))
+                        betas=(0.9, 0.999), etas=(0.8, 2.4), step_sizes=(1e-5, 100))
         params_1 = dict(lr=0.2, momentum=0.9, weight_decay=0.1, nesterov=False,
-                        etas=(0.25, 1.75), step_sizes=(1e-7, 5))
+                        betas=(0.8, 0.9), etas=(0.25, 1.75), step_sizes=(1e-7, 5))
 
         def create_model(opt_class):
             model = torch.nn.Sequential(
@@ -924,8 +924,16 @@ def create_model(opt_class):
                     p_actual = [p_actual]
                     p = [p]
                 for i in range(len(p)):
+                    self.assertEqual(type(p_actual[i]), type(p[i]))
                     self.assertAlmostEqual(p_actual[i], p[i], delta=1e-5)
 
+            # Ensure that the parameter option types are compatible with ops
+            y_pred = model(x)
+            loss = F.mse_loss(y_pred, y, size_average=False)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
     def test_compression_fp16(self):
         valid_dtypes = [torch.float32, torch.float64]
         invalid_dtypes = [torch.uint8, torch.int8, torch.int16,