设置激活检查点

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

用于对传入模型应用激活检查点的实用工具。

Parameters:

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

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

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