Repository: pytorch
Issue: FSDP Forward order differs from that of first run
Labels: oncall: distributed,module: fsdp

### 🐛 Describe the bug

Hi,

I am using fsdp(integrated with hf accelerate) to extend support for the [transformer reinforcement learning library](https://github.com/lvwerra/trl) to multi-gpu. This requires me to run multiple .generate calls, forward passes, and then a backwards pass. 

I am getting a warning: `UserWarning: Forward order differs from that of the first iteration on rank 0 -- collectives are unchecked and may give incorrect results or hang`. Some insight would be appreciated.

Minimum code to reproduce:

```
import torch
import transformers
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer
)
from accelerate import Accelerator
from accelerate.logging import get_logger
from trl.gpt2 import GPT2HeadWithValueModel
from torch.optim import Adam
from transformers import DataCollatorForLanguageModeling
import torch.nn.functional as F


def main():
    accelerator = Accelerator()
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model = GPT2HeadWithValueModel.from_pretrained('gpt2')
    ref_model = GPT2HeadWithValueModel.from_pretrained('gpt2')
    optimizer = Adam(model.parameters(), lr=1.41e-5)
    model.config.pad_token_id = model.config.eos_token_id
    model = accelerator.prepare(model)
    model.to(accelerator.device)
    print(accelerator.state)


    optimizer = accelerator.prepare(optimizer)

    rank = torch.distributed.get_rank()
    if rank == 0:
        text_in = "The purpose of life is "
    elif rank == 1:
        text_in = "Are you human? "

    query_tensors = tokenizer(text_in, return_tensors="pt").to(accelerator.device)["input_ids"]

    # had to run this 1 time at the start else was giving device mismatch error.
    # So, before directly using `model.generate` pass a batch with dummy data through the model
    outputs = model(query_tensors)

    print(query_tensors)
    gen_kwargs = {
        "max_length": 64,
        "min_length": 20,
    }
    with torch.no_grad():
        unwrapped_model = accelerator.unwrap_model(model)
        # synced_gpus was necessary else resulted into indefinite hang
        response_tensors = unwrapped_model.generate(query_tensors, synced_gpus=True, **gen_kwargs)

    text_out = tokenizer.decode(response_tensors[0], skip_special_tokens=True)
    # Arbitrarily score generation
    score = torch.tensor([1.0]).to(accelerator.device)
    print(f"\nrank{rank}:\n   in={text_in}\n  out={text_out}")

    # Now compute ppo loss
    ## First compute logprobs and ref_logprobs
    collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
    for i in range(10):
        input_ids = collator([torch.cat([q, r]) for q, r in zip(query_tensors, response_tensors)])["input_ids"]

        with torch.no_grad():
            logits, _, v = model(input_ids)
            #print('values', v)
            ref_logits, _, _ = ref_model(input_ids.cpu())
            ref_logits = ref_logits.to(accelerator.device)

        logprobs = logprobs_from_logits(logits[:,:-1,:], input_ids[:,1:])
        ref_logprobs = logprobs_from_logits(ref_logits[:,:-1,:], input_ids[:,1:])

        # Only care about logprobs for generated text
        start = query_tensors.size()[-1] - 1
        end = query_tensors.size()[-1] + response_tensors.size()[-1] - 1
        logprobs = logprobs[:, start:end]
        ref_logprobs = ref_logprobs[:, start:end]
        v = v[:, start-1: end-1]
        print('logprob sizes', logprobs.size(), ref_logprobs.size(), v.size())


        ## Compute rewards
        kl = logprobs - ref_logprobs
        non_score_reward = .2 * kl
        reward = non_score_reward.clone()
        reward[-1] += score

        ## Compute losses
        lastgaelam = 0
        advantages_reversed = []
        gen_len = response_tensors.shape[1]
        for t in reversed(range(gen_len)):
            nextvalues = v[:, t+1] if t < gen_len - 1 else 0.0
            delta = reward[:, t] + 1.00 * nextvalues - v[:, t]
            lastgaelam = delta + 1.00 * .99 * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

        returns = advantages + v
        advantages = advantages.detach()

        ### With grad this time
        logits, _, vpred = model(input_ids)
        logprob = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
        logprob, vpred = logprob[:, -gen_len:], vpred[:, -gen_len-1:-1]
        vf_loss = torch.mean((vpred - returns)**2)

        # Backpropagate
        optimizer.zero_grad()
        accelerator.backward(vf_loss)
        optimizer.step()


def logprobs_from_logits(logits, labels):
    """
    See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591
    """
    logp = F.log_softmax(logits, dim=2)
    logpy = torch.gather(logp, 2, labels.unsqueeze(2)).squeeze(-1)
    return logpy


if __name__ == "__main__":
    main()

```

A similar issue has already been opened on accelerate's repo with some discussion [here](https://github.com/huggingface/accelerate/issues/570)

### Versions

PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Amazon Linux release 2 (Karoo) (x86_64)
GCC version: (GCC) 7.3.1 20180712 (Red Hat 7.3.1-15)
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.26

Python version: 3.8.5 (default, Feb 18 2021, 01:24:20)  [GCC 7.3.1 20180712 (Red Hat 7.3.1-12)] (64-bit runtime)
Python platform: Linux-5.10.126-117.518.amzn2.x86_64-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.129.06
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.2
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchvision==0.13.1+cu113
[conda] Could not collect

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501