FullyShardedDataParallel

class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, sharding_strategy=None, cpu_offload=None, auto_wrap_policy=None, backward_prefetch=BackwardPrefetch.BACKWARD_PRE, mixed_precision=None, ignored_modules=None, param_init_fn=None, device_id=None, sync_module_states=False, forward_prefetch=False, limit_all_gathers=True, use_orig_params=False, ignored_states=None, device_mesh=None)

跨数据并行工作程序对模块参数进行分片的包装器。

这受到 Xu 等人 以及 DeepSpeed 中的 ZeRO 阶段 3 的启发。FullyShardedDataParallel 通常缩写为 FSDP。

要了解 FSDP 内部机制，请参阅 FSDP 笔记。

示例

>>> import torch
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> torch.cuda.set_device(device_id)
>>> sharded_module = FSDP(my_module)
>>> optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
>>> x = sharded_module(x, y=3, z=torch.Tensor([1]))
>>> loss = x.sum()
>>> loss.backward()
>>> optim.step()

使用 FSDP 包含包装您的模块，然后在之后初始化您的优化器。这是必需的，因为 FSDP 会更改参数变量。

设置 FSDP 时，您需要考虑目标 CUDA 设备。如果设备具有 ID (dev_id)，则您有三种选择

• 将模块放置在该设备上

• 使用 torch.cuda.set_device(dev_id) 设置设备

• 将 dev_id 传递到 device_id 构造函数参数中。

这将确保 FSDP 实例的计算设备是目标设备。对于选项 1 和 3，FSDP 初始化始终发生在 GPU 上。对于选项 2，FSDP 初始化发生在模块的当前设备上，这可能是 CPU。

如果您使用的是 sync_module_states=True 标志，您需要确保模块位于 GPU 上，或者使用 device_id 参数指定 FSDP 将在 FSDP 构造函数中将模块移动到的 CUDA 设备。这是必需的，因为 sync_module_states=True 需要 GPU 通信。

FSDP 还负责将输入张量移动到转发方法的目标 GPU 计算设备上，因此您不需要手动将它们从 CPU 移动。

对于 use_orig_params=True，ShardingStrategy.SHARD_GRAD_OP 会公开未分片的参数，而不是转发后分片的参数，这与 ShardingStrategy.FULL_SHARD 不同。如果您想检查梯度，可以使用 summon_full_params 方法，其中 with_grads=True。

对于 limit_all_gathers=True，您可能会在 FSDP 预转发中看到一个差距，其中 CPU 线程没有发出任何内核。这是故意的，它表明速率限制器生效。以这种方式同步 CPU 线程可以防止为后续的全部收集过度分配内存，并且它实际上不应该延迟 GPU 内核执行。

FSDP 会在转发和反向计算期间使用 torch.Tensor 视图替换受管模块的参数，这是由于与 autograd 相关的缘故。如果您的模块的转发依赖于对参数的保存引用，而不是每次迭代都重新获取引用，那么它将看不到 FSDP 新创建的视图，并且 autograd 将无法正常工作。

最后，当使用sharding_strategy=ShardingStrategy.HYBRID_SHARD，且分片进程组为节点内，复制进程组为节点间时，设置NCCL_CROSS_NIC=1可以帮助提高某些集群设置下复制进程组的all-reduce时间。

限制

使用 FSDP 时，需要注意以下几个限制。

• FSDP 当前不支持在使用 CPU 卸载时，在 no_sync() 之外进行梯度累积。这是因为 FSDP 使用新减少的梯度，而不是与任何现有梯度累积，这会导致结果不正确。

• FSDP 不支持运行包含在 FSDP 实例中的子模块的前向传递。这是因为子模块的参数将被分片，但子模块本身不是 FSDP 实例，因此它的前向传递不会适当地将所有参数全部收集。

• FSDP 由于其注册后向钩子的方式，不支持双重反向传播。

• FSDP 在冻结参数时有一些约束。对于 use_orig_params=False，每个 FSDP 实例必须管理所有冻结或所有非冻结的参数。对于 use_orig_params=True，FSDP 支持混合冻结和非冻结参数，但建议避免这样做，以防止比预期更高的梯度内存使用率。

• 截至 PyTorch 1.12，FSDP 对共享参数的支持有限。如果您的用例需要增强共享参数支持，请在此问题中发帖。

• 您应该避免在不使用 summon_full_params 上下文的情况下，在正向传播和反向传播之间修改参数，因为修改可能不会持久化。

参数

• module (nn.Module) – 这是要使用 FSDP 包装的模块。

• process_group (Optional[Union[ProcessGroup, Tuple[ProcessGroup, ProcessGroup]]]) – 这是模型被分片的进程组，因此是 FSDP 的全收集和归约散射集体通信使用的进程组。如果为 None，则 FSDP 使用默认进程组。对于混合分片策略，例如 ShardingStrategy.HYBRID_SHARD，用户可以传入一个进程组元组，分别表示要分片和复制的组。如果为 None，则 FSDP 为用户构建进程组以在节点内分片并在节点间复制。（默认：None）

• sharding_strategy (Optional[ShardingStrategy]) – 这配置了分片策略，这可能会权衡内存节省和通信开销。有关详细信息，请参阅ShardingStrategy。（默认：FULL_SHARD）

• cpu_offload (Optional[CPUOffload]) – 这配置了 CPU 卸载。如果设置为 None，则不会进行 CPU 卸载。有关详细信息，请参阅CPUOffload。（默认：None）

• auto_wrap_policy (Optional[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy, CustomPolicy]]) –

  这指定了一个策略，将 FSDP 应用于 module 的子模块，这是通信和计算重叠所必需的，因此会影响性能。如果为 None，则 FSDP 只应用于 module，用户应自己手动将 FSDP 应用于父模块（自下而上进行）。为了方便起见，这直接接受 ModuleWrapPolicy，这允许用户指定要包装的模块类（例如，transformer 块）。否则，这应该是一个可调用对象，它接受三个参数 module: nn.Module、recurse: bool 和 nonwrapped_numel: int，并且应该返回一个 bool，指定是否应该将 FSDP 应用于传入的 module（如果 recurse=False）或如果遍历应该继续进入模块的子树（如果 recurse=True）。用户可以向可调用对象添加其他参数。 torch.distributed.fsdp.wrap.py 中的 size_based_auto_wrap_policy 给出了一个示例可调用对象，该对象如果子树中的参数超过 100M 个 numel，则会将 FSDP 应用于该模块。我们建议在应用 FSDP 后打印模型并根据需要进行调整。

  示例

  >>> def custom_auto_wrap_policy(
  >>>     module: nn.Module,
  >>>     recurse: bool,
  >>>     nonwrapped_numel: int,
  >>>     # Additional custom arguments
  >>>     min_num_params: int = int(1e8),
  >>> ) -> bool:
  >>>     return nonwrapped_numel >= min_num_params
  >>> # Configure a custom `min_num_params`
  >>> my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
  

