多 GPU 基础

跨多个 GPU 分布可以极大地提高训练速度。然而，即使在单台机器上，这也不是默认设置。要启用多 GPU 训练，我们强烈建议您使用分布式数据并行（DDP）。

使用分布式数据并行 (DDP) 进行多 GPU 训练

DDP 通过为每个 GPU 启动一个进程来实现数据并行。如果需要，DDP 允许您将工作分配到同一台机器上或网络上的多台机器上的 GPU。

使用 CUDA 时（本文档中假定如此），PyTorch 在后台使用 NCCL 同步所有内容。PyTorch 文档进一步详细介绍了分布式后端。

在 SpeechBrain 中编写 DDP 安全代码

DDP 要求您的训练例程编写为 DDP 安全的，因为您的脚本将并发运行多次（可能跨多台机器）。标准的 SpeechBrain recipe 将与 DDP 一起工作。我们还提供了有助于编写 DDP 安全脚本的功能。

run_on_main 确保特定函数仅在一个进程中执行一次，强制其他进程等待。它经常用于在 recipe 中运行数据集准备步骤。

许多函数（如 Brain.fit）都编写为 DDP 感知的。实际上，您无需做太多即可使您的代码 DDP 安全，但这是您应该记住的一点。

注意：使用 DDP 时，批处理大小是为单个进程/GPU 定义的。这与数据并行 (DP) 不同，数据并行根据 GPU 数量分割批处理。例如，使用 DDP 时，如果您指定批处理大小为 16，则每个 GPU/进程都将使用大小为 16 的批处理，无论您有多少个 GPU。

单节点设置

这涵盖了您希望在单台机器（节点）上将训练分布到多个 GPU 的情况。

使用 SpeechBrain，这将看起来像

cd recipes/<dataset>/<task>/
torchrun --standalone --nproc_per_node=4 experiment.py hyperparams.yaml

... 其中 nproc_per_node 是要启动的进程/要使用的 GPU 数量。

多节点设置

这涵盖了您希望在网络上的多台机器上将训练分布的情况，每台机器可以有任意数量的 GPU。

请注意，跨多台机器使用 DDP 会引入通信开销，这可能会显著减慢训练速度，有时甚至比在单个节点上训练还要慢！这很大程度上取决于节点之间的网络速度。请确保您确实从跨机器分配工作中获得了任何好处。

虽然 DDP 比 DataParallel 更高效，但它有时容易出现意想不到的 bug。DDP 相当依赖于服务器设置，因此某些设置可能会遇到问题。如果您遇到问题，请确保 PyTorch 已更新到最新版本。

基础与手动多节点设置

让我们从一个简单的示例开始，用户能够直接连接到每个节点。假设我们有 2 个节点，每个节点有 2 个 GPU（总共 4 个 GPU）。

我们在每台机器上使用 torchrun 一次，参数如下

  • --nproc_per_node=2 意味着我们将为每个节点启动 2 个进程，即每个节点使用 2 个 GPU。

  • --nnodes=2 意味着我们将总共使用两个节点。

  • --node_rank=0 和 --node_rank=1 指代我们分配给节点/机器的排名/“索引”。

  • --master_addr/--master_port 定义了“主”机器的 IP 地址和端口。在这种情况下，我们任意选择第一台机器作为其他所有机器（我们案例中的第二台机器）的“主节点”。请注意，如果您不走运或在该节点上运行多个不同的训练脚本，5555 端口可能被其他进程占用，因此您可能需要选择一个不同的空闲端口。

因此，我们得到

# Machine 1
cd recipes/<dataset>/<task>/
torchrun --nproc_per_node=2 --nnodes=2 --node_rank=0 --master_addr machine_1_address --master_port 5555 experiment.py hyperparams.yaml

# Machine 2
cd recipes/<dataset>/<task>/
torchrun --nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr machine_1_address --master_port 5555 experiment.py hyperparams.yaml

在此设置中

  • 机器 1 将有 2 个子进程

    • 子进程 #1: local_rank=0, rank=0

    • 子进程 #2: local_rank=1, rank=1

  • 机器 2 将有 2 个子进程

    • 子进程 #1: local_rank=0, rank=2

    • 子进程 #2: local_rank=1, rank=3

实际上，使用 torchrun 可以确保设置正确的环境变量（LOCAL_RANK 和 RANK），因此您无需为此烦恼。

使用 Slurm 进行多节点设置

如果您可以访问使用 Slurm 的计算集群，则可以自动化此过程。我们将创建两个脚本

  • 一个 SBATCH 脚本，它将请求节点配置并调用第二个脚本。

  • 一个 SRUN 脚本，它将在每个节点上调用训练。

sbatch.sh:

#SBATCH --nodes=2 # We want two nodes (servers)
#SBATCH --ntasks-per-node=1 # we will run once the next srun per node
#SBATCH --gres=gpu:4 # we want 4 GPUs per node #cspell:ignore gres
#SBATCH --job-name=SBisSOcool
#SBATCH --cpus-per-task=10 # the only task will request 10 cores
#SBATCH --time=20:00:00 # Everything will run for 20H.

# We jump into the submission dir
cd ${SLURM_SUBMIT_DIR}

# And we call the srun that will run --ntasks-per-node times (once here) per node
srun srun_script.sh

srun_script.sh:

#!/bin/bash

# We jump into the submission dir
cd ${SLURM_SUBMIT_DIR}

# We activate our env
conda activate super_cool_sb_env

# We extract the master node address (the one that every node must connects to)
LISTNODES=`scontrol show hostname $SLURM_JOB_NODELIST`
MASTER=`echo $LISTNODES | cut -d" " -f1`

# here --nproc_per_node=4 because we want torchrun to spawn 4 processes (4 GPUs). Then we give the total amount of nodes requested (--nnodes) and then --node_rank that is necessary to dissociate the node that we are calling this from.
torchrun --nproc_per_node=4 --nnodes=${SLURM_JOB_NUM_NODES} --node_rank=${SLURM_NODEID} --master_addr=${MASTER} --master_port=5555 train.py hparams/myrecipe.yaml

(已弃用) 使用数据并行进行单节点多 GPU 训练

我们强烈建议不要使用 DataParallel，即使是用于单节点设置！请改用 DistributedDataParallel。我们不再为 DataParallel 提供支持。未来的 PyTorch 版本甚至可能完全删除 DataParallel。

在单台机器上使用数据并行进行多 GPU 训练的常见模式是

cd recipes/<dataset>/<task>/
python experiment.py params.yaml --data_parallel_backend

如果您想使用特定的 GPU 设备集，请考虑如下使用 CUDA_VISIBLE_DEVICES

cd recipes/<dataset>/<task>/
CUDA_VISIBLE_DEVICES=1,5 python experiment.py params.yaml --data_parallel_backend

重要提示：每个 GPU 进程的批处理大小将为：batch_size / GPU 数量。因此，您应该根据需要考虑更改 batch_size 值。