Selecting GPU(s) for Prediction

[rjrich]

I am running Protenix on a local Linux machine in a conda environment. My machine has 2 x Nvidia A6000 GPUs (CUDA visible devices 0 and 1). If I wanted to select GPU1 for Protenix predictions, would I just put "CUDA_VISIBLE_DEVICES=1" on the command line in front of "protenix predict ..."? In addition, is it possible to distribute the calculations over both GPUs? If so, how? Thank you.

[zhangyuxuann]

1. Put "CUDA_VISIBLE_DEVICES=1" on the command line in front of "protenix predict ..." will work.
2. For protenix predict ..., we do not support distribute the calculations over GPUs directly. But for inference_demo.sh we support such features by torchrun

export LAYERNORM_TYPE=fast_layernorm
export USE_DEEPSPEED_EVO_ATTTENTION=true

N_sample=5
N_step=200
N_cycle=10
seed=101
input_json_path="./examples/example.json"
dump_dir="./output"

torchrun \
    --nproc_per_node $NPROC \
    --master_addr $WORKER_0_HOST \
    --master_port $WORKER_0_PORT \
    --node_rank=$ID \
    --nnodes=$WORKER_NUM \
    runner/inference.py \
    --seeds ${seed} \
    --dump_dir ${dump_dir} \
    --input_json_path ${input_json_path} \
    --model.N_cycle ${N_cycle} \
    --sample_diffusion.N_sample ${N_sample} \
    --sample_diffusion.N_step ${N_step}

[rjrich]

Thank you for your quick and informative reply. Could you please explain which of the options under "torchrun" to use for specifying the use of more than one GPU and give an example of the corresponding command? Thank you.

[zhangyuxuann]

@rjrich you can refer to https://pytorch.org/docs/stable/elastic/run.html#environment-variables, if you set those DDP Environment Variables, the above script is just a demo command