How to set ntasks or ntasks-per-node in slurm for running multi-node distributed training in pytorch?

Viewed 1137

I want to run the official video classification script here.

I plan to run it on slurm machine. I am not sure how to properly set the SBATCH argument --ntasks or --ntasks-per-node if I want to run this script on 2 nodes with 8 V100 GPUs each. like --ntasks 16 or --ntasks-per-node 8 ? Any advice? thanks

1 Answers

If the python code does not explicitly start multiple Slurm tasks, then you should go with this in your job script:

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8

This would give your script one Slurm task with access to 8 GPUs on each node. If you need more CPU or memory (maybe you do some pre/postprocessing?), you can add

#SBATCH --mem=10G
#SBATCH --cpus-per-task=8
Related