ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

RL metrics

Browse source
This commit is contained in:
Daniel Han 2025-02-05 14:59:01 -08:00
commit 1d58638f5e
2 changed files with 48 additions and 132 deletions
Download patch file Download diff file Expand all files Collapse all files

113
unsloth/models/dpo.py
View file

@ -17,115 +17,8 @@ __all__ = [
"PatchKTOTrainer",
]
try:
from transformers.utils.notebook import (
IntervalStrategy,
NotebookTrainingTracker,
NotebookProgressCallback,
)
HAS_NOTEBOOK = True
except:
HAS_NOTEBOOK = False
pass
import torch
from ._utils import torch_compile_options
import inspect
import torch.nn as nn
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
from .rl import PatchRLStatistics
def PatchDPOTrainer(): PatchRLStatistics("DPO")
DPOTrainer_metrics = [
"rewards/chosen",
"rewards/rejected",
"rewards/accuracies",
"rewards/margins",
"logps/rejected",
"logps/chosen",
"logits/rejected",
"logits/chosen",
]
set_DPOTrainer_metrics = frozenset(DPOTrainer_metrics)
def NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs):
self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
self.training_loss = 0
self.last_log = 0
column_names = [self.first_column] + ["Training Loss"]
if args.eval_strategy != IntervalStrategy.NO:
column_names.append("Validation Loss")
column_names += [x.replace("/", " / ") for x in DPOTrainer_metrics]
self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
pass
def NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs):
# Only for when there is no evaluation
if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
values = {"Training Loss": logs["loss"]}
for metric in DPOTrainer_metrics:
values[metric.replace("/", " / ")] = logs[metric]
pass
# First column is necessarily Step since we're not in epoch eval strategy
values["Step"] = state.global_step
self.training_tracker.write_line(values)
pass
pass
def NotebookTrainingTracker_write_line(self, values):
"""
Write the values in the inner table.
Args:
values (`Dict[str, float]`): The values to display.
"""
if self.inner_table is None:
self.inner_table = [list(values.keys()), list(values.values())]
else:
columns = self.inner_table[0]
new_values = {}
for key, value in values.items():
lowered = key.lower()
if lowered in set_DPOTrainer_metrics:
new_values[lowered.replace("/", " / ")] = value
else:
new_values[key] = value
pass
values = new_values
self.inner_table[0] = columns
if len(self.inner_table) > 1:
last_values = self.inner_table[-1]
first_column = self.inner_table[0][0]
if last_values[0] != values[first_column]:
# write new line
self.inner_table.append([values[c] if c in values else "No Log" for c in columns])
else:
# update last line
new_values = values
for c in columns:
if c not in new_values.keys():
new_values[c] = last_values[columns.index(c)]
self.inner_table[-1] = [new_values[c] for c in columns]
else:
# Edit for evaluation purposes
self.inner_table.append([values[c] if c in values else 0 for c in columns])
pass
pass
pass
def PatchDPOTrainer():
if HAS_NOTEBOOK:
from transformers.trainer import is_in_notebook
if is_in_notebook():
# Patch DPO notebook printing
NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line
from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin
DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log
pass
pass
pass
PatchKTOTrainer = PatchDPOTrainer
def PatchKTOTrainer(): PatchRLStatistics("KTO")

67
unsloth/models/rl.py
View file

@ -29,7 +29,10 @@ except:
HAS_NOTEBOOK = False
pass
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
import inspect
import os
import re
import functools
def PatchRL(FastLanguageModel):
@ -94,7 +97,7 @@ def NotebookProgressCallback_on_log(Trainer_metrics):
# Only for when there is no evaluation
if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
values = {"Training Loss": logs["loss"]}
for metric in DPOTrainer_metrics:
for metric in Trainer_metrics:
values[metric.replace("/", " / ")] = logs[metric]
pass
# First column is necessarily Step since we're not in epoch eval strategy
@ -167,27 +170,47 @@ def _PatchRLStatistics(metrics):
pass
@functools.cache
def get_trl_metrics():
# Gets metrics so we can output them in notebooks
import trl.trainer
trainers = dir(trl.trainer)
trainers = [x for x in trainers if x.endswith("_trainer")]
filepath = inspect.getfile(trl.trainer)
filepath = os.path.split(filepath)[0]
all_metrics = dict()
for trainer in trainers:
filename = os.path.join(filepath, f"{trainer}.py")
if not os.path.exists(filename): continue
with open(filename, "r") as file: file = file.read()
# Get metrics['kl'] or stats['kl']
metrics = re.findall(r"metrics\[[\"\']([^\"\']{1,})[\"\']\]", file)
stats = re.findall(r"stats\[[\"\']([^\"\']{1,})[\"\']\]", file)
metrics = metrics + stats
# Get optional f-strings
metrics_f = re.findall(r"metrics\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file)
stats_f = re.findall(r"stats\[f[\"\']\{[^\}]{1,}\}([^\"\']{1,})[\"\']\]", file)
metrics_f = metrics_f + stats_f
# Filter out prefixes if seen
# metrics[f"{prefix}rewards/chosen"]
left_prefix = 'prefix = "eval_" if train_eval == "eval" else ""' in file
if left_prefix: metrics += metrics_f
all_metrics[trainer[:trainer.find("_")].upper()] = metrics
pass
return all_metrics
pass
def PatchRLStatistics(algorithm = "GRPO"):
algorithm = algorithm.upper()
if algorithm == "GRPO":
metrics = [
"completion_length",
"reward",
"reward_std",
"kl",
]
elif algorithm == "DPO" or algorithm == "KTO":
metrics = [
"rewards/chosen",
"rewards/rejected",
"rewards/accuracies",
"rewards/margins",
"logps/rejected",
"logps/chosen",
"logits/rejected",
"logits/chosen",
]
else:
all_metrics = get_trl_metrics()
if algorithm not in all_metrics:
print(f"Unsloth for {algorithm.upper()} is not yet implemented! Just ignore this function.")
_PatchRLStatistics(metrics)
pass
_PatchRLStatistics(all_metrics[algorithm])
pass