Repository: facebookresearch/fairscale

Issue:
FSDP fails with multiple forward and then multiple backward calls
I have a use case where I need to do multiple forward calls and then backwards call instead of usual pattern but FSDP throws following error when calling multiple backwards:

```
Asserting FSDP instance is: FullyShardedDataParallel(
  world_size=2, flatten_parameters=True, mixed_precision=False,
  (_fsdp_wrapped_module_test): FlattenParamsWrapper(
    (_fpw_module): FFN(
      (fc1): Linear(in_features=10, out_features=10, bias=True)
      (fc2): Linear(in_features=10, out_features=10, bias=True)
      (relu): ReLU()
    )
  )
)
ERROR: expected to be in states [<TrainingState.BACKWARD_POST: 4>] but current state is TrainingState.BACKWARD_PRE
```

Code to repro error: 
```
import torch
import torch.nn as nn
import argparse
from fairscale.nn.data_parallel import FullyShardedDataParallel
from fairscale.nn import enable_wrap
from fairscale.nn import wrap


class FFN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))


def main(i, args):
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    print(i)
    torch.cuda.set_device(i)
    torch.distributed.init_process_group(
        backend='nccl',
        init_method='tcp://localhost:22345',
        world_size=args.world_size,
        rank=i,
    )
    ffn = FFN()

    with enable_wrap(
        wrapper_cls=FullyShardedDataParallel
    ):
        model = wrap(
            ffn,
            process_group=torch.distributed.new_group(),
            flatten_parameters=True,
            compute_dtype=torch.float16
        ).half()

    model = model.train()

    x1 = torch.rand((10, 10)).cuda().half()
    out1 = model(x1)
    loss1 = out1.sum()
    
    x2 = torch.rand((10, 10)).cuda().half()
    out2 = model(x2)
    loss2 = out2.sum()
    
    with model.no_sync():
        loss1.backward()
        loss2.backward()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--world-size', '--ws', type=int)
    args = parser.parse_args()
    torch.multiprocessing.spawn(
        fn=main, 
        args=(args, ),
        nprocs=min(
            torch.cuda.device_count(),
            args.world_size,
        ),
        join=True,
    )
```



