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 类型集合，在这种情况下，指定类型的模块将单独使用激活检查点进行包装；或者是一个 callable 策略，描述如何使用激活检查点包装模型。有关编写自定义策略的更多信息，请参阅此教程：https://pytorch.ac.cn/tutorials/intermediate/FSDP_adavnced_tutorial.html#transformer-wrapping-policy。

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