• backward_prefetch (Optional[BackwardPrefetch]) – 这配置了对所有收集的显式后向预取。如果为 None，则 FSDP 不会进行后向预取，并且在后向传递中没有通信和计算重叠。有关详细信息，请参阅BackwardPrefetch。（默认：BACKWARD_PRE）

• mixed_precision (Optional[MixedPrecision]) – 这配置了 FSDP 的原生混合精度。如果设置为 None，则不使用混合精度。否则，可以设置参数、缓冲区和梯度缩减数据类型。有关详细信息，请参阅MixedPrecision。（默认：None）

• ignored_modules (Optional[Iterable[torch.nn.Module]]) – 忽略其自身参数和子模块参数和缓冲区的模块。ignored_modules 中的任何模块都不应该是FullyShardedDataParallel 实例，并且如果它们嵌套在此实例下，任何已经是FullyShardedDataParallel 实例的子模块都不会被忽略。此参数可用于在使用 auto_wrap_policy 或参数的分片未由 FSDP 管理时，避免在模块粒度上分片特定参数。（默认：None）

• param_init_fn (Optional[Callable[[nn.Module], None]]) –

  一个 Callable[torch.nn.Module] -> None，它指定了当前在元设备上的模块如何初始化到实际设备上。截至 v1.12，FSDP 通过 is_meta 检测具有元设备上的参数或缓冲区的模块，并应用 param_init_fn（如果指定）或调用 nn.Module.reset_parameters()（否则）。对于这两种情况，实现都只应初始化模块的参数/缓冲区，而不应初始化其子模块的参数/缓冲区。这是为了避免重新初始化。此外，FSDP 还支持通过 torchdistX 的 (https://github.com/pytorch/torchdistX) deferred_init() API 进行延迟初始化，其中延迟模块通过调用 param_init_fn（如果指定）或 torchdistX 的默认 materialize_module()（否则）进行初始化。如果指定了 param_init_fn，则它将应用于所有元设备模块，这意味着它可能应该根据模块类型进行判断。FSDP 在参数扁平化和分片之前调用初始化函数。

  示例

  >>> module = MyModule(device="meta")
  >>> def my_init_fn(module: nn.Module):
  >>>     # E.g. initialize depending on the module type
  >>>     ...
  >>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
  >>> print(next(fsdp_model.parameters()).device) # current CUDA device
  >>> # With torchdistX
  >>> module = deferred_init.deferred_init(MyModule, device="cuda")
  >>> # Will initialize via deferred_init.materialize_module().
  >>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)
  

• device_id (Optional[Union[int, torch.device]]) – 一个 int 或 torch.device，给出 FSDP 初始化发生在其上的 CUDA 设备，包括模块初始化（如果需要）和参数分片。如果 module 在 CPU 上，则应指定此参数以提高初始化速度。如果设置了默认 CUDA 设备（例如，通过 torch.cuda.set_device），则用户可以将 torch.cuda.current_device 传递给它。（默认：None）

• sync_module_states (bool) – 如果为 True，则每个 FSDP 模块将从排名 0 广播模块参数和缓冲区，以确保它们在排名之间复制（在该构造函数中添加通信开销）。这可以帮助以内存高效的方式通过 load_state_dict 加载 state_dict 检查点。有关示例，请参阅FullStateDictConfig。（默认：False）

• forward_prefetch (bool) – 如果为 True，则 FSDP 会显式地在当前前向计算之前预取下一个前向传递的全部聚合。这仅对 CPU 密集型工作负载有用，在这种情况下，更早地发出下一个全部聚合可能会提高重叠。这应该只用于静态图模型，因为预取遵循第一次迭代的执行顺序。（默认：False）

• limit_all_gathers (bool) – 如果为 True，则 FSDP 会显式地同步 CPU 线程以确保仅来自两个连续 FSDP 实例的 GPU 内存使用量（当前实例正在运行计算和下一个实例的全部聚合已预取）。如果为 False，则 FSDP 允许 CPU 线程在没有任何额外同步的情况下发出全部聚合。（默认：True）我们通常将此功能称为“速率限制器”。此标志应该只在内存压力低的特定 CPU 密集型工作负载的情况下设置为 False，在这种情况下，CPU 线程可以积极地发出所有内核，而无需担心 GPU 内存使用情况。

• use_orig_params (bool) – 将此设置为 True 将使 FSDP 使用 module 的原始参数。FSDP 通过 nn.Module.named_parameters() 向用户公开这些原始参数，而不是 FSDP 的内部 FlatParameter。这意味着优化器步骤在原始参数上运行，从而启用每个原始参数的超参数。FSDP 保留原始参数变量并在非分片形式和分片形式之间操作其数据，其中它们始终分别是底层非分片或分片 FlatParameter 的视图。使用当前算法，分片形式始终为 1D，从而丢失了原始张量结构。对于给定的秩，原始参数可能包含全部、部分或不包含其数据。在没有数据的情况下，其数据将类似于大小为 0 的空张量。用户不应编写依赖于给定原始参数在其分片形式中存在哪些数据的程序。需要使用 True 来使用 torch.compile()。将此设置为 False 会通过 nn.Module.named_parameters() 向用户公开 FSDP 的内部 FlatParameter。（默认：False）

