torchrun (彈性啟動)¶

是 torch.distributed.launch 的超集合。

torchrun 提供的功能是 torch.distributed.launch 的超集合，具有以下額外功能：

Worker 故障會透過重新啟動所有 workers 來優雅地處理。
Worker RANK 和 WORLD_SIZE 會自動分配。
節點數量允許在最小值和最大值之間變化 (彈性)。

注意

torchrun 是一個 python console script，指向 torch.distributed.run 這個主模組，該模組在 setup.py 中的 entry_points 設定中被宣告。它等同於調用 python -m torch.distributed.run。

從 torch.distributed.launch 遷移到 torchrun¶

torchrun 支援與 torch.distributed.launch 相同的參數，除了已棄用的 --use-env 之外。若要從 torch.distributed.launch 遷移到 torchrun，請按照下列步驟操作：

如果您的訓練腳本已經從 LOCAL_RANK 環境變數讀取 local_rank，那麼您只需要省略 --use-env 標誌，例如：
torch.distributed.launch

torchrun
$ python -m torch.distributed.launch --use-env train_script.py
$ torchrun train_script.py

如果您的訓練腳本從 --local-rank 命令列參數讀取 local rank，請更改您的訓練腳本，從 LOCAL_RANK 環境變數讀取，如下面的程式碼片段所示：

torch.distributed.launch

torchrun

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local-rank", type=int)
args = parser.parse_args()

local_rank = args.local_rank

import os
local_rank = int(os.environ["LOCAL_RANK"])

版本變更於 2.0.0: 啟動器會將 --local-rank=<rank> 參數傳遞給您的腳本。從 PyTorch 2.0.0 開始，帶破折號的 --local-rank 優先於先前使用的帶下劃線的 --local_rank。

為了向後相容，使用者可能需要在引數解析程式碼中處理這兩種情況。這意味著在引數解析器中同時包含 "--local-rank" 和 "--local_rank"。如果僅提供 "--local_rank"，啟動器將觸發錯誤：“error: unrecognized arguments: –local-rank=<rank>”。對於僅支援 PyTorch 2.0.0+ 的訓練程式碼，包含 "--local-rank" 應該就足夠了。

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> parser.add_argument("--local-rank", "--local_rank", type=int)
>>> args = parser.parse_args()

上述變更足以從 torch.distributed.launch 遷移到 torchrun。若要利用 torchrun 的新功能，例如彈性、容錯和錯誤報告，請參閱：

訓練腳本，以取得有關撰寫符合 torchrun 規範的訓練腳本的更多資訊。
本頁的其餘部分提供有關 torchrun 功能的更多資訊。

用法¶

單節點多工作站 (Single-node multi-worker)¶

torchrun
    --standalone
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

堆疊式單節點多工作站 (Stacked single-node multi-worker)¶

若要在同一主機上執行單節點、多工作站的多個實例（獨立作業），我們需要確保每個實例（作業）都設定在不同的埠上，以避免埠衝突（或更糟的情況，兩個作業合併為一個作業）。為此，您必須使用 --rdzv-backend=c10d 執行，並透過設定 --rdzv-endpoint=localhost:$PORT_k 來指定不同的埠。對於 --nodes=1，通常讓 torchrun 自動選擇一個空閒的隨機埠，而不是手動為每次執行指定不同的埠，會更方便。

torchrun
    --rdzv-backend=c10d
    --rdzv-endpoint=localhost:0
    --nnodes=1
    --nproc-per-node=$NUM_TRAINERS
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

容錯 (固定大小的工作站數量，無彈性，可容忍 3 個故障)¶

torchrun
    --nnodes=$NUM_NODES
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

HOST_NODE_ADDR，格式為 <host>[:<port>] (例如 node1.example.com:29400)，指定應例項化和託管 C10d rendezvous 後端的節點和埠。它可以是您訓練叢集中的任何節點，但理想情況下，您應該選擇一個具有高頻寬的節點。

注意

如果未指定埠號，HOST_NODE_ADDR 預設為 29400。

彈性 (`min=1`, `max=4`, 最多可容忍 3 次成員資格變更或故障)¶

torchrun
    --nnodes=1:4
    --nproc-per-node=$NUM_TRAINERS
    --max-restarts=3
    --rdzv-id=$JOB_ID
    --rdzv-backend=c10d
    --rdzv-endpoint=$HOST_NODE_ADDR
    YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

HOST_NODE_ADDR，格式為 <host>[:<port>] (例如 node1.example.com:29400)，指定應例項化和託管 C10d rendezvous 後端的節點和埠。它可以是您訓練叢集中的任何節點，但理想情況下，您應該選擇一個具有高頻寬的節點。

注意

如果未指定埠號，HOST_NODE_ADDR 預設為 29400。

