配置 Scale 和 GPUs#

增加 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
)

使用 GPUs#

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

from ray.train import ScalingConfig

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

在训练函数中使用 GPUs#

当设置 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 分配多个 GPUs#

有时您可能希望为 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 Distributed 支持多种 backends 来在 worker 之间通信 tensors。默认情况下，当 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 runtime environment 中传递它。

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 个使用 2 个 CPU 的训练 worker，这将不起作用，因为所需的 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,
    }
)