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
视图,原因与自动微分相关。如果模块的前向依赖于保存的参数引用而不是每次迭代重新获取引用,那么它将看不到 FSDP 新创建的视图,并且自动微分将无法正常工作。最后,当使用
sharding_strategy=ShardingStrategy.HYBRID_SHARD
且分片进程组为节点内,复制进程组为节点间时,设置NCCL_CROSS_NIC=1
可以帮助改善某些集群设置中复制进程组的所有归约时间。限制
使用 FSDP 时需要注意一些限制
当使用 CPU 卸载时,FSDP 目前不支持在
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 (可选[Union[ProcessGroup, Tuple[ProcessGroup, ProcessGroup]]]) – 这是模型被分片的进程组,因此是 FSDP 的全局聚合和规约散射集体通信所使用的进程组。如果为
None
,则 FSDP 使用默认进程组。对于混合分片策略(例如ShardingStrategy.HYBRID_SHARD
),用户可以传入一个进程组元组,分别表示用于分片和复制的组。如果为None
,则 FSDP 为用户构建进程组以进行节点内分片和节点间复制。(默认值:None
)sharding_strategy (可选[ShardingStrategy]) – 此配置分片策略,该策略可能权衡内存节省和通信开销。有关详细信息,请参阅
ShardingStrategy
。(默认值:FULL_SHARD
)cpu_offload (可选[CPUOffload]) – 此配置 CPU 卸载。如果将其设置为
None
,则不会发生 CPU 卸载。有关详细信息,请参阅CPUOffload
。(默认值:None
)auto_wrap_policy (可选[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
给出了一个示例可调用对象,如果其子树中的参数超过 1 亿个元素,则会将 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 (可选[BackwardPrefetch]) – 此配置全局聚合的反向显式预取。如果为
None
,则 FSDP 不会进行反向预取,并且反向传递中没有通信和计算重叠。有关详细信息,请参阅BackwardPrefetch
。(默认值:BACKWARD_PRE
)mixed_precision (可选[MixedPrecision]) – 此配置 FSDP 的原生混合精度。如果将其设置为
None
,则不使用混合精度。否则,可以设置参数、缓冲区和梯度缩减数据类型。有关详细信息,请参阅MixedPrecision
。(默认值:None
)ignored_modules (可选[Iterable[torch.nn.Module]]) – 此实例忽略其自身参数和子模块的参数和缓冲区的模块。
ignored_modules
中的任何模块都不应为FullyShardedDataParallel
实例,并且如果任何已构造的FullyShardedDataParallel
实例嵌套在此实例下,则不会忽略这些实例。当使用auto_wrap_policy
或参数的分片未由 FSDP 管理时,此参数可用于避免以模块粒度分片特定参数。(默认值:None
)param_init_fn (可选[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 (可选[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 (Optional[Iterable[torch.nn.Parameter]], Optional[Iterable[torch.nn.Module]]) – 将不会由此 FSDP 实例管理的忽略的参数或模块,这意味着参数未分片并且它们的梯度不会跨秩减少。此参数与现有的
ignored_modules
参数统一,我们可能会很快弃用ignored_modules
。为了向后兼容,我们同时保留ignored_states
和 ignored_modules`,但 FSDP 仅允许其中一个被指定为非None
。device_mesh (Optional[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)[source]¶
递归地将
fn
应用于每个子模块(以及自身,如.children()
所返回的)。典型用法包括初始化模型的参数(另请参阅 torch.nn.init)。
与
torch.nn.Module.apply
相比,此版本在应用fn
之前额外收集了完整参数。它不应该从另一个summon_full_params
上下文中调用。- 参数
fn (
Module
-> None) – 要应用于每个子模块的函数- 返回值
self
- 返回类型
- clip_grad_norm_(max_norm, norm_type=2.0)[source]¶
裁剪所有参数的梯度范数。
范数是在所有参数的梯度上计算的,这些梯度被视为单个向量,并且梯度在原位修改。
- 参数
- 返回值
参数的总范数(视为单个向量)。
- 返回类型
如果每个 FSDP 实例都使用
NO_SHARD
,这意味着没有梯度跨秩分片,那么您可以直接使用torch.nn.utils.clip_grad_norm_()
。如果至少某些 FSDP 实例使用分片策略(即除
NO_SHARD
之外的任何策略),则应使用此方法而不是torch.nn.utils.clip_grad_norm_()
,因为此方法处理梯度跨秩分片的事实。返回的总范数将具有所有参数/梯度中由 PyTorch 的类型提升语义定义的“最大”dtype。例如,如果所有参数/梯度都使用低精度 dtype,则返回的范数的 dtype 将是该低精度 dtype,但如果存在至少一个使用 FP32 的参数/梯度,则返回的范数的 dtype 将是 FP32。
警告
这需要在所有秩上调用,因为它使用集体通信。
- static flatten_sharded_optim_state_dict(sharded_optim_state_dict, model, optim)[source]¶
展平分片优化器状态字典。
此 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
参数的优化器。
- 返回值
- 返回类型
- static fsdp_modules(module, root_only=False)[source]¶
返回所有嵌套的 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)[source]¶
返回完整的优化器状态字典。
在 rank 0 上整合完整的优化器状态,并将其作为
dict
返回,遵循torch.optim.Optimizer.state_dict()
的约定,即使用键"state"
和"param_groups"
。model
中包含的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) – 根模块(它可能是也可能不是
FullyShardedDataParallel
实例),其参数传递到优化器optim
中。optim (torch.optim.Optimizer) –
model
参数的优化器。optim_input (Optional[Union[List[Dict[str, Any]], Iterable[torch.nn.Parameter]]]) – 传递到优化器
optim
中的输入,表示参数组的list
或参数的可迭代对象;如果为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)[source]¶
获取以
module
为根的 FSDP 模块的 state_dict_type 和相应的配置。目标模块不必是 FSDP 模块。
- 返回值
一个
StateDictSettings
,包含当前设置的 state_dict_type 和 state_dict/optim_state_dict 配置。- 引发
AssertionError` 如果不同 FSDP 子模块的 StateDictSettings 不同。 –
FSDP 子模块的 StateDictSettings 不同。 –
- 返回类型
- named_buffers(*args, **kwargs)[source]¶
返回模块缓冲区的迭代器,同时生成缓冲区的名称和缓冲区本身。
在
summon_full_params()
上下文管理器内部拦截缓冲区名称并删除所有 FSDP 特定的扁平化缓冲区前缀。
- named_parameters(*args, **kwargs)[source]¶
返回模块参数的迭代器,同时生成参数的名称和参数本身。
在
summon_full_params()
上下文管理器内部拦截参数名称并删除所有 FSDP 特定的扁平化参数前缀。
- no_sync()[source]¶
禁用跨 FSDP 实例的梯度同步。
在此上下文中,梯度将在模块变量中累积,稍后将在退出上下文后的第一个前向-反向传递中同步。这仅应在根 FSDP 实例上使用,并将递归应用于所有子 FSDP 实例。
注意
这可能会导致更高的内存使用量,因为 FSDP 将累积完整的模型梯度(而不是梯度分片),直到最终同步。
注意
在使用 CPU 卸载时,梯度在上下文管理器内部不会卸载到 CPU。相反,它们只会在最终同步后立即卸载。
- 返回类型
- static optim_state_dict(model, optim, optim_state_dict=None, group=None)[source]¶
转换对应于分片模型的优化器的状态字典。
给定的状态字典可以转换为三种类型之一:1) 完整优化器状态字典,2) 分片优化器状态字典,3) 本地优化器状态字典。
对于完整优化器状态字典,所有状态都未展平且未分片。可以通过
state_dict_type()
指定仅限 Rank0 和仅限 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) – 根模块(它可能是也可能不是
FullyShardedDataParallel
实例),其参数传递到优化器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)[source]¶
转换优化器状态字典,以便可以将其加载到与 FSDP 模型关联的优化器中。
给定一个通过
optim_state_dict()
转换的optim_state_dict
,它会被转换为可以加载到optim
的扁平化优化器状态字典,而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) – 根模块(它可能是也可能不是
FullyShardedDataParallel
实例),其参数传递到优化器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
)
- 返回类型
- register_comm_hook(state, hook)[source]¶
注册通信钩子。
这是一个增强功能,为用户提供了灵活的钩子,他们可以在其中指定 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)[source]¶
重新为优化器状态字典
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)[source]¶
将完整的优化器状态字典从 Rank 0 分散到所有其他 Rank。
返回每个 Rank 上的分片优化器状态字典。返回值与
shard_full_optim_state_dict()
相同,在 Rank 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)
注意
可以使用
shard_full_optim_state_dict()
和scatter_full_optim_state_dict()
来获取分片优化器状态字典以加载。假设完整的优化器状态字典位于 CPU 内存中,前者要求每个进程都拥有 CPU 内存中的完整字典,每个进程独立地对字典进行分片,无需任何通信;而后者仅要求进程 0 拥有 CPU 内存中的完整字典,进程 0 将每个分片移动到 GPU 内存(对于 NCCL)并将其适当通信给其他进程。因此,前者的 CPU 内存总成本更高,而后者的通信成本更高。- 参数
full_optim_state_dict (Optional[Dict[str, Any]]) – 优化器状态字典,对应于未展平的参数,并在进程 0 上保存完整未分片的优化器状态;在非零进程上忽略此参数。
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
)
- 返回值
完整的优化器状态字典现在重新映射到展平的参数而不是未展平的参数,并且仅限于包含此进程的优化器状态的一部分。
- 返回类型
Dict[str, Any]
- static set_state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None)[source]¶
设置目标模块的所有后代 FSDP 模块的
state_dict_type
。还接收模型和优化器状态字典的(可选)配置。目标模块不必是 FSDP 模块。如果目标模块是 FSDP 模块,则其
state_dict_type
也会更改。注意
此 API 仅应针对顶级(根)模块调用。
注意
此 API 使用户能够透明地使用传统的
state_dict
API 在根 FSDP 模块被另一个nn.Module
包装的情况下获取模型检查点。例如,以下内容将确保在所有非 FSDP 实例上调用state_dict
,同时为 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]) – 优化器状态字典的配置。
- 返回值
包含模块的先前 state_dict 类型和配置的 StateDictSettings。
- 返回类型
- static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None)[source]¶
对完整的优化器状态字典进行分片。
将
full_optim_state_dict
中的状态重新映射到展平的参数而不是未展平的参数,并限制为仅包含此进程的优化器状态的一部分。第一个参数应该是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)
注意
可以使用
shard_full_optim_state_dict()
和scatter_full_optim_state_dict()
来获取分片优化器状态字典以加载。假设完整的优化器状态字典位于 CPU 内存中,前者要求每个进程都拥有 CPU 内存中的完整字典,每个进程独立地对字典进行分片,无需任何通信;而后者仅要求进程 0 拥有 CPU 内存中的完整字典,进程 0 将每个分片移动到 GPU 内存(对于 NCCL)并将其适当通信给其他进程。因此,前者的 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
)
- 返回值
完整的优化器状态字典现在重新映射到展平的参数而不是未展平的参数,并且仅限于包含此进程的优化器状态的一部分。
- 返回类型
Dict[str, Any]
- static sharded_optim_state_dict(model, optim, group=None)[source]¶
以其分片形式返回优化器状态字典。
此 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
使用。
- static state_dict_type(module, state_dict_type, state_dict_config=None, optim_state_dict_config=None)[source]¶
设置目标模块的所有后代 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
配置。
- 返回类型
- static summon_full_params(module, recurse=True, writeback=True, rank0_only=False, offload_to_cpu=False, with_grads=False)[source]¶
使用此上下文管理器公开 FSDP 实例的完整参数。
在模型的前向/后向传播之后,可以使用此方法获取参数以进行额外的处理或检查。它可以接受一个非 FSDP 模块,并根据
recurse
参数调用所有包含的 FSDP 模块及其子模块的完整参数。注意
这可以用于内部 FSDP。
注意
这**不能**在正向或反向传递过程中使用。同样,也不允许在此上下文中启动前向和后向传播。
注意
在上下文管理器退出后,参数将恢复到其本地分片,存储行为与前向传播相同。
注意
可以修改完整参数,但只有对应于本地参数分片的部分会在上下文管理器退出后保留(除非
writeback=False
,在这种情况下,更改将被丢弃)。在 FSDP 未对参数进行分片的情况下,目前仅当world_size == 1
或NO_SHARD
配置时,无论writeback
如何,修改都会被持久化。注意
此方法适用于本身不是 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, 可选) – 递归调用嵌套 FSDP 实例的所有参数(默认值:True)。
**writeback** (bool, 可选) – 如果
False
,则在上下文管理器退出后丢弃对参数的修改;禁用此选项可以稍微提高效率(默认值:True)。**rank0_only** (bool, 可选) – 如果
True
,则仅在全局秩 0 上实现完整参数。这意味着在上下文中,只有秩 0 将拥有完整参数,而其他秩将拥有分片参数。请注意,设置rank0_only=True
和writeback=True
是不受支持的,因为模型参数形状在上下文中的各个秩之间将不同,并且写入它们会导致在上下文退出时各秩之间出现不一致。**offload_to_cpu** (bool, 可选) – 如果
True
,则将完整参数卸载到 CPU。请注意,此卸载目前仅在参数被分片时发生(仅在world_size = 1
或NO_SHARD
配置时不发生)。建议将offload_to_cpu
与rank0_only=True
一起使用,以避免模型参数的冗余副本被卸载到相同的 CPU 内存中。**with_grads** (bool, 可选) – 如果
True
,则梯度也会与参数一起取消分片。目前,仅当将use_orig_params=True
传递给 FSDP 构造函数并将offload_to_cpu=False
传递给此方法时才支持此功能。(默认值:False
)
- 返回类型
- class torch.distributed.fsdp.BackwardPrefetch(value)[source]¶
这配置了显式反向预取,它通过在反向传递中启用通信和计算重叠来提高吞吐量,但代价是略微增加了内存使用量。
BACKWARD_PRE
:这能够实现最大的重叠,但也会最大程度地增加内存使用量。这会在当前参数集的梯度计算之前预取下一组参数。这会使下一个所有收集和当前梯度计算重叠,并且在峰值时,它会在内存中保存当前参数集、下一组参数以及当前梯度集。BACKWARD_POST
:这能够实现较少的重叠,但需要较少的内存使用量。这会在当前参数集的梯度计算之后预取下一组参数。这会使当前减少散射和下一个梯度计算重叠,并且它会在为下一组参数分配内存之前释放当前参数集,在峰值时仅在内存中保存下一组参数和当前梯度集。FSDP 的
backward_prefetch
参数接受None
,这完全禁用反向预取。这没有重叠,也不会增加内存使用量。通常,我们不建议使用此设置,因为它可能会显著降低吞吐量。
更多技术背景:对于使用 NCCL 后端的单个进程组,任何集体操作,即使是从不同的流中发出,也会争用同一个设备的 NCCL 流,这意味着集体操作发出的相对顺序对于重叠很重要。这两个反向预取值对应于不同的发出顺序。
- class torch.distributed.fsdp.ShardingStrategy(value)[source]¶
这指定了
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'>, ))[source]¶
这配置了 FSDP 原生的混合精度训练。
- 变量
**param_dtype** (可选[torch.dtype]) – 这指定了模型参数在正向和反向传播期间的数据类型,因此也指定了正向和反向传播计算的数据类型。在正向和反向传播之外,分片参数以全精度保留(例如,用于优化器步骤),并且对于模型检查点,参数始终以全精度保存。(默认值:
None
)**reduce_dtype** (可选[torch.dtype]) – 这指定了梯度规约(即减少散射或全部归约)的数据类型。如果它是
None
但param_dtype
不是None
,则它采用param_dtype
的值,仍然以低精度运行梯度规约。允许它与param_dtype
不同,例如,强制梯度规约以全精度运行。(默认值:None
)**buffer_dtype** (可选[torch.dtype]) – 这指定了缓冲区的数据类型。FSDP 不会对缓冲区进行分片。相反,FSDP 在第一次正向传递中将它们转换为
buffer_dtype
,并在之后一直保持这种数据类型。对于模型检查点,缓冲区以全精度保存,除了LOCAL_STATE_DICT
。(默认值:None
)keep_low_precision_grads (bool) – 如果
False
,则 FSDP 在反向传播后将梯度升级到全精度,以准备优化器步骤。如果True
,则 FSDP 保留用于梯度约简的数据类型中的梯度,如果使用支持低精度运行的自定义优化器,则可以节省内存。(默认值:False
)cast_forward_inputs (bool) – 如果
True
,则此 FSDP 模块将其前向参数和关键字参数转换为param_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,)
)
注意
此 API 处于实验阶段,可能会发生变化。
注意
仅浮点张量被转换为指定的数据类型。
注意
在
summon_full_params
中,参数强制转换为全精度,但缓冲区不会。注意
即使层归一化和批归一化的输入是低精度,例如
float16
或bfloat16
,它们也会在float32
中累积。禁用 FSDP 对这些归一化模块的混合精度仅意味着仿射参数保留在float32
中。但是,这会为这些归一化模块产生单独的全收集和归约散射,这可能效率低下,因此,如果工作负载允许,用户应更倾向于继续对这些模块应用混合精度。注意
默认情况下,如果用户传递包含任何
_BatchNorm
模块的模型并指定auto_wrap_policy
,则批归一化模块将单独应用 FSDP,并且混合精度禁用。请参阅_module_classes_to_ignore
参数。注意
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
。注意
对于具有不同
MixedPrecision
配置的嵌套 FSDP 实例,我们建议设置各个cast_forward_inputs
值以配置在每个实例的前向传递之前是否转换输入。在这种情况下,由于转换发生在每个 FSDP 实例的前向传递之前,因此父 FSDP 实例应在其 FSDP 子模块之前运行其非 FSDP 子模块,以避免由于不同的MixedPrecision
配置而导致激活数据类型发生更改。示例
>>> 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)[source]¶
这配置了 CPU 卸载。
- 变量
offload_params (bool) – 这指定了在不参与计算时是否将参数卸载到 CPU。如果
True
,则将梯度卸载到 CPU,这意味着优化器步骤在 CPU 上运行。
- class torch.distributed.fsdp.StateDictConfig(offload_to_cpu=False)[source]¶
StateDictConfig
是所有state_dict
配置类的基类。用户应实例化子类(例如FullStateDictConfig
)以配置 FSDP 支持的相应state_dict
类型的设置。- 变量
offload_to_cpu (bool) – 如果
True
,则 FSDP 将状态字典值卸载到 CPU,如果False
,则 FSDP 将其保留在 GPU 上。(默认值:False
)
- class torch.distributed.fsdp.FullStateDictConfig(offload_to_cpu=False, rank0_only=False)[source]¶
FullStateDictConfig
是一个配置类,旨在与StateDictType.FULL_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 保存完整状态字典,非零排名保存空字典。如果False
,则所有排名都保存完整状态字典。(默认值:False
)
- class torch.distributed.fsdp.ShardedStateDictConfig(offload_to_cpu=False, _use_dtensor=False)[source]¶
ShardedStateDictConfig
是一个配置类,旨在与StateDictType.SHARDED_STATE_DICT
一起使用。- 变量
_use_dtensor (bool) – 如果
True
,则 FSDP 将状态字典值保存为DTensor
,如果False
,则 FSDP 将其保存为ShardedTensor
。(默认值:False
)
警告
_use_dtensor
是ShardedStateDictConfig
的一个私有字段,FSDP 使用它来确定状态字典值的类型。用户不应手动修改_use_dtensor
。
- class torch.distributed.fsdp.OptimStateDictConfig(offload_to_cpu=True)[source]¶
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)[source]¶
- 变量
rank0_only (bool) – 如果
True
,则只有排名 0 保存完整状态字典,非零排名保存空字典。如果False
,则所有排名都保存完整状态字典。(默认值:False
)
- class torch.distributed.fsdp.ShardedOptimStateDictConfig(offload_to_cpu=True, _use_dtensor=False)[source]¶
ShardedOptimStateDictConfig
是一个配置类,旨在与StateDictType.SHARDED_STATE_DICT
一起使用。- 变量
_use_dtensor (bool) – 如果
True
,则 FSDP 将状态字典值保存为DTensor
,如果False
,则 FSDP 将其保存为ShardedTensor
。(默认值:False
)
警告
_use_dtensor
是ShardedOptimStateDictConfig
的一个私有字段,FSDP 使用它来确定状态字典值的类型。用户不应手动修改_use_dtensor
。
- 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)[source]¶