code: def customized_split_pass_for_gpt2(
      gm: torch.fx.GraphModule, pp_size: int, partition_list: List[int]
  ):
      """
      This pass is only used to do the gpt2 performance test, it may move into adding_split_node_pass.py, and will be deprecated in future.
      """
      mod_graph = gm.graph
      valid_children_size = 0
      valid_children = []
      for node in mod_graph.nodes:
          if node.op == "call_module":
              valid_children_size += 1
              valid_children.append(node.target)
      if valid_children_size < pp_size:
          # If valid children is not enough to shard, we will use balanced policy instead of uniform policy.
          return balanced_split_pass(gm, pp_size)
      accumulate_layer_amount = 0
      list_of_part = partition_list
      part_index = 0
      for node in mod_graph.nodes:
          if pp_size <= 1:
              break
          if node.op == "call_module":
              if node.target in valid_children:
                  accumulate_layer_amount += 1
          if accumulate_layer_amount == list_of_part[part_index]:
              part_index += 1
              pp_size -= 1
              with mod_graph.inserting_after(node):
                  split_node = mod_graph.create_node(
                      "call_function", pipe_split
                  )
  
      gm.recompile()
      return gm
  
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 1 for calling a global function with example usage: customized_split_pass_for_gpt2(gm, pp_size, partition_list) and returns: gm, balanced_split_pass