Repository: pytorch
Issue: Using retain_graph in backward() with FSDP
Labels: oncall: distributed,triaged,module: fsdp

### 🐛 Describe the bug

We need to support a use case where for the same model output, we need to compute loss (and corresponding gradients) twice, each time with a different set of target labels. To do the above in one forward pass, we do not want the `backward()` to clean up the computation graph and hence we use `backward(retain_graph=True)`. 

**Without FSDP**
When the model is not wrapped in FSDP, we can make successive calls `loss1.backward(retain_graph=True), loss2.backward()` and they are executed successfully.
```
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer = nn.Linear(2, 2, bias=False)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        return self.softmax(self.layer(x))

loss_fn = nn.CrossEntropyLoss()
model = Model()
optimizer = torch.optim.AdamW(opt_model.parameters(), lr=0.01)
dummy_data = torch.ones((2, 2))
out = opt_model(dummy_data)

target1 = torch.zeros(2, dtype=torch.long)
optimizer.zero_grad()
loss1 = loss_fn(out, target1)
loss1.backward(retain_graph=True)

optimizer.zero_grad()
target2 = torch.ones(2, dtype=torch.long)
loss2 = loss_fn(out, target2)
loss2.backward()
```
**With FSDP (Script to reproduce error)**
However, when the model is wrapped using FSDP, multiple calls to `backward()` with the retain_graph option leads to error. The code to reproduce the error is given below:
```
import functools
import os

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
from torch.distributed.fsdp.wrap import enable_wrap, transformer_auto_wrap_policy, wrap


local_rank = int(os.getenv("LOCAL_RANK"))

torch.manual_seed(42 + local_rank)
torch.cuda.manual_seed_all(42 + local_rank)

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

wrapping_policy = functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={})

model_sharding_strategy = ShardingStrategy.FULL_SHARD

torch.cuda.set_device(int(os.getenv("RANK")))
dist.init_process_group("nccl")

v = 2048
d = 1024

model = nn.Sequential(nn.Embedding(v, d), nn.Linear(d, v))

model = FSDP(
    model,
    auto_wrap_policy=wrapping_policy,
    mixed_precision=mp_policy,
    sharding_strategy=model_sharding_strategy,
    device_id=local_rank,
    limit_all_gathers=True,
    use_orig_params=True,
)

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)
criterion = nn.CrossEntropyLoss()

# TRAINING LOOP
for i in range(10):
    inp = torch.randint(v, (8, 256), device=local_rank)
    label = inp.add(1) % v
    print("Input and labels created")
    optimizer.zero_grad()
    pred = model(inp)
    loss = criterion(pred.view(-1, pred.size(-1)), label.view(-1))
    print("Loss computed. Calling backward() with retain_graph..")
    loss.backward(retain_graph=True)
    print("Gradient computed once. Called backward() again without retain_graph..")
    loss.backward()
    print("Gradient computed twice")
    optimizer.step()
    print("Update step complete")
    if local_rank == 0 and (i + 1) % 10 == 0:
        print(f"Step {i+1}, loss {loss.item()}")
```
The relevant error log is given below:
```
Input and labels created
NCCL version 2.17.1+cuda11.8
Loss computed. Calling backward() with retain_graph..
Gradient computed once. Called backward() again without retain_graph..
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/workspace/foundation-model-stack/nlp/scripts/pretraining/repro.py", line 59, in <module>
    loss.backward()
  File "/opt/conda/lib/python3.10/site-packages/torch/_tensor.py", line 491, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/__init__.py", line 204, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 690, in _pre_backward_hook
    if state._needs_pre_backward_unshard[_handles_key]:
KeyError: (<torch.distributed.fsdp.flat_param.FlatParamHandle object at 0x7f8889313910>,)
[2023-08-04 20:41:22,975] torch.distributed.elastic.multiprocessing.api: [ERROR] failed (exitcode: 1) local_rank: 0 (pid: 135) of binary: /opt/conda/bin/python
Traceback (most recent call last):
  File "/opt/conda/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/run.py", line 797, in main
    run(args)
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/run.py", line 788, in run
    elastic_launch(
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 134, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/opt/conda/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
```

### Versions

Pytorch nightly image from 07-17-2023.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin