FullyShardedDataParallel¶

类 torch.distributed.fsdp 中。FullyShardedDataParallel（module， process_group=无， sharding_strategy=无， cpu_offload=无， auto_wrap_policy=无，backward_prefetch=BackwardPrefetch.BACKWARD_PRE，mixed_precision=无，ignored_modules=无，param_init_fn=无、device_id=无、sync_module_states=False、forward_prefetch=False、limit_all_gathers=True、use_orig_params=False， ignored_states=无， device_mesh=无）[来源]¶

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

这受到了 Xu 等人以及 DeepSpeed 的 ZeRO Stage 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)
传入 constructor 参数。dev_iddevice_id

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

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

FSDP 还负责将输入张量移动到 forward 方法添加到 GPU 计算设备，因此您无需手动移动它们从 CPU。

对于，公开未分片的 parameters 而不是 forward 之后的分片参数，这与 .如果需要帮助，要检查梯度，可以将方法与 .use_orig_params=TrueShardingStrategy.SHARD_GRAD_OPShardingStrategy.FULL_SHARDsummon_full_paramswith_grads=True

使用，您可能会在 FSDP 中看到一个缺口 pre-forward，其中 CPU 线程没有发出任何内核。这是 intentional 并显示有效的速率限制器。同步 CPU thread 可以防止为后续 all-gathers，它实际上不应该延迟 GPU 内核的执行。limit_all_gathers=True

FSDP 在与 autograd 相关的 autograd 的正向和后向计算期间将托管模块的参数替换为视图原因。如果你的模块的 forward 依赖于对参数，而不是重新获取每个引用 iteration 中，则不会看到 FSDP 新创建的视图并且 autograd 将无法正常工作。torch.Tensor

最后，当分片进程组在节点内且复制进程组是节点间，设置有助于提高 all-reduce 时间某些集群设置的复制进程组。sharding_strategy=ShardingStrategy.HYBRID_SHARDNCCL_CROSS_NIC=1

局限性

使用 FSDP 时需要注意几个限制：

FSDP 目前不支持在使用 CPU 卸载时在外部进行梯度累积。这是因为 FSDP 使用新降低的渐变，而不是与任何现有梯度，这可能会导致不正确的结果。no_sync()
FSDP 不支持运行子模块的 forward pass 包含在 FSDP 实例中。这是因为子模块的参数将被分片，但子模块本身不是 FSDP 实例，因此其 forward pass 不会 all-gather full 参数。
FSDP 不能与 double backwards 一起使用，因为它的方式注册向后钩子。
FSDP 在冻结参数时有一些限制。对于，每个 FSDP 实例必须管理全部冻结或全部未冻结的参数。对于，FSDP 支持混合 frozen 和 non-frozen 参数，但建议避免这样做，以防止梯度内存使用率高于预期。use_orig_params=Falseuse_orig_params=True
从 PyTorch 1.12 开始，FSDP 对共享参数。如果需要增强的共享参数支持您的用例，请在本期中发布。
您应该避免修改 forward 和 backward 而不使用上下文，因为修改可能不会持续存在。summon_full_params

参数

模块（nn.Module） – 这是要使用 FSDP 包装的模块。
process_group （Optional[Union[ProcessGroup， Tuple[ProcessGroup， ProcessGroup]]]） – 这是模型分片的进程组，因此用于 FSDP 的 all-gather 和 reduce-scatter 集合的集合通信。如果，则 FSDP 使用默认进程群。对于混合分片策略（如），用户可以传入一个进程组，表示要分片的组，以及 replicate。如果，则 FSDP 构造进程组，供用户分片节点内和复制节点间。（默认：NoneShardingStrategy.HYBRID_SHARDNoneNone)
sharding_strategy （Optional[ShardingStrategy]） – 这将配置分片策略，该策略可能会牺牲内存节省和通信开销。有关详细信息，请参阅。（默认：FULL_SHARD)
cpu_offload （Optional[CPUOffload]） – 这将配置 CPU 卸载。如果此设置为，则不会发生 CPU 卸载。有关详细信息，请参阅。（默认：NoneNone)
auto_wrap_policy （可选[Union[Callable[[nn.模块、bool、int]、bool]、ModuleWrapPolicy、CustomPolicy]]） –
这指定了一个策略，用于将 FSDP 应用于的子模块，这是通信和计算重叠所必需的，因此影响性能。如果，则 FSDP 仅适用于，用户应手动将 FSDP 应用于父模块他们自己（自下而上进行）。为方便起见，this 直接接受，这允许用户指定要包装的 module 类（例如 transformer 块）。否则这应该是一个可调用对象，它接受三个参数，，并且并且应该返回一个指定传入的遍历是否应应用 FSDP，或者是否应继续遍历 module 的子树 if .用户可以添加额外的参数传递给 callable。in 给出了一个 callable 示例，如果模块的子树中的参数超过 100M 数字。我们建议在应用 FSDP 后打印模型并根据需要进行调整。moduleNonemoduleModuleWrapPolicymodule: nn.Modulerecurse: boolnonwrapped_numel: intboolmodulerecurse=Falserecurse=Truesize_based_auto_wrap_policytorch.distributed.fsdp.wrap.py

例：
```
>>> 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]） – 这将配置所有收集的显式向后预取。如果，则 FSDP 不会向后预取，并且没有通信和计算在向后传递中重叠。有关详细信息，请参阅。（默认：NoneBACKWARD_PRE)
mixed_precision （Optional[MixedPrecision]） – 这将为 FSDP 配置本机混合精度。如果设置为，则不使用混合精度。否则，参数 buffer 和 gradient reduction dtypes 的 d类型。有关详细信息，请参阅。（默认：NoneNone)
ignored_modules （Optional[Iterable[torch.nn.Module]]） – 其自己的 parameters 和子模块的 parameters 和 buffer 是被此实例忽略。直接进入的模块都不应该是实例，并且任何已经构造的子模块都不会被忽略，如果它们嵌套在此实例下。此参数可用于在使用 OR 参数的分片未由 FSDP.（默认：ignored_modulesauto_wrap_policyNone)
param_init_fn （可选[Callable[[nn.模块]，无]]） –
一个指定当前在 meta 设备上的模块应如何初始化到实际设备上。从 v1.12 开始，FSDP 检测到在 meta device via 上带有参数或缓冲区的模块，如果指定则适用，否则调用对于这两种情况，实现应该只初始化模块，而不是其子模块的子模块。这是为了避免 re-initialization 重新初始化。此外，FSDP 还支持 deferred 通过 torchdistX 的（https://github.com/pytorch/torchdistX） API 进行初始化，其中延迟的模块被初始化否则调用 torchdistX 的默认值。如果是 specified，则它被应用于所有元设备模块，这意味着它可能应该在 Module 类型上出现。FSDP 调用参数扁平化和分片之前的初始化函数。Callable[torch.nn.Module] -> Noneis_metaparam_init_fnnn.Module.reset_parameters()deferred_init()param_init_fnmaterialize_module()param_init_fn

例：
```
>>> 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]]） – 一个或给出 FSDP 的 CUDA 设备进行初始化，包括模块初始化（如果需要）和参数 sharding。这应该指定为提高 CPU 上的初始化速度。如果设置了默认 CUDA 设备（例如通过），然后用户可以传递给 this。（默认：inttorch.devicemoduletorch.cuda.set_devicetorch.cuda.current_deviceNone)
sync_module_states （bool） – 如果为，则每个 FSDP 模块都将 broadcast 模块参数和缓冲区，以确保它们跨等级复制（将通信开销添加到 this 构造函数）。这有助于加载 checkpoint via 以节省内存的方式。有关此示例，请参阅。（默认：Truestate_dictload_state_dictFalse)
forward_prefetch （bool） – 如果，则 FSDP 显式预取当前 forward 之前的下一个 forward-pass all-gather 计算。这仅对 CPU 密集型工作负载有用，其中提前发出下一个 all-gather 可能会改善重叠。这应该只用于静态图模型，因为预取遵循第一次迭代的执行顺序。（默认：TrueFalse)
limit_all_gathers （bool） – 如果，则显式 FSDP 同步 CPU 线程，以确保仅从两个连续的 FSDP 实例（当前运行 computation 和下一个 all-gather 被预取的实例）。如果，则 FSDP 允许 CPU 线程发出 all-gather 无需任何额外的同步。（默认值：）我们经常将此功能称为 “Rate Limiter”。此标志应仅设置为特定的 CPU 密集型工作负载，且内存压力，在这种情况下，CPU 线程可能会主动发出所有内核，而不关心 GPU 内存使用。TrueFalseTrueFalse
use_orig_params （bool） – 将此项设置为 FSDP 将使用的原始参数。FSDP 公开了那些原始的参数传递给用户。这意味着优化器步骤在原始参数上运行，启用 per-original-parameter 超参数。FSDP 保留原始 parameter 变量并在 unsharded 之间操作其数据和分片形式，其中它们始终是底层的视图 unsharded 或 sharded 。使用 current 算法，分片形式始终为 1D，失去了原始张量结构。原始参数可以具有 all ，一些，或者不存在给定排名的数据。在 none 的情况下，其数据将类似于一个大小为 0 的空张量。用户不应根据给定数据编写程序 original 参数的分片形式。需要用。将此项设置为公开 FSDP 的 internal s 传递给用户。（默认：Truemodulenn.Module.named_parameters()FlatParameterFlatParameterTruetorch.compile()FalseFlatParameternn.Module.named_parameters()False)
ignored_states （Optional[Iterable[torch.nn.Parameter]]， Optional[Iterable[torch.nn.Module]]） – 忽略的参数或模块，这些参数或模块将不由此 FSDP 管理 instance，这意味着参数不会被分片，并且它们的渐变不会跨等级降低。此参数与 existing 参数，我们可能很快就会弃用。为了向后兼容，我们保留了 and ignored_modules'，但 FSDP 只允许一个指定为 not 。ignored_modulesignored_modulesignored_statesNone
device_mesh （Optional[DeviceMesh]） – DeviceMesh 可用作 process_group。当 device_mesh 通过时，FSDP 将使用底层进程用于 All-Gather 和 Reduce-Scatter 集体通信的组。因此这两个 args 需要互斥。对于混合分片策略，用户可以传入 2D DeviceMesh 进程组元组。对于 2D FSDP + TP，用户需要传入 device_mesh而不是 process_group。有关更多 DeviceMesh 信息，请访问： https://pytorch.org/tutorials/recipes/distributed_device_mesh.htmlShardingStrategy.HYBRID_SHARD

apply（fn）[来源]¶

递归地应用于每个 submodule （由返回）以及 self 。fn.children()

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

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

参数

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

返回

自我

返回类型

模块

check_is_root（）[来源]¶

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

返回类型: 布尔

clip_grad_norm_（max_norm， norm_type=2.0）[来源]¶

裁剪所有参数的梯度范数。

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

参数

max_norm （float or int） - 梯度的最大范数
norm_type （float or int） - 使用的 p-norm 的类型。可以是无穷大范数。'inf'

返回

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

返回类型

张肌

如果每个 FSDP 实例都使用，则表示没有梯度跨等级分片，那么您可以直接使用。NO_SHARD

如果至少有一些 FSDP 实例使用分片策略（即一个 else than ），那么你应该使用这个方法而不是因为此方法处理渐变跨等级分片的事实。NO_SHARD

返回的总范数将具有“最大”的 dtype PyTorch 的类型提升定义的所有参数/梯度语义学。例如，如果所有参数/梯度都使用低 precision dtype 的 vtype 中，则返回的 norm 的 dtype 将如此之低 precision dtype，但如果至少存在一个 parameter/ gradient 一起使用 FP32，则返回的 norm 的 dtype 将为 FP32。

警告

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

static flatten_sharded_optim_state_dict（sharded_optim_state_dict， model， optim）[来源]¶

展平分片优化器 state-dict。

API 类似于 .唯一的区别在于 input 应该是从返回。因此，将在每个等级上调用 All-gather 以收集 S。sharded_optim_state_dictShardedTensor

参数

sharded_optim_state_dict （Dict[str， Any]） – 优化器状态 dict 对应于未展平的参数，并按住 sharded optimizer 状态。
model （torch.nn.Module） – 请参阅。
optim （torch.optim.Optimizer） – 的 Optimizer 参数。model

返回

请参阅。

返回类型

Dict[str，任意]

forward（*args， **kwargs）[来源]¶

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

返回类型: 任何

static fsdp_modules（module， root_only=False）[来源]¶

返回所有嵌套的 FSDP 实例。

这可能包括自身，并且仅在 .moduleroot_only=True

参数

module （torch.nn.Module） – 根模块，可以是模块，也可以是模块。FSDP
root_only （bool） – 是否仅返回 FSDP 根模块。（默认：False)

返回

嵌套在输入 .module

返回类型

列表[FullyShardedDataParallel]

static full_optim_state_dict（model， optim， optim_input=None， rank0_only=True， group=None）[来源]¶

返回完整的优化器 state-dict。

合并排名 0 上的完整优化器状态并返回它遵循的约定，即使用键和。模块中扁平化的参数包含在中，则映射回其未拼合的参数。"state""param_groups"FSDPmodel

这需要在所有等级上调用，因为它使用集体通信。但是，如果，则状态 dict 仅在排名 0 和所有其他排名上填充返回空的 .rank0_only=True

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

与中一样，张量包含在优化器状态 dict 中，因此可能会有是别名惊喜。对于最佳实践，请考虑将立即返回 Optimizer state dict，例如使用 .torch.save()

参数

model （torch.nn.Module） – 根模块（可以是也可能不是实例），其参数传递到 Optimizer 中。optim
optim （torch.optim.Optimizer） – 的 Optimizer 参数。model
optim_input （Optional[Union[List[Dict[str， Any]]， Iterable[torch.nn.Parameter]]]） – 传入优化器的输入，表示of 参数组或参数的可迭代对象; 如果，则此方法假定输入为。此参数已弃用，并且无需再传入。（默认：optimNonemodel.parameters()None)
rank0_only （bool） – 如果，则仅保存排名 0 上填充的内容;if ，则将其保存在所有等级上。（默认：TrueFalseTrue)
组（dist.ProcessGroup） – 模型的进程组，或者如果使用默认进程组。（默认：NoneNone)

返回

A 包含的原始未拼合参数的优化器状态，并包含键 “state” 和 “param_groups” 遵循 .如果则非零排名返回空。modelrank0_only=True

返回类型

Dict[str，任意]

静态get_state_dict_type（模块）[来源]¶

获取根为的 FSDP 模块的 state_dict_type 和相应的配置。module

目标模块不必是 FSDP 模块。

返回

A 中包含 state_dict_type 和 state_dict / optim_state_dict当前设置的配置。StateDictSettings

提高

AssertionError“ 如果 StateDictSettings 为 differen –
FSDP 子模块不同。–

返回类型

StateDict设置

property module：模块¶: 返回包装的模块。

