Fairscale FSDP fix model save
# What does this PR do?

This PR fixes the fact training with fairscale fully-sharded wrapper was hanging: it looks like recent changes in fairscale make a synchronization during the model state dict call, which results in the training hanging if we don't call that state_dict method on all processes. This PR addresses that.