• ignored_states (可选[Iterable[torch.nn.Parameter]], 可选[Iterable[torch.nn.Module]]) – 被忽略的参数或模块，它们将不会由此 FSDP 实例管理，这意味着这些参数不会被分片，它们的梯度不会在秩之间减少。此参数与现有的 ignored_modules 参数统一，我们可能很快就会弃用 ignored_modules。为了向后兼容性，我们同时保留了 ignored_states 和 ignored_modules`，但 FSDP 仅允许其中一个被指定为非 None。

• device_mesh (可选[DeviceMesh]) – DeviceMesh 可用作 process_group 的替代方案。当传递 device_mesh 时，FSDP 将使用底层进程组进行全部聚合和减少散射集体通信。因此，这两个参数需要互斥。对于混合分片策略，例如 ShardingStrategy.HYBRID_SHARD，用户可以传递一个 2D DeviceMesh 而不是进程组元组。对于 2D FSDP + TP，用户需要传递 device_mesh 而不是 process_group。有关 DeviceMesh 的更多信息，请访问：https://pytorch.ac.cn/tutorials/recipes/distributed_device_mesh.html

apply(fn)

将 fn 递归地应用于每个子模块（如 .children() 所返回的），以及自身。

典型用法包括初始化模型的参数（另请参见 torch.nn.init）。

与 torch.nn.Module.apply 相比，此版本在应用 fn 之前还会收集完整的参数。不应在另一个 summon_full_params 上下文中调用它。

参数

fn (Module -> None) – 要应用于每个子模块的函数

返回值

self

返回类型

Module

check_is_root()

检查此实例是否为根 FSDP 模块。

返回类型

bool

clip_grad_norm_(max_norm, norm_type=2.0)

对所有参数的梯度范数进行裁剪。

范数是在所有参数的梯度上计算的，这些梯度被视为单个向量，并且梯度在原地被修改。

参数

• max_norm (float 或 int) – 梯度的最大范数

• norm_type (float 或 int) – 使用的 p 范数的类型。可以是 'inf' 表示无穷范数。

返回值

参数的总范数（被视为单个向量）。

返回类型

Tensor

如果每个 FSDP 实例都使用 NO_SHARD，这意味着没有梯度在秩之间分片，那么您可以直接使用 torch.nn.utils.clip_grad_norm_().

如果至少一些 FSDP 实例使用分片策略（即除了 NO_SHARD 之外的任何策略），那么您应该使用此方法而不是 torch.nn.utils.clip_grad_norm_()，因为此方法处理了梯度在秩之间分片的事实。

返回的总范数将具有所有参数/梯度中按 PyTorch 类型提升语义定义的“最大”数据类型。例如，如果所有参数/梯度都使用低精度数据类型，则返回的范数的数据类型将是该低精度数据类型，但如果存在至少一个使用 FP32 的参数/梯度，则返回的范数的数据类型将是 FP32。

警告

这需要在所有秩上调用，因为它使用集体通信。

static flatten_sharded_optim_state_dict(sharded_optim_state_dict, model, optim)

展平分片优化器状态字典。

API 类似于 shard_full_optim_state_dict()。唯一的区别是输入 sharded_optim_state_dict 应该从 sharded_optim_state_dict() 返回。因此，每个秩上都将有全部聚合调用来收集 ShardedTensor。

参数

• sharded_optim_state_dict (Dict[str, Any]) – 与未展平的参数相对应的优化器状态字典，并保存分片优化器状态。

• model (torch.nn.Module) – 请参阅 shard_full_optim_state_dict()。

• optim (torch.optim.Optimizer) – model 参数的优化器。

返回值

请参阅 shard_full_optim_state_dict()。

返回类型

Dict[str, Any]

forward(*args, **kwargs)

运行包装模块的前向传递，插入 FSDP 特定的前向和后向分片逻辑。

返回类型

Any

static fsdp_modules(module, root_only=False)

返回所有嵌套的 FSDP 实例。

这可能包括 module 本身，并且仅在 root_only=True 时才包括 FSDP 根模块。

参数

• module (torch.nn.Module) – 根模块，它可能是一个 FSDP 模块，也可能不是。

• root_only (bool) – 是否仅返回 FSDP 根模块。（默认值：False)

返回值

输入 module 中嵌套的 FSDP 模块。

返回类型

List[FullyShardedDataParallel]

static full_optim_state_dict(model, optim, optim_input=None, rank0_only=True, group=None)

返回完整的优化器状态字典。

在 rank 0 上整合完整的优化器状态，并将其作为 dict 返回，遵循 torch.optim.Optimizer.state_dict() 的约定，即使用 "state" 和 "param_groups" 作为键。 FSDP 模块中包含的扁平化参数映射回其未扁平化的参数。

这需要在所有 rank 上调用，因为它使用集体通信。但是，如果 rank0_only=True，则状态字典仅在 rank 0 上填充，所有其他 rank 返回一个空 dict。

与 torch.optim.Optimizer.state_dict() 不同，此方法使用完整的参数名称作为键，而不是参数 ID。

与 torch.optim.Optimizer.state_dict() 一样，优化器状态字典中包含的张量不会被克隆，因此可能会出现别名意外情况。为确保最佳实践，请考虑立即保存返回的优化器状态字典，例如使用 torch.save()。

参数

• model (torch.nn.Module) – 根模块（它可能是一个 FSDP 实例，也可能不是），其参数被传递到优化器 optim 中。

• optim (torch.optim.Optimizer) – model 参数的优化器。

• optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递到优化器 optim 的输入，表示一个参数组列表或一个参数迭代器；如果为 None，则此方法假定输入为 model.parameters()。此参数已弃用，不再需要传递它。（默认值：None)

• rank0_only (bool) – 如果为 True，则仅在 rank 0 上保存填充的 dict；如果为 False，则在所有 rank 上保存。（默认值：True)

• group (dist.ProcessGroup) – 模型的进程组，如果使用默认进程组，则为 None。（默认值：None)

返回值

一个 dict，其中包含 model 的原始未扁平化参数的优化器状态，并包含 “state” 和 “param_groups” 键，遵循 torch.optim.Optimizer.state_dict() 的约定。如果 rank0_only=True，则非零 rank 返回一个空的 dict。

返回类型

Dict[str, Any]

static get_state_dict_type(module)

获取以 module 为根的 FSDP 模块的 state_dict_type 及其相应的配置。

目标模块不必是 FSDP 模块。

返回值

一个 StateDictSettings，其中包含 state_dict_type 和当前设置的 state_dict/optim_state_dict 配置。

Raises

• AssertionError` if the StateDictSettings for differen –

• FSDP submodules differ. –

返回类型

StateDictSettings

property module: Module

返回包装的模块。

named_buffers(*args, **kwargs)

返回模块缓冲区的迭代器，同时生成缓冲区的名称和缓冲区本身。

拦截缓冲区名称，并在 summon_full_params() 上下文管理器内删除所有 FSDP 特定的扁平化缓冲区前缀。

返回类型

Iterator[Tuple[str, Tensor]]

named_parameters(*args, **kwargs)

