5746 tokens job: RESOURCE_EXHAUSTED (H100 80G, despite unified memory, jax flag and increase bucket size)

[GitCeliniHub]

Hi @Augustin-Zidek

It was lovely and very helpful to meet you guys on Thursday :)

As you suggested and in order to succeed with my 5746 token jobs, I threw all the additional flags at it (with the help of Claude, as this all way beyond my skills and knowledge), and I still run into ressource_exhausted.

#!/bin/bash

#SBATCH --job-name=AF3-4GPU                 
#SBATCH --mail-type=END,FAIL                
#SBATCH --mail-user=xxx@crick.ac.uk     
#SBATCH --partition=gh100
#SBATCH --reservation=h100
#SBATCH --nodes=1
#SBATCH --ntasks=1                          
#SBATCH --gres=gpu:4
#SBATCH --cpus-per-task=100                 
#SBATCH --mem=0                         
#SBATCH --time=72:00:00                     
#SBATCH --output=sbatch%j.log   

ml purge
ml Singularity

# Create JAX cache directory
mkdir -p /xxx/jax_cache
chmod 777 /xxx/jax_cache

# Create temporary directory for model_config
TEMP_DIR=$(mktemp -d)
chmod 777 $TEMP_DIR

# Create model_config.py with documented sharding strategy
cat > $TEMP_DIR/model_config.py << 'EOL'
from typing import Sequence
from typing_extensions import TypeAlias

_Shape2DType: TypeAlias = tuple[int | None, int | None]

pair_transition_shard_spec: Sequence[_Shape2DType] = (
    (2048, None),
    (3072, 1024),
    (None, 512),
)
EOL

# Settings from documentation
export XLA_PYTHON_CLIENT_PREALLOCATE=false
export TF_FORCE_UNIFIED_MEMORY=true
export XLA_CLIENT_MEM_FRACTION=3.2
export XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"

singularity exec \
    --nv \
    --bind /xxx/:/root/af_input \
    --bind /xxx/:/root/af_output \
    --bind /flask/reference/Alphafold3_dataset/model_parameters:/root/models \
    --bind /flask/reference/Alphafold3_dataset/datasets:/root/public_databases \
    --bind /xxx/jax_cache:/root/jax_cache \
    --bind $TEMP_DIR/model_config.py:/app/alphafold/alphafold3/model/model_config.py \
    /flask/apps/containers/Alphafold/3.0.0/alphafold3.sif \
    python /app/alphafold/run_alphafold.py \
    --json_path=/root/af_input/fold_input.json \
    --model_dir=/root/models \
    --db_dir=/root/public_databases \
    --output_dir=/root/af_output \
    --buckets 5746 \
    --jax_compilation_cache_dir=/root/jax_cache

# Cleanup temporary directory
rm -rf $TEMP_DIR

It runs for a few hours and returns

