[perf] [FSDP] micro-optimization for memory usage (#533)

fcbf1ea3 · Myle Ott · GitHub · 7b127ccb · fcbf1ea3
Unverified Commit fcbf1ea3 authored Mar 18, 2021 by Myle Ott Committed by GitHub Mar 18, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

fairscale/nn/data_parallel/fully_sharded_data_parallel.py fairscale/nn/data_parallel/fully_sharded_data_parallel.py +2 -0

No files found.
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -815,6 +815,8 @@ class FullyShardedDataParallel(nn.Module):

        if self.reshard_after_forward:
            self._free_full_params()
+            if self.mixed_precision:
+                self._free_fp16_param_shard()

        # Switch to main FP32 param shard. We maintain this invariant throughout
        # the code, i.e., ``p.data == p._fp32_shard`` after each function. This