From 83a2d4282c50c8a7f51386f2c6ac13ef9edeed51 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 12 Feb 2025 18:50:45 -0800 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0f342ec86e..781f5984d4 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -146,3 +146,17 @@ def grpo_trainer__move_model_to_vllm(function_name, function): return " "*function.find("def") + replacement pass RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__move_model_to_vllm) + + +# Edit _get_per_token_logps +def grpo_trainer__get_per_token_logps(function_name, function): + if function_name != "_get_per_token_logps": return function + + # Set attention_mask to boolean + function = function.replace( + "attention_mask=attention_mask", + "attention_mask=attention_mask.to(torch.bool)" + ) + return function +pass +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps)