【教程】PyTorch多機多卡分布式訓練的參數說明

【教程】PyTorch多機多卡分布式訓練的參數說明 | 附通用啟動腳本

轉載請注明出處：小鋒學長生活大爆炸[xfxuezhagn.cn]

如果本文幫助到了你，歡迎[點贊、收藏、關注]哦~

torchrun

一、什么是 torchrun

二、torchrun 的核心參數講解

三、torchrun 會自動設置的環境變量

四、torchrun 啟動過程舉例

機器 A（node_rank=0）上運行

機器 B（node_rank=1）上運行

五、小結表格

PyTorch

一、背景回顧

二、init_process_group

三、腳本中通常的典型寫法

通用啟動腳本

torchrun 與 torch.multiprocessing.spawn 的對比可以看這篇：

【知識】torchrun 與 torch.multiprocessing.spawn 的對比

`torchrun`

一、什么是 `torchrun`

torchrun 是 PyTorch 官方推薦的分布式訓練啟動器，它的作用是：

啟動 多進程分布式訓練（支持多 GPU，多節點）
自動設置每個進程的環境變量
協調節點之間建立通信

二、`torchrun` 的核心參數講解

torchrun \--nnodes=2 \--nproc_per_node=2 \--node_rank=0 \--master_addr=192.168.5.228 \--master_port=29400 \xxx.py

🔹 1. --nnodes（Number of Nodes）

表示參與訓練的總機器數。
你有幾臺服務器，就寫幾。
在分布式訓練中，一個 node 就是一臺物理或虛擬的主機。
node的編號從0開始。

? 例子：你用 2 臺機器 → --nnodes=2

🔹 2. --nproc_per_node（Processes Per Node）

表示每臺機器上要啟動幾個訓練進程。
一個進程對應一個 GPU，因通常設置為你機器上要用到的GPU數。
因此，整個分布式環境下，總訓練進程數 = nnodes * nproc_per_node

? 例子：每臺機器用了?2 張 GPU → --nproc_per_node=2

🔹 3. --node_rank

表示當前機器是第幾臺機器。
從 0 開始編號，必須每臺機器都不同！

? 例子：

機器 IP	node_rank
192.168.5.228	0
192.168.5.229	1

🔹 4. --master_addr 和 --master_port

指定主節點的 IP 和端口，用于 rendezvous（進程對齊）和通信初始化。
所有機器必須填寫相同的值！

? 建議：

master_addr 就是你指定為主節點的那臺機器的 IP
master_port 選一個未被占用的端口，比如 29400

三、torchrun 會自動設置的環境變量

當用 torchrun 啟動后，它會自動給每個進程設置這些環境變量：

環境變量	含義
`RANK`	當前進程在全局中的編號（0 ~ world_size - 1）
`LOCAL_RANK`	當前進程在本機中的編號（0 ~ nproc_per_node - 1）
`WORLD_SIZE`	總進程數 = `nnodes * nproc_per_node`

你可以在訓練腳本里用 os.environ["RANK"] 來讀取這些信息：

import os
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])

示例分配圖：

四、torchrun 啟動過程舉例

假設：

有 2 臺機器
每臺機器有 2 個 GPU
總共會啟動 4 個進程

機器 A（node_rank=0）上運行

torchrun \--nnodes=2 \--nproc_per_node=2 \--node_rank=0 \--master_addr=192.168.5.228 \--master_port=29400 \xxx.py

機器 B（node_rank=1）上運行

torchrun \--nnodes=2 \--nproc_per_node=2 \--node_rank=1 \--master_addr=192.168.5.228 \--master_port=29400 \xxx.py

torchrun 給每個進程編號的順序（分配 RANK / LOCAL_RANK）

torchrun 按照每臺機器上 node_rank 的順序，并在每臺機器上依次啟動 LOCAL_RANK=0, 1, ..., n-1，最后合成 RANK。

RANK = node_rank × nproc_per_node + local_rank

Step 1：按 node_rank 升序處理（node 0 → node 1）

Step 2：每個 node 內部從 local_rank=0 開始遞增

本質上：torchrun 是主從結構調度的

所有 node 啟動后，都會和 master_addr 通信。

master 會統一收集所有 node 的狀態。

每個 node 根據你給的 node_rank 自行派生 local_rank=0~n-1

所有節點通過 RANK = node_rank * nproc_per_node + local_rank 得到自己的全局編號。

這個機制是 可預測、可控、可復現 的。

📦 node_rank=0 (機器 1)
? ? ├── local_rank=0 → RANK=0
? ? └── local_rank=1 → RANK=1

📦 node_rank=1 (機器 2)
? ? ├── local_rank=0 → RANK=2
? ? └── local_rank=1 → RANK=3

最終分配：

Node Rank	Local Rank	Global Rank (`RANK`)	使用 GPU
0	0	0	0
0	1	1	1
1	0	2	0
1	1	3	1

五、小結表格

參數	作用	設置方式
`--nnodes`	總節點數	你寫在命令里
`--nproc_per_node`	每臺節點的進程數（= GPU 數）	你寫在命令里
`--node_rank`	當前機器編號（0開始）	每臺機器唯一
`--master_addr`	主節點 IP（所有節點需一致）	你設置
`--master_port`	主節點端口（所有節點需一致）	你設置
`RANK`	當前進程在所有進程中的編號	`torchrun` 自動設置
`LOCAL_RANK`	當前進程在本節點上的編號	`torchrun` 自動設置
`WORLD_SIZE`	總進程數 = `nnodes * nproc_per_node`	自動設置

