在多个 GPU 上进行有效训练
如果在一个 GPU 上训练模型速度太慢,或者模型的权重无法放入单个 GPU 的内存中,那么转向多 GPU 设置可能是一个可行的选择。在进行此转换之前,请彻底探索在单个 GPU 上进行高效训练的方法和工具中介绍的所有策略,因为这些策略普遍适用于任何数量的 GPU 上的模型训练。一旦您已经采用了这些策略并发现它们对于您在单个 GPU 上的情况不足,请考虑迁移到多个 GPU。
从单个 GPU 迁移到多个 GPU 需要引入某种形式的并行性,因为工作负载必须分布在资源上。可以使用多种技术来实现并行性,例如数据并行、张量并行和流水线并行。重要的是要注意,没有一劳永逸的解决方案,最佳设置取决于您使用的特定硬件配置。
本指南深入介绍了各种并行类型的概述,以及如何组合
技术和选择合适的方法的指导。有关分布式训练的分步教程,请参考🤗 Accelerate 文档。
虽然本指南中讨论的主要概念可能适用于所有框架,但这里我们将重点介绍基于 PyTorch 的实现。
在更深入地了解每种技术的具体细节之前,让我们先了解一下在大型基础设施上训练大型模型时的粗略决策过程。
可扩展性策略
首先估计训练模型需要多少 vRAM。对于托管在 🤗 Hub 上的模型,请使用我们的模型内存计算器,它可以在几百分比的误差范围内为您提供准确的计算结果。
单个节点/多 GPU 设置的并行化策略
在使用多个 GPU 的单个节点上训练模型时,您选择的并行化策略会对性能产生重大影响。以下是您的选项的细分
情况 1:您的模型适合单个 GPU
如果您的模型可以舒适地放入单个 GPU,您有两个主要选择
- DDP - 分布式数据并行
- 零冗余优化器 (ZeRO) - 根据使用情况和配置,此方法可能更快也可能更慢,但值得尝试。
情况 2:您的模型不适合单个 GPU
如果您的模型太大而无法放入单个 GPU,您有几种备选方案可供考虑
如果节点间连接速度很快(例如,NVLINK 或 NVSwitch),这三种策略(PP、ZeRO、TP)都应该产生类似的性能。但是,如果没有这些连接,PP 将比 TP 或 ZeRO 更快。TP 的程度也可能会有所不同。最好针对您的特定设置进行实验,以确定最合适的策略。
TP 几乎总是用在一个节点内。也就是说 TP 大小 <= 每个节点的 GPU 数量。
情况 3:模型的最大层无法放入单个 GPU
- 如果您没有使用 ZeRO,则必须使用张量并行 (TP),因为仅使用流水线并行 (PP) 无法容纳大型层。
- 如果您正在使用 ZeRO,请另外采用来自在单个 GPU 上进行高效训练的方法和工具中的技术。
多节点/多 GPU 设置的并行化策略
当您具有快速节点间连接(例如 NVLINK 或 NVSwitch)时,请考虑使用以下选项之一
- ZeRO - 因为它几乎不需要对模型进行修改
- 流水线并行 (PP) 与张量并行 (TP) 和数据并行 (DP) 的组合 - 这种方法将导致更少的通信,但需要对模型进行重大更改
当您具有缓慢的节点间连接并且 GPU 内存仍然不足时
- 采用数据并行 (DP) 与流水线并行 (PP)、张量并行 (TP) 和 ZeRO 的组合。
在本指南的后续部分,我们将更深入地探讨这些不同的并行方法的工作原理。
数据并行
即使只有 2 个 GPU,您也可以轻松利用 PyTorch 内置功能提供的加速训练功能,例如DataParallel
(DP) 和DistributedDataParallel
(DDP)。请注意,PyTorch 文档建议优先使用DistributedDataParallel
(DDP) 而不是DataParallel
(DP) 进行多 GPU 训练,因为它适用于所有模型。让我们来看看这两种方法是如何工作的,以及它们之间的区别。
数据并行与分布式数据并行
为了了解这两种方法之间 GPU 间通信开销的主要区别,让我们回顾一下每个批次的处理过程
DDP:
- 在开始时,主进程从 GPU 0 将模型复制到其他 GPU 上
- 然后,对于每个批次
- 每个 GPU 直接消耗其数据的小批量。
- 在
backward
期间,一旦本地梯度准备好,它们将在所有进程中进行平均。
DP:
对于每个批次
- GPU 0 读取数据批次,然后将一个小批量发送到每个 GPU。
- 最新模型从 GPU 0 复制到每个 GPU。
- 执行
forward
,每个 GPU 的输出被发送到 GPU 0 以计算损失。 - 损失从 GPU 0 分布到所有 GPU,并运行
backward
。 - 来自每个 GPU 的梯度被发送到 GPU 0 并进行平均。
主要区别包括
- DDP 每个批次只执行一次通信 - 发送梯度,而 DP 每个批次执行五次不同的数据交换。DDP 使用torch.distributed复制数据,而 DP 通过 Python 线程(引入了与 GIL 相关的限制)在进程中复制数据。因此,
DistributedDataParallel
(DDP) 通常比DataParallel
(DP) 更快,除非您的 GPU 卡互连速度很慢。 - 在 DP 下,GPU 0 的工作量明显高于其他 GPU,导致 GPU 利用率不足。
- DDP 支持跨多台机器的分布式训练,而 DP 则不支持。
这不是 DP 和 DDP 之间差异的完整列表,但是,其他细微差别超出了本指南的范围。您可以通过阅读这篇文章来更深入地了解这些方法。
让我们用一个实验来说明 DP 和 DDP 之间的区别。我们将对 DP 和 DDP 之间的差异进行基准测试,并添加 NVLink 的存在。
- 硬件:2x TITAN RTX 24GB,每个 GPU 都有 2 个 NVLink(在
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}
DDP 不带 NVlink
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) 在以下来自这篇博客文章的图表中进行了说明。
虽然它看起来很复杂,但它与DataParallel
(DP) 的概念非常相似。区别在于,ZeRO-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。
这样,3 个 GPU 中的每一个都获得重建的完整张量,并使用它自己的小批量进行前向传递。计算完成后,不再需要的数据就会被丢弃 - 它只在计算过程中使用。重建是通过预取高效地完成的。
然后,整个过程针对层 Lb 重复,然后针对层 Lc 前向传播,然后反向传播 Lc -> Lb -> La。
这种机制类似于一种高效的团队背包策略:A 负责携带帐篷,B 负责携带炉子,C 负责携带斧头。每晚他们都会与他人分享各自的物品,并从他人那里获得他们没有的东西,第二天早上,他们会收拾好分配给自己的物品类型并继续前进。这就是 ZeRO DP/Sharded DDP。将这种策略与每个成员都必须携带自己的帐篷、炉子和斧头的简单策略进行比较(类似于 PyTorch 中的 DataParallel (DP 和 DDP)),这种简单策略效率要低得多。
在阅读有关该主题的文献时,你可能会遇到以下同义词:分片、分区。如果你仔细观察 ZeRO 对模型权重的分区方式 - 它看起来非常类似于张量并行,这一点将在后面讨论。这是因为它对每一层的权重进行分区/分片,这与接下来要讨论的垂直模型并行不同。
实现
- DeepSpeed ZeRO-DP 阶段 1+2+3
Accelerate
集成transformers
集成
从朴素模型并行到流水线并行
为了解释流水线并行,我们将首先了解朴素模型并行 (MP),也称为垂直 MP。这种方法涉及将模型层组分布在多个 GPU 上,通过使用 .to()
将特定层分配给特定 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 的内存容量增加四倍,而忽略了其他硬件。
- 设备之间数据传输的开销:例如,4 块 6GB 的卡能够容纳与 1 块 24GB 卡使用朴素 MP 时相同的容量,但是一块 24GB 卡会更快地完成训练,因为它没有数据复制开销。但是,假设你拥有 40GB 的卡,需要容纳一个 45GB 的模型,那么你可以使用 4 块 40GB 的卡(但由于梯度和优化器状态的原因,几乎无法容纳)。
- 复制共享嵌入:共享嵌入可能需要在 GPU 之间来回复制。
现在你已经了解了朴素模型并行的工作原理及其缺点,让我们来看看流水线并行 (PP)。PP 几乎与朴素 MP 相同,但它通过将传入的批次分成微批次并人为地创建一个流水线来解决 GPU 闲置问题,这允许不同的 GPU 同时参与计算过程。
以下来自 GPipe 论文 的插图展示了顶部的朴素 MP,以及底部的 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_degree
(8 * 32 * 4 = 1024
)。当 chunks=1
时,你最终得到朴素 MP,效率低下。当 chunks
的值很大时,你最终得到微小的微批次大小,这也效率低下。因此,我们鼓励你尝试不同的 chunks
值,找到最能提高 GPU 利用率的最佳值。
你可能会注意到图表上有一个“死”时间的气泡,它无法并行化,因为最后一个 forward
阶段必须等待 backward
完成流水线。寻找 chunks
的最佳值的目的是实现所有参与 GPU 的高并发利用率,这转化为最小化气泡的大小。
流水线 API 解决方案已在以下软件中实现
- PyTorch
- DeepSpeed
- Megatron-LM
这些方法有一些缺点
- 它们必须对模型进行相当大的修改,因为流水线需要将模型的正常模块流重写为相同的
nn.Sequential
序列,这可能需要更改模型的设计。 - 目前,流水线 API 非常有限。如果你在流水线的第一个阶段中有一堆 Python 变量,你需要找到解决办法。目前,流水线接口要求输入和输出必须是单个张量或张量元组。这些张量必须以批次大小作为第一个维度,因为流水线将把小批量分成微批次。此处正在讨论可能的改进 https://github.com/pytorch/pytorch/pull/50693
- 在流水线阶段级别上无法实现条件控制流 - 例如,T5 等编码器-解码器模型需要特殊的解决方法来处理条件编码器阶段。
- 它们必须安排每一层,使得一层的输出成为另一层的输入。
较新的解决方案包括
- Varuna
- 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 使用 交错流水线 的概念
这里通过优先处理反向传递来进一步最小化气泡(闲置时间)。Varuna 进一步尝试通过使用模拟来发现最有效的调度方案来改进调度。
OSLO 具有基于 Transformers 的流水线并行实现,无需进行 nn.Sequential
转换。
张量并行
在张量并行中,每个 GPU 处理张量的一个切片,并且只在需要时将完整张量聚合起来。为了描述这种方法,本指南的这一部分依赖于来自 Megatron-LM 论文:Efficient Large-Scale Language Model Training on GPU Clusters 的概念和图表。
任何 Transformer 的主要构建块都是一个全连接的 nn.Linear
,后面跟着一个非线性激活函数 GeLU
。根据 Megatron 论文的符号,它的点积部分可以写成 Y = GeLU(XA)
,其中 X
是输入向量,Y
是输出向量,A
是权重矩阵。
如果我们以矩阵形式查看计算,你可以看到矩阵乘法如何在多个 GPU 之间进行分割。
如果我们将权重矩阵 A
按列分割到 N
个 GPU 上,并并行执行 XA_1
到 XA_n
的矩阵乘法,那么我们将得到 N
个输出向量 Y_1, Y_2, ..., Y_n
,它们可以独立地馈送到 GeLU
中。
利用这一原理,我们可以更新任意深度的多层感知器,而无需在 GPU 之间进行任何同步,直到最后我们需要从分片中重建输出向量。Megatron-LM 论文的作者对此提供了一个有用的说明。
多头注意力层的并行化更加简单,因为它们本身就具有并行性,因为它们有多个独立的头!
特殊注意事项:TP 需要非常快的网络,因此不建议跨多个节点进行 TP。实际上,如果一个节点有 4 个 GPU,那么最高的 TP 度数就是 4。如果你需要 8 度的 TP,你需要使用至少有 8 个 GPU 的节点。
本节基于最初更 详细的 TP 概述。作者:@anton-l.
其他名称
- DeepSpeed 将其称为 张量切片
实现
- Megatron-LM 具有内部实现,因为它非常特定于模型
- parallelformers(目前仅推理)
- SageMaker - 这是一个专有解决方案,只能在 AWS 上使用。
- OSLO 具有基于 Transformers 的张量并行实现。
SageMaker 将 TP 与 DP 相结合,以实现更高效的处理。
🤗 Transformers 状态
- 核心:尚未在核心代码中实现
- 但如果你想进行推理,parallelformers 为我们的大多数模型提供了这种支持。因此,在核心代码中实现之前,你可以使用它们。并且希望训练模式也将得到支持。
- DeepSpeed-Inference 还支持我们的 BERT、GPT-2 和 GPT-Neo 模型,它们以其超快的基于 CUDA 内核的推理模式运行,了解更多信息请 点击这里
🤗 Accelerate 与来自 Megatron-LM 的 TP 集成。
数据并行 + 管道并行
以下来自 DeepSpeed 管道教程 的图表展示了如何将 DP 与 PP 相结合。
这里重要的是要看到 DP 排名 0 看不到 GPU2,DP 排名 1 看不到 GPU3。对于 DP 来说,只有 GPU 0 和 GPU 1,它像只有 2 个 GPU 一样将数据馈送进去。GPU0 “秘密地” 通过使用 PP 将部分负载卸载到 GPU2。GPU1 也通过使用 GPU3 来协助自己做同样的事情。
由于每个维度都需要至少 2 个 GPU,因此这里至少需要 4 个 GPU。
实现
🤗 Transformers 状态:尚未实现
数据并行 + 管道并行 + 张量并行
为了获得更高效的训练,使用 3D 并行,其中 PP 与 TP 和 DP 相结合。这可以在以下图表中看到。
此图来自一篇博客文章 3D 并行:扩展到万亿参数模型,这篇文章也很值得一读。
由于每个维度都需要至少 2 个 GPU,因此这里至少需要 8 个 GPU。
实现
- DeepSpeed - DeepSpeed 还包含一个更高效的 DP,称为 ZeRO-DP。
- Megatron-LM
- Varuna
- SageMaker
- OSLO
🤗 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 上。
实现
- Megatron-DeepSpeed 和来自 BigScience 的 Megatron-Deepspeed,它是前者仓库的分支。
- OSLO
重要论文
🤗 Transformers 状态:尚未实现,因为我们没有 PP 和 TP。
FlexFlow
FlexFlow 也以稍微不同的方法解决了并行化问题。
它对样本-操作-属性-参数执行一种 4D 并行。
- 样本 = 数据并行(样本级并行)
- 操作 = 将单个操作并行化为多个子操作
- 属性 = 数据并行(长度级并行)
- 参数 = 模型并行(无论维度如何 - 水平或垂直)
示例
- 样本
我们以 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。
- 参数
它类似于张量模型并行或简单的层级模型并行。
该框架的重要性在于它考虑了 (1) GPU/TPU/CPU vs. (2) RAM/DRAM vs. (3) 快速内部连接/慢速外部连接等资源,并且它会自动优化所有这些资源,以算法方式决定在何处使用哪种并行策略。
一个非常重要的方面是,FlexFlow 被设计为优化具有静态和固定工作负载的模型的 DNN 并行,因为具有动态行为的模型可能更喜欢跨迭代的不同并行策略。
因此,这个承诺非常有吸引力 - 它在您选择的集群上运行 30 分钟的模拟,并提出最佳策略来利用这个特定环境。如果您添加/删除/替换任何部分,它将运行并重新优化该计划。然后您可以进行训练。不同的设置将有自己的自定义优化。
🤗 Transformers 状态:Transformers 模型可以通过 transformers.utils.fx 进行 FX 跟踪,这是 FlexFlow 的先决条件,但是需要对 FlexFlow 进行更改以使其与 Transformers 模型一起使用。
GPU 选择
在多 GPU 上训练时,您可以指定要使用的 GPU 数量及其顺序。例如,当您拥有具有不同计算能力的 GPU,并且希望首先使用更快的 GPU 时,这将很有用。该选择过程适用于 DistributedDataParallel 和 DataParallel,以仅使用可用 GPU 的子集,并且您不需要 Accelerate 或者 DeepSpeed 集成。
GPU 数量
例如,如果您有 4 个 GPU,而您只想使用前 2 个
使用 --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:0
和 cuda:1
。您也可以反转 GPU 顺序以首先使用 2。现在,映射是 cuda:1
用于 GPU 0,而 cuda:0
用于 GPU 2。
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 顺序的另一个环境变量。您可以按以下顺序排列它们:
- PCIe 总线 ID,与
nvidia-smi
和rocm-smi
分别匹配 NVIDIA 和 AMD GPU 的顺序
export CUDA_DEVICE_ORDER=PCI_BUS_ID
- GPU 计算能力
export CUDA_DEVICE_ORDER=FASTEST_FIRST
CUDA_DEVICE_ORDER
特别有用,如果您的训练设置包括旧 GPU 和新 GPU,其中旧 GPU 首先出现,但您无法物理交换卡以使新 GPU 首先出现。在这种情况下,设置 CUDA_DEVICE_ORDER=FASTEST_FIRST
以始终首先使用较新且更快的 GPU(nvidia-smi
或 rocm-smi
仍然按其 PCIe 顺序报告 GPU)。或者您也可以设置 export CUDA_VISIBLE_DEVICES=1,0
。