We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[feat] Sync OSS.param_groups and the sharded optimizer param_groups
## 🚀 Feature
Make sure that all attributes (not just LR) are in sync in between the OSS.param_groups and the actual optimizer. 
 
## Motivation
Some frameworks make it possible to alter any attribute here, not just LR (momentum, beta for Adam, ..). We do not currently support this and silently fail

## Pitch
Part of the not really well defined PyTorch optimizer features which are nice to have

## Alternatives
- At least add a warning when we're out of sync
- Nuke the .param_groups access
- Force a new OSS optimizer to be created everytime the user wants to change an attribute in param_groups

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->


--- 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, 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     def partition_parameters(self) -> List[List[dict]]:
71         """Partitions parameters across distributed ranks.
72 
73         Returns a list of param_groups (which is a list of dict) where each
74         element of the list contains the param_groups for a rank. Element 0
75         corresponds to rank 0, etc. We need all the ranks for the broadcast
76         inside step().
77         """
78         world_size = dist.get_world_size(self.group)
79         param_groups: List[List] = [list() for _ in range(world_size)]
80         sizes = [0] * world_size
81         for param_group in self.param_groups:
82             param_lists: List[List] = [list() for _ in range(world_size)]
83             for param in param_group["params"]:
84                 # Add this param to rank with smallest size.
85                 rank = sizes.index(min(sizes))
86                 param_lists[rank].append(param)
87                 sizes[rank] += param.numel()
88             for rank, params in enumerate(param_lists):
89                 param_group_rank = copy.copy(param_group)
90                 param_group_rank["params"] = params
91                 param_groups[rank].append(param_group_rank)
92         return param_groups
93 
94     # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.
95     # For example, the apex library contains fused optimizers with a step that supports extra kwargs.
96     def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:
97         # Sync lr in case its been update by an LRScheduler.
98         self._sync_lr()
99 
100         # Run the optimizer step on this shard only
101         loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
102 
103         # Sync all the states
104         for rank, param_groups in enumerate(self.partition_parameters()):
105             for param_group in param_groups:
106                 for param in param_group["params"]:
107                     dist.broadcast(tensor=param, src=rank, group=self.group)
108         return loss
109 
110     def local_state_dict(self) -> dict:
111         """ Gets this rank's state_dict. """
112         return self.optim.state_dict()
113 
114     def consolidate_state_dict(self, recipient_rank: int = 0) -> None:
115         """ Update the consolidated state_dict list, one per rank.
116 
117         This needs to be called on all replicas """
118 
119         # Sync lr in case its been update by an LRScheduler.
120         self._sync_lr()
121 
122         if self.rank == recipient_rank:
123             # Pull the sharded state from all the other replicas
124             # Store all the states in order, rank by rank
125             logging.debug("Pulling the sharded optimizer state from all replicas")
126             self._all_states = self._collect_sharded_states()
127         else:
128             # Acknowledge broadcasts, and send this rank's shard when needed
129             self._broadcast_state_dict()
130 
131     def state_dict(self) -> Dict[str, Any]:
132         """
133         Return the last known global optimizer state, which consist of a list of the shards.
134 
135         NOTE: This is limited to the replica which was responsible for the consolidation.
136         The state may also not be up to date, depending on when `consolidate_state_dict` was last called.
137         """
138 
139         assert (
140             len(self._all_states) > 0
141         ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
142 
143         return {
144             "state": [s["state"] for s in self._all_states],
145             "param_groups": [s["param_groups"] for s in self._all_states],
146         }
147 
148     def load_local_state_dict(self, state_dict: dict) -> None:
149         """ Loads this rank's state_dict. """
150 
151         self.optim.load_state_dict(state_dict)
152 
153         # Workaround PyTorch bug that casts state (https://github.com/pytorch/pytorch/issues/43706)
154         # Copied from https://github.com/pytorch/fairseq/blob/v0.9.0/fairseq/optim/fp16_optimizer.py#L251-L268
155         groups = self.optim.param_groups
156         saved_groups = state_dict["param_groups"]
157         id_map = {
158             old_id: p
159             for old_id, p in zip(chain(*(g["params"] for g in saved_groups)), chain(*(g["params"] for g in groups)))
160         }
161         for k, v in state_dict["state"].items():
162             if k in id_map:
163                 param = id_map[k]
164                 self.optim.state[param] = recursive_copy_to_device(v, non_blocking=True, device=param.device)
165 
166         # Restore the global param_groups (the params themselves are already correct)
167         for global_group, local_group in zip(self.param_groups, groups):
168             for k, v in local_group.items():
169                 if k != "params":
170                     global_group[k] = v
171 
172     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
173         """ Restore the global parameter groups as well as the shard """
174         # Dispatch this rank's state dictionary to the wrapped shard optimizer
175         self.load_local_state_dict(
176             {"state": state_dict["state"][self.rank], "param_groups": state_dict["param_groups"][self.rank]}
177         )
178 
179         # Update the param_groups attribute for this instance
180         # TODO(ben)
181 
182     def add_param_group(self, param_group: dict) -> None:
183         super().add_param_group(param_group)
184         if not self.in_super_constructor:
185             param_groups = self.partition_parameters()[self.rank]
186             if len(param_groups) == len(self.optim.param_groups) + 1:
187                 self.optim.add_param_group(param_groups[-1])
188 
189     def _sync_lr(self) -> None:
190         """Sync learning rate (needed to support LRScheduler)."""
191         for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
192             local_group["lr"] = global_group["lr"]
193 
194     def _collect_sharded_states(self) -> List[Dict[str, Any]]:
195         """
196         Collect all the state shards, in CPU memory.
197         """
198         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
199         all_states: List[Dict[str, Any]] = []
200 
201         for rank in range(dist.get_world_size(group=self.group)):
202             if rank == self.rank:
203                 logging.debug("Saving self state")
204                 all_states.append(
205                     recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))
206                 )
207 
208                 # Sync with other replicas
209                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
210             else:
211                 # Fetch the optim state from the other replicas
212                 logging.debug("Receiving state from rank %s ", rank)
213                 replica_state = broadcast_object(
214                     empty_buffer, src_rank=rank, group=self.group, dist_device=self._device
215                 )
216 
217                 all_states.append(
218                     recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))
219                 )
220 
221                 logging.debug("State from rank %s received", rank)
222 
223         return all_states
224 
225     def _broadcast_state_dict(self) -> None:
226         """
227         Broadcast this rank's state shard, discard others
228         """
229         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
230 
231         for rank in range(dist.get_world_size(group=self.group)):
232             if rank == self.rank:
233                 # Send the state to the reference replica
234                 logging.debug(
235                     "Sending the sharded optimizer state to the reference replica from rank %s", rank,
236                 )
237                 broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
238             else:
239                 # Discard this tensor/rank, broadcast necessary for syncing
240                 logging.debug("Discarding broadcast from rank %s", rank)
241                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
242
```
--- 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
@@ -67,6 +67,12 @@
         # Current device is set by the parameters allocated to this rank
         self._device = split_param_groups[self.rank][0]["params"][0].device
 
+        # Sync local and global param_groups keys
+        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
+            for k, v in local_group.items():
+                if k != "params":
+                    global_group[k] = v
+
     def partition_parameters(self) -> List[List[dict]]:
         """Partitions parameters across distributed ranks.
 
@@ -94,8 +100,8 @@
     # NOTE(msb) We add a kwargs in order to support Optimizer sub-classes that support extra kwargs.
     # For example, the apex library contains fused optimizers with a step that supports extra kwargs.
     def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) -> Optional[float]:
-        # Sync lr in case its been update by an LRScheduler.
-        self._sync_lr()
+        # 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
@@ -116,8 +122,8 @@
 
         This needs to be called on all replicas """
 
-        # Sync lr in case its been update by an LRScheduler.
-        self._sync_lr()
+        # Sync lr and other attributes in case its been updated
+        self._sync_param_groups()
 
         if self.rank == recipient_rank:
             # Pull the sharded state from all the other replicas
@@ -176,9 +182,6 @@
             {"state": state_dict["state"][self.rank], "param_groups": state_dict["param_groups"][self.rank]}
         )
 
-        # Update the param_groups attribute for this instance
-        # TODO(ben)
-
     def add_param_group(self, param_group: dict) -> None:
         super().add_param_group(param_group)
         if not self.in_super_constructor:
@@ -186,10 +189,13 @@
             if len(param_groups) == len(self.optim.param_groups) + 1:
                 self.optim.add_param_group(param_groups[-1])
 
-    def _sync_lr(self) -> None:
-        """Sync learning rate (needed to support LRScheduler)."""
+    def _sync_param_groups(self) -> None:
+        """Sync learning rate and other optimizer attributes (needed to support schedulers)."""
         for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
-            local_group["lr"] = global_group["lr"]
+            for k in local_group.keys():
+                if k != "params":
+                    # Params have been sharded and should not be synced here
+                    local_group[k] = global_group[k]
 
     def _collect_sharded_states(self) -> List[Dict[str, Any]]:
         """

{"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@@ -67,6 +67,12 @@\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@@ -94,8 +100,8 @@\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 lr in case its been update by an LRScheduler.\n-        self._sync_lr()\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@@ -116,8 +122,8 @@\n \n         This needs to be called on all replicas \"\"\"\n \n-        # Sync lr in case its been update by an LRScheduler.\n-        self._sync_lr()\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@@ -176,9 +182,6 @@\n             {\"state\": state_dict[\"state\"][self.rank], \"param_groups\": state_dict[\"param_groups\"][self.rank]}\n         )\n \n-        # Update the param_groups attribute for this instance\n-        # TODO(ben)\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@@ -186,10 +189,13 @@\n             if len(param_groups) == len(self.optim.param_groups) + 1:\n                 self.optim.add_param_group(param_groups[-1])\n \n-    def _sync_lr(self) -> None:\n-        \"\"\"Sync learning rate (needed to support LRScheduler).\"\"\"\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-            local_group[\"lr\"] = global_group[\"lr\"]\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", "issue": "[feat] Sync OSS.param_groups and the sharded optimizer param_groups\n## \ud83d\ude80 Feature\r\nMake sure that all attributes (not just LR) are in sync in between the OSS.param_groups and the actual optimizer. \r\n \r\n## Motivation\r\nSome frameworks make it possible to alter any attribute here, not just LR (momentum, beta for Adam, ..). We do not currently support this and silently fail\r\n\r\n## Pitch\r\nPart of the not really well defined PyTorch optimizer features which are nice to have\r\n\r\n## Alternatives\r\n- At least add a warning when we're out of sync\r\n- Nuke the .param_groups access\r\n- Force a new OSS optimizer to be created everytime the user wants to change an attribute in param_groups\r\n\r\n## Additional context\r\n\r\n<!-- Add any other context or screenshots about the feature request here. -->\r\n\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, 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    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 lr in case its been update by an LRScheduler.\n        self._sync_lr()\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 in case its been update by an LRScheduler.\n        self._sync_lr()\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        return {\n            \"state\": [s[\"state\"] for s in self._all_states],\n            \"param_groups\": [s[\"param_groups\"] for s in self._all_states],\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        # Dispatch this rank's state dictionary to the wrapped shard optimizer\n        self.load_local_state_dict(\n            {\"state\": state_dict[\"state\"][self.rank], \"param_groups\": state_dict[\"param_groups\"][self.rank]}\n        )\n\n        # Update the param_groups attribute for this instance\n        # TODO(ben)\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_lr(self) -> None:\n        \"\"\"Sync learning rate (needed to support LRScheduler).\"\"\"\n        for global_group, local_group in zip(self.param_groups, self.optim.param_groups):\n            local_group[\"lr\"] = global_group[\"lr\"]\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, 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        return {\n            \"state\": [s[\"state\"] for s in self._all_states],\n            \"param_groups\": [s[\"param_groups\"] for s in self._all_states],\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        # Dispatch this rank's state dictionary to the wrapped shard optimizer\n        self.load_local_state_dict(\n            {\"state\": state_dict[\"state\"][self.rank], \"param_groups\": state_dict[\"param_groups\"][self.rank]}\n        )\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"}]}