From 9741989886bd18a1d92ada4ffa5c74c75fbfb1e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Jun 2025 01:12:21 -0700 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 9c3d219c79..597ef5b8d3 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -261,13 +261,13 @@ def grpo_trainer__get_per_token_logps(function_name, function): os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype): # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded - hidden_states = model( + logits = 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 - return hidden_states + return logits # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770