We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `horovod/torch/__init__.py`
Content:
```
1 # Copyright 2018 Uber Technologies, Inc. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 
20 from horovod.common import check_extension
21 
22 try:
23     check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',
24                     __file__, 'mpi_lib_v2')
25 except:
26     check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',
27                     __file__, 'mpi_lib', '_mpi_lib')
28 
29 from horovod.torch.compression import Compression
30 from horovod.torch.mpi_ops import allreduce, allreduce_async, allreduce_, allreduce_async_
31 from horovod.torch.mpi_ops import allgather, allgather_async
32 from horovod.torch.mpi_ops import broadcast, broadcast_async, broadcast_, broadcast_async_
33 from horovod.torch.mpi_ops import poll, synchronize
34 from horovod.torch.mpi_ops import init, shutdown
35 from horovod.torch.mpi_ops import size, local_size, rank, local_rank
36 from horovod.torch.mpi_ops import mpi_threads_supported
37 
38 import torch
39 import collections
40 
41 
42 class _DistributedOptimizer(torch.optim.Optimizer):
43     def __init__(self, params, named_parameters, compression):
44         super(self.__class__, self).__init__(params)
45         self._compression = compression
46 
47         if named_parameters is not None:
48             named_parameters = list(named_parameters)
49         else:
50             named_parameters = []
51 
52         # make sure that named_parameters are tuples
53         if any([not isinstance(p, tuple) for p in named_parameters]):
54             raise ValueError('named_parameters should be a sequence of '
55                              'tuples (name, parameter), usually produced by '
56                              'model.named_parameters().')
57 
58         if len(named_parameters) > 0:
59             self._parameter_names = {v: k for k, v
60                                      in sorted(named_parameters)}
61         else:
62             self._parameter_names = {v: 'allreduce.noname.%s' % i
63                                      for param_group in self.param_groups
64                                      for i, v in enumerate(param_group['params'])}
65 
66         self._handles = {}
67         self._grad_accs = []
68         self._requires_update = set()
69         if size() > 1:
70             self._register_hooks()
71 
72     def _register_hooks(self):
73         for param_group in self.param_groups:
74             for p in param_group['params']:
75                 if p.requires_grad:
76                     p.grad = p.data.new(p.size()).zero_()
77                     self._requires_update.add(p)
78                     p_tmp = p.expand_as(p)
79                     grad_acc = p_tmp.grad_fn.next_functions[0][0]
80                     grad_acc.register_hook(self._make_hook(p))
81                     self._grad_accs.append(grad_acc)
82 
83     def _allreduce_grad(self, p):
84         name = self._parameter_names.get(p)
85         tensor = p.grad.data
86         tensor_compressed, ctx = self._compression.compress(tensor)
87 
88         handle = allreduce_async_(tensor_compressed, average=True, name=name)
89         return handle, ctx
90 
91     def _make_hook(self, p):
92         def hook(*ignore):
93             assert p not in self._handles
94             assert not p.grad.requires_grad
95             handle, ctx = self._allreduce_grad(p)
96             self._handles[p] = (handle, ctx)
97         return hook
98 
99     def synchronize(self):
100         missing_p = self._requires_update - set(self._handles.keys())
101         for p in missing_p:
102             self._allreduce_grad(p)
103 
104         for p, value in self._handles.items():
105             handle, ctx = value
106             output = synchronize(handle)
107             p.grad.data.set_(self._compression.decompress(output, ctx))
108         self._handles.clear()
109 
110     def step(self, closure=None):
111         self.synchronize()
112         return super(self.__class__, self).step(closure)
113 
114 
115 def DistributedOptimizer(optimizer, named_parameters=None, compression=Compression.none):
116     """
117     An optimizer that wraps another torch.optim.Optimizer, using an allreduce to
118     average gradient values before applying gradients to model weights.
119 
120     Allreduce operations are executed after each gradient is computed by `loss.backward()`
121     in parallel with each other. The `step()` method ensures that all allreduce operations are
122     finished before applying gradients to the model.
123 
124     DistributedOptimizer exposes the `synchronize()` method, which forces allreduce operations
125     to finish before continuing the execution. It's useful in conjunction with gradient
126     clipping, or other operations that modify gradients in place before `step()` is executed.
127 
128     Example of gradient clipping:
129     ```
130     output = model(data)
131     loss = F.nll_loss(output, target)
132     loss.backward()
133     optimizer.synchronize()
134     torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)
135     optimizer.step()
136     ```
137 
138     Arguments:
139         optimizer: Optimizer to use for computing gradients and applying updates.
140         named_parameters: A mapping between parameter names and values. Used for naming of
141                           allreduce operations. Typically just `model.named_parameters()`.
142         compression: Compression algorithm used during allreduce to reduce the amount
143                      of data sent during the each parameter update step.  Defaults to
144                      not using compression.
145     """
146     # We dynamically create a new class that inherits from the optimizer that was passed in.
147     # The goal is to override the `step()` method with an allreduce implementation.
148     cls = type(optimizer.__class__.__name__, (optimizer.__class__,),
149                dict(_DistributedOptimizer.__dict__))
150     return cls(optimizer.param_groups, named_parameters, compression)
151 
152 
153 def broadcast_parameters(params, root_rank):
154     """
155     Broadcasts the parameters from root rank to all other processes.
156     Typical usage is to broadcast the `model.state_dict()`,
157     `model.named_parameters()`, or `model.parameters()`.
158 
159     Arguments:
160         params: One of the following:
161             - list of parameters to broadcast
162             - dict of parameters to broadcast
163         root_rank: The rank of the process from which parameters will be
164                    broadcasted to all other processes.
165     """
166     if isinstance(params, dict):
167         params = sorted(params.items())
168     elif isinstance(params, list):
169         # support both named_parameters() and regular parameters()
170         params = [p if isinstance(p, tuple) else (None, p) for p in params]
171     else:
172         raise ValueError('invalid params of type: %s' % type(params))
173 
174     # Run asynchronous broadcasts.
175     handles = []
176     for name, p in params:
177         handle = broadcast_async_(p, root_rank, name)
178         handles.append(handle)
179 
180     # Wait for completion.
181     for handle in handles:
182         synchronize(handle)
183 
184 
185 def broadcast_optimizer_state(optimizer, root_rank):
186     """
187     Broadcasts an optimizer state from root rank to all other processes.
188 
189     Arguments:
190         optimizer: An optimizer.
191         root_rank: The rank of the process from which the optimizer will be
192                    broadcasted to all other processes.
193     """
194     if isinstance(optimizer, torch.optim.LBFGS):
195         # TODO(travis): L-BFGS cannot be easily supported without serializing
196         # the entire state_dict, as its structure is deeply nested and contains
197         # None type parameter values
198         raise ValueError('cannot broadcast torch.optim.LBFGS state')
199 
200     state_dict = optimizer.state_dict()
201 
202     # Newly created optimizers will not have their state initialized, so
203     # do that initialization here
204     if len(state_dict['state']) == 0:
205         for group in optimizer.param_groups:
206             for p in group['params']:
207                 p.grad = p.data.new(p.size()).zero_()
208         # This function accepts a torch.optim.Optimizer or a DistributedOptimizer
209         # wrapped around a torch optimizer. Calling step() with a DistributedOptimizer
210         # forces allreduce on all model parameters, which will result in deadlock
211         # unless every rank calls step(). Therefore, to finish state initialization
212         # only call optimizer.step() with a torch.optim.Optimizer.
213         if optimizer.__module__ == DistributedOptimizer.__module__:
214             super(optimizer.__class__, optimizer).step()
215         else:
216             optimizer.step()
217         state_dict = optimizer.state_dict()
218 
219     # If the state_dict is still empty after initialization, then
220     # the optimizer is stateless, and there is nothing to broadcast.
221     # Furthermore, attempting to access the state dict would result in
222     # an error.
223     if len(state_dict['state']) == 0:
224         return
225 
226     params = []
227     callbacks = {}
228     occurrences = collections.defaultdict(int)
229 
230     # Some optimizer parameters may be represented as scalars instead of
231     # tensors.  In such cases, we need to wrap the scalar in a tensor, then
232     # broadcast, then update the appropriate value in the state_dict with the
233     # new unwrapped scalar value via a callback.
234     def _create_callback(pid, name, t, p):
235         def _from_tensor():
236             state_dict['state'][pid][name] = t(p.numpy()[0])
237         return _from_tensor
238 
239     def _create_option_callback(index, option_key, option_tensor, dtype):
240         def _from_tensor():
241             optimizer.param_groups[index][option_key] = dtype(option_tensor.numpy()[0])
242         return _from_tensor
243 
244     # Param groups are an ordered list, normally there is only one per model,
245     # but users can add additional param groups for example to train
246     # previously frozen layers
247     for index, group in enumerate(state_dict['param_groups']):
248         # Broadcast options like learning rate
249         for option_key, option_value in group.items():
250             if option_key == 'params':
251                 continue
252 
253             # Options like the learning rate are scalar, and need to be wrapped in tensors
254             key = '%s.%d' % (option_key, index)
255             dtype = type(option_value)
256             option_tensor = torch.Tensor([option_value])
257             callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtype)
258             params.append((key, option_tensor))
259 
260         # The params list here is ordered by the layers in the model
261         for pid in group['params']:
262             param_state = state_dict['state'][pid]
263             for name, p in param_state.items():
264                 # Some parameter names may appear more than once, in which
265                 # case we ensure they have a unique identifier defined by
266                 # their order
267                 occurrences[name] += 1
268                 key = '%s.%d' % (str(name), occurrences[name])
269 
270                 if not torch.is_tensor(p):
271                     # Wrap the scalar in a FloatTensor, and remember its type
272                     # so we can cast it back after unwrapping
273                     t = type(p)
274                     p = torch.Tensor([p])
275                     callbacks[key] = _create_callback(pid, name, t, p)
276 
277                 params.append((key, p))
278 
279     # Synchronized broadcast of all parameters
280     broadcast_parameters(params, root_rank)
281 
282     # Post-broadcast clenaup for non-tensor parameters
283     for key, p in params:
284         if key in callbacks:
285             callbacks[key]()
286
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
             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 @@
 
             # 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

{"golden_diff": "diff --git a/horovod/torch/__init__.py b/horovod/torch/__init__.py\n--- a/horovod/torch/__init__.py\n+++ b/horovod/torch/__init__.py\n@@ -227,6 +227,22 @@\n     callbacks = {}\n     occurrences = collections.defaultdict(int)\n \n+    # Returns the full type structure of the possibly nested objects for recursive casting back\n+    def _get_types(x):\n+        if isinstance(x, collections.Iterable):\n+            return type(x), [_get_types(xi) for xi in x]\n+        else:\n+            return type(x)\n+\n+    # Casts an object encoded in a tensor back into its original type and subtypes\n+    def _recursive_cast(x, dtype):\n+        if isinstance(dtype, tuple):\n+            t, dtypes = dtype\n+            x = t(x)\n+            return t([_recursive_cast(x[i], dtypes[i]) for i in range(len(x))])\n+        else:\n+            return dtype(x)\n+\n     # Some optimizer parameters may be represented as scalars instead of\n     # tensors.  In such cases, we need to wrap the scalar in a tensor, then\n     # broadcast, then update the appropriate value in the state_dict with the\n@@ -236,9 +252,9 @@\n             state_dict['state'][pid][name] = t(p.numpy()[0])\n         return _from_tensor\n \n-    def _create_option_callback(index, option_key, option_tensor, dtype):\n+    def _create_option_callback(index, option_key, option_tensor, dtypes):\n         def _from_tensor():\n-            optimizer.param_groups[index][option_key] = dtype(option_tensor.numpy()[0])\n+            optimizer.param_groups[index][option_key] = _recursive_cast(option_tensor.numpy()[0], dtypes)\n         return _from_tensor\n \n     # Param groups are an ordered list, normally there is only one per model,\n@@ -252,9 +268,9 @@\n \n             # Options like the learning rate are scalar, and need to be wrapped in tensors\n             key = '%s.%d' % (option_key, index)\n-            dtype = type(option_value)\n+            dtypes = _get_types(option_value)\n             option_tensor = torch.Tensor([option_value])\n-            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtype)\n+            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtypes)\n             params.append((key, option_tensor))\n \n         # The params list here is ordered by the layers in the model\n", "issue": "pytorch + horovod 0.15.1 distributed optimizer not working anymore\nI 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\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"train.py\", line 641, in <module>\r\n    train_images(hps)\r\n  File \"train.py\", line 444, in train_images\r\n    train_step(batch, batch_idx, epoch, hps, model, opt, train_logger)\r\n  File \"train.py\", line 457, in train_step\r\n    opt.step()\r\n  File \"/opt/conda/lib/python3.6/site-packages/horovod/torch/__init__.py\", line 97, in step\r\n    return super(self.__class__, self).step(closure)\r\n  File \"/opt/conda/lib/python3.6/site-packages/torch/optim/adamax.py\", line 75, in step\r\n    exp_avg.mul_(beta1).add_(1 - beta1, grad)\r\nTypeError: mul_() received an invalid combination of arguments - got (numpy.float32), but expected one of:\r\n * (Tensor other)\r\n      didn't match because some of the arguments have invalid types: (numpy.float32)\r\n * (float other)\r\n      didn't match because some of the arguments have invalid types: (numpy.float32)\r\n```\r\n\r\nDowngrading to 0.15.0 fixes the issue. The behavior is independent of CPU, GPU or MultipleGPU training.\n", "before_files": [{"content": "# Copyright 2018 Uber Technologies, Inc. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom horovod.common import check_extension\n\ntry:\n    check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',\n                    __file__, 'mpi_lib_v2')\nexcept:\n    check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',\n                    __file__, 'mpi_lib', '_mpi_lib')\n\nfrom horovod.torch.compression import Compression\nfrom horovod.torch.mpi_ops import allreduce, allreduce_async, allreduce_, allreduce_async_\nfrom horovod.torch.mpi_ops import allgather, allgather_async\nfrom horovod.torch.mpi_ops import broadcast, broadcast_async, broadcast_, broadcast_async_\nfrom horovod.torch.mpi_ops import poll, synchronize\nfrom horovod.torch.mpi_ops import init, shutdown\nfrom horovod.torch.mpi_ops import size, local_size, rank, local_rank\nfrom horovod.torch.mpi_ops import mpi_threads_supported\n\nimport torch\nimport collections\n\n\nclass _DistributedOptimizer(torch.optim.Optimizer):\n    def __init__(self, params, named_parameters, compression):\n        super(self.__class__, self).__init__(params)\n        self._compression = compression\n\n        if named_parameters is not None:\n            named_parameters = list(named_parameters)\n        else:\n            named_parameters = []\n\n        # make sure that named_parameters are tuples\n        if any([not isinstance(p, tuple) for p in named_parameters]):\n            raise ValueError('named_parameters should be a sequence of '\n                             'tuples (name, parameter), usually produced by '\n                             'model.named_parameters().')\n\n        if len(named_parameters) > 0:\n            self._parameter_names = {v: k for k, v\n                                     in sorted(named_parameters)}\n        else:\n            self._parameter_names = {v: 'allreduce.noname.%s' % i\n                                     for param_group in self.param_groups\n                                     for i, v in enumerate(param_group['params'])}\n\n        self._handles = {}\n        self._grad_accs = []\n        self._requires_update = set()\n        if size() > 1:\n            self._register_hooks()\n\n    def _register_hooks(self):\n        for param_group in self.param_groups:\n            for p in param_group['params']:\n                if p.requires_grad:\n                    p.grad = p.data.new(p.size()).zero_()\n                    self._requires_update.add(p)\n                    p_tmp = p.expand_as(p)\n                    grad_acc = p_tmp.grad_fn.next_functions[0][0]\n                    grad_acc.register_hook(self._make_hook(p))\n                    self._grad_accs.append(grad_acc)\n\n    def _allreduce_grad(self, p):\n        name = self._parameter_names.get(p)\n        tensor = p.grad.data\n        tensor_compressed, ctx = self._compression.compress(tensor)\n\n        handle = allreduce_async_(tensor_compressed, average=True, name=name)\n        return handle, ctx\n\n    def _make_hook(self, p):\n        def hook(*ignore):\n            assert p not in self._handles\n            assert not p.grad.requires_grad\n            handle, ctx = self._allreduce_grad(p)\n            self._handles[p] = (handle, ctx)\n        return hook\n\n    def synchronize(self):\n        missing_p = self._requires_update - set(self._handles.keys())\n        for p in missing_p:\n            self._allreduce_grad(p)\n\n        for p, value in self._handles.items():\n            handle, ctx = value\n            output = synchronize(handle)\n            p.grad.data.set_(self._compression.decompress(output, ctx))\n        self._handles.clear()\n\n    def step(self, closure=None):\n        self.synchronize()\n        return super(self.__class__, self).step(closure)\n\n\ndef DistributedOptimizer(optimizer, named_parameters=None, compression=Compression.none):\n    \"\"\"\n    An optimizer that wraps another torch.optim.Optimizer, using an allreduce to\n    average gradient values before applying gradients to model weights.\n\n    Allreduce operations are executed after each gradient is computed by `loss.backward()`\n    in parallel with each other. The `step()` method ensures that all allreduce operations are\n    finished before applying gradients to the model.\n\n    DistributedOptimizer exposes the `synchronize()` method, which forces allreduce operations\n    to finish before continuing the execution. It's useful in conjunction with gradient\n    clipping, or other operations that modify gradients in place before `step()` is executed.\n\n    Example of gradient clipping:\n    ```\n    output = model(data)\n    loss = F.nll_loss(output, target)\n    loss.backward()\n    optimizer.synchronize()\n    torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)\n    optimizer.step()\n    ```\n\n    Arguments:\n        optimizer: Optimizer to use for computing gradients and applying updates.\n        named_parameters: A mapping between parameter names and values. Used for naming of\n                          allreduce operations. Typically just `model.named_parameters()`.\n        compression: Compression algorithm used during allreduce to reduce the amount\n                     of data sent during the each parameter update step.  Defaults to\n                     not using compression.\n    \"\"\"\n    # We dynamically create a new class that inherits from the optimizer that was passed in.\n    # The goal is to override the `step()` method with an allreduce implementation.\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,),\n               dict(_DistributedOptimizer.__dict__))\n    return cls(optimizer.param_groups, named_parameters, compression)\n\n\ndef broadcast_parameters(params, root_rank):\n    \"\"\"\n    Broadcasts the parameters from root rank to all other processes.\n    Typical usage is to broadcast the `model.state_dict()`,\n    `model.named_parameters()`, or `model.parameters()`.\n\n    Arguments:\n        params: One of the following:\n            - list of parameters to broadcast\n            - dict of parameters to broadcast\n        root_rank: The rank of the process from which parameters will be\n                   broadcasted to all other processes.\n    \"\"\"\n    if isinstance(params, dict):\n        params = sorted(params.items())\n    elif isinstance(params, list):\n        # support both named_parameters() and regular parameters()\n        params = [p if isinstance(p, tuple) else (None, p) for p in params]\n    else:\n        raise ValueError('invalid params of type: %s' % type(params))\n\n    # Run asynchronous broadcasts.\n    handles = []\n    for name, p in params:\n        handle = broadcast_async_(p, root_rank, name)\n        handles.append(handle)\n\n    # Wait for completion.\n    for handle in handles:\n        synchronize(handle)\n\n\ndef broadcast_optimizer_state(optimizer, root_rank):\n    \"\"\"\n    Broadcasts an optimizer state from root rank to all other processes.\n\n    Arguments:\n        optimizer: An optimizer.\n        root_rank: The rank of the process from which the optimizer will be\n                   broadcasted to all other processes.\n    \"\"\"\n    if isinstance(optimizer, torch.optim.LBFGS):\n        # TODO(travis): L-BFGS cannot be easily supported without serializing\n        # the entire state_dict, as its structure is deeply nested and contains\n        # None type parameter values\n        raise ValueError('cannot broadcast torch.optim.LBFGS state')\n\n    state_dict = optimizer.state_dict()\n\n    # Newly created optimizers will not have their state initialized, so\n    # do that initialization here\n    if len(state_dict['state']) == 0:\n        for group in optimizer.param_groups:\n            for p in group['params']:\n                p.grad = p.data.new(p.size()).zero_()\n        # This function accepts a torch.optim.Optimizer or a DistributedOptimizer\n        # wrapped around a torch optimizer. Calling step() with a DistributedOptimizer\n        # forces allreduce on all model parameters, which will result in deadlock\n        # unless every rank calls step(). Therefore, to finish state initialization\n        # only call optimizer.step() with a torch.optim.Optimizer.\n        if optimizer.__module__ == DistributedOptimizer.__module__:\n            super(optimizer.__class__, optimizer).step()\n        else:\n            optimizer.step()\n        state_dict = optimizer.state_dict()\n\n    # If the state_dict is still empty after initialization, then\n    # the optimizer is stateless, and there is nothing to broadcast.\n    # Furthermore, attempting to access the state dict would result in\n    # an error.\n    if len(state_dict['state']) == 0:\n        return\n\n    params = []\n    callbacks = {}\n    occurrences = collections.defaultdict(int)\n\n    # Some optimizer parameters may be represented as scalars instead of\n    # tensors.  In such cases, we need to wrap the scalar in a tensor, then\n    # broadcast, then update the appropriate value in the state_dict with the\n    # new unwrapped scalar value via a callback.\n    def _create_callback(pid, name, t, p):\n        def _from_tensor():\n            state_dict['state'][pid][name] = t(p.numpy()[0])\n        return _from_tensor\n\n    def _create_option_callback(index, option_key, option_tensor, dtype):\n        def _from_tensor():\n            optimizer.param_groups[index][option_key] = dtype(option_tensor.numpy()[0])\n        return _from_tensor\n\n    # Param groups are an ordered list, normally there is only one per model,\n    # but users can add additional param groups for example to train\n    # previously frozen layers\n    for index, group in enumerate(state_dict['param_groups']):\n        # Broadcast options like learning rate\n        for option_key, option_value in group.items():\n            if option_key == 'params':\n                continue\n\n            # Options like the learning rate are scalar, and need to be wrapped in tensors\n            key = '%s.%d' % (option_key, index)\n            dtype = type(option_value)\n            option_tensor = torch.Tensor([option_value])\n            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtype)\n            params.append((key, option_tensor))\n\n        # The params list here is ordered by the layers in the model\n        for pid in group['params']:\n            param_state = state_dict['state'][pid]\n            for name, p in param_state.items():\n                # Some parameter names may appear more than once, in which\n                # case we ensure they have a unique identifier defined by\n                # their order\n                occurrences[name] += 1\n                key = '%s.%d' % (str(name), occurrences[name])\n\n                if not torch.is_tensor(p):\n                    # Wrap the scalar in a FloatTensor, and remember its type\n                    # so we can cast it back after unwrapping\n                    t = type(p)\n                    p = torch.Tensor([p])\n                    callbacks[key] = _create_callback(pid, name, t, p)\n\n                params.append((key, p))\n\n    # Synchronized broadcast of all parameters\n    broadcast_parameters(params, root_rank)\n\n    # Post-broadcast clenaup for non-tensor parameters\n    for key, p in params:\n        if key in callbacks:\n            callbacks[key]()\n", "path": "horovod/torch/__init__.py"}], "after_files": [{"content": "# Copyright 2018 Uber Technologies, Inc. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom horovod.common import check_extension\n\ntry:\n    check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',\n                    __file__, 'mpi_lib_v2')\nexcept:\n    check_extension('horovod.torch', 'HOROVOD_WITH_PYTORCH',\n                    __file__, 'mpi_lib', '_mpi_lib')\n\nfrom horovod.torch.compression import Compression\nfrom horovod.torch.mpi_ops import allreduce, allreduce_async, allreduce_, allreduce_async_\nfrom horovod.torch.mpi_ops import allgather, allgather_async\nfrom horovod.torch.mpi_ops import broadcast, broadcast_async, broadcast_, broadcast_async_\nfrom horovod.torch.mpi_ops import poll, synchronize\nfrom horovod.torch.mpi_ops import init, shutdown\nfrom horovod.torch.mpi_ops import size, local_size, rank, local_rank\nfrom horovod.torch.mpi_ops import mpi_threads_supported\n\nimport torch\nimport collections\n\n\nclass _DistributedOptimizer(torch.optim.Optimizer):\n    def __init__(self, params, named_parameters, compression):\n        super(self.__class__, self).__init__(params)\n        self._compression = compression\n\n        if named_parameters is not None:\n            named_parameters = list(named_parameters)\n        else:\n            named_parameters = []\n\n        # make sure that named_parameters are tuples\n        if any([not isinstance(p, tuple) for p in named_parameters]):\n            raise ValueError('named_parameters should be a sequence of '\n                             'tuples (name, parameter), usually produced by '\n                             'model.named_parameters().')\n\n        if len(named_parameters) > 0:\n            self._parameter_names = {v: k for k, v\n                                     in sorted(named_parameters)}\n        else:\n            self._parameter_names = {v: 'allreduce.noname.%s' % i\n                                     for param_group in self.param_groups\n                                     for i, v in enumerate(param_group['params'])}\n\n        self._handles = {}\n        self._grad_accs = []\n        self._requires_update = set()\n        if size() > 1:\n            self._register_hooks()\n\n    def _register_hooks(self):\n        for param_group in self.param_groups:\n            for p in param_group['params']:\n                if p.requires_grad:\n                    p.grad = p.data.new(p.size()).zero_()\n                    self._requires_update.add(p)\n                    p_tmp = p.expand_as(p)\n                    grad_acc = p_tmp.grad_fn.next_functions[0][0]\n                    grad_acc.register_hook(self._make_hook(p))\n                    self._grad_accs.append(grad_acc)\n\n    def _allreduce_grad(self, p):\n        name = self._parameter_names.get(p)\n        tensor = p.grad.data\n        tensor_compressed, ctx = self._compression.compress(tensor)\n\n        handle = allreduce_async_(tensor_compressed, average=True, name=name)\n        return handle, ctx\n\n    def _make_hook(self, p):\n        def hook(*ignore):\n            assert p not in self._handles\n            assert not p.grad.requires_grad\n            handle, ctx = self._allreduce_grad(p)\n            self._handles[p] = (handle, ctx)\n        return hook\n\n    def synchronize(self):\n        missing_p = self._requires_update - set(self._handles.keys())\n        for p in missing_p:\n            self._allreduce_grad(p)\n\n        for p, value in self._handles.items():\n            handle, ctx = value\n            output = synchronize(handle)\n            p.grad.data.set_(self._compression.decompress(output, ctx))\n        self._handles.clear()\n\n    def step(self, closure=None):\n        self.synchronize()\n        return super(self.__class__, self).step(closure)\n\n\ndef DistributedOptimizer(optimizer, named_parameters=None, compression=Compression.none):\n    \"\"\"\n    An optimizer that wraps another torch.optim.Optimizer, using an allreduce to\n    average gradient values before applying gradients to model weights.\n\n    Allreduce operations are executed after each gradient is computed by `loss.backward()`\n    in parallel with each other. The `step()` method ensures that all allreduce operations are\n    finished before applying gradients to the model.\n\n    DistributedOptimizer exposes the `synchronize()` method, which forces allreduce operations\n    to finish before continuing the execution. It's useful in conjunction with gradient\n    clipping, or other operations that modify gradients in place before `step()` is executed.\n\n    Example of gradient clipping:\n    ```\n    output = model(data)\n    loss = F.nll_loss(output, target)\n    loss.backward()\n    optimizer.synchronize()\n    torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)\n    optimizer.step()\n    ```\n\n    Arguments:\n        optimizer: Optimizer to use for computing gradients and applying updates.\n        named_parameters: A mapping between parameter names and values. Used for naming of\n                          allreduce operations. Typically just `model.named_parameters()`.\n        compression: Compression algorithm used during allreduce to reduce the amount\n                     of data sent during the each parameter update step.  Defaults to\n                     not using compression.\n    \"\"\"\n    # We dynamically create a new class that inherits from the optimizer that was passed in.\n    # The goal is to override the `step()` method with an allreduce implementation.\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,),\n               dict(_DistributedOptimizer.__dict__))\n    return cls(optimizer.param_groups, named_parameters, compression)\n\n\ndef broadcast_parameters(params, root_rank):\n    \"\"\"\n    Broadcasts the parameters from root rank to all other processes.\n    Typical usage is to broadcast the `model.state_dict()`,\n    `model.named_parameters()`, or `model.parameters()`.\n\n    Arguments:\n        params: One of the following:\n            - list of parameters to broadcast\n            - dict of parameters to broadcast\n        root_rank: The rank of the process from which parameters will be\n                   broadcasted to all other processes.\n    \"\"\"\n    if isinstance(params, dict):\n        params = sorted(params.items())\n    elif isinstance(params, list):\n        # support both named_parameters() and regular parameters()\n        params = [p if isinstance(p, tuple) else (None, p) for p in params]\n    else:\n        raise ValueError('invalid params of type: %s' % type(params))\n\n    # Run asynchronous broadcasts.\n    handles = []\n    for name, p in params:\n        handle = broadcast_async_(p, root_rank, name)\n        handles.append(handle)\n\n    # Wait for completion.\n    for handle in handles:\n        synchronize(handle)\n\n\ndef broadcast_optimizer_state(optimizer, root_rank):\n    \"\"\"\n    Broadcasts an optimizer state from root rank to all other processes.\n\n    Arguments:\n        optimizer: An optimizer.\n        root_rank: The rank of the process from which the optimizer will be\n                   broadcasted to all other processes.\n    \"\"\"\n    if isinstance(optimizer, torch.optim.LBFGS):\n        # TODO(travis): L-BFGS cannot be easily supported without serializing\n        # the entire state_dict, as its structure is deeply nested and contains\n        # None type parameter values\n        raise ValueError('cannot broadcast torch.optim.LBFGS state')\n\n    state_dict = optimizer.state_dict()\n\n    # Newly created optimizers will not have their state initialized, so\n    # do that initialization here\n    if len(state_dict['state']) == 0:\n        for group in optimizer.param_groups:\n            for p in group['params']:\n                p.grad = p.data.new(p.size()).zero_()\n        # This function accepts a torch.optim.Optimizer or a DistributedOptimizer\n        # wrapped around a torch optimizer. Calling step() with a DistributedOptimizer\n        # forces allreduce on all model parameters, which will result in deadlock\n        # unless every rank calls step(). Therefore, to finish state initialization\n        # only call optimizer.step() with a torch.optim.Optimizer.\n        if optimizer.__module__ == DistributedOptimizer.__module__:\n            super(optimizer.__class__, optimizer).step()\n        else:\n            optimizer.step()\n        state_dict = optimizer.state_dict()\n\n    # If the state_dict is still empty after initialization, then\n    # the optimizer is stateless, and there is nothing to broadcast.\n    # Furthermore, attempting to access the state dict would result in\n    # an error.\n    if len(state_dict['state']) == 0:\n        return\n\n    params = []\n    callbacks = {}\n    occurrences = collections.defaultdict(int)\n\n    # Returns the full type structure of the possibly nested objects for recursive casting back\n    def _get_types(x):\n        if isinstance(x, collections.Iterable):\n            return type(x), [_get_types(xi) for xi in x]\n        else:\n            return type(x)\n\n    # Casts an object encoded in a tensor back into its original type and subtypes\n    def _recursive_cast(x, dtype):\n        if isinstance(dtype, tuple):\n            t, dtypes = dtype\n            x = t(x)\n            return t([_recursive_cast(x[i], dtypes[i]) for i in range(len(x))])\n        else:\n            return dtype(x)\n\n    # Some optimizer parameters may be represented as scalars instead of\n    # tensors.  In such cases, we need to wrap the scalar in a tensor, then\n    # broadcast, then update the appropriate value in the state_dict with the\n    # new unwrapped scalar value via a callback.\n    def _create_callback(pid, name, t, p):\n        def _from_tensor():\n            state_dict['state'][pid][name] = t(p.numpy()[0])\n        return _from_tensor\n\n    def _create_option_callback(index, option_key, option_tensor, dtypes):\n        def _from_tensor():\n            optimizer.param_groups[index][option_key] = _recursive_cast(option_tensor.numpy()[0], dtypes)\n        return _from_tensor\n\n    # Param groups are an ordered list, normally there is only one per model,\n    # but users can add additional param groups for example to train\n    # previously frozen layers\n    for index, group in enumerate(state_dict['param_groups']):\n        # Broadcast options like learning rate\n        for option_key, option_value in group.items():\n            if option_key == 'params':\n                continue\n\n            # Options like the learning rate are scalar, and need to be wrapped in tensors\n            key = '%s.%d' % (option_key, index)\n            dtypes = _get_types(option_value)\n            option_tensor = torch.Tensor([option_value])\n            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtypes)\n            params.append((key, option_tensor))\n\n        # The params list here is ordered by the layers in the model\n        for pid in group['params']:\n            param_state = state_dict['state'][pid]\n            for name, p in param_state.items():\n                # Some parameter names may appear more than once, in which\n                # case we ensure they have a unique identifier defined by\n                # their order\n                occurrences[name] += 1\n                key = '%s.%d' % (str(name), occurrences[name])\n\n                if not torch.is_tensor(p):\n                    # Wrap the scalar in a FloatTensor, and remember its type\n                    # so we can cast it back after unwrapping\n                    t = type(p)\n                    p = torch.Tensor([p])\n                    callbacks[key] = _create_callback(pid, name, t, p)\n\n                params.append((key, p))\n\n    # Synchronized broadcast of all parameters\n    broadcast_parameters(params, root_rank)\n\n    # Post-broadcast clenaup for non-tensor parameters\n    for key, p in params:\n        if key in callbacks:\n            callbacks[key]()\n", "path": "horovod/torch/__init__.py"}]}