Skip to content

GitLab

Unverified Commit fcbf1ea3 authored by Myle Ott's avatar Myle Ott Committed by GitHub
Browse files

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

parent 7b127ccb
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
fairscale/nn/data_parallel/fully_sharded_data_parallel.py
View file @ fcbf1ea3
...@@ -815,6 +815,8 @@ class FullyShardedDataParallel(nn.Module): ...@@ -815,6 +815,8 @@ class FullyShardedDataParallel(nn.Module):
if self.reshard_after_forward: if self.reshard_after_forward:
self._free_full_params() self._free_full_params()
if self.mixed_precision:
self._free_fp16_param_shard()
# Switch to main FP32 param shard. We maintain this invariant throughout # Switch to main FP32 param shard. We maintain this invariant throughout
# the code, i.e., ``p.data == p._fp32_shard`` after each function. This # the code, i.e., ``p.data == p._fp32_shard`` after each function. This
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment