torchrun (弹性启动)¶

torch.distributed.launch 的超集。

torchrun 提供的功能是 torch.distributed.launch 的超集，并具有以下附加功能

Worker 故障可通过重启所有 worker 来优雅地处理。
Worker RANK 和 WORLD_SIZE 会自动分配。
节点数允许在最小和最大尺寸之间变化（弹性）。

注意

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，请按照以下步骤操作

如果您的训练脚本已经在从 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 参数中读取本地 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 特性的信息。

用法¶

单节点多 worker¶

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

堆叠式单节点多 worker¶

要在同一主机上运行单节点多 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...)

容错（固定大小的 worker 数量，无弹性，容忍 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 - 执行相同功能（例如训练器）的 worker 集合。
LocalWorkerGroup - 在同一节点上运行的 worker 组中的 worker 子集。
RANK - worker 在 worker 组中的 rank。
WORLD_SIZE - worker 组中 worker 的总数。
LOCAL_RANK - worker 在本地 worker 组中的 rank。
LOCAL_WORLD_SIZE - 本地 worker 组的大小。
rdzv_id - 用户定义的 ID，用于唯一标识作业的 worker 组。每个节点都使用此 ID 加入特定 worker 组的成员。

rdzv_backend - rendezvous 的后端（例如 c10d）。这通常是一个强一致性的键值存储。
rdzv_endpoint - rendezvous 后端端点；通常格式为 <host>:<port>。

一个 Node 运行 LOCAL_WORLD_SIZE 个 worker，这些 worker 构成一个 LocalWorkerGroup。作业中节点中所有 LocalWorkerGroups 的并集构成了 WorkerGroup。

环境变量¶

以下环境变量可在您的脚本中使用

LOCAL_RANK - 本地 rank。
RANK - 全局 rank。
GROUP_RANK - worker 组的 rank。一个介于 0 和 max_nnodes 之间的数字。当每个节点运行单个 worker 组时，这是节点的 rank。
ROLE_RANK - worker 在所有具有相同角色的 worker 中的 rank。worker 的角色在 WorkerSpec 中指定。
LOCAL_WORLD_SIZE - 本地 world size（例如，本地运行的 worker 数量）；等于在 torchrun 上指定的 --nproc-per-node。
WORLD_SIZE - world size（作业中 worker 的总数）。
ROLE_WORLD_SIZE - 使用 WorkerSpec 中指定的相同角色启动的 worker 总数。
MASTER_ADDR - 运行 rank 0 worker 的主机的 FQDN；用于初始化 Torch 分布式后端。
MASTER_PORT - MASTER_ADDR 上的端口，可用于托管 C10d TCP 存储。
TORCHELASTIC_RESTART_COUNT - 迄今为止 worker 组重启的次数。
TORCHELASTIC_MAX_RESTARTS - 配置的最大重启次数。
TORCHELASTIC_RUN_ID - 等于 rendezvous run_id（例如，唯一的作业 ID）。
PYTHON_EXEC - 系统可执行文件覆盖。如果提供，python 用户脚本将使用 PYTHON_EXEC 的值作为可执行文件。默认情况下使用 sys.executable。

部署¶

（C10d 后端不需要）启动 rendezvous 后端服务器并获取端点（作为 --rdzv-endpoint 传递给启动器脚本）
单节点多 worker：在主机上启动启动器以启动 agent 进程，该进程创建和监视本地 worker 组。
多节点多 worker：在参与训练的所有节点上使用相同的参数启动启动器。

当使用作业/集群管理器时，多节点作业的入口点命令应为该启动器。

故障模式¶

Worker 故障：对于具有 n 个 worker 的训练作业，如果 k<=n 个 worker 发生故障，则所有 worker 都将停止并重启，最多 max_restarts 次。
Agent 故障：Agent 故障会导致本地 worker 组故障。这取决于作业管理器是否使整个作业失败（gang 语义）或尝试替换节点。agent 支持这两种行为。
节点故障：与 agent 故障相同。

成员变更¶

节点离开（缩减规模）：agent 会收到离开通知，所有现有 worker 都将停止，形成新的 WorkerGroup，并且所有 worker 都将以新的 RANK 和 WORLD_SIZE 启动。
节点加入（扩大规模）：新节点被允许加入作业，所有现有 worker 都将停止，形成新的 WorkerGroup，并且所有 worker 都将以新的 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。这通常是进程的本地 rank。换句话说，device_ids 需要是 [int(os.environ("LOCAL_RANK"))]，并且 output_device 需要是 int(os.environ("LOCAL_RANK")) 才能使用此实用程序

在发生故障或成员变更时，所有幸存的 worker 都会立即被杀死。请务必检查点您的进度。检查点的频率应取决于您的作业对工作丢失的容忍度。
此模块仅支持同构的 LOCAL_WORLD_SIZE。也就是说，假设所有节点运行相同数量的本地 worker（每个角色）。
RANK 是不稳定的。在重启之间，节点上的本地 worker 可以被分配与之前不同的 rank 范围。永远不要硬编码任何关于 rank 稳定性的假设，或者 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 错误时，此工具将总结错误详情（例如，时间、rank、主机、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()

torchrun (弹性启动)¶

从 torch.distributed.launch 迁移到 torchrun¶

用法¶

单节点多 worker¶

堆叠式单节点多 worker¶

容错（固定大小的 worker 数量，无弹性，容忍 3 次故障）¶

弹性（`min=1`，`max=4`，容忍最多 3 次成员变更或故障）¶

关于 rendezvous 后端的说明¶

定义¶

环境变量¶

部署¶

故障模式¶

成员变更¶

重要提示¶

文档

教程

资源

torchrun (弹性启动)¶

从 torch.distributed.launch 迁移到 torchrun¶

用法¶

单节点多 worker¶

堆叠式单节点多 worker¶

容错（固定大小的 worker 数量，无弹性，容忍 3 次故障）¶

弹性（min=1，max=4，容忍最多 3 次成员变更或故障）¶

关于 rendezvous 后端的说明¶

定义¶

环境变量¶

部署¶

故障模式¶

成员变更¶

重要提示¶

文档

教程

资源

弹性（`min=1`，`max=4`，容忍最多 3 次成员变更或故障）¶