code: def run_all_reduce(local_rank, args):
      if args.dist == "torch":
          import torch.distributed as dist
      elif args.dist == "deepspeed":
          import deepspeed.comm as dist
  
      # Prepare benchmark header
      print_header(args, "all_reduce")
  
      world_size = dist.get_world_size()
      global_rank = dist.get_rank()
  
      if args.scan:
          M_LIST = []
          for x in (2**p for p in range(1, args.maxsize)):
              M_LIST.append(x)
  
          sync_all()
          # loop over various tensor sizes
          for M in M_LIST:
              global_rank = dist.get_rank()
              try:
                  mat = torch.ones(
                      world_size, M, dtype=getattr(torch, args.dtype)
                  ).cuda(local_rank)
                  sync_all()
                  input = (mat.mul_(float(global_rank))).view(-1)
              except RuntimeError as e:
                  if "out of memory" in str(e):
                      if dist.get_rank() == 0:
                          print(
                              "WARNING: Ran out of GPU memory. Exiting comm op."
                          )
                      sync_all()
                      break
              sync_all()
              timed_all_reduce(input, args)
      else:
          # Send the biggest message size our GPUs can fit. If you're facing OOM errors, reduce the mem_factor
          # Don't need output tensor, so we double mem_factor
          elements_per_gpu = max_numel(
              comm_op="all_reduce",
              dtype=getattr(torch, args.dtype),
              mem_factor=args.mem_factor * 2,
              local_rank=local_rank,
              args=args,
          )
          try:
              mat = torch.ones(
                  elements_per_gpu, dtype=getattr(torch, args.dtype)
              ).cuda(local_rank)
              input = (mat.mul_(float(global_rank))).view(-1)
          except RuntimeError as e:
              if "out of memory" in str(e):
                  if dist.get_rank() == 0:
                      print(
                          "WARNING: Ran out of GPU memory. Try to reduce the --mem-factor argument!"
                      )
                  sync_all()
                  return
          sync_all()
          timed_all_reduce(input, args)
  
imports: from benchmarks.communication.utils import *
  from benchmarks.communication.constants import *
  import time
label_desc: Use image node_id 2 for calling a global function with example usage: run_all_reduce(local_rank, args) without return types