Repository: facebookresearch/fairscale

Issue:
FSDP state_dict returned data type when using mixed precision
The docstring for `FSDP.state_dict()` says that returned parameters will be in FP32: https://github.com/facebookresearch/fairscale/blob/b84b9146aca8fd575622462c1b311243e40b115c/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L647-L648.

But this is not the case when `mixed_precision=True`. So what's the best way to get the FP32 parameters when using mixed precision?



Discussion:
That's interesting. It should be full precision since we use summon_full_params with this parameter here:

https://github.com/facebookresearch/fairscale/blob/b84b9146aca8fd575622462c1b311243e40b115c/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L824

but maybe we don't have a unit test for it and it got broken. Can you double check you indeed not getting full precision params?
I was able to narrow down the issue. Apparently it's just buffers in the state dict that will be in half precision.
That’s great! Please share more details of what you found if you can. It can help others in the future who run into the same issue and help us improve the code to avoid this if possible. 

Resolving patch:
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -653,17 +653,21 @@ def state_dict(self, *args: Any, **kwargs: Any) -> Any:
         if torch.cuda.is_available():
             torch.cuda.synchronize()
         self._lazy_init()
-        if self.mixed_precision:
-            # Buffers dtype stays consistent with parameters.
-            self._cast_buffers(dtype=torch.float32)
+
+        def maybe_cast_buffers(dtype: Optional[torch.dtype] = None) -> None:
+            if self.mixed_precision:
+                self._cast_buffers(dtype=dtype)
 
         if self._return_full_state_dict:
             if self.training_state != TrainingState.SUMMON_FULL_PARAMS:
                 with self.summon_full_params(recurse=False, volatile=True):
+                    maybe_cast_buffers(torch.float32)
                     state_dict = super().state_dict(*args, **kwargs)
             else:
+                maybe_cast_buffers(torch.float32)
                 state_dict = super().state_dict(*args, **kwargs)
         else:
+            maybe_cast_buffers(torch.float32)
             if self.flatten_parameters:
                 assert isinstance(self.module, FlattenParamsWrapper)
                 state_dict = self.module.flat_state_dict(*args, **kwargs)
@@ -674,9 +678,8 @@ def state_dict(self, *args: Any, **kwargs: Any) -> Any:
             for k in state_dict.keys():
                 state_dict[k] = state_dict[k].cpu()
 
-        if self.mixed_precision:
-            # In case we are in mixed precision, restore buffers back to buffer_dtype.
-            self._cast_buffers()
+        # In case we are in mixed precision, restore buffers back to buffer_dtype.
+        maybe_cast_buffers()
         return state_dict
 
     @typing.overload
@@ -860,7 +863,7 @@ def _reset_lazy_init(self) -> None:
 
     def _lazy_init(self) -> None:
         """Initialization steps that should happen lazily, typically right
-           before the first forward pass.
+        before the first forward pass.
         """
         # Initialize param attributes lazily, in case the param's dtype or
         # device changes after __init__.


Test patch:
diff --git a/tests/nn/data_parallel/test_fsdp_state_dict.py b/tests/nn/data_parallel/test_fsdp_state_dict.py
--- a/tests/nn/data_parallel/test_fsdp_state_dict.py
+++ b/tests/nn/data_parallel/test_fsdp_state_dict.py
@@ -101,8 +101,9 @@ def test_state_dict_before_forward(self, mixed_precision):
     def _test_state_dict_before_forward(cls, config, rank, group):
         ddp_model = cls.get_wrapped_model(group, cuda_first=False, config=config)
         sd = ddp_model.state_dict()
-        wt = sd["embed_tokens.weight"]
-        assert wt.dtype == torch.float32, f"got dtype {wt.dtype} expected torch.float32"
+        for param_name in ("embed_tokens.weight", "vocab_bias"):
+            wt = sd[param_name]
+            assert wt.dtype == torch.float32, f"got dtype {wt.dtype} for {param_name}, expected torch.float32"
         cls._train_for_several_steps(ddp_model, 1, ddp_model.mixed_precision)
 
     @classmethod
@@ -232,16 +233,10 @@ def _test_state_dict_device(self, config, rank, group, pure_fp16=False, **model_
                 assert v.device.type == "cuda", v.device.type
 
         expected_dtype = torch.float16 if pure_fp16 else torch.float32
-        buffers = {
-            k.replace("_fsdp_wrapped_module.", "").replace("_fpw_module.", "") for k, _ in fsdp_model.named_buffers()
-        }
         for k, v in sd.items():
             if not torch.is_floating_point(v):
                 continue
-            if k in buffers:
-                assert v.dtype == fsdp_model.buffer_dtype, f"{v.dtype} != {fsdp_model.buffer_dtype}"
-            else:
-                assert v.dtype == expected_dtype, f"{v.dtype} != {expected_dtype}"
+            assert v.dtype == expected_dtype, f"{v.dtype} != {expected_dtype}"
 
 
 @skip_if_cuda