使用 PyTorch Fully Sharded Data Parallel 加速大型模型训练

在这篇文章中，我们将探讨如何利用 Accelerate 库来训练大型模型，该库使用户能够利用 PyTorch FullyShardedDataParallel (FSDP) 的最新功能。

动机 🤗

随着机器学习 (ML) 模型规模、大小和参数的不断增长，ML 从业人员发现很难在其硬件上训练甚至加载如此大的模型。一方面，已经发现大型模型学习速度快（数据和计算效率高），并且与小型模型相比性能显著提高 [1]；另一方面，在大多数可用硬件上训练此类模型变得 prohibitive。

分布式训练是实现大型 ML 模型训练的关键。大规模分布式训练领域最近取得了重大进展。以下是一些最显著的进展：

1. 使用 ZeRO - Zero Redundancy Optimizer [2] 的数据并行
   1. 阶段 1：在数据并行 worker/GPU 之间分片优化器状态
   2. 阶段 2：在数据并行 worker/GPU 之间分片优化器状态 + 梯度
   3. 阶段 3：在数据并行 worker/GPU 之间分片优化器状态 + 梯度 + 模型参数
   4. CPU 卸载：在 ZERO 阶段 2 [3] 的基础上将梯度 + 优化器状态卸载到 CPU
2. 张量并行 [4]：一种模型并行形式，通过巧妙的方式在加速器/GPU 之间分片具有大量参数的单个层的参数，以实现并行计算，同时避免昂贵的通信同步开销。
3. 流水线并行 [5]：一种模型并行形式，将模型的不同层放置在不同的加速器/GPU 上，并采用流水线技术使所有加速器同时运行。例如，第二个加速器/GPU 计算第一个微批次，而第一个加速器/GPU 计算第二个微批次。
4. 3D 并行 [3]：结合使用 ZERO 的数据并行 + 张量并行 + 流水线并行来训练数千亿参数的庞大模型。例如，BigScience 176B 参数语言模型就采用了这种技术 [6]。

在这篇文章中，我们将探讨使用 ZeRO 进行数据并行，更具体地说是最新的 PyTorch 功能 FullyShardedDataParallel (FSDP)。 DeepSpeed 和 FairScale 已经实现了 ZERO 论文的核心思想。这些功能已集成到 `transformers` Trainer 中，并附有精彩的博客文章 通过 DeepSpeed 和 FairScale 使用 ZeRO 适应更多模型并更快地训练 [10]。 PyTorch 最近将 Fairscale FSDP 上游集成到 PyTorch Distributed 中，并进行了额外的优化。

Accelerate 🚀：无需任何代码更改即可利用 PyTorch FSDP

我们将研究使用 GPT-2 Large (762M) 和 XL (1.5B) 模型变体进行因果语言建模的任务。

下面是 GPT-2 模型预训练的代码。它类似于官方的因果语言建模示例 此处，增加了 2 个参数 `n_train` (2000) 和 `n_val` (500) 以防止在整个数据上进行预处理/训练，从而进行快速概念验证基准测试。

run_clm_no_trainer.py

运行命令 `accelerate config` 后的示例 FSDP 配置

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
fsdp_config:
  min_num_params: 2000
  offload_params: false
  sharding_strategy: 1
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
use_cpu: false

多 GPU FSDP

在这里，我们对单节点多 GPU 设置进行了实验。我们比较了分布式数据并行 (DDP) 和 FSDP 在各种配置下的性能。首先，使用 GPT-2 Large (762M) 模型，其中 DDP 在特定批大小下可以正常工作，而不会抛出内存不足 (OOM) 错误。接下来，使用 GPT-2 XL (1.5B) 模型，其中 DDP 即使在批大小为 1 的情况下也会因 OOM 错误而失败。我们观察到 FSDP 允许 GPT-2 Large 模型使用更大的批大小，并且它能够以不错的批大小训练 GPT-2 XL 模型，这与 DDP 不同。

硬件设置：2 块 24GB NVIDIA Titan RTX GPU。

训练 GPT-2 Large 模型（762M 参数）的命令

export BS=#`try with different batch sizes till you don't get OOM error,
#i.e., start with larger batch size and go on decreasing till it fits on GPU`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-large \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

FSDP 运行示例：

表 1：GPT-2 Large (762M) 模型上的 FSDP 基准测试

