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

Update rl_replacements.py

Browse source
This commit is contained in:
Daniel Han 2025-06-22 03:44:56 -07:00
commit f8484e7a19
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/rl_replacements.py
View file

@ -263,7 +263,7 @@ def grpo_trainer__get_per_token_logps(function_name, function):
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
hidden_states = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
#logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
print("##############", hidden_states.shape)
print("##############, hidden_states", hidden_states.shape)
return hidden_states
# input_ids = input_ids[:, -logits_to_keep:]
# For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
@ -336,6 +336,7 @@ def grpo_trainer_compute_loss(function_name, function):
else:
old_hidden_states = None
print("$$$$$$$$$$$ input_ids", input_ids.shape)
input_ids = input_ids[:, -logits_to_keep:]
if per_token_logps is not None: