code: def prepare_model_for_kbit_training(
      model: Any,
      use_gradient_checkpointing: bool = True,
      use_reentrant: Optional[bool] = True,
  ) -> Any:
      """
      Calculates where to place the gradient checkpoints given n_layers.
      We also freeze all other layers's gradients
  
      Args:
          model: Any LlamaModel with layers.
          use_gradient_checkpointing (`bool`, *optional*):
              Default enabled. Provides memory savings by not saving all activations,
              but only some.
          use_reentrant (`bool`, *optional*):
              https://github.com/pytorch/pytorch/blob/main/torch/utils/checkpoint.py#L354
              Optimal gradient checkpointing algorithm which will be the default in
              future Pytorch versions.
      """
  
      # Freeze all parameters
      for param in model.parameters():
          param.requires_grad_(False)
  
      if use_gradient_checkpointing:
          model.gradient_checkpointing_enable()
  
      # If use_reentrant = True which is the Pytorch default, we just make the input requires_grad.
      if use_reentrant:
          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
              )
  
      return model
  
file_path: unsloth/unsloth/models/_utils.py
imports: import torch
  from typing import Union, Optional, List, Any, Callable
  import numpy
  import warnings
  import gc
  import bitsandbytes
  from transformers.models.llama.modeling_llama import logger
  from transformers import AutoTokenizer
  from platform import system
  import xformers.ops.fmha
  from xformers import __version__
  from peft.tuners.lora.layer import LoraLayer
  from peft.tuners.lora import LoraLayer
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_model_for_kbit_training(model, use_gradient_checkpointing, use_reentrant) and returns: model