与 DDP 相比，从表 1 中我们可以观察到 FSDP 能够实现更大的批处理大小，在不使用 CPU 卸载和使用 CPU 卸载的情况下分别达到 2 倍到 3 倍。在训练时间方面，DDP 混合精度最快，其次是使用 ZERO 阶段 2 和阶段 3 的 FSDP。由于因果语言建模任务始终具有固定的上下文序列长度（--block_size），因此 FSDP 在训练时间方面的提速并不明显。对于动态批处理的应用，FSDP 能够实现更大的批处理大小，因此在训练时间方面可能会有显著的提速。FSDP 混合精度支持目前在 Transformer 上存在一些 问题。一旦支持，训练时间将进一步大幅提高。

CPU 卸载以实现无法适应 GPU 内存的庞大模型训练

训练 GPT-2 XL 模型（1.5B 参数）的命令

export BS=#`try with different batch sizes till you don't get OOM error,
#i.e., start with larger batch size and go on decreasing till it fits on GPU`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-xl \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

表 2：GPT-2 XL (1.5B) 模型上的 FSDP 基准测试

从表 2 中，我们可以观察到 DDP（有和没有 fp16）甚至无法以批大小为 1 的方式运行，并导致 CUDA OOM 错误。FSDP 与 Zero-Stage 3 能够在 2 个 GPU 上以批大小 5（有效批大小 = 10 (5 X 2)）运行。FSDP 与 CPU 卸载可以在使用 2 个 GPU 时进一步将每个 GPU 的最大批大小增加到 14。FSDP 与 CPU 卸载能够在单个 GPU 上以批大小 10 训练 GPT-2 1.5B 模型。这使得计算资源最少的 ML 从业人员能够训练此类大型模型，从而使大型模型训练民主化。

FSDP 集成的功能和限制

让我们深入了解 Accelerate 为 FSDP 集成提供的当前支持和已知限制。

FSDP 支持所需的 PyTorch 版本：PyTorch Nightly（或 1.12.0，如果您在此发布后阅读此内容），因为 FSDP 激活后的模型保存仅在最近的修复后才可用。

通过 CLI 配置

1. 分片策略：[1] FULL_SHARD，[2] SHARD_GRAD_OP
2. 最小参数数量：FSDP 默认自动包装的最小参数数量。
3. 卸载参数：决定是否将参数和梯度卸载到 CPU。

为了获得更多控制，用户可以利用 `FullyShardedDataParallelPlugin`，其中他们可以指定 `auto_wrap_policy`、`backward_prefetch` 和 `ignored_modules`。

创建此类的实例后，用户可以在创建 Accelerator 对象时将其传入。

有关这些选项的更多信息，请参阅 PyTorch FullyShardedDataParallel 代码。

接下来，我们将看到 `min_num_params` 配置的重要性。下面是 [8] 中关于 FSDP 自动包装策略重要性的摘录。

(来源： 链接)

当使用 `default_auto_wrap_policy` 时，如果层中的参数数量超过 `min_num_params`，则该层将被包装在 FSDP 模块中。用于在 GLUE MRPC 任务上微调 BERT-Large (330M) 模型的代码是官方完整的 NLP 示例，概述了如何正确使用 FSDP 功能，并增加了用于跟踪峰值内存使用量的实用程序。

fsdp_with_peak_mem_tracking.py

我们利用 Accelerate 中的跟踪功能支持来记录训练和评估的峰值内存使用情况以及评估指标。下面是来自 wandb 运行 的图表快照。

我们可以观察到，DDP 消耗的内存是 FSDP 自动包装的两倍。FSDP 不带自动包装的内存消耗比带自动包装的 FSDP 多，但远低于 DDP。带自动包装且 `min_num_params=2k` 的 FSDP 内存消耗略低于 `min_num_params=1M` 的设置。这突出显示了 FSDP 自动包装策略的重要性，用户应该尝试使用 `min_num_params` 来找到能够显著节省内存且不会导致大量通信开销的设置。PyTorch 团队正在开发此配置的自动调优工具，如 [8] 中所述。

需要注意的一些注意事项

• PyTorch FSDP 自动包装子模块，扁平化参数并就地分片参数。因此，在模型包装之前创建的任何优化器都会损坏并占用更多内存。因此，强烈建议在创建优化器之前准备模型，这会更有效率。在单个模型的情况下，`Accelerate` 将自动包装模型并为您创建一个优化器，并发出警告消息。

  FSDP 警告：使用 FSDP 时，建议在创建优化器之前为模型调用 prepare，这样更有效率。

然而，以下是使用 FSDP 时准备模型和优化器的推荐方法

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+ model = accelerator.prepare(model)

optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)

- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
-    )

+ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+         optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )

