AlphaFold3

From Alliance Doc
Jump to navigation Jump to search
This page is a translated version of the page AlphaFold3 and the translation is 100% complete.
Other languages:

Nous décrivons ici l'utilisation de AlphaFold v3.0.

Le code source et la documentation se trouvent sur leur page GitHub. Toute publication qui divulgue des résultats découlant de l'utilisation de ce code source ou des paramètres du modèle doit citer le document AlphaFold3.

Versions disponibles

Sur nos grappes, AlphaFold3 est disponible sous forme de paquets préconstruits (wheels) Pour les lister, utilisez avail_wheels.

Question.png 
[name@server ~]$ avail_wheels alphafold3

AlphaFold2 est encore disponible (voir la documentation).

Créer un fichier des dépendances requises

1. Chargez les dépendances de AlphaFold3.

Question.png 
[name@server ~]$ module load StdEnv/2023 hmmer/3.4 rdkit/2024.03.5 python/3.12

2. Téléchargez le script d'exécution.

Question.png 
[name@server ~]$ wget https://raw.githubusercontent.com/google-deepmind/alphafold3/23e3d46d4ca126e8731e8c0cbb5673e9a848ceb5/run_alphafold.py

3. Créez et activez un environnement virtuel Python. 

[name@server ~]$ virtualenv --no-download ~/alphafold3_env
[name@server ~]$ source ~/alphafold3_env/bin/activate


4. Installez une version de AlphaFold3 ainsi que ses dépendances Python 

(alphafold3_env) [name@server ~] pip install --no-index --upgrade pip
(alphafold3_env) [name@server ~] pip install --no-index alphafold3==X.Y.Z

où X.Y.Z est la version spécifique, par exemple 3.0.0. N'entrez pas le numéro de la version si vous voulez installer la plus récente.

5. Compilez les données nécessaires.

Question.png 
(alphafold3_env) [name@server ~] build_data

Ceci crée des fichiers de données dans l'environnement virtuel.

6. Validez.

Question.png 
(alphafold3_env) [name@server ~] python run_alphafold.py --help

7. Gelez l'environnement et l'ensemble des requis.

Question.png 
(alphafold3_env) [name@server ~] pip freeze > ~/alphafold3-requirements.txt

8. Désactivez l'environnement.

Question.png 
(alphafold3_env) [name@server ~] deactivate

9. Nettoyez et supprimez l'environnement virtuel.

Question.png 
[name@server ~]$ rm -r ~/alphafold3_env

L'environnement virtuel sera plutôt créé dans votre tâche.

Modèle

Le modèle peut être obtenu de Google habituellement dans les 2 ou 3 jours ouvrables suivants. Voyez ici Obtaining Model Parameters.

Bases de données

AlphaFold3 nécessite un ensemble de bases de données.

Important : Les bases de données doivent résider dans le répertoire $SCRATCH.

1. Téléchargez le script de téléchargement.

Question.png 
[name@server ~]$ wget https://raw.githubusercontent.com/google-deepmind/alphafold3/refs/heads/main/fetch_databases.sh

2. Téléchargez les bases de données. 

[name@server ~]$ mkdir -p $SCRATCH/alphafold/dbs
[name@server ~]$ bash fetch_databases.sh $SCRATCH/alphafold/dbs


Exécution par étapes

Alphafold3 doit être exécuté par étapes, c'est-à-dire

  1. séparer le pipeline de données pour CPU seulement et le modèle d'inférence (qui demande un GPU) pour optimiser les coûts et l'utilisation des ressources;
  2. cacher les résultats de la recherche de MSA/modèle, pour ensuite réutiliser le JSON augmenté pour plusieurs différentes inférences ou pour des variations d'autres fonctionnalités (par exemple un ligand).

Pour des références, voir

1. Pipeline de données (CPU)

Modifiez le script suivant selon vos besoins.

File : alphafold3-data.sh

#!/bin/bash

#SBATCH --job-name=alphafold3-data
#SBATCH --account=def-someprof    # adjust this to match the accounting group you are using to submit jobs
#SBATCH --time=08:00:00           # adjust this to match the walltime of your job
#SBATCH --cpus-per-task=8         # a MAXIMUM of 8 core, AlphaFold has no benefit to use more
#SBATCH --mem=64G                 # adjust this according to the memory you need

# Load modules dependencies.
module load StdEnv/2023 hmmer/3.4 rdkit/2024.03.5 python/3.12

DOWNLOAD_DIR=$SCRATCH/alphafold/dbs    # set the appropriate path to your downloaded data
INPUT_DIR=$SCRATCH/alphafold/input     # set the appropriate path to your input data
OUTPUT_DIR=$SLURM_TMPDIR/alphafold/output   # set the appropriate path to your output data

# Generate your virtual environment in $SLURM_TMPDIR.
virtualenv --no-download $SLURM_TMPDIR/env
source $SLURM_TMPDIR/env/bin/activate

# Install AlphaFold and its dependencies.
pip install --no-index --upgrade pip
pip install --no-index --requirement ~/alphafold3-requirements.txt

# build data in $VIRTUAL_ENV
build_data

# https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md#compilation-time-workaround-with-xla-flags
export XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"

# Edit with the proper arguments and run your commands.
# run_alphafold.py --help
python run_alphafold.py \
    --db_dir=$DOWNLOAD_DIR \
    --input_dir=$INPUT_DIR \
    --output_dir=$OUTPUT_DIR \
    --jax_compilation_cache_dir=$HOME/.cache \
    --nhmmer_n_cpu=$SLURM_CPUS_PER_TASK \
    --jackhmmer_n_cpu=$SLURM_CPUS_PER_TASK \
    --norun_inference  # Run data stage

# copy back
mkdir $SCRATCH/alphafold/output
cp -vr $OUTPUT_DIR $SCRATCH/alphafold/output


2. Inférence de modèle

Modifiez le script suivant selon vos besoins.


Compatibilité

Alphafold3 prend en charge seulement la capacité de calcul des versions 8.0 et plus récentes, soit pour des A100s et plus performants.



File : alphafold3-inference.sh

#!/bin/bash

#SBATCH --job-name=alphafold3-inference
#SBATCH --account=def-someprof    # adjust this to match the accounting group you are using to submit jobs
#SBATCH --time=08:00:00           # adjust this to match the walltime of your job
#SBATCH --cpus-per-task=1         # AlphaFold has no benefit to use more for the inference stage
#SBATCH --gpus=a100:1             # Alphafold3 inference only runs on ONE A100 or greater.
#SBATCH --mem=20G                 # adjust this according to the memory you need

# Load modules dependencies.
module load StdEnv/2023 hmmer/3.4 rdkit/2024.03.5 python/3.12 cuda/12 cudnn/9.2

DOWNLOAD_DIR=$SCRATCH/alphafold/dbs    # set the appropriate path to your downloaded data
INPUT_DIR=$SCRATCH/alphafold/input     # set the appropriate path to your input data, following the data stage.
OUTPUT_DIR=$SCRATCH/alphafold/output   # set the appropriate path to your output data

# Generate your virtual environment in $SLURM_TMPDIR.
virtualenv --no-download $SLURM_TMPDIR/env
source $SLURM_TMPDIR/env/bin/activate

# Install AlphaFold and its dependencies.
pip install --no-index --upgrade pip
pip install --no-index --requirement ~/alphafold3-requirements.txt

# build data in $VIRTUAL_ENV
build_data

# https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md#compilation-time-workaround-with-xla-flags
export XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"

# https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md#gpu-memory
export XLA_PYTHON_CLIENT_PREALLOCATE=true
export XLA_CLIENT_MEM_FRACTION=0.95

# Edit with the proper arguments and run your commands.
# run_alphafold.py --help
python run_alphafold.py \
    --db_dir=$DOWNLOAD_DIR \
    --input_dir=$INPUT_DIR \
    --output_dir=$OUTPUT_DIR \
    --jax_compilation_cache_dir=$HOME/.cache \
    --norun_data_pipeline  # Run inference stage


3. Soumettre une tâche

Soumettez la tâche à l'ordonnanceur.

Tâches indépendantes

Question.png 
[name@server ~]$ sbatch alphafold3-data.sh

Attendez la fin et soumettez ensuite la deuxième étape.

Question.png 
[name@server ~]$ sbatch alphafold3-inference.sh

Tâches dépendantes

[name@server ~]$ jid1=$(sbatch alphafold3-data.sh)
[name@server ~]$ jid2=$(sbatch --dependency=afterok:$jid1 alphafold3-inference.sh)
[name@server ~]$ sq

Si la première étape échoue, vous devez annuler manuellement la deuxième étape.

Question.png 
[name@server ~]$ scancel -u $USER -n alphafold3-inference

Dépannage

Mémoire insuffisante (GPU)

Si vous voulez exécuter AlphaFold3 avec plus de 5120 jetons ou sur un GPU de mémoire moindre, (par exemple sur un A100 avec 40Go de mémoire), vous pouvez activer la fonctionnalité de mémoire unifiée.

Dans le script de soumission à l'étape d'inférence, ajoutez les variables d'environnement suivantes : 

export XLA_PYTHON_CLIENT_PREALLOCATE=false
export TF_FORCE_UNIFIED_MEMORY=true
export XLA_CLIENT_MEM_FRACTION=2.0  # 2 x 40GB = 80 GB


et ajuster en conséquence la quantité de mémoire allouée à la tâche, par exemple #SBATCH --mem=80G.

Retrieved from "https://docs.alliancecan.ca/mediawiki/index.php?title=AlphaFold3/fr&oldid=173812"