error running with a nvidia rtx pro 4500/5000

[SaberLi110]

Hi, everyone

When I was using the Nvidia RTX PRO 4500/5000, I modified the Dockerfile, and the relevant information is as follows:

1. using the nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04 image
2. using jax[cuda12]==0.6.0 and 0.6.1 and 0.8
3. upgrading to dm-haiku==0.0.14
4. installing tensorflow==2.12.0 and latest

I can successfully compile it, and AlphaFold runs, but only using the CPU; it cannot utilize the GPU (nvidia-smi shows no processes). However, if I enter the AlphaFold container image, JAX can detect the GPU.

any suggestions?

[sahilprajapti-ai]

Hi @SaberLi110
error solve
--> Correct, Working Solution (Proven)
--> 1. Use CUDA 12.1 (NOT 12.8)

Change your base image:

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

--> CUDA 12.8 will not work reliably with AlphaFold as of now.

--> 2. Pin correct versions (CRITICAL)
pip install
jax[cuda12]==0.4.26
jaxlib==0.4.26
dm-haiku==0.0.14
tensorflow==2.12.0

--> These versions are known-good for AlphaFold

--> JAX 0.6+ and 0.8 are NOT AlphaFold-compatible yet

--> 3. Force XLA to use GPU

Add before running AlphaFold:

export XLA_FLAGS="--xla_gpu_cuda_data_dir=/usr/local/cuda"
export TF_FORCE_GPU_ALLOW_GROWTH=true
export CUDA_VISIBLE_DEVICES=0

Without this, AlphaFold often defaults to CPU silently

--> 4. Run Docker with proper GPU flags
docker run --gpus all \ --runtime=nvidia
-e NVIDIA_VISIBLE_DEVICES=all
-e NVIDIA_DRIVER_CAPABILITIES=compute,utility
alphafold_image

Verify inside container:

nvidia-smi

--> 5. Verify before running AlphaFold

Inside container:

python3 - <<EOF
import jax
print(jax.devices())
EOF

Expected output:

[GpuDevice(id=0)]

Then test:

python3 - <<EOF
import jax.numpy as jnp
x = jnp.ones((4096,4096))
print(x.device())
EOF

[ionfwsrijan]

Hi there,

I’m Srijan Jaiswal, and I’m eager to contribute to this project as part of my open-source contribution. I have reviewed this issue and feel confident that I can provide a solution since I have experience working with similar technologies in my past projects. Could you please assign this issue to me? I’d love to work on it and contribute to the project.

Best regards,
Srijan Jaiswal