PyTorch

PyTorch 的分布式通信是如何通過 init_process_group 與 torchrun 生成的環境變量配合起來工作的。

一、背景回顧

你已經用 torchrun 啟動了多個訓練進程，并且 torchrun 為每個進程自動設置了這些環境變量：

變量名	含義
`RANK`	當前進程的全局編號（從 0 開始）
`LOCAL_RANK`	本機上的編號（一般等于 GPU ID）
`WORLD_SIZE`	總進程數
`MASTER_ADDR`	主節點的 IP
`MASTER_PORT`	主節點用于通信的端口

那么 這些變量是如何參與進程通信初始化的？ 這就涉及到 PyTorch 的核心函數：

二、`init_process_group`

torch.distributed.init_process_group?是 PyTorch 初始化分布式通信的入口：

torch.distributed.init_process_group(backend="nccl",  # 或者 "gloo"、"mpi"init_method="env://",  # 通過環境變量讀取設置
)

關鍵點：

backend="nccl"：推薦用于 GPU 分布式通信（高性能）
init_method="env://"：表示通過環境變量來初始化

你不需要自己設置 RANK / WORLD_SIZE / MASTER_ADDR，只要寫：

import torch.distributed as distdist.init_process_group(backend="nccl", init_method="env://")

PyTorch 會自動去環境中讀這些變量：

RANK → 當前進程編號
WORLD_SIZE → 總進程數
MASTER_ADDR、MASTER_PORT → 主節點 IP 和端口

然后就能正確初始化所有通信進程。

三、腳本中通常的典型寫法

import os
import torch# 初始化 PyTorch 分布式通信環境
torch.distributed.init_process_group(backend="nccl", init_method="env://")# 獲取全局/本地 rank、world size
rank = int(os.environ.get("RANK", -1))
local_rank = int(os.environ.get("LOCAL_RANK", -1))
world_size = int(os.environ.get("WORLD_SIZE", -1))# 設置 GPU 顯卡綁定
torch.cuda.set_device(local_rank)
device = torch.device("cuda")# 打印綁定信息
print(f"[RANK {rank} | LOCAL_RANK {local_rank}] Using CUDA device {torch.cuda.current_device()}: {torch.cuda.get_device_name(torch.cuda.current_device())} | World size: {world_size}")

這段代碼在所有進程中都一樣寫，但每個進程啟動時帶的環境變量不同，所以最終 rank、local_rank、world_size 就自然不同了。

通用啟動腳本

#!/bin/bash# 設置基本參數
MASTER_ADDR=192.168.5.228           # 主機IP
MASTER_PORT=29400                   # 主機端口
NNODES=2                            # 參與訓練的總機器數
NPROC_PER_NODE=2                    # 每臺機器上的進程數# 所有網卡的IP地址，用于篩選
ALL_LOCAL_IPS=$(hostname -I)
# 根據本機 IP 配置通信接口
if [[ "$ALL_LOCAL_IPS" == *"192.168.5.228"* ]]; thenNODE_RANK=0                       # 表示當前機器是第0臺機器IFNAME=ens1f1np1  mytorchrun=~/anaconda3/envs/dglv2/bin/torchrun
elif [[ "$ALL_LOCAL_IPS" == *"192.168.5.229"* ]]; thenNODE_RANK=1                       # 表示當前機器是第1臺機器IFNAME=ens2f1np1mytorchrun=/opt/software/anaconda3/envs/dglv2/bin/torchrun
elseexit 1
fi# 設置 RDMA 接口
export NCCL_IB_DISABLE=0            # 是否禁用InfiniBand
export NCCL_IB_HCA=mlx5_1           # 使用哪個RDMA接口進行通信
export NCCL_SOCKET_IFNAME=$IFNAME   # 使用哪個網卡進行通信
export NCCL_DEBUG=INFO              # 可選：調試用
export GLOO_IB_DISABLE=0            # 是否禁用InfiniBand
export GLOO_SOCKET_IFNAME=$IFNAME   # 使用哪個網卡進行通信
export PYTHONUNBUFFERED=1           # 實時輸出日志# 啟動分布式任務
$mytorchrun \--nnodes=$NNODES \--nproc_per_node=$NPROC_PER_NODE \--node_rank=$NODE_RANK \--master_addr=$MASTER_ADDR \--master_port=$MASTER_PORT \cluster.py## 如果想獲取準確報錯位置，可以加以下內容，這樣可以同步所有 CUDA 操作，錯誤不會“延遲觸發”，你會看到確切是哪一行代碼出了問題：
## CUDA_LAUNCH_BLOCKING=1 torchrun ...

本文來自互聯網用戶投稿，該文觀點僅代表作者本人，不代表本站立場。本站僅提供信息存儲空間服務，不擁有所有權，不承擔相關法律責任。
如若轉載，請注明出處：http://www.pswp.cn/pingmian/76906.shtml
繁體地址，請注明出處：http://hk.pswp.cn/pingmian/76906.shtml
英文地址，請注明出處：http://en.pswp.cn/pingmian/76906.shtml

如若內容造成侵權/違法違規/事實不符，請聯系多彩編程網進行投訴反饋email:809451989@qq.com，一經查實，立即刪除！