Avoid wrapping LightningModule in *DataParallel overrides when not fitting
🚀 Feature
For distributed testing or prediction, we don't need to wrap the LightningModule inside of DistributedDataParallel or DataParallel for testing as there are no gradients we need to synchronize. We only need this during the fit stage when model training occurs
Motivation
This can reduce overhead with distributed inference in Lightning. We can also use torchscript modules or models without any trainable parameters purely for inference.
Pitch
We'd need the training type plugins to be aware of the Trainer state somehow. Then we could only apply the wrapper here in case the trainer is set to fit: 
  
    
      pytorch-lightning/pytorch_lightning/plugins/training_type/ddp.py
    
    
        Lines 249 to 256
      in
      80c5293
    
  
  
    

        
          
           def configure_ddp(self): 
        

        
          
               self.pre_configure_ddp() 
        

        
          
               self._model = DistributedDataParallel( 
        

        
          
                   LightningDistributedModule(self.model), 
        

        
          
                   device_ids=self.determine_ddp_device_ids(), 
        

        
          
                   **self._ddp_kwargs, 
        

        
          
               ) 
        

        
          
               self._register_ddp_hooks() 
        
    
  


Alternatives

Additional context