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 激活检查点的其他参数。