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

Merge pull request #2929 from rolandtannous/fix/fix-grpo-get-per-token-logps-argument-mismatch

Browse source 
Fix argument mismatch in GRPO _get_per_token_logps lambda function
This commit is contained in:
Michael Han 2025-07-10 14:30:05 -07:00 • committed by GitHub
commit 54916d9215
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/rl_replacements.py
View file

@ -176,7 +176,7 @@ def grpo_trainer__prepare_inputs(function_name, function):
import re
# This matches the function signature, decorators and any comments immediately following
pattern = r"(\s*@profiling_decorator\s*\n\s*def _prepare_inputs\s*\([^\)]*\)\s*(->\s*[^:]+)?\s*:\s*\n(?:[ ]*#[^\n]*\n)*)"
match = re.search(pattern, function)
insert = (
" if hasattr(self, 'llm'):\n"
@ -196,7 +196,7 @@ def grpo_trainer__prepare_inputs(function_name, function):
rest_of_function,
flags=re.DOTALL | re.MULTILINE
)
# We also need to remove the old wake up call from the beginning of the function
# since it's injected before the comments.
header_and_comments = re.sub(
@ -373,7 +373,7 @@ def grpo_trainer_compute_loss(function_name, function):
_input_ids = input_ids
_logits_to_keep = logits_to_keep
get_logps_func = lambda model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False: self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep, batch_size) if hasattr(self, "_get_per_token_logps") else self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)['logps']
get_logps_func = lambda model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False: self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) if hasattr(self, "_get_per_token_logps") else self._get_per_token_logps_and_entropies(model, input_ids, attention_mask, logits_to_keep, batch_size, compute_entropy)['logps']
per_token_logps = get_logps_func(model, input_ids, attention_mask, logits_to_keep)