From 11769cf006ef349fdd451cf777b213614e259792 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Tue, 2 Jun 2026 10:59:09 +0000 Subject: [PATCH] Use context instead of manual setting for UNSLOTH_RETURN_HIDDEN_STATES --- unsloth/models/rl.py | 66 +++++++++++++++++++++++++++++-- unsloth/models/rl_replacements.py | 31 +++++++++++---- 2 files changed, 87 insertions(+), 10 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 82ef7b48d2..0768915a73 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -597,10 +597,14 @@ def _grpo_owns_lm_head(module): def _grpo_causal_head(model): # The module that owns lm_head / output embeddings (whose forward emits `.logits`). + if model is None: + return None + if _grpo_owns_lm_head(model): + return model get_base_model = getattr(model, "get_base_model", None) if callable(get_base_model): base_model = get_base_model() - if base_model is not None: + if base_model is not None and _grpo_owns_lm_head(base_model): return base_model return model @@ -608,6 +612,8 @@ def _grpo_causal_head(model): def _grpo_hidden_states_wrap_target(model): if model is None: return None + if _grpo_owns_lm_head(model): + return model get_base_model = getattr(model, "get_base_model", None) if callable(get_base_model): base_model = get_base_model() @@ -624,6 +630,53 @@ def _grpo_hidden_states_wrap_target(model): return model +def _grpo_config_value(config, name, default = None): + if config is None: + return default + value = getattr(config, name, default) + if value is not default: + return value + text_config = getattr(config, "text_config", None) + if text_config is not None: + value = getattr(text_config, name, default) + if value is not default: + return value + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_config = get_text_config() + return getattr(text_config, name, default) + except Exception: + pass + return default + + +def _grpo_has_post_lm_head_transform(module): + config = getattr(module, "config", None) + if config is None: + return False + + final_logit_softcapping = _grpo_config_value( + config, "final_logit_softcapping", None + ) + if final_logit_softcapping not in (None, 0, 0.0): + return True + + logit_scale = _grpo_config_value(config, "logit_scale", None) + if logit_scale not in (None, 0, 0.0, 1, 1.0): + return True + + logits_scaling = _grpo_config_value(config, "logits_scaling", None) + if logits_scaling not in (None, 0, 0.0, 1, 1.0): + return True + + lm_head_multiplier = _grpo_config_value(config, "lm_head_multiplier", None) + if lm_head_multiplier not in (None, 0, 0.0, 1, 1.0): + return True + + return False + + def _model_supports_unsloth_return_hidden_states(model): target_model = _grpo_hidden_states_wrap_target(model) for candidate in (model, target_model): @@ -724,6 +777,8 @@ def _install_grpo_lm_head_passthrough(model): # weight is untouched, and the accelerate-managed top-level forward is not wrapped, so there is # no bound-self collision. No-op when the flag is 0. head = _grpo_causal_head(model) + if _grpo_has_post_lm_head_transform(head): + return False lm_head = getattr(head, "lm_head", None) if lm_head is None: get_output_embeddings = getattr(head, "get_output_embeddings", None) @@ -738,9 +793,14 @@ def _install_grpo_lm_head_passthrough(model): original_lm_head_forward = lm_head.forward def passthrough_forward(*args, **kwargs): + forward_args = args[1:] if len(args) > 0 and args[0] is lm_head else args if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": - return args[0] if args else next(iter(kwargs.values())) - return original_lm_head_forward(*args, **kwargs) + if len(forward_args) > 0: + return forward_args[0] + if len(kwargs) > 0: + return next(iter(kwargs.values())) + raise TypeError("forward() missing 1 required positional argument: 'input'") + return original_lm_head_forward(*forward_args, **kwargs) lm_head.forward = passthrough_forward lm_head._unsloth_grpo_passthrough = True diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 31d54675c9..a0e088af49 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -25,6 +25,7 @@ import re import torch import inspect import linecache +from contextlib import contextmanager from collections import defaultdict from unsloth_zoo.rl_replacements import ( RL_REPLACEMENTS, @@ -74,6 +75,19 @@ RL_CONFIG_CHANGES = defaultdict(list) RL_METRICS_CHANGES = defaultdict(list) RL_ADDITIONAL_FUNCTIONS = defaultdict(list) + +@contextmanager +def _temporary_unsloth_return_hidden_states(): + old_value = os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES") + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + try: + yield + finally: + if old_value is None: + os.environ.pop("UNSLOTH_RETURN_HIDDEN_STATES", None) + else: + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = old_value + _DPO_VISION_KEYS = ( "pixel_position_ids", "image_position_ids", @@ -1100,8 +1114,9 @@ def grpo_trainer__get_per_token_logps(function_name, function): if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": self._autocast_dtype = torch.float16 - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - with torch.amp.autocast(device_type = DEVICE_TYPE, dtype = self._autocast_dtype): + with _temporary_unsloth_return_hidden_states(), torch.amp.autocast( + device_type = DEVICE_TYPE, dtype = self._autocast_dtype + ): # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded logits = model( input_ids = input_ids, @@ -1379,9 +1394,9 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): token_type_ids_chunks, mm_token_type_ids_chunks, ) - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" - - with _get_inference_mode_context_manager(model): + with _temporary_unsloth_return_hidden_states(), _get_inference_mode_context_manager( + model + ): for ( input_ids_chunk, attention_mask_chunk, @@ -1478,8 +1493,6 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): logprobs = torch.cat(all_logprobs_list, dim = 0) entropies = None - os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" - return logprobs.detach(), entropies # logps, entropies # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. @@ -1536,6 +1549,10 @@ grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] RL_PRE_ITEMS["grpo_trainer"].append( inspect.getsource(_unsloth_get_final_logit_softcapping) ) +RL_PRE_ITEMS["grpo_trainer"].append("from contextlib import contextmanager") +RL_PRE_ITEMS["grpo_trainer"].append( + inspect.getsource(_temporary_unsloth_return_hidden_states) +) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_clear_stateful_mrope))