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

Update rl.py

Browse source
This commit is contained in:
Daniel Han 2025-02-14 04:44:03 -08:00
commit 705fd769a3
1 changed files with 3 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/models/rl.py
View file

@ -86,7 +86,7 @@ pass
# https://github.com/huggingface/trl/blob/main/trl/trainer/utils.py#L1674
@torch.compile(dynamic = True, fullgraph = True, options = torch_compile_options,)
def _selective_log_softmax(logits, index):
def selective_log_softmax(logits, index):
logits = logits.to(torch.float32)
selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
# loop to reduce peak mem consumption
@ -96,10 +96,6 @@ def _selective_log_softmax(logits, index):
return per_token_logps
pass
def selective_log_softmax(logits, index):
return _selective_log_softmax(logits, index)
pass
RLTrainer_replacement = '''
import os
@ -423,10 +419,8 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
pass
# Selective log softmax
selective_log_softmax_code = \
inspect.getsource(_selective_log_softmax) + "\n" + \
inspect.getsource(selective_log_softmax) + "\n"
selective_log_softmax_code = inspect.getsource(selective_log_softmax)
# Get final source code
RLTrainer_source = RLTrainer_replacement.format(
RLTrainer_name = RLTrainer_name,