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

Generic efficient GRPO

Browse source
This commit is contained in:
Daniel Han 2025-06-26 00:10:07 -07:00
commit c4283d771d
2 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/llama.py
View file

@ -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,

7
unsloth/models/rl_replacements.py
View file

@ -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'):