code: def run_test_initialize_model_parallel(
      rank, model_parallel_size, filename, filename_rpc
  ):
      dist_init(rank, model_parallel_size, filename, filename_rpc)
  
      if torch.distributed.get_rank() == 0:
          print(
              "> testing initialize_model_parallel with size {} ...".format(
                  model_parallel_size
              )
          )
      model_parallel_size_ = min(
          model_parallel_size, torch.distributed.get_world_size()
      )
      assert not mpu.model_parallel_is_initialized()
      mpu.initialize_model_parallel(model_parallel_size_)
      assert mpu.model_parallel_is_initialized()
  
      # Checks.
      def check(group, world_size, rank):
          assert world_size == torch.distributed.get_world_size(
              group=group
          )
          assert rank == torch.distributed.get_rank(group=group)
  
      # Model parallel.
      world_size = model_parallel_size_
      rank = torch.distributed.get_rank() % model_parallel_size_
      assert world_size == mpu.get_model_parallel_world_size()
      assert rank == mpu.get_model_parallel_rank()
      check(mpu.get_model_parallel_group(), world_size, rank)
  
      # Data parallel.
      world_size = (
          torch.distributed.get_world_size() // model_parallel_size_
      )
      rank = torch.distributed.get_rank() // model_parallel_size
      assert world_size == mpu.get_data_parallel_world_size()
      assert rank == mpu.get_data_parallel_rank()
      check(mpu.get_data_parallel_group(), world_size, rank)
  
      # Reset groups
      mpu.destroy_model_parallel()
  
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(">> passed the test :-)")
  
imports: import torch
  from fairscale.fair_dev.testing.testing import dist_init, spawn_for_all_world_sizes
  from fairscale.nn.model_parallel import initialize
label_desc: Use image node_id 1 for calling a global function with example usage: run_test_initialize_model_parallel(rank, model_parallel_size, filename, filename_rpc) without return types