ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove debugging

Browse source
This commit is contained in:
Daniel Han 2025-06-26 00:41:37 -07:00
commit f41bfbc092
1 changed files with 9 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

38
unsloth/models/rl_replacements.py
View file

@ -261,11 +261,12 @@ def grpo_trainer__get_per_token_logps(function_name, function):
os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
with torch.amp.autocast(device_type = 'cuda', dtype = self._autocast_dtype):
# We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
print("input_ids Unsloth 264", input_ids.shape)
print("logits_to_keep Unsloth 264", logits_to_keep)
hidden_states = model(input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=logits_to_keep + 1).logits
print("hidden_states Unsloth 264", hidden_states.shape)
#logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
hidden_states = model(
input_ids = input_ids,
attention_mask = attention_mask,
logits_to_keep = logits_to_keep + 1,
).logits
# logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
return hidden_states
# input_ids = input_ids[:, -logits_to_keep:]
# For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves.
@ -318,14 +319,8 @@ def grpo_trainer_compute_loss(function_name, function):
logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens
_input_ids = input_ids
_logits_to_keep = logits_to_keep
print("prompt_mask Unsloth 320", prompt_mask.shape)
print("completion_mask Unsloth 320", completion_mask.shape)
print("input_ids Unsloth 320", input_ids.shape)
print("logits_to_keep Unsloth 320", logits_to_keep)
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
if per_token_logps is not None:
print("per_token_logps Unsloth 320", per_token_logps.shape)
# Compute the KL divergence between the model and the reference model
# _prepare_inputs doesn't return reference log probs anymore. We need to calculate it ourselves.
@ -333,25 +328,16 @@ def grpo_trainer_compute_loss(function_name, function):
if self.beta != 0.0:
with torch.inference_mode(), model.disable_adapter():
ref_per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
if ref_per_token_logps is not None:
print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
else:
ref_per_token_logps = None
# per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
# x - x.detach() allows for preserving gradients from x
advantages = inputs["advantages"]
print("advantages Unsloth 320", advantages.shape)
# per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
# per_token_loss = -(per_token_loss - self.beta * per_token_kl)
# loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
old_hidden_states = inputs.get("old_per_token_logps", None)
if old_hidden_states is not None:
print("old_hidden_states Unsloth 320", old_hidden_states.shape)
print("input_ids Unsloth 320", input_ids.shape)
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) # Gemma
@ -365,14 +351,9 @@ def grpo_trainer_compute_loss(function_name, function):
if per_token_logps is not None:
if ref_per_token_logps is not None:
print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
ref_per_token_logps = ref_per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
print("ref_per_token_logps Unsloth 320", ref_per_token_logps.shape)
print("per_token_logps Unsloth 320", per_token_logps.shape)
per_token_logps = per_token_logps[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
print("per_token_logps Unsloth 320", per_token_logps.shape)
loss, completion_length, mean_kl = grpo_compute_loss_slow(
ref_per_token_logps,
per_token_logps,
@ -428,13 +409,12 @@ def grpo_trainer_compute_loss(function_name, function):
logit_scale_divide = logit_scale_divide,
attention_mask = attention_mask,
)
pass
pass
# Log the metrics
# completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
# mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
# self._metrics["kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
if "train" in self._metrics:
mode = "eval" if self.control.should_evaluate else "train"
self._metrics[mode]["completion_length"].append(completion_length.item())