def prepare_model_for_finetuning(model: "PreTrainedModel", training_args: BaseHFTrainingArguments):
    """Prepare the model for fine-tuning.

    Freeze base model's layers and prepare model for gradient checkpointing if necessary.
    Similar to peft.prepare_model_for_kbit_training but no casting to fp32 and gradient checkpointing is
    also supported for non-quantized models.

    :param model: The Hugging Face PyTorch model to prepare for fine-tuning.
    :param training_args: The training arguments for the model.
    """
    for param in model.parameters():
        # freeze base model's layers
        param.requires_grad = False

    if training_args.gradient_checkpointing and not model.supports_gradient_checkpointing:
        logger.warning(
            "gradient_checkpointing is True, but model does not support gradient checkpointing! Setting"
            " gradient_checkpoing to False"
        )
        training_args.gradient_checkpointing = False
    elif training_args.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()

    logger.debug("The number of trainable parameters in the original model: %s", count_trainable_parameters(model))
Prepare the model for fine-tuning.

    Freeze base model's layers and prepare model for gradient checkpointing if necessary.
    Similar to peft.prepare_model_for_kbit_training but no casting to fp32 and gradient checkpointing is
    also supported for non-quantized models.

    :param model: The Hugging Face PyTorch model to prepare for fine-tuning.
    :param training_args: The training arguments for the model.
    
olive/passes/pytorch/train_utils.py