code: def prepare_model_for_kbit_training(
      model,
      use_gradient_checkpointing=True,
      gradient_checkpointing_kwargs=None,
  ):
      r"""
      Note this method only works for `transformers` models.
  
      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`
          use_gradient_checkpointing (`bool`, *optional*, defaults to `True`):
              If True, use gradient checkpointing to save memory at the expense of slower backward pass.
          gradient_checkpointing_kwargs (`dict`, *optional*, defaults to `None`):
              Keyword arguments to pass to the gradient checkpointing function, please refer to the documentation of
              `torch.utils.checkpoint.checkpoint` for more details about the arguments that you can pass to that method.
              Note this is only available in the latest transformers versions (> 4.34.1).
      """
      loaded_in_kbit = getattr(
          model, "is_loaded_in_8bit", False
      ) or getattr(model, "is_loaded_in_4bit", False)
      is_gptq_quantized = (
          getattr(model, "quantization_method", None) == "gptq"
      )
      if gradient_checkpointing_kwargs is None:
          gradient_checkpointing_kwargs = {}
  
      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)
  
      if (
          loaded_in_kbit or is_gptq_quantized
      ) and use_gradient_checkpointing:
          # When having `use_reentrant=False` + gradient_checkpointing, there is no need for this hack
          if (
              "use_reentrant" not in gradient_checkpointing_kwargs
              or gradient_checkpointing_kwargs["use_reentrant"]
          ):
              # 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
                  )
  
          # To support older transformers versions, check if the model supports gradient_checkpointing_kwargs
          _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list(
              inspect.signature(
                  model.gradient_checkpointing_enable
              ).parameters
          )
  
          if (
              not _supports_gc_kwargs
              and len(gradient_checkpointing_kwargs) > 0
          ):
              warnings.warn(
                  "gradient_checkpointing_kwargs is not supported in this version of transformers. The passed kwargs will be ignored."
                  " if you want to use that feature, please upgrade to the latest version of transformers.",
                  FutureWarning,
              )
  
          gc_enable_kwargs = (
              {}
              if not _supports_gc_kwargs
              else {
                  "gradient_checkpointing_kwargs": gradient_checkpointing_kwargs
              }
          )
  
          # enable gradient checkpointing for memory efficiency
          model.gradient_checkpointing_enable(**gc_enable_kwargs)
      return model
  
imports: import copy
  import inspect
  import warnings
  from typing import Optional, Tuple
  import accelerate
  import torch
  from accelerate.hooks import add_hook_to_module, remove_hook_from_module
  from accelerate.utils import is_npu_available, is_xpu_available
  from safetensors.torch import storage_ptr, storage_size
  from ..import_utils import is_auto_gptq_available, is_torch_tpu_available
  from .constants import CONFIG_NAME, EMBEDDING_LAYER_NAMES, INCLUDE_LINEAR_LAYERS_SHORTHAND, SAFETENSORS_WEIGHTS_NAME, TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING, TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING, TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING, TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, bloom_model_postprocess_past_key_value, starcoder_model_postprocess_past_key_value
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_model_for_kbit_training(model, use_gradient_checkpointing, gradient_checkpointing_kwargs) and returns: model