code: def module_params_sharding_pass(
      gm: torch.fx.GraphModule, device_mesh: DeviceMesh, overlap=False
  ):
      """
      Apply the sharding action to the module parameters and buffers following the
      instructions of solver solution.
      """
      mod_graph = gm.graph
      nodes = tuple(mod_graph.nodes)
      # This stream is created for overlapping the communication and computation.
      reduction_stream = torch.cuda.Stream()
  
      def _add_hook_for_grad_communication(node, param, name=None):
          comm_actions = node.best_strategy.communication_actions
  
          def _filter_param_to_hook(node, op_data, comm_action, name):
              if (
                  node.op == "call_module"
                  and op_data.type == OperationDataType.PARAM
                  and op_data.name == name
                  and comm_action.comm_type == CommType.HOOK
              ):
                  return True
              if (
                  node.op == "get_attr"
                  and isinstance(
                      node._meta_data, torch.nn.parameter.Parameter
                  )
                  and comm_action.comm_type == CommType.HOOK
              ):
                  return True
              return False
  
          for operation_data, comm_action in comm_actions.items():
              comm_spec_to_use = comm_action.comm_spec
              # register hook to the parameters
              if _filter_param_to_hook(
                  node, operation_data, comm_action, name=name
              ):
  
                  def wrapper(param, comm_spec, stream, overlap):
                      def hook_fn(grad):
                          if overlap:
                              with torch.cuda.stream(stream):
                                  _all_reduce(
                                      grad, comm_spec, async_op=True
                                  )
                          else:
                              _all_reduce(
                                  grad, comm_spec, async_op=False
                              )
  
                      param.register_hook(hook_fn)
  
                  wrapper(
                      param,
                      comm_spec_to_use,
                      reduction_stream,
                      overlap=overlap,
                  )
  
      def _shard_param(param, target_sharding_spec):
          # apply the sharding spec of parameters
          if target_sharding_spec.dim_partition_dict != {}:
              origin_sharding_spec = ShardingSpec(
                  device_mesh, param.shape, {}
              )
              setattr(param, "sharding_spec", origin_sharding_spec)
              # TODO: build a ColoParameter class to manager the distributed parameters
              # we could use .data here, because all the operations just happen before the real training
              # loop, so we don't need to track these operations in the autograd graph.
              param = torch.nn.Parameter(
                  shape_consistency_manager.apply_for_autoparallel_runtime(
                      param.data,
                      param.sharding_spec,
                      target_sharding_spec,
                  )
                  .detach()
                  .clone()
              )
          return param
  
      for node in nodes:
          if node.op == "call_module":
              target_module = node.graph.owning_module.get_submodule(
                  node.target
              )
              # TODO: we need to do more actions to take care of the shared parameters.
              if (
                  hasattr(target_module, "processed")
                  and target_module.processed
              ):
                  continue
              setattr(target_module, "processed", True)
              for name, param in target_module.named_parameters():
                  target_sharding_spec = (
                      node.best_strategy.get_sharding_spec_by_name(name)
                  )
                  param = _shard_param(param, target_sharding_spec)
  
                  setattr(target_module, name, param)
                  _add_hook_for_grad_communication(node, param, name)
  
              sharded_buffer_dict = {}
              # apply the sharding spec of buffers
              for name, buffer in target_module.named_buffers():
                  origin_sharding_spec = ShardingSpec(
                      device_mesh, buffer.shape, {}
                  )
                  setattr(buffer, "sharding_spec", origin_sharding_spec)
                  target_sharding_spec = (
                      node.best_strategy.get_sharding_spec_by_name(name)
                  )
                  buffer_sharded = shape_consistency_manager.apply(
                      buffer, target_sharding_spec
                  )
                  sharded_buffer_dict[name] = buffer_sharded
  
              for name, buffer_sharded in sharded_buffer_dict.items():
                  setattr(
                      target_module,
                      name,
                      buffer_sharded.detach().clone(),
                  )
  
          if node.op == "get_attr":
              root = node.graph.owning_module
              atoms = node.target.split(".")
              attr_len = len(atoms)
              if attr_len == 1:
                  target_module = root
                  target = getattr(root, atoms[0])
              else:
                  target_module = root
                  for atom in atoms[:-1]:
                      target_module = getattr(target_module, atom)
                  target = getattr(target_module, atoms[-1])
  
              target_sharding_spec = node.sharding_spec
              target = _shard_param(target, target_sharding_spec)
  
              assert hasattr(target_module, atoms[-1])
              setattr(target_module, atoms[-1], target)
              _add_hook_for_grad_communication(node, target)
  
      return gm
  
file_path: colossalai/colossalai/auto_parallel/passes/runtime_preparation_pass.py
imports: import operator
  from typing import Dict, List, Union
  import torch
  from torch.fx.node import Node
  from colossalai._analyzer.fx.node_util import MetaInfo
  from colossalai.auto_parallel.tensor_shard.constants import RESHAPE_FUNC_OP
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommType, OperationDataType
  from colossalai.auto_parallel.tensor_shard.solver.strategies_constructor import StrategiesConstructor
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.tensor.comm_spec import _all_reduce
  from colossalai.tensor.shape_consistency import ShapeConsistencyManager
  from colossalai.tensor.sharding_spec import ShardingSpec
  from .constants import SHAPE_ARGUMENT_OPS
label_desc: Use image node_id 5 for calling a global function with example usage: module_params_sharding_pass(gm, device_mesh, overlap) and returns: gm, param, False, True, True