class_docstr: A context manager to split the model into pipeline stages.
class_docstr_tok: ["A","context","manager","to","split","the","model","into","pipeline","stages","."]
docstr: The Callback function when entering the context
docstr_tok: ["The","Callback","function","when","entering","the","context"]
code: def _pre_context_exec(self):
      """
      The Callback function when entering the context
      """
      # reserve rng states
      self.cpu_rng_state = torch.get_rng_state()
      self.cuda_rng_state = torch.cuda.get_rng_state()
  
imports: import torch
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.utils import CheckpointModule
  from colossalai.tensor import ColoParameter
  from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses
  from .layer_spec import LayerSpec
  from .utils import build_kwargs_for_module, call_module, customized_partition, exec_funcs_with_kwargs, partition_balanced, partition_uniform
label_desc: Use image node_id 6 for calling the PipelinableContext obj's underlying member method code with example usage: obj._pre_context_exec() without return types