Support for sharded optimizers when dumping checkpoints outside of the DDP sharded training type plugin
🐛 Bug
Using fairscale distributed optimizers without DDP sharded leads to crashes or inconsistent stae in the trainer when checkpointing. This will also occur with PyTorch's latest prototype version of the ZeroRedundancy optimizer. @SeanNaren

Imagine this scenario:

Someone wraps their optimizer with Fairscale OSS inside their lightning module. The user does not use the DDP sharded plugin.
At checkpoint time, when the trainer dumps the checkpoint dict, it looks up the optimizer state
The optimizer state goes through the training type plugin
The training type plugin calls optimizer.state_dict() - For fairscale/pytorch distributed optimizers, we need to consolidate the state dict on a rank. Afterwards, we should look up the state dict only from that rank.

One could add a callback which implements on_save_checkpoint to call consolidate_state_dict() on the optimizer across all ranks. However, the trainer calls state_dict on all ranks, leading to the exception here: https://github.com/facebookresearch/fairscale/blob/1204c7cf54ec301d46a0d3f3fd703da6b306f8f5/fairscale/optim/oss.py#L354-L358
This error occurs only on non-zero ranks. As a result, the error in checkpointing is compounded by the comment here around the exception handling logic for training and its interaction with checkpointing: #6343 (comment)
Proposal to fix:

We upstream this optimizer state on all training type plugins: 
  
    
      pytorch-lightning/pytorch_lightning/plugins/training_type/sharded.py
    
    
        Lines 56 to 68
      in
      966184a
    
  
  
    

        
          
               def optimizer_state(self, optimizer: 'OSS') -> Optional[dict]: 
        

        
          
                   if is_lightning_optimizer(optimizer): 
        

        
          
                       optimizer = optimizer._optimizer 
        

        
          
                   optimizer.consolidate_state_dict() 
        

        
          
                   return self._optim_state_dict(optimizer) 
        

        
          
            
        

        
          
               @rank_zero_only 
        

        
          
               def _optim_state_dict(self, optimizer): 
        

        
          
                   """ 
        

        
          
                   Retrieves state dict only on rank 0, which contains the entire optimizer state after calling 
        

        
          
                   :meth:`consolidate_state_dict`. 
        

        
          
                   """ 
        

        
          
                   return optimizer.state_dict() 
        
    
  


This way, until PyTorch has a consistent interface for distributed/sharded optimizers, we check if the optimizer is fairscale OSS/PyTorch ZeroRedundancy. If so, we call consolidate_state_dict on all ranks, and then fetch the optimizer state only from rank 0.
plugins for FSDP or fancier things can continue to override this to handle the optimizer states however they need to.
If the optimizer is neither of these, we use the existing flow

cc @awaelchli @ananthsub @ninginthecloud @rohitgr7 @tchaton @akihironitta @blefaudeux