What is the best practice to track TPU memory usage?

[ayaka14732]

For GPU, it is common to use commands like watch -n2 nvidia-smi to easily monitor the GPU memory usage. For TPU, however, it seems that there is no such command.

Using the profiling tool seems to be the only way to achieve this. However, it seems that JAX support for the profiling tool is relatively poorer than Tensorflow. For example, there is no data for profiles captured by JAX jax.profiler.start_trace when viewing the tensorflow_stats and tf_data_bottleneck_analysis on Tensorboard.

Therefore, I would like to know what is the best practice to track TPU memory usage?

[ayaka14732]

See https://github.com/ayaka14732/jax-smi

[sagelywizard]

There's also a new CLI for checking TPU utilization metrics, which isn't specific to JAX, called tpu-info. You can install it via pip with

pip install git+https://github.com/google/cloud-accelerator-diagnostics/#subdirectory=tpu_info

[ayaka14732]

Please make it an independent package! It's complicated for users to depend on a subdirectory in requirements.txt

[steveepreston]

@sagelywizard Thank you for sharing this.

Usage Example

[steveepreston]

@ayaka14732 It' accessible through pypi

pip install tpu-info

https://pypi.org/project/tpu-info/