From c4283d771dabb423afb77d3735c631ec02c77911 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 26 Jun 2025 00:10:07 -0700 Subject: [PATCH] Generic efficient GRPO --- unsloth/models/llama.py | 1 + unsloth/models/rl_replacements.py | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9db8abdd43..d0e3aeb144 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1131,6 +1131,7 @@ def CausalLM_fast_forward(fast_forward_inference): if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": if num_logits_to_keep != 0: hidden_states = hidden_states[:, -num_logits_to_keep:, :] + hidden_states.__is_hidden_state = True return CausalLMOutputWithPast( loss = None, logits = hidden_states, diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index ebc3549273..0056f27509 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -248,10 +248,11 @@ RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__move_model_to_vllm) # Edit _get_per_token_logps to handle mixed precision def grpo_trainer__get_per_token_logps(function_name, function): - if function_name != "_get_per_token_logps": return function + if function_name != "_get_per_token_logps": return function - def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep, calc_logprob_flag = None): - if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0' and not calc_logprob_flag: + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep): + if True: #os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0': + print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") return None # Unsloth efficient GRPO # Otherwise, calculate normally: if not hasattr(self, '_autocast_dtype'):