named_buffers（*args， **kwargs）[来源]¶

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

拦截缓冲区名称并删除所有出现的特定于 FSDP 的扁平化缓冲区前缀当进入 Context Manager 中时。

返回类型: Iterator[Tuple[str，张量]]

named_parameters（*args， **kwargs）[来源]¶

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

截取参数名称并删除所有出现的特定于 FSDP 的扁平化参数前缀当进入 Context Manager 中时。

返回类型: Iterator[Tuple[str，参数]]

no_sync（）[来源]¶

禁用 FSDP 实例之间的梯度同步。

在此上下文中，梯度将在 module 中累积变量，稍后将在第一个 forward-backward 传递。这应该只是在根 FSDP 实例上使用，并将递归地应用于所有子 FSDP 实例。

注意

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

注意

当与 CPU 卸载一起使用时，梯度不会在 Context Manager 中卸载到 CPU。相反，他们只会在最终同步后立即卸载。

返回类型: 发电机

static optim_state_dict（model， optim， optim_state_dict=None， group=None）[来源]¶

转换与分片模型对应的优化器的 state-dict。

给定的 state-dict 可以转换为以下三种类型之一： 1）全量优化器 state_dict，2）分片优化器 state_dict，3）本地优化器 state_dict。

对于完整的优化器state_dict，所有状态都是未展平的，而不是分片的。仅 Rank0 和 CPU 只能通过指定避免 OOM。

对于分片优化器 state_dict，所有状态都是未展平的，但已分片。只能通过 CPU 指定以进一步保存记忆。

对于本地state_dict，不会执行任何变换。但是一个 state 将从 nn.Tensor 设置为 ShardedTensor 来表示其分片 nature（尚不支持此功能）。

例：

>>> 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） – 根模块（可以是也可能不是实例），其参数传递到 Optimizer 中。optim
optim （torch.optim.Optimizer） – 的 Optimizer 参数。model
optim_state_dict （Dict[str， Any]） – state_dict的目标优化器变换。如果值为 None，则将使用 optim.state_dict（）。( 违约：None)
组（dist.ProcessGroup） – 模型的流程组，参数分片，或者使用默认进程组。( 违约：NoneNone)

返回

A 包含的优化器状态。优化器状态的分片基于。modelstate_dict_type

返回类型

Dict[str，任意]

static optim_state_dict_to_load（model， optim， optim_state_dict， is_named_optimizer=False， load_directly=False， group=无）[来源]¶

转换优化器 state-dict，以便可以将其加载到与 FSDP 模型关联的优化器中。

给定通过转换的 a，它将被转换为扁平化的优化器 state_dict 可以加载到的是的优化器。必须由 FullyShardedDataParallel 分片。optim_state_dictoptimmodelmodel

>>> 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） – 根模块（可以是也可能不是实例），其参数传递到 Optimizer 中。optim
optim （torch.optim.Optimizer） – 的 Optimizer 参数。model
optim_state_dict （Dict[str， Any]） – 要加载的优化器状态。
is_named_optimizer （bool） – 这个优化器是 NamedOptimizer 还是 KeyedOptimizer 的 API 中。仅当 TorchRec 的 KeyedOptimizer 或 torch.distributed 的 NamedOptimizer 的 NamedOptimizer 中。optim
load_directly （bool） – 如果此 API 设置为 True，则此 API 也将调用 optim.load_state_dict（result）后，再返回结果。否则，用户负责调用（默认：optim.load_state_dict()False)
组（dist.ProcessGroup） – 模型的流程组，参数分片，或者使用默认进程组。( 违约：NoneNone)

返回类型

Dict[str，任意]

register_comm_hook（state， hook）[来源]¶

注册通信钩子。

这是一项增强功能，为用户提供了一个灵活的挂钩，用户可以在其中指定 FSDP 的聚合方式跨多个工作线程的梯度。这个钩子可以用来实现多种算法，比如 GossipGrad 和梯度压缩其中涉及不同的沟通策略参数同步。

警告

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

参数

状态（对象） –
传递给 hook 以在训练过程中维护任何状态信息。示例包括梯度压缩中的误差反馈、 peer 来与 gossipGrad 中的 next 通信，等等。它由每个 worker 本地存储并由 worker 上的所有梯度张量共享。
hook （Callable） – 可调用，具有以下签名之一： 1) : 此函数接受一个 Python 张量，该张量表示相对于所有变量的完整、扁平化、未分片的梯度对应于此 FSDP 单元正在包装的模型（未被其他 FSDP 子单位包裹）。然后，它执行所有必要的处理并返回 ; 2) : 此函数接受两个 Python 张量，第一个表示相对于所有变量的完整、扁平化、未分片的梯度对应于此 FSDP 单元正在包装的模型（未被其他 FSDP 子单位包裹）。后者表示一个预先调整大小的张量，用于在减少。在这两种情况下， callable 都会执行所有必要的处理并返回。签名为 1 的可调用对象应处理 NO_SHARD 情况下的梯度通信。签名为 2 的可调用对象应处理分片案例的梯度通信。hook: Callable[torch.Tensor] -> NoneNonehook: Callable[torch.Tensor, torch.Tensor] -> NoneNone

