Transformers 文档

在多个GPU上进行高效训练

多GPU上的高效训练

如果在单个GPU上训练模型太慢,或者模型的权重无法适应单个GPU的内存,那么转向多GPU设置可能是一个可行的选择。在做出这一转变之前,请彻底探索在单个GPU上进行高效训练的方法和工具中涵盖的所有策略,因为它们适用于任何数量的GPU上的模型训练。一旦你采用了这些策略并发现它们在单个GPU上对你的情况不足,考虑转向多个GPU。

从单个GPU过渡到多个GPU需要引入某种形式的并行性,因为工作负载必须分布在资源之间。可以采用多种技术来实现并行性,例如数据并行性、张量并行性和管道并行性。需要注意的是,没有一种适用于所有情况的解决方案,最佳设置取决于您使用的特定硬件配置。

本指南提供了对各种并行类型的深入概述,以及如何结合
技术和选择适当方法的指导。有关分布式训练的逐步教程,请参阅 🤗 Accelerate 文档

虽然本指南中讨论的主要概念可能适用于各种框架,但在这里我们主要关注基于PyTorch的实现。

在深入了解每种技术的具体细节之前,让我们先回顾一下在大型基础设施上训练大型模型时的粗略决策过程。

可扩展性策略

首先估计训练模型需要多少vRAM。对于托管在🤗 Hub上的模型,请使用我们的 模型内存计算器,它可以在几个百分点的误差范围内提供准确的估算。

单节点/多GPU设置的并行化策略

在单个节点上使用多个GPU训练模型时,您选择的并行化策略可以显著影响性能。以下是您的选项的详细说明:

案例1:您的模型适合单个GPU

如果你的模型可以轻松适应单个GPU,你有两个主要选择:

  1. DDP - 分布式数据并行
  2. Zero Redundancy Optimizer (ZeRO) - 根据使用的情况和配置,这种方法可能会更快,也可能不会,但值得尝试。

案例2:您的模型无法适应单个GPU:

如果你的模型对于单个GPU来说太大,你有几种替代方案可以考虑:

  1. 管道并行 (PP)
  2. ZeRO
  3. TensorParallel (TP)

在节点间连接非常快的情况下(例如,NVLINK或NVSwitch),所有三种策略(PP、ZeRO、TP)应该会带来相似的性能。然而,如果没有这些,PP将比TP或ZeRO更快。TP的程度也可能产生影响。最好通过实验来确定最适合您特定设置的策略。

TP 几乎总是在单个节点内使用。也就是说,TP 大小 <= 每个节点的 GPU 数量。

案例3:模型中最大的层无法适应单个GPU

  1. 如果您没有使用ZeRO,您必须使用TensorParallel (TP),因为仅使用PipelineParallel (PP)不足以容纳大层。
  2. 如果您正在使用ZeRO,另外采用单GPU高效训练的方法和工具中的技术。

多节点/多GPU设置的并行化策略

  • 当您拥有快速的节点间连接(例如,NVLINK 或 NVSwitch)时,请考虑使用以下选项之一:

    1. ZeRO - as it requires close to no modifications to the model
    2. A combination of PipelineParallel(PP) with TensorParallel(TP) and DataParallel(DP) - this approach will result in fewer communications, but requires significant changes to the model
  • 当节点间连接速度较慢且GPU内存仍然不足时:

    1. Employ a combination of DataParallel(DP) with PipelineParallel(PP), TensorParallel(TP), and ZeRO.

在本指南的以下部分中,我们将深入探讨这些不同的并行方法的工作原理。

数据并行

即使只有2个GPU,你也可以轻松利用PyTorch内置功能提供的加速训练能力,例如DataParallel (DP) 和 DistributedDataParallel (DDP)。请注意,PyTorch文档建议在多GPU训练中优先使用DistributedDataParallel (DDP) 而不是 DataParallel (DP),因为它适用于所有模型。让我们来看看这两种方法的工作原理以及它们之间的区别。

