{"record":{"id":"dfc230e942eab9af","repo":"Lightning-AI/pytorch-lightning","slug":"swa-does-not-currently-support-sharded-models","errorCode":null,"errorMessage":"SWA does not currently support sharded models.","messagePattern":"SWA does not currently support sharded models\\.","errorType":"exception","errorClass":"MisconfigurationException","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/callbacks/stochastic_weight_avg.py","lineNumber":153,"sourceCode":"    @property\n    def swa_start(self) -> int:\n        assert isinstance(self._swa_epoch_start, int)\n        return max(self._swa_epoch_start - 1, 0)  # 0-based\n\n    @property\n    def swa_end(self) -> int:\n        if self._max_epochs == -1:\n            return float(\"inf\")  # type: ignore[return-value]\n        return self._max_epochs - 1  # 0-based\n\n    @staticmethod\n    def pl_module_contains_batch_norm(pl_module: \"pl.LightningModule\") -> bool:\n        return any(isinstance(module, nn.modules.batchnorm._BatchNorm) for module in pl_module.modules())\n\n    @override\n    def setup(self, trainer: \"pl.Trainer\", pl_module: \"pl.LightningModule\", stage: str) -> None:\n        if isinstance(trainer.strategy, (FSDPStrategy, DeepSpeedStrategy)):\n            raise MisconfigurationException(\"SWA does not currently support sharded models.\")\n\n        # copy the model before moving it to accelerator device.\n        self._average_model = deepcopy(pl_module)\n\n    @override\n    def on_fit_start(self, trainer: \"pl.Trainer\", pl_module: \"pl.LightningModule\") -> None:\n        if len(trainer.optimizers) != 1:\n            raise MisconfigurationException(\"SWA currently works with 1 `optimizer`.\")\n\n        if len(trainer.lr_scheduler_configs) > 1:\n            raise MisconfigurationException(\"SWA currently not supported for more than 1 `lr_scheduler`.\")\n\n        assert trainer.max_epochs is not None\n        if isinstance(self._swa_epoch_start, float):\n            if trainer.max_epochs == -1:\n                raise MisconfigurationException(\n                    \"SWA with `swa_epoch_start` as a float is not supported when `max_epochs=-1`. \"\n                    \"Please provide `swa_epoch_start` as an integer.\"","sourceCodeStart":135,"sourceCodeEnd":171,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/callbacks/stochastic_weight_avg.py#L135-L171","documentation":"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.","triggerScenarios":"Trainer(strategy='fsdp', ...) or Trainer(strategy='deepspeed_stage_2' / a DeepSpeedStrategy instance, ...) together with callbacks=[SWA(...)].","commonSituations":"Adding SWA to a large-model FSDP/DeepSpeed training script; upgrading a single-GPU SWA recipe to multi-node sharded training.","solutions":["Remove the SWA callback when using FSDP or DeepSpeed","For FSDP-style averaging, use a custom averaging approach or ecosystem tools supporting sharded weights (e.g. compile averaged weights offline from checkpoints)","Switch to a non-sharded strategy (DDP) if SWA is essential and memory permits"],"exampleFix":"# before\ntrainer = Trainer(strategy=\"fsdp\", callbacks=[SWA(swa_epoch_start=0.7)])\n# after\ntrainer = Trainer(strategy=\"ddp\", callbacks=[SWA(swa_epoch_start=0.7)])\n# or simply drop the SWA callback under fsdp","handlingStrategy":"validation","validationCode":"from lightning.pytorch.strategies import FSDPStrategy, DeepSpeedStrategy\nfrom lightning.pytorch.callbacks import StochasticWeightAveraging\n\ndef callbacks_for(strategy):\n    if isinstance(strategy, (FSDPStrategy, DeepSpeedStrategy)):\n        return []  # SWA unsupported\n    return [StochasticWeightAveraging(swa_epoch_start=0.75)]","typeGuard":"def swa_supported(trainer) -> bool:\n    from lightning.pytorch.strategies import FSDPStrategy, DeepSpeedStrategy\n    return not isinstance(trainer.strategy, (FSDPStrategy, DeepSpeedStrategy))","tryCatchPattern":null,"preventionTips":["Check trainer.strategy before enabling SWA","Keep a strategy->allowed-callbacks mapping in multi-backend training scripts"],"tags":["swa","fsdp","deepspeed","sharding","incompatible-strategy"],"backgroundTag":"callback-strategy-incompatible","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}