发起 MPI 分布式训练

创建任务时 训练框架 选 MPI。通用字段见 创建任务配置说明。

用 MPI 启动的训练走这一页。Horovod、DeepSpeed 等基于 MPI 的启动方式，也可以按 MPI 框架提交。

角色与节点

角色	节点数	资源
master	固定 1	可以是 CPU 或 GPU
worker	可以为 0	通常是 GPU

两边都选 GPU 时，卡型和每节点卡数必须一致。

切换框架后启动命令不会自动变。打开 查看示例，确认后点 填入该示例。

平台写入的环境变量

不要在 环境变量 里覆盖下面这些名称，也不要改平台为 RDMA 写入的 NCCL_IB_* 一类变量。与角色无关的通用变量见 通用环境变量列表。

变量	在示例里的用法
GJ_TRAINING_ROLE	当前节点角色。worker 只启动 sshd，不执行 mpiexec
GJ_TRAINING_VC_JOB_NAME	和 Pod 名拼出 master 主机名
GJ_TRAINING_POD_NAME	当前 Pod 名
GJ_TRAINING_MASTER_CPU_LIMIT	master 的 CPU 核数，双方都不用 GPU 时用作 slot
GJ_TRAINING_MASTER_GPU_LIMIT	master 的 GPU 卡数
GJ_TRAINING_MASTER_REPLICAS	master 节点数
GJ_TRAINING_WORKER_CPU_LIMIT	worker 的 CPU 核数，双方都不用 GPU 时用作 slot
GJ_TRAINING_WORKER_GPU_LIMIT	worker 的 GPU 卡数
GJ_TRAINING_WORKER_REPLICAS	worker 节点数
MPI_HOST	平台写入的 worker 主机列表，示例用它拼 --host

示例脚本自己计算 MPI_NP 和 MPI_HOSTS，不要在环境变量里预先写这两个名称。

启动命令

同一段运行参数会下发到每个节点。worker 只启动 sshd；master 计算主机列表和进程数后执行 mpiexec。运行命令为 bash。运行参数与控制台 查看示例 中的 MPI 示例一致：

set -euo pipefail

ROLE="${GJ_TRAINING_ROLE}"
JOB_NAME="${GJ_TRAINING_VC_JOB_NAME}"
POD_NAME="${GJ_TRAINING_POD_NAME}"

MASTER_CPU="${GJ_TRAINING_MASTER_CPU_LIMIT}"
MASTER_GPU="${GJ_TRAINING_MASTER_GPU_LIMIT}"
MASTER_REPLICAS="${GJ_TRAINING_MASTER_REPLICAS}"

WORKER_CPU="${GJ_TRAINING_WORKER_CPU_LIMIT}"
WORKER_GPU="${GJ_TRAINING_WORKER_GPU_LIMIT}"
WORKER_REPLICAS="${GJ_TRAINING_WORKER_REPLICAS}"

if [[ "${ROLE}" == "worker" ]]; then
  echo "Starting MPI worker"
  mkdir -p /var/run/sshd
  exec /usr/sbin/sshd -D
fi

build_worker_hosts(){
  local slots="$1"

  if [[ -z "${MPI_HOST:-}" ]]; then
    return 0
  fi

  echo "${MPI_HOST}" | awk -v slots="${slots}" -F',' '
  {
    sep=""
    for(i=1;i<=NF;i++){
      if($i!=""){
        printf "%s%s:%s",sep,$i,slots
        sep=",";
      }
    }
  }'
}

MASTER_HOST="${POD_NAME}.${JOB_NAME}"

if [[ "${MASTER_GPU}" -eq 0 && "${WORKER_GPU}" -eq 0 ]]; then

  WORKER_HOSTS=$(build_worker_hosts "${WORKER_CPU}")

  MPI_HOSTS="${MASTER_HOST}:${MASTER_CPU}"

  if [[ -n "${WORKER_HOSTS}" ]]; then
    MPI_HOSTS="${MPI_HOSTS},${WORKER_HOSTS}"
  fi

  MPI_NP=$((MASTER_REPLICAS*MASTER_CPU + WORKER_REPLICAS*WORKER_CPU))

elif [[ "${MASTER_GPU}" -gt 0 && "${WORKER_GPU}" -gt 0 ]]; then

  WORKER_HOSTS=$(build_worker_hosts "${WORKER_GPU}")

  MPI_HOSTS="${MASTER_HOST}:${MASTER_GPU}"

  if [[ -n "${WORKER_HOSTS}" ]]; then
    MPI_HOSTS="${MPI_HOSTS},${WORKER_HOSTS}"
  fi

  MPI_NP=$((MASTER_REPLICAS*MASTER_GPU + WORKER_REPLICAS*WORKER_GPU))

elif [[ "${MASTER_GPU}" -eq 0 && "${WORKER_GPU}" -gt 0 ]]; then

  MPI_HOSTS=$(build_worker_hosts "${WORKER_GPU}")

  MPI_NP=$((WORKER_REPLICAS*WORKER_GPU))

else

  MPI_HOSTS="${MASTER_HOST}:${MASTER_GPU}"

  MPI_NP=$((MASTER_REPLICAS*MASTER_GPU))

fi

if [[ "${MPI_NP}" -le 0 ]]; then
  echo "ERROR: invalid MPI_NP=${MPI_NP}"
  exit 1
fi

mkdir -p /var/run/sshd
/usr/sbin/sshd -D &

echo "MPI_HOST=${MPI_HOST:-}"
echo "MPI_HOSTS=${MPI_HOSTS}"
echo "MPI_NP=${MPI_NP}"

exec mpiexec \
  --allow-run-as-root \
  --host "${MPI_HOSTS}" \
  -np "${MPI_NP}" \
  --map-by slot \
 <代码文件的绝对路径>

把 <代码文件的绝对路径> 换成容器内的入口文件绝对路径。镜像里需要有 /usr/sbin/sshd 和 mpiexec。