Pytorch transformers memory leak in inference

Viewed 1216

I am using Pytorch and a pretrained model from the transformers library. However, while finetuning it runs out of GPU memory very quickly and I wonder why.

I've found out that there is a memory leak in forward pass.

show_memory('before call')
outputs = model(inputs, masked_lm_labels=labels) if args.mlm else model(inputs, labels=labels)
show_memory('after call')

del outputs
gc.collect()
model.zero_grad()
torch.cuda.empty_cache()
show_memory('after clearing')

Here show_memory function is defined as following:

def show_memory(text=''):
  t = torch.cuda.get_device_properties(0).total_memory
  c = torch.cuda.memory_cached(0)
  a = torch.cuda.memory_allocated(0)
  f = c-a  # free inside cache
  print(f'\n\n{text}\nTotal: {t}\nCached: {c} \nAllocated: {a} \nFree in cache: {f}\n\n')

And the output is

before call 
Total: 17071734784  
Cached: 3441426432  
Allocated: 3324039680  
Free in cache: 117386752


after call 
Total: 17071734784 
Cached: 6830424064  
Allocated: 6720267264  
Free in cache: 110156800


after clearing 
Total: 17071734784 
Cached: 6109003776  
Allocated: 4876882944  
Free in cache: 1232120832

So despite I clear everything, there are still almost two gigabytes more data in memory than there was originally. I wonder why it happens and if there is a way to inspect such cases in detail.

0 Answers
Related