diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 38973fe30e..aa53ed21fa 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -216,6 +216,24 @@ def grpo_trainer__prepare_inputs(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__prepare_inputs) +# Remove reload_weights rpc call +def grpo_trainer__generate_single_turn(function_name, function): + if function_name != "_generate_single_turn": + return function + + # Remove reload_weights rpc call + function = re.sub( + r"^\s*self\.llm\.collective_rpc\([\"']reload_weights[\"']\)\s*\n", + "", + function, + flags = re.MULTILINE, + ) + return function + + +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_single_turn) + + # Fix incorrect special tokens handling and truncation in older TRL versions def grpo_trainer__generate_and_score_completions(function_name, function): if function_name != "_generate_and_score_completions":