關於 rendezvous 後端的注意事項¶

對於多節點訓練，您需要指定：

--rdzv-id：一個唯一的作業 ID（由參與作業的所有節點共享）
--rdzv-backend：torch.distributed.elastic.rendezvous.RendezvousHandler 的一個實作
--rdzv-endpoint：rendezvous 後端正在執行的端點；通常採用 host:port 格式。

目前，開箱即用支援 c10d (建議)、etcd-v2 和 etcd (舊版) rendezvous 後端。若要使用 etcd-v2 或 etcd，請設定一個啟用 v2 API 的 etcd 伺服器 (例如 --enable-v2)。

警告

etcd-v2 和 etcd rendezvous 使用 etcd API v2。您必須在 etcd 伺服器上啟用 v2 API。我們的測試使用 etcd v3.4.3。

警告

對於基於 etcd 的 rendezvous，我們建議使用 etcd-v2 而不是 etcd，後者在功能上是等效的，但使用了修訂後的實作。etcd 處於維護模式，將在未來的版本中移除。

定義¶

Node - 一個物理實例或容器；映射到作業管理器使用的單元。
Worker - 分散式訓練中的一個工作站 (worker)。
WorkerGroup - 執行相同功能的 workers (例如： trainers) 的集合。
LocalWorkerGroup - 在同一個節點上運行的 worker group 中的 workers 子集。
RANK - worker 在 worker group 中的排名。
WORLD_SIZE - worker group 中 workers 的總數量。
LOCAL_RANK - worker 在 local worker group 中的排名。
LOCAL_WORLD_SIZE - local worker group 的大小。
rdzv_id - 用戶定義的 ID，用於唯一識別 job 的 worker group。每個節點都使用此 ID 加入特定的 worker group。

rdzv_backend - rendezvous 的 backend (例如：c10d)。這通常是一個強一致性的鍵值儲存。
rdzv_endpoint - rendezvous backend 的端點；通常格式為 <host>:<port>。

一個 Node 運行 LOCAL_WORLD_SIZE 個 workers，這些 workers 組成一個 LocalWorkerGroup。 job 中所有節點的 LocalWorkerGroups 的聯集構成 WorkerGroup。

環境變數¶

以下環境變數可在您的腳本中使用：

LOCAL_RANK - 本地排名 (local rank)。
RANK - 全域排名 (global rank)。
GROUP_RANK - worker group 的排名。一個介於 0 和 max_nnodes 之間的數字。當每個節點運行單個 worker group 時，這是節點的排名。
ROLE_RANK - 在所有具有相同角色的 workers 中，該 worker 的排名。worker 的角色在 WorkerSpec 中指定。
LOCAL_WORLD_SIZE - 本地 world size (例如，本地運行的 workers 數量)；等於在 torchrun 上指定的 --nproc-per-node。
WORLD_SIZE - world size (job 中 workers 的總數)。
ROLE_WORLD_SIZE - 使用 WorkerSpec 中指定的相同角色啟動的 workers 的總數。
MASTER_ADDR - 運行排名為 0 的 worker 的主機的 FQDN；用於初始化 Torch Distributed backend。
MASTER_PORT - MASTER_ADDR 上可用於託管 C10d TCP store 的端口。
TORCHELASTIC_RESTART_COUNT - 到目前為止，worker group 重新啟動的次數。
TORCHELASTIC_MAX_RESTARTS - 配置的最大重新啟動次數。
TORCHELASTIC_RUN_ID - 等於 rendezvous run_id (例如，唯一的 job id)。
PYTHON_EXEC - 系統可執行文件覆蓋。如果提供，python 用户腳本將使用 PYTHON_EXEC 的值作爲可執行文件。默認情況下使用 sys.executable 。

部署¶

（C10d backend 不需要）啟動 rendezvous backend 伺服器並獲取端點 (將作為 --rdzv-endpoint 傳遞給啟動腳本)
單節點多 worker：在主機上啟動 launcher 以啟動 agent 進程，該進程創建並監控本地 worker group。
多節點多 worker：在參與訓練的所有節點上使用相同的引數啟動 launcher。

當使用 job/cluster 管理器時，多節點 job 的入口點命令應該是這個 launcher。

失敗模式¶

Worker 失敗：對於具有 n 個 workers 的訓練 job，如果 k<=n 個 workers 失敗，則所有 workers 都會停止並重新啟動，最多達到 max_restarts。
Agent 失敗：Agent 失敗會導致 local worker group 失敗。是否由 job 管理器終止整個 job (gang semantics) 或嘗試替換節點取決於 job 管理器。Agent 同時支持這兩種行為。
節點失敗：與 Agent 失敗相同。

成員變更¶

