code: def backward(ctx, *args):
      if not torch.autograd._is_checkpoint_valid():
          raise RuntimeError(
              "Checkpointing is not compatible with .grad() or when an `inputs` parameter is "
              "passed to .backward(). Please use .backward() and do not pass its `inputs` argument."
          )
      # Copy the list to avoid modifying original list.
      inputs = list(ctx.inputs)
      tensor_indices = ctx.tensor_indices
  
      if ctx.activation_offload:
          tensors = ctx.tensor_inputs
      else:
          tensors = ctx.saved_tensors
  
      # store the current states
      bwd_cpu_rng_state = torch.get_rng_state()
      sync_states()
      bwd_seed_states = get_states(copy=True)
      bwd_current_mode = get_current_mode()
  
      # set the states to what it used to be
      torch.set_rng_state(ctx.fwd_cpu_rng_state)
      for parallel_mode, state in ctx.fwd_seed_states.items():
          set_seed_states(parallel_mode, state)
      set_mode(ctx.fwd_current_mode)
      if ctx.activation_offload:
          tensors = copy_to_device(tensors, ctx.device)
  
      # Fill in inputs with appropriate saved tensors.
      for i, idx in enumerate(tensor_indices):
          inputs[idx] = tensors[i]
      detached_inputs = detach_variable(tuple(inputs))
      if ctx.had_autocast_in_fwd:
          with torch.enable_grad(), autocast():
              outputs = ctx.run_function(*detached_inputs)
      else:
          with torch.enable_grad():
              outputs = ctx.run_function(*detached_inputs)
  
      if isinstance(outputs, torch.Tensor):
          outputs = (outputs,)
      # recover the rng states
      torch.set_rng_state(bwd_cpu_rng_state)
      for parallel_mode, state in bwd_seed_states.items():
          set_seed_states(parallel_mode, state)
      set_mode(bwd_current_mode)
  
      # run backward() with only tensor that requires grad
      outputs_with_grad = []
      args_with_grad = []
      for i in range(len(outputs)):
          if torch.is_tensor(outputs[i]) and outputs[i].requires_grad:
              outputs_with_grad.append(outputs[i])
              args_with_grad.append(args[i])
      if len(outputs_with_grad) == 0:
          raise RuntimeError(
              "none of output has requires_grad=True,"
              " this checkpoint() is not necessary"
          )
      torch.autograd.backward(outputs_with_grad, args_with_grad)
      grads = tuple(
          inp.grad if isinstance(inp, torch.Tensor) else None
          for inp in detached_inputs
      )
      return (None, None) + grads
  
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 2 for calling the CheckpointFunction obj's underlying member method code with example usage: obj.backward(ctx) and returns: unknown