This commit is contained in:
Daniel Han 2025-05-28 05:56:19 -07:00
commit 59f185eb2f

View file

@ -363,13 +363,27 @@ RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size)
def grpo_trainer_metrics(RLTrainer_source, RLConfig_source):
if "reward_funcs" not in RLTrainer_source: return ""
# For new TRL we have /mean and /std
use_mean = "rewards/{reward_func_name}/mean" in RLTrainer_source
use_std = "rewards/{reward_func_name}/std" in RLTrainer_source
if not use_mean:
use_normal = "rewards/{reward_func_name}" in RLTrainer_source
else:
use_normal = False
pass
log_metrics = \
"if not isinstance(reward_funcs, list): _reward_funcs = [reward_funcs]\n"\
"else: _reward_funcs = reward_funcs\n"\
"for reward_func in _reward_funcs:\n"\
" try:\n"\
" reward_func_name = reward_func.__name__\n"\
" other_metrics.append(f'rewards/{reward_func_name}')\n"\
f" if {use_mean}:\n"\
" other_metrics.append(f'rewards/{reward_func_name}/mean')\n"\
f" if {use_std}:\n"\
" other_metrics.append(f'rewards/{reward_func_name}/std')\n"\
f" if {use_normal}:\n"\
" other_metrics.append(f'rewards/{reward_func_name}')\n"\
" except: pass\n"
return log_metrics
pass