API for clearing XLA:CUDA cached memory.

[ysiraichi]

PyTorch/XLA currently has no means to clear cached memory, i.e. something similar to torch.cuda.empty_cache(). This is relevant for benchmarking a model on both PyTorch CUDA and XLA:CUDA.

For comparison, using PyTorch CUDA, we have:

# Memory usage rises to 4G
>>> a = torch.rand(1024, 1024, 1024, device="cuda")
>>> del a
>>> gc.collect()
# Memory goes back to ~200M
>>> torch.cuda.empty_cache()

In contrast, using XLA:CUDA with PJRT_ALLOCATOR_PREALLOCATE=false, we have:

# Memory usage rises to 4G
>>> a = torch.rand(1024, 1024, 1024, device=xm.xla_device())
>>> a.cpu()
>>> del a
>>> gc.collect()
# Memory usage never goes back...

Not only that, but it would be nice to have an API call for deleting the pre-allocated memory.

cc @miladm @JackCaoG

[miladm]

Thank you @ysiraichi!
do you plan to work on this issue?

[ysiraichi]

It's actually not something that's blocking me, so I wasn't going to. But, if you think this is worth looking into, I could do it.