Gpu memory leak with self.log on_epoch=True
pl 1.0.5
Using new logging api I want to log a metric in LightningModule
self.log(";;;;;;;;;;;;;;;;;;;", 1, on_step=False, on_epoch=True)

This is a dummy example but it is sufficient to add to LightningModule's training_step to cause a memory leak on gpu.
What could go wrong? We want to log a metric which is not even a cuda tensor. How could it lead to a gpu memory leak?
Well thanks to the magic of metric epoch aggregation stuff
Let's dig in and take a look at here

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_loop.py
    
    
        Lines 550 to 569
      in
      b3db197
    
  
  
    

        
          
           # ------------------------------------ 
        

        
          
           # TRAINING_STEP + TRAINING_STEP_END 
        

        
          
           # ------------------------------------ 
        

        
          
           batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx) 
        

        
          
            
        

        
          
           # when returning -1 from train_step, we end epoch early 
        

        
          
           if batch_output.signal == -1: 
        

        
          
               break 
        

        
          
            
        

        
          
           # only track outputs when user implements training_epoch_end 
        

        
          
           # otherwise we will build up unnecessary memory 
        

        
          
           epoch_end_outputs = self.process_train_step_outputs( 
        

        
          
               batch_output.training_step_output_for_epoch_end, 
        

        
          
               self.early_stopping_accumulator, 
        

        
          
               self.checkpoint_accumulator, 
        

        
          
           ) 
        

        
          
            
        

        
          
           # hook 
        

        
          
           # TODO: add outputs to batches 
        

        
          
           self.on_train_batch_end(epoch_output, epoch_end_outputs, batch, batch_idx, dataloader_idx) 
        
    
  


Here we run batch, convert batch_output to epoch_end_outputs if on_epoch was set and append epoch_end_outputs to epoch_output inside on_train_batch_end
epoch_output is defined here

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_loop.py
    
    
         Line 540
      in
      b3db197
    
  
  
    

        
          
           epoch_output = [[] for _ in range(self.num_optimizers)] 
        
    
  


Everything seems normal, but there is a problem inside batch_output there is a surprise - loss value stored on gpu.

I think you can guess by now what could go wrong if we store a lot of separate cuda tensors in a long long epoch_output
Yeah the gpu memory is going to end and you'll get a famous
RuntimeError: CUDA out of memory. Tried to allocate 114.00 MiB (GPU 1; 10.92 GiB total capacity; 9.39 GiB already allocated; 27.38 MiB free; 10.24 GiB reserved in total by PyTorch)

Where is the loss appended to output? Here

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_loop.py
    
    
        Lines 396 to 427
      in
      b3db197
    
  
  
    

        
          
           def _process_training_step_output_1_0(self, training_step_output, split_batch): 
        

        
          
               result = self.trainer.get_model()._results 
        

        
          
            
        

        
          
               loss = None 
        

        
          
               hiddens = None 
        

        
          
            
        

        
          
               # handle dict return 
        

        
          
               if isinstance(training_step_output, dict): 
        

        
          
                   loss = training_step_output.pop("loss", None) 
        

        
          
                   hiddens = training_step_output.pop("hiddens", None) 
        

        
          
                   result["extra"] = training_step_output 
        

        
          
            
        

        
          
               # handle scalar return 
        

        
          
               elif isinstance(training_step_output, torch.Tensor): 
        

        
          
                   loss = training_step_output 
        

        
          
                   result["extra"] = {} 
        

        
          
            
        

        
          
               # map to results under the hood 
        

        
          
               result.minimize = loss 
        

        
          
               result.hiddens = hiddens 
        

        
          
            
        

        
          
               # track batch for manual reduction with result 
        

        
          
               result.track_batch_size(len(split_batch)) 
        

        
          
            
        

        
          
               # track metrics without grads for epoch reduction 
        

        
          
               training_step_output_for_epoch_end = copy(result) 
        

        
          
               training_step_output_for_epoch_end.detach() 
        

        
          
            
        

        
          
               # what flows back into the system 
        

        
          
               training_step_output = result 
        

        
          
            
        

        
          
               return training_step_output_for_epoch_end, training_step_output 
        
    
  


In the first line we get a pretty result without the loss in it, and in line 414 the loss get appended and we start our memory leak chain of events
How is it affecting the training? It can lead to error only on the first epoch of training. If you've got enough memory to hold a list of gpu losses during the 1st epoch there won't be any exceptions as subsequent epochs will have the same list of losses, if not you'll get it somewhere in the middle of 1st epoch. And of course the more steps you have in an epoch the more memory this list of gpu losses will require as one loss is stored per step
Here is the comparison for my task. My gpu could hold 2k steps before memory error
With self.log

Without self.log

You can see how there is a rapid growth in the first minute in both as the model is loaded and feeded the 1st batch.
The difference is in subsequent minutes where in the former case the list of losses eats 7gb of gpu memory and leads to crash, and in the latter nothing happens and training goes on
Pretty cool how one self.log could eat 2 times more gpu memory more than actual training process