Enabling XLA in tensorflow 2.16 causes memory leaks about tensorflow HOT 9 OPEN

Hi @Di-Is ,

I have replicated the reported memory leak with XLA. Attached gist for reference.

from tensorflow.

This document from NVIDIA on tweaking environment variables for XLA memory could help. We're yet to test this out though and will try to keep this thread updated on any findings:
https://docs.nvidia.com/deeplearning/frameworks/tensorflow-user-guide/index.html#xla-best-practices

Btw, the formation of clusters / Operation fusing / compilation using XLA increases the memory for each different shape of input that the framework comes across. If we keep the shape constant, we've found that after X amount of time, the memory growth stabilizes even with TF 2.15.

from tensorflow.

@NBCBM
Please do not post low-quality LLM-generated text.

from tensorflow.

Hi @Di-Is ,
I tried to run your code on Colab using TF v2.16 and faced the same issue. Please find the gist here for reference.

Thank you!

from tensorflow.

@Venkat6871
Thank you for your reply!
I am using tensorflow compiled for Nvidia GPU(tensorflow[and-cuda] 2.16.1).
It seems there is also a difference in the Python version you confirmed.

from tensorflow.

I observed memory leaks in Google Colab with T4 GPU.

from tensorflow.

@Venkat6871, do you have any progress on this issue.

I'm also affected because of this issue, I tried will all tf versions from 2.11 up to 2.16 and it seems like this happens since tf 2.12. It seems this was not happening for tf 2.11

thank you

from tensorflow.

I face the exact problem. Using Ubuntu 20.04, NVIDIA RTX3060 & python=3.11

@Venkat6871, do you have any progress on this issue?

from tensorflow.

I'm also encountering the exact same problem using Ubuntu 20.04 NVIDIA GTX1070 and python 3.10. Would be great to get this fixed since the memory leakage can get pretty astronomical if left unchecked. I've had it grow up to 12gB until it crashed my kernel.

from tensorflow.