Update train.py

train_and_evaluate/train.py

@@ -126,7 +126,7 @@ def train(opt):
     # enable gradient checkpointing to save GPU memory, but this action would slowdown the training speed 20-30%.
     # in addition, gradient_checkpointing can not be enabled when using deepspeed ZERO-3
     model.gradient_checkpointing_enable()
-
+    model.enable_input_require_grads()
     dataset = SFTDataset(opt.sft_data_dir, tokenizer, opt.block_size, opt.mode)
     if accelerator.is_main_process:
         sanity_check(dataset[0]["input_ids"], dataset[0]["labels"], tokenizer)
@@ -207,4 +207,4 @@ def train(opt):
 
 if __name__ == "__main__":
     opt = parse_option()
-    train(opt)
\ No newline at end of file
+    train(opt)