FSDP Fails with floating nn.Parameter

Bug description

I'm training an adversarial model with PyTorch Lightning, similar to a GAN.:

• The generator and discriminator each have their own optimizer.
• We train the generator for G steps consecutively with a frozen discriminator and then train the discriminator for D steps consecutively with a frozen generator.

When I try training the model with FSDP strategy -- I receive errors during backprop:

• The first step of training is successful in optimizing model 1 according to step 1.
• This error emerges on the second step of training, when I toggle optimizers and take a step to optimize model 2 according to loss function 2. This leads me to believe that my error is tied to incorrect handling of the unused parameters.

What version are you seeing the problem on?

master

How to reproduce the bug

"""
File: test_fsdp.py
Description: Minimal example of FSDP failure with unused parameters
"""

import os

import torch
from torch import nn
from pytorch_lightning import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset


class AdversarialModel(LightningModule):

    def __init__(self):
        super().__init__()
        self.generator = GeneratorNetwork()
        self.discriminator = DiscriminatorNetwork()
        self.automatic_optimization = False

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        opts = self.optimizers()
        current_cycle = batch_idx % len(opts)

        if current_cycle == 0:
            #  compute loss from generator
            self.computed_loss = self.generator(batch).mean()
        else:
            # compute loss from discriminator
            self.computed_loss = self.discriminator(batch).mean()

    def on_train_batch_end(self, outputs, batch, batch_idx):
        opts = self.optimizers()
        current_cycle = batch_idx % len(opts)
        opt = opts[current_cycle]

        with opt.toggle_model():
            self.manual_backward(self.computed_loss)
            opt.step()
            opt.zero_grad()

    def validation_step(self, batch, batch_idx):
        generator_loss = self.generator(batch).mean()
        discriminator_loss = self.discriminator(batch).mean()
        self.log("valid_generator_loss", generator_loss)
        self.log("valid_discriminator_loss", discriminator_loss)

    def configure_optimizers(self):
        return [
            torch.optim.SGD(self.generator.parameters(), lr=0.1),
            torch.optim.SGD(self.discriminator.parameters(), lr=0.1)
        ]


class GeneratorNetwork(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def parameters(self, recurse: bool = True):
        return self.layer.parameters(recurse=recurse)


class DiscriminatorNetwork(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def parameters(self, recurse: bool = True):
        return self.layer.parameters(recurse=recurse)


class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = AdversarialModel()
    trainer = Trainer(default_root_dir=os.getcwd(),
                      limit_train_batches=10,
                      limit_val_batches=10,
                      num_sanity_val_steps=0,
                      max_epochs=1,
                      enable_model_summary=False,
                      num_nodes=1,
                      devices=8,
                      strategy='fsdp',
                      enable_progress_bar=True)
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)


if __name__ == "__main__":
    run()

Error messages and logs

ERROR: expected to be in states [<TrainingState.FORWARD_BACKWARD: 2>] but current state is TrainingState.IDLE
  File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 708, in _post_backward_hook
    _assert_in_training_states(state, [TrainingState.FORWARD_BACKWARD])
  File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_common_utils.py", line 471, in _assert_in_training_states
    traceback.print_stack()
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/sahil/project/tests/test_fsdp.py", line 119, in <module>
[rank0]:     run()
[rank0]:   File "/home/sahil/project/tests/test_fsdp.py", line 115, in run
[rank0]:     trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 544, in fit
[rank0]:     call._call_and_handle_interrupt(
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 580, in _fit_impl
[rank0]:     self._run(model, ckpt_path=ckpt_path)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 990, in _run
[rank0]:     results = self._run_stage()
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1034, in _run_stage
[rank0]:     self.fit_loop.run()
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
[rank0]:     self.advance()
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
[rank0]:     self.epoch_loop.run(self._data_fetcher)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run
[rank0]:     self.advance(data_fetcher)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 270, in advance
[rank0]:     call._call_lightning_module_hook(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 159, in _call_lightning_module_hook
[rank0]:     output = fn(*args, **kwargs)
[rank0]:   File "/home/sahil/project/tests/test_fsdp.py", line 44, in on_train_batch_end
[rank0]:     self.manual_backward(self.computed_loss)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1084, in manual_backward
[rank0]:     self.trainer.strategy.backward(loss, None, *args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 212, in backward
[rank0]:     self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision.py", line 72, in backward
[rank0]:     model.backward(tensor, *args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1103, in backward
[rank0]:     loss.backward(*args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/_tensor.py", line 525, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/autograd/__init__.py", line 267, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/autograd/graph.py", line 744, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 708, in _post_backward_hook
[rank0]:     _assert_in_training_states(state, [TrainingState.FORWARD_BACKWARD])
[rank0]:   File "/home/sahil/.cache/pypoetry/virtualenvs/project-auw7Hy33-py3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_common_utils.py", line 472, in _assert_in_training_states
[rank0]:     raise ValueError(msg)
[rank0]: ValueError: expected to be in states [<TrainingState.FORWARD_BACKWARD: 2>] but current state is TrainingState.IDLE

Environment

More info

Originally, I thought it was a PyTorch issue -- https://github.com/pytorch/pytorch/issues/132068

But after converting my reproduction script to raw torch, the error went away -- so my hunch is that there is an edge case in the PyTorch Lightning FSDP wrapper.

cc @awaelchli @carmocca

Start with the reproducer in test_fsdp.py, especially training_step, on_train_batch_end, and the manual_backward call, then trace the FSDP failure at torch/distributed/fsdp/_runtime_utils.py:708. Run the two-optimizer example with the stated FSDP setup and compare the first and second optimization steps. Done means the discriminator step completes without the TrainingState.IDLE error.