We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Async broadcast request in OSS
## 🚀 Feature
Make sure that the broadcast request to sync shards are async, and not blocking the main thread

## Motivation
OSS is currently too slow

## Pitch
Easy fix, no compromise

## Alternatives
Stay slow for no reason

## Additional context
Clear bottleneck when tracing execution

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `fairscale/optim/oss.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
2 #
3 # This source code is licensed under the BSD license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 import copy
7 from itertools import chain
8 import logging
9 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type
10 
11 import torch
12 import torch.distributed as dist
13 from torch.optim import SGD, Optimizer
14 
15 from .utils import broadcast_object, recursive_copy_to_device
16 
17 if TYPE_CHECKING:  # pragma: no cover
18     from torch.optim.optimizer import _params_t
19 else:
20     _params_t = Any
21 
22 
23 class OSS(Optimizer):
24     """Wraps an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>`
25     optimizer and shards its state as described by ZeRO_.
26     ::
27         opt = OSS(params, optim=torch.optim.Adam, lr=0.01)
28 
29     .. _ZeRO: https://arxiv.org/abs/1910.02054
30 
31     We use a greedy algorithm to pack a number of parameters
32     at each rank. Each parameter belongs to a single rank and
33     is not divided among rank.
34 
35     After each rank completed their parameter update, they broadcast
36     the new version of the parameters to all other ranks to synchronize
37     the parameters for next round forward/backward computation.
38 
39     Args:
40         params (list of tensors):
41             parameters to be optimized
42     Keyword Args:
43         optim (torch.nn.Optimizer):
44             optimizer to shard (default: SGD)
45         group (group):
46             torch.distributed group (default: group.WORLD)
47     """
48 
49     optim: Optimizer
50     in_super_constructor: bool
51 
52     def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any = dist.group.WORLD, **defaults: Any):
53         # Hold all the model params in the root .param_groups
54         self.in_super_constructor = True
55         super().__init__(params, defaults)
56         self.in_super_constructor = False
57 
58         # Build the wrapped optimizer, responsible for a shard of the params
59         self.group = group
60         self.rank = dist.get_rank(group)
61         split_param_groups = self.partition_parameters()
62         self.optim = optim(split_param_groups[self.rank], **defaults)
63 
64         # Optional consolidated optimizer state
65         self._all_states: List[Dict[str, Any]] = []
66 
67         # Current device is set by the parameters allocated to this rank
68         self._device = split_param_groups[self.rank][0]["params"][0].device
69 
70         # Sync local and global param_groups keys
71         for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
72             for k, v in local_group.items():
73                 if k != "params":
74                     global_group[k] = v
75 
76     def partition_parameters(self) -> List[List[dict]]:
77         """Partitions parameters across distributed ranks.
78 
79         Returns a list of param_groups (which is a list of dict) where each
80         element of the list contains the param_groups for a rank. Element 0
81         corresponds to rank 0, etc. We need all the ranks for the broadcast
82         inside step().
83         """
84         world_size = dist.get_world_size(self.group)
85         param_groups: List[List] = [list() for _ in range(world_size)]
86         sizes = [0] * world_size
87         for param_group in self.param_groups:
88             param_lists: List[List] = [list() for _ in range(world_size)]
89             for param in param_group["params"]:
90                 # Add this param to rank with smallest size.
91                 rank = sizes.index(min(sizes))
92                 param_lists[rank].append(param)
93                 sizes[rank] += param.numel()
94             for rank, params in enumerate(param_lists):
95                 param_group_rank = copy.copy(param_group)
96                 param_group_rank["params"] = params
97                 param_groups[rank].append(param_group_rank)
98         return param_groups
99 
100     # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.
101     # For example, the apex library contains fused optimizers with a step that supports extra kwargs.
102     def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:
103         # Sync oss param_groups attributes in case they've been updated by a scheduler.
104         self._sync_param_groups()
105 
106         # Run the optimizer step on this shard only
107         loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
108 
109         # Sync all the states
110         for rank, param_groups in enumerate(self.partition_parameters()):
111             for param_group in param_groups:
112                 for param in param_group["params"]:
113                     dist.broadcast(tensor=param, src=rank, group=self.group)
114         return loss
115 
116     def local_state_dict(self) -> dict:
117         """ Gets this rank's state_dict. """
118         return self.optim.state_dict()
119 
120     def consolidate_state_dict(self, recipient_rank: int = 0) -> None:
121         """ Update the consolidated state_dict list, one per rank.
122 
123         This needs to be called on all replicas """
124 
125         # Sync lr and other attributes in case its been updated
126         self._sync_param_groups()
127 
128         if self.rank == recipient_rank:
129             # Pull the sharded state from all the other replicas
130             # Store all the states in order, rank by rank
131             logging.debug("Pulling the sharded optimizer state from all replicas")
132             self._all_states = self._collect_sharded_states()
133         else:
134             # Acknowledge broadcasts, and send this rank's shard when needed
135             self._broadcast_state_dict()
136 
137     def state_dict(self) -> Dict[str, Any]:
138         """
139         Return the last known global optimizer state, which consist of a list of the shards.
140 
141         NOTE: This is limited to the replica which was responsible for the consolidation.
142         The state may also not be up to date, depending on when `consolidate_state_dict` was last called.
143         """
144 
145         assert (
146             len(self._all_states) > 0
147         ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
148 
149         # Flatten the param_groups, save the partition which logs the rank <> shard correspondence
150         partition: List[Tuple[int, int]] = []
151         param_groups: List[Dict[Any, Any]] = []
152 
153         start = 0
154         for i, s in enumerate(self._all_states):
155             param_groups.extend(s["param_groups"])
156             end = start + len(s["param_groups"])
157             partition.append((start, end))
158             start = end
159 
160         return {
161             "state": [s["state"] for s in self._all_states],
162             "param_groups": param_groups,
163             "partition": partition,
164         }
165 
166     def load_local_state_dict(self, state_dict: dict) -> None:
167         """ Loads this rank's state_dict. """
168 
169         self.optim.load_state_dict(state_dict)
170 
171         # Workaround PyTorch bug that casts state (https://github.com/pytorch/pytorch/issues/43706)
172         # Copied from https://github.com/pytorch/fairseq/blob/v0.9.0/fairseq/optim/fp16_optimizer.py#L251-L268
173         groups = self.optim.param_groups
174         saved_groups = state_dict["param_groups"]
175         id_map = {
176             old_id: p
177             for old_id, p in zip(chain(*(g["params"] for g in saved_groups)), chain(*(g["params"] for g in groups)))
178         }
179         for k, v in state_dict["state"].items():
180             if k in id_map:
181                 param = id_map[k]
182                 self.optim.state[param] = recursive_copy_to_device(v, non_blocking=True, device=param.device)
183 
184         # Restore the global param_groups (the params themselves are already correct)
185         for global_group, local_group in zip(self.param_groups, groups):
186             for k, v in local_group.items():
187                 if k != "params":
188                     global_group[k] = v
189 
190     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
191         """ Restore the global parameter groups as well as the shard """
192 
193         # Get this optimizer's param_groups shard
194         param_groups = state_dict["param_groups"][
195             state_dict["partition"][self.rank][0] : state_dict["partition"][self.rank][1]
196         ]
197         # Dispatch this rank's state dictionary to the wrapped shard optimizer
198         self.load_local_state_dict({"state": state_dict["state"][self.rank], "param_groups": param_groups})
199 
200     def add_param_group(self, param_group: dict) -> None:
201         super().add_param_group(param_group)
202         if not self.in_super_constructor:
203             param_groups = self.partition_parameters()[self.rank]
204             if len(param_groups) == len(self.optim.param_groups) + 1:
205                 self.optim.add_param_group(param_groups[-1])
206 
207     def _sync_param_groups(self) -> None:
208         """Sync learning rate and other optimizer attributes (needed to support schedulers)."""
209         for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
210             for k in local_group.keys():
211                 if k != "params":
212                     # Params have been sharded and should not be synced here
213                     local_group[k] = global_group[k]
214 
215     def _collect_sharded_states(self) -> List[Dict[str, Any]]:
216         """
217         Collect all the state shards, in CPU memory.
218         """
219         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
220         all_states: List[Dict[str, Any]] = []
221 
222         for rank in range(dist.get_world_size(group=self.group)):
223             if rank == self.rank:
224                 logging.debug("Saving self state")
225                 all_states.append(
226                     recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))
227                 )
228 
229                 # Sync with other replicas
230                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
231             else:
232                 # Fetch the optim state from the other replicas
233                 logging.debug("Receiving state from rank %s ", rank)
234                 replica_state = broadcast_object(
235                     empty_buffer, src_rank=rank, group=self.group, dist_device=self._device
236                 )
237 
238                 all_states.append(
239                     recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))
240                 )
241 
242                 logging.debug("State from rank %s received", rank)
243 
244         return all_states
245 
246     def _broadcast_state_dict(self) -> None:
247         """
248         Broadcast this rank's state shard, discard others
249         """
250         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
251 
252         for rank in range(dist.get_world_size(group=self.group)):
253             if rank == self.rank:
254                 # Send the state to the reference replica
255                 logging.debug(
256                     "Sending the sharded optimizer state to the reference replica from rank %s", rank,
257                 )
258                 broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
259             else:
260                 # Discard this tensor/rank, broadcast necessary for syncing
261                 logging.debug("Discarding broadcast from rank %s", rank)
262                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
263
```
--- 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/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -106,11 +106,14 @@
         # Run the optimizer step on this shard only
         loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
 
-        # Sync all the states
+        # Sync all the states. Broadcast requests are issued async, we check completeness before moving on
+        requests = []
         for rank, param_groups in enumerate(self.partition_parameters()):
             for param_group in param_groups:
                 for param in param_group["params"]:
-                    dist.broadcast(tensor=param, src=rank, group=self.group)
+                    requests.append(dist.broadcast(tensor=param, src=rank, group=self.group, async_op=True))
+
+        _ = list(map(lambda x: x.wait(), requests))
         return loss
 
     def local_state_dict(self) -> dict:

{"golden_diff": "diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py\n--- a/fairscale/optim/oss.py\n+++ b/fairscale/optim/oss.py\n@@ -106,11 +106,14 @@\n         # Run the optimizer step on this shard only\n         loss = self.optim.step(closure=closure, **kwargs)  # type: ignore\n \n-        # Sync all the states\n+        # Sync all the states. Broadcast requests are issued async, we check completeness before moving on\n+        requests = []\n         for rank, param_groups in enumerate(self.partition_parameters()):\n             for param_group in param_groups:\n                 for param in param_group[\"params\"]:\n-                    dist.broadcast(tensor=param, src=rank, group=self.group)\n+                    requests.append(dist.broadcast(tensor=param, src=rank, group=self.group, async_op=True))\n+\n+        _ = list(map(lambda x: x.wait(), requests))\n         return loss\n \n     def local_state_dict(self) -> dict:\n", "issue": "Async broadcast request in OSS\n## \ud83d\ude80 Feature\r\nMake sure that the broadcast request to sync shards are async, and not blocking the main thread\r\n\r\n## Motivation\r\nOSS is currently too slow\r\n\r\n## Pitch\r\nEasy fix, no compromise\r\n\r\n## Alternatives\r\nStay slow for no reason\r\n\r\n## Additional context\r\nClear bottleneck when tracing execution\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport copy\nfrom itertools import chain\nimport logging\nfrom typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type\n\nimport torch\nimport torch.distributed as dist\nfrom torch.optim import SGD, Optimizer\n\nfrom .utils import broadcast_object, recursive_copy_to_device\n\nif TYPE_CHECKING:  # pragma: no cover\n    from torch.optim.optimizer import _params_t\nelse:\n    _params_t = Any\n\n\nclass OSS(Optimizer):\n    \"\"\"Wraps an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>`\n    optimizer and shards its state as described by ZeRO_.\n    ::\n        opt = OSS(params, optim=torch.optim.Adam, lr=0.01)\n\n    .. _ZeRO: https://arxiv.org/abs/1910.02054\n\n    We use a greedy algorithm to pack a number of parameters\n    at each rank. Each parameter belongs to a single rank and\n    is not divided among rank.\n\n    After each rank completed their parameter update, they broadcast\n    the new version of the parameters to all other ranks to synchronize\n    the parameters for next round forward/backward computation.\n\n    Args:\n        params (list of tensors):\n            parameters to be optimized\n    Keyword Args:\n        optim (torch.nn.Optimizer):\n            optimizer to shard (default: SGD)\n        group (group):\n            torch.distributed group (default: group.WORLD)\n    \"\"\"\n\n    optim: Optimizer\n    in_super_constructor: bool\n\n    def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any = dist.group.WORLD, **defaults: Any):\n        # Hold all the model params in the root .param_groups\n        self.in_super_constructor = True\n        super().__init__(params, defaults)\n        self.in_super_constructor = False\n\n        # Build the wrapped optimizer, responsible for a shard of the params\n        self.group = group\n        self.rank = dist.get_rank(group)\n        split_param_groups = self.partition_parameters()\n        self.optim = optim(split_param_groups[self.rank], **defaults)\n\n        # Optional consolidated optimizer state\n        self._all_states: List[Dict[str, Any]] = []\n\n        # Current device is set by the parameters allocated to this rank\n        self._device = split_param_groups[self.rank][0][\"params\"][0].device\n\n        # Sync local and global param_groups keys\n        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):\n            for k, v in local_group.items():\n                if k != \"params\":\n                    global_group[k] = v\n\n    def partition_parameters(self) -> List[List[dict]]:\n        \"\"\"Partitions parameters across distributed ranks.\n\n        Returns a list of param_groups (which is a list of dict) where each\n        element of the list contains the param_groups for a rank. Element 0\n        corresponds to rank 0, etc. We need all the ranks for the broadcast\n        inside step().\n        \"\"\"\n        world_size = dist.get_world_size(self.group)\n        param_groups: List[List] = [list() for _ in range(world_size)]\n        sizes = [0] * world_size\n        for param_group in self.param_groups:\n            param_lists: List[List] = [list() for _ in range(world_size)]\n            for param in param_group[\"params\"]:\n                # Add this param to rank with smallest size.\n                rank = sizes.index(min(sizes))\n                param_lists[rank].append(param)\n                sizes[rank] += param.numel()\n            for rank, params in enumerate(param_lists):\n                param_group_rank = copy.copy(param_group)\n                param_group_rank[\"params\"] = params\n                param_groups[rank].append(param_group_rank)\n        return param_groups\n\n    # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.\n    # For example, the apex library contains fused optimizers with a step that supports extra kwargs.\n    def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:\n        # Sync oss param_groups attributes in case they've been updated by a scheduler.\n        self._sync_param_groups()\n\n        # Run the optimizer step on this shard only\n        loss = self.optim.step(closure=closure, **kwargs)  # type: ignore\n\n        # Sync all the states\n        for rank, param_groups in enumerate(self.partition_parameters()):\n            for param_group in param_groups:\n                for param in param_group[\"params\"]:\n                    dist.broadcast(tensor=param, src=rank, group=self.group)\n        return loss\n\n    def local_state_dict(self) -> dict:\n        \"\"\" Gets this rank's state_dict. \"\"\"\n        return self.optim.state_dict()\n\n    def consolidate_state_dict(self, recipient_rank: int = 0) -> None:\n        \"\"\" Update the consolidated state_dict list, one per rank.\n\n        This needs to be called on all replicas \"\"\"\n\n        # Sync lr and other attributes in case its been updated\n        self._sync_param_groups()\n\n        if self.rank == recipient_rank:\n            # Pull the sharded state from all the other replicas\n            # Store all the states in order, rank by rank\n            logging.debug(\"Pulling the sharded optimizer state from all replicas\")\n            self._all_states = self._collect_sharded_states()\n        else:\n            # Acknowledge broadcasts, and send this rank's shard when needed\n            self._broadcast_state_dict()\n\n    def state_dict(self) -> Dict[str, Any]:\n        \"\"\"\n        Return the last known global optimizer state, which consist of a list of the shards.\n\n        NOTE: This is limited to the replica which was responsible for the consolidation.\n        The state may also not be up to date, depending on when `consolidate_state_dict` was last called.\n        \"\"\"\n\n        assert (\n            len(self._all_states) > 0\n        ), \"The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand\"\n\n        # Flatten the param_groups, save the partition which logs the rank <> shard correspondence\n        partition: List[Tuple[int, int]] = []\n        param_groups: List[Dict[Any, Any]] = []\n\n        start = 0\n        for i, s in enumerate(self._all_states):\n            param_groups.extend(s[\"param_groups\"])\n            end = start + len(s[\"param_groups\"])\n            partition.append((start, end))\n            start = end\n\n        return {\n            \"state\": [s[\"state\"] for s in self._all_states],\n            \"param_groups\": param_groups,\n            \"partition\": partition,\n        }\n\n    def load_local_state_dict(self, state_dict: dict) -> None:\n        \"\"\" Loads this rank's state_dict. \"\"\"\n\n        self.optim.load_state_dict(state_dict)\n\n        # Workaround PyTorch bug that casts state (https://github.com/pytorch/pytorch/issues/43706)\n        # Copied from https://github.com/pytorch/fairseq/blob/v0.9.0/fairseq/optim/fp16_optimizer.py#L251-L268\n        groups = self.optim.param_groups\n        saved_groups = state_dict[\"param_groups\"]\n        id_map = {\n            old_id: p\n            for old_id, p in zip(chain(*(g[\"params\"] for g in saved_groups)), chain(*(g[\"params\"] for g in groups)))\n        }\n        for k, v in state_dict[\"state\"].items():\n            if k in id_map:\n                param = id_map[k]\n                self.optim.state[param] = recursive_copy_to_device(v, non_blocking=True, device=param.device)\n\n        # Restore the global param_groups (the params themselves are already correct)\n        for global_group, local_group in zip(self.param_groups, groups):\n            for k, v in local_group.items():\n                if k != \"params\":\n                    global_group[k] = v\n\n    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:\n        \"\"\" Restore the global parameter groups as well as the shard \"\"\"\n\n        # Get this optimizer's param_groups shard\n        param_groups = state_dict[\"param_groups\"][\n            state_dict[\"partition\"][self.rank][0] : state_dict[\"partition\"][self.rank][1]\n        ]\n        # Dispatch this rank's state dictionary to the wrapped shard optimizer\n        self.load_local_state_dict({\"state\": state_dict[\"state\"][self.rank], \"param_groups\": param_groups})\n\n    def add_param_group(self, param_group: dict) -> None:\n        super().add_param_group(param_group)\n        if not self.in_super_constructor:\n            param_groups = self.partition_parameters()[self.rank]\n            if len(param_groups) == len(self.optim.param_groups) + 1:\n                self.optim.add_param_group(param_groups[-1])\n\n    def _sync_param_groups(self) -> None:\n        \"\"\"Sync learning rate and other optimizer attributes (needed to support schedulers).\"\"\"\n        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):\n            for k in local_group.keys():\n                if k != \"params\":\n                    # Params have been sharded and should not be synced here\n                    local_group[k] = global_group[k]\n\n    def _collect_sharded_states(self) -> List[Dict[str, Any]]:\n        \"\"\"\n        Collect all the state shards, in CPU memory.\n        \"\"\"\n        empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)\n        all_states: List[Dict[str, Any]] = []\n\n        for rank in range(dist.get_world_size(group=self.group)):\n            if rank == self.rank:\n                logging.debug(\"Saving self state\")\n                all_states.append(\n                    recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device(\"cpu\"))\n                )\n\n                # Sync with other replicas\n                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)\n            else:\n                # Fetch the optim state from the other replicas\n                logging.debug(\"Receiving state from rank %s \", rank)\n                replica_state = broadcast_object(\n                    empty_buffer, src_rank=rank, group=self.group, dist_device=self._device\n                )\n\n                all_states.append(\n                    recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device(\"cpu\"))\n                )\n\n                logging.debug(\"State from rank %s received\", rank)\n\n        return all_states\n\n    def _broadcast_state_dict(self) -> None:\n        \"\"\"\n        Broadcast this rank's state shard, discard others\n        \"\"\"\n        empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)\n\n        for rank in range(dist.get_world_size(group=self.group)):\n            if rank == self.rank:\n                # Send the state to the reference replica\n                logging.debug(\n                    \"Sending the sharded optimizer state to the reference replica from rank %s\", rank,\n                )\n                broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)\n            else:\n                # Discard this tensor/rank, broadcast necessary for syncing\n                logging.debug(\"Discarding broadcast from rank %s\", rank)\n                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)\n", "path": "fairscale/optim/oss.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport copy\nfrom itertools import chain\nimport logging\nfrom typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type\n\nimport torch\nimport torch.distributed as dist\nfrom torch.optim import SGD, Optimizer\n\nfrom .utils import broadcast_object, recursive_copy_to_device\n\nif TYPE_CHECKING:  # pragma: no cover\n    from torch.optim.optimizer import _params_t\nelse:\n    _params_t = Any\n\n\nclass OSS(Optimizer):\n    \"\"\"Wraps an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>`\n    optimizer and shards its state as described by ZeRO_.\n    ::\n        opt = OSS(params, optim=torch.optim.Adam, lr=0.01)\n\n    .. _ZeRO: https://arxiv.org/abs/1910.02054\n\n    We use a greedy algorithm to pack a number of parameters\n    at each rank. Each parameter belongs to a single rank and\n    is not divided among rank.\n\n    After each rank completed their parameter update, they broadcast\n    the new version of the parameters to all other ranks to synchronize\n    the parameters for next round forward/backward computation.\n\n    Args:\n        params (list of tensors):\n            parameters to be optimized\n    Keyword Args:\n        optim (torch.nn.Optimizer):\n            optimizer to shard (default: SGD)\n        group (group):\n            torch.distributed group (default: group.WORLD)\n    \"\"\"\n\n    optim: Optimizer\n    in_super_constructor: bool\n\n    def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any = dist.group.WORLD, **defaults: Any):\n        # Hold all the model params in the root .param_groups\n        self.in_super_constructor = True\n        super().__init__(params, defaults)\n        self.in_super_constructor = False\n\n        # Build the wrapped optimizer, responsible for a shard of the params\n        self.group = group\n        self.rank = dist.get_rank(group)\n        split_param_groups = self.partition_parameters()\n        self.optim = optim(split_param_groups[self.rank], **defaults)\n\n        # Optional consolidated optimizer state\n        self._all_states: List[Dict[str, Any]] = []\n\n        # Current device is set by the parameters allocated to this rank\n        self._device = split_param_groups[self.rank][0][\"params\"][0].device\n\n        # Sync local and global param_groups keys\n        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):\n            for k, v in local_group.items():\n                if k != \"params\":\n                    global_group[k] = v\n\n    def partition_parameters(self) -> List[List[dict]]:\n        \"\"\"Partitions parameters across distributed ranks.\n\n        Returns a list of param_groups (which is a list of dict) where each\n        element of the list contains the param_groups for a rank. Element 0\n        corresponds to rank 0, etc. We need all the ranks for the broadcast\n        inside step().\n        \"\"\"\n        world_size = dist.get_world_size(self.group)\n        param_groups: List[List] = [list() for _ in range(world_size)]\n        sizes = [0] * world_size\n        for param_group in self.param_groups:\n            param_lists: List[List] = [list() for _ in range(world_size)]\n            for param in param_group[\"params\"]:\n                # Add this param to rank with smallest size.\n                rank = sizes.index(min(sizes))\n                param_lists[rank].append(param)\n                sizes[rank] += param.numel()\n            for rank, params in enumerate(param_lists):\n                param_group_rank = copy.copy(param_group)\n                param_group_rank[\"params\"] = params\n                param_groups[rank].append(param_group_rank)\n        return param_groups\n\n    # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.\n    # For example, the apex library contains fused optimizers with a step that supports extra kwargs.\n    def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:\n        # Sync oss param_groups attributes in case they've been updated by a scheduler.\n        self._sync_param_groups()\n\n        # Run the optimizer step on this shard only\n        loss = self.optim.step(closure=closure, **kwargs)  # type: ignore\n\n        # Sync all the states. Broadcast requests are issued async, we check completeness before moving on\n        requests = []\n        for rank, param_groups in enumerate(self.partition_parameters()):\n            for param_group in param_groups:\n                for param in param_group[\"params\"]:\n                    requests.append(dist.broadcast(tensor=param, src=rank, group=self.group, async_op=True))\n\n        _ = list(map(lambda x: x.wait(), requests))\n        return loss\n\n    def local_state_dict(self) -> dict:\n        \"\"\" Gets this rank's state_dict. \"\"\"\n        return self.optim.state_dict()\n\n    def consolidate_state_dict(self, recipient_rank: int = 0) -> None:\n        \"\"\" Update the consolidated state_dict list, one per rank.\n\n        This needs to be called on all replicas \"\"\"\n\n        # Sync lr and other attributes in case its been updated\n        self._sync_param_groups()\n\n        if self.rank == recipient_rank:\n            # Pull the sharded state from all the other replicas\n            # Store all the states in order, rank by rank\n            logging.debug(\"Pulling the sharded optimizer state from all replicas\")\n            self._all_states = self._collect_sharded_states()\n        else:\n            # Acknowledge broadcasts, and send this rank's shard when needed\n            self._broadcast_state_dict()\n\n    def state_dict(self) -> Dict[str, Any]:\n        \"\"\"\n        Return the last known global optimizer state, which consist of a list of the shards.\n\n        NOTE: This is limited to the replica which was responsible for the consolidation.\n        The state may also not be up to date, depending on when `consolidate_state_dict` was last called.\n        \"\"\"\n\n        assert (\n            len(self._all_states) > 0\n        ), \"The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand\"\n\n        # Flatten the param_groups, save the partition which logs the rank <> shard correspondence\n        partition: List[Tuple[int, int]] = []\n        param_groups: List[Dict[Any, Any]] = []\n\n        start = 0\n        for i, s in enumerate(self._all_states):\n            param_groups.extend(s[\"param_groups\"])\n            end = start + len(s[\"param_groups\"])\n            partition.append((start, end))\n            start = end\n\n        return {\n            \"state\": [s[\"state\"] for s in self._all_states],\n            \"param_groups\": param_groups,\n            \"partition\": partition,\n        }\n\n    def load_local_state_dict(self, state_dict: dict) -> None:\n        \"\"\" Loads this rank's state_dict. \"\"\"\n\n        self.optim.load_state_dict(state_dict)\n\n        # Workaround PyTorch bug that casts state (https://github.com/pytorch/pytorch/issues/43706)\n        # Copied from https://github.com/pytorch/fairseq/blob/v0.9.0/fairseq/optim/fp16_optimizer.py#L251-L268\n        groups = self.optim.param_groups\n        saved_groups = state_dict[\"param_groups\"]\n        id_map = {\n            old_id: p\n            for old_id, p in zip(chain(*(g[\"params\"] for g in saved_groups)), chain(*(g[\"params\"] for g in groups)))\n        }\n        for k, v in state_dict[\"state\"].items():\n            if k in id_map:\n                param = id_map[k]\n                self.optim.state[param] = recursive_copy_to_device(v, non_blocking=True, device=param.device)\n\n        # Restore the global param_groups (the params themselves are already correct)\n        for global_group, local_group in zip(self.param_groups, groups):\n            for k, v in local_group.items():\n                if k != \"params\":\n                    global_group[k] = v\n\n    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:\n        \"\"\" Restore the global parameter groups as well as the shard \"\"\"\n\n        # Get this optimizer's param_groups shard\n        param_groups = state_dict[\"param_groups\"][\n            state_dict[\"partition\"][self.rank][0] : state_dict[\"partition\"][self.rank][1]\n        ]\n        # Dispatch this rank's state dictionary to the wrapped shard optimizer\n        self.load_local_state_dict({\"state\": state_dict[\"state\"][self.rank], \"param_groups\": param_groups})\n\n    def add_param_group(self, param_group: dict) -> None:\n        super().add_param_group(param_group)\n        if not self.in_super_constructor:\n            param_groups = self.partition_parameters()[self.rank]\n            if len(param_groups) == len(self.optim.param_groups) + 1:\n                self.optim.add_param_group(param_groups[-1])\n\n    def _sync_param_groups(self) -> None:\n        \"\"\"Sync learning rate and other optimizer attributes (needed to support schedulers).\"\"\"\n        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):\n            for k in local_group.keys():\n                if k != \"params\":\n                    # Params have been sharded and should not be synced here\n                    local_group[k] = global_group[k]\n\n    def _collect_sharded_states(self) -> List[Dict[str, Any]]:\n        \"\"\"\n        Collect all the state shards, in CPU memory.\n        \"\"\"\n        empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)\n        all_states: List[Dict[str, Any]] = []\n\n        for rank in range(dist.get_world_size(group=self.group)):\n            if rank == self.rank:\n                logging.debug(\"Saving self state\")\n                all_states.append(\n                    recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device(\"cpu\"))\n                )\n\n                # Sync with other replicas\n                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)\n            else:\n                # Fetch the optim state from the other replicas\n                logging.debug(\"Receiving state from rank %s \", rank)\n                replica_state = broadcast_object(\n                    empty_buffer, src_rank=rank, group=self.group, dist_device=self._device\n                )\n\n                all_states.append(\n                    recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device(\"cpu\"))\n                )\n\n                logging.debug(\"State from rank %s received\", rank)\n\n        return all_states\n\n    def _broadcast_state_dict(self) -> None:\n        \"\"\"\n        Broadcast this rank's state shard, discard others\n        \"\"\"\n        empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)\n\n        for rank in range(dist.get_world_size(group=self.group)):\n            if rank == self.rank:\n                # Send the state to the reference replica\n                logging.debug(\n                    \"Sending the sharded optimizer state to the reference replica from rank %s\", rank,\n                )\n                broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)\n            else:\n                # Discard this tensor/rank, broadcast necessary for syncing\n                logging.debug(\"Discarding broadcast from rank %s\", rank)\n                broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)\n", "path": "fairscale/optim/oss.py"}]}