评价此页

★ ★ ★ ★ ★

intermediate/FSDP_tutorial

在 Google Colab 中运行

Fully Sharded Data Parallel (FSDP2) 入门#

创建日期：2022 年 3 月 17 日 | 最后更新：2025 年 9 月 2 日 | 最后验证：2024 年 11 月 5 日

作者: Wei Feng, Will Constable, Yifan Mao

注意

查看本教程的代码，请访问 pytorch/examples。FSDP1 已弃用。FSDP1 教程已归档至 [1] 和 [2]

FSDP2 的工作原理#

在 DistributedDataParallel (DDP) 训练中，每个进程（rank）拥有一个模型副本并处理一批数据，最后通过 all-reduce 在各个进程间同步梯度。

与 DDP 相比，FSDP 通过对模型参数、梯度和优化器状态进行分片（sharding）来减少 GPU 内存占用。这使得训练无法装入单个 GPU 的大模型成为可能。如下图所示：

在正向和反向计算之外，参数被完全分片
在正向和反向计算之前，分片后的参数被 all-gather 成非分片状态
在反向计算过程中，本地的非分片梯度被 reduce-scatter 成分片梯度
优化器使用分片梯度更新分片参数，从而产生分片后的优化器状态

FSDP 可以被看作是将 DDP 的 all-reduce 操作分解为 reduce-scatter 和 all-gather 操作

与 FSDP1 相比，FSDP2 具有以下优势：

将分片参数表示为在 dim-i 上分片的 DTensor，从而易于操作单个参数、实现无需通信的分片状态字典（state dict），以及更简单的 meta-device 初始化流程。
改进了内存管理系统，通过避免 recordStream (文档) 实现了更低且确定性的 GPU 内存占用，且无需任何 CPU 同步。
提供了一个张量子类扩展点以自定义 all-gather，例如用于 float8 线性层的 float8 all-gather (文档)，以及用于 QLoRA 的 NF4 (文档)。
可以在同一个通信组中混合冻结和非冻结参数，而无需使用额外的内存。

如何使用 FSDP2#

模型初始化#

在子模块上应用 fully_shard：与 DDP 不同，我们应该在子模块以及根模型上应用 fully_shard。在下面的 Transformer 示例中，我们先在每个层上应用 fully_shard，然后再在根模型上应用。

在 layers[i] 的正向计算过程中，其余层保持分片状态以减少内存占用。
在 fully_shard(model) 内部，FSDP2 将 model.layers 中的参数排除，并将剩余参数归类到一个参数组中，以便进行高效的 all-gather 和 reduce-scatter。
fully_shard 会将分片后的模型移动到实际的训练设备（例如 cuda）。

命令: torchrun --nproc_per_node 2 train.py

from torch.distributed.fsdp import fully_shard, FSDPModule
model = Transformer()
for layer in model.layers:
    fully_shard(layer)
fully_shard(model)

assert isinstance(model, Transformer)
assert isinstance(model, FSDPModule)
print(model)
#  FSDPTransformer(
#    (tok_embeddings): Embedding(...)
#    ...
#    (layers): 3 x FSDPTransformerBlock(...)
#    (output): Linear(...)
#  )

我们可以通过 print(model) 来检查嵌套包装的情况。FSDPTransformer 是 Transformer 和 FSDPModule 的联合类。FSDPTransformerBlock 也是如此。所有 FSDP2 公共 API 均通过 FSDPModule 公开。例如，用户可以调用 model.unshard() 来手动控制 all-gather 调度。详情请参见下文的“显式预取”。

model.parameters() 作为 DTensor：fully_shard 会在各进程间分片参数，并将 model.parameters() 从普通的 torch.Tensor 转换为 DTensor 以表示分片参数。FSDP2 默认在 dim-0 上进行分片，因此 DTensor 的放置策略为 Shard(dim=0)。假设有 N 个进程，参数在分片前有 N 行。分片后，每个进程拥有该参数的 1 行。我们可以使用 param.to_local() 查看分片后的参数。

from torch.distributed.tensor import DTensor
for param in model.parameters():
    assert isinstance(param, DTensor)
    assert param.placements == (Shard(0),)
    # inspect sharded parameters with param.to_local()

optim = torch.optim.Adam(model.parameters(), lr=1e-2)

