FullyShardedDataParallel

class （module， process_group=无， cpu_offload=无， fsdp_auto_wrap_policy=无， backward_prefetch=无）torch.distributed.fsdp.FullyShardedDataParallel

用于跨数据并行工作程序对 Module 参数进行分片的包装器。这 的灵感来自 Xu 等人以及 DeepSpeed 的 ZeRO Stage 3。 FullyShardedDataParallel 通常简称为 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 将就地分片参数，这将破坏任何 以前初始化的优化器。

参数

• 模块 （nn.Module） – 要用 FSDP 包装的模块。

• process_group （Optional[ProcessGroup]） – 用于分片的进程组

• cpu_offload （可选 [CPUOffload]） – CPU 卸载配置。目前只有 parameter 和 gradient CPU 支持卸载。可以通过传入来启用它。请注意，此 当前隐式启用梯度卸载到 CPU，以便 params 和 grads 位于同一设备上才能与 Optimizer 一起使用。这 API 可能会发生更改。默认是在这种情况下有 将不卸载。cpu_offload=CPUOffload(offload_params=True)None

• fsdp_auto_wrap_policy –

  （可选 [callable]）： 一个可调用对象，指定一个策略以 FSDP 递归方式包装层。 请注意，此策略目前仅适用于 传入的模块。其余模块始终包装在 返回的 FSDP 根实例。 写入 is Callable 的一个示例，此策略包装层 参数大小大于 100M。用户可以提供应接受以下参数的自定义可调用对象：、、、 额外的自定义参数也可以添加到自定义的可调用对象中。default_auto_wrap_policytorch.distributed.fsdp.wrapfsdp_auto_wrap_policyfsdp_auto_wrap_policymodule: nn.Modulerecurse: boolunwrapped_params: intfsdp_auto_wrap_policy

  例：

  >>> def custom_auto_wrap_policy(
  >>>     module: nn.Module,
  >>>     recurse: bool,
  >>>     unwrapped_params: int,
  >>>     # These are customizable for this policy function.
  >>>     min_num_params: int = int(1e8),
  >>> ) -> bool:
  >>>     return unwrapped_params >= min_num_params
  

• backward_prefetch -（可选[BackwardPrefetch]）： 这是一项实验性功能，可能会在 不久的将来。它允许用户启用两种不同的backward_prefetch 算法来帮助反向通信和计算重叠。 每种算法的优缺点在 类 中进行了解释。BackwardPrefetch

财产 module

使 model.module 可访问，就像 DDP 一样。