fsdp_checkpointing_base() — pytorch Function Reference

Architecture documentation for the fsdp_checkpointing_base() function in dist_util.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ fsdp_checkpointing_base() — pytorch Function Reference

Dependency Diagram

graph TD
  aa6a976a_820e_3084_1e1e_2603805389df["fsdp_checkpointing_base()"]
  ff4518e9_afbb_35ef_3d8f_c605ec26b945["apply_fsdp()"]
  ff4518e9_afbb_35ef_3d8f_c605ec26b945 -->|calls| aa6a976a_820e_3084_1e1e_2603805389df
  style aa6a976a_820e_3084_1e1e_2603805389df fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/dist_util.py lines 108–123

def fsdp_checkpointing_base(model, blocks):
    """apply activation checkpointing to model
    returns None as model is updated directly
    """
    non_reentrant_wrapper = functools.partial(
        checkpoint_wrapper,
        offload_to_cpu=False,
        checkpoint_impl=CheckpointImpl.NO_REENTRANT,
    )

    def check_fn(submodule):
        return isinstance(submodule, blocks)

    apply_activation_checkpointing(
        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
    )