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

本文，我们将了解如何基于 PyTorch 最新的 完全分片数据并行 (Fully Sharded Data Parallel，FSDP) 功能用 Accelerate 库来训练大模型。

动机 🤗

随着机器学习 (ML) 模型的规模、大小和参数量的不断增加，ML 从业者发现在自己的硬件上训练甚至加载如此大的模型变得越来越难。 一方面，人们发现大模型与较小的模型相比，学习速度更快 (数据和计算效率更高) 且会有显著的提升 [1]; 另一方面，在大多数硬件上训练此类模型变得令人望而却步。

分布式训练是训练这些机器学习大模型的关键。大规模分布式训练 领域最近取得了不少重大进展，我们将其中一些最突出的进展总结如下:

1. 使用 ZeRO 数据并行 - 零冗余优化器 [2]

2. 阶段 1: 跨数据并行进程 / GPU 对优化器状态 进行分片

3. 阶段 2: 跨数据并行进程/ GPU 对优化器状态 + 梯度 进行分片

4. 阶段 3: 跨数据并行进程 / GPU 对优化器状态 + 梯度 + 模型参数 进行分片

5. CPU 卸载: 进一步将 ZeRO 阶段 2 的优化器状态 + 梯度 卸载到 CPU 上 [3]

6. 张量并行 [4]: 模型并行的一种形式，通过对各层参数进行精巧的跨加速器 / GPU 分片，在实现并行计算的同时避免了昂贵的通信同步开销。

7. 流水线并行 [5]: 模型并行的另一种形式，其将模型的不同层放在不同的加速器 / GPU 上，并利用流水线来保持所有加速器同时运行。举个例子，在第 2 个加速器 / GPU 对第 1 个 micro batch 进行计算的同时，第 1 个加速器 / GPU 对第 2 个 micro batch 进行计算。

8. 3D 并行 [3]: 采用 ZeRO 数据并行 + 张量并行 + 流水线并行 的方式来训练数百亿参数的大模型。例如，BigScience 176B 语言模型就采用了该并行方式 [6]。

本文我们主要关注 ZeRO 数据并行，更具体地讲是 PyTorch 最新的 完全分片数据并行 (Fully Sharded Data Parallel，FSDP) 功能。DeepSpeed 和 FairScale 实现了 ZeRO 论文的核心思想。我们已经将其集成到了 transformers 的 Trainer 中，详见博文 通过 DeepSpeed 和 FairScale 使用 ZeRO 进行更大更快的训练[10]。最近，PyTorch 已正式将 Fairscale FSDP 整合进其 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 尚能够支持其中某些