code: def _checkpoint_without_reentrant(
      function, activation_offload=False, *args
  ):
      # 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(x):
          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
              if activation_offload:
                  for arg in args:
                      if torch.is_tensor(arg):
                          arg = arg.to(device=device)
  
              # rerun forward, the inner_pack will store all the activations in storage
              if has_autocast_in_fwd:
                  with torch.enable_grad(), autocast(), torch.autograd.graph.saved_tensors_hooks(
                      inner_pack, inner_unpack
                  ):
                      _unused = function(*args)
              else:
                  with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(
                      inner_pack, inner_unpack
                  ):
                      _unused = function(*args)
  
          if x not in storage:
              raise RuntimeError(
                  "Attempt to retrieve a tensor saved by autograd multiple times without checkpoint"
                  " recomputation being triggered in between, this is not currently supported. Please"
                  " open an issue with details on your use case so that we can prioritize adding this."
              )
  
          return storage[x]
  
      # get device if we need to offload the activation
      if activation_offload:
          device = get_current_device()
  
      # run function with pack and unpack as saved_tensors_hooks
      with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
          output = function(*args)
  
          # offload activation if needed
          if activation_offload:
              for arg in args:
                  if torch.is_tensor(arg):
                      arg = arg.to(device="cpu")
  
      return output
  
file_path: ColossalAI/colossalai/legacy/utils/activation_checkpoint.py
imports: import weakref
  import torch
  from torch.utils.checkpoint import check_backward_validity, detach_variable
  from colossalai.legacy.context.random import get_current_mode, get_states, set_mode, set_seed_states, sync_states
  from colossalai.utils.device import autocast, get_current_device
label_desc: Use image node_id 3 for calling a global function with example usage: _checkpoint_without_reentrant(function, activation_offload) and returns: output, res, storage