使用PyTorch完全分片的数据并行加速大型模型训练

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

随着机器学习（ML）模型的规模、大小和参数的不断增长，ML从业者发现在他们的硬件上训练甚至加载这样大的模型变得困难。一方面，已经发现大型模型学习速度快（数据和计算效率高），与较小的模型相比具有显著的性能优势[1]；另一方面，使用现有的大多数硬件来训练这样的模型变得难以实现。

分布式训练是实现训练这样大型ML模型的关键。在大规模分布式训练领域，近年来取得了重大进展。以下是其中最显著的几个进展：

1. 使用ZeRO（Zero Redundancy Optimizer）的数据并行性[2]
   1. 阶段1：将优化器状态在数据并行工作器/ GPU 之间分片
   2. 阶段2：将优化器状态 + 梯度在数据并行工作器/ GPU 之间分片
   3. 阶段3：将优化器状态 + 梯度 + 模型参数在数据并行工作器/ GPU 之间分片
   4. CPU 卸载：在 ZERO 阶段 2 的基础上将梯度 + 优化器状态卸载到 CPU 上 [3]
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 分布式中，并进行了额外的优化。

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

下面是预训练 GPT-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 时也会出现内存不足错误。我们观察到，相比于 DDP，FSDP 可以实现更大的批量大小用于 GPT-2 Large 模型，并且可以使用较大的批量大小来训练 GPT-2 XL 模型。

硬件设置：2X24GB NVIDIA Titan RTX 显卡。

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

export BS=#`尝试不同的批大小，直到不再出现内存溢出错误，
#即，从较大的批大小开始逐渐减小直到适合显卡`

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：基准测试 FSDP 在 GPT-2 Large（762M）模型上

从表格 1 可以看出，相对于 DDP，FSDP 可以使用更大的批大小 ，分别为无 CPU offload 设置和有 CPU offload 设置时的2倍至3倍。在训练时间方面，DDP 使用混合精度是最快的，其次是使用 ZERO Stage 2 和 Stage 3 的 FSDP。由于因果语言建模任务的上下文序列长度（–block_size）始终是固定的，因此 FSDP 的训练时间加速并不是很大。对于具有动态批处理的应用程序，FSDP 允许更大的批大小，很可能在训练时间上有相当大的加速。目前，FSDP 对 Transformer 模型的混合精度支持存在一些问题。一旦这个问题得到解决，训练时间的加速将进一步显著提高。

通过 CPU 卸载来训练无法适应 GPU 内存的庞大模型

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

export BS=#`尝试不同的批大小，直到不再出现内存溢出错误，
#即，从较大的批大小开始逐渐减小直到适合显卡`

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：基准测试 FSDP 在 GPT-2 XL（1.5B）模型上

从表格 2 可以看出，DDP（使用和不使用 fp16）甚至无法在批大小为 1 的情况下运行，导致 CUDA 内存溢出错误。FSDP 使用 Zero-Stage 3 可以在 2 个显卡上以批大小为 5（有效批大小为 10（5 X 2））运行。当使用 2 个显卡时，FSDP 在 CPU offload 的情况下可以将最大批大小进一步增加到每个显卡 14。这使得 ML 从业者在计算资源有限的情况下能够训练这样大的模型，从而使大型模型训练民主化。

FSDP 集成的能力和限制

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

FSDP 支持所需的 PyTorch 版本：PyTorch Nightly（或者如果您在发布后阅读此文档，则为 1.12.0），因为只有最近的修复才支持激活 FSDP 的模型保存。

通过命令行界面进行配置：

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 Auto Wrap Policy 重要性的摘录。

（图片链接）

（来源：链接）

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

fsdp_with_peak_mem_tracking.py

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

我们可以观察到 DDP 消耗的内存是 FSDP 自动包装的两倍。没有自动包装的 FSDP 消耗的内存比 DDP 少，但比自动包装的 FSDP 多。当将 min_num_params 设置为 2k 时，自动包装的 FSDP 消耗的内存比将 min_num_params 设置为 1M 时稍微少一些。这突显了 FSDP Auto Wrap Policy 的重要性，用户应该尝试不同的 min_num_params 设置，以节省内存并避免过多的通信开销。正如[8]中所提到的，PyTorch 团队正在开发针对该配置的自动调优工具。

需要注意的几个问题：

– PyTorch FSDP 自动包装子模块，将参数扁平化并原地分片。由于这一点，在模型包装之前创建的任何优化器将被破坏并占用更多内存。因此，强烈建议在创建优化器之前准备好模型，这样效率更高。在单个模型的情况下，Accelerate 会自动包装模型并创建优化器，并在有警告消息的情况下提醒您。

– 以下是使用 FSDP 时准备模型和优化器的推荐方式：

“`python
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
+ )
“`

– 如果您使用多个参数组创建了优化器，并在一起调用了准备方法，则参数组将被合并，并显示以下警告：

“`
FSDP Warning: When using FSDP, several parameter groups will be conflated into a single one due to nested module wrapping and parameter flattening.
“`

这是因为在包装之前创建的参数组在包装后将失去意义，因为嵌套的 FSDP 模块被扁平化为1D数组（可能包含许多层）。例如，下面是在 GPU 0 上的 FSDP 模型的命名参数（使用 2 个 GPU 的情况下。在1D数组中有约 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中，每个工作器/加速器/GPU都有整个模型参数、梯度和优化器状态的副本。每个工作器获得不同的数据批次，经过前向传递，计算损失，然后进行反向传递以生成梯度。现在，进行全局梯度的聚合操作，每个工作器获取来自其他工作器的梯度并进行平均。这样，每个工作器现在都具有相同的全局梯度，优化器使用这些梯度来更新模型参数。我们可以看到，完整的副本在每个GPU上消耗了大量冗余内存，这限制了批次大小以及模型的大小。

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

有关PyTorch FSDP工作原理和使用此功能进行的广泛实验的详细信息，请参考[7, 8, 9]。

如果在PyTorch FSDP的集成部分遇到任何问题，请在accelerate中打开一个Issue。

但是，如果您在PyTorch FSDP配置和部署方面遇到问题，您需要向各自领域的专家请教，请打开一个PyTorch Issue。

[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

[3] DeepSpeed: Extreme-scale model training for everyone – Microsoft Research

[4] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

[5] Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models

[6] Which hardware do you need to train a 176B parameters model?

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

[8] Getting Started with Fully Sharded Data Parallel(FSDP) — PyTorch Tutorials 1.11.0+cu102 documentation

[9] Training a 1 Trillion Parameter Model With PyTorch Fully Sharded Data Parallel on AWS | by PyTorch | PyTorch | Mar, 2022 | VoAGI

[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale