code: def prepare_model_for_kbit_training(
      model, use_gradient_checkpointing=True
  ):
      r"""
      This method wraps the entire protocol for preparing a model before running a training.
      This includes:
          1- Cast the layernorm in fp32
          2- making output embedding layer require grads
          3- Add the upcasting of the lm head to fp32
  
      Args:
          model, (`transformers.PreTrainedModel`):
              The loaded model from `transformers`
      """
  
      is_gptq_quantized = (
          getattr(model, "quantization_method", None) == "gptq"
      )
      for name, param in model.named_parameters():
          # freeze base model's layers
          param.requires_grad = False
  
      if not is_gptq_quantized:
          # cast all non INT8 parameters to fp32
          # for param in model.parameters():
          #     if (param.dtype == torch.float16) or (param.dtype == torch.bfloat16):
          #         param.data = param.data.to(torch.float32)
  
          # change to below way to reduce memory for Linear
          # otherwise lora finetuning on arc may OOM at this convert
          for module in model.modules():
              if list(module.children()) == []:
                  # leaf module
                  if not isinstance(module, (Linear, Embedding)):
                      for param in module.parameters():
                          if (param.dtype == torch.float16) or (
                              param.dtype == torch.bfloat16
                          ):
                              param.data = param.data.to(torch.float32)
  
      if use_gradient_checkpointing:
          # For backward compatibility
          if hasattr(model, "enable_input_require_grads"):
              model.enable_input_require_grads()
          else:
  
              def make_inputs_require_grad(module, input, output):
                  output.requires_grad_(True)
  
              model.get_input_embeddings().register_forward_hook(
                  make_inputs_require_grad
              )
  
          # enable gradient checkpointing for memory efficiency
          model.gradient_checkpointing_enable()
  
      return model
  
imports: import torch
  from torch.nn import Linear, Embedding
  from bigdl.llm.transformers.low_bit_linear import LowBitLinear, BF16Linear, get_qk_size
  from peft.tuners.lora import LoraLayer
  from bigdl.llm.utils.common import invalidInputError
  from bigdl.llm.transformers.utils import get_autocast_dtype
  from bigdl.llm.ggml.quantize import ggml_tensor_qtype
  import functools
  from bigdl.llm.transformers import training_patch
  from peft.tuners.lora import LoraModel
  from peft.tuners.lora import LoraConfig
  from dataclasses import dataclass, field
  from peft.mapping import PEFT_TYPE_TO_CONFIG_MAPPING
label_desc: Use image node_id 3 for calling a global function with example usage: prepare_model_for_kbit_training(model, use_gradient_checkpointing) and returns: model