启动分布式程序

启动分布式程序#

安装 MLX python 包会提供一个辅助脚本 mlx.launch，可用于在多个节点上分布式运行 python 脚本。它允许使用 MPI 后端或 ring 后端启动。请参阅分布式文档了解不同后端。

使用方法#

mlx.launch 的最小使用示例如下

mlx.launch --hosts ip1,ip2 my_script.py

或者在 localhost 上进行测试

mlx.launch -n 2 my_script.py

mlx.launch 命令连接到提供的主机，并在每个主机上启动输入的脚本。它监控每个启动的进程，如果其中一个进程意外失败或 mlx.launch 终止，则会终止其余进程。它还会负责将每个远程进程的输出分别转发到 stdout 和 stderr。

提供主机#

主机可以像上面那样作为命令行参数提供，但完全定义主机列表的方式是通过 JSON hostfile。hostfile 的模式非常简单。它只是一个对象列表，通过 ssh 连接的主机名和用于通信的 IP 列表来定义每个主机。

[
    {"ssh": "hostname1", "ips": ["123.123.1.1", "123.123.2.1"]},
    {"ssh": "hostname2", "ips": ["123.123.1.2", "123.123.2.2"]}
]

您可以使用 mlx.distributed_config --over ethernet 创建一个包含与 en0 接口对应的 IP 的 hostfile。

设置远程主机#

为了能够在每个主机上启动脚本，我们需要能够通过 ssh 连接。此外，输入的脚本和 python 可执行文件需要在每个主机上并且位于相同的路径。以下是一个良好的错误调试清单

  • ssh hostname 工作正常，无需密码或主机确认

  • python 可执行文件在所有主机上都可用且位于相同的路径。您可以使用 mlx.launch --print-python 查看该路径是什么。

  • 您想运行的脚本在所有主机上都可用且位于相同的路径

MPI 细节#

可以通过将 --backend mpi 传递给 mlx.launch 来使用 MPI。在这种情况下，mlx.launch 是 mpirun 的一个轻量级包装器。此外，

  • hostfile 中的 IP 将被忽略

  • ssh 连接要求更强，因为每个节点都需要能够连接到所有其他节点

  • mpirun 需要在每个节点上可用且位于相同的路径

最后，可以使用 --mpi-arg 将参数传递给 mpirun。例如，要为 MPI 的字节传输层选择一个特定的接口，我们可以按如下方式调用 mlx.launch

mlx.launch --backend mpi --mpi-arg '--mca btl_tcp_if_include en0' --hostfile hosts.json my_script.py

Ring 细节#

Ring 后端也是默认后端，可以使用参数 --backend ring 明确选择。Ring 后端有一些与 MPI 不同的特定要求和参数

  • 参数 --hosts 只接受 IP，不接受主机名。如果需要 ssh 到与我们想要绑定的 IP 不对应的主机名，则必须提供 hostfile。

  • --starting-port 定义了在远程主机上绑定的端口。具体来说，第一个 IP 的 rank 0 将使用此端口，后续的每个 IP 或 rank 将在此端口基础上加 1。

  • --connections-per-ip 允许我们增加相邻节点之间的连接数。这对应于 mpirun 的 --mca btl_tcp_links 2。

目录