static rekey_optim_state_dict（optim_state_dict， optim_state_key_type， model， optim_input=None， optim=None）[来源]¶

将优化器 state dict 重新生成密钥以使用 key type 。optim_state_dictoptim_state_key_type

这可用于实现来自 FSDP 的模型的优化器 state dict 之间的兼容性实例和没有的实例。

要重新键入 FSDP 完整优化器状态 dict（即 from ）以使用参数 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)

要将普通优化器 state dict 从未包装模型重新生成密钥，请将其设置为 loadable to a wrapped model 的

>>> 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)

返回: 优化器状态 dict 使用由指定的参数键。optim_state_key_type
返回类型: Dict[str，任意]

static scatter_full_optim_state_dict（full_optim_state_dict， model， optim_input=None， optim=None， group=None）[来源]¶

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

返回每个等级的分片优化器状态字典。返回值与相同，并且 on rank 0，则第一个参数应为的返回值。

例：

>>> 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)

注意

Both 和都可用于获取分片优化器状态 dict 来加载。假设 full optimizer state dict 驻留在 CPU 内存中，前者需要每个 rank 在 CPU 内存中拥有完整的 dict，其中每个 rank 单独对 dict 进行分片而不进行任何通信，而后者只需要 rank 0 即可在 CPU 内存中拥有完整的 dict，其中，排名 0 将每个分片移动到 GPU 内存（对于 NCCL），并且适当地将其传达给 Rank。因此，前者具有更高的总 CPU 内存成本，而后者具有更高的通信成本。

参数

full_optim_state_dict （Optional[Dict[str， Any]]） – 优化器状态 dict 对应的未扁平化参数，并按住如果处于 rank 0 上，则为完整的非分片优化器状态;参数在非零等级上被忽略。
model （torch.nn.Module） – 根模块（可以是也可能不是实例），其参数对应于中的优化器状态。full_optim_state_dict
optim_input （Optional[Union[List[Dict[str， Any]]， Iterable[torch.nn.Parameter]]]） – 传入优化器的输入，表示of 参数组或参数的可迭代对象; 如果，则此方法假定输入为。此参数已弃用，并且无需再传入。（默认：Nonemodel.parameters()None)
optim （Optional[torch.optim.Optimizer]） – 将加载的优化器此方法返回的 state dict。这是首选参数以用于。（默认：optim_inputNone)
组（dist.ProcessGroup） – 模型的进程组，或者如果使用默认进程组。（默认：NoneNone)

返回

完整的优化器状态 dict 现在重新映射到展平参数而不是未展平参数，以及 restricted 以仅包含此 rank 的 Optimizer 状态部分。

返回类型

Dict[str，任意]

static set_state_dict_type（module， state_dict_type， state_dict_config=None， optim_state_dict_config=None）[来源]¶

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

还需要对模型和优化器的状态 dict 进行（可选）配置。目标模块不必是 FSDP 模块。如果目标 module 是 FSDP 模块，它也会被更改。state_dict_type

注意

此 API 应仅针对顶级（root）调用模块。

注意

此 API 使用户能够透明地使用传统 API 来获取模型检查点，在这种情况下，根 FSDP 模块由另一个 .例如以下内容将确保在所有非 FSDP 上调用实例，同时分派到sharded_state_dict实现对于 FSDP：state_dictnn.Modulestate_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]） – 配置对于优化器状态 dict。

返回

一个 StateDictSettings，其中包含前面的 state_dict 类型和模块的配置。

返回类型

StateDict设置

static shard_full_optim_state_dict（full_optim_state_dict， model， optim_input=None， optim=None）[来源]¶

分片一个完整的优化器 state-dict。

将状态重新映射到展平参数，而不是未展平参数，并限制为仅 Optimizer 状态的此 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)

注意

Both 和都可用于获取分片优化器状态 dict 来加载。假设 full optimizer state dict 驻留在 CPU 内存中，前者需要每个 rank 在 CPU 内存中拥有完整的 dict，其中每个 rank 单独对 dict 进行分片而不进行任何通信，而后者只需要 rank 0 即可在 CPU 内存中拥有完整的 dict，其中，排名 0 将每个分片移动到 GPU 内存（对于 NCCL），并且适当地将其传达给 Rank。因此，前者具有更高的总 CPU 内存成本，而后者具有更高的通信成本。

参数

full_optim_state_dict （Dict[str， Any]） – 优化器状态 dict 对应于未展平的参数，并按住 full non-sharded optimizer 状态。
model （torch.nn.Module） – 根模块（可以是也可能不是实例），其参数对应于中的优化器状态。full_optim_state_dict
optim_input （Optional[Union[List[Dict[str， Any]]， Iterable[torch.nn.Parameter]]]） – 传入优化器的输入，表示of 参数组或参数的可迭代对象; 如果，则此方法假定输入为。此参数已弃用，并且无需再传入。（默认：Nonemodel.parameters()None)
optim （Optional[torch.optim.Optimizer]） – 将加载的优化器此方法返回的 state dict。这是首选参数以用于。（默认：optim_inputNone)

