code: def benchmark_inference(args):
      config = CONFIG_MAP[args.model]
      model = transformers.LlamaForCausalLM(config)
      if dist.get_rank() == 0:
          print("Model loaded")
      engine = InferenceEngine(
          pp_size=args.pp_size,
          tp_size=args.tp_size,
          dtype=args.dtype,
          micro_batch_size=args.mb_size,
          model=model,
          verbose=args.verbose,
          max_batch_size=args.batch_size,
          max_input_len=args.seq_len,
          max_output_len=args.output_len,
      )
      data = data_gen(args.batch_size, args.seq_len)
  
      N_WARMUP_STEPS = 2
  
      for _ in range(N_WARMUP_STEPS):
          engine.generate(data)
  
      torch.cuda.synchronize()
      whole_end2end = time.time()
      outputs = engine.generate(data)
      torch.cuda.synchronize()
      whole_end2end = time.time() - whole_end2end
  
      print_details_info(outputs, model.config, args, whole_end2end)
  
file_path: ColossalAI/examples/inference/benchmark_llama.py
imports: import argparse
  import time
  import torch
  import torch.distributed
  import transformers
  import colossalai
  import colossalai.utils.device
  from colossalai.inference import InferenceEngine
  from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn
  from colossalai.utils.device import get_current_device
label_desc: Use image node_id 3 for calling a global function with example usage: benchmark_inference(args) without return types