Function _checkpoint_without_reentrant

internlm/utils/checkpoint.py:173–269

(function, activation_offload=False, *args)

def _checkpoint_without_reentrant(function, activation_offload=False, *args):  # pylint: disable=W1113
    # store rng_state
    fwd_cpu_state = torch.get_rng_state()
    sync_states()
    fwd_seed_states = get_states(copy=True)
    fwd_current_mode = get_current_mode()

    # check if use autocast
    if hasattr(torch, "is_autocast_enabled"):
        has_autocast_in_fwd = torch.is_autocast_enabled()
    else:
        has_autocast_in_fwd = False

    # using WeakKeyDictionary to store all the activation the first time we call unpack
    storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
    weak_holder_list = []

    # class for weakref.ref
    class Holder:
        pass

    # return a Holder object for later unpack process
    def pack():
        res = Holder()
        weak_holder_list.append(weakref.ref(res))
        return res

    # unpack hook
    def unpack(x):
        unpack_counter = 0

        # re-compute all the activation inside the function when we first call unpack
        if len(storage) == 0:

            def inner_pack(inner):
                nonlocal unpack_counter
                unpack_counter += 1

                # If the holder went out of scope, the SavedVariable is dead and so
                # the value will never be read from the storage. Skip filling it.
                if weak_holder_list[unpack_counter - 1]() is None:
                    return

                # Use detach here to ensure we don't keep the temporary autograd
                # graph created during the second forward
                storage[weak_holder_list[unpack_counter - 1]()] = inner.detach()
                return

            def inner_unpack(packed):
                raise RuntimeError("You are calling backwards on a tensor that is never exposed. Please open an issue.")

            # restore rng state
            torch.set_rng_state(fwd_cpu_state)
            for parallel_mode, state in fwd_seed_states.items():
                set_seed_states(parallel_mode, state)
            set_mode(fwd_current_mode)

            # reload arg into device if needed