code: def initialize_torch_distributed():
      if torch.cuda.is_available():
          from torch.distributed import ProcessGroupNCCL
  
          # Set the device id.
          assert (
              WORLD_SIZE <= torch.cuda.device_count()
          ), "Each process is one gpu"
          device = RANK % torch.cuda.device_count()
          torch.cuda.set_device(device)
          torch.cuda.set_per_process_memory_fraction(
              MEMORY_FRACTION, device
          )
          backend = "nccl"
          options = ProcessGroupNCCL.Options()
          options.is_high_priority_stream = True
          options._timeout = timedelta(seconds=60)
      else:
          backend = "gloo"
          options = None
  
      if WORLD_SIZE == 1:
          return FakeGroup(RANK, WORLD_SIZE), RANK, WORLD_SIZE
      else:
          if os.getenv("DEBUG", None) == "1":
              return FakeGroup(RANK, WORLD_SIZE), RANK, WORLD_SIZE
  
          if not torch.distributed.is_initialized():
              # Call the init process.
              torch.distributed.init_process_group(
                  backend=backend,
                  world_size=WORLD_SIZE,
                  rank=RANK,
                  timeout=timedelta(seconds=60),
                  pg_options=options,
              )
          else:
              logger.warning(
                  "torch.distributed is already initialized."
              )
  
          return torch.distributed.group.WORLD, RANK, WORLD_SIZE
  
file_path: text-generation-inference/server/text_generation_server/utils/dist.py
imports: import os
  import torch
  from datetime import timedelta
  from loguru import logger
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_torch_distributed() and returns: FakeGroup,  RANK,  WORLD_SIZE, torch,  RANK,  WORLD_SIZE, FakeGroup,  RANK,  WORLD_SIZE