配置规模和GPU#

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

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

增加工作者的数量#

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

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8
)

使用GPU#

要使用GPU，请将 use_gpu=True 传递给 ScalingConfig。这将请求每个训练工作器使用一个GPU。在下面的示例中，训练将在8个GPU上运行（8个工作器，每个使用一个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()

为工作者分配多个GPU#

有时你可能想为一个工作进程分配多个GPU。例如，如果你想为每个工作进程分配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 允许你为每个工作器指定加速器类型。如果你想为模型训练使用特定的加速器类型，这非常有用。在一个异构的 Ray 集群中，这意味着你的训练工作器将被强制运行在指定的 GPU 类型上，而不是任意一个 GPU 节点上。你可以从可用的加速器类型获取支持的 accelerator_type 列表。

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

小技巧

确保您的集群具有指定加速器类型的实例，或者能够自动扩展以满足请求。

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

(PyTorch) 设置通信后端#

PyTorch Distributed 支持多个后端用于在工作者之间通信张量。默认情况下，当 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 间通信的网络接口卡。

为了确保所有训练工作者的环境变量都已设置，您可以在 Ray 运行时环境中传递它：

import ray

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

trainer = TorchTrainer(...)

设置每个工作者的资源#

如果你想为每个训练工作器分配超过一个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。在这种情况下，多个工作进程将被分配到同一个 CUDA 设备。

from ray.train import ScalingConfig

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

培训资源#

到目前为止，我们已经为每个训练工作者配置了资源。从技术上讲，每个训练工作者都是一个 Ray 角色。当你调用 Trainer.fit() 时，Ray Train 还会为 Trainer 对象调度一个角色。

这个对象通常只管理训练工作者之间的轻量级通信。你仍然可以指定它的资源，如果你实现了一个进行更重处理的自己的训练器，这可能会很有用。

from ray.train import ScalingConfig

scaling_config = ScalingConfig(
    num_workers=8,
    trainer_resources={
        "CPU": 4,
        "GPU": 1,
    }
)

默认情况下，训练器使用1个CPU。如果你有一个拥有8个CPU的集群，并且想要启动4个每个使用2个CPU的训练工作器，这将无法工作，因为所需的CPU总数将是9（4 * 2 + 1）。在这种情况下，你可以指定训练器资源使用0个CPU：

from ray.train import ScalingConfig

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