code: def initialize_distributed(
      head_ip, port, world_size, rank, local_rank, comm_device
  ):
      print(
          f"Initializing distributed environment at {head_ip}:{port}, "
          f"world_size={world_size}, rank={rank}, local_rank={local_rank}."
      )
  
      # Initialize distributed environment
      torch.cuda.set_device(local_rank)
      distributed_init_method = f"tcp://{head_ip}:{port}"
      global _COMM_DEVICE
      _COMM_DEVICE = comm_device
      if comm_device == "cpu":
          backend = "gloo"
      elif comm_device == "gpu":
          backend = "nccl"
      else:
          raise ValueError(f"Unknown comm_device: {comm_device}")
      dist.init_process_group(
          backend=backend,
          init_method=distributed_init_method,
          world_size=world_size,
          rank=rank,
      )
  
      # Create groups for pipeline parallelism
      global _PIPELINE_PARALLEL_PRED_GROUP, _PIPELINE_PARALLEL_SUCC_GROUP
      if world_size > 1:
          for pred in range(world_size):
              succ = (pred + 1) % world_size
              group = dist.new_group([pred, succ])
              if pred == rank:
                  _PIPELINE_PARALLEL_PRED_GROUP = group
              if succ == rank:
                  _PIPELINE_PARALLEL_SUCC_GROUP = group
  
      suppress_output(rank)
      print("Finished initializing distributed environment")
  
imports: import torch
  import torch.distributed
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_distributed(head_ip, port, world_size, rank, local_rank, comm_device) without return types