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

Fix for grpo_compute_loss_slow (#2702)

Browse source 
* slice last logit

* move slicing
This commit is contained in:
simpissa 2025-06-21 23:58:06 -05:00 • committed by GitHub
commit 3e75f21086
1 changed files with 2 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/rl_replacements.py
View file

@ -337,6 +337,8 @@ def grpo_trainer_compute_loss(function_name, function):
old_hidden_states = None
input_ids = input_ids[:, -logits_to_keep:]
if per_token_logps is not None:
ref_per_token_logps = ref_per_token_logps[:, :-1, :]
per_token_logps = per_token_logps[:, :-1, :]
loss, completion_length, mean_kl = grpo_compute_loss_slow(
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
loss_type = self.args.loss_type,