set_activation_checkpointing¶
- torchtune.utils.set_activation_checkpointing(model: Module, auto_wrap_policy: Optional[Set[Module]] = None, **kwargs) None[source]¶
Utility to setup activation checkpointing and wrap the model for checkpointing.
- Parameters:
model (nn.Module) – Model to setup activation checkpointing.
auto_wrap_policy (Optional[Set[nn.Module]]) – Policy to wrap module.
**kwargs – additional arguments to pass to torch.distributed activation checkpointing.
