Lightning-AI/pytorch-lightning · error · MisconfigurationException

SWA does not currently support sharded models.

Error message

SWA does not currently support sharded models.

What it means

StochasticWeightAveraging deep-copies the full model to accumulate weights, which is incompatible with model-weights sharding. At setup() it rejects FSDPStrategy and DeepSpeedStrategy with this MisconfigurationException.

Source

Thrown at src/lightning/pytorch/callbacks/stochastic_weight_avg.py:153

    @property
    def swa_start(self) -> int:
        assert isinstance(self._swa_epoch_start, int)
        return max(self._swa_epoch_start - 1, 0)  # 0-based

    @property
    def swa_end(self) -> int:
        if self._max_epochs == -1:
            return float("inf")  # type: ignore[return-value]
        return self._max_epochs - 1  # 0-based

    @staticmethod
    def pl_module_contains_batch_norm(pl_module: "pl.LightningModule") -> bool:
        return any(isinstance(module, nn.modules.batchnorm._BatchNorm) for module in pl_module.modules())

    @override
    def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
        if isinstance(trainer.strategy, (FSDPStrategy, DeepSpeedStrategy)):
            raise MisconfigurationException("SWA does not currently support sharded models.")

        # copy the model before moving it to accelerator device.
        self._average_model = deepcopy(pl_module)

    @override
    def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        if len(trainer.optimizers) != 1:
            raise MisconfigurationException("SWA currently works with 1 `optimizer`.")

        if len(trainer.lr_scheduler_configs) > 1:
            raise MisconfigurationException("SWA currently not supported for more than 1 `lr_scheduler`.")

        assert trainer.max_epochs is not None
        if isinstance(self._swa_epoch_start, float):
            if trainer.max_epochs == -1:
                raise MisconfigurationException(
                    "SWA with `swa_epoch_start` as a float is not supported when `max_epochs=-1`. "
                    "Please provide `swa_epoch_start` as an integer."

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Remove the SWA callback when using FSDP or DeepSpeed
2. For FSDP-style averaging, use a custom averaging approach or ecosystem tools supporting sharded weights (e.g. compile averaged weights offline from checkpoints)
3. Switch to a non-sharded strategy (DDP) if SWA is essential and memory permits

Example fix

# before
trainer = Trainer(strategy="fsdp", callbacks=[SWA(swa_epoch_start=0.7)])
# after
trainer = Trainer(strategy="ddp", callbacks=[SWA(swa_epoch_start=0.7)])
# or simply drop the SWA callback under fsdp

from lightning.pytorch.strategies import FSDPStrategy, DeepSpeedStrategy
from lightning.pytorch.callbacks import StochasticWeightAveraging

def callbacks_for(strategy):
    if isinstance(strategy, (FSDPStrategy, DeepSpeedStrategy)):
        return []  # SWA unsupported
    return [StochasticWeightAveraging(swa_epoch_start=0.75)]

def swa_supported(trainer) -> bool:
    from lightning.pytorch.strategies import FSDPStrategy, DeepSpeedStrategy
    return not isinstance(trainer.strategy, (FSDPStrategy, DeepSpeedStrategy))

When it happens

Trigger: Trainer(strategy='fsdp', ...) or Trainer(strategy='deepspeed_stage_2' / a DeepSpeedStrategy instance, ...) together with callbacks=[SWA(...)].

Common situations: Adding SWA to a large-model FSDP/DeepSpeed training script; upgrading a single-GPU SWA recipe to multi-node sharded training.

Related errors

• The hybrid sharding strategy requires you to pass at least o
• When saving the DeepSpeed Stage 3 checkpoint, each worker wi
• No models were set up for backward. Did you forget to call `
• When using multiple models + deepspeed, please provide the m
• The optimizer has references to the model's meta-device para

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