Repository: facebookresearch/fairscale

Issue:
[feat] OSS: make it compatible with nvidia's AMP
## 🚀 Feature
AMP is using a load(get()) trick to make sure that the state tensors are cast in the right format, which breaks OSS assumptions, see https://github.com/NVIDIA/apex/blob/2ec84ebdca59278eaf15e8ddf32476d9d6d8b904/apex/amp/_initialize.py#L205

## Motivation
nvidia's AMP is getting deprecated in favour of Pytorch, but some of its features are not yet there. It brings a lot of speed to the table which is necessary for big jobs

## Pitch
Make it happen (tm)

## Alternatives
Do not support AMP

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:
cc @mannatsingh 

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -141,13 +141,18 @@ def state_dict(self) -> Dict[str, Any]:
         """
         Return the last known global optimizer state, which consist of a list of the shards.
 
-        NOTE: This is limited to the replica which was responsible for the consolidation.
+        NOTE:
+        - If the state has not been consolidated, this returns a shard's worth, not the global state.
+        - Returning the global state is limited to the replica which was responsible for the consolidation.
         The state may also not be up to date, depending on when `consolidate_state_dict` was last called.
         """
 
-        assert (
-            len(self._all_states) > 0
-        ), "The optimizer state is not materialized, please call consolidate_state_dict on every replica beforehand"
+        if len(self._all_states) == 0:
+            logging.warning("Optimizer state has not been consolidated. Returning the local state")
+            logging.warning("Please call `consolidate_state_dict()` beforehand if you meant to save the global state")
+            state_dict = self.local_state_dict()
+            state_dict["local_state_dict"] = True
+            return state_dict
 
         # Flatten the param_groups, save the partition which logs the rank <> shard correspondence
         partition: List[Tuple[int, int]] = []
@@ -164,6 +169,7 @@ def state_dict(self) -> Dict[str, Any]:
             "state": [s["state"] for s in self._all_states],
             "param_groups": param_groups,
             "partition": partition,
+            "local_state_dict": False,
         }
 
     def load_local_state_dict(self, state_dict: dict) -> None:
@@ -193,12 +199,16 @@ 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 """
 
-        # Get this optimizer's param_groups shard
-        param_groups = state_dict["param_groups"][
-            state_dict["partition"][self.rank][0] : state_dict["partition"][self.rank][1]
-        ]
-        # Dispatch this rank's state dictionary to the wrapped shard optimizer
-        self.load_local_state_dict({"state": state_dict["state"][self.rank], "param_groups": param_groups})
+        # Check whether we got a local or global dict
+        if state_dict["local_state_dict"]:
+            self.load_local_state_dict(state_dict)
+        else:
+            # Get this optimizer's param_groups shard
+            param_groups = state_dict["param_groups"][
+                state_dict["partition"][self.rank][0] : state_dict["partition"][self.rank][1]
+            ]
+            # Dispatch this rank's state dictionary to the wrapped shard optimizer
+            self.load_local_state_dict({"state": state_dict["state"][self.rank], "param_groups": param_groups})
 
     def add_param_group(self, param_group: dict) -> None:
         super().add_param_group(param_group)


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
@@ -133,6 +133,20 @@ def test_local_state_dict():
     assert x == torch.tensor([0.9], device=DEVICE)
 
 
+def test_implicit_local_state_dict():
+    x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
+    o = optim.OSS([x], lr=0.1)
+    local_state_dict = o.state_dict()
+    o = optim.OSS([x], lr=0.01)
+    o.load_state_dict(local_state_dict)
+    # We should now be using a lr of 0.1.
+    assert o.optim.param_groups[0]["lr"] == 0.1
+    assert o.param_groups[0]["lr"] == 0.1
+    x.backward()
+    o.step()
+    assert x == torch.tensor([0.9], device=DEVICE)
+
+
 def run_test_add_param_group(rank, world_size):
     dist_init(rank, world_size)
     params = []