• 在单个模型的情况下，如果您已经使用多个参数组创建了优化器并一起调用了 prepare，那么参数组将丢失并显示以下警告：

  FSDP 警告：使用 FSDP 时，由于嵌套模块包装和参数扁平化，多个参数组将被合并为单个参数组。

这是因为在包装之前创建的参数组在包装之后将失去意义，因为嵌套 FSDP 模块的参数将被扁平化为一维数组（这可能会占用多层）。例如，以下是 GPU 0 上 FSDP 模型的命名参数（当使用 2 个 GPU 时。一维数组中大约有 55M（110M/2）个参数，因为这将是参数的第一个分片）。在这里，如果一个模型对未包装的 BERT-Base 模型的 [bias, LayerNorm.weight] 命名参数应用了不权重衰减，则不能将其应用于下面的 FSDP 包装模型，因为没有这些字符串的命名参数，并且这些层的参数与各种其他层的参数连接在一起。更多细节请参阅此 问题（`原始模型参数的 .grads 未设置，这意味着它们无法单独优化（这就是我们不支持多个参数组的原因）`）。

```
{
'_fsdp_wrapped_module.flat_param': torch.Size([494209]),

'_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]),

'_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
}
```

• 在多个模型的情况下，必须在创建优化器之前准备模型，否则会抛出错误。

• FSDP 目前不支持混合精度，我们正在等待 PyTorch 修复对其的支持。

工作原理 📝

(来源： 链接)

上述工作流程概述了 FSDP 激活后幕后发生的事情。让我们首先了解 DDP 的工作原理以及 FSDP 如何改进它。在 DDP 中，每个 worker/加速器/GPU 都有一个完整的模型参数、梯度和优化器状态副本。每个 worker 获得不同的数据批次，它通过前向传播，计算损失，然后是反向传播以生成梯度。现在，执行一次 all-reduce 操作，其中每个 worker 从其余 worker 获取梯度并进行平均。这样，每个 worker 现在拥有相同的全局梯度，优化器使用这些梯度更新模型参数。我们可以看到，拥有完整的副本会在每个 GPU 上消耗大量冗余内存，这限制了批处理大小以及模型的大小。

FSDP 精确地通过在数据并行 worker 之间分片优化器状态、梯度和模型参数来解决这个问题。它进一步促进了所有这些张量的 CPU 卸载，从而能够加载无法适应可用 GPU 内存的大型模型。与 DDP 类似，每个 worker 获得不同的数据批次。在前向传播期间，如果启用了 CPU 卸载，则本地分片的参数首先复制到 GPU/加速器。然后，每个 worker 对给定的 FSDP 包装模块/层执行 all-gather 操作以获取所有所需的参数，执行计算，然后释放/清空其他 worker 的参数分片。这适用于所有 FSDP 模块。在前向传播之后计算损失，在反向传播期间，再次执行 all-gather 操作以获取给定 FSDP 模块所需的所有参数，执行计算以获取局部梯度，然后释放其他 worker 的分片。现在，局部梯度被平均并通过 reduce-scatter 操作分片到每个相关 worker。这允许每个 worker 更新其本地分片的参数。如果激活了 CPU 卸载，梯度将传递到 CPU 以直接在 CPU 上更新参数。

请参阅 [7, 8, 9] 以获取有关 PyTorch FSDP 工作原理和使用此功能进行的广泛实验的所有深入细节。

问题

如果您在使用 PyTorch FSDP 集成部分遇到任何问题，请在 accelerate 中提出问题。

但是，如果您在 PyTorch FSDP 配置和部署方面遇到问题，则需要咨询相关领域的专家，因此，请改为提出 PyTorch 问题。

参考

[1] 先训练大模型，再压缩：重新思考 Transformer 高效训练和推理的模型大小

[2] ZeRO：内存优化以训练万亿参数模型

[3] DeepSpeed：人人可用的极限规模模型训练 - 微软研究院

[4] Megatron-LM：使用模型并行训练数十亿参数语言模型

[5] 隆重推出 GPipe，一个用于高效训练大规模神经网络模型的开源库

[6] 训练 176B 参数模型需要什么硬件？

[7] PyTorch Fully Sharded Data Parallel (FSDP) API 介绍 | PyTorch

[8] Fully Sharded Data Parallel (FSDP) 入门 — PyTorch 教程 1.11.0+cu102 文档

[9] 使用 PyTorch Fully Sharded Data Parallel 在 AWS 上训练万亿参数模型 | PyTorch | PyTorch | 2022 年 3 月 | Medium

[10] 通过 DeepSpeed 和 FairScale 使用 ZeRO 适应更多模型并更快地训练