We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
OSS to expose classic optimizer compliant param_groups
## 🚀 Feature
Expose the "classic" param_groups from the OSS optimizer wrapper, either via a new API (not as nice) or if possible via the classic .param_groups attribute (get and set) 

## Motivation
Some frameworks like classy vision use that to alter the LR per layer or over time (or both). Exposing this would make OSS integration a lot easier

## Pitch
Either use dedicated getter and setter properties which change behaviour depending on whether the call comes from the outside or the super constructor, or add a specific API to get the consolidated param_groups and set the consolidated param_groups

## Alternatives
- Destroy the optimizer every time there's an adjustment
- Only expose adjustments which apply to the whole param_groups ("set all LR to a given value")


--- 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/utils.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 io
7 from typing import Any, Dict
8 
9 import torch
10 from torch._six import container_abcs
11 import torch.distributed as dist
12 
13 
14 # Credits:  classy_vision/generic/distributed_util.py
15 def recursive_copy_to_device(value: Any, non_blocking: bool, device: torch.device) -> Any:
16     """
17     Recursively searches lists, tuples, dicts and copies tensors to device if
18     possible. Non-tensor values are passed as-is in the result.
19 
20     NOTE:  These are all copies, so if there are two objects that reference
21     the same object, then after this call, there will be two different objects
22     referenced on the device.
23     """
24 
25     if isinstance(value, torch.Tensor):
26         return value.to(device, non_blocking=non_blocking)
27 
28     if isinstance(value, (list, tuple)):
29         values = []
30         for val in value:
31             values.append(recursive_copy_to_device(val, non_blocking=non_blocking, device=device))
32 
33         return values if isinstance(value, list) else tuple(values)
34 
35     if isinstance(value, container_abcs.Mapping):
36         device_val: Dict[str, Any] = {}
37         for key, val in value.items():
38             device_val[key] = recursive_copy_to_device(val, non_blocking=non_blocking, device=device)
39 
40         return device_val
41 
42     return value
43 
44 
45 def broadcast_object(
46     obj: Any, src_rank: int, group: object = dist.group.WORLD, dist_device: torch.device = torch.device("cpu")
47 ) -> Any:
48     """
49     Either broadcast from master to the fleet (default),
50     or use the src setting as the original rank.
51     """
52 
53     if dist.get_rank() == src_rank:
54         # Emit data
55         buffer = io.BytesIO()
56         torch.save(obj, buffer)  # type: ignore
57         data = bytearray(buffer.getbuffer())
58         length_tensor = torch.LongTensor([len(data)]).to(dist_device)
59         data_send_tensor = torch.ByteTensor(data).to(dist_device)
60         dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)
61         dist.broadcast(data_send_tensor, src=src_rank, group=group, async_op=False)
62     else:
63         # Fetch from the source
64         length_tensor = torch.LongTensor([0]).to(dist_device)
65         dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)
66         data_recv_tensor = torch.empty([int(length_tensor.item())], dtype=torch.uint8, device=dist_device)
67         dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)
68         buffer = io.BytesIO(data_recv_tensor.cpu().numpy())
69         obj = torch.load(buffer, map_location=dist_device)  # type: ignore
70     return obj
71
```

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 import logging
8 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type
9 
10 import torch
11 import torch.distributed as dist
12 from torch.optim import SGD, Optimizer
13 
14 from .utils import broadcast_object, recursive_copy_to_device
15 
16 if TYPE_CHECKING:
17     from torch.optim.optimizer import _params_t
18 else:
19     _params_t = Any
20 
21 
22 class OSS(Optimizer):
23     """Wraps an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>`
24     optimizer and shards its state as described by ZeRO_.
25     ::
26         opt = OSS(params, optim=torch.optim.Adam, lr=0.01)
27 
28     .. _ZeRO: https://arxiv.org/abs/1910.02054
29 
30     We use a greedy algorithm to pack a number of parameters
31     at each rank. Each parameter belongs to a single rank and
32     is not divided among rank.
33 
34     After each rank completed their parameter update, they broadcast
35     the new version of the parameters to all other ranks to synchronize
36     the parameters for next round forward/backward computation.
37 
38     Args:
39         params (list of tensors):
40             parameters to be optimized
41     Keyword Args:
42         optim (torch.nn.Optimizer):
43             optimizer to shard (default: SGD)
44         group (group):
45             torch.distributed group (default: group.WORLD)
46     """
47 
48     optim: Optimizer
49     in_super_constructor: bool
50 
51     def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any = dist.group.WORLD, **defaults: Any):
52         self.in_super_constructor = True
53         super().__init__(params, defaults)
54         self.in_super_constructor = False
55 
56         self.group = group
57         self.rank = dist.get_rank(group)
58         param_groups = self.partition_parameters()
59         self.optim = optim(param_groups[self.rank], **defaults)
60 
61         # Optional consolidated optimizer state
62         self._all_states: List[Dict[str, Any]] = []
63 
64         # Current device is set by the parameters allocated to this rank
65         self._device = self.partition_parameters()[self.rank][0]["params"][0].device
66 
67     def partition_parameters(self) -> List[List[dict]]:
68         """Partitions parameters across distributed ranks.
69 
70         Returns a list of param_groups (which is a list of dict) where each
71         element of the list contains the param_groups for a rank. Element 0
72         corresponds to rank 0, etc. We need all the ranks for the broadcast
73         inside step().
74         """
75         world_size = dist.get_world_size(self.group)
76         param_groups: List[List] = [list() for _ in range(world_size)]
77         sizes = [0] * world_size
78         for param_group in self.param_groups:
79             param_lists: List[List] = [list() for _ in range(world_size)]
80             for param in param_group["params"]:
81                 # Add this param to rank with smallest size.
82                 rank = sizes.index(min(sizes))
83                 param_lists[rank].append(param)
84                 sizes[rank] += param.numel()
85             for rank, params in enumerate(param_lists):
86                 if len(params) > 0:
87                     param_group_rank = copy.copy(param_group)
88                     param_group_rank["params"] = params
89                     param_groups[rank].append(param_group_rank)
90         return param_groups
91 
92     def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
93         loss = self.optim.step(closure=closure)
94         for rank, param_groups in enumerate(self.partition_parameters()):
95             for param_group in param_groups:
96                 for param in param_group["params"]:
97                     dist.broadcast(param, rank, group=self.group)
98         return loss
99 
100     def local_state_dict(self) -> dict:
101         """ Gets this rank's state_dict. """
102         return self.optim.state_dict()
103 
104     def consolidate_state_dict(self, recipient_rank: int = 0) -> None:
105         """ Update the consolidated state_dict list, one per rank.
106 
107         This needs to be called on all replicas """
108 
109         if self.rank == recipient_rank:
110             # Pull the sharded state from all the other replicas
111             # Store all the states in order, rank by rank
112             logging.debug("Pulling the sharded SGD state from all replicas")
113             self._all_states = self._collect_sharded_states()
114         else:
115             # Acknowledge broadcasts, and send this rank's shard when needed
116             self._broadcast_state_dict()
117 
118     def state_dict(self) -> Dict[str, Any]:
119         """
120         Return the last known global optimizer state, which consist of a list of the shards.
121 
122         NOTE: This is limited to the replica which was responsible for the consolidation.
123         The state may also not be up to date, depending on when `consolidate_state_dict` was last called.
124         """
125 
126         assert (
127             len(self._all_states) > 0
128         ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
129 
130         return {"state": self._all_states, "param_groups": self.param_groups}
131 
132     def load_local_state_dict(self, state_dict: dict) -> None:
133         """ Loads this rank's state_dict. """
134 
135         # Make sure that the state is on the appropriate device
136         state_dict_ondevice = recursive_copy_to_device(state_dict, non_blocking=False, device=self._device)
137 
138         self.optim.load_state_dict(state_dict_ondevice)
139 
140     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
141         """ Loads this rank's optimizer state_dict, given the global optimizer state. """
142         # Dispatch this rank's state dictionary to the local load
143         self.load_local_state_dict(state_dict["state"][self.rank])
144 
145     def add_param_group(self, param_group: dict) -> None:
146         super().add_param_group(param_group)
147         if not self.in_super_constructor:
148             param_groups = self.partition_parameters()[self.rank]
149             if len(param_groups) == len(self.optim.param_groups) + 1:
150                 self.optim.add_param_group(param_groups[-1])
151 
152     def _collect_sharded_states(self) -> List[Dict[str, Any]]:
153         """
154         Collect all the state shards, in CPU memory.
155         """
156         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
157         all_states: List[Dict[str, Any]] = []
158 
159         for rank in range(dist.get_world_size(group=self.group)):
160             if rank == self.rank:
161                 logging.debug("Saving self state")
162                 all_states.append(
163                     recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))
164                 )
165 
166                 # Sync with other replicas
167                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
168             else:
169                 # Fetch the optim state from the other replicas
170                 logging.debug("Receiving state from rank %s ", rank)
171                 replica_state = broadcast_object(
172                     empty_buffer, src_rank=rank, group=self.group, dist_device=self._device
173                 )
174 
175                 all_states.append(
176                     recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))
177                 )
178 
179                 logging.debug("State from rank %s received", rank)
180 
181         return all_states
182 
183     def _broadcast_state_dict(self) -> None:
184         """
185         Broadcast this rank's state shard, discard others
186         """
187         empty_buffer = torch.tensor([0], dtype=torch.uint8, device=self._device)
188 
189         for rank in range(dist.get_world_size(group=self.group)):
190             if rank == self.rank:
191                 # Send the state to the reference replica
192                 logging.debug(
193                     "Sending the sharded SGD state to the reference replica from rank %s", rank,
194                 )
195                 broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
196             else:
197                 # Discard this tensor/rank, broadcast necessary for syncing
198                 logging.debug("Discarding broadcast from rank %s", rank)
199                 broadcast_object(empty_buffer, src_rank=rank, group=self.group, dist_device=self._device)
200
```
--- 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
@@ -13,7 +13,7 @@
 
 from .utils import broadcast_object, recursive_copy_to_device
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from torch.optim.optimizer import _params_t
 else:
     _params_t = Any
@@ -49,10 +49,12 @@
     in_super_constructor: bool
 
     def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any = dist.group.WORLD, **defaults: Any):
+        # Hold all the nmodel params in the root .param_groups
         self.in_super_constructor = True
         super().__init__(params, defaults)
         self.in_super_constructor = False
 
+        # Build the wrapped optimizer, responsible for a shard of the params
         self.group = group
         self.rank = dist.get_rank(group)
         param_groups = self.partition_parameters()
@@ -90,11 +92,14 @@
         return param_groups
 
     def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
+        # Run the optimizer step on this shard only
         loss = self.optim.step(closure=closure)
+
+        # Sync all the states
         for rank, param_groups in enumerate(self.partition_parameters()):
             for param_group in param_groups:
                 for param in param_group["params"]:
-                    dist.broadcast(param, rank, group=self.group)
+                    dist.broadcast(tensor=param, src=rank, group=self.group)
         return loss
 
     def local_state_dict(self) -> dict:
@@ -138,10 +143,13 @@
         self.optim.load_state_dict(state_dict_ondevice)
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        """ Loads this rank's optimizer state_dict, given the global optimizer state. """
-        # Dispatch this rank's state dictionary to the local load
+        """ Restore the global parameter groups as well as the shard """
+        # Dispatch this rank's state dictionary to the wrapped shard optimizer
         self.load_local_state_dict(state_dict["state"][self.rank])
 
+        # Restore the global param_groups
+        self.param_groups = state_dict["param_groups"]
+
     def add_param_group(self, param_group: dict) -> None:
         super().add_param_group(param_group)
         if not self.in_super_constructor:
diff --git a/fairscale/optim/utils.py b/fairscale/optim/utils.py
--- a/fairscale/optim/utils.py
+++ b/fairscale/optim/utils.py
@@ -53,7 +53,7 @@
     if dist.get_rank() == src_rank:
         # Emit data
         buffer = io.BytesIO()
-        torch.save(obj, buffer)  # type: ignore
+        torch.save(obj, buffer)
         data = bytearray(buffer.getbuffer())
         length_tensor = torch.LongTensor([len(data)]).to(dist_device)
         data_send_tensor = torch.ByteTensor(data).to(dist_device)
@@ -66,5 +66,5 @@
         data_recv_tensor = torch.empty([int(length_tensor.item())], dtype=torch.uint8, device=dist_device)
         dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)
         buffer = io.BytesIO(data_recv_tensor.cpu().numpy())
-        obj = torch.load(buffer, map_location=dist_device)  # type: ignore
+        obj = torch.load(buffer, map_location=dist_device)
     return obj

{"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@@ -13,7 +13,7 @@\n \n from .utils import broadcast_object, recursive_copy_to_device\n \n-if TYPE_CHECKING:\n+if TYPE_CHECKING:  # pragma: no cover\n     from torch.optim.optimizer import _params_t\n else:\n     _params_t = Any\n@@ -49,10 +49,12 @@\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 nmodel 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         param_groups = self.partition_parameters()\n@@ -90,11 +92,14 @@\n         return param_groups\n \n     def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:\n+        # Run the optimizer step on this shard only\n         loss = self.optim.step(closure=closure)\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(param, rank, group=self.group)\n+                    dist.broadcast(tensor=param, src=rank, group=self.group)\n         return loss\n \n     def local_state_dict(self) -> dict:\n@@ -138,10 +143,13 @@\n         self.optim.load_state_dict(state_dict_ondevice)\n \n     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:\n-        \"\"\" Loads this rank's optimizer state_dict, given the global optimizer state. \"\"\"\n-        # Dispatch this rank's state dictionary to the local load\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(state_dict[\"state\"][self.rank])\n \n+        # Restore the global param_groups\n+        self.param_groups = state_dict[\"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:\ndiff --git a/fairscale/optim/utils.py b/fairscale/optim/utils.py\n--- a/fairscale/optim/utils.py\n+++ b/fairscale/optim/utils.py\n@@ -53,7 +53,7 @@\n     if dist.get_rank() == src_rank:\n         # Emit data\n         buffer = io.BytesIO()\n-        torch.save(obj, buffer)  # type: ignore\n+        torch.save(obj, buffer)\n         data = bytearray(buffer.getbuffer())\n         length_tensor = torch.LongTensor([len(data)]).to(dist_device)\n         data_send_tensor = torch.ByteTensor(data).to(dist_device)\n@@ -66,5 +66,5 @@\n         data_recv_tensor = torch.empty([int(length_tensor.item())], dtype=torch.uint8, device=dist_device)\n         dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)\n         buffer = io.BytesIO(data_recv_tensor.cpu().numpy())\n-        obj = torch.load(buffer, map_location=dist_device)  # type: ignore\n+        obj = torch.load(buffer, map_location=dist_device)\n     return obj\n", "issue": "OSS to expose classic optimizer compliant param_groups\n## \ud83d\ude80 Feature\r\nExpose the \"classic\" param_groups from the OSS optimizer wrapper, either via a new API (not as nice) or if possible via the classic .param_groups attribute (get and set) \r\n\r\n## Motivation\r\nSome frameworks like classy vision use that to alter the LR per layer or over time (or both). Exposing this would make OSS integration a lot easier\r\n\r\n## Pitch\r\nEither use dedicated getter and setter properties which change behaviour depending on whether the call comes from the outside or the super constructor, or add a specific API to get the consolidated param_groups and set the consolidated param_groups\r\n\r\n## Alternatives\r\n- Destroy the optimizer every time there's an adjustment\r\n- Only expose adjustments which apply to the whole param_groups (\"set all LR to a given value\")\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 io\nfrom typing import Any, Dict\n\nimport torch\nfrom torch._six import container_abcs\nimport torch.distributed as dist\n\n\n# Credits:  classy_vision/generic/distributed_util.py\ndef recursive_copy_to_device(value: Any, non_blocking: bool, device: torch.device) -> Any:\n    \"\"\"\n    Recursively searches lists, tuples, dicts and copies tensors to device if\n    possible. Non-tensor values are passed as-is in the result.\n\n    NOTE:  These are all copies, so if there are two objects that reference\n    the same object, then after this call, there will be two different objects\n    referenced on the device.\n    \"\"\"\n\n    if isinstance(value, torch.Tensor):\n        return value.to(device, non_blocking=non_blocking)\n\n    if isinstance(value, (list, tuple)):\n        values = []\n        for val in value:\n            values.append(recursive_copy_to_device(val, non_blocking=non_blocking, device=device))\n\n        return values if isinstance(value, list) else tuple(values)\n\n    if isinstance(value, container_abcs.Mapping):\n        device_val: Dict[str, Any] = {}\n        for key, val in value.items():\n            device_val[key] = recursive_copy_to_device(val, non_blocking=non_blocking, device=device)\n\n        return device_val\n\n    return value\n\n\ndef broadcast_object(\n    obj: Any, src_rank: int, group: object = dist.group.WORLD, dist_device: torch.device = torch.device(\"cpu\")\n) -> Any:\n    \"\"\"\n    Either broadcast from master to the fleet (default),\n    or use the src setting as the original rank.\n    \"\"\"\n\n    if dist.get_rank() == src_rank:\n        # Emit data\n        buffer = io.BytesIO()\n        torch.save(obj, buffer)  # type: ignore\n        data = bytearray(buffer.getbuffer())\n        length_tensor = torch.LongTensor([len(data)]).to(dist_device)\n        data_send_tensor = torch.ByteTensor(data).to(dist_device)\n        dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)\n        dist.broadcast(data_send_tensor, src=src_rank, group=group, async_op=False)\n    else:\n        # Fetch from the source\n        length_tensor = torch.LongTensor([0]).to(dist_device)\n        dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)\n        data_recv_tensor = torch.empty([int(length_tensor.item())], dtype=torch.uint8, device=dist_device)\n        dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)\n        buffer = io.BytesIO(data_recv_tensor.cpu().numpy())\n        obj = torch.load(buffer, map_location=dist_device)  # type: ignore\n    return obj\n", "path": "fairscale/optim/utils.py"}, {"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\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:\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        self.in_super_constructor = True\n        super().__init__(params, defaults)\n        self.in_super_constructor = False\n\n        self.group = group\n        self.rank = dist.get_rank(group)\n        param_groups = self.partition_parameters()\n        self.optim = optim(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 = self.partition_parameters()[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                if len(params) > 0:\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    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:\n        loss = self.optim.step(closure=closure)\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(param, 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        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 SGD 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 {\"state\": self._all_states, \"param_groups\": self.param_groups}\n\n    def load_local_state_dict(self, state_dict: dict) -> None:\n        \"\"\" Loads this rank's state_dict. \"\"\"\n\n        # Make sure that the state is on the appropriate device\n        state_dict_ondevice = recursive_copy_to_device(state_dict, non_blocking=False, device=self._device)\n\n        self.optim.load_state_dict(state_dict_ondevice)\n\n    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:\n        \"\"\" Loads this rank's optimizer state_dict, given the global optimizer state. \"\"\"\n        # Dispatch this rank's state dictionary to the local load\n        self.load_local_state_dict(state_dict[\"state\"][self.rank])\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 _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 SGD 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 io\nfrom typing import Any, Dict\n\nimport torch\nfrom torch._six import container_abcs\nimport torch.distributed as dist\n\n\n# Credits:  classy_vision/generic/distributed_util.py\ndef recursive_copy_to_device(value: Any, non_blocking: bool, device: torch.device) -> Any:\n    \"\"\"\n    Recursively searches lists, tuples, dicts and copies tensors to device if\n    possible. Non-tensor values are passed as-is in the result.\n\n    NOTE:  These are all copies, so if there are two objects that reference\n    the same object, then after this call, there will be two different objects\n    referenced on the device.\n    \"\"\"\n\n    if isinstance(value, torch.Tensor):\n        return value.to(device, non_blocking=non_blocking)\n\n    if isinstance(value, (list, tuple)):\n        values = []\n        for val in value:\n            values.append(recursive_copy_to_device(val, non_blocking=non_blocking, device=device))\n\n        return values if isinstance(value, list) else tuple(values)\n\n    if isinstance(value, container_abcs.Mapping):\n        device_val: Dict[str, Any] = {}\n        for key, val in value.items():\n            device_val[key] = recursive_copy_to_device(val, non_blocking=non_blocking, device=device)\n\n        return device_val\n\n    return value\n\n\ndef broadcast_object(\n    obj: Any, src_rank: int, group: object = dist.group.WORLD, dist_device: torch.device = torch.device(\"cpu\")\n) -> Any:\n    \"\"\"\n    Either broadcast from master to the fleet (default),\n    or use the src setting as the original rank.\n    \"\"\"\n\n    if dist.get_rank() == src_rank:\n        # Emit data\n        buffer = io.BytesIO()\n        torch.save(obj, buffer)\n        data = bytearray(buffer.getbuffer())\n        length_tensor = torch.LongTensor([len(data)]).to(dist_device)\n        data_send_tensor = torch.ByteTensor(data).to(dist_device)\n        dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)\n        dist.broadcast(data_send_tensor, src=src_rank, group=group, async_op=False)\n    else:\n        # Fetch from the source\n        length_tensor = torch.LongTensor([0]).to(dist_device)\n        dist.broadcast(length_tensor, src=src_rank, group=group, async_op=False)\n        data_recv_tensor = torch.empty([int(length_tensor.item())], dtype=torch.uint8, device=dist_device)\n        dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False)\n        buffer = io.BytesIO(data_recv_tensor.cpu().numpy())\n        obj = torch.load(buffer, map_location=dist_device)\n    return obj\n", "path": "fairscale/optim/utils.py"}, {"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\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 nmodel 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        param_groups = self.partition_parameters()\n        self.optim = optim(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 = self.partition_parameters()[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                if len(params) > 0:\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    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:\n        # Run the optimizer step on this shard only\n        loss = self.optim.step(closure=closure)\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        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 SGD 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 {\"state\": self._all_states, \"param_groups\": self.param_groups}\n\n    def load_local_state_dict(self, state_dict: dict) -> None:\n        \"\"\" Loads this rank's state_dict. \"\"\"\n\n        # Make sure that the state is on the appropriate device\n        state_dict_ondevice = recursive_copy_to_device(state_dict, non_blocking=False, device=self._device)\n\n        self.optim.load_state_dict(state_dict_ondevice)\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(state_dict[\"state\"][self.rank])\n\n        # Restore the global param_groups\n        self.param_groups = state_dict[\"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 _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 SGD 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"}]}