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 forward(ctx, run_function, distribute_saved_activations, *args):
      ctx.run_function = run_function
      ctx.distribute_saved_activations = distribute_saved_activations
  
      # Copy the rng states.
      ctx.fwd_cpu_rng_state = torch.get_rng_state()
      ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
      ctx.fwd_cuda_rng_state_tracker = (
          get_cuda_rng_tracker().get_states()
      )
  
      with torch.no_grad():
          outputs = run_function(*args)
  
      # Divide hidden states across model parallel group and only keep
      # the chunk corresponding to the current rank.
      if distribute_saved_activations:
          ctx.input_0_shape = args[0].data.shape
          safely_set_viewless_tensor_data(
              args[0],
              split_tensor_into_1d_equal_chunks(
                  args[0].data, new_buffer=True
              ),
          )
  
      # Store everything.
      ctx.save_for_backward(*args)
  
      return outputs
  
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 1 for calling the CheckpointFunction obj's underlying member method code with example usage: obj.forward(ctx, run_function, distribute_saved_activations) and returns: outputs