set_activation_checkpointing

torchtune.training.set_activation_checkpointing(model: Module, auto_wrap_policy: Union[Set[Type], Callable[[Module, bool, int], bool]], **kwargs) → None

将激活检查点应用于传入模型的实用程序。

参数：

• model (nn.Module) – 要应用激活检查点的模型。

• auto_wrap_policy (ACWrapPolicyType) – 包装模块的策略。这可以是 nn.Module 类型的集合，在这种情况下，指定类型(s)的模块将分别使用激活检查点进行包装，或者是一个 callable 策略，描述如何使用激活检查点包装模型。有关编写自定义策略的更多信息，请参阅本教程：https://pytorch.ac.cn/tutorials/intermediate/FSDP_adavnced_tutorial.html#transformer-wrapping-policy。

• **kwargs – 传递给 torch.distributed 激活检查点的其他参数。