DataParallel 与 DistributedDataParallel

为了理解两种方法在GPU间通信开销上的关键差异,让我们回顾一下每批次的处理过程:

DDP:

  • 在开始时,主进程将模型从GPU 0复制到其他GPU
  • 然后对于每个批次:
    1. 每个GPU直接消耗其小批量数据。
    2. backward期间,一旦本地梯度准备就绪,它们将在所有进程之间进行平均。

DP:

对于每个批次:

  1. GPU 0 读取数据批次,然后将每个小批次发送到每个 GPU。
  2. 最新的模型从GPU 0复制到每个GPU。
  3. forward 被执行,每个 GPU 的输出被发送到 GPU 0 以计算损失。
  4. 损失从GPU 0分配到所有GPU,并运行backward
  5. 每个GPU的梯度被发送到GPU 0并进行平均。

主要区别包括:

  1. DDP 每批次仅执行一次通信 - 发送梯度,而 DP 每批次执行五次不同的数据交换。 DDP 使用 torch.distributed 复制数据,而 DP 通过 Python 线程在进程内复制数据(这引入了与 GIL 相关的限制)。因此,DistributedDataParallel (DDP) 通常比 DataParallel (DP) 更快,除非你的 GPU 卡之间的连接速度较慢。
  2. 在DP下,GPU 0执行的工作量显著多于其他GPU,导致GPU利用率不足。
  3. DDP 支持跨多台机器的分布式训练,而 DP 不支持。

这并不是DP和DDP之间差异的详尽列表,然而,其他细微差别超出了本指南的范围。 您可以通过阅读这篇文章来更深入地了解这些方法。

让我们通过一个实验来说明DP和DDP之间的区别。我们将在NVLink存在的背景下对DP和DDP之间的差异进行基准测试:

  • 硬件:2块TITAN RTX 24GB,每块通过NVlink连接,使用2个NVLinks(在nvidia-smi topo -m中显示为NV2)。
  • 软件: pytorch-1.8-to-be + cuda-11.0 / transformers==4.3.0.dev0.

要禁用其中一个基准测试中的NVLink功能,我们使用NCCL_P2P_DISABLE=1

以下是基准测试代码和输出:

DP

rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \
python examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 110.5948, 'train_samples_per_second': 1.808, 'epoch': 0.69}

DDP 带 NVlink

rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \
torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 101.9003, 'train_samples_per_second': 1.963, 'epoch': 0.69}

无NVlink的DDP

rm -r /tmp/test-clm; NCCL_P2P_DISABLE=1 CUDA_VISIBLE_DEVICES=0,1 \
torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 131.4367, 'train_samples_per_second': 1.522, 'epoch': 0.69}

为了方便起见,以下是收集在表格中的相同基准测试结果:

类型 NVlink 时间
2:DP Y 110秒
2:DDP Y 101秒
2:DDP N 131秒

如你所见,在这种情况下,DP比使用NVlink的DDP慢约10%,但比不使用NVlink的DDP快约15%。 实际差异将取决于每个GPU需要与其他GPU同步的数据量——需要同步的数据越多,慢速链接对整体运行时间的影响就越大。

ZeRO 数据并行

ZeRO驱动的数据并行(ZeRO-DP)在以下图表中展示,该图表来自这篇博客文章

DeepSpeed-Image-1

虽然它可能看起来很复杂,但它与DataParallel (DP) 的概念非常相似。不同之处在于,每个GPU不是复制完整的模型参数、梯度和优化器状态,而是只存储其中的一部分。然后,在运行时,当给定层需要完整的层参数时,所有GPU会同步以互相提供它们缺失的部分。

为了说明这个想法,考虑一个具有3层(La、Lb和Lc)的简单模型,其中每层有3个参数。 例如,层La有权重a0、a1和a2:

La | Lb | Lc
---|----|---
a0 | b0 | c0
a1 | b1 | c1
a2 | b2 | c2

