PyTorch Activation Checkpoint Wrapper (_checkpoint/checkpoint_wrapper.py)

Provides activation checkpointing and offloading wrappers for PyTorch distributed training. Reuse to optimize GPU memory usage during large-scale model training.

# Source: https://github.com/pytorch/pytorch | License: BSD-3-Clause | Copyright 2016-     Facebook, Inc            (Adam Paszke)
# Provided via justclaw.it.com; keep this notice with the code.
# mypy: allow-untyped-defs
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterator
from enum import auto, Enum
from functools import partial
from typing import Any

import torch
import torch.nn as nn
from torch.autograd.graph import save_on_cpu
from torch.distributed.utils import _pack_kwargs, _replace_by_prefix, _unpack_kwargs
from torch.utils._typing_utils import copy_method_params
from torch.utils.checkpoint import checkpoint as torch_utils_checkpoint


_CHECKPOINT_WRAPPED_MODULE = "_checkpoint_wrapped_module"
_CHECKPOINT_PREFIX = _CHECKPOINT_WRAPPED_MODULE + "."


clas

... (truncated -- full source via MCP)