返回模块参数的迭代器，同时生成参数的名称和参数本身。

拦截参数名称，并在 summon_full_params() 上下文管理器内删除所有 FSDP 特定的扁平化参数前缀。

返回类型

Iterator[Tuple[str, Parameter]]

no_sync()

禁用跨 FSDP 实例的梯度同步。

在此上下文中，梯度将在模块变量中累积，这些变量将在退出上下文后的第一个前向-反向传递中同步。这应该只在根 FSDP 实例上使用，并且会递归地应用于所有子 FSDP 实例。

Note

这可能会导致更高的内存使用，因为 FSDP 将在最终同步之前累积完整的模型梯度（而不是梯度分片）。

Note

在使用 CPU 卸载时，梯度在上下文管理器内不会卸载到 CPU。相反，它们只会在最终同步之后立即卸载。

返回类型

Generator

static optim_state_dict(model, optim, optim_state_dict=None, group=None)

转换与分片模型对应的优化器的状态字典。

给定的状态字典可以转换为三种类型之一：1) 全优化器状态字典，2) 分片优化器状态字典，3) 局部优化器状态字典。

对于全优化器状态字典，所有状态均未扁平化且未分片。可以通过 state_dict_type() 指定仅排名 0 且仅 CPU 以避免 OOM。

对于分片优化器状态字典，所有状态均未扁平化，但已分片。可以通过 state_dict_type() 指定仅 CPU 以进一步节省内存。

对于本地状态字典，将不会执行任何转换。但是，状态将从 nn.Tensor 转换为 ShardedTensor 以表示其分片性质（此功能尚不支持）。

示例

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> from torch.distributed.fsdp import FullStateDictConfig
>>> from torch.distributed.fsdp import FullOptimStateDictConfig
>>> # Save a checkpoint
>>> model, optim = ...
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> state_dict = model.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(model, optim)
>>> save_a_checkpoint(state_dict, optim_state_dict)
>>> # Load a checkpoint
>>> model, optim = ...
>>> state_dict, optim_state_dict = load_a_checkpoint()
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> model.load_state_dict(state_dict)
>>> optim_state_dict = FSDP.optim_state_dict_to_load(
>>>     model, optim, optim_state_dict
>>> )
>>> optim.load_state_dict(optim_state_dict)

参数

• model (torch.nn.Module) – 根模块（它可能是一个 FSDP 实例，也可能不是），其参数被传递到优化器 optim 中。

• optim (torch.optim.Optimizer) – model 参数的优化器。

• optim_state_dict (Dict[str, Any]) – 要转换的目标优化器状态字典。如果值为 None，将使用 optim.state_dict()。（默认值：None）

• group (dist.ProcessGroup) – 模型的进程组，参数在其中分片或 None 如果使用默认进程组。（默认值：None）

返回值

包含 model 的优化器状态的 dict。优化器状态的分片基于 state_dict_type。

返回类型

Dict[str, Any]

static optim_state_dict_to_load(model, optim, optim_state_dict, is_named_optimizer=False, load_directly=False, group=None)

转换优化器状态字典，以便将其加载到与 FSDP 模型关联的优化器中。

给定通过 optim_state_dict() 转换的 optim_state_dict，它将被转换为可以加载到 optim（model 的优化器）中的扁平优化器状态字典。 model 必须由 FullyShardedDataParallel 分片。

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.distributed.fsdp import StateDictType
>>> from torch.distributed.fsdp import FullStateDictConfig
>>> from torch.distributed.fsdp import FullOptimStateDictConfig
>>> # Save a checkpoint
>>> model, optim = ...
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> state_dict = model.state_dict()
>>> original_osd = optim.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(
>>>     model,
>>>     optim,
>>>     optim_state_dict=original_osd
>>> )
>>> save_a_checkpoint(state_dict, optim_state_dict)
>>> # Load a checkpoint
>>> model, optim = ...
>>> state_dict, optim_state_dict = load_a_checkpoint()
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.FULL_STATE_DICT,
>>>     FullStateDictConfig(rank0_only=False),
>>>     FullOptimStateDictConfig(rank0_only=False),
>>> )
>>> model.load_state_dict(state_dict)
>>> optim_state_dict = FSDP.optim_state_dict_to_load(
>>>     model, optim, optim_state_dict
>>> )
>>> optim.load_state_dict(optim_state_dict)

参数

• model (torch.nn.Module) – 根模块（它可能是一个 FSDP 实例，也可能不是），其参数被传递到优化器 optim 中。

• optim (torch.optim.Optimizer) – model 参数的优化器。

• optim_state_dict (Dict[str, Any]) – 要加载的优化器状态。

• is_named_optimizer (bool) – 此优化器是 NamedOptimizer 还是 KeyedOptimizer。仅在 optim 是 TorchRec 的 KeyedOptimizer 或 torch.distributed 的 NamedOptimizer 时设置为 True。

• load_directly (bool) – 如果设置为 True，此 API 还会在返回结果之前调用 optim.load_state_dict(result)。否则，用户负责调用 optim.load_state_dict()（默认值：False）

• group (dist.ProcessGroup) – 模型的进程组，参数在其中分片或 None 如果使用默认进程组。（默认值：None）

返回类型

Dict[str, Any]

register_comm_hook(state, hook)

注册通信钩子。

这是一个增强功能，为用户提供了一个灵活的钩子，他们可以在其中指定 FSDP 如何跨多个工作进程聚合梯度。此钩子可用于实现多种算法，例如 GossipGrad 和梯度压缩，它们涉及在使用 FullyShardedDataParallel 训练期间进行参数同步的不同通信策略。

警告

应在运行初始正向传递之前且仅运行一次注册 FSDP 通信钩子。

参数

• state (object) –

  传递给钩子以在训练过程中维护任何状态信息。例如，梯度压缩中的错误反馈、GossipGrad 中要与其通信的同伴等等。它由每个工作进程本地存储，并由工作进程上的所有梯度张量共享。

• hook (Callable) – 可调用函数，具有以下签名之一：1) hook: Callable[torch.Tensor] -> None：此函数接收一个 Python 张量，它表示与此 FSDP 单元包装的模型（未由其他 FSDP 子单元包装）相对应的所有变量的全扁平化未分片梯度。然后执行所有必要的处理并返回 None； 2) hook: Callable[torch.Tensor, torch.Tensor] -> None：此函数接收两个 Python 张量，第一个张量表示与此 FSDP 单元包装的模型（未由其他 FSDP 子单元包装）相对应的所有变量的全扁平化未分片梯度。第二个张量表示一个预先大小的张量，用于存储减少后的分片梯度的块。在这两种情况下，可调用函数执行所有必要的处理并返回 None。签名为 1 的可调用函数应处理 NO_SHARD 案例的梯度通信。签名为 2 的可调用函数应处理分片案例的梯度通信。

