目錄
捷徑

torchrun（彈性啟動）

torch.distributed.launch 的超集。

torchrun 提供了 torch.distributed.launch 的功能超集，並具有以下額外功能

  1. 透過重新啟動所有工作者，可以妥善處理工作者失敗。

  2. 工作者 RANK 和 WORLD_SIZE 會自動分配。

  3. 允許節點數量在最小和最大大小之間變化（彈性）。

注意

torchrun 是一個 Python 主控台腳本，指向 setup.py 中 entry_points 設定中宣告的主要模組 torch.distributed.run。它相當於呼叫 python -m torch.distributed.run。

從 torch.distributed.launch 轉換到 torchrun

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

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

    torch.distributed.launch

    torchrun

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

    $ torchrun train_script.py

  2. 如果您的訓練腳本從 --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"，啟動器將觸發錯誤：「錯誤：無法辨識的參數：–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 功能的資訊。

用法

單節點多工作者

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

堆疊式單節點多工作者

要在同一台主機上執行多個單節點、多工作者的執行個體（個別作業），我們需要確保每個執行個體（作業）都設定在不同的埠上，以避免埠衝突（或者更糟的情況是，兩個作業被合併為一個作業）。為此，您必須使用 --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...)

以 <主機>[:<埠>] 形式（例如 node1.example.com:29400）的 HOST_NODE_ADDR 指定要建立和託管 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...)

以 <主機>[:<埠>] 形式（例如 node1.example.com:29400）的 HOST_NODE_ADDR 指定要建立和託管 C10d 會合後端的節點和埠。它可以是您的訓練叢集中的任何節點，但理想情況下，您應該選擇具有高頻寬的節點。

注意

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

關於會合後端的注意事項

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

  1. --rdzv-id：一個唯一的工作 ID（由參與工作的所有節點共用）

  2. --rdzv-backend：torch.distributed.elastic.rendezvous.RendezvousHandler 的實作

  3. --rdzv-endpoint：會合後端執行的端點；通常採用 主機:埠 的形式。

目前支援開箱即用的 c10d（推薦）、etcd-v2 和 etcd（舊版）會合後端。若要使用 etcd-v2 或 etcd，請設定啟用 v2 API 的 etcd 伺服器（例如 --enable-v2）。

警告

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

警告

對於基於 etcd 的會合，我們建議使用 etcd-v2，而不是功能相同的 etcd，但使用的是修改過的實作。etcd 處於維護模式，並將在未來版本中移除。

定義

  1. 節點 - 一個實體執行個體或容器；對應到工作管理器處理的單位。

  2. 工作器 - 分散式訓練上下文中的工作器。

  3. 工作器群組 - 執行相同函數（例如訓練器）的工作器集合。

  4. 本地工作器群組 - 在同一個節點上運行的部分工作器群組工作器。

  5. RANK - 工作器在工作器群組中的排名。

  6. WORLD_SIZE - 工作器群組中的工作器總數。

  7. LOCAL_RANK - 工作器在本地工作器群組中的排名。

  8. LOCAL_WORLD_SIZE - 本地工作器群組的大小。

  9. rdzv_id - 一個使用者定義的 ID，用於唯一識別工作的工作器群組。每個節點都使用此 ID 作為特定工作器群組的成員加入。

  1. rdzv_backend - 會合的後端（例如 c10d）。這通常是一個強一致性的鍵值儲存庫。

  2. rdzv_endpoint - 會合後端端點；通常採用 <主機>:<埠> 的形式。

一個 節點 運行 LOCAL_WORLD_SIZE 個工作器，這些工作器構成一個 本地工作器群組。工作中所有節點中所有 本地工作器群組 的聯合構成 工作器群組。

環境變數

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

  1. LOCAL_RANK - 本地排名。

  2. RANK - 全域排名。

  3. GROUP_RANK - 工作器群組的排名。介於 0 到 max_nnodes 之間的數字。當每個節點運行一個工作器群組時，這就是節點的排名。

  4. ROLE_RANK - 工作器在具有相同角色的所有工作器中的排名。工作器的角色在 WorkerSpec 中指定。

  5. LOCAL_WORLD_SIZE - 本地世界大小（例如，本地運行的 worker 數量）；等於在 torchrun 上指定的 --nproc-per-node。

  6. WORLD_SIZE - 世界大小（工作中的工作器總數）。

  7. ROLE_WORLD_SIZE - 使用 WorkerSpec 中指定的相同角色啟動的工作器總數。

  8. MASTER_ADDR - 運行排名為 0 的工作器的 FQDN；用於初始化 Torch Distributed 後端。

  9. MASTER_PORT - MASTER_ADDR 上可用於託管 C10d TCP 儲存庫的埠。

  10. TORCHELASTIC_RESTART_COUNT - 到目前為止，工作器群組重新啟動的次數。

  11. TORCHELASTIC_MAX_RESTARTS - 配置的最大重新啟動次數。

  12. TORCHELASTIC_RUN_ID - 等於會合 run_id（例如，唯一的工作 ID）。

  13. PYTHON_EXEC - 系統可執行檔覆寫。如果提供，則 Python 使用者腳本將使用 PYTHON_EXEC 的值作為可執行檔。預設情況下使用 sys.executable。

部署

  1. （C10d 後端不需要）啟動會合後端伺服器並取得端點（作為 --rdzv-endpoint 傳遞給啟動器腳本）

  2. 單節點多工作器：在主機上啟動啟動器以啟動代理進程，該進程會建立並監控本地工作器群組。

  3. 多節點多工作器：在參與訓練的所有節點上使用相同的參數啟動啟動器。

當使用工作/叢集管理器時，多節點工作的進入點命令應該是這個啟動器。

故障模式

  1. 工作器故障：對於具有 n 個工作器的訓練工作，如果 k<=n 個工作器故障，則所有工作器都會停止並重新啟動，最多 max_restarts 次。

  2. 代理故障：代理故障會導致本地工作器群組故障。由工作管理器決定是使整個工作失敗（群組語義）還是嘗試替換節點。代理支援這兩種行為。

  3. 節點故障：與代理故障相同。

成員變更

  1. 節點離開（縮減）：代理會收到離開通知，所有現有的工作器都會停止，會形成一個新的 工作器群組，並且所有工作器都會使用新的 RANK 和 WORLD_SIZE 啟動。

  2. 節點到達（擴展）：新的節點會被允許加入工作，所有現有的工作器都會停止，會形成一個新的 工作器群組，並且所有工作器都會使用新的 RANK 和 WORLD_SIZE 啟動。

重要注意事項

  1. 這個工具和多進程分散式（單節點或多節點）GPU 訓練目前只有在使用 NCCL 分散式後端時才能達到最佳效能。因此，NCCL 後端是 GPU 訓練的推薦後端。

  2. 初始化 Torch 進程群組所需的環境變數由本模組提供，您無需手動傳遞 RANK。若要在您的訓練腳本中初始化進程群組，只需運行

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

  1. 在您的訓練程序中，您可以使用常規的分散式函數，也可以使用 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"))

  1. 發生故障或成員變更時，所有存活的工作器都會立即被終止。請確保檢查點您的進度。檢查點的頻率應該取決於您的工作對丟失工作的容忍度。

  2. 本模組僅支援同質的 LOCAL_WORLD_SIZE。也就是說，假設所有節點都運行相同數量的本地工作器（每個角色）。

  3. RANK 不穩定。在重新啟動之間，節點上的本地工作器可能會被分配到與之前不同的排名範圍。永遠不要對排名的穩定性或 RANK 和 LOCAL_RANK 之間的某些關聯性進行硬編碼假設。

  4. 當使用彈性（min_size!=max_size）時，不要對 WORLD_SIZE 進行硬編碼假設，因為世界大小可能會隨著節點被允許離開和加入而改變。

  5. 建議您的腳本具有以下結構

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)

  1. （推薦）在工作器出錯時，此工具將摘要錯誤的詳細資訊（例如，時間、排名、主機、PID、回溯等）。在每個節點上，第一個錯誤（按時間戳記）會被啟發式地報告為「根本原因」錯誤。若要取得回溯作為此錯誤摘要列印輸出的一部分，您必須在訓練腳本中裝飾您的主要進入點函數，如下例所示。如果沒有裝飾，則摘要將不包含異常的回溯，並且只會包含退出代碼。有關 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()

文件

取得 PyTorch 的完整開發人員文件

查看文件

教學課程

取得適用於初學者和進階開發人員的深入教學課程

查看教學課程

資源

尋找開發資源並取得問題解答

查看資源

為了分析流量並優化您的體驗，我們在本網站上使用 Cookie。點擊或瀏覽即表示您同意我們使用 Cookie。作為本網站的現任維護者，Facebook 的 Cookie 政策適用。進一步瞭解，包括關於可用控制措施的資訊：Cookie 政策。