code: def initialize_distributed(backend="nccl"):
      """Initialize torch.distributed."""
      # Get local rank in case it is provided.
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--local_rank",
          type=int,
          default=None,
          help="local rank passed from distributed launcher",
      )
      args = parser.parse_args()
      local_rank = args.local_rank
  
      # Get rank and world size.
      rank = int(os.getenv("RANK", "0"))
      world_size = int(os.getenv("WORLD_SIZE", "1"))
  
      print(
          "> initializing torch.distributed with local rank: {}, "
          "rank: {}, world size: {}".format(
              local_rank, rank, world_size
          )
      )
  
      # Set the device id.
      device = rank % torch.cuda.device_count()
      if local_rank is not None:
          device = local_rank
      torch.cuda.set_device(device)
  
      # Call the init process.
      init_method = "tcp://"
      master_ip = os.getenv("MASTER_ADDR", "localhost")
      master_port = os.getenv("MASTER_PORT", "6000")
      init_method += master_ip + ":" + master_port
      torch.distributed.init_process_group(
          backend=backend,
          world_size=world_size,
          rank=rank,
          init_method=init_method,
      )
  
file_path: megatron-lm/megatron/mpu/tests/commons.py
imports: import argparse
  import os
  import random
  import numpy
  import torch
  import mpu
label_desc: Use image node_id 2 for calling a global function with example usage: initialize_distributed(backend) without return types