Lightning-AI/pytorch-lightning · error · TypeError

Found modules that are wrapped with `torch.distributed.fsdp.

Error message

Found modules that are wrapped with `torch.distributed.fsdp.FullyShardedDataParallel`. The `{self.__class__.__name__}` only supports the new FSDP2 APIs in PyTorch >= 2.4.

What it means

ModelParallelStrategy in current Lightning supports only FSDP2 (torch.distributed.fsdp.fully_shard, PyTorch >= 2.4). If it detects modules wrapped in the legacy class torch.distributed.fsdp.FullyShardedDataParallel, it raises TypeError, since mixing legacy FSDP with FSDP2-based plans is unsupported.

Source

Thrown at src/lightning/pytorch/strategies/model_parallel.py:175

        # Users can access device mesh in `LightningModule.configure_model()`
        assert self.lightning_module is not None
        self.lightning_module._device_mesh = self._device_mesh

    @override
    def setup(self, trainer: "pl.Trainer") -> None:
        from torch.distributed.fsdp import FullyShardedDataParallel

        assert self.model is not None
        assert self.accelerator is not None
        self.accelerator.setup(trainer)

        if not is_overridden("configure_model", self.lightning_module):
            raise TypeError(
                f"When using the {type(self).__name__}, you are required to override the `configure_model()` hook in"
                f" the LightningModule and apply parallelization there."
            )
        if any(isinstance(mod, FullyShardedDataParallel) for mod in self.model.modules()):
            raise TypeError(
                "Found modules that are wrapped with `torch.distributed.fsdp.FullyShardedDataParallel`."
                f" The `{self.__class__.__name__}` only supports the new FSDP2 APIs in PyTorch >= 2.4."
            )

        _materialize_distributed_module(self.model, self.root_device)

        self.model = self.precision_plugin.convert_module(self.model)
        self.model_to_device()  # move all remaining layers if any left on CPU.

        self.barrier()

        if trainer.state.fn == TrainerFn.FITTING:
            self.setup_optimizers(trainer)
        self.setup_precision_plugin()
        if trainer.state.fn == TrainerFn.FITTING:
            _optimizers_to_device(self.optimizers, self.root_device)

    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Replace FullyShardedDataParallel(module, ...) wrapping with FSDP2's fully_shard(module) applied per-submodule
2. Require torch >= 2.4 and Lightning versions where FSDP2 is the path
3. If you must use legacy FSDP, use an older Lightning release's FSDPStrategy

Example fix

# before
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
self.model = FSDP(self.model, auto_wrap_policy=policy)

# after
from torch.distributed.fsdp import fully_shard
for block in self.model.blocks:
    fully_shard(block)
fully_shard(self.model)

import torch
from torch.distributed.fsdp import FullyShardedDataParallel
assert not any(isinstance(m, FullyShardedDataParallel) for m in model.modules()), "use fully_shard (FSDP2) instead"

def uses_legacy_fsdp(model) -> bool:
    from torch.distributed.fsdp import FullyShardedDataParallel
    return any(isinstance(m, FullyShardedDataParallel) for m in model.modules())

When it happens

Trigger: Manually wrapping modules with FullyShardedDataParallel(m) inside configure_model (or anywhere) and then using ModelParallelStrategy / FSDP2-based strategies; old code migrated from FSDPStrategy (which used legacy FSDP1).

Common situations: Upgrading from PyTorch < 2.4 era FSDP code or Lightning's older FSDPStrategy; copy-pasted FSDP1 tutorials.

Related errors

• The optimizer has references to the model's meta-device para
• The optimizer has references to the model's meta-device para
• `precision={precision!r})` is not supported in FSDP. `precis
• `precision={precision!r}` does not use a scaler, found {scal
• Gradient clipping is not implemented for optimizers handling

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/18235f9cfb6272c0. Report an issue: GitHub.