如果我们有3个GPU,ZeRO-DP会将模型分割到3个GPU上,如下所示:

GPU0:
La | Lb | Lc
---|----|---
a0 | b0 | c0

GPU1:
La | Lb | Lc
---|----|---
a1 | b1 | c1

GPU2:
La | Lb | Lc
---|----|---
a2 | b2 | c2

在某种程度上,这与张量并行中的水平切片相同,与垂直切片相反,垂直切片是将整个层组放在不同的GPU上。现在让我们看看这是如何工作的:

每个GPU在DP中工作时都会获得通常的小批量数据:

x0 => GPU0
x1 => GPU1
x2 => GPU2

输入未经修改地传递,就像它们将由原始模型处理一样。

首先,输入到达层 La。此时会发生什么?

在GPU0上:x0小批量需要a0、a1、a2参数来完成其通过层的正向路径,但GPU0只有a0。 它将从GPU1获取a1,从GPU2获取a2,将模型的所有部分组合在一起。

同时,GPU1 获取另一个小批量数据 - x1。GPU1 拥有 a1 参数,但需要 a0 和 a2,因此它从 GPU0 和 GPU2 获取这些参数。 同样的情况发生在 GPU2 上,它获取小批量数据 x2。它从 GPU0 和 GPU1 获取 a0 和 a1。

这样,每个GPU都会获得完整的张量重建,并使用自己的小批量进行前向传递。 一旦计算完成,不再需要的数据就会被丢弃——它只在计算期间使用。 重建是通过预取高效完成的。

然后整个过程在层Lb中重复,接着是Lc向前,然后向后Lc -> Lb -> La。

这种机制类似于一种高效的团队背包策略:A携带帐篷,B携带炉子,C携带斧头。每天晚上他们都会与其他人分享他们所拥有的东西,并从其他人那里得到他们没有的东西,早上他们打包分配好的装备类型并继续前进。这就是ZeRO DP/Sharded DDP。将这种策略与每个人必须携带自己的帐篷、炉子和斧头的简单策略(类似于PyTorch中的DataParallel(DP和DDP))进行比较,后者效率要低得多。

在阅读关于这个主题的文献时,你可能会遇到以下同义词:分片、分区。 如果你仔细观察ZeRO如何分区模型的权重 - 它看起来与张量并行非常相似 这将在后面讨论。这是因为它分区/分片每一层的权重,与接下来讨论的垂直模型并行不同。

实现:

从朴素模型并行到流水线并行

为了解释管道并行性,我们首先来看一下朴素模型并行性(MP),也称为垂直MP。这种方法涉及通过使用.to()将特定层分配给特定GPU,从而将模型层组分布在多个GPU上。当数据流经这些层时,它会被移动到与层相同的GPU,而其他层保持不变。

我们将这种模型并行称为“垂直”并行,因为模型通常是这样可视化的。例如,下图显示了一个8层模型垂直分割成两部分,将第0-3层放在GPU0上,第4-7层放在GPU1上:

================
| Layer |      |
|   0   |      |
|   1   | GPU0 |
|   2   |      |
|   3   |      |
================
| Layer |      |
|   4   |      |
|   5   | GPU1 |
|   6   |      |
|   7   |      |
================

在这个例子中,当数据从第0层移动到第3层时,它与常规的前向传递没有区别。然而,将数据从第3层传递到第4层需要将其从GPU0移动到GPU1,这会引入通信开销。如果参与的GPU位于同一计算节点上(例如同一台物理机器),这种复制是快速的,但如果GPU分布在不同的计算节点上(例如多台机器),通信开销可能会显著增加。

随后,第4层到第7层的工作方式与原始模型中的相同。在第7层完成后,通常需要将数据发送回第0层,那里有标签(或者将标签发送到最后一层)。现在可以计算损失,优化器可以开始工作。

