code: def split_with_split_nodes_pass_for_gp2_test(
      annotated_gm: torch.fx.GraphModule,
  ):
      """
      This pass will be used in gpt2 test, only a part of changes may be added into
      split_with_split_nodes_pass, and it will be deprecated in future.
      """
      part_idx = 0
  
      def eliminate_unused_placeholders(gm):
          for node in gm.graph.nodes:
              if node.op == "placeholder":
                  if not len(node.users):
                      gm.graph.erase_node(node)
          gm.recompile()
          return gm
  
      def refill_outputs_and_placeholders(
          gm, next_partition_placeholders
      ):
          """
          This method is used to eliminate the outputs in previous partition which is unused in next partition.
          In split module pass, it treats partitions as a DAG, but we need treat them as a single direction linked list in pipeline parallel.
          The difference is if a output from partition 0 is an input argument of partition 3, the DAG will not transfer it
          to partition 1 and partition 2. However, in single direction linked list, we need to do so.
          """
          output_type = None
          output_args = []
          non_output_list = []
          new_placeholder_list = []
          for node in gm.graph.nodes:
              if node.op == "output":
                  if isinstance(node.args[0], (tuple, list)):
                      output_type = node.args[0].__class__
                      output_args.extend([n.name for n in node.args[0]])
                  else:
                      output_args.append(node.args[0].name)
                  rm_list = []
                  for name in output_args:
                      if (
                          next_partition_placeholders
                          and name not in next_partition_placeholders
                      ):
                          rm_list.append(name)
                  for name in rm_list:
                      output_args.remove(name)
                  gm.graph.erase_node(node)
              else:
                  non_output_list.append(node.name)
  
          for name in next_partition_placeholders:
              if name not in output_args:
                  output_args.append(name)
  
          for name in output_args:
              if name not in non_output_list:
                  gm.graph.placeholder(name)
  
          # convert name to node for output_args
          for index, name in enumerate(output_args):
              for n in gm.graph.nodes:
                  if n.name == name:
                      output_args[index] = n
                      continue
  
          # reorder the output args to make sure
          # output args has same order as next partition placeholder
          reorder_output_args = []
          if next_partition_placeholders:
              for name in next_partition_placeholders:
                  for node in output_args:
                      if node.name == name:
                          reorder_output_args.append(node)
                          continue
  
          for node in gm.graph.nodes:
              if node.op == "placeholder":
                  new_placeholder_list.append(node.name)
          if output_type is not None:
              gm.graph.output(output_type(output_args))
          else:
              gm.graph.output(output_args)
          gm.recompile()
          return gm, new_placeholder_list
  
      def split_callback(n: torch.fx.Node):
          nonlocal part_idx
          if (n.op, n.target) == ("call_function", pipe_split):
              part_idx += 1
          return part_idx
  
      split_mod = split_module_for_gpt2_test(
          annotated_gm, None, split_callback
      )
      split_submodules = []
      for name, submodule in split_mod.named_modules():
          if isinstance(submodule, torch.fx.GraphModule):
              for node in submodule.graph.nodes:
                  if (node.op, node.target) == (
                      "call_function",
                      pipe_split,
                  ):
                      submodule.graph.erase_node(node)
              submodule.recompile()
              split_submodules.append(submodule)
  
      submodules = list(split_mod.children())
      placeholder_dict = {}
      for submodule in submodules:
          submodule = eliminate_unused_placeholders(submodule)
          placeholder_dict[submodule] = []
      submodules.reverse()
      for index, submodule in enumerate(submodules):
          if index == 0:
              placeholder_list = []
          else:
              placeholder_list = placeholder_dict[submodules[index - 1]]
          (
              submodule,
              placeholder_dict[submodule],
          ) = refill_outputs_and_placeholders(
              submodule, placeholder_list
          )
          submodule.recompile()
  
      split_mod.recompile()
  
      return split_mod, split_submodules
  
file_path: ColossalAI/colossalai/fx/passes/passes_for_gpt2_test.py
imports: import inspect
  from typing import Any, Callable, Dict, List, Optional
  import torch
  from packaging import version
  from torch.fx._compatibility import compatibility
  from torch.fx.graph_module import GraphModule
  from colossalai.fx.passes.adding_split_node_pass import balanced_split_pass, pipe_split
  from colossalai.fx.passes.meta_info_prop import TensorMetadata
  from colossalai.fx.passes.split_module import Partition
label_desc: Use image node_id 2 for calling a global function with example usage: split_with_split_nodes_pass_for_gp2_test(annotated_gm) and returns: split_mod,  split_submodules, gm, gm,  new_placeholder_list, part_idx