Avoid calling _free_fp16_param_shard() too early

Author: jiecaoyu

fairscale/nn/data_parallel/fully_sharded_data_parallel.py

@@ -1706,7 +1706,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
             # bandwidth but uses more GPU memory.
             self._free_full_params([param])
 
-        if self.mixed_precision:
+        if self.mixed_precision and (self._require_backward_grad_sync or self.reshard_after_forward):
             # This is a no-op if reshard_after_forward is True, since we already
             # free the param shard when rebuilding the full params in the
             # pre_backward_hook.