在 Opacus 中启用完全分片数据并行（FSDP2）

引言与背景

Opacus 凭借其最新的增强功能，在支持大规模模型的私有训练方面取得了显著进展。最近，我们引入了快速梯度裁剪 (FGC) 和幽灵裁剪 (GC)，使开发人员和研究人员无需实例化每样本梯度即可执行梯度裁剪。与依赖钩子的 Opacus 原生实现相比，这些方法减少了 DP-SGD 的内存占用。

即使有了这些进步，训练大型模型（如大型语言模型 (LLM)）对 Opacus 来说仍然是一个重大挑战。随着对大规模模型私有训练的需求持续增长，Opacus 支持数据并行和模型并行技术至关重要。目前，Opacus 支持差分隐私分布式数据并行 (DPDDP) 以实现多 GPU 训练。虽然 DPDDP 有效地将模型训练扩展到多个 GPU 和节点，但它要求每个 GPU 存储模型和优化器状态的副本，导致内存需求高，特别是对于大型模型。

这一限制凸显了对替代并行化技术的需求，例如完全分片数据并行 (FSDP)，它可以通过模型、梯度和优化器状态分片提供更高的内存效率和可扩展性。在训练 Llama 或其他大型语言模型的背景下，通常会根据模型大小采用不同的并行策略来扩展训练。

• 1D 并行：对于小型模型（<100 亿参数），使用 DDP 或 FSDP。
• 2D 并行：对于中型模型（100-1000 亿参数），FSDP 与张量并行 (TP) 结合。
• 4D 并行：对于大型模型（>1000 亿参数），FSDP 与 TP、流水线并行 (PP) 和上下文并行 (CP) 结合。

通过采用 FSDP（示例），Opacus 正在增强其能力，以促进 LLM 更高效、更可扩展的私有训练或微调。这一发展标志着在满足机器学习社区不断变化的需求方面迈出了充满希望的一步，为 2D 和 4D 并行等高级并行策略支持中大型模型的私有训练铺平了道路。

FSDP 与 FGC 和 GC

完全分片数据并行 (FSDP) 是一种强大的数据并行技术，它通过高效管理多个 GPU 工作器之间的内存使用，实现更大模型的训练。FSDP 允许在工作器之间分片模型参数、梯度和优化器状态，从而显著减少训练所需的内存占用。尽管这种方法会由于训练期间的参数收集和丢弃而产生额外的通信开销，但通常可以通过将其与计算重叠来缓解成本。

在 FSDP 中，即使每个数据微批次的计算仍然是每个 GPU 工作器的本地计算，也会一次性收集一个块（例如，一个层）的完整参数，从而降低内存占用。一旦一个块被处理，每个 GPU 工作器会丢弃从其他工作器收集的参数分片，只保留其本地分片。因此，峰值内存使用由每层参数+梯度+优化器状态的最大大小以及激活的总大小决定，这取决于每设备批次大小。有关 FSDP 的更多详细信息，请参阅 PyTorch FSDP 论文。

FSDP 与 FGC 或 GC 的流程如下：

• 前向传播
  • 对于层中的每个层
    • [FSDP 钩子] all_gather 层的完整参数
    • 层的前向传播
    • [FSDP 钩子] 丢弃层的完整参数
    • [Opacus 钩子] 存储层的激活
• 重置 optimizer.zero_grad()
• 第一次反向传播
  • 对于层中的每个层
    • [FSDP 钩子] all_gather 层的完整参数
    • 层的反向传播
    • [Opacus 钩子] 使用 FGC 或 GC 计算每样本梯度范数
    • [FSDP 钩子] 丢弃层的完整参数
    • [FSDP 钩子] reduce_scatter 层的梯度 → 非必要
• 使用每样本梯度范数重新缩放损失函数
• 重置 optimizer.zero_grad()
• 第二次反向传播
  • 对于层中的每个层
    • [FSDP 钩子] all_gather 层的完整参数
    • 层的反向传播
    • [FSDP 钩子] 丢弃层的完整参数
    • [FSDP 钩子] reduce_scatter 层的梯度
• 在每个设备上对相应的参数分片添加噪声
• 在每个设备上对相应的参数分片应用优化器步骤

图 1：Opacus 中基于 FSDP 的快速梯度裁剪或幽灵裁剪的工作流程。请注意，计算和通信之间存在重叠——1) 在前向传播中：当前层 (l) 的计算与下一层 (l+1) 参数的 all_gather 重叠。2) 在反向传播中：当前层 (l) 的梯度计算与上一层 (l+1) 梯度的 reduce_scatter 和下一层 (l-1) 参数的 all_gather 重叠。

如何在 Opacus 中使用 FSDP

训练循环与标准 PyTorch 循环相同。与以前的 Opacus 一样，我们使用

• PrivacyEngine()，它配置模型和优化器以运行 DP-SGD。
• 要启用 FSDP 的幽灵裁剪，使用参数grad_sample_mode="ghost_fsdp"。
• 此外，我们在初始化优化器并调用make_private()
  之前用 FSDP2Wrapper 封装模型

