FSDP: Added fix + Regression test for incorrect params_group mapping in optimizer state sharding

fairscale/nn/data_parallel/fsdp_optim_utils.py

@@ -52,13 +52,15 @@ def flatten_optim_state_dict(sd: Dict) -> Dict:
 
     # Now make a new param_groups copy and update it.
     new_sd_pg = copy.deepcopy(sd["param_groups"])
+    local_param_ids = sorted(
+            {lid for lid in param_id_map.values() if lid is not None}
+        )
     # add pointers from the `params` dict.
     for pg_id, _ in enumerate(sd["param_groups"]):
         # The values() list may look like [0,0,None,None,2,2]. We use
         # groupby to remove the duplicates and then count the length of
         # resulting iter.
-        num_local_params = sum(1 for _ in groupby(param_id_map.values()))
-        new_sd_pg[pg_id]["params"] = list(range(num_local_params))
+        new_sd_pg[pg_id]["params"] = local_param_ids
 
     # update the original sd so that we don't lose extra keys, like loss_scale.
     sd["state"] = new_state

tests/nn/data_parallel/test_fsdp_optim_utils.py

@@ -0,0 +1,23 @@
+import pytest
+
+from fairscale.nn.data_parallel.fsdp_optim_utils import flatten_optim_state_dict
+
+def test_flatten_optim_state_dict_param_group_ids_are_unique():
+    sd = {
+        "state": {},
+        "param_groups": [{"params": [0, 1]}],
+        "param_id_map": {
+            0: 0,
+            1: 0,
+            2: None,
+            3: None,
+            4: 2,
+            5: 2,
+        },
+        "uncollected_local_ids": [],
+    }
+
+    out = flatten_optim_state_dict(sd)
+    params = out["param_groups"][0]["params"]
+
+    assert params == [0, 2]