code: def init_colo_module(
      module: torch.nn.Module,
      compute_spec: ComputeSpec,
      pg: ProcessGroup,
      recursive=True,
      mode="default",
  ):
      compute_pattern = compute_spec.compute_pattern
      if is_colo_module(module):
          # for each param
          # set its process_group, dist_spec and compute_spec
          colo_module = get_colo_module(module)
          colo_module.register(compute_pattern, pg)
          if not colo_module.has_compute_pattern_with_mode(
              compute_pattern, mode=mode
          ):
              raise NotImplementedError
          # a set for modules which update at least one param in the init process.
          # these modules need to be checked whether all params still match one of the valid compute pattern.
          modules_update_param = {module}
          for (
              param_name,
              dist_spec,
          ) in colo_module.get_dist_specs_with_mode(
              compute_pattern, mode=mode
          ).items():
              if dist_spec is None:
                  continue
              param = module.get_parameter(param_name)
              if isinstance(param, ColoParameter):
                  param.set_process_group(pg)
                  param.set_dist_spec(dist_spec)
                  param.compute_spec = compute_spec
                  for mod in param.shared_param_modules:
                      modules_update_param.add(mod)
          for mod in modules_update_param:
              check_colo_module(mod, pg, recursive=False)
      if recursive == True:
          for submodule in module.children():
              init_colo_module(
                  submodule,
                  compute_spec,
                  pg=pg,
                  recursive=True,
                  mode=mode,
              )
  
file_path: colossalai/colossalai/legacy/nn/parallel/layers/module_utils.py
imports: from typing import Dict
  import torch
  from colossalai.legacy.tensor import ComputeSpec, ProcessGroup
  from colossalai.tensor import ColoParameter
  from .None import ColoModule
label_desc: Use image node_id 5 for calling a global function with example usage: init_colo_module(module, compute_spec, pg, recursive, mode) without return types