快捷方式

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

Callable 的别名[[Module, bool, int], bool]

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源