朴素模型并行存在几个缺点:

  • 在任何给定时刻,除了一个GPU外,其他GPU都处于空闲状态:如果使用了4个GPU,这几乎等同于将单个GPU的内存容量增加了四倍,而忽略了其他硬件。
  • 设备间数据传输的开销:例如,使用简单的MP(模型并行),4张6GB的显卡可以容纳与1张24GB显卡相同大小的模型,但单张24GB显卡将更快完成训练,因为它没有数据复制的开销。但是,如果你有40GB的显卡并且需要适应一个45GB的模型,你可以使用4张40GB的显卡(但由于梯度和优化器状态的存在,这几乎是不可能的)
  • 复制共享嵌入: 共享嵌入可能需要在GPU之间来回复制。

现在你已经熟悉了模型并行的朴素方法及其缺点,让我们来看看管道并行(PP)。 PP与朴素MP几乎相同,但它通过将传入的批次分成微批次并人为地创建一个管道来解决GPU闲置问题, 这使得不同的GPU可以同时参与计算过程。

以下插图来自GPipe论文,展示了顶部的朴素MP和底部的PP:

MP vs PP

在图表的底部,您可以观察到管道并行(PP)方法最小化了空闲的GPU区域,这些区域被称为“气泡”。图表的两个部分都显示了并行度为4,这意味着有4个GPU参与了管道。您可以看到有4个管道阶段的前向路径(F0、F1、F2和F3),随后是反向顺序的后向路径(B3、B2、B1和B0)。

PP 引入了一个新的超参数来调整 - chunks,它决定了有多少数据块通过同一个管道阶段按顺序发送。例如,在下面的图表中,你可以看到 chunks=4。GPU0 对块 0、1、2 和 3 执行相同的前向路径(F0,0, F0,1, F0,2, F0,3),然后等待其他 GPU 完成它们的工作。只有当其他 GPU 开始完成它们的工作时,GPU0 才会再次开始工作,对块 3、2、1 和 0 执行反向路径(B0,3, B0,2, B0,1, B0,0)。

请注意,这与梯度累积步骤的概念相同。PyTorch 使用 chunks,而 DeepSpeed 将相同的超参数称为梯度累积步骤。

由于分块的存在,PP引入了微批次(MBS)的概念。DP将全局数据批次大小分割成小批次,因此如果你有4个DP度,一个全局批次大小为1024的数据将被分割成4个小批次,每个批次大小为256(1024/4)。如果chunks(或GAS)的数量是32,我们最终得到的微批次大小为8(256/32)。每个管道阶段一次处理一个微批次。要计算DP + PP设置的全局批次大小,请使用公式:mbs * chunks * dp_degree8 * 32 * 4 = 1024)。 使用chunks=1时,你将得到原始的MP,这是低效的。使用较大的chunks值时,你将得到非常小的微批次大小,这也是低效的。因此,我们鼓励尝试不同的chunks值,以找到最能有效利用GPU的值。

你可能会注意到图表上有一个“死”时间的气泡,这部分时间无法并行化,因为最后一个forward阶段必须等待backward完成管道。找到chunks的最佳值的目的是在所有参与的GPU上实现高并发的GPU利用率,这相当于最小化气泡的大小。

Pipeline API 解决方案已在以下地方实施:

  • PyTorch
  • DeepSpeed
  • Megatron-LM

这些带有一些缺点:

  • 他们必须对模型进行大量修改,因为Pipeline要求将模块的正常流程重写为nn.Sequential序列,这可能需要改变模型的设计。
  • 目前,Pipeline API 非常受限。如果你有一堆 Python 变量在 Pipeline 的第一阶段传递,你将不得不找到一种解决方法。目前,Pipeline 接口要求输入和输出只能是单个张量或张量元组。这些张量的第一维必须是批量大小,因为 Pipeline 会将小批量分成微批量。可能的改进正在这里讨论 https://github.com/pytorch/pytorch/pull/50693
  • 在管道阶段的级别上无法实现条件控制流 - 例如,像T5这样的编码器-解码器模型需要特殊的解决方法以处理条件编码器阶段。
  • 他们必须安排每一层,使得一层的输出成为另一层的输入。

