配置规模和 GPU#

扩展 Ray Train 训练运行的规模非常简单，只需几行代码即可完成。主要接口是 ScalingConfig，它配置了 worker 的数量以及它们应使用的资源。

本指南中，worker 指的是 Ray Train 分布式训练 worker，它是一个运行你的训练函数的 Ray Actor。

增加 worker 的数量#

控制训练代码并行性的主要接口是设置 worker 的数量。这可以通过将 num_workers 属性传递给 ScalingConfig 来完成。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8
)

使用 GPU#

要使用 GPU，将 use_gpu=True 传递给 ScalingConfig。这将为每个训练 worker 请求一个 GPU。在下面的示例中，训练将在 8 个 GPU 上运行（8 个 worker，每个使用一个 GPU）。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8,
    use_gpu=True
)

在训练函数中使用 GPU#

设置 use_gpu=True 后，Ray Train 会在你的训练函数中自动设置环境变量，以便检测和使用 GPU（例如 CUDA_VISIBLE_DEVICES）。

你可以使用 ray.train.torch.get_device() 获取关联的设备。

import torch
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, get_device


def train_func():
    assert torch.cuda.is_available()

    device = get_device()
    assert device == torch.device("cuda:0")

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=1,
        use_gpu=True
    )
)
trainer.fit()

为 worker 分配多个 GPU#

有时你可能想为 worker 分配多个 GPU。例如，如果你想为每个 worker 分配 2 个 GPU，可以在 ScalingConfig 中指定 resources_per_worker={"GPU": 2}。

你可以使用 ray.train.torch.get_devices() 获取关联的设备列表。

import torch
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, get_device, get_devices


def train_func():
    assert torch.cuda.is_available()

    device = get_device()
    devices = get_devices()
    assert device == torch.device("cuda:0")
    assert devices == [torch.device("cuda:0"), torch.device("cuda:1")]

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=1,
        use_gpu=True,
        resources_per_worker={"GPU": 2}
    )
)
trainer.fit()

设置 GPU 类型#

Ray Train 允许你为每个 worker 指定加速器类型。如果你想使用特定类型的加速器进行模型训练，这非常有用。在异构 Ray 集群中，这意味着你的训练 worker 将被迫在指定的 GPU 类型上运行，而不是在任何任意的 GPU 节点上。你可以从可用加速器类型中获取支持的 accelerator_type 列表。

例如，如果你想为每个 worker 分配一个 NVIDIA A100 GPU，可以在 ScalingConfig 中指定 accelerator_type="A100"。

提示

确保你的集群包含具有指定加速器类型的实例，或者能够自动伸缩以满足请求。

ScalingConfig(
    num_workers=1,
    use_gpu=True,
    accelerator_type="A100"
)

(PyTorch) 设置通信后端#

PyTorch 分布式支持多种后端用于在 worker 之间通信张量。默认情况下，当设置 use_gpu=True 时，Ray Train 将使用 NCCL，否则使用 Gloo。

如果你想明确覆盖此设置，可以配置一个 TorchConfig 并将其传递给 TorchTrainer。

from ray.train.torch import TorchConfig, TorchTrainer

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(
        num_workers=num_training_workers,
        use_gpu=True, # Defaults to NCCL
    ),
    torch_config=TorchConfig(backend="gloo"),
)

(NCCL) 设置通信网络接口#

当使用 NCCL 进行分布式训练时，可以通过设置 NCCL_SOCKET_IFNAME 环境变量来配置用于 GPU 之间通信的网络接口卡。

为了确保为所有训练 worker 设置该环境变量，你可以在 Ray 运行时环境中传递它。

import ray

runtime_env = {"env_vars": {"NCCL_SOCKET_IFNAME": "ens5"}}
ray.init(runtime_env=runtime_env)

trainer = TorchTrainer(...)

设置每个 worker 的资源#

如果你想为每个训练 worker 分配多个 CPU 或 GPU，或者你定义了自定义集群资源，请设置 resources_per_worker 属性。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8,
    resources_per_worker={
        "CPU": 4,
        "GPU": 2,
    },
    use_gpu=True,
)

注意

如果你在 resources_per_worker 中指定了 GPU，你还需要设置 use_gpu=True。

你还可以指示 Ray Train 使用小数 GPU。在这种情况下，多个 worker 将被分配到同一个 CUDA 设备。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8,
    resources_per_worker={
        "CPU": 4,
        "GPU": 0.5,
    },
    use_gpu=True,
)

(已弃用) Trainer 资源#

重要

此 API 已弃用。请参阅此迁移指南了解更多详情。

至此，我们配置了每个训练 worker 的资源。从技术上讲，每个训练 worker 都是一个 Ray Actor。当你调用 trainer.fit() 时，Ray Train 也会为 trainer 对象调度一个 actor。

这个对象通常只管理训练 worker 之间的轻量级通信。默认情况下，一个 trainer 使用 1 个 CPU。如果你有一个拥有 8 个 CPU 的集群，并想启动 4 个训练 worker，每个 worker 使用 2 个 CPU，这将无法工作，因为所需的 CPU 总数将是 9 个 (4 * 2 + 1)。在这种情况下，你可以将 trainer 的资源指定为使用 0 个 CPU。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=4,
    resources_per_worker={
        "CPU": 2,
    },
    trainer_resources={
        "CPU": 0,
    }
)