code: def dist_init(
      distributed_rank: int,
      world_size: int,
      init_method: str,
      device_id: int,
      backend: str = "nccl",
      gpu_streams: int = 1,
  ):
      """
      1. After spawn process per GPU, we want all workers to call init_process_group
      around the same time or times out.
      2. After dist_init, we want all workers to start calling all_reduce/barrier
      around the same time or NCCL timeouts.
      """
      global _round_robin_process_group
  
      if init_method and world_size > 1 and torch.cuda.is_available():
          # providing a large process group timeout to prevent errors during
          # initialization.
          dist_c10d.init_process_group(
              backend=backend,
              init_method=init_method,
              world_size=world_size,
              rank=distributed_rank,
              timeout=timedelta(minutes=90),
          )
          # calling all_reduce for synchronzing all workers
          dist_tensor = torch.tensor(
              [1],
              dtype=torch.float32,
              device="cuda:{}".format(device_id),
          )
          dist_c10d.all_reduce(dist_tensor)
  
          if gpu_streams >= 1:
              _round_robin_process_group = (
                  dist_c10d._round_robin_process_groups(
                      [
                          dist_c10d.new_group(backend=backend)
                          for _ in range(gpu_streams)
                      ]
                  )
              )
  
              for _ in range(gpu_streams):
                  dist_tensor = torch.tensor(
                      [1],
                      dtype=torch.float32,
                      device="cuda:{}".format(device_id),
                  )
                  _round_robin_process_group.allreduce(dist_tensor)
              print(
                  f"Using {gpu_streams} GPU streams for gradient sync."
              )
  
          if distributed_rank != 0:
              suppress_output()
  
imports: from datetime import timedelta
  import pytext.utils.cuda
  import torch
  import torch.distributed
label_desc: Use image node_id 1 for calling a global function with example usage: dist_init(distributed_rank, world_size, init_method, device_id, backend, gpu_streams) without return types