static rekey_optim_state_dict(optim_state_dict, optim_state_key_type, model, optim_input=None, optim=None)

重新键入优化器状态字典 optim_state_dict 以使用键类型 optim_state_key_type。

这可用于实现来自具有 FSDP 实例的模型和不具有 FSDP 实例的模型的优化器状态字典之间的兼容性。

要重新键入 FSDP 全优化器状态字典（即来自 full_optim_state_dict()）以使用参数 ID 并可加载到未包装的模型中

>>> wrapped_model, wrapped_optim = ...
>>> full_osd = FSDP.full_optim_state_dict(wrapped_model, wrapped_optim)
>>> nonwrapped_model, nonwrapped_optim = ...
>>> rekeyed_osd = FSDP.rekey_optim_state_dict(full_osd, OptimStateKeyType.PARAM_ID, nonwrapped_model)
>>> nonwrapped_optim.load_state_dict(rekeyed_osd)

要重新键入来自未包装模型的普通优化器状态字典，以便将其加载到包装模型中

>>> nonwrapped_model, nonwrapped_optim = ...
>>> osd = nonwrapped_optim.state_dict()
>>> rekeyed_osd = FSDP.rekey_optim_state_dict(osd, OptimStateKeyType.PARAM_NAME, nonwrapped_model)
>>> wrapped_model, wrapped_optim = ...
>>> sharded_osd = FSDP.shard_full_optim_state_dict(rekeyed_osd, wrapped_model)
>>> wrapped_optim.load_state_dict(sharded_osd)

返回值

使用 optim_state_key_type 指定的参数键重新键入的优化器状态字典。

返回类型

Dict[str, Any]

static scatter_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None, group=None)

将完整优化器状态字典从排名 0 分散到所有其他排名。

返回每个排名上的分片优化器状态字典。返回值与 shard_full_optim_state_dict() 相同，并且在排名 0 上，第一个参数应为 full_optim_state_dict() 的返回值。

示例

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> model, optim = ...
>>> full_osd = FSDP.full_optim_state_dict(model, optim)  # only non-empty on rank 0
>>> # Define new model with possibly different world size
>>> new_model, new_optim, new_group = ...
>>> sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, new_model, group=new_group)
>>> new_optim.load_state_dict(sharded_osd)

Note

无论是 shard_full_optim_state_dict() 还是 scatter_full_optim_state_dict() 都可以用来获取要加载的分片优化器状态字典。假设完整的优化器状态字典位于 CPU 内存中，前者要求每个 rank 都在 CPU 内存中拥有完整的字典，其中每个 rank 独立地对字典进行分片，无需任何通信，而后者仅要求 rank 0 在 CPU 内存中拥有完整的字典，其中 rank 0 将每个分片移动到 GPU 内存中（对于 NCCL），并将其适当地通信给其他 rank。因此，前者的 CPU 内存总成本更高，而后者的通信成本更高。

参数

• full_optim_state_dict (Optional[Dict[str, Any]]) – 与未扁平化参数相对应的优化器状态字典，并在 rank 0 上保存完整的非分片优化器状态；该参数在非零 rank 上被忽略。

• model (torch.nn.Module) – 根模块（可能是也可能不是 FullyShardedDataParallel 实例），其参数与 full_optim_state_dict 中的优化器状态对应。

• optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递给优化器的输入，表示参数组的 list 或参数的迭代器；如果为 None，则此方法假设输入为 model.parameters()。此参数已弃用，不再需要传递它。（默认值：None)

• optim (Optional[torch.optim.Optimizer]) – 将加载此方法返回的状态字典的优化器。这是比 optim_input 更推荐使用的参数。（默认值：None)

• group (dist.ProcessGroup) – 模型的进程组，如果使用默认进程组，则为 None。（默认值：None)

返回值

完整的优化器状态字典现在被重新映射到扁平化的参数而不是未扁平化的参数，并且仅限于包含此 rank 的优化器状态的一部分。

返回类型

Dict[str, Any]

static set_state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None)

设置目标模块所有后代 FSDP 模块的 state_dict_type。

还为模型和优化器的状态字典提供（可选的）配置。目标模块不必是 FSDP 模块。如果目标模块是 FSDP 模块，它的 state_dict_type 也会被更改。

Note

此 API 应该只针对顶级（根）模块调用。

Note

此 API 使用户能够在根 FSDP 模块被另一个 nn.Module 包装的情况下透明地使用传统的 state_dict API 来获取模型检查点。例如，以下将确保 state_dict 在所有非 FSDP 实例上被调用，同时将 FSDP 的 sharded_state_dict 实现分派到其中。

示例

>>> model = DDP(FSDP(...))
>>> FSDP.set_state_dict_type(
>>>     model,
>>>     StateDictType.SHARDED_STATE_DICT,
>>>     state_dict_config = ShardedStateDictConfig(offload_to_cpu=True),
>>>     optim_state_dict_config = OptimStateDictConfig(offload_to_cpu=True),
>>> )
>>> param_state_dict = model.state_dict()
>>> optim_state_dict = FSDP.optim_state_dict(model, optim)

参数

• module (torch.nn.Module) – 根模块。

• state_dict_type (StateDictType) – 要设置的所需 state_dict_type。

• state_dict_config (Optional[StateDictConfig]) – 目标 state_dict_type 的配置。

• optim_state_dict_config (Optional[OptimStateDictConfig]) – 优化器状态字典的配置。

返回值

包含模块的先前状态字典类型和配置的 StateDictSettings。

返回类型

StateDictSettings

static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None)

分片完整的优化器状态字典。

将 full_optim_state_dict 中的状态重新映射到扁平化的参数而不是未扁平化的参数，并将状态限制为仅包含此 rank 的优化器状态的一部分。第一个参数应该是 full_optim_state_dict() 的返回值。

示例

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> model, optim = ...
>>> full_osd = FSDP.full_optim_state_dict(model, optim)
>>> torch.save(full_osd, PATH)
>>> # Define new model with possibly different world size
>>> new_model, new_optim = ...
>>> full_osd = torch.load(PATH)
>>> sharded_osd = FSDP.shard_full_optim_state_dict(full_osd, new_model)
>>> new_optim.load_state_dict(sharded_osd)

