code: def get_ds_opt_model(
      model_name,
      dtype,
      cpu_offload,
      disk_offload,
      offload_dir,
      dummy_weights,
  ):
      import deepspeed
      import torch.distributed as dist
      from transformers.deepspeed import HfDeepSpeedConfig
  
      config = get_model_config(model_name)
      hidden_size = config.hidden_size
      deepspeed.init_distributed("nccl")
      rank = dist.get_rank()
      pin_memory = bool(args.pin_memory)
  
      ds_config = {
          "fp16": {
              "enabled": dtype == torch.float16,
          },
          "bf16": {
              "enabled": dtype == torch.bfloat16,
          },
          "zero_optimization": {
              "stage": 3,
              "stage3_prefetch_bucket_size": hidden_size * hidden_size,
              "stage3_param_persistence_threshold": 0,
          },
          "steps_per_print": 2000,
          "train_batch_size": args.batch_size,
          "wall_clock_breakdown": False,
      }
  
      if cpu_offload:
          ds_config["zero_optimization"]["offload_param"] = dict(
              device="cpu", pin_memory=pin_memory
          )
  
      if disk_offload:
          ds_config["zero_optimization"]["offload_param"] = dict(
              device="nvme",
              pin_memory=True,
              nvme_path=offload_dir,
              buffer_count=5,
              buffer_size=2 * GB,
          )
          ds_config["aio"] = {
              "block_size": 1048576,
              "queue_depth": 8,
              "thread_count": 1,
              "single_submit": False,
              "overlap_events": True,
          }
  
      dschf = HfDeepSpeedConfig(ds_config)
  
      model = OPTForCausalLM.from_pretrained(
          dummy_weights or model_name, torch_dtype=dtype
      )
      model = model.eval()
      ds_engine = deepspeed.initialize(
          model=model, config_params=ds_config
      )[0]
      ds_engine.module.eval()
      model = ds_engine.module
  
      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 5 for calling a global function with example usage: get_ds_opt_model(model_name, dtype, cpu_offload, disk_offload, offload_dir, dummy_weights) and returns: model