code: def get_hf_opt_model(
      model_name,
      dtype,
      cpu_offload,
      disk_offload,
      offload_dir,
      num_gpus,
      dummy_weights,
  ):
      if num_gpus == 1 and dtype != torch.int8:
          # Here we use a custom device_map instead of device_map == "auto"
          # becase we want to offload as many as possible weights out of GPU
          # to allow a larger batch size.
          if cpu_offload:
              # NOTE: We must put some weights on GPU. Otherwise, huggingface reports errors.
              device_map = {
                  "model.decoder.embed_tokens.weight": 0,
                  "model.decoder.embed_positions.weight": 0,
                  "model.decoder.final_layer_norm": "cpu",
                  "model.decoder.layers": "cpu",
                  "lm_head.weight": 0,
              }
          elif disk_offload:
              device_map = {
                  "model.decoder.embed_tokens.weight": 0,
                  "model.decoder.embed_positions.weight": 0,
                  "model.decoder.final_layer_norm": "disk",
                  "model.decoder.layers": "disk",
                  "lm_head.weight": 0,
              }
          else:
              device_map = None
          max_memory = None
      else:
          # Here we use device_map == "auto", but set a low `max_memory` threshold
          # becase we want to offload as many as possible weights out of GPU
          # to allow a larger batch size.
          device_map = "auto"
          if cpu_offload:
              # `max_memory` should be larger than the embedding.
              # We use 2GB here because the embeding of opt-175b is 1.2GB.
              max_memory = {k: "2GB" for k in range(num_gpus)}
          elif disk_offload:
              max_memory = {k: "2GB" for k in range(num_gpus)}
          else:
              max_memory = {k: "14GB" for k in range(num_gpus)}
          max_memory["cpu"] = "160GB"
  
      if dtype == torch.int8:
          kwargs = {"load_in_8bit": True}
      else:
          kwargs = {"torch_dtype": dtype}
  
      disable_torch_init()
      model = OPTForCausalLM.from_pretrained(
          dummy_weights or model_name,
          device_map=device_map,
          max_memory=max_memory,
          offload_folder=offload_dir,
          **kwargs,
      )
      if device_map is None:
          model.cuda()
  
      model.eval()
      return model
  
file_path: H2O/h2o_flexgen/benchmark/hf_ds/hf_opt.py
imports: import argparse
  import multiprocessing
  import os
  import pickle
  import time
  import numpy
  from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch
  from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
  from transformers import OPTForCausalLM
  import torch
  from flexgen.timer import timers
  from flexgen.utils import GB, project_decode_latency, write_benchmark_log
  from flexgen.opt_config import get_opt_config, disable_torch_init, disable_hf_opt_init
label_desc: Use image node_id 6 for calling a global function with example usage: get_hf_opt_model(model_name, dtype, cpu_offload, disk_offload, offload_dir, num_gpus, dummy_weights) and returns: model