google-deepmind / google-deepmind/acme
Examples won't run right away with GPU installation of JAX
- Dominant language
- Python
- Stars
- 4.1k
- Forks
- 553
- PR merge metrics
- No merged PRs in 30d
Description
Hi Acme developers,
I would like to point out some issues with the current examples. Right now, the instructions on the README install the CPU version of JAX by default. If the user installs a GPU version of JAX, then the examples won't run without adding an additional line that turns off GPU usage in tensorflow, the following line should work.
```python
# Adding this in the beginning to turn off GPU usage in tensorflow, which is the default installed by Acme.
tf.config.set_visible_devices([], "GPU")
```
It would be great if the examples can run with GPU version of JAX straight away. I suppose adding this line to the examples should fix the problem.
However, I would also like to mention another potential issue with tf and jax competing for GPU usage. When launchpad is used to launch distributed agents based on DistributedLayout, it currently requires quite a little bit of configurations to get right to ensure that only the JAX learner is using the GPU.
For example, I have to add the same line to the learner method in distributed layout to ensure that in the learner node tf is not using GPU
```python
def learner(
self,
random_key: networks_lib.PRNGKey,
replay: reverb.Client,
counter: counting.Counter,
):
"""The learning part of the agent."""
# TODO(yl) this is a hack to ensure that TF does not use GPU
# Find a way to do this outside layouts
tf.config.experimental.set_visible_devices([], "GPU")
...
```
I suspect a similar issue would be present if we are using GPUs in the actor since adding transitions to reverb would trigger GPU memory allocation for TensorFlow which would clash with JAX.
It would be great if there's a easy way to set these things up so that users would have to worry less about running the LP programs with GPUs.
Contributor guide
Assessment
This issue has not been assessed yet.