Note

无论是 shard_full_optim_state_dict() 还是 scatter_full_optim_state_dict() 都可以用来获取要加载的分片优化器状态字典。假设完整的优化器状态字典位于 CPU 内存中，前者要求每个 rank 都在 CPU 内存中拥有完整的字典，其中每个 rank 独立地对字典进行分片，无需任何通信，而后者仅要求 rank 0 在 CPU 内存中拥有完整的字典，其中 rank 0 将每个分片移动到 GPU 内存中（对于 NCCL），并将其适当地通信给其他 rank。因此，前者的 CPU 内存总成本更高，而后者的通信成本更高。

参数

• full_optim_state_dict (Dict[str, Any]) – 与未扁平化参数相对应的优化器状态字典，保存完整的非分片优化器状态。

• model (torch.nn.Module) – 根模块（可能是也可能不是 FullyShardedDataParallel 实例），其参数与 full_optim_state_dict 中的优化器状态对应。

• optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递给优化器的输入，表示参数组的 list 或参数的迭代器；如果为 None，则此方法假设输入为 model.parameters()。此参数已弃用，不再需要传递它。（默认值：None)

• optim (Optional[torch.optim.Optimizer]) – 将加载此方法返回的状态字典的优化器。这是比 optim_input 更推荐使用的参数。（默认值：None)

返回值

完整的优化器状态字典现在被重新映射到扁平化的参数而不是未扁平化的参数，并且仅限于包含此 rank 的优化器状态的一部分。

返回类型

Dict[str, Any]

static sharded_optim_state_dict(model, optim, group=None)

以分片形式返回优化器状态字典。

该 API 与 full_optim_state_dict() 类似，但此 API 将所有非零维状态块分为 ShardedTensor 以节省内存。此 API 仅应在模型 state_dict 使用上下文管理器 with state_dict_type(SHARDED_STATE_DICT): 推导出时使用。

有关详细用法，请参阅 full_optim_state_dict()。

警告

返回的状态字典包含 ShardedTensor，不能直接由常规 optim.load_state_dict 使用。

返回类型

Dict[str, Any]

static state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None)

设置目标模块所有后代 FSDP 模块的 state_dict_type。

此上下文管理器具有与 set_state_dict_type() 相同的功能。请阅读 set_state_dict_type() 文档以了解详细信息。

示例

>>> model = DDP(FSDP(...))
>>> with FSDP.state_dict_type(
>>>     model,
>>>     StateDictType.SHARDED_STATE_DICT,
>>> ):
>>>     checkpoint = model.state_dict()

参数

• module (torch.nn.Module) – 根模块。

• state_dict_type (StateDictType) – 要设置的所需 state_dict_type。

• state_dict_config (Optional[StateDictConfig]) – 目标 state_dict_type 的模型 state_dict 配置。

• optim_state_dict_config (Optional[OptimStateDictConfig]) – 目标 state_dict_type 的优化器 state_dict 配置。

返回类型

Generator

static summon_full_params(module, recurse=True, writeback=True, rank0_only=False, offload_to_cpu=False, with_grads=False)

使用此上下文管理器公开 FSDP 实例的完整参数。

在模型进行前向/反向传播之后可能很有用，用于获取参数以进行额外处理或检查。它可以接受非 FSDP 模块，并将为所有包含的 FSDP 模块及其子模块（取决于 recurse 参数）调用完整参数。

Note

这可以在内部 FSDP 上使用。

Note

此方法不能在正向或反向传递中使用。正向和反向也不能从此上下文中启动。

Note

参数将在上下文管理器退出后恢复到其本地分片，存储行为与正向相同。

Note

可以修改完整参数，但只有对应于本地参数分片的那些部分将在上下文管理器退出后保留（除非writeback=False，在这种情况下更改将被丢弃）。在 FSDP 不分片参数的情况下，目前只有在world_size == 1或NO_SHARD配置的情况下，修改才会被保留，而与writeback无关。

Note

此方法适用于不是 FSDP 本身的模块，但可能包含多个独立的 FSDP 单位。在这种情况下，给定的参数将应用于所有包含的 FSDP 单位。

警告

请注意，rank0_only=True与writeback=True结合目前不支持，并将引发错误。这是因为模型参数形状在上下文中的各个等级之间将不同，写入它们会导致在退出上下文时各个等级之间不一致。

警告

请注意，offload_to_cpu和rank0_only=False将导致完整参数被冗余地复制到同一个机器上的 GPU 的 CPU 内存中，这可能会带来 CPU 内存不足的风险。建议将offload_to_cpu与rank0_only=True一起使用。

参数

• recurse (bool, Optional) – 递归地调用所有嵌套 FSDP 实例的参数（默认值：True）。

• writeback (bool, Optional) – 如果False，则在上下文管理器退出后丢弃对参数的修改；禁用此功能可以略微提高效率（默认值：True）

• rank0_only (bool, Optional) – 如果True，则完整参数只在全局等级 0 上实现。这意味着在上下文中，只有等级 0 将拥有完整参数，而其他等级将拥有分片参数。请注意，将rank0_only=True与writeback=True一起设置是不支持的，因为模型参数形状在上下文中的各个等级之间将不同，写入它们会导致在退出上下文时各个等级之间不一致。

• offload_to_cpu (bool, Optional) – 如果True，则完整参数将被卸载到 CPU。请注意，这种卸载目前只发生在参数被分片的情况下（这种情况只有在 world_size = 1 或NO_SHARD配置的情况下才不会发生）。建议将offload_to_cpu与rank0_only=True一起使用，以避免将模型参数的冗余副本卸载到同一个 CPU 内存中。

• with_grads (bool, Optional) – 如果True，则梯度也与参数一起被取消分片。目前，这只有在将use_orig_params=True传递给 FSDP 构造函数并将offload_to_cpu=False传递给此方法时才支持。（默认值：False）

返回类型

Generator

class torch.distributed.fsdp.BackwardPrefetch(value)

这配置了显式反向预取，它通过在反向传递中启用通信和计算重叠来提高吞吐量，但代价是略微增加内存使用量。

• BACKWARD_PRE: 这启用了最大的重叠，但内存使用量也最大。这在当前参数集的梯度计算之前预取下一组参数。这将重叠下一组全收集和当前梯度计算，在峰值时，它将当前参数集、下一组参数和当前梯度集都存储在内存中。