FSDP2Wrapper 将 FSDP2（FSDP 的第二个版本）应用于根模块以及每个不需要 functorch 计算每样本梯度范数的 torch.nn 层。依赖 functorch 的层类型不会单独使用 FSDP2 封装，因此会落入根模块的通信组中。连接到根模块通信组的层将首先被解片（在前向/反向传播开始时），最后被重新分片（在整个前向/反向传播之后）。这将影响峰值内存，因为连接到根模块的层在执行后不会立即重新分片。

我们在实现中使用 FSDP2，因为以前的版本 (FSDP) 不兼容幽灵裁剪的两次反向传播设置。

from opacus import PrivacyEngine
from opacus.utils.fsdp_utils import FSDP2Wrapper

def launch(rank, world_size):
   torch.cuda.set_device(rank)
   setup_distributed_env(rank, world_size)
   criterion = nn.CrossEntropyLoss() # example loss function
   model     = ToyModel()
   model     = FSDP2Wrapper(model) # different from DPDDP wrapper
   optimizer = optim.SGD(model.parameters(), lr=args.lr)

   privacy_engine = PrivacyEngine()
   model_gc, optimizer_gc, criterion_gc, train_loader, = privacy_engine.make_private(
           module=model,
           optimizer=optimizer,
           data_loader=train_loader,
           noise_multiplier=noise_multiplier
           max_grad_norm=max_grad_norm,
           criterion=criterion,
           grad_sample_mode="ghost_fsdp",)

   # The training loop below is identical to that of PyTorch
   for input_data, target_data in train_loader:
       input_data, target_data = input_data.to(rank), target_data.to(rank)
       output_gc = model_gc(input_data) # Forward pass
       optimizer_gc.zero_grad()
       loss = criterion_gc(output_gc, target_data)
       loss.backward()
       optimizer_gc.step()  # Add noise and update the model

world_size = torch.cuda.device_count()
mp.spawn(
       launch,
       args=(world_size,),
       nprocs=world_size,
       join=True,
   )

内存分析

我们提供了使用 GPT2 系列模型进行完全微调的内存消耗结果。我们只训练 Transformer 块；嵌入层（token 嵌入层、位置嵌入层）被冻结。

图 2 报告了使用幽灵裁剪方法训练模型时 FSDP2 与 DPDDP 支持的最大批次大小。对于 1.5B 参数的 GPT2 模型，在 1x8 A100 40GB GPU 上，与使用 DPDDP 相比，FSDP2 可以实现 2.6 倍的更大批次大小。对于大于 1B 的模型，参数和优化器状态的大小占主导地位，FSDP2 显示出显著的改进。

图 2：在 1x8 A100 40GB GPU 节点上，使用基于幽灵裁剪的 DP-SGD 训练一系列 GPT2 模型时的最大批次大小。我们使用莎士比亚数据集，最大序列长度为 1024，并使用 float32 AdamW 优化器。

表 1 显示了给定步骤的峰值内存和步骤执行后占用的总内存。值得注意的是，模型初始化后 FSDP2 的总内存比 DPDDP 低 8 倍，因为 FSDP2 将模型分片到 8 个 GPU。FSDP2 和 DPDDP 的前向传播大致将峰值内存增加约 10GB，因为两种并行类型中激活都没有分片。对于反向传播和优化器步骤，DPDDP 的峰值内存与模型大小成比例，而 FSDP2 的峰值内存与模型大小除以工作器数量成比例。通常，随着模型大小的增加，FSDP2 的优势变得更加明显。

表 1：在 1x8 A100 40GB GPU 节点上，使用基于幽灵裁剪的 DP-SGD (AdamW 优化器，等效批次大小为 16) 训练 GPT2-xl 模型 (1.5B) 时，在 rank 0 上的内存估计。峰值内存表示步骤期间分配的最大内存，总内存是执行给定步骤后占用的内存量。

延迟分析

表 2 显示了使用 DP-DDP 和 FSDP2 对 Llama-3 8B 模型进行 LoRA 微调的最大批次大小和延迟数据。我们观察到，对于使用幽灵裁剪的 DP-SGD，FSDP2 支持的批次大小几乎是 DP-DDP 的两倍，但在相同有效批次大小下，吞吐量 (0.6x) 低于使用钩子的 DP-DDP。在这种特定的 LoRA 微调设置中，使用 FSDP2 没有带来任何显著的改进。但是，如果数据集的样本序列长度为 4096，而 DP-DDP 无法容纳，那么 FSDP2 就变得必要。

表 2a：在 Tiny Shakespeare 数据集上使用 LoRA 微调 Llama-3 8B（可训练参数：6.8M），AdamW 优化器，32 位精度，最大序列长度 512，1x8 A100 80GB GPU。此处不使用任何梯度累积。

