class_docstr: Checkpoint Function 
  
  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`
  2) the states in the model parallel tracker are also properly tracked/set/reset.
class_docstr_tok: ["Checkpoint","Function","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","`","2",")","the","states","in","the","model","parallel","tracker","are","also","properly","tracked\/set\/reset","."]
code: def backward(ctx, *args):
      if not torch.autograd._is_checkpoint_valid():
          raise RuntimeError(
              "Checkpointing is not compatible with .grad(), "
              "please use .backward() if possible"
          )
      inputs = ctx.saved_tensors
      if ctx.distribute_saved_activations:
          safely_set_viewless_tensor_data(
              inputs[0],
              gather_split_1d_tensor(inputs[0].data).view(
                  ctx.input_0_shape
              ),
          )
  
      # Store the current states.
      bwd_cpu_rng_state = torch.get_rng_state()
      bwd_cuda_rng_state = torch.cuda.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)
  
      # Compute the forward pass.
      detached_inputs = detach_variable(inputs)
      with torch.enable_grad():
          outputs = ctx.run_function(*detached_inputs)
  
      # 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 for backward pass
      outputs, args = zip(
          *filter(lambda x: torch.is_tensor(x[0]), zip(outputs, args))
      )
      torch.autograd.backward(outputs, args)
      grads = tuple(
          inp.grad if isinstance(inp, torch.Tensor) else inp
          for inp in detached_inputs
      )
      return (None, None) + grads
  
imports: import contextlib
  import torch
  from torch import _C
  from torch.cuda import _lazy_call
  from torch.cuda import device
  from torch.utils.checkpoint import detach_variable
  from megatron.core.parallel_state import get_data_parallel_rank, get_expert_model_parallel_rank, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from megatron.core.utils import safely_set_viewless_tensor_data
  from .utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks
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