FSDP1 与 FSDP2 对比

本指南解释了 `FSDP1` 和 `FSDP2` 之间的主要区别，并帮助您以最小的改动将现有代码迁移到 `FSDP2`。

FSDP2 比 FSDP1 好在哪里？

首先，我们需要了解 `FSDP1` 和 `FSDP2` 的内部工作原理，以便理解它们之间的差异。这也有助于我们理解 `FSDP1` 的局限性以及 `FSDP2` 如何解决这些问题。

我们将讨论一个场景：我们有一个包含 3 个 `Linear` 层的 `Layer`，并使用 `FSDP` 将其包装以在 2 个 GPU 上进行分片。

FSDP1

首先，我们必须了解最初的 `FSDP1` 及其带来的局限性。它将每个 `FSDP` 模块表示为单个 `FlatParameter`，这是一个包含所有模块参数的一维张量，然后在不同的进程（ranks）之间进行分片。也就是说，如果用 `FSDP1` 包装 `Layer`，会得到如下结果：

您可能会注意到一个问题。整个 `Layer` 被展平成一个 `FlatParameter`，然后在各个进程之间分片。但如果它是一个单独的 `FlatParameter` 对象，我们如何存储元数据呢？这是其局限性之一。如果不使用一些不优雅的技巧，就无法正确存储每个参数的元数据，例如 `dtype`、`requires_grad` 等。

FSDP2

这就是为什么引入了 `FSDP2`。它不使用 `FlatParameter`，而是使用 `DTensor`，即“分布式张量”（Distributed Tensor）。每个 `DTensor` 基本上代表一个已在不同进程间分片的普通 `torch.Tensor`。它包含关于原始 `torch.Tensor` 的元数据以及它是如何分片的，放置类型是什么等等。这就是为什么它被称为“按参数分片”（per-parameter sharding）。下图显示了差异：

原始 `Layer` 的每个参数都在第 0 维上进行分片，并在 2 个 GPU 之间分割。现在，每个 `Linear` 层都是一个单独的 `DTensor`，按参数存储元数据变得可能且直接。

在上图中，为了适应屏幕显示，张量是在第 1 维上分片的，实际上，如上所述，它们是在第 0 维上分片的。

FSDP2 提供了什么？

`FSDP2` 是 PyTorch 完全分片数据并行训练 API 的新改进版本。它的主要优势是使用 `DTensor` 来表示分片参数。与 `FSDP1` 相比，它提供了：

更简单的内部实现，其中每个 `Parameter` 都是一个单独的 `DTensor`
由于上述原因，可以轻松实现部分参数冻结，这使得像 `LORA` 这样的方法可以直接使用
通过 `DTensor`，`FSDP2` 支持在同一模型中混合使用 `fp8` 和其他参数类型
使用 `SHARDED_STATE_DICT` 和 `torch.distributed.checkpoint`，可以实现更快、更简单的检查点，无需跨进程进行额外通信，这样，每个进程只保存自己的分片和相应的元数据
对于加载，它使用分片模型的 `state_dict` 直接加载分片参数
支持异步检查点，其中参数首先被复制到 CPU 内存，之后主线程继续训练，而另一个线程将参数存储到磁盘
内存效率和确定性的内存使用，`FSDP2` 不再使用 `recordStream`，而是使用流到流的同步（更多技术细节请参阅此论坛帖子和此 issue）
未来，计划通过 `torch.compile` 优化通信模式，进一步提高性能和内存效率

API 差异

我们已经讨论了内部差异，现在让我们讨论一下您作为用户需要了解的差异。

以下是通过 `accelerate` CLI 使用 `FSDP2` 时配置选项的主要变化：

旧版 (`FSDP1`)	新版 (`FSDP2`)	变化说明
`--fsdp_sharding_strategy`	`--fsdp_reshard_after_forward`	取代 `--fsdp_sharding_strategy`，改为 `true`（之前是 `FULL_SHARD`）或 `false`（之前是 `SHARD_GRAD_OP`）
`--fsdp_backward_prefetch`	**已移除**	`FSDP2` 默认使用之前的 `BACKWARD_PRE` 选项，因为只有这样才能实现通信和计算的重叠
`--fsdp_forward_prefetch`	**尚未实现**	如何实现此功能正在积极讨论中，目前在 `FSDP2` 中尚不支持
`--fsdp_sync_module_states`	**已移除**	使用 `FSDP2` 后，此参数变得多余
`--fsdp_cpu_ram_efficient_loading`	`--fsdp_cpu_ram_efficient_loading`	如果为 `true`，`FSDP2` 同样只在 rank 0 上加载模型，然后参数同步到其他 rank，这与 `FSDP1` 的行为相同，但是不再需要设置 `--fsdp_sync_module_states`
`--fsdp_state_dict_type`	`--fsdp_state_dict_type`	`LOCAL_STATE_DICT` 已过时，在 `FSDP2` 中 `SHARDED_STATE_DICT` 是默认选项，这导致没有额外的通信，并且每个 rank 保存自己的分片；另一个可能的选项是 `FULL_STATE_DICT`，它会导致额外的通信和内存使用激增，但会从 rank 0 保存完整的模型。
`--fsdp_use_orig_params`	**已移除**	`FSDP2` 在后台使用 `DTensor` 类，这意味着它默认总是使用原始参数
新增	`--fsdp_version`	默认选项为 `1`，以免破坏现有代码，设置为 `2` 以使用 `FSDP2`

对于所有其他未更改的选项，请参阅 `FSDP` 文档。

如何切换到 FSDP2

如果使用 Python 代码：

在创建插件时，只需设置 `fsdp_version=2`，并根据上表替换选项即可。

from accelerate import FullyShardedDataParallelPlugin, Accelerator

fsdp_plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2
    # other options...
)
accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

如果使用 YAML 配置：

使用我们的转换工具

accelerate to-fsdp2 --config_file config.yaml --output_file new_config.yaml

这将自动将所有 FSDP1 设置转换为其 FSDP2 的等效设置。使用 `--overwrite` 来更新现有文件，而不是创建一个新文件。

< > 在 GitHub 上更新

Accelerate