節點離開 (縮減)：Agent 會收到離開的通知，所有現有的 workers 都會停止，形成一個新的 WorkerGroup，並且所有 workers 都會以新的 RANK 和 WORLD_SIZE 啟動。
節點加入 (擴展)：新節點被允許加入 job，所有現有的 workers 都會停止，形成一個新的 WorkerGroup，並且所有 workers 都會以新的 RANK 和 WORLD_SIZE 啟動。

重要注意事項¶

目前，此工具和多進程分散式（單節點或多節點）GPU 訓練僅在使用 NCCL 分散式 backend 時才能達到最佳性能。因此，建議將 NCCL backend 用於 GPU 訓練。
初始化 Torch 進程組所需的環境變數由此模組提供給您，無需手動傳遞 RANK。要在您的訓練腳本中初始化進程組，只需運行：

>>> import torch.distributed as dist
>>> dist.init_process_group(backend="gloo|nccl")

在您的訓練程式中，您可以使用常規的分散式函數或使用 torch.nn.parallel.DistributedDataParallel() 模組。如果您的訓練程式使用 GPU 進行訓練，並且您想使用 torch.nn.parallel.DistributedDataParallel() 模組，以下是如何配置它。

local_rank = int(os.environ["LOCAL_RANK"])
model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[local_rank],
                                                  output_device=local_rank)

請確保將 device_ids 參數設定為程式碼將運作的唯一 GPU 裝置 ID。這通常是程序的本地排名 (local rank)。換句話說，為了使用此工具，device_ids 必須是 [int(os.environ("LOCAL_RANK"))]，而 output_device 必須是 int(os.environ("LOCAL_RANK"))。

發生故障或成員變更時，所有倖存的 worker 都會立即被終止。請務必對您的進度進行檢查點 (checkpoint)。檢查點的頻率應取決於您對遺失工作的容忍度。
此模組僅支援同質的 LOCAL_WORLD_SIZE。也就是說，假設所有節點都運行相同數量的本地 worker（每個角色）。
RANK 並不穩定。在重新啟動之間，節點上的本地 worker 可能會被分配到與之前不同的排名範圍。永遠不要硬編碼關於排名穩定性或 RANK 和 LOCAL_RANK 之間相關性的任何假設。
當使用彈性 (min_size!=max_size) 時，請勿硬編碼關於 WORLD_SIZE 的假設，因為隨著允許節點離開和加入，world size 可能會改變。
建議您的腳本具有以下結構

def main():
  load_checkpoint(checkpoint_path)
  initialize()
  train()

def train():
  for batch in iter(dataset):
    train_step(batch)

    if should_checkpoint:
      save_checkpoint(checkpoint_path)

(建議) 發生 worker 錯誤時，此工具將總結錯誤的詳細資訊 (例如，時間、排名、主機、PID、追溯等)。在每個節點上，第一個錯誤（按時間戳記）會被啟發式地報告為「根本原因 (Root Cause)」錯誤。要取得作為此錯誤摘要一部分的追溯列印輸出，您必須按照以下範例所示，裝飾您訓練腳本中的主要進入點函式。如果沒有裝飾，則摘要將不包含異常的追溯，僅包含 exitcode。有關 torchelastic 錯誤處理的詳細資訊，請參閱：https://pytorch.dev.org.tw/docs/stable/elastic/errors.html

from torch.distributed.elastic.multiprocessing.errors import record

@record
def main():
    # do train
    pass

if __name__ == "__main__":
    main()

torchrun (彈性啟動)¶

從 torch.distributed.launch 遷移到 torchrun¶

用法¶

單節點多工作站 (Single-node multi-worker)¶

堆疊式單節點多工作站 (Stacked single-node multi-worker)¶

容錯 (固定大小的工作站數量，無彈性，可容忍 3 個故障)¶

彈性 (`min=1`, `max=4`, 最多可容忍 3 次成員資格變更或故障)¶

關於 rendezvous 後端的注意事項¶

定義¶

環境變數¶

部署¶

失敗模式¶

成員變更¶

重要注意事項¶

文件

教學

資源

torchrun (彈性啟動)¶

從 torch.distributed.launch 遷移到 torchrun¶

用法¶

單節點多工作站 (Single-node multi-worker)¶

堆疊式單節點多工作站 (Stacked single-node multi-worker)¶

容錯 (固定大小的工作站數量，無彈性，可容忍 3 個故障)¶

彈性 (min=1, max=4, 最多可容忍 3 次成員資格變更或故障)¶

關於 rendezvous 後端的注意事項¶

定義¶

環境變數¶

部署¶

失敗模式¶

成員變更¶

重要注意事項¶

文件

教學

資源

彈性 (`min=1`, `max=4`, 最多可容忍 3 次成員資格變更或故障)¶