Class ActivationRecomputation

llm_analysis/analysis.py:39–53

class ActivationRecomputation(Enum):
    NONE = 0
    """No activation recomputation; requires the most amount of memory."""

    SELECTIVE = 1
    """Selectively checkpoints and recomputes only parts of each transformer layer that
    take up a considerable amount of memory but are not computationally expensive to
    recompute, i.e. QK^T matrix multiply, softmax, softmax dropout, and attention over
    V."""

    FULL = 2
    """Full activation recomputation stores the input to EVERY transformer layer, which
    is sharded across the tensor parallel group, thus requiring an extra all-gather
    (ignored for now) per layer and add communication overhead; requires the lease
    amount of memory; requires an extra forward pass."""


@total_ordering