最近的解决方案包括:

  • 瓦鲁纳
  • Sagemaker

我们尚未对Varuna和SageMaker进行实验,但他们的论文报告称,他们已经克服了上述提到的问题,并且对用户模型的改动较小。

实现:

  • PyTorch(在pytorch-1.8中初步支持,并在1.9中逐步改进,在1.10中更是如此)。一些示例
  • DeepSpeed
  • Megatron-LM 有一个内部实现 - 没有API。
  • Varuna
  • SageMaker - 这是一个专有解决方案,只能在AWS上使用。
  • OSLO - 这是基于 Hugging Face Transformers 实现的。

🤗 Transformers 状态:截至撰写本文时,没有任何模型支持完整的PP。GPT2和T5模型具有简单的MP支持。 主要障碍是无法将模型转换为nn.Sequential并使所有输入为张量。这是因为目前模型包含许多功能,使得转换非常复杂,需要移除这些功能才能实现。

DeepSpeed 和 Megatron-LM 的集成可在 🤗 Accelerate 中找到

其他方法:

DeepSpeed、Varuna 和 SageMaker 使用了 Interleaved Pipeline 的概念

Interleaved pipeline execution

在这里,通过优先考虑反向传递,进一步最小化了气泡(空闲时间)。Varuna 进一步尝试通过使用模拟来发现最有效的调度,从而改进调度。

OSLO 实现了基于 Transformers 的管道并行,无需进行 nn.Sequential 转换。

张量并行

在张量并行中,每个GPU处理张量的一部分,并且仅在需要时聚合完整的张量以进行操作。 为了描述这种方法,本指南的这一部分依赖于Megatron-LM 论文中的概念和图表:在GPU集群上高效的大规模语言模型训练

任何变压器的主要构建块都是一个全连接的nn.Linear,后面跟着一个非线性激活GeLU。 根据Megatron论文的符号表示,其点积部分可以写成Y = GeLU(XA),其中X是 输入向量,Y是输出向量,A是权重矩阵。

如果我们以矩阵形式查看计算,你可以看到矩阵乘法如何在多个GPU之间拆分:

Parallel GEMM

如果我们将权重矩阵 A 按列分割到 N 个 GPU 上,并并行执行矩阵乘法 XA_1XA_n, 那么我们将得到 N 个输出向量 Y_1, Y_2, ..., Y_n,这些向量可以独立地输入到 GeLU 中:

Independent GeLU

利用这一原理,我们可以更新任意深度的多层感知器,而不需要在GPU之间进行任何同步,直到最后,我们需要从分片中重建输出向量。Megatron-LM论文的作者为此提供了一个有用的图示:

Parallel shard processing

并行化多头注意力层甚至更简单,因为它们本身已经是并行的,因为有多个独立的头!

Parallel self-attention

特别注意事项:TP需要非常快的网络,因此不建议在超过一个节点上进行TP。 实际上,如果一个节点有4个GPU,那么最高的TP度数为4。如果你需要8的TP度数,你需要使用至少拥有8个GPU的节点。

本节基于原始的更详细的TP概述。 由@anton-l提供。

替代名称:

实现:

  • Megatron-LM 有一个内部实现,因为它非常特定于模型
  • parallelformers (目前仅支持推理)
  • SageMaker - 这是一个专有解决方案,只能在AWS上使用。
  • OSLO 具有基于 Transformers 的张量并行实现。

SageMaker 将 TP 与 DP 结合以实现更高效的处理。

🤗 Transformers 状态:

  • 核心:尚未在核心中实现
  • 但如果你想要推理,parallelformers为我们的大多数模型提供了这种支持。因此,在核心中实现此功能之前,你可以使用他们的。希望训练模式也能得到支持。
  • Deepspeed-Inference 还支持我们的 BERT、GPT-2 和 GPT-Neo 模型在其超快的基于 CUDA 内核的推理模式下运行,更多信息请参见 这里

