From 07a6d7f7737494fcebad99e4f9de5a4f6b214f94 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Jun 2025 18:50:45 -0700 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 38258e4fad..9afe1a49cc 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -352,6 +352,13 @@ def grpo_trainer_compute_loss(function_name, function): print("logits_to_keep Unsloth 320", logits_to_keep) input_ids = input_ids[:, -logits_to_keep:] print("input_ids Unsloth 320", input_ids.shape) + + # Get logit softcapping and logit scale + logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) + if logit_softcapping is None: logit_softcapping = 0 + logit_scale_multiply = getattr(model.config, "logit_scale", 0) + if logit_scale_multiply is None: logit_scale_multiply = 0 + if per_token_logps is not None: if ref_per_token_logps is not None: @@ -377,6 +384,8 @@ def grpo_trainer_compute_loss(function_name, function): max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, ) else: if hasattr(self.args, "loss_type"): @@ -394,6 +403,8 @@ def grpo_trainer_compute_loss(function_name, function): max_completion_length = self.args.max_completion_length, delta = self.args.delta, temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, ) else: # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17 @@ -406,6 +417,8 @@ def grpo_trainer_compute_loss(function_name, function): old_hidden_states, n_chunks = self.args.unsloth_num_chunks, temperature = self.args.temperature, + logit_softcapping = logit_softcapping, + logit_scale_multiply = logit_scale_multiply, ) # Log the metrics