注意：优化器是在应用 fully_shard 之后构建的。模型和优化器的状态字典均由 DTensor 表示。

DTensor 促进了优化器、梯度裁剪和检查点保存。

torch.optim.Adam 和 torch.nn.utils.clip_grad_norm_ 可直接用于 DTensor 参数。这使得单设备训练和分布式训练之间的代码保持一致。
我们可以使用 DTensor 和 DCP API 来操作参数以获取完整的状态字典；详情请见下文的“状态字典”部分。对于分布式状态字典，我们可以保存/加载检查点 (文档)，无需额外的通信。

带预取的正向/反向传播#

命令: torchrun --nproc_per_node 2 train.py

for _ in range(epochs):
    x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    loss = model(x).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

fully_shard 会注册正向/反向钩子，以便在计算前 all-gather 参数，并在计算后重新分片。为了使 all-gather 与计算重叠，FSDP2 提供了隐式预取（开箱即用）和显式预取（供高级用户手动控制 all-gather 调度）。

隐式预取：CPU 线程在第 i 层之前发起 all-gather i。All-gather 被加入到其专有的 CUDA 流中，而第 i 层的计算则在默认流中进行。对于非 CPU 受限的工作负载（例如大 Batch Size 的 Transformer），all-gather i+1 可以与第 i 层的计算重叠。反向传播中的隐式预取逻辑相似，区别在于 all-gather 的发起顺序是正向传播的逆序。

我们建议用户从隐式预取开始，以了解其开箱即用的性能。

显式预取：用户可以使用 set_modules_to_forward_prefetch 指定正向顺序，使用 set_modules_to_backward_prefetch 指定反向顺序。如下面的代码所示，CPU 线程在第 i 层发起 all-gather i + 1 和 i + 2。

显式预取在以下情况下效果显著：

CPU 受限的工作负载：如果使用隐式预取，当第 i 层的内核执行时，CPU 线程可能太慢，无法为第 i+1 层发起 all-gather。此时必须在运行第 i 层正向传播前显式发起 all-gather i+1。

多层预取：隐式预取一次仅 all-gather 下一层，以保持内存占用最小。使用显式预取可以一次 all-gather 多个层，从而在增加内存占用的同时可能获得更好的性能。参见代码中的 layers_to_prefetch。

更早发起第一次 all-gather：隐式预取发生在调用 model(x) 时。第一次 all-gather 会被暴露出来。我们可以显式调用 model.unshard() 来更早地发起第一次 all-gather。

命令: torchrun --nproc_per_node 2 train.py --explicit-prefetching

num_to_forward_prefetch = 2
for i, layer in enumerate(model.layers):
    if i >= len(model.layers) - num_to_forward_prefetch:
        break
    layers_to_prefetch = [
        model.layers[i + j] for j in range(1, num_to_forward_prefetch + 1)
    ]
    layer.set_modules_to_forward_prefetch(layers_to_prefetch)

num_to_backward_prefetch = 2
for i, layer in enumerate(model.layers):
    if i < num_to_backward_prefetch:
        continue
    layers_to_prefetch = [
        model.layers[i - j] for j in range(1, num_to_backward_prefetch + 1)
    ]
    layer.set_modules_to_backward_prefetch(layers_to_prefetch)

for _ in range(epochs):
    # trigger 1st all-gather earlier
    # this overlaps all-gather with any computation before model(x)
    model.unshard()
    x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    loss = model(x).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

启用混合精度#

FSDP2 提供了一种灵活的混合精度策略以加速训练。一个典型的用例是：

将 float32 参数转换为 bfloat16 以进行正向/反向计算，参见 param_dtype=torch.bfloat16。
将梯度提升（Upcasting）到 float32 以进行 reduce-scatter 以保持精度，参见 reduce_dtype=torch.float32。

与 torch.amp 相比，FSDP2 混合精度具有以下优势：

高性能且灵活的参数转换：FSDPModule 内部的所有参数都在模块边界（正向/反向传播之前和之后）进行转换。我们可以为每一层设置不同的混合精度策略。例如，前几层可以使用 float32，而其余层可以使用 bfloat16。
float32 梯度归约 (reduce-scatter)：不同进程间的梯度可能会有很大差异。在 float32 中进行梯度归约对于数值稳定性至关重要。

