Skip to content

activation_checkpoint

Activation checkpointing policies for the modular training framework.

The modular trainer owns these policies under fastvideo.train, which keeps model plugins within one training package.

Classes

fastvideo.train.utils.activation_checkpoint.CheckpointType

Bases: str, Enum

Supported activation checkpointing policies.

Functions:

fastvideo.train.utils.activation_checkpoint.apply_activation_checkpointing 

apply_activation_checkpointing(module: Module, checkpointing_type: str = FULL, n_layer: int = 1) -> Module

Apply the selected activation checkpointing policy to a module.

Source code in fastvideo/train/utils/activation_checkpoint.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def apply_activation_checkpointing(
    module: torch.nn.Module,
    checkpointing_type: str = CheckpointType.FULL,
    n_layer: int = 1,
) -> torch.nn.Module:
    """Apply the selected activation checkpointing policy to a module."""
    if checkpointing_type == CheckpointType.FULL:
        module = _apply_activation_checkpointing_blocks(module)
    elif checkpointing_type == CheckpointType.OPS:
        module = _apply_activation_checkpointing_ops(
            module,
            _SELECTIVE_ACTIVATION_CHECKPOINTING_OPS,
        )
    elif checkpointing_type == CheckpointType.BLOCK_SKIP:
        module = _apply_activation_checkpointing_blocks(module, n_layer)
    else:
        raise ValueError(f"Checkpointing type '{checkpointing_type}' not supported. "
                         f"Supported types are {CheckpointType.__members__.keys()}")
    return module