Skip to content

Fix VLM GRPO matmul shape mismatch in grpo_accumulated_loss #546

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
danielhanchen merged 4 commits into from
Mar 15, 2026
Merged

Fix VLM GRPO matmul shape mismatch in grpo_accumulated_loss #546

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
eb1d34b
Fix VLM GRPO matmul shape mismatch in chunked_hidden_states_selective…
danielhanchen Mar 15, 2026
96dce6c
Add supports_tower_connector_lora to WorkerLoRAManager
danielhanchen Mar 15, 2026
3d08551
Remove redundant scaling in logits fallback path
danielhanchen Mar 15, 2026
3d22fa3
Move temperature into chunked_selective_log_softmax to avoid full flo…
danielhanchen Mar 15, 2026
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
42 changes: 30 additions & 12 deletions unsloth_zoo/rl_replacements.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -940,6 +940,17 @@ def efficient_log_softmax(hidden_states, lm_head, index, chunks=32,

new_hidden_states_chunk = new_hidden_states_chunk[:, -(logits_to_keep + max_left_pad + 1): , :]
new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :]
logprobs_chunk = efficient_log_softmax(
new_hidden_states_chunk,
lm_head,
completion_ids,
chunks=input_ids_chunk.shape[0]*multiplier,
logit_scale_multiply=logit_scale_multiply,
logit_scale_divide=logit_scale_divide,
logit_softcapping=logit_softcapping,
temperature=temperature,
batch_size = B
)
else:
new_hidden_states_chunk = unwrapped_model(
input_ids = input_ids_chunk,
Expand All @@ -952,18 +963,25 @@ def efficient_log_softmax(hidden_states, lm_head, index, chunks=32,
).logits

new_hidden_states_chunk = new_hidden_states_chunk[:, :-1, :]

logprobs_chunk = efficient_log_softmax(
new_hidden_states_chunk,
lm_head,
completion_ids,
chunks=input_ids_chunk.shape[0]*multiplier,
logit_scale_multiply=logit_scale_multiply,
logit_scale_divide=logit_scale_divide,
logit_softcapping=logit_softcapping,
temperature=temperature,
batch_size = B
)
# Guard: check if model returned hidden states or logits
if new_hidden_states_chunk.shape[-1] == lm_head.shape[1]:
logprobs_chunk = efficient_log_softmax(
new_hidden_states_chunk,
lm_head,
completion_ids,
chunks=input_ids_chunk.shape[0]*multiplier,
logit_scale_multiply=logit_scale_multiply,
logit_scale_divide=logit_scale_divide,
logit_softcapping=logit_softcapping,
temperature=temperature,
batch_size = B
)
else:
# Model returned logits directly - scaling/softcapping already applied by model forward
new_hidden_states_chunk = new_hidden_states_chunk.to(torch.float32)
if temperature != 1.0:
new_hidden_states_chunk = new_hidden_states_chunk / temperature
logprobs_chunk = chunked_selective_log_softmax(new_hidden_states_chunk, completion_ids)

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Apply logit scaling and softcapping in logits fallback

In the new VLM fallback branch, when new_hidden_states_chunk is already logits, the code only applies temperature before chunked_selective_log_softmax, but it skips logit_scale_multiply, logit_scale_divide, and logit_softcapping that are still applied in efficient_log_softmax. For runs where any of these kwargs are non-zero, this silently changes the computed token log-probabilities and therefore the GRPO loss for the fallback path, leading to inconsistent training behavior across models.

Useful? React with 👍 / 👎.

#This is needed to avoid race conditions with GPT OSS offload_embbed=True
#However, it seems that this line does not slow down or disrupt models.
device_synchronize()
Expand Down
8 changes: 8 additions & 0 deletions unsloth_zoo/vllm_lora_worker_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,14 @@ def set_active_adapters(self, requests: Set[Any],
if mapping is not None:
self._adapter_manager.set_adapter_mapping(mapping)

def supports_tower_connector_lora(self) -> bool:
manager = getattr(self, '_adapter_manager', None)
if manager is None:
return False
return (
getattr(manager, 'supports_mm', False)
and getattr(manager, 'supports_tower_connector_lora', False)
)

def _apply_adapters(self, adapter_requests: Set[Any]) -> None:
if apply_adapters_worker:
Expand Down
Loading