From 4d880cc9ab87e1cd46f9a3bd3d4345dd8e85317d Mon Sep 17 00:00:00 2001
From: "Ye, Xinyu" <xinyu.ye@intel.com>
Date: Fri, 8 Dec 2023 02:49:00 -0500
Subject: [PATCH] fix gradient issue for qlora on seq2seq.

Signed-off-by: Ye, Xinyu <xinyu.ye@intel.com>
---
 intel_extension_for_transformers/llm/finetuning/finetuning.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/intel_extension_for_transformers/llm/finetuning/finetuning.py b/intel_extension_for_transformers/llm/finetuning/finetuning.py
index b08d3b71842..151aca21df7 100644
--- a/intel_extension_for_transformers/llm/finetuning/finetuning.py
+++ b/intel_extension_for_transformers/llm/finetuning/finetuning.py
@@ -821,6 +821,8 @@ def preprocess_logits_for_metrics(logits, labels):
                 model = prepare_model_for_kbit_training(
                     model, use_gradient_checkpointing=training_args.gradient_checkpointing
                 )
+            if training_args.gradient_checkpointing:
+                model.gradient_checkpointing_enable()
 
             if not finetune_args.full_finetune:
                 # PEFT settings