code: def run_flexgen_dist(args):
      t_name = args.model.replace("175b", "66b")
      tokenizer = AutoTokenizer.from_pretrained(
          t_name, padding_side="left"
      )
      num_inner_iterations = (
          args.num_inner_iterations
          if args.num_inner_iterations is not None
          else args.world_size
      )
      num_prompts = (
          args.num_gpu_batches
          * args.gpu_batch_size
          * num_inner_iterations
          * 1
      )
      prompt_len, gen_len, cut_gen_len = (
          args.prompt_len,
          args.gen_len,
          args.cut_gen_len,
      )
  
      # Task and policy
      warmup_inputs = get_test_inputs(32, num_prompts, tokenizer)
      inputs = get_test_inputs(prompt_len, num_prompts, tokenizer)
  
      gpu = TorchDevice(f"cuda:{args.local_rank}")
      cpu = TorchDevice("cpu")
      disk = TorchDisk(args.offload_dir, None, args.local_rank)
      env = ExecutionEnv(
          gpu=gpu,
          cpu=cpu,
          disk=disk,
          mixed=TorchMixedDevice([gpu, cpu, disk]),
      )
      TorchTensor.name_count = count(
          start=args.rank, step=args.world_size
      )
  
      comm_test(gpu.dev if args.comm_device == "gpu" else cpu.dev)
  
      policy = Policy(
          args.gpu_batch_size,
          args.num_gpu_batches,
          args.percent[0],
          args.percent[1],
          args.percent[2],
          args.percent[3],
          args.percent[4],
          args.percent[5],
          args.overlap,
          args.sep_layer,
          args.pin_weight,
          args.cpu_cache_compute,
          args.attn_sparsity,
          args.compress_weight,
          CompressionConfig(
              num_bits=4, group_size=64, group_dim=0, symmetric=False
          ),
          args.compress_cache,
          CompressionConfig(
              num_bits=4, group_size=64, group_dim=2, symmetric=False
          ),
      )
      assert not (
          args.compress_cache and args.attn_sparsity < 1.0
      ), "Not implemented"
  
      opt_config = get_opt_config(args.model)
      model = DistOptLM(
          opt_config,
          env,
          args.path,
          policy,
          args.rank,
          args.world_size,
          args.comm_device,
          num_inner_iterations=num_inner_iterations,
          async_comm=args.async_comm,
      )
      cache_size = opt_config.cache_bytes(
          num_prompts, prompt_len + gen_len
      )
      hidden_size = opt_config.hidden_bytes(
          num_prompts, prompt_len + gen_len
      )
      print(
          f"model size: {opt_config.model_bytes()/GB:.3f} GB, "
          f"cache size: {cache_size/GB:.3f} GB, "
          f"hidden size (prefill): {hidden_size/GB:.3f} GB"
      )
  
      try:
          print("warmup - generate")
          output_ids = model.generate(
              warmup_inputs, max_new_tokens=2, verbose=args.verbose
          )
  
          print("benchmark - generate")
          for timer_name in ["generate-prompt", "generate"]:
              timers(timer_name).reset()
          output_ids = model.generate(
              inputs,
              max_new_tokens=args.gen_len,
              debug_mode=args.debug_mode,
              cut_gen_len=cut_gen_len,
              verbose=args.verbose,
          )
          prompt_costs = timers("generate-prompt").costs
          generate_costs = timers("generate").costs
      finally:
          env.close_copy_threads()
  
      if args.rank != args.world_size - 1:
          return
  
      # Log output
      prefill_latency = sum(prompt_costs)
      prefill_throughput = num_prompts * prompt_len / prefill_latency
      if cut_gen_len:  # project latency of cut_gen_len to gen_len
          costs = (
              np.array(generate_costs)
              .reshape(-1, cut_gen_len - 1)
              .sum(axis=0)
              .tolist()
          )
          decode_latency = project_decode_latency(
              [None] + costs, prompt_len, gen_len
          )
      else:
          decode_latency = sum(generate_costs)
      decode_throughput = (
          num_prompts * (gen_len - 1) / max(decode_latency, 1e-10)
      )
      num_generated_tokens = num_prompts * gen_len
      total_latency = prefill_latency + decode_latency
      total_throughput = num_generated_tokens / total_latency
      _, gpu_peak_mem = gpu.mem_stats()
      _, cpu_peak_mem = cpu.mem_stats()
  
      if DUMMY_WEIGHT not in args.path:
          outputs = tokenizer.batch_decode(
              output_ids, skip_special_tokens=True
          )
          show_str = "Outputs:\n" + 70 * "-" + "\n"
          for i in [0, len(outputs) - 1]:
              show_str += f"{i}: {outputs[i]}\n"
              show_str += "-" * 70 + "\n"
          print(show_str)
  
      gpu.print_stats()
      cpu.print_stats()
      projected = args.debug_mode or cut_gen_len
  
      log_str = (
          f"model size: {opt_config.model_bytes()/GB:.3f} GB\t"
          f"cache size: {cache_size/GB:.3f} GB\t"
          f"hidden size (prefill): {hidden_size/GB:.3f} GB\n"
          f"peak gpu mem: {gpu_peak_mem / GB:.3f} GB\n"
          f"prefill latency: {prefill_latency:.2f} s\t"
          f"prefill throughput: {prefill_throughput:.2f} token/s\n"
          f"decode latency: {decode_latency:.2f} s\t"
          f"decode throughput: {decode_throughput:.2f} token/s\n"
          f"total latency: {total_latency:.2f} s\t"
          f"total throughput: {total_throughput:.2f} token/s"
      )
      print(log_str)
  
      if not args.no_log:
          if args.log_file == "auto":
              basename = f"rank-{args.rank}-{get_filename(args)}"
              log_filename = basename + ".log"
          else:
              log_filename = args.log_file
          with open(log_filename, "a") as fout:
              fout.write(log_str + "\n")
  
imports: import argparse
  from itertools import count
  import os
  import pickle
  import traceback
  from typing import Union, List, Optional
  import numpy
  import torch
  import torch.distributed
  from transformers import AutoTokenizer
  from flexgen.compression import CompressionConfig
  from flexgen.dist_utils import initialize_distributed
  from flexgen.flex_opt import Policy, InputEmbed, OutputEmbed, SelfAttention, MLP, TransformerLayer, OptLM, get_filename, add_parser_arguments, get_test_inputs, DUMMY_WEIGHT
  from flexgen.opt_config import get_opt_config
  from flexgen.pytorch_backend import TorchDevice, TorchDisk, TorchLink, TorchMixedDevice, TorchTensor
  from flexgen.timer import timers
  from flexgen.utils import Task, ExecutionEnv, GB, T, ValueHolder, array_1d, array_2d, array_3d, array_4d, str2bool, project_decode_latency
label_desc: Use image node_id 2 for calling a global function with example usage: run_flexgen_dist(args) without return types