啟動分散式程式

啟動分散式程式#

MLX Python 軟體包提供兩個工具，協助你設定 Mac 以進行分散式計算，並在多個節點或單一節點的多個行程上啟動分散式程式。這兩個工具分別是

  • mlx.launch

  • mlx.distributed_config

各種後端的入門與新手指南，請參閱 distributed 文件。

mlx.distributed_config#

除非你只是在本機用於開發或多 GPU 的 CUDA 環境啟動分散式工作，否則你會有多台 Mac 需要設定為可與 MLX 進行分散式通訊。

mlx.distributed_config 旨在自動化設定網路介面（特別是 Thunderbolt 通訊）以及建立供 mlx.launch 使用的 hostfile。

我們將分析使用 mlx.distributed_config 的 3 種情況

  1. 使用 JACCL 的 Thunderbolt RDMA

  2. 使用 ring 後端的 Thunderbolt TCP/IP

  3. 使用 ring 後端的 Ethernet TCP/IP

JACCL#

依照 啟用 RDMA 的步驟 完成後，你可以執行以下命令來設定節點並建立 hostfile。

mlx.distributed_config --verbose --backend jaccl \
     --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 --over thunderbolt \
     --auto-setup --output m3-ultra-jaccl.json

讓我們逐步看看這個腳本設定節點時會做哪些步驟。

  1. ssh 到所有節點以確認可以連線

  2. 取得 Thunderbolt 連線拓撲，也就是在每個節點上執行命令以計算節點之間的連線關係。

  3. 確認我們有有效的完整連通 mesh

  4. 確認 RDMA 已啟用

  5. 從介面 en0 取得 Ethernet IP

  6. 停用 Thunderbolt bridge，並為每條 Thunderbolt 纜線建立點對點網路

  7. 寫入 hostfile

了解上述步驟可以讓你手動設定節點，並排除任何設定問題。例如，你可以在設定檔中把 Ethernet IP 改成其他介面（只要所有節點都能連到即可）。

--auto-setup 參數需要每個節點具備無密碼 sudo。若不可用，設定腳本會列印需要在各節點執行的命令。

Thunderbolt 上的 Ring#

在 Thunderbolt 上設定 ring 後端，只需將 --backend 從 jaccl 改為 ring。

步驟非常相似，主要差異在於腳本不再驗證節點是否完整連通，而是嘗試辨識 ring 拓撲（或多個 ring）。

Ethernet 上的 Ring#

在 Ethernet 上設定 ring 後端不需要設定網路介面，因此它只會從每個節點取得 en0 的 IP 並寫入 hostfile。

除錯纜線連線#

mlx.distributed_config 可藉由匯出連線圖來協助你除錯 Thunderbolt 節點的連線情況。

執行

mlx.distributed_config --verbose \
     --hosts host1,host2,host3,host4 \
     --over thunderbolt --dot

會輸出節點間連線的 GraphViz 表示法，讓你很容易判斷哪條纜線沒有正確連接。

範例請見 JACCL 章節。

mlx.launch#

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

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

或用於在 localhost 測試

mlx.launch -n 2 my_script.py

mlx.launch 會連線到指定的主機並在每台主機上啟動輸入腳本。它會監控每個已啟動的行程，若其中一個意外失敗或 mlx.launch 被終止，則會結束其餘行程。它也會負責將每個遠端行程的輸出分別轉送到 stdout 與 stderr。

更重要的是，它也會將 stdin 廣播到每個行程，讓互動式程式能在分散式模式下運作，並可使用互動式除錯器進行除錯。

提供主機#

主機可以像上面一樣以命令列參數提供，但要完整定義主機清單，建議使用 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 查看該路徑。

  • 要執行的腳本在所有主機上的路徑相同

如果你從與實際執行節點環境完全不同的主機啟動，可以指定 --no-verify-script，讓 mlx.launch 不在啟動分散式工作前檢查本機是否存在可執行檔與腳本。

Ring 的注意事項#

ring 後端也是預設後端，可用 --backend ring 明確指定。ring 後端有一些與其他後端不同的特殊需求與參數：

  • --hosts 參數只接受 IP，不接受主機名稱。若需要 ssh 到的主機名稱不對應你要綁定的 IP，就必須提供 hostfile。

  • --starting-port 定義在遠端主機上要綁定的連接埠。第一個 IP 的 rank 0 會使用此連接埠，後續每個 IP 或 rank 會在此基礎上加 1。

  • --connections-per-ip 可增加相鄰節點之間的連線數量。這相當於 mpirun 的 --mca btl_tcp_links 2。

JACCL 的注意事項#

JACCL 後端可用 --backend jaccl 指定。使用此後端啟動必須提供 hostfile，因為其中需要包含各節點彼此連線的 RDMA 裝置資訊。

NCCL 的注意事項#

NCCL 後端是 CUDA 環境的預設後端。當從 Mac 啟動到具 CUDA 的 Linux 主機時，應使用 --backend nccl 指定後端。

--repeat-hosts, -n 參數應用於啟動多節點、多 GPU 的工作。例如：

mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 --no-verify-script -- ./my-job.sh

會嘗試啟動 16 個行程，每個節點 8 個，全部執行 my-job.sh。

MPI 的注意事項#

可透過對 mlx.launch 傳入 --backend mpi 來使用 MPI。在此情況下，mlx.launch 只是 mpirun 的薄包裝。此外：

  • hostfile 中的 IP 會被忽略

  • ssh 連線要求更嚴格，因為每個節點都必須能連到所有其他節點

  • mpirun 必須在每個節點上以相同路徑存在

最後，你可以使用 --mpi-arg 將參數傳給 mpirun。例如，若要為 MPI 的 byte-transfer-layer 選擇特定介面，可以如下呼叫 mlx.launch：

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