命令: torchrun --nproc_per_node 2 train.py --mixed-precision

model = Transformer(model_args)
fsdp_kwargs = {
    "mp_policy": MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )
}
for layer in model.layers:
    fully_shard(layer, **fsdp_kwargs)
fully_shard(model, **fsdp_kwargs)

# sharded parameters are float32
for param in model.parameters():
    assert param.dtype == torch.float32

# unsharded parameters are bfloat16
model.unshard()
for param in model.parameters(recurse=False):
    assert param.dtype == torch.bfloat16
model.reshard()

# optimizer states are in float32
optim = torch.optim.Adam(model.parameters(), lr=1e-2)

# training loop
# ...

DTensor 的梯度裁剪与优化器#

命令: torchrun --nproc_per_node 2 train.py

# optim is constructed base on DTensor model parameters
optim = torch.optim.Adam(model.parameters(), lr=1e-2)
for _ in range(epochs):
    x = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
    loss = model(x).sum()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)
    optim.step()
    optim.zero_grad()

优化器是在模型应用 fully_shard 后初始化的，并持有 model.parameters() 的 DTensor 引用。对于梯度裁剪，torch.nn.utils.clip_grad_norm_ 可用于 DTensor 参数。张量操作将在 DTensor 内部被正确调度，以便在不同进程间通信部分张量，从而保持单设备语义。

使用 DTensor API 处理状态字典#

我们展示了如何将完整状态字典转换为 DTensor 状态字典进行加载，以及如何将其转换回完整状态字典进行保存。

命令: torchrun --nproc_per_node 2 train.py

第一次时，它为模型和优化器创建检查点。
第二次时，它从之前的检查点加载以恢复训练。

加载状态字典：我们在 meta 设备下初始化模型，并调用 fully_shard 将 model.parameters() 从普通的 torch.Tensor 转换为 DTensor。在从 torch.load 读取完整状态字典后，我们可以调用 distribute_tensor，使用来自 model.state_dict() 的相同放置策略和设备网格，将普通 torch.Tensor 转换为 DTensor。最后，我们调用 model.load_state_dict 将 DTensor 状态字典加载到模型中。

from torch.distributed.tensor import distribute_tensor

# mmap=True reduces CPU memory usage
full_sd = torch.load(
    "checkpoints/model_state_dict.pt",
    mmap=True,
    weights_only=True,
    map_location='cpu',
)
meta_sharded_sd = model.state_dict()
sharded_sd = {}
for param_name, full_tensor in full_sd.items():
    sharded_meta_param = meta_sharded_sd.get(param_name)
    sharded_tensor = distribute_tensor(
        full_tensor,
        sharded_meta_param.device_mesh,
        sharded_meta_param.placements,
    )
    sharded_sd[param_name] = nn.Parameter(sharded_tensor)
# `assign=True` since we cannot call `copy_` on meta tensor
model.load_state_dict(sharded_sd, assign=True)

保存状态字典：model.state_dict() 返回一个 DTensor 状态字典。我们可以通过调用 full_tensor() 将 DTensor 转换为普通的 torch.Tensor。它在内部执行跨进程的 all-gather，以获取非分片的普通 torch.Tensor。对于 rank 0，full_param.cpu() 会逐个将张量卸载到 CPU，以避免因未分片的参数导致 GPU 内存峰值。

sharded_sd = model.state_dict()
cpu_state_dict = {}
for param_name, sharded_param in sharded_sd.items():
    full_param = sharded_param.full_tensor()
    if torch.distributed.get_rank() == 0:
        cpu_state_dict[param_name] = full_param.cpu()
    else:
        del full_param
torch.save(cpu_state_dict, "checkpoints/model_state_dict.pt")

优化器状态字典的处理方式类似 (代码)。用户可以自定义上述 DTensor 脚本以适配第三方检查点。

如果无需自定义，我们可以直接使用 DCP API，它同时支持单节点和多节点训练。

使用 DCP API 处理状态字典#

命令: torchrun --nproc_per_node 2 train.py --dcp-api

第一次时，它为模型和优化器创建检查点。
第二次时，它从之前的检查点加载以恢复训练。

