多GPU基础

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

使用分布式数据并行（DDP）进行多GPU训练

DDP通过为每个GPU生成一个进程来实现数据并行。DDP允许你在同一台机器上的GPU之间或网络上的多台机器之间分配工作。

在使用CUDA时（本文档中我们将假设使用CUDA），PyTorch在后台使用NCCL来同步所有内容。PyTorch文档进一步详细说明了分布式后端。

在SpeechBrain中编写DDP安全代码

DDP要求您的训练程序编写为DDP安全的，因为您的脚本将同时运行多次（可能跨多台机器）。标准的SpeechBrain配方将与DDP一起工作。我们还提供了帮助编写DDP安全脚本的功能。

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

许多像Brain.fit这样的函数被编写为DDP感知的。实际上，要使你的代码DDP安全并不需要做很多工作，但这是你应该记住的事情。

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

单节点设置

这涵盖了您希望在单个机器（节点）上跨多个GPU进行训练的情况。

使用SpeechBrain，这将看起来像：

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

… 其中 nproc_per_node 是要生成的进程数/要使用的GPU数量。

多节点设置

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

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

虽然DDP比DataParallel更高效，但它有时容易出现意外的错误。DDP非常依赖于服务器，因此某些设置可能会遇到问题。如果您遇到问题，请确保PyTorch是最新的。

基础与手动多节点设置

让我们从一个简单的例子开始，用户可以连接到每个节点。假设我们有两个节点，每个节点有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 / Number of GPUs。因此，您应根据需要考虑更改batch_size值。