code: def check_linear_module(rank, world_size, port):
      disable_existing_loggers()
      launch(
          config={},
          rank=rank,
          world_size=world_size,
          host="localhost",
          port=port,
          backend="nccl",
      )
      model = LinearModel(4, 8).cuda()
      input = torch.rand(4, 4).cuda()
      output_compare = model(input)
      physical_mesh_id = torch.arange(0, 4)
      mesh_shape = (2, 2)
      # [[0, 1]
      #  [2, 3]]
      device_mesh = DeviceMesh(
          physical_mesh_id, mesh_shape, init_process_group=True
      )
      meta_args = {"x": torch.rand(4, 4).to("meta")}
      gm = initialize_model(
          model, meta_args=meta_args, device_mesh=device_mesh
      )
      output = gm(input)
      assert_close(output, output_compare)
  
file_path: colossalai/tests/test_auto_parallel/test_tensor_shard/test_bias_addition_forward.py
imports: import pytest
  import torch
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.initialize import launch
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing import assert_close, rerun_if_address_is_in_use, run_on_environment_flag, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: check_linear_module(rank, world_size, port) without return types