Training on SLURM with multiple GPUs

[matval]

Hi,
I'm trying to train a model using Pytorch Lightining version 1.9.5 with DDPStrategy and use 2x V100.

If I run sbatch as:

#!/bin/bash

#SBATCH --nodes=1
#SBATCH --gres=gpu:v100:2
#SBATCH --ntasks-per-node=2
#SBATCH --mem-per-cpu=1200
#SBATCH --time=12:00:00

srun python3 train.py

The code runs fine, but it's using only two CPUs for data loading, which makes the training too slow. What is the correct way to enable the training on multiple GPUs with a different number of CPUs?
Thank you.

[changspencer]

Shouldn't you have the option to specify --cpus-per-task in your SLURM batch script? That's where I normally specify a higher number of CPUs for multiple workers.

[GuangShuaiWang]

Add #SBATCH --ntasks-per-node=<gpu_num> in slumr script, it work for me