CPU memory leak when using PL on GCP TPUs

[Natithan]

Hi,

I've been stuck for quite some time on this.
I am training a VilBERT-like model, and because each training run takes quite a long time, I am running it on google cloud TPUs in the hope of speeding it up.

It doesn't work however, as during training the CPU memory slowly increases until it's killed by the VMs OOM-process-killer.
The TPU memory is not filling up however. (See image below).
I have 64GB CPU RAM available, and am using a v3-8 TPU.
As I didn't observe this increase in CPU memory while training on the GPUs in my lab, I think it might have to do with how I am using the TPUs.
I'm using pytorch-lightning to move my code to TPUs, hence why I'm hoping someone on this forum can help me :).

I also checked the memory usage when using one TPU core (rather than 8), and there too the memory increases per step (see image).
For that run, I also tracked memory usage using mprof. (see image)
What is odd, is that this gave me a really different graph however, one in which there is no leak apparently. I feel like this should give me some hint, but I'm not sure what :).

I've also tried to use memory_profiler to see the memory-changes line-by-line.
Because the program is multiprocess however, this gives quite a bit of seemingly unrelated memory fluctuation.

I did a bit of googling around, here are some related threads that didn't solve my problem:

• TPU Trainer memory leak and memory requirements huggingface/transformers#5509: this person didn't really fix their memory leak, but mitigated it by reducing how fast it leaked: they used a redis server to not load data in each subprocess.
• https://stackoverflow.com/questions/65107933/pytorch-model-training-cpu-memory-leak-issue in this thread, the issue was not detaching a tensor, causing it to keep the whole computation graph. As far as I can tell however, my code doesn't keep any tensors between steps that could 'build up' in this way.
• Memory (CPU and GPU) leaks during the 1st epoch #1510 this issue seems related as well, but it should be solved in pytorch-lightning according to the thread.

If anyone had a similar problem, and could point me in the right direction (or even how to best find the cause), that would be great!

[kaushikb11]

Hi @Natithan, thanks for posting this, will give it a look early next week. Did you give the new TPU VMs a try? And let me know if the problem still persists?

[kaushikb11]

For sure. @Natithan Could you confirmed the memory leak isn't an issue while training on GPUs or CPUs?

[Natithan]

Hey @kaushikb11 ,
I'm currently just running it on GPUs and will be busy trying to finish my paper for the next month or so.
I'll check the memory usage on GPU after I'm done with that :)

[Natithan]

Just to follow up on this, I won't be investigating this more in the foreseeable future as it's no longer a blocking issue for me (moved on to next project). Thanks for the help in any case!

[tchaton]

Dear @Natithan,

Any chance you could just check this out using the Lightning master ? We resolved several memory leaks issues and would appreciated if you could confirm it still exists.

Best,
T.C

[Natithan]

Hi @tchaton , it will be tricky to do so, as I already deleted the VMs on GCP: I ended up doing everything on GPUs locally. Sorry for that!