Skip to content

Fix AttributeError in GRPO trainer for models without llm attribute #2780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 21, 2025
Merged

Fix AttributeError in GRPO trainer for models without llm attribute #2780

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8a71802
Additional tests for unsloth-zoo PR#174
rolandtannous Jun 21, 2025
88b2a9c
Fix Gemma3ForCausalLm does not have attribute self.llm
rolandtannous Jun 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 13 additions & 10 deletions unsloth/models/rl_replacements.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,8 +184,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest)
rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest)
insert = (
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.wake_up()\n"
" if hasattr(self, 'llm'):\n"
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.wake_up()\n"
)
function = function[:sig_end] + insert + rest
else:
Expand All @@ -199,8 +200,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
rest = re.sub(r"^[ \t]*free, total = torch.cuda.mem_get_info\(\)\s*\n", "", rest)
rest = re.sub(r"^[ \t]*print\(f?\".*cuda.*\"\)\s*\n", "", rest)
insert = (
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.wake_up()\n"
" if (hasattr(self, 'llm'):\n"
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.wake_up()\n"
)
function = header_and_comments + insert + rest

Expand All @@ -218,8 +220,9 @@ def grpo_trainer__prepare_inputs(function_name, function):
"self.accelerator.unwrap_model(self.model, keep_fp32_wrapper = False)",
)
sleep_and_cache = (
"if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n"
"if hasattr(self, 'llm'):\n"
" if getattr(self.llm.llm_engine.vllm_config.model_config, 'enable_sleep_mode', False):\n"
" self.llm.sleep(os.environ.get('VLLM_SLEEP_MODE', 1))\n"
" "
)
if re.search(r"\n\s*return ", function):
Expand Down Expand Up @@ -310,7 +313,7 @@ def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch
logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens
_input_ids = input_ids
_logits_to_keep = logits_to_keep

per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep)

# Compute the KL divergence between the model and the reference model
Expand All @@ -330,12 +333,12 @@ def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch
# loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
if "old_per_token_logps" in inputs.keys():
old_hidden_states = inputs["old_per_token_logps"]
else:
else:
old_hidden_states = None
input_ids = input_ids[:, -logits_to_keep:]
if per_token_logps is not None:
loss, completion_length, mean_kl = grpo_compute_loss_slow(
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
ref_per_token_logps, per_token_logps, old_hidden_states, input_ids, completion_mask, self.beta, advantages,
loss_type = self.args.loss_type,
epsilon_low = self.epsilon_low, epsilon_high = self.epsilon_high,
max_completion_length = self.args.max_completion_length,
Expand All @@ -356,7 +359,7 @@ def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch
loss, completion_length, mean_kl = grpo_accumulated_loss(
self, _input_ids, logits_to_keep, completion_mask, advantages, old_hidden_states,
n_chunks = self.args.unsloth_num_chunks,
)
)

# Log the metrics
# completion_length = self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
Expand Down