Shortcuts

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.

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources