Update rl.py

This commit is contained in:
Daniel Han 2025-03-05 23:17:51 -08:00
commit 894ea97496

View file

@ -284,6 +284,18 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
extra_args += eval_changes
pass
# Force logits to be produced if preprocess_logits_for_metrics or compute_metrics is used
if "model" in call_args:
logits_check = \
"_output_logits = False"\
"if locals().get('compute_metrics', None) is not None: _output_logits = True\n"\
"if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True\n"\
"if _output_logits:\n"\
" import os\n"\
" os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"
extra_args += logits_check
pass
# Check max_seq_length
if "model" in call_args:
length_check = \