Repository: facebookresearch/fairscale

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")



Discussion:
cc @mannatsingh @vreis , FYI, I'm trying to think of an elegant way to handle that. The crux of the problem right now with Classy integration is that ClassyOptimizer is a wrapper, OSS also is a wrapper, and the LR schedulers expect to interact with 'self.optimizer.param_groups'. This is unambiguous for a "normal" optimizer (there's a single state in flight at any time), but it becomes ambiguous when the state is sharded because it could be the replica's param_groups or the whole optimizer param_groups (which is virtual at this point). Classy assumes the latter, OSS silently exposes the former. 
Yeah, that's a good point. At the same time, I think we can let Classy think it's overriding the whole params when in fact it would just be overriding the sharded params. I don't think Classy actually *needs* the params to be the complete set. At the end of the day, what Classy is doing is taking some param groups and splitting them into subsets to configure different optimizer settings. But I am not sure how involved this is on the OSS side.
@mannatsingh there are two possible practical issues I think, at least that's what I'm seeing in some of the unit test suite:
- Classy pulls a rank state and alters some keys, only to find that at times they are not there (->`KeyError`), if `param_groups` only exposes a shard
- Classy pulls the whole state and "pushes" the new one, but only the reference replica is affected, if `param_groups` exposes the whole virtual state but does not sync back
> Classy pulls a rank state and alters some keys, only to find that at times they are not there (->KeyError), if param_groups only exposes a shard

Hmm - can the shards actually change? Otherwise I'm not sure why the keys can go missing?

> Classy pulls the whole state and "pushes" the new one, but only the reference replica is affected, if param_groups exposes the whole virtual state but does not sync back

Right now Classy should ideally push the new state on all replicas, but yeah, this can be tricky. On Classy's side if we show Classy the whole virtual state, Classy should act on it the same way on all replicas. And then the work is offloaded to the virtualization logic to actually deal with this properly.
> > Classy pulls a rank state and alters some keys, only to find that at times they are not there (->KeyError), if param_groups only exposes a shard
> 
> Hmm - can the shards actually change? Otherwise I'm not sure why the keys can go missing?

I mean that you can have something [like this ](https://pytorch.org/docs/stable/optim.html#torch.optim.lr_scheduler.LambdaLR), where you have one scheduler per param_group for instance. Now if everything is done from the reference rank, then it does not have all the groups in its own `.param_groups` attribute
Actually, so that it's easier for me to understand, do you mind writing a simple example on how parameter groups are set in OSS right now across different ranks? 
Right now, it looks like (pseudo code)

```
zero = OSS(model_params)

# through the torch.optimizer default constructor,
#  so this has nothing specific to the wrapped optimizer, no LR for instance
#  *same value for all ranks*
zero.param_groups : model_params

# this is probably what classy is looking for, but *rank specific* (subset of the parameters)
zero.optim.param_groups :  rank_and_optim_specific_param_groups
```

Note that classy_optimizer adds another level of redirection, but it does not change the logic
Thanks @blefaudeux ! What I was actually asking was slightly different -
```
param_groups = [{"params": trunk_params, "lr": 0.1}, {"params": head_params, "lr": 0.2}]

optimizer = SGD(param_groups)

# is the same supported for OSS?
zero = OSS(param_groups)
```
> Thanks @blefaudeux ! What I was actually asking was slightly different -
> 
> ```
> param_groups = [{"params": trunk_params, "lr": 0.1}, {"params": head_params, "lr": 0.2}]
> 
> optimizer = SGD(param_groups)
> 
> # is the same supported for OSS?
> zero = OSS(param_groups)
> ```

yes, the same is supported by OSS !
Ah, now I'm getting what the issue is. Sorry for being daft!
Just to summarize, here's what's happening -

```python
param_groups = [{"params": trunk_params, "lr": 0.1}, {"params": head_params, "lr": 0.2}]
classy_oss = ClassyOSS()
classy_oss.set_param_groups(param_groups) # this calls OSS(param_groups), which works

classy_oss.step() # this needs self.optimizer.param_groups to update the dicts, but that isn't available  
```
> Ah, now I'm getting what the issue is. Sorry for being daft!
> Just to summarize, here's what's happening -
> 
> ```python
> param_groups = [{"params": trunk_params, "lr": 0.1}, {"params": head_params, "lr": 0.2}]
> classy_oss = ClassyOSS()
> classy_oss.set_param_groups(param_groups) # this calls OSS(param_groups), which works
> 
> classy_oss.step() # this needs self.optimizer.param_groups to update the dicts, but that isn't available  
> ```

my fault, I was not clear enough.. It's basically a case of russian dolls exposing similar but not identical attributes, and outside accessors who are misled. As shown in a diff this can be "fixed" on the client side, but it's not very elegant

Resolving patch:
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 @@ class OSS(Optimizer):
     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 @@ def partition_parameters(self) -> List[List[dict]]:
         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 @@ def load_local_state_dict(self, state_dict: dict) -> None:
         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 @@ def broadcast_object(
     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 @@ def broadcast_object(
         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


Test patch:
diff --git a/tests/optim/test_oss.py b/tests/optim/test_oss.py
--- a/tests/optim/test_oss.py
+++ b/tests/optim/test_oss.py
@@ -44,9 +44,20 @@ def test_state_dict():
     o = optim.OSS([x], lr=0.1)
     o.consolidate_state_dict()  # Sync state dict in between replicas - even if there are none
     state_dict = o.state_dict()
+
+    # Check that the pulled state is what we expect
+    assert state_dict["param_groups"][0]["lr"] == 0.1
+
+    # Check that the pulled state and the .param_groups attribute are in sync
+    assert state_dict["param_groups"][0]["lr"] == o.param_groups[0]["lr"]
+
+    # Check that it's correctly loaded
     o = optim.OSS([x], lr=0.01)
     o.load_state_dict(state_dict)
-    # We should now be using a lr of 0.1.
+
+    # We should now be using a lr of 0.1, both within the optimizer
+    # and as exposed by the .param_groups attribute
+    assert o.param_groups[0]["lr"] == 0.1
     x.backward()
     o.step()
     assert x == torch.tensor([0.9], device=DEVICE)