分布式训练ddp和slurm

cv-daily2024-08-23 18:51

使用ddp：https://blog.51cto.com/u_16213675/9632480

使用slurm：https://blog.csdn.net/LittleNyima/article/details/136813418

使用 Slurm 管理多机多卡训练

对于一般的用户来说，ddp多机多卡训练方式已经基本上够用了。然而对于需要进行更大规模训练的人来说，在每个节点上依次运行命令比较繁琐并且容易出错。同时，大规模 GPU 集群需要有效的管理方式，来提高资源利用率。为了做到这一点，Slurm 是一个比较好的选择。Slurm 主要的作用在于任务调度，其可以为用户分配计算机节点来执行任务，并且支持任务队列，可以比较高效地分配资源。

在编写训练脚本时，无论启动方式如何，我们关心的都是 master 节点地址、local rank、进程总数等信息，我们可以参考 mmcv 的方式对这些内容进行初始化：

python 复制代码

def _init_dist_slurm(backend: str, port: Optional[int] = None) -> None:
    proc_id = int(os.environ['SLURM_PROCID'])
    ntasks = int(os.environ['SLURM_NTASKS'])
    node_list = os.environ['SLURM_NODELIST']
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(proc_id % num_gpus)
    addr = subprocess.getoutput(
        f'scontrol show hostname {node_list} | head -n1')
    # specify master port
    if port is not None:
        os.environ['MASTER_PORT'] = str(port)
    elif 'MASTER_PORT' in os.environ:
        pass  # use MASTER_PORT in the environment variable
    else:
        # if torch.distributed default port(29500) is available
        # then use it, else find a free port
        if _is_free_port(29500):
            os.environ['MASTER_PORT'] = '29500'
        else:
            os.environ['MASTER_PORT'] = str(_find_free_port())
    # use MASTER_ADDR in the environment variable if it already exists
    if 'MASTER_ADDR' not in os.environ:
        os.environ['MASTER_ADDR'] = addr
    os.environ['WORLD_SIZE'] = str(ntasks)
    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
    os.environ['RANK'] = str(proc_id)
    dist.init_process_group(backend=backend)

在任务启动时，使用 Slurm 提供的工具：

python 复制代码

srun \
    -p ${PARTITION} \
    --job-name=${JOB_NAME} \
    --gres=${GPUS_PER_NODE} \
    --ntasks=${GPUS} \
    --ntasks-per-node=${GPUS_PER_NODE} \
    --cpus-per-task=${CPUS_PER_TASK} \
    --kill-on-bad-exit=1 \
    python train.py