• BACKWARD_POST: 这启用了较少的重叠，但需要较少的内存使用量。这在当前参数集的梯度计算之后预取下一组参数。这将重叠当前减少散射和下一组梯度计算，并且它在为下一组参数分配内存之前释放当前参数集，在峰值时只将下一组参数和当前梯度集存储在内存中。

• FSDP 的backward_prefetch参数接受None，这完全禁用反向预取。这没有重叠，也不会增加内存使用量。一般来说，我们不推荐这种设置，因为它可能会显著降低吞吐量。

更多技术背景：对于使用 NCCL 后端的单个进程组，任何集体操作，即使是从不同的流中发出，也会争用同一个每设备 NCCL 流，这意味着集体操作发出的相对顺序对于重叠很重要。两个反向预取值对应于不同的发出顺序。

class torch.distributed.fsdp.ShardingStrategy(value)

这指定了要由FullyShardedDataParallel用于分布式训练的分片策略。

• FULL_SHARD: 参数、梯度和优化器状态被分片。对于参数，此策略在正向之前取消分片（通过全收集），在正向之后重新分片，在反向计算之前取消分片，并在反向计算之后重新分片。对于梯度，它在反向计算之后同步并分片它们（通过减少散射）。分片优化器状态在每个等级上本地更新。

• SHARD_GRAD_OP: 梯度和优化器状态在计算过程中被分片，此外，参数在计算之外被分片。对于参数，此策略在正向之前取消分片，在正向之后不重新分片，并且只在反向计算之后重新分片。分片优化器状态在每个等级上本地更新。在no_sync()内部，参数在反向计算之后不会被重新分片。

• NO_SHARD: 参数、梯度和优化器状态不会被分片，而是类似于 PyTorch 的DistributedDataParallel API 在各个等级之间被复制。对于梯度，此策略在反向计算之后同步它们（通过全归约）。未分片的优化器状态在每个等级上本地更新。

• HYBRID_SHARD: 在节点内应用FULL_SHARD，并在各个节点之间复制参数。这会导致通信量减少，因为代价高昂的全收集和减少散射只在节点内完成，对于中等规模的模型，这可能更高效。

• _HYBRID_SHARD_ZERO2: 在节点内应用SHARD_GRAD_OP，并在各个节点之间复制参数。这与HYBRID_SHARD类似，只是这可能会提供更高的吞吐量，因为未分片的参数在正向传递之后不会被释放，从而节省了反向传递之前的全收集。

class torch.distributed.fsdp.MixedPrecision(param_dtype=None, reduce_dtype=None, buffer_dtype=None, keep_low_precision_grads=False, cast_forward_inputs=False, cast_root_forward_inputs=True, _module_classes_to_ignore=(<class 'torch.nn.modules.batchnorm._BatchNorm'>, ))

这配置了 FSDP 原生的混合精度训练。

变量

• param_dtype (Optional[torch.dtype]) – 这指定了模型参数在正向和反向传递期间的数据类型，因此也是正向和反向计算的数据类型。在正向和反向传递之外，分片参数以全精度保留（例如，用于优化器步骤），并且对于模型检查点，参数始终以全精度保存。（默认值：None）

• reduce_dtype (Optional[torch.dtype]) – 这指定了梯度缩减（即减少散射或全归约）的数据类型。如果它为None，但param_dtype不为None，则它将采用param_dtype值，仍然以低精度运行梯度缩减。它允许与param_dtype不同，例如，强制梯度缩减以全精度运行。（默认值：None）

• buffer_dtype (Optional[torch.dtype]) – 这指定了缓冲区的数据类型。FSDP 不会分片缓冲区。相反，FSDP 在第一次正向传递中将它们转换为buffer_dtype，并在之后将它们保留在该数据类型中。对于模型检查点，缓冲区将以全精度保存，除了LOCAL_STATE_DICT。（默认值：None）

• keep_low_precision_grads (bool) – 如果 False，则 FSDP 在反向传播后将梯度向上转换为全精度，以备优化器步骤使用。如果 True，则 FSDP 会将梯度保留在用于梯度缩减的 dtype 中，这可以在使用支持低精度运行的自定义优化器时节省内存。 (默认值：False)

• cast_forward_inputs (bool) – 如果 True，则此 FSDP 模块将其正向参数和关键字参数转换为 param_dtype。这是为了确保参数和输入 dtype 匹配正向计算，正如许多操作所需的。当仅对一些 FSDP 模块（而不是所有模块）应用混合精度时，可能需要将其设置为 True，在这种情况下，混合精度 FSDP 子模块需要重新转换其输入。 (默认值：False)

• cast_root_forward_inputs (bool) – 如果 True，则根 FSDP 模块将其正向参数和关键字参数转换为 param_dtype，覆盖 cast_forward_inputs 的值。对于非根 FSDP 模块，此操作无效。 (默认值：True)

• _module_classes_to_ignore (Sequence[Type[torch.nn.modules.module.Module]]) – (Sequence[Type[nn.Module]]): 这指定了使用 auto_wrap_policy 时要忽略的混合精度模块类：这些类的模块将单独应用 FSDP，混合精度禁用（这意味着最终的 FSDP 构造将偏离指定的策略）。如果未指定 auto_wrap_policy，则此操作无效。此 API 处于实验阶段，可能会发生变化。 (默认值：(_BatchNorm,))

Note

此 API 处于实验阶段，可能会发生变化。

Note

仅浮点张量将转换为指定的 dtype。

Note

在 summon_full_params 中，参数被强制转换为全精度，但缓冲区不会。

Note

层归一化和批归一化即使其输入是低精度（如 float16 或 bfloat16），也会在 float32 中累积。为这些归一化模块禁用 FSDP 的混合精度仅意味着仿射参数将保留在 float32 中。但是，这会导致为这些归一化模块执行单独的全局收集和缩减散射，这可能效率低下，因此，如果工作负载允许，用户应该优先考虑继续对这些模块应用混合精度。

Note

默认情况下，如果用户传递一个包含任何 _BatchNorm 模块的模型，并指定一个 auto_wrap_policy，则批归一化模块将单独应用 FSDP，混合精度禁用。请参阅 _module_classes_to_ignore 参数。

Note

