评价此页
★ ★ ★ ★ ★

torchrun (弹性启动)#

创建于：2021年5月4日 | 最后更新于：2021年8月26日

模块 torch.distributed.run。

torch.distributed.run 是一个模块，用于在每个训练节点上启动多个分布式训练进程。

torchrun 是一个 Python 命令行脚本，指向 entry_points 配置中声明的主要模块 torch.distributed.run 在 setup.py 中。它等同于调用 python -m torch.distributed.run。

torchrun 可用于单节点分布式训练，其中将在每个节点上启动一个或多个进程。它可用于 CPU 训练或 GPU 训练。如果用于 GPU 训练，每个分布式进程将操作于单个 GPU。这可以实现单节点训练性能的显著提升。torchrun 也可用于多节点分布式训练，通过在每个节点上启动多个进程来提高多节点分布式训练的性能。这对于具有直接 GPU 支持的多个 Infiniband 接口的系统尤其有利，因为所有接口都可以用于聚合通信带宽。

在单节点分布式训练或多节点分布式训练这两种情况下，torchrun 都将启动每个节点上指定的进程数（--nproc-per-node）。如果用于 GPU 训练，这个数字需要小于或等于当前系统的 GPU 数量（nproc_per_node），并且每个进程将操作于从 *GPU 0 到 GPU (nproc_per_node - 1)* 的单个 GPU。

版本 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。

关于集合后端的说明#

对于多节点训练，您需要指定

  1. --rdzv-id：一个唯一的作业 ID（所有参与作业的节点共享）

  2. --rdzv-backend：torch.distributed.elastic.rendezvous.RendezvousHandler 的一个实现

  3. --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 处于维护模式，将在未来版本中移除。

定义#

  1. Node - 一个物理实例或容器；映射到作业管理器工作的单元。

  2. Worker - 在分布式训练的上下文中，指一个工作进程。

  3. WorkerGroup - 执行相同功能的（例如，训练器）工作进程的集合。

  4. LocalWorkerGroup - 同一节点上运行的工作进程组的子集。

  5. RANK - 工作进程在工作进程组内的秩。

  6. WORLD_SIZE - 工作进程组中工作进程的总数。

  7. LOCAL_RANK - 工作进程在本地工作进程组内的秩。

  8. LOCAL_WORLD_SIZE - 本地工作进程组的大小。

  9. rdzv_id - 一个用户定义的 ID，用于唯一标识作业的工作进程组。每个节点使用此 ID 作为特定工作进程组的成员。

  1. rdzv_backend - 集合后端（例如，c10d）。这通常是一个强一致性的键值存储。

  2. rdzv_endpoint - 集合后端端点；通常形式为 <host>:<port>。

一个 Node 运行 LOCAL_WORLD_SIZE 个工作进程，它们构成一个 LocalWorkerGroup。作业中所有节点上的 LocalWorkerGroups 的联合构成 WorkerGroup。

环境变量#

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

  1. LOCAL_RANK - 本地秩。

  2. RANK - 全局秩。

  3. GROUP_RANK - 工作进程组的秩。一个介于 0 和 max_nnodes 之间的数字。当每个节点运行一个工作进程组时，这就是节点的秩。

  4. ROLE_RANK - 在所有具有相同角色的工作进程中的秩。工作进程的角色在 WorkerSpec 中指定。

  5. LOCAL_WORLD_SIZE - 本地全局大小（例如，本地运行的工作进程数）；等于 torchrun 上指定的 --nproc-per-node。

  6. WORLD_SIZE - 全局大小（作业中工作进程的总数）。

  7. ROLE_WORLD_SIZE - 使用 WorkerSpec 中指定的相同角色启动的工作进程总数。

  8. MASTER_ADDR - 运行秩为 0 的工作进程的主机的完全限定域名 (FQDN)；用于初始化 Torch 分布式后端。

  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 传递给 torchrun）

  2. 单节点多工作进程：在主机上启动 torchrun 以启动代理进程，该进程创建并监视本地工作进程组。

  3. 多节点多工作进程：在所有参与训练的节点上使用相同的参数启动 torchrun。

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

故障模式#

  1. 工作进程故障：对于具有 n 个工作进程的训练作业，如果 k<=n 个工作进程发生故障，所有工作进程将被停止并最多重启 max_restarts 次。

  2. 代理故障：代理故障会导致本地工作进程组故障。由作业管理器决定是终止整个作业（同组语义）还是尝试替换节点。代理支持这两种行为。

  3. 节点故障：与代理故障相同。

成员变更#

  1. 节点离开（缩减规模）：代理将收到节点离开的通知，所有现有工作进程将被停止，形成一个新的 WorkerGroup，并且所有工作进程将以新的 RANK 和 WORLD_SIZE 启动。

  2. 节点加入（扩容）：新节点将被加入作业，所有现有工作进程将被停止，形成一个新的 WorkerGroup，并且所有工作进程将以新的 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.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()
在此页面上
在 GitHub 上编辑
显示源代码
PyTorch 库