Merge remote-tracking branch 'upstream/main'

red-hat-data-services · Aug 17, 2024 · 610a3da · 610a3da
2 parents fd86b4f + a6d093e
commit 610a3da
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -468,7 +468,7 @@ def test_run_causallm_ft_save_with_save_model_dir_save_strategy_no():
         # validate that no checkpoints created
         assert not any(x.startswith("checkpoint-") for x in os.listdir(tempdir))
 
-        sft_trainer.save(tempdir, trainer)
+        sft_trainer.save(tempdir, trainer, "debug")
         assert any(x.endswith(".safetensors") for x in os.listdir(tempdir))
         _test_run_inference(checkpoint_path=tempdir)
 

diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py
@@ -373,13 +373,15 @@ def save(path: str, trainer: SFTTrainer, log_level="WARNING"):
             Path to save the model to.
         trainer: SFTTrainer
             Instance of SFTTrainer used for training to save the model.
+        log_level: str
+            Optional threshold to set save save logger to, default warning.
     """
     logger = logging.getLogger("sft_trainer_save")
     # default value from TrainingArguments
     if log_level == "passive":
         log_level = "WARNING"
 
-    logger.setLevel(log_level)
+    logger.setLevel(log_level.upper())
 
     if not os.path.exists(path):
         os.makedirs(path, exist_ok=True)