🤗 Accelerate 集成了 TP from Megatron-LM

数据并行 + 管道并行

以下图表来自DeepSpeed的管道教程,展示了如何将DP与PP结合使用。

DP + PP-2d

这里重要的是要看到DP rank 0看不到GPU2,DP rank 1看不到GPU3。对于DP来说,只有GPU0和GPU1,它像只有2个GPU一样提供数据。GPU0“秘密地”使用PP将其部分负载卸载到GPU2。GPU1也通过请求GPU3的帮助来执行相同的操作。

由于每个维度至少需要2个GPU,因此这里至少需要4个GPU。

实现:

🤗 Transformers 状态:尚未实现

数据并行 + 管道并行 + 张量并行

为了获得更高效的训练,使用了3D并行性,其中PP与TP和DP结合使用。这可以在下图中看到。

dp-pp-tp-3d

此图来自一篇博客文章3D并行:扩展到万亿参数模型,这也是一篇值得一读的文章。

由于每个维度至少需要2个GPU,因此这里至少需要8个GPU。

实现:

🤗 Transformers 状态:尚未实现,因为我们没有 PP 和 TP。

ZeRO 数据并行 + 管道并行 + 张量并行

DeepSpeed 的主要特性之一是 ZeRO,它是 DP 的一个超级可扩展扩展。它已经在 ZeRO 数据并行 中讨论过。通常它是一个独立的功能,不需要 PP 或 TP。但它可以与 PP 和 TP 结合使用。

当ZeRO-DP与PP(以及可选的TP)结合时,通常只启用ZeRO阶段1(优化器分片)。

虽然在理论上可以将ZeRO阶段2(梯度分片)与管道并行性结合使用,但这会对性能产生负面影响。每个微批次都需要一个额外的reduce-scatter集合来在分片之前聚合梯度,这会增加潜在的显著通信开销。由于管道并行性的本质,使用小的微批次,并且重点是尝试平衡算术强度(微批次大小)与最小化管道气泡(微批次数量)。因此,这些通信成本将影响性能。

此外,由于PP的存在,层数已经比正常情况下少,因此内存节省不会很大。PP已经将梯度大小减少了1/PP,因此在此基础上进行梯度分片节省的效果不如纯DP显著。

出于同样的原因,ZeRO 阶段 3 也不是一个好的选择 - 需要更多的节点间通信。

由于我们有ZeRO,另一个好处是ZeRO-Offload。由于这是第1阶段,优化器状态可以卸载到CPU。

实现:

重要论文:

🤗 Transformers 状态:尚未实现,因为我们没有 PP 和 TP。

FlexFlow

FlexFlow 也以一种稍微不同的方法解决了并行化问题。

论文: “超越数据和模型并行的深度神经网络” by Zhihao Jia, Matei Zaharia, Alex Aiken

它在样本-操作符-属性-参数上执行一种4D并行处理。

  1. 样本 = 数据并行(按样本并行)
  2. 操作符 = 将单个操作并行化为多个子操作
  3. 属性 = 数据并行性(长度方向的并行)
  4. 参数 = 模型并行性(无论维度 - 水平或垂直)

示例:

  • 示例

让我们取10批长度为512的序列。如果我们通过样本维度将它们并行化到2个设备中,我们得到10 x 512,这变成了5 x 2 x 512。

  • 操作符

如果我们执行层归一化,我们首先计算标准差,然后计算均值,接着我们可以对数据进行归一化。 操作符并行性允许并行计算标准差和均值。因此,如果我们通过操作符维度将它们并行化到2个设备(cuda:0, cuda:1)上,首先我们将输入数据复制到两个设备上,然后cuda:0计算标准差,cuda:1同时计算均值。

  • 属性

