code: def get_accelerate_model(args, model_class):
      n_gpus = torch.cuda.device_count()
      max_memory = f"{args.max_memory_MB}MB"
      max_memory = {i: max_memory for i in range(n_gpus)}
      device_map = "auto"
  
      # if we are in a distributed setting, we need to set the device map and max memory per device
      if os.environ.get("LOCAL_RANK") is not None:
          local_rank = int(os.environ.get("LOCAL_RANK", "0"))
          device_map = {"": local_rank}
          max_memory = {"": max_memory[local_rank]}
  
      if args.full_finetune:
          assert args.bits in [16, 32]
  
      print(f"loading base model {args.model_name_or_path}...")
      compute_dtype = (
          torch.float16
          if args.fp16
          else (torch.bfloat16 if args.bf16 else torch.float32)
      )
      logger.info(f"compute_dtype:{compute_dtype}")
      model = model_class.from_pretrained(
          args.model_name_or_path,
          load_in_4bit=args.bits == 4,
          load_in_8bit=args.bits == 8,
          device_map=device_map,
          max_memory=max_memory,
          offload_folder="./cache",  # If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
          quantization_config=BitsAndBytesConfig(
              load_in_4bit=args.bits == 4,
              load_in_8bit=args.bits == 8,
              llm_int8_threshold=6.0,
              llm_int8_has_fp16_weight=False,
              bnb_4bit_compute_dtype=compute_dtype,
              bnb_4bit_use_double_quant=args.double_quant,
              bnb_4bit_quant_type=args.quant_type,
          ),
          torch_dtype=compute_dtype,
          trust_remote_code=args.trust_remote_code,
          use_auth_token=args.use_auth_token,
      )
      if compute_dtype == torch.float16 and args.bits == 4:
          major, minor = torch.cuda.get_device_capability()
          if major >= 8:
              print("=" * 80)
              print(
                  "Your GPU supports bfloat16, you can accelerate training with the argument --bf16"
              )
              print("=" * 80)
  
      setattr(model, "model_parallel", True)
      setattr(model, "is_parallelizable", True)
      model.config.torch_dtype = compute_dtype
  
      if not args.full_finetune:
          if args.bits in [4, 8]:
              model = prepare_model_for_kbit_training(
                  model,
                  use_gradient_checkpointing=args.gradient_checkpointing,
              )
      if args.gradient_checkpointing:
          model.gradient_checkpointing_enable()
  
      if not args.full_finetune:
          print(f"adding LoRA modules...")
          config = LoraConfig(
              r=args.lora_r,
              lora_alpha=args.lora_alpha,
              target_modules=args.lora_target_modules,
              lora_dropout=args.lora_dropout,
              bias="none",
              task_type="CAUSAL_LM",
          )
          model = get_peft_model(model, config)
          if args.resume_from_checkpoint is not None:
              print("Loading adapters from checkpoint.")
              checkpoint_name = os.path.join(
                  args.resume_from_checkpoint, "pytorch_model.bin"
              )
              if not os.path.exists(checkpoint_name):
                  checkpoint_name = os.path.join(
                      args.resume_from_checkpoint, "adapter_model.bin"
                  )
                  args.resume_from_checkpoint = False
              if os.path.exists(checkpoint_name):
                  print(f"Restarting from {checkpoint_name}")
                  adapters_weights = torch.load(checkpoint_name)
                  set_peft_model_state_dict(model, adapters_weights)
              else:
                  print(f"Checkpoint {checkpoint_name} not found")
      """
      for name, module in model.named_modules():
          #if isinstance(module, LoraLayer):
              #if args.bf16:
                  #module = module.to(torch.bfloat16)
          #if 'norm' in name:
              #module = module.to(torch.float32)
          if 'lm_head' in name or 'embed_tokens' in name:
              if hasattr(module, 'weight'):
                  if args.bf16 and module.weight.dtype == torch.float32:
                      module = module.to(torch.bfloat16)
      """
      return model, args
  
file_path: DeepKE/example/llm/InstructKGC/src/finetune.py
imports: import os
  import sys
  import logging
  import argparse
  from datasets import Dataset
  import torch
  from transformers import HfArgumentParser, DataCollatorForSeq2Seq, GenerationConfig, BitsAndBytesConfig
  from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, set_peft_model_state_dict
  from peft.tuners.lora import LoraLayer
  from utils.general_utils import LORA_TARGET_MODULES_DICT, seed_torch, get_model_tokenizer, get_model_name, get_format_time
  from utils.prompter import Prompter
  from utils.args import DataArguments, ModelArguments, TrainingArguments, GenerationArguments
  from datamodule import COLL_FN_DICT, DataCollatorForCPMBEE
label_desc: Use image node_id 2 for calling a global function with example usage: get_accelerate_model(args, model_class) and returns: model,  args