flax.metrics.tensorboard will take most GPU memory

[twni2016]

Hi,

I found an issue that when I use from flax.metrics.tensorboard import SummaryWriter, once it starts logging data, it will take most of GPU memory (~97%). Even if I set

import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = 'false'

at the first lines of code.

I did not find this issue when I use tensorboardX for logging.

Update: the minimal codes that can reproduce the issue:

from flax.metrics.tensorboard import SummaryWriter

writer = SummaryWriter("tmp")
writer.scalar("key", 1.0, 100)
import ipdb; ipdb.set_trace() # GPU memory will be almost occupied at this point
writer.flush()

My installed packages:

flax=0.5.2=pypi_0
jax=0.3.14=pypi_0
jaxlib=0.3.14+cuda11.cudnn82=pypi_0

Do you know the reason?

[marcvanzee]

This should not happen, but it is hard to debug without any code. Could you provide us with a minimal example in a Colab where the problem occurs? Generally you should follow the best practices from our examples: use jax.device_get to transfer the metrics from the device to the host and use flush() to flush the summary writer.

[twni2016]

I tried jax.device_get but it won't help. In fact, I can reproduce the issue with the code, updated above. It is indeed weird.

[marcvanzee]

Hi @twni2016, apologies for the delay! I tried reproducing this on a public Colab with GPU (http://colab.research.google.com) but I could not reproduce it. Would you mind creating a Colab where the GPU memory problem is reproduced? Perhaps it has to do with you local setup?

[twni2016]

I found this issue in my local setup but did not test in Colab. I will close this issue as it might be caused from my side.