我们有10批长度为512的数据。如果我们按属性维度将它们并行化到2个设备上,10 x 512将变为10 x 2 x 256。

  • 参数

它与张量模型并行或简单的逐层模型并行类似。

flex-flow-soap

该框架的意义在于,它能够处理诸如(1)GPU/TPU/CPU与(2)RAM/DRAM与(3)快速内部连接/慢速外部连接等资源,并自动优化所有这些算法,决定在何处使用哪种并行化。

一个非常重要的方面是,FlexFlow 是为优化具有静态和固定工作负载的模型的 DNN 并行化而设计的,因为具有动态行为的模型可能在不同的迭代中偏好不同的并行化策略。

因此,这个承诺非常吸引人——它在选择的集群上运行30分钟的模拟,并提出利用这个特定环境的最佳策略。如果你添加/删除/替换任何部分,它将运行并重新优化该计划。然后你就可以进行训练。不同的设置将有其自己的自定义优化。

🤗 Transformers 状态:Transformers 模型可以通过 transformers.utils.fx 进行 FX 跟踪,这是 FlexFlow 的前提条件,但需要在 FlexFlow 端进行更改才能使其与 Transformers 模型一起工作。

GPU 选择

在多个GPU上进行训练时,您可以指定要使用的GPU数量及其顺序。例如,当您拥有不同计算能力的GPU并希望首先使用更快的GPU时,这可能会很有用。选择过程适用于DistributedDataParallelDataParallel,以便仅使用可用GPU的子集,并且您不需要Accelerate或DeepSpeed集成

GPU数量

例如,如果您有4个GPU,但只想使用前2个:

torchrun
Accelerate
DeepSpeed

使用 --nproc_per_node 来选择使用多少个GPU。

torchrun --nproc_per_node=2  trainer-program.py ...

GPU的顺序

现在,要选择使用哪些GPU及其顺序,您将使用CUDA_VISIBLE_DEVICES环境变量。最简单的方法是在~/bashrc或其他启动配置文件中设置环境变量。CUDA_VISIBLE_DEVICES用于映射使用的GPU。例如,如果您有4个GPU(0, 1, 2, 3),并且只想运行GPU 0和2:

CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

只有2个物理GPU(0和2)对PyTorch“可见”,并且它们分别映射到cuda:0cuda:1。你也可以反转GPU的顺序,首先使用2。现在,映射是GPU 0为cuda:1,GPU 2为cuda:0

CUDA_VISIBLE_DEVICES=2,0 torchrun trainer-program.py ...

你也可以将CUDA_VISIBLE_DEVICES环境变量设置为空值,以创建一个没有GPU的环境。

CUDA_VISIBLE_DEVICES= python trainer-program.py ...

与任何环境变量一样,它们可以被导出,而不是添加到命令行中。然而,这不推荐,因为如果你忘记了环境变量是如何设置的,最终可能会使用错误的GPU,这可能会造成混淆。相反,通常的做法是在同一命令行上为特定的训练运行设置环境变量。

CUDA_DEVICE_ORDER 是一个你可以用来控制GPU排序的替代环境变量。你可以通过以下方式对它们进行排序:

  1. PCIe总线ID与nvidia-smirocm-smi的顺序相匹配,分别用于NVIDIA和AMD GPU
export CUDA_DEVICE_ORDER=PCI_BUS_ID
  1. GPU计算能力
export CUDA_DEVICE_ORDER=FASTEST_FIRST

如果您的训练设置包含一个较旧和一个较新的GPU,且较旧的GPU首先出现,但您无法物理交换卡以使较新的GPU首先出现,那么CUDA_DEVICE_ORDER尤其有用。在这种情况下,设置CUDA_DEVICE_ORDER=FASTEST_FIRST以始终首先使用较新且更快的GPU(nvidia-smirocm-smi仍然按照它们的PCIe顺序报告GPU)。或者您也可以设置export CUDA_VISIBLE_DEVICES=1,0

< > Update on GitHub