code: def main():
      # for this trivial example peer_rank == rank and peer_group_size == world_size
  
      torch.distributed.init_process_group("nccl")
      rank = torch.distributed.get_rank()
      world_size = torch.distributed.get_world_size()
      torch.cuda.set_device(rank)
      peer_ranks = [i for i in range(world_size)]
      pool = PeerMemoryPool(0, 2 * 1024 * 1024, peer_ranks)
  
      num_steps = 100
  
      half_halo = 1
      halo_ex = PeerHaloExchanger1d(peer_ranks, rank, pool, half_halo)
  
      H_split_tests(
          1,
          64,
          336,
          200,
          half_halo,
          rank,
          world_size,
          halo_ex,
          num_steps,
      )
      W_split_tests(
          1,
          64,
          200,
          336,
          half_halo,
          rank,
          world_size,
          halo_ex,
          num_steps,
      )
  
file_path: apex/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
imports: import unittest
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 5 for calling a global function with example usage: main() without return types