返回

完整的优化器状态 dict 现在重新映射到展平参数而不是未展平参数，以及 restricted 以仅包含此 rank 的 Optimizer 状态部分。

返回类型

Dict[str，任意]

static sharded_optim_state_dict（model， optim， group=None）[来源]¶

以分片形式返回优化器 state-dict。

API 类似于此 API 块所有非零维状态以节省内存。此 API 应仅在派生模型时使用与 Context Manager 一起使用。ShardedTensorstate_dictwith state_dict_type(SHARDED_STATE_DICT):

有关详细使用方法，请参阅。

警告

返回的 state dict 包含和不能被常规直接使用。ShardedTensoroptim.load_state_dict

返回类型: Dict[str，任意]

static state_dict_type（module， state_dict_type， state_dict_config=None， optim_state_dict_config=None）[来源]¶

设置目标模块的所有后代 FSDP 模块的。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_dictstate_dict_type
optim_state_dict_config （Optional[OptimStateDictConfig]） – 目标的优化器配置。state_dictstate_dict_type

返回类型

发电机

static summon_full_params（module， recurse=True， writeback=True， rank0_only=False， offload_to_cpu=False， with_grads=False）[来源]¶

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

在前进/后退后，模型可以得到用于其他处理或检查的参数。它可以采用非 FSDP 模块，并将为所有包含的 FSDP 模块调用完整的参数作为以及他们的孩子，这取决于争论。recurse

注意

这可用于内部 FSDP。

注意

这不能在向前或向后传递中使用。也不可以从此上下文中启动 forward 和 backward。

注意

参数将在上下文之后恢复为其本地分片 manager 退出时，存储行为与 forward 相同。

注意

可以修改 full 参数，但只能修改 portion 对应的本地参数分片将在上下文管理器退出（除非，在这种情况下更改将被丢弃）。在 FSDP 不分片的情况下参数（当前仅在、或 config 时）保留修改，而不管 .writeback=Falseworld_size == 1NO_SHARDwriteback

注意

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

警告

请注意，目前不支持 with 结合使用，并且会引发错误。这是因为模型参数形状会有所不同在上下文中跨等级，并写入它们可能会导致退出上下文时等级之间的不一致。rank0_only=Truewriteback=True

警告

请注意，和 will 导致完整参数被冗余复制到 CPU 内存 GPU 位于同一台计算机上，这可能会产生 CPU OOM 的 OOM 中。建议与一起使用。offload_to_cpurank0_only=Falseoffload_to_cpurank0_only=True

参数

recurse （bool，可选） – 递归调用嵌套的所有参数 FSDP 实例（默认值：True）。
writeback （bool， Optional） – 如果，对参数的修改是在上下文管理器退出后丢弃; 禁用此选项可能会稍微更有效（默认值：True）False
rank0_only （bool，可选） – 如果，则完整参数为仅在全局排名 0 上实现。这意味着，在 context，只有排名 0 才会有完整的参数，而其他 ranks 将具有分片参数。请注意，不支持 with 的设置，因为模型参数形状会因等级而异在上下文中，写入它们可能会导致退出上下文时等级之间的不一致。Truerank0_only=Truewriteback=True
offload_to_cpu （bool， Optional） – 如果，则完整参数为卸载到 CPU。请注意，此卸载目前仅如果参数是分片的（但事实并非如此，则会出现对于 world_size = 1 或 config）。推荐与 to use 搭配 to avoid 模型参数的冗余副本被卸载到相同的 CPU 内存。TrueNO_SHARDoffload_to_cpurank0_only=True
with_grads （bool， Optional） – 如果，则梯度也是 unsharded 替换为参数。目前，这只是传递到 FSDP 时受支持 constructor 和此方法。（默认：Trueuse_orig_params=Trueoffload_to_cpu=FalseFalse)

返回类型

发电机

类 torch.distributed.fsdp 中。BackwardPrefetch（value）[来源]¶

这将配置显式向后预取，从而通过以下方式提高吞吐量在内存使用量略有增加的成本。

BACKWARD_PRE：这将启用最多的重叠，但会增加内存使用量最大。这将预取当前参数集的梯度计算。这与下一个 all-gather 和当前梯度计算，在峰值时，它保存当前参数集、下一组参数和当前内存中的梯度集。
BACKWARD_POST：这样可以减少重叠，但需要的内存更少用法。这将预取当前的 set of parameters 的梯度计算。这与当前的 reduce-scatter 和下一个梯度计算，它会释放当前参数集，然后再为下一组 parameters 的 S 参数，仅包含下一组参数和当前一组峰值时内存中的梯度。
FSDP 的参数接受，这会禁用完全向后预取。这没有重叠，也不会增加内存使用率。通常，我们不建议使用此设置，因为它可能会显著降低吞吐量。backward_prefetchNone

有关更多技术上下文：对于使用 NCCL 后端的单个进程组，任何 collective，即使它们来自不同的流，也会争夺相同的每个设备 NCCL 流，这意味着集体被发给重叠事项。二倒预取值对应于不同的问题顺序。