MixedPrecision 默认情况下具有 cast_root_forward_inputs=True 和 cast_forward_inputs=False。对于根 FSDP 实例，其 cast_root_forward_inputs 优先于其 cast_forward_inputs。对于非根 FSDP 实例，其 cast_root_forward_inputs 值将被忽略。默认设置足以满足典型情况，在这种情况下，每个 FSDP 实例都具有相同的 MixedPrecision 配置，并且只需要在模型正向传播的开始处将输入转换为 param_dtype。

Note

对于具有不同 MixedPrecision 配置的嵌套 FSDP 实例，建议设置各个 cast_forward_inputs 值以在每个实例的正向传播之前配置是否转换输入。在这种情况下，由于转换发生在每个 FSDP 实例的正向传播之前，因此父 FSDP 实例应该在其 FSDP 子模块之前运行其非 FSDP 子模块，以避免由于不同的 MixedPrecision 配置而导致激活 dtype 发生变化。

示例

>>> model = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3))
>>> model[1] = FSDP(
>>>     model[1],
>>>     mixed_precision=MixedPrecision(param_dtype=torch.float16, cast_forward_inputs=True),
>>> )
>>> model = FSDP(
>>>     model,
>>>     mixed_precision=MixedPrecision(param_dtype=torch.bfloat16, cast_forward_inputs=True),
>>> )

以上显示了一个工作示例。另一方面，如果将 model[1] 替换为 model[0]，这意味着使用不同 MixedPrecision 的子模块首先运行其正向传播，则 model[1] 将错误地看到 float16 激活，而不是 bfloat16 激活。

class torch.distributed.fsdp.CPUOffload(offload_params=False)

这将配置 CPU 卸载。

变量

offload_params (bool) – 这指定了在不参与计算时是否将参数卸载到 CPU。如果 True，则也会将梯度卸载到 CPU，这意味着优化器步骤在 CPU 上运行。

class torch.distributed.fsdp.StateDictConfig(offload_to_cpu=False)

StateDictConfig 是所有 state_dict 配置类的基类。用户应该实例化子类（例如 FullStateDictConfig）以配置 FSDP 支持的相应 state_dict 类型的设置。

变量

offload_to_cpu (bool) – 如果 True，则 FSDP 将 state dict 值卸载到 CPU，如果 False，则 FSDP 将其保留在 GPU 上。 (默认值：False)

class torch.distributed.fsdp.FullStateDictConfig(offload_to_cpu=False, rank0_only=False)

FullStateDictConfig 是一个旨在与 StateDictType.FULL_STATE_DICT 一起使用的配置类。建议在保存完整 state dict 时同时启用 offload_to_cpu=True 和 rank0_only=True，以分别节省 GPU 内存和 CPU 内存。此配置类旨在通过 state_dict_type() 上下文管理器使用，如下所示

>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> fsdp = FSDP(model, auto_wrap_policy=...)
>>> cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
>>> with FSDP.state_dict_type(fsdp, StateDictType.FULL_STATE_DICT, cfg):
>>>     state = fsdp.state_dict()
>>>     # `state` will be empty on non rank 0 and contain CPU tensors on rank 0.
>>> # To reload checkpoint for inference, finetuning, transfer learning, etc:
>>> model = model_fn() # Initialize model in preparation for wrapping with FSDP
>>> if dist.get_rank() == 0:
>>>     # Load checkpoint only on rank 0 to avoid memory redundancy
>>>     state_dict = torch.load("my_checkpoint.pt")
>>>     model.load_state_dict(state_dict)
>>> # All ranks initialize FSDP module as usual. `sync_module_states` argument
>>> # communicates loaded checkpoint states from rank 0 to rest of the world.
>>> fsdp = FSDP(model, device_id=torch.cuda.current_device(), auto_wrap_policy=..., sync_module_states=True)
>>> # After this point, all ranks have FSDP model with loaded checkpoint.

变量

rank0_only (bool) – 如果 True，则仅排名 0 保存完整 state dict，而排名非零则保存空字典。如果 False，则所有排名都保存完整 state dict。 (默认值：False)

class torch.distributed.fsdp.ShardedStateDictConfig(offload_to_cpu=False, _use_dtensor=False)

ShardedStateDictConfig 是一个旨在与 StateDictType.SHARDED_STATE_DICT 一起使用的配置类。

变量

_use_dtensor (bool) – 如果 True，则 FSDP 将 state dict 值保存为 DTensor，如果 False，则 FSDP 将其保存为 ShardedTensor。 (默认值：False)

警告

_use_dtensor 是 ShardedStateDictConfig 的私有字段，FSDP 使用它来确定 state dict 值的类型。用户不应手动修改 _use_dtensor。

class torch.distributed.fsdp.LocalStateDictConfig(offload_to_cpu: bool = False)

class torch.distributed.fsdp.OptimStateDictConfig(offload_to_cpu=True)

OptimStateDictConfig 是所有 optim_state_dict 配置类的基类。用户应该实例化子类（例如 FullOptimStateDictConfig）来配置 FSDP 支持的相应 optim_state_dict 类型的设置。

变量

offload_to_cpu (bool) – 如果为 True，则 FSDP 会将状态字典的张量值卸载到 CPU，如果为 False，则 FSDP 会将它们保留在原始设备上（除非启用了参数 CPU 卸载，否则为 GPU）。（默认值：True）

class torch.distributed.fsdp.FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False)

变量

rank0_only (bool) – 如果 True，则仅排名 0 保存完整 state dict，而排名非零则保存空字典。如果 False，则所有排名都保存完整 state dict。 (默认值：False)

class torch.distributed.fsdp.ShardedOptimStateDictConfig(offload_to_cpu=True, _use_dtensor=False)

ShardedOptimStateDictConfig 是一个与 StateDictType.SHARDED_STATE_DICT 一起使用的配置类。

变量

_use_dtensor (bool) – 如果 True，则 FSDP 将 state dict 值保存为 DTensor，如果 False，则 FSDP 将其保存为 ShardedTensor。 (默认值：False)

警告

_use_dtensor 是 ShardedOptimStateDictConfig 的一个私有字段，FSDP 使用它来确定状态字典值的类型。用户不应该手动修改 _use_dtensor。

class torch.distributed.fsdp.LocalOptimStateDictConfig(offload_to_cpu: bool = False)

class torch.distributed.fsdp.StateDictSettings(state_dict_type: torch.distributed.fsdp.api.StateDictType, state_dict_config: torch.distributed.fsdp.api.StateDictConfig, optim_state_dict_config: torch.distributed.fsdp.api.OptimStateDictConfig)