Xla: Retrieve information about TPU (memory, name, ...)

[patrickvonplaten]

❓ Questions and Help

Is it possible to retrieve information about the TPU in use? For example

• Memory capacity of device (or is it always 8 GB?)
• Current memory usage
• Name of TPU
• TPU version

[jysohn23]

While your job is running you can capture a profile of the job using the capture_tpu_profile (https://cloud.google.com/tpu/docs/cloud-tpu-tools) tool which should also give you an idea of the HBM usage, etc. Also our v2 TPUs have 8GB HBM and our v3 have 16GB HBM per core.

Also, to clarify, you want to know the name and version of the TPU while you're already training on it?

[dlibenzi]

We planned to add an XRT API to extract live memory info, but never got to it.
This might be the week 😄

[patrickvonplaten]

Thanks a lot for your answers!
Yes, having a XRT API would be super useful!

A feature I think that would be very useful would be something similar to PyTorchs max_memory_reseverd: https://pytorch.org/docs/stable/cuda.html#torch.cuda.max_memory_reserved so that one would know
the required memory for a given model on TPU and that could be called from python code

@jysohn23
Yeah, I guess the there isn't really a name for a TPU (similar to GPU TITAN RTX), so just the version v1, v2, v3 would be great :-)
By name I just meant the version, e.g. v2 TPU

[dlibenzi]

There is an in-flight internal XRT CL. This should be available somewhere next week.

[dlibenzi]

https://github.com/pytorch/xla/blob/c88725afc682f4aaa69cd69fee4cd723763f3b60/torch_xla/core/xla_model.py#L776

[patrickvonplaten]

That's awesome!

[dlibenzi]

I just noticed that this morning our nightly build chain failed, so this will be out tomorrow in nightly.