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 在 ranks 之间分片。但是如果它是一个 FlatParameter 对象，我们如何存储元数据？这是局限性之一。在没有一些丑陋的 hack 的情况下，正确存储每个参数的元数据（例如 dtype、requires_grad 等）是不可能的。

FSDP2

这就是引入 FSDP2 的原因。它不使用 FlatParameter，而是使用 DTensor，它是 “Distributed Tensor”（分布式张量）的缩写。每个 DTensor 基本上代表一个在 ranks 之间分片的原始 torch.Tensor。它包含有关原始 torch.Tensor 及其分片方式、放置类型等的元数据。这就是为什么它被称为 per-parameter sharding（逐参数分片）。下图显示了差异

原始 Layer 的每个 Parameter 都在第 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 在 ranks 之间进行额外通信，这样，每个 rank 只保存自己的分片和相应的元数据
对于加载，它使用分片模型的 state_dict 来直接加载分片参数
支持异步检查点，其中参数首先复制到 CPU 内存，之后，主线程继续训练，而另一个线程将参数存储在磁盘上
内存效率和确定性内存使用，FSDP2 不再使用 recordStream，而是使用流到流同步（有关更多技术细节，请参阅此论坛帖子和此问题）
未来，计划通过 torch.compile 优化通信模式，进一步提高性能和内存效率

API 差异

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

以下是使用 accelerate CLI 时，配置选项的主要更改

之前 (`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 上加载模型，然后参数同步到其他 ranks，这与 `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 保存完整模型。 `accelerate` 尚不支持 `FULL_STATE_DICT`。
`--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