在多个 GPU 上高效训练¶
如果你发现单个 GPU 训练模型太慢，或者模型的权重无法容纳在单个 GPU 的内存中，那么过渡到多 GPU 配置可能是一个可行的选择。在进行这一转变之前，请务必彻底探索所有在单个 GPU 上提高训练效率的方法（详见单 GPU 高效训练方法和工具），因为这些方法适用于任何数量的 GPU 训练。一旦你应用了这些方法并且发现它们在单个 GPU 上仍然不足，再考虑转向多个 GPU。
从单个 GPU 到多个 GPU 的过渡需要引入某种形式的并行处理，因为工作负载需要分布在资源上。可以采用多种技术来实现并行处理，例如数据并行处理（Data Parallelism）、张量并行处理（Tensor Parallelism）和流水线并行处理（Pipeline Parallelism）。需要注意的是，并没有一种通用的最佳方案，最佳设置取决于你所使用的具体硬件配置。
本指南深入介绍了各种类型的并行处理，并提供了如何组合技术以及选择合适方法的指导。有关分布式训练的逐步教程，请参阅🤗 Accelerate 文档。
虽然本指南讨论的主要概念可能适用于各种框架，但这里我们重点介绍基于 PyTorch 的实现。
在深入探讨每种技术的具体细节之前，让我们先了解一下在大型基础设施上训练大型模型时的决策过程。
可扩展性策略¶
首先估算训练模型所需的显存（vRAM）。对于托管在 🤗 Hub 上的模型，可以使用我们的模型内存计算器，该计算器可以提供准确的结果，误差范围在几个百分点内。
单节点/多 GPU 配置的并行化策略¶
当在一个节点上使用多个 GPU 训练模型时，选择合适的并行化策略对性能有显著影响。以下是你的选项：
情况 1：模型可以容纳在单个 GPU 上
如果模型可以舒适地容纳在一个 GPU 上，你有两个主要选项：
- DDP - 分布式数据并行（Distributed DataParallel）
- ZeRO 冗余优化器（ZeRO） - 根据具体情况和配置，这种方法可能更快或更慢，但值得一试。
情况 2：模型无法容纳在单个 GPU 上
如果模型太大无法容纳在一个 GPU 上，你有几种选择：
- 流水线并行（PipelineParallel，简称 PP）
- ZeRO
- 张量并行（TensorParallel，简称 TP）
如果有非常快速的节点间连接（例如 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 的组合。
在本指南的后续部分，我们将深入探讨这些不同的并行化方法是如何工作的。
数据并行（Data Parallelism）¶
即使只有 2 个 GPU，你也可以利用 PyTorch 内置功能（如 DataParallel (DP) 和 DistributedDataParallel (DDP)）提供的加速训练能力。请注意，PyTorch 文档建议在多 GPU 训练时优先使用 DistributedDataParallel (DDP) 而不是 DataParallel (DP)，因为它适用于所有模型。让我们看看这两种方法的工作原理及其不同之处。
DataParallel 与 DistributedDataParallel¶
为了理解这两种方法在 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 相关的限制）在进程内复制数据。因此，除非 GPU 之间的连接速度较慢，否则
DistributedDataParallel(DDP) 通常比DataParallel(DP) 更快。 - 在 DP 下，GPU 0 的工作量显著大于其他 GPU，导致 GPU 利用率低下。
- DDP 支持跨多台机器的分布式训练，而 DP 不支持。
这不是 DP 和 DDP 之间差异的详尽列表，但其他细微差别超出了本指南的范围。你可以通过阅读这篇文章来更深入地了解这些方法：使用 PyTorch 在 AWS 上进行分布式数据并行训练。
接下来，我们通过一个实验来说明 DP 和 DDP 之间的差异。我们将基准测试 DP 和 DDP 之间的差异，并添加 NVLink 是否存在的上下文：
- 硬件：2 块 TITAN RTX 24GB 每块 + NVlink，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 w/ 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 w/o 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 | 110s |
| 2:DDP | Y | 101s |
| 2:DDP | N | 131s |
如你所见，在这种情况下，DP 比带有 NVlink 的 DDP 慢约 10%，但比没有 NVlink 的 DDP 快约 15%。实际差异取决于每个 GPU 需要与其他 GPU 同步的数据量 - 需要同步的数据越多，较慢的连接对整体运行时间的影响就越大。
ZeRO 数据并行（ZeRO Data Parallelism）¶
由这篇博客文章中的图表展示了 ZeRO 优化的数据并行（ZeRO-DP）。
尽管看起来很复杂，但它与 DataParallel (DP) 的概念非常相似。不同之处在于，ZeRO 不再复制完整的模型参数、梯度和优化器状态，而是每个 GPU 只存储其中的一部分。然后，在运行时需要整个层的参数时，所有 GPU 会同步以互相传递缺失的部分。
为了说明这个概念，假设有一个简单的包含 3 层（La、Lb 和 Lc）的模型，每层有 3 个参数。例如，层 La 有参数 a0、a1 和 a2：
如果我们有 3 个 GPU，ZeRO-DP 会将模型拆分到 3 个 GPU 上，如下所示：
这种方式类似于张量并行（Tensor Parallelism）的水平切片，而不是垂直切片，后者将整个层组放在不同的 GPU 上。现在我们来看看它是如何工作的：
每个 GPU 都会像 DP 一样接收一个小批量数据：
输入数据未经修改，就像它们会被原始模型处理一样。
首先，输入数据到达层 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)），这会低效得多。
在阅读相关文献时，你可能会遇到以下同义词：Sharded、Partitioned。如果你仔细观察 ZeRO 如何分割模型的权重 - 它看起来非常类似于张量并行（Tensor Parallelism），这将在后面讨论。这是因为它分割/切片了每一层的权重，而不是垂直模型并行（Vertical Model Parallelism），这将在下一节中讨论。
实现：
- DeepSpeed ZeRO-DP 阶段 1+2+3
Accelerate集成transformers集成
从简单的模型并行到流水线并行¶
为了解释流水线并行（Pipeline Parallelism），我们先来看看简单的模型并行（Naive Model Parallelism，也称为 Vertical MP）。这种方法涉及将模型层组分布到多个 GPU 上，通过 .to() 将特定层分配给特定 GPU。随着数据流经这些层，数据会被移动到与层相同的 GPU 上，而其他层保持不变。
我们称这种模型并行为“垂直”的，因为模型通常是这样可视化的。例如，下图显示了一个 8 层模型垂直切分为两部分，将第 0 到第 3 层放在 GPU0 上，第 4 到第 7 层放在 GPU1 上：
在这个例子中，当数据从第 0 层传到第 3 层时，与常规前向传播没有区别。然而，将数据从第 3 层传到第 4 层需要将数据从 GPU0 移动到 GPU1，这会引入通信开销。如果参与的 GPU 在同一个计算节点上（例如同一台物理机），这种复制很快，但如果 GPU 分布在不同的计算节点上（例如多台机器），通信开销可能会大幅增加。
之后，第 4 到第 7 层的工作方式与原始模型相同。在第 7 层完成后，通常需要将数据发送回第 0 层（或反之，将标签发送到最后一层）。现在可以计算损失，优化器可以进行工作。
简单模型并行存在以下几个缺点：
- 任何时候只有一个 GPU 处于活跃状态：如果使用 4 个 GPU，这几乎等同于将单个 GPU 的内存容量增加四倍，而忽略其他硬件。
- 设备间的数据传输开销：例如，4 张 6GB 显卡可以容纳与一张 24GB 显卡相同大小的模型，但单张 24GB 显卡会更快完成训练，因为它没有数据复制开销。但是，如果你有 40GB 的显卡，需要容纳 45GB 的模型，可以使用 4 张 40GB 的显卡（但几乎不可能，因为还要考虑梯度和优化器状态）。
- 共享嵌入层的复制：共享嵌入层可能需要在 GPU 之间来回复制。
现在你已经熟悉了简单的模型并行如何工作及其缺点，我们来看流水线并行（Pipeline Parallelism，简称 PP）。PP 与简单的模型并行几乎相同，但它通过将传入批次切分为微批次并人工创建流水线，解决了 GPU 闲置问题，允许不同的 GPU 并行参与计算过程。
下图来自 GPipe 论文，展示了顶部的简单模型并行和底部的流水线并行：
在图的底部，可以看到流水线并行（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 称之为梯度累积步骤。
由于 chunks，PP 引入了微批次（Micro-Batches，简称 MBS）的概念。DP 将全局数据批次大小切分为小批次，所以如果你的 DP 程度为 4，全局批次大小为 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 变量，你必须找到变通方法。目前，流水线接口要求输入和输出只能是一个 Tensor 或一组 Tensors。这些 Tensors 必须在第一个维度上有批次大小，因为流水线会将小批次切分为微批次。相关改进正在讨论中：https://github.com/pytorch/pytorch/pull/50693
- 无法在管道阶段的条件控制流 - 例如，T5 编码器-解码器模型需要特殊的工作around来处理条件编码阶段。
- 它们需要安排每个层，使得一个层的输出成为另一个层的输入。
更近期的解决方案包括：
- 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 并确保所有输入都是 Tensors。这是因为在当前的模型中包含了许多特性，使得转换非常复杂，需要移除这些特性才能实现。
DeepSpeed 和 Megatron-LM 集成可用在 🤗 Accelerate 中。
其他方法：
DeepSpeed、Varuna 和 SageMaker 使用了 交错流水线 的概念。
这里通过优先进行反向传播进一步最小化了“气泡”（空闲时间）。Varuna 进一步尝试通过模拟来发现最有效的调度。
OSLO 采用了基于 Transformers 的流水线并行实现，无需进行 nn.Sequential 转换。
张量并行（Tensor Parallelism）¶
在张量并行（Tensor Parallelism）中，每个 GPU 处理张量的一个切片，并在需要时聚合完整的张量。为了描述这种方法，本指南的这一部分依赖于 Megatron-LM 论文 Efficient Large-Scale Language Model Training on GPU Clusters 中的概念和图表。
任何变换器的核心组成部分是一个全连接的 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 函数：
基于这一原则，我们可以更新任意深度的多层感知器（MLP），而无需在 GPU 之间进行同步，直到最后一步，需要从分片中重建输出向量。Megatron-LM 论文的作者提供了一个有用的插图：
多头自注意力层的并行化更简单，因为它们本身就是并行的，具有多个独立的头！
特别注意事项：TP 需要非常快的网络，因此不建议跨多个节点进行 TP。实际上，如果一个节点有 4 个 GPU，则最高 TP 程度为 4。如果你需要 TP 程度为 8，需要使用至少有 8 个 GPU 的节点。
本节基于 @anton-l 的 详细 TP 概述。
其他名称：
- 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 来说，只有 GPU0 和 GPU1，就像只有 2 个 GPU 一样，数据被送入这两个 GPU。GPU0 会“秘密地”将部分负载卸载到 GPU2，使用 PP。同样，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。
实现：
重要论文：
🤗 Transformers 状态：尚未实现在核心部分，因为我们没有 PP 和 TP。
FlexFlow¶
FlexFlow 以稍微不同的方法解决了并行化问题。
论文：“超越深度神经网络的数据并行和模型并行” by Zhihao Jia, Matei Zaharia, Alex Aiken
它在样本、算子、属性和参数四个维度上进行 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 适用于静态和固定工作负载的模型并行化优化，因为动态行为的模型可能在不同迭代中需要不同的并行化策略。
因此，该框架的前景非常诱人 - 它可以在选定的集群上运行 30 分钟的模拟，然后提出最佳策略以利用该特定环境。如果你添加、移除或替换任何部分，它会重新优化计划。然后你可以进行训练。不同的设置会有自己的定制优化。
🤗 Transformers 状态：Transformers 模型可以通过 transformers.utils.fx 进行 FX 跟踪，这是 FlexFlow 的先决条件，但需要在 FlexFlow 方面进行更改以使其与 Transformers 模型兼容。
GPU 选择¶
在多个 GPU 上训练时，你可以指定使用的 GPU 数量和顺序。这在你有不同计算能力的 GPU 且希望先使用较快的 GPU 时非常有用。选择过程适用于 DistributedDataParallel 和 DataParallel，使用它们只需要部分可用的 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:0 和 cuda:1。你也可以反转 GPU 的顺序，先使用 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 和 AMD GPU 的
nvidia-smi和rocm-smi的顺序匹配：
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。