JAX Synergistic Memory Inspector
jax-smi is a tool for real-time inspection of the memory usage of a JAX process. It is similar to nvidia-smi for GPU, but works on multiple platforms including CPU, GPU and TPU.
On TPU platforms, jax-smi is the only way to monitor TPU memory usage. On GPU platforms, jax-smi is also preferable to nvidia-smi. The latter is unable to report real-time memory usage of JAX processes, as JAX always pre-allocates 90% of the GPU memory by default.
This project is supported by Cloud TPUs from Google's TPU Research Cloud (TRC).
Installation
Install go. On Ubuntu, this is usually done by:
sudo apt-get install golangIf you followed tpu-starter to set up the TPU environment, go should be already installed.
Then install jax-smi with:
pip install jax-smiUsage
In your JAX script:
from jax_smi import initialise_tracking
initialise_tracking()
# some computation...Open a shell and run:
jax-smiApproach
Save the memory profile to /dev/shm/memory.prof in a separate thread every 1 second using jax.profiler.save_device_memory_profile().
Inspect the memory profile with go tool pprof -tags /dev/shm/memory.prof.
See https://twitter.com/ayaka14732/status/1565013139594551296 for more details.
Limitations
Tracing can only be performed by one process at a time. If tracing is performed by multiple JAX processes, they will write the memory profiles to the same file, which will lead to conflicts.
The jax-smi command line tool cannot detect if a memory profile file is out of date. Therefore, even if no JAX process is running, the tool will still read the outdated memory profile and report outdated memory usage information.