Running model inference for seed 1...
Traceback (most recent call last):
  File "/app/alphafold/run_alphafold.py", line 699, in <module>
    app.run(main)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 684, in main
    process_fold_input(
  File "/app/alphafold/run_alphafold.py", line 556, in process_fold_input
    all_inference_results = predict_structure(
                            ^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 373, in predict_structure
    result = model_runner.run_inference(example, rng_key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 311, in run_inference
    result = self._model(rng_key, featurised_example)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 84554026840 bytes.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Any insights?
Cheers!
Celine

[Augustin-Zidek]

Hello Celine, great talking to you as well!

I think the problem is that the environment variables you are setting are not being passed to the Singularity container. Could you try adding the following flags:

--env 'XLA_PYTHON_CLIENT_PREALLOCATE=false'
--env 'TF_FORCE_UNIFIED_MEMORY=true'
--env 'XLA_CLIENT_MEM_FRACTION=3.2'
--env 'XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"'

in the singularity exec command just below the --bind flags? I.e. this:

singularity exec \
    --nv \
    --bind /xxx/:/root/af_input \
    --bind /xxx/:/root/af_output \
    --bind /flask/reference/Alphafold3_dataset/model_parameters:/root/models \
    --bind /flask/reference/Alphafold3_dataset/datasets:/root/public_databases \
    --bind /xxx/jax_cache:/root/jax_cache \
    --bind $TEMP_DIR/model_config.py:/app/alphafold/alphafold3/model/model_config.py \
    --env 'XLA_PYTHON_CLIENT_PREALLOCATE=false' \
    --env 'TF_FORCE_UNIFIED_MEMORY=true' \
    --env 'XLA_CLIENT_MEM_FRACTION=3.2' \
    --env 'XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"' \
    /flask/apps/containers/Alphafold/3.0.0/alphafold3.sif \
    python /app/alphafold/run_alphafold.py \
    --json_path=/root/af_input/fold_input.json \
    --model_dir=/root/models \
    --db_dir=/root/public_databases \
    --output_dir=/root/af_output \
    --buckets 5746 \
    --jax_compilation_cache_dir=/root/jax_cache

See https://docs.sylabs.io/guides/latest/user-guide/environment_and_metadata.html#environment-from-the-host for more details.

Moreover, what is the host RAM on the machine you are running this on?

[GitCeliniHub]

Thanks, I'll try this :)
My understanding is that we have 2TB/node; 4GPU/node.

[Augustin-Zidek]

My understanding is that we have 2TB/node; 4GPU/node.

Perfect, then once we get this working, memory should not be an issue. :)

[GitCeliniHub]

Something has definitely changed!
It stopped at the same step, but with a different error.

Running model inference for seed 1...
Traceback (most recent call last):
  File "/app/alphafold/run_alphafold.py", line 699, in <module>
    app.run(main)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 684, in main
    process_fold_input(
  File "/app/alphafold/run_alphafold.py", line 556, in process_fold_input
    all_inference_results = predict_structure(
                            ^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 373, in predict_structure
    result = model_runner.run_inference(example, rng_key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 311, in run_inference
    result = self._model(rng_key, featurised_example)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the function failed to launch on the GPU
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[Augustin-Zidek]

Hmm, I wonder if all environment variables have to be provided in one --env flag. Could you try this?

singularity exec \
    --nv \
    --bind /xxx/:/root/af_input \
    --bind /xxx/:/root/af_output \
    --bind /flask/reference/Alphafold3_dataset/model_parameters:/root/models \
    --bind /flask/reference/Alphafold3_dataset/datasets:/root/public_databases \
    --bind /xxx/jax_cache:/root/jax_cache \
    --bind $TEMP_DIR/model_config.py:/app/alphafold/alphafold3/model/model_config.py \
    --env 'XLA_PYTHON_CLIENT_PREALLOCATE=false,TF_FORCE_UNIFIED_MEMORY=true,XLA_CLIENT_MEM_FRACTION=3.2,XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"' \
    /flask/apps/containers/Alphafold/3.0.0/alphafold3.sif \
    python /app/alphafold/run_alphafold.py \
    --json_path=/root/af_input/fold_input.json \
    --model_dir=/root/models \
    --db_dir=/root/public_databases \
    --output_dir=/root/af_output \
    --buckets 5746 \
    --jax_compilation_cache_dir=/root/jax_cache

[GitCeliniHub]

This is now back to ressource exhausted.

Running model inference for seed 1...
Traceback (most recent call last):
  File "/app/alphafold/run_alphafold.py", line 699, in <module>
    app.run(main)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/alphafold3_venv/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 684, in main
    process_fold_input(
  File "/app/alphafold/run_alphafold.py", line 556, in process_fold_input
    all_inference_results = predict_structure(
                            ^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 373, in predict_structure
    result = model_runner.run_inference(example, rng_key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/app/alphafold/run_alphafold.py", line 311, in run_inference
    result = self._model(rng_key, featurised_example)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 80.08GiB (85989088472B) on device ordinal 0
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Maybe I could ask our HPC team if this bit "jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the function failed to launch on the GPU" is coming from a specific setting on our cluster (?)

[Augustin-Zidek]

Yes, worth checking with them, this looks more infrastructure specific. What CUDA and NVIDIA driver version are you running with?

[GitCeliniHub]

NVIDIA-SMI 560.28.03 Driver Version: 560.28.03 CUDA Version: 12.6

[Augustin-Zidek]

One more thing to try is removing the custom pair_transition_shard_spec. I.e. removing the

cat > $TEMP_DIR/model_config.py << 'EOL'
from typing import Sequence
from typing_extensions import TypeAlias

_Shape2DType: TypeAlias = tuple[int | None, int | None]

pair_transition_shard_spec: Sequence[_Shape2DType] = (
    (2048, None),
    (3072, 1024),
    (None, 512),
)
EOL

and the associated:

--bind $TEMP_DIR/model_config.py:/app/alphafold/alphafold3/model/model_config.py \

[Augustin-Zidek]

I am also seeing that you are still running with v3.0.0 - we released v3.0.1 in the meanwhile where we fixed a bunch of issues, it might be worth updating to that and trying whether that addresses the issue you are seeing.

[GitCeliniHub]

I asked HPC to install 3.0.1; in the emantime I tried the lighter script without the pair_transition_shard_spec and associated lines and still run into "jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the function failed to launch on the GPU"

[Augustin-Zidek]

A few questions:

1. Are you able to fold shorter proteins with the unified memory flags set, or do you still get the XlaRuntimeError?
2. Are you the only user running on machine with the GPU?
3. What operating system runs on the machine with the GPU?

[GitCeliniHub]

I asked our HPC team to install 3.0.1 on our cluster; they have, but I sill run into "jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the function failed to launch on the GPU"

1. I have not tried; I will :)
2. All the AF3 users are running on our H100 cluster; none of us have issues running smaller jobs
3. My understanding of HPC is quite poor, but I believe the information is here:
   NAME="Rocky Linux"
   VERSION="8.7 (Green Obsidian)"
   ID="rocky"
   ID_LIKE="rhel centos fedora"
   VERSION_ID="8.7"
   PLATFORM_ID="platform:el8"
   PRETTY_NAME="Rocky Linux 8.7 (Green Obsidian)"
   ANSI_COLOR="0;32"
   LOGO="fedora-logo-icon"
   CPE_NAME="cpe:/o:rocky:rocky:8:GA"
   HOME_URL="https://rockylinux.org/"
   BUG_REPORT_URL="https://bugs.rockylinux.org/"
   ROCKY_SUPPORT_PRODUCT="Rocky-Linux-8"
   ROCKY_SUPPORT_PRODUCT_VERSION="8.7"
   REDHAT_SUPPORT_PRODUCT="Rocky Linux"
   REDHAT_SUPPORT_PRODUCT_VERSION="8.7"