目录
  • 文档 >
  • 火炬分布式弹性 >
  • torchrun (弹性启动)
快捷键

torchrun (弹性启动) ¶

模块 torch.distributed.run 。

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

torchrun 是一个 Python 控制台脚本，用于在 setup.py 中声明的 entry_points 配置中的 torch.distributed.run 主模块。它相当于调用 python -m torch.distributed.run 。

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

在单节点分布式训练或多节点分布式训练的两种情况下， 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 将触发错误：“错误：未识别的参数：–local-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 ，格式为 [:]（例如 node1.example.com:29400），指定 C10d rendezvous 后端应该实例化和托管在哪个节点和端口上。它可以是您的训练集群中的任何节点，但理想情况下，您应该选择一个具有高带宽的节点。

注意

如果未指定端口号，则默认为 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 ，格式为 [:]（例如 node1.example.com:29400），指定 C10d rendezvous 后端应该实例化和托管在哪个节点和端口上。它可以是您的训练集群中的任何节点，但理想情况下，您应该选择一个具有高带宽的节点。

注意

如果未指定端口号，则默认为 29400。

关于 rendezvous 后端说明 ¶

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

  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. 执行相同函数的工人集合（例如，训练师）。

  4. 在同一节点上运行的工人组中的工人子集。

  5. 工人组内工人的排名。

  6. 工人组中工人的总数。

  7. LOCAL_RANK - 工作组内工人的排名。

  8. LOCAL_WORLD_SIZE - 本地工作组的规模。

  9. rdzv_id - 用户定义的唯一标识符，用于标识作业的工人组。该 ID 由每个节点使用，以加入特定的工人组。

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

  2. 预约后端端点；通常形式为。

一个 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 - 本地世界大小（例如本地运行的 worker 数量）；等于在 torchrun 中指定的 --nproc-per-node 。

  6. WORLD_SIZE - 世界大小（作业中所有 worker 的总数）。

  7. ROLE_WORLD_SIZE - 与 WorkerSpec 中指定的相同角色一起启动的总 worker 数量。

  8. MASTER_ADDR - 运行 rank 0 worker 的主机的 FQDN；用于初始化 Torch 分布式后端。

  9. MASTER_PORT - 可以用于托管 C10d TCP 存储的端口号。

  10. TORCHELASTIC_RESTART_COUNT - 到目前为止的 worker 组重启次数。

  11. TORCHELASTIC_MAX_RESTARTS - 配置的最大重启次数。

  12. TORCHELASTIC_RUN_ID - 等于 rendezvous run_id （例如，唯一的作业 ID）。

  13. 系统可执行文件覆盖。如果提供，Python 用户脚本将使用 PYTHON_EXEC 的值作为可执行文件。默认使用 sys.executable。

部署

  1. 对于 C10d 后端不需要。启动 rendezvous 后端服务器并获取端点（将作为 --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。这通常是进程的本地 rank。换句话说， 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.org/docs/stable/elastic/errors.html

from torch.distributed.elastic.multiprocessing.errors import record


@record
def main():
    # do train
    pass


if __name__ == "__main__":
    main()

© 版权所有 PyTorch 贡献者。

使用 Sphinx 构建，主题由 Read the Docs 提供。

文档

查看 PyTorch 的全面开发者文档

查看文档

教程

深入了解初学者和高级开发者的教程

查看教程

资源

查找开发资源并获得您的疑问解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookie。通过点击或导航，您同意允许我们使用 cookie。作为本网站的当前维护者，Facebook 的 cookie 政策适用。了解更多信息，包括可用的控制选项：cookie 政策。