We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[feat] OSS: Support nvidia's LARC
## 🚀 Feature
Make it possible to support LARC with OSS

## Motivation
LARC is a must have for large batch jobs, right now OSS will break on LARC because of the closure() being passed

## Pitch
Should be doable to gracefully handle optimizers with do not support closures in step()

## Alternatives
Not supporting LARC, reduces a lot of OSS interest

## Additional context

cc @mannatsingh @prigoyal @msbaines 

--- 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. Broadcast requests are issued async, we check completeness before moving on
110         requests = []
111         for rank, param_groups in enumerate(self.partition_parameters()):
112             for param_group in param_groups:
113                 for param in param_group["params"]:
114                     requests.append(dist.broadcast(tensor=param, src=rank, group=self.group, async_op=True))
115 
116         _ = list(map(lambda x: x.wait(), requests))
117         return loss
118 
119     def local_state_dict(self) -> dict:
120         """ Gets this rank's state_dict. """
121         return self.optim.state_dict()
122 
123     def consolidate_state_dict(self, recipient_rank: int = 0) -> None:
124         """ Update the consolidated state_dict list, one per rank.
125 
126         This needs to be called on all replicas """
127 
128         # Sync lr and other attributes in case its been updated
129         self._sync_param_groups()
130 
131         if self.rank == recipient_rank:
132             # Pull the sharded state from all the other replicas
133             # Store all the states in order, rank by rank
134             logging.debug("Pulling the sharded optimizer state from all replicas")
135             self._all_states = self._collect_sharded_states()
136         else:
137             # Acknowledge broadcasts, and send this rank's shard when needed
138             self._broadcast_state_dict()
139 
140     def state_dict(self) -> Dict[str, Any]:
141         """
142         Return the last known global optimizer state, which consist of a list of the shards.
143 
144         NOTE: This is limited to the replica which was responsible for the consolidation.
145         The state may also not be up to date, depending on when `consolidate_state_dict` was last called.
146         """
147 
148         assert (
149             len(self._all_states) > 0
150         ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
151 
152         # Flatten the param_groups, save the partition which logs the rank <> shard correspondence
153         partition: List[Tuple[int, int]] = []
154         param_groups: List[Dict[Any, Any]] = []
155 
156         start = 0
157         for i, s in enumerate(self._all_states):
158             param_groups.extend(s["param_groups"])
159             end = start + len(s["param_groups"])
160             partition.append((start, end))
161             start = end
162 
163         return {
164             "state": [s["state"] for s in self._all_states],
165             "param_groups": param_groups,
166             "partition": partition,
167         }
168 
169     def load_local_state_dict(self, state_dict: dict) -> None:
170         """ Loads this rank's state_dict. """
171 
172         self.optim.load_state_dict(state_dict)
173 
174         # Workaround PyTorch bug that casts state (https://github.com/pytorch/pytorch/issues/43706)
175         # Copied from https://github.com/pytorch/fairseq/blob/v0.9.0/fairseq/optim/fp16_optimizer.py#L251-L268
176         groups = self.optim.param_groups
177         saved_groups = state_dict["param_groups"]
178         id_map = {
179             old_id: p
180             for old_id, p in zip(chain(*(g["params"] for g in saved_groups)), chain(*(g["params"] for g in groups)))
181         }
182         for k, v in state_dict["state"].items():
183             if k in id_map:
184                 param = id_map[k]
185                 self.optim.state[param] = recursive_copy_to_device(v, non_blocking=True, device=param.device)
186 
187         # Restore the global param_groups (the params themselves are already correct)
188         for global_group, local_group in zip(self.param_groups, groups):
189             for k, v in local_group.items():
190                 if k != "params":
191                     global_group[k] = v
192 
193     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
194         """ Restore the global parameter groups as well as the shard """
195 
196         # Get this optimizer's param_groups shard
197         param_groups = state_dict["param_groups"][
198             state_dict["partition"][self.rank][0] : state_dict["partition"][self.rank][1]
199         ]
200         # Dispatch this rank's state dictionary to the wrapped shard optimizer
201         self.load_local_state_dict({"state": state_dict["state"][self.rank], "param_groups": param_groups})
202 
203     def add_param_group(self, param_group: dict) -> None:
204         super().add_param_group(param_group)
205         if not self.in_super_constructor:
206             param_groups = self.partition_parameters()[self.rank]
207             if len(param_groups) == len(self.optim.param_groups) + 1:
208                 self.optim.add_param_group(param_groups[-1])
209 
210     def _sync_param_groups(self) -> None:
211         """Sync learning rate and other optimizer attributes (needed to support schedulers)."""
212         for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
213             for k in local_group.keys():
214                 if k != "params":
215                     # Params have been sharded and should not be synced here
216                     local_group[k] = global_group[k]
217 
218     def _collect_sharded_states(self) -> List[Dict[str, Any]]:
219         """
220         Collect all the state shards, in CPU memory.
221         """
222         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
223         all_states: List[Dict[str, Any]] = []
224 
225         for rank in range(dist.get_world_size(group=self.group)):
226             if rank == self.rank:
227                 logging.debug("Saving self state")
228                 all_states.append(
229                     recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))
230                 )
231 
232                 # Sync with other replicas
233                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
234             else:
235                 # Fetch the optim state from the other replicas
236                 logging.debug("Receiving state from rank %s ", rank)
237                 replica_state = broadcast_object(
238                     empty_buffer, src_rank=rank, group=self.group, dist_device=self._device
239                 )
240 
241                 all_states.append(
242                     recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))
243                 )
244 
245                 logging.debug("State from rank %s received", rank)
246 
247         return all_states
248 
249     def _broadcast_state_dict(self) -> None:
250         """
251         Broadcast this rank's state shard, discard others
252         """
253         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
254 
255         for rank in range(dist.get_world_size(group=self.group)):
256             if rank == self.rank:
257                 # Send the state to the reference replica
258                 logging.debug(
259                     "Sending the sharded optimizer state to the reference replica from rank %s", rank,
260                 )
261                 broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
262             else:
263                 # Discard this tensor/rank, broadcast necessary for syncing
264                 logging.debug("Discarding broadcast from rank %s", rank)
265                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
266
```
--- 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
@@ -103,8 +103,11 @@
         # Sync oss param_groups attributes in case they've been updated by a scheduler.
         self._sync_param_groups()
 
-        # Run the optimizer step on this shard only
-        loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
+        # Run the optimizer step on this shard only:
+        if closure is not None:
+            loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
+        else:
+            loss = self.optim.step(**kwargs)
 
         # Sync all the states. Broadcast requests are issued async, we check completeness before moving on
         requests = []

{"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@@ -103,8 +103,11 @@\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+        # Run the optimizer step on this shard only:\n+        if closure is not None:\n+            loss = self.optim.step(closure=closure, **kwargs)  # type: ignore\n+        else:\n+            loss = self.optim.step(**kwargs)\n \n         # Sync all the states. Broadcast requests are issued async, we check completeness before moving on\n         requests = []\n", "issue": "[feat] OSS: Support nvidia's LARC\n## \ud83d\ude80 Feature\r\nMake it possible to support LARC with OSS\r\n\r\n## Motivation\r\nLARC is a must have for large batch jobs, right now OSS will break on LARC because of the closure() being passed\r\n\r\n## Pitch\r\nShould be doable to gracefully handle optimizers with do not support closures in step()\r\n\r\n## Alternatives\r\nNot supporting LARC, reduces a lot of OSS interest\r\n\r\n## Additional context\r\n\r\ncc @mannatsingh @prigoyal @msbaines \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. 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"}], "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        if closure is not None:\n            loss = self.optim.step(closure=closure, **kwargs)  # type: ignore\n        else:\n            loss = self.optim.step(**kwargs)\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"}]}