code: def register_module(self, module: nn.Module):
      assert (
          self.ranks_in_group is not None
      ), f"Rank {gpc.get_local_rank(ParallelMode.PIPELINE)} is not in pipeline_ranks {self.pipeline_ranks}"
      src = self.ranks_in_group[self.pipeline_ranks[0]]
      for p in module.parameters():
          setattr(p, "pipeline_shared_module_pg", self.group)
          dist.broadcast(p, src, group=self.group)
  
file_path: colossalai/colossalai/legacy/nn/layer/wrapper/pipeline_wrapper.py
imports: from typing import List, Tuple, Union
  import torch.distributed
  import torch.nn
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
label_desc: Use image node_id 3 for calling the PipelineSharedModuleWrapper obj's underlying member method code with example usage: obj.register_module(module) without return types