FullyShardedDataParallel

class torch.distributed.fsdp.FullyShardedDataParallel(module, process_group=None, cpu_offload=None, fsdp_auto_wrap_policy=None, backward_prefetch=None)

A wrapper for sharding Module parameters across data parallel workers. This is inspired by Xu et al. as well as the ZeRO Stage 3 from DeepSpeed. FullyShardedDataParallel is commonly shorten to FSDP。

Example:

>>> 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 会原地对参数进行分片，这将破坏任何 之前初始化的优化器。

Parameters

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

• process_group (Optional[ProcessGroup]) – 分片用的进程组

• cpu_offload (Optional[CPUOffload]) – CPU卸载配置。目前仅支持参数和梯度的CPU卸载。可以通过传入 cpu_offload=CPUOffload(offload_params=True)来启用。请注意，这 目前会隐式地启用梯度卸载到CPU，以便 参数和梯度位于同一设备上以与优化器配合使用。这个 API可能会发生变化。默认值是 None，在这种情况下 不会进行任何卸载。

• fsdp_auto_wrap_policy –

  (可选 [可调用对象]): 一个可调用对象，用于指定递归地将层用FSDP包装的策略。 请注意，此策略目前只会应用于传入模块的子模块。 其余模块始终会被返回的FSDP根实例所包装。 default_auto_wrap_policy 写在 torch.distributed.fsdp.wrap 中是 一个 fsdp_auto_wrap_policy 可调用对象的例子，该策略会包装参数大小大于100M的层。用户可以提供自定义的 fsdp_auto_wrap_policy 可调用对象，该对象应接受以下参数: module: nn.Module, recurse: bool, unwrapped_params: int, 还可以将额外的自定义参数添加到自定义的 fsdp_auto_wrap_policy 可调用对象中。

  Example:

  >>> 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 – (Optional[BackwardPrefetch]): 这是一个实验性功能，未来可能会发生变化。它允许用户启用两种不同的 backward_prefetch 算法，以帮助反向传播中的通信和计算重叠。每种算法的优缺点在类 BackwardPrefetch 中进行了说明。

property module

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