torchrun (彈性啟動)¶
模組 torch.distributed.run。
torch.distributed.run 是一個模組，用於在每個訓練節點上啟動多個分散式訓練程序。
torchrun 是一個 Python 控制檯指令碼，指向在 setup.py 的 entry_points 配置中宣告的主模組 torch.distributed.run。它等同於呼叫 python -m torch.distributed.run。
torchrun 可用於單節點分散式訓練，其中每個節點將啟動一個或多個程序。它可用於 CPU 訓練或 GPU 訓練。如果用於 GPU 訓練，每個分散式程序將在單個 GPU 上執行。這可以顯著提高單節點訓練效能。torchrun 也可用於多節點分散式訓練，透過在每個節點上啟動多個程序來提高多節點分散式訓練效能。對於具有直接 GPU 支援的多個 Infiniband 介面的系統而言，這尤其有利，因為所有介面都可以用於聚合通訊頻寬。
在單節點分散式訓練或多節點分散式訓練這兩種情況下，torchrun 將按給定數量在每個節點上啟動程序 (--nproc-per-node)。如果用於 GPU 訓練，此數量需要小於或等於當前系統上的 GPU 數量 (nproc_per_node)，並且每個程序將在單個 GPU 上執行，從 GPU 0 到 GPU (nproc_per_node - 1)。
2.0.0 版本中有所更改: torchrun 將把引數 --local-rank=<rank> 傳遞給您的指令碼。從 PyTorch 2.0.0 開始，帶短劃線的 --local-rank 優先於先前使用的帶下劃線的 --local_rank。
為了向後相容，使用者可能需要在其引數解析程式碼中處理這兩種情況。這意味著在引數解析器中同時包含 "--local-rank" 和 "--local_rank"。如果僅提供了 "--local_rank"，torchrun 將會觸發錯誤：“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()
用法¶
單節點多工作程序¶
torchrun
--standalone
--nnodes=1
--nproc-per-node=$NUM_TRAINERS
YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)
注意
--nproc-per-node 可以是 "gpu"（為每個 GPU 啟動一個程序）、"cpu"（為每個 CPU 啟動一個程序）、"auto"（如果 CUDA 可用則等同於 "gpu"，否則等同於 "cpu"），或者是一個指定程序數量的整數。更多詳情請參見 torch.distributed.run.determine_local_world_size。
堆疊式單節點多工作程序¶
要在同一主機上執行多個單節點多工作程序例項（獨立作業），我們需要確保每個例項（作業）設定在不同的埠上，以避免埠衝突（或更糟的是，兩個作業合併為一個作業）。為此，您必須使用 --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 匯合後端的節點和埠。它可以是您訓練叢集中的任何節點，但理想情況下應選擇具有高頻寬的節點。
注意
如果未指定埠號，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 匯合後端的節點和埠。它可以是您訓練叢集中的任何節點，但理想情況下應選擇具有高頻寬的節點。
注意
如果未指定埠號，HOST_NODE_ADDR 預設為 29400。
關於匯合後端的注意事項¶
對於多節點訓練，您需要指定
--rdzv-id: 一個唯一的作業 ID（由參與作業的所有節點共享）--rdzv-backend:torch.distributed.elastic.rendezvous.RendezvousHandler的實現--rdzv-endpoint: 匯合後端執行的端點；通常形式為host:port。
目前內建支援 c10d（推薦）、etcd-v2 和 etcd（遺留）匯合後端。要使用 etcd-v2 或 etcd，請設定一個啟用 v2 API（例如 --enable-v2）的 etcd 伺服器。
警告
etcd-v2 和 etcd 匯合後端使用 etcd API v2。您必須在 etcd 伺服器上啟用 v2 API。我們的測試使用 etcd v3.4.3。
警告
對於基於 etcd 的匯合，我們建議使用 etcd-v2 而非 etcd，前者功能上等同，但使用了改進的實現。etcd 已進入維護模式，並將在未來版本中移除。
定義¶
節點 (Node)- 一個物理例項或容器；對應於作業管理器處理的單元。工作程序 (Worker)- 分散式訓練中的一個工作程序。工作程序組 (WorkerGroup)- 執行相同功能（例如，訓練器）的工作程序集合。本地工作程序組 (LocalWorkerGroup)- 執行在同一節點上的工作程序組中的一個子集。全域性秩 (RANK)- 工作程序在工作程序組中的秩。全域性工作程序數 (WORLD_SIZE)- 工作程序組中的工作程序總數。本地秩 (LOCAL_RANK)- 工作程序在本地工作程序組中的秩。本地工作程序數 (LOCAL_WORLD_SIZE)- 本地工作程序組的大小。rdzv_id- 使用者定義的 ID，唯一標識一個作業的工作程序組。此 ID 由每個節點用於加入特定工作程序組作為成員。
rdzv_backend- 匯合的後端（例如c10d）。這通常是一個強一致性的鍵值儲存。rdzv_endpoint- 匯合後端端點；通常形式為<host>:<port>。
一個 節點 (Node) 執行 本地工作程序數 (LOCAL_WORLD_SIZE) 個工作程序，這些工作程序構成一個 本地工作程序組 (LocalWorkerGroup)。作業中所有節點的 本地工作程序組 (LocalWorkerGroup) 的並集構成了 工作程序組 (WorkerGroup)。
環境變數¶
以下環境變數在您的指令碼中可用
LOCAL_RANK- 本地秩。RANK- 全域性秩。GROUP_RANK- 工作程序組的秩。一個介於 0 和max_nnodes之間的數字。當每個節點執行一個工作程序組時，這表示節點的秩。ROLE_RANK- 在所有具有相同角色的工作程序中的秩。工作程序的角色在WorkerSpec中指定。LOCAL_WORLD_SIZE- 本地工作程序數（例如，本地執行的工作程序數量）；等於在torchrun上指定的--nproc-per-node。WORLD_SIZE- 全域性工作程序數（作業中的工作程序總數）。ROLE_WORLD_SIZE- 以WorkerSpec中指定的相同角色啟動的工作程序總數。MASTER_ADDR- 執行秩為 0 的工作程序的主機的 FQDN（完全限定域名）；用於初始化 Torch 分散式後端。MASTER_PORT-MASTER_ADDR上可用於託管 C10d TCP store 的埠。TORCHELASTIC_RESTART_COUNT- 工作程序組到目前為止的重啟次數。TORCHELASTIC_MAX_RESTARTS- 配置的最大重啟次數。TORCHELASTIC_RUN_ID- 等於匯合run_id（例如，唯一的作業 ID）。PYTHON_EXEC- 系統可執行檔案覆蓋。如果提供，Python 使用者指令碼將使用PYTHON_EXEC的值作為可執行檔案。預設使用 sys.executable。
部署¶
（C10d 後端不需要）啟動匯合後端伺服器並獲取端點（作為
--rdzv-endpoint傳遞給torchrun）單節點多工作程序：在主機上啟動
torchrun以啟動代理程序，該程序建立並監控本地工作程序組。多節點多工作程序：在所有參與訓練的節點上使用相同的引數啟動
torchrun。
使用作業/叢集管理器時，多節點作業的入口點命令應該是 torchrun。
故障模式¶
工作程序故障：對於具有
n個工作程序的訓練作業，如果k<=n個工作程序發生故障，所有工作程序將被停止並重啟，最多重啟max_restarts次。代理故障：代理故障導致本地工作程序組故障。作業管理器可以選擇使整個作業失敗（組語義）或嘗試替換節點。這兩種行為都由代理支援。
節點故障：與代理故障相同。
成員變更¶
節點離開（縮容）：代理收到離開通知，所有現有工作程序停止，形成新的
WorkerGroup，所有工作程序使用新的RANK和WORLD_SIZE啟動。節點加入（擴容）：新節點被新增到作業中，所有現有工作程序停止，形成新的
WorkerGroup，所有工作程序使用新的RANK和WORLD_SIZE啟動。
重要注意事項¶
該工具和多程序分散式（單節點或多節點）GPU 訓練目前僅在使用 NCCL 分散式後端時能獲得最佳效能。因此，NCCL 後端是 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。這通常是程序的本地秩。換句話說，為了使用此工具，device_ids 需要設定為 [int(os.environ("LOCAL_RANK"))]，並且 output_device 需要設定為 int(os.environ("LOCAL_RANK"))。
在故障或成員變更發生時，所有存活的工作程序都會立即終止。請務必定期儲存檢查點。檢查點的頻率應取決於您的作業對丟失工作的容忍度。
此模組僅支援同構的
本地工作程序數 (LOCAL_WORLD_SIZE)。也就是說，假定所有節點執行相同數量的本地工作程序（按角色）。全域性秩 (RANK)是不穩定的。在重啟之間，節點上的本地工作程序可能被分配與之前不同的秩範圍。切勿硬編碼關於秩穩定性或RANK與LOCAL_RANK之間相關性的任何假設。使用彈性功能（
min_size!=max_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)
（推薦）當工作程序發生錯誤時，此工具將彙總錯誤詳情（例如，時間、秩、主機、程序 ID、堆疊跟蹤等）。在每個節點上，第一個錯誤（按時間戳）會被啟發式地報告為“根本原因”錯誤。要將堆疊跟蹤作為錯誤摘要輸出的一部分，您必須按照下面的示例所示，修飾訓練指令碼中的主入口點函式。如果未修飾，則摘要將不包含異常的堆疊跟蹤，僅包含退出碼。有關 torchelastic 錯誤處理的詳細資訊，請參見：https://pytorch.com.tw/docs/stable/elastic/errors.html
from torch.distributed.elastic.multiprocessing.errors import record
@record
def main():
# do train
pass
if __name__ == "__main__":
main()