类 torch.distributed.fsdp 中。ShardingStrategy（value）[来源]¶

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

FULL_SHARD：参数、梯度和优化器状态被分片。对于参数，此策略在 forward，在 forward 之后重新分片，在 backward 之前取消分片计算，并在向后计算后重新分片。对于渐变，它会在 backward 计算。分片优化器状态按排名在本地更新。
SHARD_GRAD_OP：梯度和优化器状态在计算，此外，参数在外部分片计算。对于参数，此策略在 forward 的 Forward 函数，不会在 forward 之后重新分片它们，而只对它们进行重新分片在反向计算之后。分片优化器状态已更新本地每个等级。在 Inside 中，参数不会重新分片在反向计算之后。no_sync()
NO_SHARD：参数、梯度和优化器状态未分片而是跨等级复制，类似于 PyTorch 的 API。对于梯度，此策略在反向计算后同步它们（通过 all-reduce）。这未分片的优化器状态按排名在本地更新。DistributedDataParallel
HYBRID_SHARD：在节点内应用，并跨节点复制参数节点。这会导致通信量减少，因为昂贵的 all-gathers 和 reduce-scatters 仅在节点内完成，这对于 medium 来说可能性能更高 -大小的模型。FULL_SHARD
_HYBRID_SHARD_ZERO2：在节点内应用，并跨节点复制参数节点。这与类似，不同之处在于这可能会提供更高的吞吐量由于未分片的参数在正向传递后不会释放，因此将 all-gathers 在 pre-backward 中。SHARD_GRAD_OPHYBRID_SHARD

类 torch.distributed.fsdp 中。MixedPrecision（param_dtype=无， reduce_dtype=无， buffer_dtype=无， keep_low_precision_grads=False， cast_forward_inputs=False， cast_root_forward_inputs=True， _module_classes_to_ignore=（<类 'torch.nn.modules.batchnorm._BatchNorm'>，）））[来源]¶

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

变量

param_dtype （Optional[torch.dtype]） – 指定模型的 dtype 参数，因此前向和后向计算。在 forward 和 backward 之外，分片参数保持完全精确（例如，对于 optimizer 步骤），对于模型检查点，参数为始终以完全精度保存。（默认：None)
reduce_dtype （Optional[torch.dtype]） – 指定梯度减少（即 Reduce-Scatter 或 All-Reduce）。如果 this is but is not ，则 this 将采用该值，仍在低精度。这允许与不同，例如强制梯度缩减以全精度运行。（默认：Noneparam_dtypeNoneparam_dtypeparam_dtypeNone)
buffer_dtype （Optional[torch.dtype]） – 指定缓冲区。FSDP 不分片缓冲区。相反，FSDP 在第一次向前传球中将他们投射到并保持在那个 dtype 的 ID 调用。对于模型检查点，缓冲区被保存除外，完全精确。（默认：buffer_dtypeLOCAL_STATE_DICTNone)
keep_low_precision_grads （bool） – 如果，则 FSDP 向上转换在准备中向后传递后将梯度调整为全精度对于优化器步骤。如果，则 FSDP 保留梯度在用于梯度减少的 dtype 中，如果使用支持以低精度运行的自定义优化器。（默认：FalseTrueFalse)
cast_forward_inputs （bool） – 如果，则此 FSDP 模块将强制转换其正向 args 和 kwargs 为。这是为了确保该参数和输入 dtypes 与前向计算匹配，因为许多 Ops 都需要。这可能需要设置为仅当将混合精度应用于某些但不是全部 FSDP 模块，其中如果一个混合精度的 FSDP 子模块需要重新转换它的 inputs。（默认：Trueparam_dtypeTrueFalse)
cast_root_forward_inputs （bool） – 如果，则根 FSDP 模块将其前向 args 和 kwargs 转换为，覆盖的值。对于非根 FSDP 模块，这不会执行任何操作。（默认：Trueparam_dtypecast_forward_inputsTrue)
_module_classes_to_ignore （Sequence[Type[torch.nn.modules.module.Module]]） – （Sequence[Type[nn.Module]]）：此指定使用时要忽略混合精度的模块类：这些类的模块将具有 FSDP 在禁用混合精度的情况下单独应用于它们（意味着最终的 FSDP 结构将偏离指定的策略）。如果未指定，则执行什么都不做。此 API 是实验性的，可能会发生更改。（默认：auto_wrap_policyauto_wrap_policy(_BatchNorm,))

注意

此 API 是实验性的，可能会发生更改。

注意

只有浮点张量才会强制转换为其指定的 dtype。

注意

在中，参数被强制为 full 精度，但缓冲区不是。summon_full_params

注意

层范数和批次范数累积，即使它们的输入精度较低，如或。为这些 norm 模块禁用 FSDP 的混合精度仅意味着仿射参数保存在中。但是，这会导致为那些 norm 模块分隔 all-gather 和 reduce-scatters，这些可能效率低下，因此如果工作负载允许，用户应该更喜欢仍将混合精度应用于这些模块。float32float16bfloat16float32

