test: add a unit test for additional_callbacks param of train()

Signed-off-by: Vassilis Vassiliadis <vassilis.vassiliadis@ibm.com>

Author: VassilisVassiliadis

tests/test_sft_trainer.py

@@ -616,3 +616,22 @@ def test_bad_torch_dtype():
 
         with pytest.raises(ValueError):
             sft_trainer.train(model_args, DATA_ARGS, train_args, PEFT_PT_ARGS)
+
+
+def test_run_with_additional_callbacks():
+    """Ensure that train() can work with additional_callbacks"""
+    # Third Party
+    from transformers.trainer_callback import TrainerCallback
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        train_args.output_dir = tempdir
+        model_args = copy.deepcopy(MODEL_ARGS)
+
+        sft_trainer.train(
+            model_args,
+            DATA_ARGS,
+            train_args,
+            PEFT_PT_ARGS,
+            additional_callbacks=[TrainerCallback()],
+        )