Repository: facebookresearch/fairscale

Issue:
[refactor] enhance wrap and auto_wrap
- Two things were done in this PR
  1. We don't need to import FSDP in wrap.py since the wrapper class
     type is stored in the context now.
  2. We can use a `should_wrap` function to customize wrapping policy
     for auto_wrap, including size of module, blacklist, exclude list
- The auto_wrap function got simplified a bit as a minor side effect.

# Before submitting

- [ ] Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
- [ ] Did you read the [contributor guideline](https://github.com/facebookresearch/fairscale/blob/master/CONTRIBUTING.md)?
- [ ] Did you make sure to update the docs?
- [ ] Did you write any new necessary tests?

## What does this PR do?
Fixes # (issue).

## PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃



Discussion:
cc @SeanNaren. I will comment on the diff to make it easier to review.
> 
> 
> cc @SeanNaren. I will comment on the diff to make it easier to review.

I think this may cause some grief on your integration effort with lightning. Apologize in advance. I didn't think about those changes until today. I hope they are small enough changes and longer term better for the API in terms of being more generic (wrapper is more pure) and more customizable (more flexible should_wrap).
Sorry for the delay. I think this is ready.
Thanks for review again, Myle! Even better now. Take another look?
Great job @min-xu-ai!!
@myleott this seems like an unrelated to this change?

https://app.circleci.com/pipelines/github/facebookresearch/fairscale/2064/workflows/3de30f4d-5c84-4b8a-bb51-8f8e6dc231db/jobs/12695/tests#failed-test-0

It is fsdp & moe fail to compare with DDP.
Yes, it seems to be failing intermittently, which is concerning. Probably something wrong with stream synchronization, I'll look into it

Resolving patch:
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -1127,7 +1127,7 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
         if params is None:
             params = self.params
         self.has_full_params = False
-        current_stream = torch.cuda.current_stream()
+        self._streams["all_gather"].wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(self._streams["all_gather"]):
             for p in params:
                 if not p._is_sharded:  # e.g., world_size == 1
@@ -1140,7 +1140,6 @@ def _free_full_params(self, params: Optional[List[Parameter]] = None) -> None:
                 # unshard parameters, we should reuse the original Tensor
                 # Storage object and unshard it in-place. For now, just resize
                 # the Storage to 0 to save memory.
-                p._full_param_padded.record_stream(current_stream)
                 free_storage_(p._full_param_padded)
 
     @torch.no_grad()


Test patch:
diff --git a/tests/nn/data_parallel/test_fsdp.py b/tests/nn/data_parallel/test_fsdp.py
--- a/tests/nn/data_parallel/test_fsdp.py
+++ b/tests/nn/data_parallel/test_fsdp.py
@@ -287,6 +287,20 @@ def test_mixture_of_experts(self, moe_config):
         )
         spawn_and_init(test_fn)
 
+    @parameterized.expand([[{"checkpoint_act": False}], [{"checkpoint_act": True}]], name_func=rename_test)
+    def test_mixture_of_experts_with_delay_before_free(self, moe_config):
+        fsdp_config = {"mixed_precision": True}
+        test_fn = functools.partial(
+            self._test_identical_outputs,
+            functools.partial(MixtureOfExperts, delay_before_free_ms=250, **moe_config),
+            fsdp_config,
+            # MixtureOfExperts implements custom reduce logic, so the reference
+            # behavior should use that logic instead of PyTorch DDP.
+            ref_ddp_fn=self._dummy_ddp_fn,
+            norm_type=None,
+        )
+        spawn_and_init(test_fn)
+
     def test_mixture_of_experts_grad_clip_breaks(self):
         config = {"mixed_precision": True}
         test_fn = functools.partial(
@@ -760,9 +774,10 @@ def forward(self, *args, **kwargs):
 
 
 class MixtureOfExperts(NestedWrappedModule):
-    def __init__(self, group, wrapper_config, checkpoint_act=False):
+    def __init__(self, group, wrapper_config, checkpoint_act=False, delay_before_free_ms=0):
         super().__init__(group, wrapper_config)
         self.group = group
+        self.delay_before_free_ms = delay_before_free_ms
 
         # "expert" params are different on each rank
         torch.manual_seed(42 + group.rank())
@@ -787,6 +802,22 @@ def __init__(self, group, wrapper_config, checkpoint_act=False):
 
         self.module = nn.Sequential(nn.Linear(8, 4), shared, expert, nn.Linear(4, 8))
 
+    def forward(self, x):
+        if self.delay_before_free_ms > 0:
+            expert = self.module[2]
+            if isinstance(expert, FullyShardedDataParallel):
+                orig_free_full_params = self.module[2]._free_full_params
+
+                def _free_full_params_with_delay(*args):
+                    torch.cuda._sleep(int(self.delay_before_free_ms * get_cycles_per_ms()))
+                    return orig_free_full_params(*args)
+
+                assert hasattr(expert, "_free_full_params")
+                with mock.patch.object(expert, "_free_full_params", _free_full_params_with_delay):
+                    return self.module(x)
+
+        return self.module(x)
+
     def run_backward(self, loss):
         loss.backward()