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

Update rl_replacements.py

Browse source
This commit is contained in:
Daniel Han 2025-03-18 04:59:01 -07:00
commit 3851ce78f9
1 changed files with 8 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

11
unsloth/models/rl_replacements.py
View file

@ -208,8 +208,8 @@ def grpo_trainer__get_per_token_logps(function_name, function):
if function_name != "_get_per_token_logps": return function
def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep):
if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
return None # Unsloth efficient GRPO
# if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
# return None # Unsloth efficient GRPO
# Otherwise, calculate normally:
if not hasattr(self, '_autocast_dtype'):
self._autocast_dtype = torch.float16 if os.environ.get('ACCELERATE_MIXED_PRECISION', 'fp16') == 'fp16' else torch.bfloat16
@ -260,8 +260,13 @@ 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
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
if os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '1':
attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)
else:
per_token_logps = None
# Compute the KL divergence between the model and the reference model
ref_per_token_logps = inputs["ref_per_token_logps"]
# per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1