Lightning-AI/pytorch-lightning · error · TypeError

When using the {type(self).__name__}, you are required to ov

Error message

When using the {type(self).__name__}, you are required to override the `configure_model()` hook in the LightningModule and apply parallelization there.

What it means

ModelParallelStrategy requires that model parallelization (e.g. FSDP2 fully_shard calls) happen inside LightningModule.configure_model(), which is invoked on each rank with the right device context. If that hook is not overridden, Lightning cannot apply parallelization and raises TypeError.

Source

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

        if self._tensor_parallel_size == "auto":
            self._tensor_parallel_size = self.num_processes
        self._device_mesh = _setup_device_mesh(
            self._data_parallel_size, self._tensor_parallel_size, self.world_size, self.root_device
        )
        # 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)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Override configure_model() in your LightningModule and apply parallelization there (e.g. torch.distributed.fsdp.fully_shard on submodules)
2. Do the wrapping lazily there rather than in __init__
3. See the ModelParallel / FSDP2 examples in the Lightning repo/docs

Example fix

# before
class LitModel(L.LightningModule):
    def __init__(self):
        self.model = Transformer()  # no configure_model

# after
class LitModel(L.LightningModule):
    def configure_model(self):
        self.model = Transformer()
        for block in self.model.blocks:
            fully_shard(block)
        fully_shard(self.model)

from lightning.pytorch.utilities.model_helpers import is_overridden
assert is_overridden("configure_model", model), "override configure_model() when using ModelParallelStrategy"

When it happens

Trigger: Using ModelParallelStrategy (or subclasses like FSDP2Strategy/ModelParallelStrategy-based strategies) with a LightningModule that does not override configure_model; wrapping applied in __init__ or never applied at all.

Common situations: Migrating a single-device model to a model-parallel strategy; new users expecting automatic FSDP wrapping (as older FSDPStrategy auto-wrapped).

Related errors

• Accessing the device mesh before processes have initialized
• The `{type(self).__name__}` does not use the `CheckpointIO`
• Could not find a distributed model in the provided checkpoin
• Found multiple distributed models in the given state. Loadin
• The sizes `data_parallel_size={data_parallel_size}` and `ten

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