From 705fd769a31639b480c704049442c27914a7641d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 14 Feb 2025 04:44:03 -0800 Subject: [PATCH] Update rl.py --- unsloth/models/rl.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 1ac511e833..128725a0a9 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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,