apply_selective_activation_checkpointing¶
- torchtune.training.apply_selective_activation_checkpointing(model: Module, ac_mode: str, ac_option: Optional[Union[int, str]]) None[source]¶
Utility to setup activation checkpointing and wrap the model for checkpointing.
- Parameters:
model (nn.Module) – Model to setup activation checkpointing.
ac_mode (str) – Activation checkpointing mode. [‘none’, ‘full’, ‘selective’]
ac_option (Optional[Union[int, str]]) – Activation checkpointing option. If ac_mode is “selective”, ac_option can be an integer or a string representing the number of layers to checkpoint. If ac_mode is “selective” and ac_option is “op”, then selective op ac is run. If ac_mode is “none” or “full”, ac_option is ignored.
