code: def _comm_spec_apply(gm: torch.fx.GraphModule):
      """
      This pass is used to add the comm spec apply node to the origin graph.
      """
      mod_graph = gm.graph
      nodes = tuple(mod_graph.nodes)
  
      (
          _,
          _,
          comm_actions_dict_node,
          node_to_index_dict,
      ) = _preprocess_graph(nodes)
  
      for node in nodes:
          if not hasattr(node, "best_strategy") or node.op == "output":
              continue
  
          comm_actions = node.best_strategy.communication_actions
          for op_data, comm_action in comm_actions.items():
              if comm_action.comm_type == CommType.HOOK:
                  continue
              if comm_action.comm_type == CommType.BEFORE:
                  if op_data.type == OperationDataType.OUTPUT:
                      comm_object = node
                  elif comm_action.key_for_kwarg is not None:
                      comm_object = node.kwargs[
                          comm_action.key_for_kwarg
                      ]
                  else:
                      comm_object = node.args[comm_action.arg_index]
                  with mod_graph.inserting_before(node):
                      comm_spec_apply_node = mod_graph.create_node(
                          "call_function",
                          runtime_comm_spec_apply,
                          args=(
                              comm_object,
                              comm_actions_dict_node,
                              node_to_index_dict[node],
                              op_data.name,
                          ),
                      )
                  # the origin node may be a positional argument or key word argument of user node
                  if comm_action.key_for_kwarg is not None:
                      # substitute the origin node with comm_spec_apply_node
                      new_kwargs = dict(node.kwargs)
                      new_kwargs[
                          comm_action.key_for_kwarg
                      ] = comm_spec_apply_node
                      node.kwargs = new_kwargs
                  else:
                      # substitute the origin node with comm_spec_apply_node
                      new_args = list(node.args)
                      new_args[
                          comm_action.arg_index
                      ] = comm_spec_apply_node
                      node.args = tuple(new_args)
  
              elif comm_action.comm_type == CommType.AFTER:
                  with mod_graph.inserting_after(node):
                      comm_spec_apply_node = mod_graph.create_node(
                          "call_function",
                          runtime_comm_spec_apply,
                          args=(
                              node,
                              comm_actions_dict_node,
                              node_to_index_dict[node],
                              op_data.name,
                          ),
                      )
                  user_list = list(node.users.keys())
                  for user in user_list:
                      if user == comm_spec_apply_node:
                          continue
                      new_args = list(user.args)
                      new_kwargs = dict(user.kwargs)
                      # the origin node may be a positional argument or key word argument of user node
                      if node in new_args:
                          # substitute the origin node with comm_spec_apply_node
                          new_args[
                              new_args.index(node)
                          ] = comm_spec_apply_node
                          user.args = tuple(new_args)
                      elif str(node) in new_kwargs:
                          # substitute the origin node with comm_spec_apply_node
                          new_kwargs[str(node)] = comm_spec_apply_node
                          user.kwargs = new_kwargs
              if hasattr(node.meta["info"], "activation_checkpoint"):
                  MetaInfo(
                      comm_spec_apply_node,
                      mod_dir=node.meta["info"].mod_dir,
                      activation_checkpoint=tuple(
                          node.meta["info"].activation_checkpoint
                      ),
                  )
  
      return gm
  
file_path: ColossalAI/colossalai/auto_parallel/passes/runtime_apply_pass.py
imports: from typing import Dict, List
  import torch
  from torch.fx.node import Node
  from colossalai._analyzer.fx.node_util import MetaInfo
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommType, OperationDataType
  from colossalai.tensor.comm_spec import CommSpec
  from colossalai.tensor.shape_consistency import ShapeConsistencyManager
  from colossalai.tensor.sharding_spec import ShardingSpec
label_desc: Use image node_id 6 for calling a global function with example usage: _comm_spec_apply(gm) and returns: gm