How to free the gpu memory of tensor list obtained by all_gather_object api?

distributed
#1

Following the example:

>>> # Note: Process group initialization omitted on each rank.
>>> import torch.distributed as dist
>>> # Assumes world_size of 3.
>>> gather_objects = ["foo", 12, {1: 2}] # any picklable object
>>> output = [None for _ in gather_objects]
>>> dist.all_gather_object(output, gather_objects[dist.get_rank()])
>>> output
['foo', 12, {1: 2}]

I used similar way to gather tensors into an output list during the training. These tensors occupied to much gpu memory and made CUDA OOM in the next steps. I tried to use del and torch.cuda.empty_cache(), but the memory is not released. Does it has another way to free the gpu memory?

#2

@roy_zhang could you please show how did you detect that memory was not released?

#3

I just watch the nvidia-smi. But, this problem is solved, I use all_gather in a complex scenario, the cuda tensor are not actually transfer to the target gpu even the target process could get all tensors, I guess it should be mapping? I am sure that each process creates context in all gpus making the gpu memory increasing. I tried to move tensors on cpu in gather_objects, then it is running…