torchrun (弹性启动)¶
torch.distributed.launch
的超集。
torchrun
提供了 torch.distributed.launch
的所有功能,并增加了以下功能
通过重新启动所有工作器来优雅地处理工作器故障。
自动分配工作器
RANK
和WORLD_SIZE
。允许在最小和最大大小之间更改节点数量(弹性)。
注意
torchrun
是一个 python 控制台脚本,指向 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
cmd 参数中读取本地排名。请将您的训练脚本更改为从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
的功能。
用法¶
单节点多工作器¶
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...)
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 的 etcd 服务器(例如 --enable-v2
)。
警告
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
。作业中所有节点的 LocalWorkerGroups
的并集构成 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 的工作器的宿主机的主机名;用于初始化 Torch 分布式后端。MASTER_PORT
-MASTER_ADDR
上的端口,可用于托管 C10d TCP 存储。TORCHELASTIC_RESTART_COUNT
- 工作器组重启次数。TORCHELASTIC_MAX_RESTARTS
- 配置的最大重启次数。TORCHELASTIC_RUN_ID
- 等于会合的run_id
(例如,唯一的作业 ID)。PYTHON_EXEC
- 系统可执行文件覆盖。如果提供,Python 用户脚本将使用PYTHON_EXEC
的值作为可执行文件。默认情况下使用 sys.executable。
部署¶
(C10d 后端不需要)启动会合后端服务器并获取端点(将作为
--rdzv-endpoint
传递给启动器脚本)单节点多工作器:在主机上启动启动器,以启动创建并监控本地工作器组的代理进程。
多节点多工作器:在参与训练的所有节点上使用相同的参数启动启动器。
当使用作业/集群管理器时,多节点作业的入口点命令应为此启动器。
故障模式¶
工作器故障:对于具有
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)
(推荐) 当工作进程出现错误时,此工具会汇总错误的详细信息(例如时间、排名、主机、pid、回溯等)。在每个节点上,第一个错误(按时间戳排序)会被启发式地报告为“根本原因”错误。要将回溯作为此错误摘要输出的一部分,您必须像下面示例中所示那样装饰训练脚本中的主要入口点函数。如果未装饰,则摘要将不包含异常的回溯,而只会包含退出代码。有关 torchelastic 错误处理的详细信息,请参阅:https://pytorch.ac.cn/docs/stable/elastic/errors.html
from torch.distributed.elastic.multiprocessing.errors import record
@record
def main():
# do train
pass
if __name__ == "__main__":
main()