Avoid configuring SyncBatchNorm when not fitting

[four4fish]

Proposed refactoring or deprecation

Conditionally configure syncbatchnorm in distributed plugins

Motivation

This issue is closely related to #6977
Carrying forward discussion from #9096
Related issue in PyTorch: pytorch/pytorch#48988

SyncBatchNorm doesn't sync stats when used under eval mode. We can conditionally check for whether to configure this if we're fitting vs validating/testing/predicting.

Pitch

Conditionally determine whether to configure syncbatchnorm in the module here: https://github.com/PyTorchLightning/pytorch-lightning/blob/a451997c4da89be3b1e4f7f79b52015bd32f2ea4/pytorch_lightning/plugins/training_type/ddp.py#L384-L387

Essentially rewrite this

        if self.sync_batchnorm:
            self.model = self.configure_sync_batchnorm(self.model)

        # skip wrapping the model if we are not fitting as no gradients need to be exchanged
        trainer_fn = self.lightning_module.trainer.state.fn
        if trainer_fn == TrainerFn.FITTING:
            self.configure_ddp()

as this

        # skip wrapping the model if we are not fitting as no gradients need to be exchanged
        trainer_fn = self.lightning_module.trainer.state.fn
        if trainer_fn == TrainerFn.FITTING:
            if self.sync_batchnorm:
                self.model = self.configure_sync_batchnorm(self.model)
            self.configure_ddp()

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning

• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[tchaton]

Hey @four4fish,

Do you see any situations where users might want to update their BatchNorm stats on the validation dataset in a distributed way. If not, I think it is a good proposal.

Best,
T.C