how to set a specific gpu to use in tapir code

[google-deepmind]

how to set a specific gpu to use in tapir code

[CHYjeremy]

thank you very much @m43

[m43]

To select a specific GPU, you can specify the device when using jax.jit. Here's an example:

gpus = jax.devices('gpu')
device = gpus[1] # If you want to use "cuda:1"
model_apply = jax.jit(model.apply, device=device)

To confirm that you are utilizing the GPU, execute nvidia-smi in the terminal. It will display GPU usage statistics and the associated processes. For a real-time monitoring tool, consider using nvtop which provides a "top-like" interface.

As for avoiding recompilation, JAX compiles the function and caches the XLA code when a jax.jit function is called for the first time. From what I understand, subsequent calls use the cached version unless there are changes in arguments labeled as static in the call to jax.jit.

Lastly, to prevent JAX from preallocating substantial GPU memory, consider setting this environment variable:

export XLA_PYTHON_CLIENT_PREALLOCATE=false

Hope this helps!

[CHYjeremy]

Thank you very much for the amazing answer.

but let me follow up with your suggestion on setting environment variable, may i ask, to what extent will it hurt the performance(speed) of the algorithm when you prevent jax from preallocating much gpu memory?

thanks.

[m43]

Allocating memory on-the-fly might not have a significant impact on speed compared to preallocating, though this isn't explicitly detailed in the documentation. Disabling preallocation can reduce overall memory usage, especially if JAX doesn't fully utilize the preallocated memory. This becomes particularly relevant if TensorFlow or PyTorch models run concurrently on the same GPU. However, without preallocation, there's an increased risk of memory fragmentation, which could lead to OOM issues for JAX programs that need to consume most of the GPU memory. You can find more details on this topic here. If you're not facing any OOM issues, the default memory settings might be well-suited for your needs.