Discussion:
I see 2 losses that are different from the use case that we have covered in our tests [here](https://github.com/facebookresearch/fairscale/blob/0a526bcbe212b043a00a5bdfae73c7ef532c88b4/tests/nn/data_parallel/test_fsdp_grad_acc.py#L96). Does the error happen on the second BW call? Can you also share a complete stack trace of the error?

Also following up from our offline thread, can you tell me a bit more about the use case?
Full stack trace: 

I might have some local changes so the line numbers may not lineup with whats in main
```
ERROR: expected to be in states [<TrainingState.BACKWARD_POST: 4>] but current state is TrainingState.BACKWARD_PRE
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/namangoyal/rsc/fairseq_gshard/fairseq-py/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1473, in _wait_for_post_backward
    self.assert_state(TrainingState.BACKWARD_POST)
  File "/home/namangoyal/rsc/fairseq_gshard/fairseq-py/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 1891, in assert_state
    traceback.print_stack()
1 True False TrainingState.IDLE
1 True False TrainingState.IDLE
Traceback (most recent call last):
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/namangoyal/rsc/fairseq_gshard/fairseq-py/scripts/test_fsdp/multiple_backward.py", line 69, in <module>
    torch.multiprocessing.spawn(
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 240, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 198, in start_processes
    while not context.join():
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/namangoyal/rsc/fairseq_gshard/fairseq-py/scripts/test_fsdp/multiple_backward.py", line 59, in main
    loss2.backward()
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/_tensor.py", line 352, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/namangoyal/.conda/envs/fairseq-20220105/lib/python3.9/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
SystemError: <built-in method run_backward of torch._C._EngineBase object at 0x7fd5ac2e4c40> returned NULL without setting an error
```



cc @zhaojuanmao for this use case as well.

we tested

```
loss1 = fwd()
loss1.backward()
loss2 = fwd()
loss2.backward()
```

This new use case is different. The assert msg means that FSDP wrapper expects to see the backward hooks firing in certain order but it was not.

does this use case work with DDP? Also, why delay calling loss1.backward()? It seems OK to backward on it first and then do the second fwd on x2, right? Doing it in the example you have seems to only increase the GPU memory needed for the activations. Am I missing something?
@min-xu-ai yes it works with DDP and its required currently for my usecase of combining FSDP with pipeline parallel cause the backward of a microbatch is delayed.   
I suspect that the problem is with pre-backward hooks, which is installed on the output tensor using register_hook. For models that runs part of a module's fwd function multiple times, we have a guard that prevents pre-backward hook to run multiple times.

<img width="923" alt="image" src="https://user-images.githubusercontent.com/24926999/150874309-84fc5112-206f-4331-8818-92a4eedd07a5.png">


This guard is reset at the end of each backward, but I suspect the use of `no_sync` context prevented that. What if you remove the no_sync context? no_sync + grad_accumulation is not ideal for FSDP IMHO since it consumes more memory due to the un-synced gradients.
It still fails without the no_sync context. 
`no_sync + grad_accumulation is not ideal for FSDP`, I am not sure about this, basically the goal is to do adam state sharding without doing model weights sharding, which is deepspeed zero stage 2.
Ah, I see. That makes a lot of sense. 

I think you meant zero stage 1 though. Since in stage 2 you would shard the gradients. For sharping only the Adam state, perhaps using ddp + oss.py would be simpler for you? Just an idea. 
@ngoyal2707  Is there an action item for us to follow up here or will you be trying what Min suggested? Want to make sure we unblock your use case if needed.
@anj-s I still think action item is to basically support above use-case with FSDP. I understand that gradients wont be sharded due to `no_sync` but I dont see any technical reason to not support this use-case.


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
@@ -586,7 +586,7 @@ def append_shared_param(self, p: Parameter) -> None:
 
             p must to be already sharded by the owning module.
 
-            Check the corresponding unit test to see how is it used and tested.
+            Check the corresponding unit tests to see how is it used and tested.
             In particular, the sharing FSDP wrappers are "siblings" not "parent"
             and "child" of each other in the nested module structure.
 
@@ -1700,35 +1700,38 @@ def _register_post_backward_hooks(self) -> None:
         3. If it fires once but too early or doesn't fire, we leave gradients
         unsharded. (could lead to dimension too large)
 
-        Due to multiple-pass forward, this function can be called on
-        the same parameter multiple times in a single forward pass. If we register
-        the hook multiple time, we end up getting called multiple times. We
-        could try to get a new hook every time and delete the previous one
-        registered. However, due to *unknown reason* (I have debugged it for
-        a long time!), in mixed precision mode, we get two different ``grad_acc``
-        objects below during different calls of this function (in the same
-        forward pass). If we keep the last one, the hook end up firing too
-        early. In full precision mode, we luckily get the *same* ``grad_acc``
-        object, so deleting and re-registering still ensured the hook fire
-        once after all gradients are generated.
-
-        Empirically, keep the first hook register per forward pass seems to
-        work the best. We do need to remove the hook at the end of the
-        backward pass. Otherwise, the next forward pass will not register
-        a new hook, which is needed for a new forward pass.
+        There are several cases here:
+        1. We can call the same module multiple times in a single outer forward
+           pass. We register multiple hooks but autograd should fire the last
+           one after the total gradient is computed and accumulated. If it does
+           fire multiple times, we may have a crash due to gradient being already
+           sharded and shape mismatch.
+           On the other hand, due to _saved_grad_shard, this case may also work
+           but with extra grad scatter-gather.
+        2. With activation checkpointing and case 1.
+        3. The same outer forward can be called multiple times before any backward
+           is called (within the no_sync context) for a special way of gradient
+           accumulation. (see test_fsdp_fwd_fwd_bwd_bwd.py)
+        4. When a param is shared by multiple FSDP wrapper instances, this can
+           register multiple times. (See test_fsdp_shared_weights.py)
+
+        It appears that registering the hook everytime and let them fire and
+        hook being removed/freed automatically is the correct thing to do. But this
+        is purely based on experiments.
         """
         if not torch.is_grad_enabled():
             return  # don't register grad hooks if grad isn't enabled
         for p in self.params:
             if p.requires_grad:
-                if hasattr(p, "_shard_bwd_hook"):
-                    continue
-                # Register a hook on the first call, empirically, autograd
-                # fires it at the end for this param, which makes sense.
+                # Register a hook.
                 p_tmp = p.expand_as(p)  # Get a grad_fn on p_tmp.
                 assert p_tmp.grad_fn is not None
                 grad_acc = p_tmp.grad_fn.next_functions[0][0]  # Gets its GradAccumulation object.
                 handle = grad_acc.register_hook(functools.partial(self._post_backward_hook, p))
+                # Important, we need to save the hook, otherwise, it appears to be
+                # deleted/freed/unregistered.
+                # However, we don't free/unhook at the end of bwd (as we used to do it
+                # in _finalize_parameters below). If we do, that may unregister the wrong hook.
                 p._shard_bwd_hook = (grad_acc, handle)
 
     @torch.no_grad()
@@ -1756,20 +1759,26 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None:
         # then subsequent hook callbacks will see POST state.
         self.assert_state([TrainingState.BACKWARD_PRE, TrainingState.BACKWARD_POST])
         self.training_state = TrainingState.BACKWARD_POST
-        if param.grad is None:
-            return
 
         if hasattr(param, "_linked_param"):
-            # This links to a shared param. We should finalize the linked param here.
-            assert param.shape == (1,), param.shape
+            # This links to a shared param. We should try to finalize the linked param here.
+            # This is done by module code to ensure correct gradient computation.
+            # p._is_shared and p._linked_param are closely related but not the same.
+            # See fairscale/experimental/nn/mevo.py.
+            assert param.shape == (1,), param.shape  # This param should have this special dim.
             # If the _is_shared flag is set, then this shared weight is indeed being
             # shared between different FSDP wrappers. Otherwise, they are linked but
             # likely in the same FSDP wrapper, which means we shouldn't finalize the
             # linked param..
             if hasattr(param._linked_param, "_is_shared") and param._linked_param._is_shared:
+                # param._linked_param may or may not have .grad since this callback
+                # could happen multiple times to support #918. Since we check `if param.grad is None`
+                # below anyway, this is OK.
                 param = param._linked_param
 
-        assert param.grad is not None, param.shape
+        if param.grad is None:
+            return
+
         if param.grad.requires_grad:
             raise RuntimeError("FSDP only works with gradients that don't require gradients")
 
@@ -1950,10 +1959,6 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
             for p in fsdp_module.params:
                 if not p.requires_grad:
                     continue
-                if hasattr(p, "_shard_bwd_hook"):
-                    p_assert(len(p._shard_bwd_hook) == 2, f"WFPB: incorrect hook num: {len(p._shard_bwd_hook)}")
-                    p._shard_bwd_hook[1].remove()
-                    delattr(p, "_shard_bwd_hook")
 
                 # Leave the gradient accumulation state as-is if not synchronizing this pass. This ensures p.grad
                 # remains the unsharded gradient accumulated from prior no-sync passes, and p._saved_grad_shard
@@ -1969,11 +1974,13 @@ def _finalize_parameters(fsdp_module: FullyShardedDataParallel) -> None:
                 elif hasattr(p, "_saved_grad_shard"):
                     p_assert(
                         p.device == p._saved_grad_shard.device,
-                        f"WFPB: incorrect saved_grad_shard device {p.device} vs {p._saved_grad_shard.device}",
+                        f"WFPB: incorrect saved_grad_shard device p.device={p.device} "
+                        f"vs p._saved_grad_shard.device={p._saved_grad_shard.device}",
                     )
                     p_assert(
                         p.shape == p._saved_grad_shard.shape,
-                        f"WFPB: incorrect saved_grad_shard shape {p.shape} vs {p._saved_grad_shard.shape}",
+                        f"WFPB: incorrect saved_grad_shard shape p.shape={p.shape} "
+                        f"vs p._saved_grad_shard.shape={p._saved_grad_shard.shape}",
                     )
                     p.grad = p._saved_grad_shard
 


Test patch:
diff --git a/tests/ci_test_list_2.txt b/tests/ci_test_list_2.txt
--- a/tests/ci_test_list_2.txt
+++ b/tests/ci_test_list_2.txt
@@ -51,3 +51,4 @@ tests/nn/pipe/test_stream.py
 tests/nn/moe/test_moe_layer.py
 tests/nn/moe/test_top2gating.py
 tests/nn/data_parallel/test_fsdp_offload.py
+tests/nn/data_parallel/test_fsdp_fwd_fwd_bwd_bwd.py
diff --git a/tests/nn/data_parallel/test_fsdp_fwd_fwd_bwd_bwd.py b/tests/nn/data_parallel/test_fsdp_fwd_fwd_bwd_bwd.py
new file mode 100644
--- /dev/null
+++ b/tests/nn/data_parallel/test_fsdp_fwd_fwd_bwd_bwd.py
@@ -0,0 +1,72 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from fairscale.fair_dev.testing.testing import skip_if_single_gpu, temp_files_ctx
+from fairscale.nn import enable_wrap, wrap
+from fairscale.nn.data_parallel import FullyShardedDataParallel
+
+
+class FFN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(10, 10)
+        self.fc2 = nn.Linear(10, 10)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        return self.fc2(self.relu(self.fc1(x)))
+
+
+def main(rank, sync_file):
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+    torch.cuda.set_device(rank)
+    torch.distributed.init_process_group(
+        backend="nccl",
+        init_method=f"file://{sync_file}",
+        world_size=2,
+        rank=rank,
+    )
+    ffn = FFN().cuda().half()
+
+    with enable_wrap(wrapper_cls=FullyShardedDataParallel):
+        model = wrap(
+            ffn,
+            process_group=torch.distributed.new_group(),
+            flatten_parameters=True,
+            compute_dtype=torch.float16,
+        )
+
+    model = model.train()
+
+    # We test this behavior because it might be used by pipelining.
+    # However, we don't check if the speed (compute/comm overlapping)
+    # and memory (necessary all-gather & free) are optimal.
+    losses = []
+    for _ in range(3):
+        x = torch.rand((10, 10)).cuda().half()
+        out = model(x)
+        loss = out.sum()
+        losses.append(loss)
+
+    # Only the last bwd can be outside of no_sync context.
+    with model.no_sync():
+        losses[0].backward()
+        losses[1].backward()
+    losses[2].backward()
+
+
+@skip_if_single_gpu
+def test_fwd_fwd_bwd_bwd():
+    with temp_files_ctx(num=1) as temp_files:
+        torch.multiprocessing.spawn(
+            fn=main,
+            nprocs=2,
+            args=(temp_files[0],),
+            join=True,
+        )