[GRPO] Try returning hidden statex for GRPO

[Datta0]

unslothai/unsloth-zoo#602 fixed an important issue where when some models return logits we were failing with shape mismatch. This is because for GRPO we generally expect hidden states to be returned with our wrappers (for most models ofc) adn lm_head is applied chunk wise to avoid materialising full large logits which are much larger than hidden states (4K vs 256K ish for eg)
This is an effort to make more models return hidden states for efficiency reasons. Orthogonal to the above mentioned PR :)

Ref: unslothai/unsloth-zoo#609

[gemini-code-assist]

Code Review

This pull request implements a fallback mechanism for GRPO to retrieve hidden states from models that do not natively support it by wrapping the forward method. Feedback suggests optimizing performance by pre-computing function signatures during installation and using them to robustly handle positional arguments. It was also recommended to log exceptions during argument binding to aid debugging.

[gemini-code-assist]

The _get_num_logits_to_keep function currently only checks kwargs. If logits_to_keep or num_logits_to_keep are passed as positional arguments, they will be missed. It's better to use the model's signature to bind the arguments and extract the values robustly. Per repository rules, ensure exceptions are logged rather than silently ignored.

Suggested change

	def _get_num_logits_to_keep(kwargs):
	num_logits_to_keep = kwargs.get("num_logits_to_keep", 0) or 0
	logits_to_keep = kwargs.get("logits_to_keep", 0) or 0
	def _get_num_logits_to_keep(sig, args, kwargs):
	try:
	bound = sig.bind_partial(*args, **kwargs)
	return max(bound.arguments.get("num_logits_to_keep", 0) or 0,
	bound.arguments.get("logits_to_keep", 0) or 0)
	except Exception as e:
	import logging
	logging.debug(f"Error binding signature: {e}")
	return max(kwargs.get("num_logits_to_keep", 0) or 0,
	kwargs.get("logits_to_keep", 0) or 0)

References

1. Avoid using broad, silent exception handlers like except Exception: pass. Instead, log the exception, even if at a debug level, to aid in future debugging.

[gemini-code-assist]

Use the pre-computed signature and original arguments to correctly extract num_logits_to_keep, accounting for positional arguments.

Suggested change

	hidden_states = hidden_states[-1]
	num_logits_to_keep = _get_num_logits_to_keep(sig, args, kwargs)

[gemini-code-assist]

To improve performance, _drop_forward_kwargs_consumed_positionally should accept a pre-computed signature instead of calling inspect.signature() on every forward pass. This avoids redundant introspection overhead during training.

Suggested change

	def _drop_forward_kwargs_consumed_positionally(sig, args, kwargs):

References

1. To improve efficiency, avoid redundant data iterations. Combine checks and transformations into a single loop and return computed values for callers to reuse.

[gemini-code-assist]

Use the passed signature instead of re-computing it.

Suggested change

	consumed_names = []
	for parameter in sig.parameters.values():

[gemini-code-assist]

Pre-compute the signature once during installation to avoid the overhead of inspect.signature() on every forward pass.

Suggested change

	return False
	original_forward = target_model.forward
	sig = inspect.signature(original_forward)
	model_name = type(target_model).__name__

References

1. To improve efficiency, avoid redundant data iterations. Combine checks and transformations into a single loop and return computed values for callers to reuse.

[gemini-code-assist]

Pass the pre-computed signature to the helper function.

Suggested change

	return original_forward(*args, **kwargs)
	forward_kwargs = _drop_forward_kwargs_consumed_positionally(
	forward_kwargs = _drop_forward_kwargs_consumed_positionally(
	sig, args, kwargs
	)