表 2b：在 Tiny Shakespeare 数据集上使用 LoRA 微调 Llama-3 8B（可训练参数：6.8M），AdamW 优化器，32 位精度，最大序列长度 512，1x8 A100 80GB GPU。此处启用梯度累积以将总批次大小增加到 256。 

表 3 显示了 Llama-3 8B 完全微调的吞吐量数据。目前，FSDP2 与幽灵裁剪不兼容绑定的参数（嵌入层）。我们在微调期间冻结这些层，这将可训练参数从 8B 减少到 7.5B。如表 3 所示，DP-DDP 即使每设备批次大小为 1 也会抛出 OOM 错误。而使用 FSDP2，每个设备可以容纳批次大小为 8 的数据，从而实现 Llama-3 8B 的完全微调。

为了比较 FSDP2 和 DP-DDP 的完全微调，我们将 AdamW 优化器切换到不带动量的 SGD，并将可训练参数从 7.5B 减少到 5.1B，通过冻结归一化层和门投影层的权重。这使得 DP-DDP 可以以 2 的批次大小运行（如果启用梯度累积则为 1）。在这种设置下，我们观察到 FSDP2 对于等效批次大小比 DP-DDP 快 1.65 倍。

表 3：基于幽灵裁剪 DP-SGD 的 Llama-3 8B 在 Tiny Shakespeare 数据集上的完全微调，最大序列长度 512，1x8 A100 80GB GPU。

正确性验证

我们对幽灵裁剪 DP-SGD 与 FSDP 进行了内部 Meta 用例的集成测试，该用例包括使用 LoRA 微调的 Llama-3 8B 模型，用于下一个词预测任务。我们的结果表明，使用 FSDP 的幽灵裁剪与 DP-DDP 的训练损失大致相同（差异可忽略不计）。以前的结果以及单元测试（链接）已经证明了实现的正确性。

图 3：使用幽灵裁剪 DP-SGD 进行下一个词预测任务的 Llama-3 8B LoRA 微调的训练损失（y 轴）与迭代次数（x 轴）。

限制

当前版本的 FSDP 不支持以下场景：

1. 具有绑定参数的层。
2. 在训练阶段冻结/解冻可训练参数。

以下是当前 FSDP2 幽灵裁剪梯度累积实现的主要两个限制。

• • 延迟
    • 当前 FSDP2 梯度累积的实现会在每次反向传播后同步梯度。由于幽灵裁剪有两次反向传播，因此对于 k 个梯度累积步骤，我们有 2k 次梯度同步调用 (reduce_scatter)。
    • 这是因为当每次前向传播有两次反向传播时，不能直接使用 no_sync。
    • 理想情况下，对于 k 个梯度累积步骤，我们应该只有 1 次梯度同步调用。
    • 在 LoRA 微调的情况下，reduce_scatter 的延迟可以忽略不计。此外，通过合理的计算/通信重叠，这种开销可以被掩盖。

• 内存

• • 梯度累积使用一个额外的缓冲区来存储累积的梯度（分片），无论梯度累积步骤的数量如何。
  • 当梯度累积步骤的数量等于 1 时，我们希望避免使用额外的缓冲区。这并非 FSDP2 特有，而是 Opacus 库的普遍瓶颈。

主要结论

• 对于可训练参数量较小的模型，例如 LoRA 微调
  • 尽可能推荐使用带梯度累积的 DP-DDP。
  • 如果 DP-DDP 在所需序列长度或模型大小下抛出 OOM 错误，则切换到 FSDP2。
• 对于具有相当数量可训练参数的完全微调
  • 推荐使用 FSDP2，因为它比 DP-DDP 具有更高的吞吐量
  • 在大多数情况下，FSDP2 是唯一的选择，因为 DP-DDP 即使批次大小为 1 也会触发 OOM。
• 上述观察结果适用于私有和非私有情况。

结论

在这篇文章中，我们介绍了完全分片数据并行 (FSDP) 与快速梯度裁剪 (FGC) 和幽灵裁剪 (GC) 在 Opacus 中的集成，展示了其扩展具有超过 10 亿可训练参数的大规模模型私有训练的潜力。通过利用 FSDP，我们展示了可以完全微调 Llama-3 8B 模型，这是由于内存限制而无法通过差分隐私分布式数据并行 (DP-DDP) 实现的壮举。

FSDP 在 Opacus 中的引入标志着 Opacus 库的重大进步，为 LLM 的私有训练提供了可扩展且内存高效的解决方案。这一发展不仅增强了 Opacus 处理大规模模型的能力，也为未来集成其他模型并行策略奠定了基础。

展望未来，我们的重点将是使用幽灵裁剪启用 2D 并行，并将 FSDP 与使用钩子的原生 Opacus 集成。这些努力旨在进一步优化训练过程，减少延迟，并将 Opacus 的适用性扩展到更大、更复杂的模型。我们对这些进步将解锁的可能性感到兴奋，并致力于突破私有机器学习的极限。此外，我们邀请开发人员、研究人员和爱好者加入我们，您的贡献和见解在我们不断增强 Opacus 的过程中弥足珍贵。

鸣谢