From 30c59b1e34cdb8e217bd32ffcf1e4f963319c1d5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 17 Feb 2025 00:05:32 -0800 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 99dba9b9a3..0f1c81bb8d 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -233,15 +233,15 @@ def grpo_trainer_compute_loss(function_name, function): loss, completion_length, mean_kl = grpo_compute_loss( ref_per_token_logps, per_token_logps, input_ids, completion_mask, self.beta, advantages, ) + RL_REPLACEMENTS["data"] = ( + ref_per_token_logps.detach(), per_token_logps.detach(), _input_ids, completion_mask, self.beta, advantages, + loss.detach(), completion_length, mean_kl, completion_ids, _logits_to_keep, + ) from unsloth_zoo.rl_replacements import RL_REPLACEMENTS if "count" in RL_REPLACEMENTS: RL_REPLACEMENTS["count"] += 1 if RL_REPLACEMENTS["count"] == 10: raise else: RL_REPLACEMENTS["count"] = 1 - RL_REPLACEMENTS["data"] = ( - ref_per_token_logps, per_token_logps.detach(), _input_ids, completion_mask, self.beta, advantages, - loss, completion_length, mean_kl, completion_ids, _logits_to_keep, - ) # Log the metrics # completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item() self._metrics["completion_length"].append(completion_length.item())