code: def balanced_split_pass_v2(gm: torch.fx.GraphModule, pp_size: int):
      """
      In balanced_split_pass_v12, we split module by the size of nodes(weights+bias+outputs).
      """
      mod_graph = gm.graph
      # To use balanced_split_pass_v2, we need run meta_info_prop interpreter first.
      # If nodes don't have meta info, this pass will fall back to normal balanced split pass.
      check_node = list(mod_graph.nodes)[0]
      if "tensor_meta" not in check_node.meta:
          return balanced_split_pass(gm, pp_size)
  
      total_element_size = 0
      for node in mod_graph.nodes:
          total_element_size += node.node_size
  
      partition_size = total_element_size // pp_size
      accumulate_node_size = 0
      for node in mod_graph.nodes:
          if pp_size <= 1:
              break
          if "pipe_split" in node.name:
              continue
          accumulate_node_size += node.node_size
          if accumulate_node_size >= partition_size:
              total_element_size = (
                  total_element_size - accumulate_node_size
              )
              accumulate_node_size = 0
              pp_size -= 1
              partition_size = total_element_size // pp_size
              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/adding_split_node_pass.py
imports: import numpy
  import torch
  import tqdm
  from colossalai.fx.passes.split_module import split_module
label_desc: Use image node_id 12 for calling a global function with example usage: balanced_split_pass_v2(gm, pp_size) and returns: gm, balanced_split_pass