分布式¶

Torchx 帮助您运行分布式训练任务。请查看 torchx.components.train 在单个训练任务运行示例中的内容。在这里我们将使用相同的 Trainer App 示例。但将以分布式方式运行它。

Torchx 使用 Torch 分布式运行来启动用户进程，并期望用户应用程序将以分布式数据并行的方式编写

分布式训练执行¶

为了在单个或多个进程上运行您的训练器，无论是在远程还是本地，您所需要做的就是编写一个分布式torchx组件。我们将使用的示例是分布式训练器组件

该组件定义了用户应用程序的启动方式，torchx 将负责将此转换为特定调度器的定义。

单节点，多训练器（桌面）¶

尝试在您的桌面启动单节点、多训练器示例：

torchx run -s local_cwd     ./torchx/examples/apps/lightning_classy_vision/component.py:trainer_dist     --nnodes 1 --nproc_per_node 2     --rdzv_backend c10d --rdzv_endpoint localhost:29500

0 是对组件函数的引用：分布式训练器组件

单节点，多训练器（kubernetes）¶

现在让我们在Kubernetes集群上启动相同的组件。请查看torchx.schedulers.kubernetes_scheduler，了解在使用kubernetes调度器运行之前需要安装的依赖项。

我们可以使用以下命令在 Kubernetes 上启动应用程序：

torchx run -s kubernetes --scheduler_args namespace=default,queue=default     ./torchx/examples/apps/lightning_classy_vision/component.py:trainer_dist     --nnodes 1 --nproc_per_node 2

namespaces 参数对应您要启动的 Kubernetes 命名空间。 queue 参数是 Volcano 队列。

输出示例：

kubernetes://torchx/default:cv-trainer-pa2a7qgee9zng
torchx 2021-10-18 18:46:55 INFO     Launched app: kubernetes://torchx/default:cv-trainer-pa2a7qgee9zng
torchx 2021-10-18 18:46:55 INFO     AppStatus:
  msg: <NONE>
  num_restarts: -1
  roles: []
  state: PENDING (2)
  structured_error_msg: <NONE>
  ui_url: null

torchx 2021-10-18 18:46:55 INFO     Job URL: None

你可以使用作业 URL 查询作业的状态或日志：

torchx status kubernetes://torchx/default:cv-trainer-pa2a7qgee9zng

torchx 2021-10-18 18:47:44 INFO     AppDef:
  State: SUCCEEDED
  Num Restarts: -1
Roles:
 *worker[0]:SUCCEEDED

尝试运行 torchx log kubernetes://torchx/default:cv-trainer-pa2a7qgee9zng。

多个节点，多个训练器（Kubernetes）¶

在 Kubernetes 中启动多个节点训练器很简单：

torchx run -s kubernetes --scheduler_args namespace=default,queue=default     ./torchx/examples/apps/lightning_classy_vision/component.py:trainer_dist     --nnodes 2 --nproc_per_node 2

上面的命令将在 Kubernetes default 命名空间中使用 Volcano default 队列启动分布式训练任务。它将使用可从 etcd-server:2379 访问的 etcd 服务来执行 etcd 联合。

您可以覆盖会合点端点：

torchx run -s kubernetes --scheduler_args namespace=default,queue=default     ./torchx/examples/apps/lightning_classy_vision/component.py:trainer_dist     --nnodes 2 --nproc_per_node 1     --rdzv_endpoint etcd-server.default.svc.cluster.local:2379

注意

对于GPU训练，保持nproc_per_node等于实例拥有的GPU数量。

上面的命令将在 Kubernetes default 命名空间中使用 Volcano default 队列启动分布式训练任务。它将使用可从 etcd-server:2379 访问的 etcd 服务来执行 etcd 联合。

内置分布式组件¶

在上面的例子中，我们使用了自定义组件来启动用户应用程序。并不总是需要用户自己编写组件。

Torchx 自带一组内置组件，这些组件描述了典型的执行模式。

dist.ddp¶

dist.ddp 是用于以DDP方式分布式运行的应用程序的组件。您可以使用它快速迭代您的应用程序，而无需编写自己的组件。

注意

dist.ddp 是一个通用组件，因此它非常适合快速迭代，但不适合用于生产环境。如果您的应用程序需要投入生产，请建议您编写自己的组件。了解更多关于编写自定义组件的信息。

我们将使用 dist.ddp 来执行以下示例：

# main.py
import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
import torch.distributed.run

def compute_world_size():
    rank = int(os.getenv("RANK", "0"))
    world_size = int(os.getenv("WORLD_SIZE", "1"))
    dist.init_process_group()
    print("successfully initialized process group")

    t = F.one_hot(torch.tensor(rank), num_classes=world_size)
    dist.all_reduce(t)
    computed_world_size = int(torch.sum(t).item())
    print(
        f"rank: {rank}, actual world_size: {world_size}, computed world_size: {computed_world_size}"
    )

if __name__ == "__main__":
    compute_world_size()

单个训练器在桌面上¶

我们可以在桌面上使用以下命令在四个进程中运行此示例：

torchx run -s local_cwd dist.ddp --entrypoint main.py --nproc_per_node 4

在Kubernetes集群上的单一训练器¶

我们可以在 Kubernetes 集群上执行它

torchx run -s kubernetes --scheduler_args namespace=default,queue=defaultdist.ddp --entrypoint main.py --nproc_per_node 4

组件 API 文档

torchx.components.dist.ddp(*script_args: str, entrypoint: str, image: str = 'ghcr.io/pytorch/torchx:0.1.0rc2', rdzv_backend: Optional[str] = None, rdzv_endpoint: Optional[str] = None, resource: Optional[str] = None, nnodes: int = 1, nproc_per_node: int = 1, name: str = 'test-name', role: str = 'worker', env: Optional[Dict[str, str]] = None) → torchx.specs.api.AppDef [source]¶

分布式数据并行风格的应用（一个角色，多个副本）。使用 torch.distributed.run 启动和协调 PyTorch 工作进程。

Parameters

script_args – 脚本参数。
图像 – 容器图像。
入口点 – 在镜像中运行的脚本或二进制文件。
rdzv_backend – 要使用的 rendezvous 后端，允许的值可以在 rdzv 注册表文档中找到。默认后端是 c10d
rdzv_endpoint – 控制器端点。如果 rdzv_backend 是 etcd，这是 etcd 端点；如果是 c10d，则是其中一个主机的端点。默认端点是 localhost:29500
资源 – 可选的命名资源标识符。在本地调度器上运行时，资源参数将被忽略。
nnodes – 节点数量。
nproc_per_node – 数目每个节点上的进程。
名称 – 应用程序的名称。
角色 – 深度学习框架的名称。
环境 – 环境变量。

Returns

Torchx 应用定义

Return type

specs.AppDef