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

Remove reload_weights rpc call from grpo trainer

Browse source 
This commit removes the `reload_weights` RPC call from `_generate_single_turn` in `GRPOTrainer` via Unsloth's runtime patching mechanism. This call, introduced in recent `trl` versions, causes a `ValueError` during training when using Unsloth (specifically with LoRA) because it attempts to reload weights that are not fully initialized or managed by vLLM in the expected way.

The fix uses `re.sub` to robustly remove the offending line from the source code at runtime.

Fixes #3673
This commit is contained in:
google-labs-jules[bot] 2025-12-06 02:50:00 +00:00
commit 39bdd266ff
1 changed files with 18 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/models/rl_replacements.py
View file

@ -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":