diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py index 51f1c9a63a..9c12abb98f 100644 --- a/unsloth/models/dpo.py +++ b/unsloth/models/dpo.py @@ -17,8 +17,8 @@ __all__ = [ "PatchKTOTrainer", ] -from .rl import PatchRLStatistics +from .rl import PatchFastRL -def PatchDPOTrainer(): PatchRLStatistics("DPO") +def PatchDPOTrainer(): PatchFastRL("DPO") -def PatchKTOTrainer(): PatchRLStatistics("KTO") +def PatchKTOTrainer(): PatchFastRL("KTO") diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 02bc10c6fe..40d68f6a74 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -13,8 +13,7 @@ # limitations under the License. __all__ = [ - "PatchRL", - "PatchRLStatistics", + "PatchFastRL", ] import torch @@ -202,6 +201,9 @@ def get_trl_metrics(): left_prefix = 'prefix = "eval_" if train_eval == "eval" else ""' in file if left_prefix: metrics += metrics_f + # Remove all eval_ things + metrics = [x for x in metrics if not x.startswith("eval_")] + all_metrics[trainer[:trainer.find("_")].upper()] = metrics pass return all_metrics @@ -219,3 +221,9 @@ def PatchRLStatistics(algorithm = "GRPO"): pass _PatchRLStatistics(all_metrics[algorithm], algorithm) pass + + +def PatchFastRL(algorithm = "GRPO", FastLanguageModel = None): + if FastLanguageModel is not None: PatchRL(FastLanguageModel) + PatchRLStatistics(algorithm) +pass