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。如果 ac_mode 为 “none” 或 “full”，则忽略 ac_option。