Xla: [Question] how to track TPU memory usage

[nosound2]

How can I access TPU (internal) memory utilization?

Thanks!

[dlibenzi]

There is nothing explicit.
But essentially, after a step barrier, all the live XLA tensors will be on device.
So the TPU memory utilization will be the amount of memory used to store those.
Note that the TPU memory is also used to store TPU programs, so the live tensors one does not provide the full picture.
Also, the layout of the tensors on the TPU can influence quite a bit the memory utilized.
We have also this, which prints a lot of data:

print(torch_xla._XLAC._xla_tensors_report(0, str(device)))

[ibeltagy]

what is a step barrier?

[dlibenzi]

This:

https://github.com/pytorch/xla/blob/4b8d8c0ada25572f2980920d9d4f4f8c42938ebe/torch_xla_py/data_parallel.py#L59

The one that turns the IR graphs accumulated on live tensors, into TPU device data.

[ibeltagy]

thanks. Another memory related question. Sometimes the training crashes and prints a long log that ends with
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
and I have to restart the TPU for this to stop happening.

• what is RunOptions?
• is this really an OOM error or something else?
• and if it is an OOM error, why does it go away after restarting the tpu?

[dlibenzi]

I think someone on our side forgot that TPUs are not only a Tensorflow thing :)
The RunOptions is a TF session run thing.
Yes, it is a OOM.
We should be getting the OOM tensor list, as we do many times. Maybe that is a path where it is not reported.
We will have to check.
Why does it go away?
Could be the device gets into a fragmentation state we cannot recover from.
We do have defrag code, but certain parts cannot be relocated.

[mruberry]

Closing because it appears the initial question was answered. Please keep asking questions!