发起 MPI 分布式训练

AIStudio 任务功能支持发起 MPI 的分布式训练任务，同时也支持常见的 Horovod、DeepSpeed 等基于 MPI 的训练框架。本文详细描述了如何配置和启动 MPI 分布式训练任务。

基本流程

在AIStudio 任务中，通过以下步骤启动 MPI 分布式训练任务：

1. 创建训练任务：配置任务 Worker 规格、Worker 数量，设置「分布式框架」为 MPI。
2. 初始化：平台为创建对应的 pod，配置各种训练角色（launcher worker），并注入必要环境变量，供用户代码获取环境信息。
3. 容错检查：训练开始前，平台发起自检，确保所有 Worker 之间网络通畅、GPU、存储工作正常。
4. 训练执行：
   • 若任意 pod 失败（退出码非 0），训练任务失败。
   • 若所有 pod 成功完成（退出码为 0），训练任务成功。

环境变量

平台注入了 Open MPI 或 MPICH 框架需要用到的环境变量，供用户在训练代码或者入口命令中使用。

• OMPI_MCA_orte_default_hostfile=/etc/mpi/hostfile：OpenMPI 的默认 hostfile 路径，用 mpirun 启动任务时无需再指定 hostfile
• MPICH_HOSTFILE=/etc/mpi/hostfile： MPICH 的默认 hostfile 路径，用 mpiexec 启动任务时无需再指定 hostfile

例如，对于 2 个 Worker，每个 Worker 1 个 GPU 的 MPI 训练任务，在任务运行后，登录 launcher 角色的节点，可看到如下环境变量：

root@jo-da4r6pellvvq4tmz-launcher-0:~# env | grep mpi
OMPI_MCA_orte_default_hostfile=/etc/mpi/hostfile
MPICH_HOSTFILE=/etc/mpi/hostfile

root@jo-da4r6pellvvq4tmz-launcher-0:~# cat /etc/mpi/hostfile 
jo-da4r6pellvvq4tmz-launcher-0 slots=1
jo-da4r6pellvvq4tmz-worker-0 slots=1

启动命令

通过网页提交训练任务时，可以在界面中填写 Bash 启动命令（entrypoint）。以下是简单实例详细说明：

/usr/local/openmpi/bin/mpirun --allow-run-as-root \
--np 16
-x ... \
-x ... \
...

NOTE

请注意添加 --allow-run-as-root 允许 root 用户运行。

故障排除

• 在任务运行过程中，可从网页端登录任务 Worker，访问 Web Terminal。在任务详情页底部可看到登录按钮。
• 任务功能提供 atlctl 命令行调试工具，您可以从 Web Terminal 登录任意的任务 Worker，执行停止任务、统一下发测试命令等调试工作。