Update fairscale version (#11567)

Co-authored-by: Aki Nitta <nitta@akihironitta.com>
Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>
Co-authored-by: Jirka <jirka.borovec@seznam.cz>
Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>

.azure-pipelines/gpu-tests.yml

@@ -52,7 +52,7 @@ jobs:
 
     - bash: |
         python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if 'horovod' not in line] ; open(fname, 'w').writelines(lines)"
-        pip install fairscale==0.4.0
+        pip install fairscale==0.4.5
         pip install deepspeed==0.5.7
         pip install bagua-cuda102==0.9.0
         pip install . --requirement requirements/devel.txt

dockers/base-cuda/Dockerfile

@@ -133,7 +133,7 @@ RUN \
 
 RUN \
     # install FairScale
-    pip install fairscale==0.4.0 && \
+    pip install fairscale==0.4.5 && \
     python -c "import fairscale; print(fairscale.__version__)"
 
 RUN \

tests/strategies/test_ddp_fully_sharded_with_full_state_dict.py

@@ -83,8 +83,7 @@ class TestFSDPModel(BoringModel):
         assert isinstance(self.layer, FullyShardedDataParallel)
         assert isinstance(self.layer.module[0], FullyShardedDataParallel)
         assert isinstance(self.layer.module[2], FullyShardedDataParallel)
-        # root should not be resharding
-        assert self.layer.reshard_after_forward is False
+
         # Assert that the nested layers are set reshard_after_forward to True
         assert self.layer.module[0].reshard_after_forward is True
         assert self.layer.module[2].reshard_after_forward is True