We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Very FATAL BUG about hvd.broadcast_optimizer_state
It seems that hvd.broadcast_optimizer_state **cannot broadcast optimizer state to other gpus or other machines gpus.**

i have running `pytorch_imagenet_resnet50.py`, when i load a pretrained checkpoint, it seems that only my 1st gpu on my rank0 's learning have been changed to pretrained checkpoint. but others are **still init learning!!!**
my experiments result are as follows:

=> loading checkpoint 'horovod_imagenet_v6/checkpoint-180.pth.tar'
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 2.0796608879436464e-11
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)
/search/speech/yh/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters
Train Epoch     #1:   0%|          | 0/5005 [00:00<?, ?it/s]SGD (
Parameter Group 0
    dampening: 0
    lr: 0.1
    momentum: 0.9
    nesterov: False
    weight_decay: 5e-05
)




As we can see above, only one gpu's lr is ` lr: 2.0796608879436464e-11`, but others is still 0.1 which is the init lr.

but in your code , `    # Horovod: broadcast parameters & optimizer state.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)`   which says : Horovod: broadcast parameters & optimizer state.  So maybe this is a very serious bug???

Please confirm it as soon as possible. Thanks a lot! :)

--- 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         self._parameter_names = {v: k for k, v
59                                  in sorted(named_parameters)}
60         self._handles = {}
61         self._grad_accs = []
62 
63         if size() > 1:
64             self._register_hooks()
65 
66     def _register_hooks(self):
67         for param_group in self.param_groups:
68             for p in param_group['params']:
69                 if p.requires_grad:
70                     p_tmp = p.expand_as(p)
71                     grad_acc = p_tmp.grad_fn.next_functions[0][0]
72                     grad_acc.register_hook(self._make_hook(p))
73                     self._grad_accs.append(grad_acc)
74 
75     def _make_hook(self, p):
76         def hook(*ignore):
77             assert p not in self._handles
78             assert not p.grad.requires_grad
79             name = self._parameter_names.get(p)
80 
81             tensor = p.grad.data
82             tensor_compressed, ctx = self._compression.compress(tensor)
83 
84             handle = allreduce_async_(tensor_compressed, average=True, name=name)
85             self._handles[p] = (handle, ctx)
86         return hook
87 
88     def synchronize(self):
89         for p, value in self._handles.items():
90             handle, ctx = value
91             output = synchronize(handle)
92             p.grad.data.set_(self._compression.decompress(output, ctx))
93         self._handles.clear()
94 
95     def step(self, closure=None):
96         self.synchronize()
97         return super(self.__class__, self).step(closure)
98 
99 
100 def DistributedOptimizer(optimizer, named_parameters=None, compression=Compression.none):
101     """
102     An optimizer that wraps another torch.optim.Optimizer, using an allreduce to
103     average gradient values before applying gradients to model weights.
104 
105     Allreduce operations are executed after each gradient is computed by `loss.backward()`
106     in parallel with each other. The `step()` method ensures that all allreduce operations are
107     finished before applying gradients to the model.
108 
109     DistributedOptimizer exposes the `synchronize()` method, which forces allreduce operations
110     to finish before continuing the execution. It's useful in conjunction with gradient
111     clipping, or other operations that modify gradients in place before `step()` is executed.
112 
113     Example of gradient clipping:
114     ```
115     output = model(data)
116     loss = F.nll_loss(output, target)
117     loss.backward()
118     optimizer.synchronize()
119     torch.nn.utils.clip_grad_norm(model.parameters(), args.clip)
120     optimizer.step()
121     ```
122 
123     Arguments:
124         optimizer: Optimizer to use for computing gradients and applying updates.
125         named_parameters: A mapping between parameter names and values. Used for naming of
126                           allreduce operations. Typically just `model.named_parameters()`.
127         compression: Compression algorithm used during allreduce to reduce the amount
128                      of data sent during the each parameter update step.  Defaults to
129                      not using compression.
130     """
131     # We dynamically create a new class that inherits from the optimizer that was passed in.
132     # The goal is to override the `step()` method with an allreduce implementation.
133     cls = type(optimizer.__class__.__name__, (optimizer.__class__,),
134                dict(_DistributedOptimizer.__dict__))
135     return cls(optimizer.param_groups, named_parameters, compression)
136 
137 
138 def broadcast_parameters(params, root_rank):
139     """
140     Broadcasts the parameters from root rank to all other processes.
141     Typical usage is to broadcast the `model.state_dict()`,
142     `model.named_parameters()`, or `model.parameters()`.
143 
144     Arguments:
145         params: One of the following:
146             - list of parameters to broadcast
147             - dict of parameters to broadcast
148         root_rank: The rank of the process from which parameters will be
149                    broadcasted to all other processes.
150     """
151     if isinstance(params, dict):
152         params = sorted(params.items())
153     elif isinstance(params, list):
154         # support both named_parameters() and regular parameters()
155         params = [p if isinstance(p, tuple) else (None, p) for p in params]
156     else:
157         raise ValueError('invalid params of type: %s' % type(params))
158 
159     # Run asynchronous broadcasts.
160     handles = []
161     for name, p in params:
162         handle = broadcast_async_(p, root_rank, name)
163         handles.append(handle)
164 
165     # Wait for completion.
166     for handle in handles:
167         synchronize(handle)
168 
169 
170 def broadcast_optimizer_state(optimizer, root_rank):
171     """
172     Broadcasts an optimizer state from root rank to all other processes.
173 
174     Arguments:
175         optimizer: An optimizer.
176         root_rank: The rank of the process from which the optimizer will be
177                    broadcasted to all other processes.
178     """
179     if isinstance(optimizer, torch.optim.LBFGS):
180         # TODO(travis): L-BFGS cannot be easily supported without serializing
181         # the entire state_dict, as its structure is deeply nested and contains
182         # None type parameter values
183         raise ValueError('cannot broadcast torch.optim.LBFGS state')
184 
185     state_dict = optimizer.state_dict()
186 
187     # Newly created optimizers will not have their state initialized, so
188     # do that initialization here
189     if len(state_dict['state']) == 0:
190         for group in optimizer.param_groups:
191             for p in group['params']:
192                 p.grad = p.data.new(p.size()).zero_()
193         optimizer.step()
194         state_dict = optimizer.state_dict()
195 
196     # If the state_dict is still empty after initialization, then
197     # the optimizer is stateless, and there is nothing to broadcast.
198     # Furthermore, attempting to access the state dict would result in
199     # an error.
200     if len(state_dict['state']) == 0:
201         return
202 
203     params = []
204     callbacks = {}
205     occurrences = collections.defaultdict(int)
206 
207     # Some optimizer parameters may be represented as scalars instead of
208     # tensors.  In such cases, we need to wrap the scalar in a tensor, then
209     # broadcast, then update the appropriate value in the state_dict with the
210     # new unwrapped scalar value via a callback.
211     def _create_callback(pid, name, t, p):
212         def _from_tensor():
213             state_dict['state'][pid][name] = t(p.numpy()[0])
214         return _from_tensor
215 
216     # Groups are unordered, but their params will be distinct
217     for group in state_dict['param_groups']:
218         # The params list here is ordered by the layers in the model
219         for pid in group['params']:
220             param_state = state_dict['state'][pid]
221             for name, p in param_state.items():
222                 # Some parameter names may appear more than once, in which
223                 # case we ensure they have a unique identifier defined by
224                 # their order
225                 occurrences[name] += 1
226                 key = '%s.%d' % (str(name), occurrences[name])
227 
228                 if not torch.is_tensor(p):
229                     # Wrap the scalar in a FloatTensor, and remember its type
230                     # so we can cast it back after unwrapping
231                     t = type(p)
232                     p = torch.Tensor([p])
233                     callbacks[key] = _create_callback(pid, name, t, p)
234 
235                 params.append((key, p))
236 
237     # Synchronized broadcast of all parameters
238     broadcast_parameters(params, root_rank)
239 
240     # Post-broadcast clenaup for non-tensor parameters
241     for key, p in params:
242         if key in callbacks:
243             callbacks[key]()
244
```
--- 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
@@ -213,8 +213,27 @@
             state_dict['state'][pid][name] = t(p.numpy()[0])
         return _from_tensor
 
-    # Groups are unordered, but their params will be distinct
-    for group in state_dict['param_groups']:
+    def _create_option_callback(index, option_key, option_tensor, dtype):
+        def _from_tensor():
+            optimizer.param_groups[index][option_key] = dtype(option_tensor.numpy()[0])
+        return _from_tensor
+
+    # Param groups are an ordered list, normally there is only one per model,
+    # but users can add additional param groups for example to train
+    # previously frozen layers
+    for index, group in enumerate(state_dict['param_groups']):
+        # Broadcast options like learning rate
+        for option_key, option_value in group.items():
+            if option_key == 'params':
+                continue
+
+            # Options like the learning rate are scalar, and need to be wrapped in tensors
+            key = '%s.%d' % (option_key, index)
+            dtype = type(option_value)
+            option_tensor = torch.Tensor([option_value])
+            callbacks[key] = _create_option_callback(index, option_key, option_tensor, dtype)
+            params.append((key, option_tensor))
+
         # The params list here is ordered by the layers in the model
         for pid in group['params']:
             param_state = state_dict['state'][pid]

{"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@@ -213,8 +213,27 @@\n             state_dict['state'][pid][name] = t(p.numpy()[0])\n         return _from_tensor\n \n-    # Groups are unordered, but their params will be distinct\n-    for group in state_dict['param_groups']:\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", "issue": "Very FATAL BUG about hvd.broadcast_optimizer_state\nIt seems that hvd.broadcast_optimizer_state **cannot broadcast optimizer state to other gpus or other machines gpus.**\r\n\r\ni have running `pytorch_imagenet_resnet50.py`, when i load a pretrained checkpoint, it seems that only my 1st gpu on my rank0 's learning have been changed to pretrained checkpoint. but others are **still init learning!!!**\r\nmy experiments result are as follows:\r\n\r\n=> loading checkpoint 'horovod_imagenet_v6/checkpoint-180.pth.tar'\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 2.0796608879436464e-11\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\nSGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\n/search/speech/yh/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\r\n  from ._conv import register_converters as _register_converters\r\nTrain Epoch     #1:   0%|          | 0/5005 [00:00<?, ?it/s]SGD (\r\nParameter Group 0\r\n    dampening: 0\r\n    lr: 0.1\r\n    momentum: 0.9\r\n    nesterov: False\r\n    weight_decay: 5e-05\r\n)\r\n\r\n\r\n\r\n\r\nAs we can see above, only one gpu's lr is ` lr: 2.0796608879436464e-11`, but others is still 0.1 which is the init lr.\r\n\r\nbut in your code , `    # Horovod: broadcast parameters & optimizer state.\r\n    hvd.broadcast_parameters(model.state_dict(), root_rank=0)\r\n    hvd.broadcast_optimizer_state(optimizer, root_rank=0)`   which says : Horovod: broadcast parameters & optimizer state.  So maybe this is a very serious bug???\r\n\r\nPlease confirm it as soon as possible. Thanks a lot! :)\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        self._parameter_names = {v: k for k, v\n                                 in sorted(named_parameters)}\n        self._handles = {}\n        self._grad_accs = []\n\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_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 _make_hook(self, p):\n        def hook(*ignore):\n            assert p not in self._handles\n            assert not p.grad.requires_grad\n            name = self._parameter_names.get(p)\n\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            self._handles[p] = (handle, ctx)\n        return hook\n\n    def synchronize(self):\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        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    # Groups are unordered, but their params will be distinct\n    for group in state_dict['param_groups']:\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        self._parameter_names = {v: k for k, v\n                                 in sorted(named_parameters)}\n        self._handles = {}\n        self._grad_accs = []\n\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_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 _make_hook(self, p):\n        def hook(*ignore):\n            assert p not in self._handles\n            assert not p.grad.requires_grad\n            name = self._parameter_names.get(p)\n\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            self._handles[p] = (handle, ctx)\n        return hook\n\n    def synchronize(self):\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        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"}]}