Specify Gradient Clipping Norm in Trainer
🚀 Feature
Allow specification of the gradient clipping norm_type, which by default is euclidean and fixed.
Motivation
We are using pytorch lightning to increase training performance in the standalone Federated Learning context (experimental setting). In this context the trained models diverge from their underlying data and get aggregated on the server side which leads to larger gradients in general. To preserve the direction of the gradient, but limit the magnitude per single dimension, we need to apply the inf norm.
Pitch
Add a parameter gradient_clipping_norm_type: float=2.0 to trainer. Pass the parameter to the _clip_gradients method. Changing the call from
_clip_gradients(optimizer, grad_clip_val)
to somewhat like
_clip_gradients(optimizer, grad_clip_val, grad_clip_norm_type)
Additional context
The impact is minimal and only effects the following function

  
    
      pytorch-lightning/pytorch_lightning/accelerators/accelerator.py
    
    
        Lines 119 to 128
      in
      861d699
    
  
  
    

        
          
           def clip_gradients(self, optimizer, clip_val=None): 
        

        
          
               # use the trainer's clip val if none passed 
        

        
          
               grad_clip_val = self.trainer.gradient_clip_val 
        

        
          
               if clip_val is not None: 
        

        
          
                   grad_clip_val = clip_val 
        

        
          
               grad_clip_val = float(grad_clip_val) 
        

        
          
            
        

        
          
               if grad_clip_val <= 0: 
        

        
          
                   return 
        

        
          
               self._clip_gradients(optimizer, grad_clip_val) 
        
    
  


Which is calling

  
    
      pytorch-lightning/pytorch_lightning/accelerators/accelerator.py
    
    
         Line 130
      in
      861d699
    
  
  
    

        
          
           def _clip_gradients(self, optimizer: Optimizer, grad_clip_val: Union[float, int], norm_type: float = 2.0):