注意

默认情况下，如果用户传递包含任何模块的模型并指定，则批处理规范模块将分别应用 FSDP，精度混合禁用。请参阅参数。_BatchNormauto_wrap_policy_module_classes_to_ignore

注意

MixedPrecisionhas 和 default 的。对于根 FSDP 实例，它的优先级高于它的。对于非根 FSDP 实例，将忽略其值。默认设置为对于每个 FSDP 实例具有相同的 configuration 并且只需要在模型的 forward pass 开始时将 input 转换为的典型情况来说就足够了。cast_root_forward_inputs=Truecast_forward_inputs=Falsecast_root_forward_inputscast_forward_inputscast_root_forward_inputsMixedPrecisionparam_dtype

注意

对于具有不同配置的嵌套 FSDP 实例，我们建议在每个实例的向前。在这种情况下，由于强制转换发生在每个 FSDP 之前实例的转发，父 FSDP 实例应具有其非 FSDP 子模块在其 FSDP 子模块之前运行，以避免激活 dtype 由于配置不同而更改。MixedPrecisioncast_forward_inputsMixedPrecision

例：

>>> 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),
>>> )

上面显示了一个工作示例。另一方面，如果被替换为，这意味着使用 different 首先向前运行，然后会错误地看到 activations 而不是 activations。model[1]model[0]MixedPrecisionmodel[1]float16bfloat16

类 torch.distributed.fsdp 中。CPUOffload（offload_params=False）[来源]¶

这将配置 CPU 卸载。

变量: offload_params （bool） – 这指定是否将参数卸载到 CPU 时。如果，则此也将梯度卸载到 CPU，这意味着优化器步骤在 CPU 上运行。True

类 torch.distributed.fsdp 中。StateDictConfig（offload_to_cpu=False）[来源]¶

StateDictConfig是所有配置的基类类。用户应实例化子类（例如），以便为 FSDP 支持的相应类型。state_dictFullStateDictConfigstate_dict

变量: offload_to_cpu （bool） – 如果，则 FSDP 卸载状态 dict 值设置为 CPU，如果为，则 FSDP 会将它们保留在 GPU 上。（默认：TrueFalseFalse)

类 torch.distributed.fsdp 中。FullStateDictConfig（offload_to_cpu=False， rank0_only=False）[来源]¶

FullStateDictConfig是一个 config 类，旨在与一起使用。我们建议在保存完整状态时同时启用和 dicts 分别节省 GPU 内存和 CPU 内存。这个配置类旨在通过上下文管理器作为遵循：StateDictType.FULL_STATE_DICToffload_to_cpu=Truerank0_only=Truestate_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） – 如果，则只有排名 0 保存完整状态 dict 和非零 ranks 保存一个空 dict。如果，则所有 ranks 保存完整的 state dict。（默认：TrueFalseFalse)

类 torch.distributed.fsdp 中。ShardedStateDictConfig（offload_to_cpu=False， _use_dtensor=False）[来源]¶

ShardedStateDictConfig是一个 config 类，旨在与一起使用。StateDictType.SHARDED_STATE_DICT

变量: _use_dtensor （bool） – 如果，则 FSDP 保存状态 dict 值作为，如果，则 FSDP 将它们保存为。（默认：TrueDTensorFalseShardedTensorFalse)

警告

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

类 torch.distributed.fsdp 中。LocalStateDictConfig（offload_to_cpu： bool = False）[来源]¶

类 torch.distributed.fsdp 中。OptimStateDictConfig（offload_to_cpu=True）[来源]¶

OptimStateDictConfig是所有配置类的基类。用户应实例化子类（例如），以便为 FSDP 支持的相应类型。optim_state_dictFullOptimStateDictConfigoptim_state_dict

变量: offload_to_cpu （bool） – 如果，则 FSDP 会卸载状态 dict 的 tensor 值添加到 CPU，如果为，则 FSDP 会将它们保留在原始设备（即 GPU，除非参数 CPU 卸载为 enabled）。（默认：TrueFalseTrue)

类 torch.distributed.fsdp 中。FullOptimStateDictConfig（offload_to_cpu=True， rank0_only=False）[来源]¶

变量: rank0_only （bool） – 如果，则只有排名 0 保存完整状态 dict 和非零 ranks 保存一个空 dict。如果，则所有 ranks 保存完整的 state dict。（默认：TrueFalseFalse)

类 torch.distributed.fsdp 中。ShardedOptimStateDictConfig（offload_to_cpu=True， _use_dtensor=False）[来源]¶

ShardedOptimStateDictConfig是一个 config 类，旨在与一起使用。StateDictType.SHARDED_STATE_DICT

变量: _use_dtensor （bool） – 如果，则 FSDP 保存状态 dict 值作为，如果，则 FSDP 将它们保存为。（默认：TrueDTensorFalseShardedTensorFalse)

警告

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

类 torch.distributed.fsdp 中。LocalOptimStateDictConfig（offload_to_cpu： bool = False）[来源]¶

类 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）[来源]¶

FullyShardedDataParallel¶

文档

教程

资源