[flyingdoubleG]

Description

KLFn.apply_kl_penalty_to_reward (trinity/algorithm/kl_fn/kl_fn.py) reads the reference log probs before checking that they exist:

logprob = experiences.batch["old_log_probs"]
ref_logprob = experiences.batch["ref_log_prob"]      # KeyError if absent

if "ref_log_prob" in experiences.batch.keys():
    ...
else:                                                 # unreachable
    kl_coef = 0.0
    kl = torch.zeros_like(response_mask, dtype=torch.float32)
    experiences.batch["token_level_rewards"] = token_level_scores

So the intended fallback (rewards unchanged, kl_coef 0) never runs; a batch without ref_log_prob raises KeyError instead. The trainer calls this for every algorithm with compute_advantage_in_trainer, including ones that do not use a reference model (on_policy_distill, jsd), so this happens when a KL penalty function is configured for them.

The change moves the two reads inside the check; behaviour with ref_log_prob present is unchanged.

Tests: tests/algorithm/kl_fn_test.py::test_kl_penalty_to_reward_without_reference (fails before with KeyError: 'ref_log_prob') and ::test_kl_penalty_to_reward_with_reference; the file passes (13 tests, CPU).

Checklist

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

🤖 Generated with Claude Code