Lightning-AI/pytorch-lightning · error · ValueError

Currently model averaging cannot work with a distributed opt

Error message

Currently model averaging cannot work with a distributed optimizer of type {optimizer.__class__.__name__}.

What it means

Post-localSGD's periodic model averaging is incompatible with optimizers that shard or distributedly update parameters. During setup, _enable_model_averaging inspects self.optimizers and raises ValueError if any is ZeroRedundancyOptimizer, PostLocalSGDOptimizer, a torch DistributedOptimizer (non-Windows), or a LightningOptimizer wrapping one of these.

Source

Thrown at src/lightning/pytorch/strategies/ddp.py:261

                ddp_comm_hook=self._ddp_comm_hook,
                ddp_comm_wrapper=self._ddp_comm_wrapper,
            )

    def _enable_model_averaging(self) -> None:
        log.debug(f"{self.__class__.__name__}: reinitializing optimizers with post localSGD")
        if self._model_averaging_period is None:
            raise ValueError(
                "Post-localSGD algorithm is used, but model averaging period is not provided to DDP strategy."
            )
        from torch.distributed.optim import DistributedOptimizer, PostLocalSGDOptimizer, ZeroRedundancyOptimizer

        for optimizer in self.optimizers:
            if isinstance(optimizer, LightningOptimizer):
                optimizer = optimizer._optimizer

            is_distributed_optimizer = isinstance(optimizer, DistributedOptimizer) if not _IS_WINDOWS else False
            if isinstance(optimizer, (ZeroRedundancyOptimizer, PostLocalSGDOptimizer)) or is_distributed_optimizer:
                raise ValueError(
                    f"Currently model averaging cannot work with a distributed optimizer of type "
                    f"{optimizer.__class__.__name__}."
                )

        assert self._ddp_comm_state is not None
        self._model_averager = torch.distributed.algorithms.model_averaging.averagers.PeriodicModelAverager(
            period=self._model_averaging_period, warmup_steps=self._ddp_comm_state.start_localSGD_iter
        )

    @override
    def optimizer_step(
        self,
        optimizer: Optimizer,
        closure: Callable[[], Any],
        model: Optional[Union["pl.LightningModule", Module]] = None,
        **kwargs: Any,
    ) -> Any:
        """Performs the actual optimizer step.

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Switch to plain torch.optim optimizers (e.g. Adam/SGD) for all optimizers when using post_local_sgd
  2. Or disable post_local_sgd / remove the comm wrapper if you need ZeRO or distributed optimizers
  3. If you intended just gradient communication compression, use a comm hook (e.g. SLOWFAST/FP16 compress) without post-localSGD instead

Example fix

# before
class LitModel(LightningModule):
    def configure_optimizers(self):
        return ZeroRedundancyOptimizer(self.parameters(), optimizer_class=torch.optim.Adam, lr=1e-3)
# + DDPStrategy(post_local_sgd=True, model_averaging_period=5)
// after
class LitModel(LightningModule):
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)
# + DDPStrategy(post_local_sgd=True, model_averaging_period=5)
Defensive patterns

Strategy: validation

Validate before calling

from torch.distributed.optim import ZeroRedundancyOptimizer, DistributedOptimizer
from torch.optim import Optimizer

bad = [o for o in model.configure_optimizers() if not isinstance(o, Optimizer)]
# ensure no ZeRO/Distributed/PostLocalSGD optimizer is returned when post_local_sgd=True

Type guard

def all_plain_optimizers(optimizers) -> bool:
    return all(isinstance(o, Optimizer) for o in optimizers)

Prevention

When it happens

Trigger: Combining DDPStrategy(post_local_sgd=True, model_averaging_period=...) with ZeroRedundancyOptimizer (ZeRO-style sharding) or torch.distributed.optim.DistributedOptimizer in your configure_optimizers.

Common situations: Trying to stack memory-saving ZeRO/ZeroRedundancy optimizations with post-localSGD communication compression; these algorithms both own parameter synchronization and conflict.

Related errors


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