Repository: facebookresearch/fairscale

Issue:
Rewrite state_dict in a more pytorch idiomatic way
## 🚀 Feature
Change the param_groups handling in the state dict, in order to follow more closely the default PyTorch assumptions
https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer.state_dict

## Motivation
- Some users may assume that the default pytorch optimizer interface with respect to the state dict is respected by fairscale/oss
- People familiar with pytorch optimizers would have an easier learning curve when peeking into OSS

## Pitch
Rewrite the exposed state dict in order to return "state" and "param_groups" in accordance to pytorch expectations, without duplications

## Alternatives
- rely on the python/pytorch memory model to remove duplicates in memory and while serializing
- add wrappers on the user side

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
cc @msbaines @mannatsingh 

Resolving patch:
diff --git a/benchmarks/oss.py b/benchmarks/oss.py
--- a/benchmarks/oss.py
+++ b/benchmarks/oss.py
@@ -50,7 +50,7 @@ def collate(inputs: List[Any]):
             "label": torch.stack([i[1] for i in inputs]).to(rank),
         }
 
-    def _print(msg):
+    def print_(msg):
         if dist.get_rank() == 0:
             print(msg)
 
@@ -91,7 +91,7 @@ def closure():
 
         epoch_end = time.monotonic()
         measurements.append(data_size / (epoch_end - epoch_start))
-        _print(f"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec")
+        print_(f"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec")
 
     torch.cuda.synchronize(rank)
     training_stop = time.monotonic()
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -58,14 +58,14 @@ def __init__(self, params: _params_t, optim: Type[Optimizer] = SGD, group: Any =
         # 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()
-        self.optim = optim(param_groups[self.rank], **defaults)
+        split_param_groups = self.partition_parameters()
+        self.optim = optim(split_param_groups[self.rank], **defaults)
 
         # Optional consolidated optimizer state
         self._all_states: List[Dict[str, Any]] = []
 
         # Current device is set by the parameters allocated to this rank
-        self._device = self.partition_parameters()[self.rank][0]["params"][0].device
+        self._device = split_param_groups[self.rank][0]["params"][0].device
 
     def partition_parameters(self) -> List[List[dict]]:
         """Partitions parameters across distributed ranks.
@@ -122,7 +122,7 @@ def consolidate_state_dict(self, recipient_rank: int = 0) -> None:
         if self.rank == recipient_rank:
             # Pull the sharded state from all the other replicas
             # Store all the states in order, rank by rank
-            logging.debug("Pulling the sharded SGD state from all replicas")
+            logging.debug("Pulling the sharded optimizer state from all replicas")
             self._all_states = self._collect_sharded_states()
         else:
             # Acknowledge broadcasts, and send this rank's shard when needed
@@ -140,7 +140,10 @@ def state_dict(self) -> Dict[str, Any]:
             len(self._all_states) > 0
         ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
 
-        return {"state": self._all_states}
+        return {
+            "state": [s["state"] for s in self._all_states],
+            "param_groups": [s["param_groups"] for s in self._all_states],
+        }
 
     def load_local_state_dict(self, state_dict: dict) -> None:
         """ Loads this rank's state_dict. """
@@ -169,7 +172,12 @@ def load_local_state_dict(self, state_dict: dict) -> None:
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         """ 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])
+        self.load_local_state_dict(
+            {"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)
@@ -224,7 +232,7 @@ def _broadcast_state_dict(self) -> None:
             if rank == self.rank:
                 # Send the state to the reference replica
                 logging.debug(
-                    "Sending the sharded SGD state to the reference replica from rank %s", rank,
+                    "Sending the sharded optimizer state to the reference replica from rank %s", rank,
                 )
                 broadcast_object(self.local_state_dict(), src_rank=rank, group=self.group, dist_device=self._device)
             else:


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
@@ -50,11 +50,15 @@ def test_state_dict():
     o.consolidate_state_dict()  # Sync state dict in between replicas - even if there are none
     state_dict = o.state_dict()
 
+    # Check that the state dict is pytorch-compliant key wise
+    assert "param_groups" in state_dict.keys()
+    assert "state" in state_dict.keys()
+
     # Check that the pulled state is what we expect
-    assert state_dict["state"][0]["param_groups"][0]["lr"] == 0.1
+    assert state_dict["param_groups"][0][0]["lr"] == 0.1
 
     # Check that the pulled state and the .param_groups attribute are in sync
-    assert state_dict["state"][0]["param_groups"][0]["lr"] == o.param_groups[0]["lr"]
+    assert state_dict["param_groups"][0][0]["lr"] == o.param_groups[0]["lr"]
 
     # Check that it's correctly loaded
     o = optim.OSS([x], lr=0.01)