class_docstr: This function is adapted from torch.utils.checkpoint with
  two main changes:
      1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`  #ignore-cuda
      2) the states in the model parallel tracker are also properly
         tracked/set/reset.
      3) Performance activation partitioning, contiguous memory optimization
      4) CPU Checkpointing
      5) Profile forward and backward functions
class_docstr_tok: ["This","function","is","adapted","from","torch.utils.checkpoint","with","two","main","changes",":","1",")","torch.cuda.set_rng_state","is","replaced","with","`","_set_cuda_rng_state","`","#","ignore-cuda","2",")","the","states","in","the","model","parallel","tracker","are","also","properly","tracked\/set\/reset",".","3",")","Performance","activation","partitioning",",","contiguous","memory","optimization","4",")","CPU","Checkpointing","5",")","Profile","forward","and","backward","functions"]
code: def backward(ctx, *grads):
      global timers
      see_memory_usage("In backward", force=False)
      # removing pointers to the contiguous buffer memory
      # so that they can be garbage collected once the checkpoints
      # have been used
      if SYNCHRONIZE:
          get_accelerator().synchronize()
      if PROFILE_TIME:
          timers("backward").start()
  
      if CONTIGUOUS_CHECKPOINTING:
          global data_offsets, size_offsets
          global contiguous_data_buffers, contiguous_size_buffers
  
          for buffers in contiguous_data_buffers:
              buffers = []
  
          # frees up all the pointers to the checkpoints except for the ones
          # stored by save for backward
          contiguous_data_buffers = []
          contiguous_size_buffers = []
          data_offsets = []
          size_offsets = []
  
      see_memory_usage("In backward checkpointing code", force=False)
      if not torch.autograd._is_checkpoint_valid():
          raise RuntimeError(
              "Checkpointing is not compatible with .grad(), "
              "please use .backward() if possible"
          )
  
      global cuda_device, transport_stream, PARTITION_ACTIVATIONS
  
      # Rebuild deepspeed_saved_tensors
      for t in ctx.deepspeed_saved_tensors:
          if (
              t is not None
              and hasattr(t, "saved_data")
              and t.saved_data is not None
          ):
              t.data = t.saved_data.to(t.device)
              t.saved_data = None
  
      if PARTITION_ACTIVATIONS:
          # with get_accelerator().stream(transport_stream):
          inputs = gather_partitioned_activations(
              ctx.deepspeed_saved_tensors,
              device=cuda_device if CPU_CHECKPOINT else None,
          )
          detached_inputs = detach_variable(inputs)
      elif CPU_CHECKPOINT:
          inputs = move_to_device(
              ctx.deepspeed_saved_tensors,
              cuda_device,
              is_activation_to_checkpoint,
          )
          detached_inputs = detach_variable(inputs)
      else:
          inputs = ctx.deepspeed_saved_tensors
          detached_inputs = detach_variable(inputs)
  
      # Add non tensor input args
      detached_inputs = merge_tensors(
          tensor_objects=detached_inputs,
          non_tensor_objects=ctx.non_tensor_args,
          tensor_flags=ctx.tensor_flags,
      )
  
      # Store the current states.
      bwd_cpu_rng_state = torch.get_rng_state()
      bwd_cuda_rng_state = get_accelerator().get_rng_state()
      bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
  
      # Set the states to what it used to be before the forward pass.
      torch.set_rng_state(ctx.fwd_cpu_rng_state)
      _set_cuda_rng_state(ctx.fwd_cuda_rng_state)
      get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
  
      # if PARTITION_ACTIVATIONS:
      #     current_stream=get_accelerator().current_stream()
      #     current_stream.wait_stream(transport_stream)
  
      see_memory_usage(
          "In backward checkpointing code before forward", force=False
      )
  
      with torch.enable_grad():
          outputs = ctx.run_function(*detached_inputs)
  
      see_memory_usage(
          "In backward checkpointing code after forward", force=False
      )
      # Set the states back to what it was at the start of this function.
      torch.set_rng_state(bwd_cpu_rng_state)
      _set_cuda_rng_state(bwd_cuda_rng_state)
      get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
  
      if isinstance(outputs, torch.Tensor):
          outputs = (outputs,)
  
      # Filter out non tensor outputs
      outputs, _, _ = extract_tensors(all_objects=outputs)
  
      # Construct arguments to autograd.backward().
      # This is usually just outputs and grads, but forward() can return tensors that
      # are not differentiable.
      output_tensors = []
      grad_tensors = []
      for out, grad in zip(outputs, grads):
          if out.requires_grad:
              output_tensors.append(out)
              grad_tensors.append(grad)
  
      see_memory_usage(
          "In backward checkpointing code before backward", force=False
      )
  
      torch.autograd.backward(output_tensors, grad_tensors)
  
      # Force clear our stashed tensors to prevent a memory leak in certain scenarios
      ctx.deepspeed_saved_tensors = None
      ctx.non_tensor_args = None
      ctx.tensor_flags = None
  
      see_memory_usage(
          "After backward checkpointing code after backward",
          force=False,
      )
  
      if PROFILE_TIME:
          timers("backward").stop()
          timers.log(["backward"])
      if SYNCHRONIZE:
          get_accelerator().synchronize()
      ret_list = [None, None]  # first None for ctx
      for inp in detached_inputs:
          if torch.is_tensor(inp):
              ret_list.append(inp.grad)
          else:
              ret_list.append(None)
  
      return tuple(ret_list)
  
imports: import copy
  import torch
  import contextlib
  from deepspeed import comm
  import weakref
  import mmap
  from torch import _C
  from deepspeed.runtime.config import DeepSpeedConfig
  from deepspeed.utils import logger
  from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank
  from deepspeed.utils.timer import SynchronizedWallClockTimer, FORWARD_GLOBAL_TIMER
  from deepspeed.accelerator import get_accelerator
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: tuple