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

引言和背景

Opacus 在支持大规模模型的隐私保护训练方面取得了重大进展。最近，我们引入了快速梯度裁剪 (FGC) 和幽灵裁剪 (GC)，使开发人员和研究人员无需实例化每样本梯度即可执行梯度裁剪。与依赖钩子 (hooks) 的原生 Opacus 实现相比，这些方法降低了 DP-SGD 的内存占用。

尽管有了这些改进，但训练大型语言模型 (LLM) 等大规模模型对 Opacus 来说仍然是一项重大挑战。随着对大规模模型隐私保护训练的需求不断增长，Opacus 支持数据并行和模型并行技术至关重要。目前，Opacus 支持差分隐私分布式数据并行 (DPDDP) 以实现多 GPU 训练。虽然 DPDDP 能有效地跨多个 GPU 和节点扩展模型训练，但它要求每个 GPU 存储一份模型和优化器状态的副本，这导致了高内存需求，对于大型模型尤其如此。

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

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

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

具有 FGC 和 GC 的 FSDP

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

在 FSDP 中，即使每个数据微批次的计算对每个 GPU 工作节点而言是局部的，但参数是在同一时间为整个块（例如一层）收集的，从而降低了内存占用。一旦处理完一个块，每个 GPU 工作节点就会丢弃从其他工作节点收集的参数分片，仅保留其本地分片。因此，峰值内存使用量由每层参数+梯度+优化器状态的最大大小以及激活的总大小决定，后者取决于每设备的批大小。有关 FSDP 的更多详细信息，请参阅 PyTorch FSDP 论文。

带有 FGC 或 GC 的 FSDP 流程如下：

• 前向传播
  • 对于各层中的每一层
    • [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 下启用幽灵裁剪 (Ghost Clipping)，使用参数 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 所支持的最大批大小。使用 FSDP2，在 1×8 A100 40GB GPU 上，对于 15 亿参数的 GPT2 模型，我们可以实现比使用 DPDDP 大 2.6 倍的批大小。FSDP2 在更大模型（>10 亿参数）上表现出显著改进，因为此时参数和优化器状态的大小占主导地位。

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

表 1 展示了给定步骤的峰值内存和步骤执行后的总占用内存。值得注意的是，模型初始化后 FSDP2 的总内存比 DPDDP 低 8 倍，因为 FSDP2 将模型分片到了 8 个 GPU 上。对于 FSDP2 和 DPDDP，前向传播的峰值内存都会增加约 10GB，因为这两种并行方式都不会对激活值进行分片。对于反向传播和优化器步骤，DPDDP 的峰值内存与模型大小成正比，而 FSDP2 的峰值内存与模型大小除以工作节点数成正比。通常情况下，随着模型规模的增加，FSDP2 的优势会更加明显。

表 1：在 1×8 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 数据集上对 Llama-3 8B 进行 LoRA 微调（可训练参数：680 万），使用 32 位精度的 AdamW 优化器，最大序列长度 512，1×8 A100 80GB GPU。此处未使用梯度累积。

表 2b：在 Tiny Shakespeare 数据集上对 Llama-3 8B 进行 LoRA 微调（可训练参数：680 万），使用 32 位精度的 AdamW 优化器，最大序列长度 512，1×8 A100 80GB GPU。此处我们启用梯度累积将总批大小增加至 256。

表 3 展示了 Llama-3 8B 全参数微调的吞吐量数据。目前，带有幽灵裁剪的 FSDP2 不支持绑定参数（嵌入层）。我们在微调过程中冻结了这些层，这使可训练参数从 80 亿减少到 75 亿。如表 3 所示，即使每设备批大小为 1，DP-DDP 也会抛出 OOM 错误。而使用 FSDP2，每个设备可以容纳批大小 8，从而实现 Llama-3 8B 的全参数微调。

为了比较 FSDP2 与 DP-DDP 的全参数微调，我们将 AdamW 优化器切换为无动量的 SGD，并通过冻结归一化层和门控投影层的权重，将可训练参数从 75 亿减少到 51 亿。这使得 DP-DDP 能够以批大小 2（如果启用梯度累积则为 1）运行。在这种设置下，我们观察到在相同批大小下，FSDP2 比 DP-DDP 快 1.65 倍。

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

正确性验证

我们在 Meta 的一个内部用例上对带有 FSDP 的幽灵裁剪 DP-SGD 进行了集成测试，该用例包含一个用于下一个词预测任务的 Llama-3 8B 模型（采用 LoRA 微调）。我们的结果表明，带有 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 中的引入标志着该库的重大进步，为 LLM 的隐私保护训练提供了可扩展且内存高效的解决方案。这一发展不仅增强了 Opacus 处理大规模模型的能力，也为未来集成其他模型并行策略奠定了基础。

展望未来，我们的工作重点将是启用带有幽灵裁剪的 2D 并行，并将 FSDP 与使用钩子的原生 Opacus 集成。这些努力旨在进一步优化训练过程、降低延迟，并将 Opacus 的适用性扩展到更大、更复杂的模型。我们对这些进步将解锁的可能性感到兴奋，并将致力于突破隐私机器学习的边界。此外，我们诚邀开发人员、研究人员和爱好者加入我们的旅程。当我们继续改进 Opacus 时，您的贡献和见解非常宝贵。

鸣谢