加载状态字典：我们可以使用 set_model_state_dict 将完整状态字典加载到 FSDP2 模型中。通过设置 broadcast_from_rank0=True，我们仅需在 rank 0 上加载完整状态字典，以避免 CPU 内存峰值。DCP 会自动对张量进行分片并广播到其他进程。

from torch.distributed.checkpoint.state_dict import set_model_state_dict
set_model_state_dict(
    model=model,
    model_state_dict=full_sd,
    options=StateDictOptions(
        full_state_dict=True,
        broadcast_from_rank0=True,
    ),
)

保存状态字典：使用 full_state_dict=True 和 cpu_offload=True 的 get_model_state_dict 会执行 all-gather 并将张量卸载到 CPU。其工作原理与 DTensor API 类似。

from torch.distributed.checkpoint.state_dict import get_model_state_dict
model_state_dict = get_model_state_dict(
    model=model,
    options=StateDictOptions(
        full_state_dict=True,
        cpu_offload=True,
    )
)
torch.save(model_state_dict, "model_state_dict.pt")

有关使用 set_optimizer_state_dict 和 get_optimizer_state_dict 加载和保存优化器状态字典的示例，请参考 pytorch/examples。

FSDP1 到 FSDP2 迁移指南#

让我们看一个 FSDP 用法和等效的 fully_shard 用法示例。我们将重点介绍关键差异并建议迁移步骤。

原始 FSDP() 用法

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
with torch.device("meta"):
    model = Transformer()
policy = ModuleWrapPolicy({TransformerBlock})
model = FSDP(model, auto_wrap_policy=policy)
def param_init_fn(module: nn.Module) -> None: ...
model = FSDP(model, auto_wrap_policy=policy, param_init_fn=param_init_fn)

新的 fully_shard() 用法

with torch.device("meta"):
    model = Transformer()
for module in model.modules():
    if isinstance(module, TransformerBlock):
        fully_shard(module)
fully_shard(model)
for tensor in itertools.chain(model.parameters(), model.buffers()):
    assert tensor.device == torch.device("meta")


# Initialize the model after sharding
model.to_empty(device="cuda")
model.reset_parameters()

迁移步骤

替换导入
直接实现您的“策略”（将 fully_shard 应用于所需的子层）
使用 fully_shard 而不是 FSDP 包装根模型
弃用 param_init_fn，改为手动调用 model.reset_parameters()
替换其他 FSDP1 关键字参数（见下文）

sharding_strategy

FULL_SHARD: reshard_after_forward=True
SHARD_GRAD_OP: reshard_after_forward=False
HYBRID_SHARD: reshard_after_forward=True (配合 2D 设备网格)
_HYBRID_SHARD_ZERO2: reshard_after_forward=False (配合 2D 设备网格)

cpu_offload

CPUOffload.offload_params=False: offload_policy=None
CPUOffload.offload_params = True: offload_policy=CPUOffloadPolicy()

backward_prefetch

BACKWARD_PRE: 始终使用
BACKWARD_POST: 不支持

mixed_precision

buffer_dtype 被省略，因为 fully_shard 不会对缓冲区进行分片
fully_shard 的 cast_forward_inputs 对应 FSDP1 中的 cast_forward_inputs 和 cast_root_forward_inputs
output_dtype 是 fully_shard 的一个新配置

device_id: 从 device_mesh 的设备中推断

sync_module_states=True/False: 移至 DCP。用户可以使用 set_model_state_dict 配合 broadcast_from_rank0=True 从 rank0 广播状态字典

forward_prefetch: 可以通过以下方式手动控制预取

手动调用 fsdp_module.unshard()
使用 API set_modules_to_forward_prefetch 和 set_modules_to_backward_prefetch 控制自动预取

limit_all_gathers: 不再需要，因为 fully_shard 移除了 CPU 同步

use_orig_params: 始终使用原始参数（不再有平铺参数 flat parameter）

no_sync(): set_requires_gradient_sync

ignored_params 和 ignored_states: ignored_params

Fully Sharded Data Parallel (FSDP2) 入门#

FSDP2 的工作原理#

如何使用 FSDP2#

模型初始化#

带预取的正向/反向传播#

启用混合精度#

DTensor 的梯度裁剪与优化器#

使用 DTensor API 处理状态字典#

使用 DCP API 处理状态字典#

FSDP1 到 FSDP2 迁移指南#

文档

教程

资源