apply_selective_activation_checkpointing

torchtune.training.apply_selective_activation_checkpointing(model: Module, ac_mode: str, ac_option: Optional[Union[int, str]]) → None

用于设置激活检查点并为模型打包以进行检查点的实用函数。

参数:

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

• ac_mode (str) – 激活检查点模式。['none', 'full', 'selective']

• ac_option (Optional[Union[int, str]]) – 激活检查点选项。如果 ac_mode 是“selective”，ac_option 可以是一个整数或表示要检查点层数的字符串。如果 ac_mode 是“selective”且 ac_option 是“op”，则运行选择性 op 激活检查点。如果 ac_mode 是“none”或“full”，则忽略 ac_option。