评价此页

★ ★ ★ ★ ★

intermediate/FSDP_tutorial

在 Google Colab 中运行

FSDP2（完全分片数据并行）入门#

创建日期：2022 年 3 月 17 日 | 最后更新：2025 年 7 月 24 日 | 最后验证：2024 年 11 月 05 日

作者: 魏锋, 威尔·康斯特布尔, 毛轶凡

注意

从 pytorch/examples 查看本教程中的代码。FSDP1 已弃用。FSDP1 教程已归档在 [1] 和 [2]

FSDP2 工作原理#

在 DistributedDataParallel (DDP) 训练中，每个秩都拥有一个模型副本并处理一批数据，最终它使用 all-reduce 在所有秩之间同步梯度。

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

在正向和反向计算之外，参数是完全分片的
在正向和反向之前，分片参数被 all-gather 成非分片参数
在反向传播内部，局部非分片梯度被 reduce-scatter 成共享梯度
优化器使用共享梯度更新共享参数，从而产生共享优化器状态

FSDP 可被视为 DDP 的 all-reduce 操作的分解，分解为 reduce-scatter 和 all-gather 操作

与 FSDP1 相比，FSDP2 具有以下优点

将分片参数表示为在 dim-i 上分片的 DTensor，从而可以轻松操作单个参数、无需通信的分片状态字典，以及更简单的元设备初始化流程。
改进的内存管理系统通过避免 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() as 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 密集型工作负载（例如大批量大小的 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

2 层以上预取：隐式预取一次只 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
将梯度上转换为 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 后初始化，并持有对 DTensor model.parameters() 的引用。对于梯度裁剪，torch.nn.utils.clip_grad_norm_ 适用于 DTensor 参数。Tensor 操作将在 DTensor 内部正确调度，以在各秩之间通信部分张量，从而保留单设备语义。

使用 DTensor API 的状态字典#

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

命令: torchrun --nproc_per_node 2 train.py

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

加载状态字典：我们在元设备下初始化模型并调用 fully_shard 将 model.parameters() 从普通的 torch.Tensor 转换为 DTensor。从 torch.load 读取完整状态字典后，我们可以调用 distribute_tensor 将普通的 torch.Tensor 转换为 DTensor，使用与 model.state_dict() 相同的放置和设备网格。最后，我们可以调用 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 形式）。对于秩 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，我们只在秩 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,
    ),
)

保存状态字典：get_model_state_dict 配合 full_state_dict=True 和 cpu_offload=True 会聚合张量并将其卸载到 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")

请参阅 pytorch/examples 以了解如何使用 set_optimizer_state_dict 和 get_optimizer_state_dict 加载和保存优化器状态字典。

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 kwargs（参见下文）

分片策略

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（CPU 卸载）

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

后向预取

BACKWARD_PRE: 总是使用
BACKWARD_POST: 不支持

混合精度

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 从秩 0 广播状态字典。

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: 始终使用原始参数（不再是扁平参数）

no_sync(): set_requires_gradient_sync

ignored_params 和 ignored_states: ignored_params

FSDP2（完全分片数据并行）入门#

FSDP2 工作原理#

如何使用 FSDP2#

模型初始化#

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

启用混合精度#

DTensor 的梯度裁剪和优化器#

使用 DTensor API 的状态字典#

使用 DCP API 的状态字典#

FSDP1 到 FSDP2 迁移指南#

文档

教程

资源