torchtune.training.FSDPPolicyType
- torchtune.training.FSDPPolicyType
用于可作为 FSDP 包装策略的函数的数据类型。特别是,此类型表示可以接受 nn.Module、布尔标志和整数并返回布尔值来指示是否应该使用 FSDP 包装模块的函数。此类型对象可以直接传递到 PyTorch FSDP 的
auto_wrap_policy
参数以指定 FSDP 如何包装子模块。以下函数是创建和返回符合
FSDPPolicyType
契约的函数的示例。def get_fsdp_policy(module: nn.Module, modules_to_wrap: Set[Type], min_num_params: int): def my_fsdp_policy(module: nn.Module, modules_to_wrap: Set[Type], recurse: bool, min_num_params: int) -> bool: if recurse: return True # Wrap layers that are of type in ``modules_to_wrap`` and layers with more than min_num_params return isinstance(module, tuple(modules_to_wrap)) or sum(p.numel() for p in module.parameters()) > 1000 return functools.partial(my_fsdp_policy, modules_to_wrap=modules_to_wrap)
有关更多详细信息,请参阅
auto_wrap_policy
的文档,网址为 https://pytorch.ac.cn/docs/stable/fsdp.html。