unsloth/unsloth
Daniel Han 11449208f4
Fix VLM GRPO matmul shape mismatch in _get_per_token_logps_and_entropies (#4301)
* Fix VLM GRPO matmul shape mismatch in _get_per_token_logps_and_entropies

VLM models (e.g. Qwen2.5-VL) can return logits [B*T, vocab_size] instead
of hidden states [B*T, hidden_dim] from their forward pass. When this
happens, chunked_hidden_states_selective_log_softmax tries to compute
logits @ lm_head.t() which fails with a shape mismatch.

Add a shape guard in the VLM branch of _get_per_token_logps_and_entropies:
check output.shape[-1] against lm_head.shape[1] (hidden_dim). When hidden
states are returned, the existing path is taken. When logits are returned,
scaling/softcapping/temperature are applied manually and
chunked_selective_log_softmax is used instead.

Also add chunked_selective_log_softmax to the import from unsloth_zoo.

The text-only branch (pixel_values is None) is unchanged.

Companion PR to unslothai/unsloth-zoo for grpo_accumulated_loss.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove redundant scaling in logits fallback path

When COMPILE_DISABLE=1 and the model returns logits directly, scaling
and softcapping are already applied by the model forward. Only
temperature (a GRPO training parameter) needs to be applied.

* Pass temperature to chunked_selective_log_softmax instead of manual cast

Use the new temperature parameter in chunked_selective_log_softmax
(added in companion zoo PR) to avoid casting the entire logits tensor
to float32 before the function call.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-16 03:54:16 -07:00
..
dataprep Guard optional vLLM imports when extension is broken (#4068) 2026-02-15 22:09:29 -08:00
kernels Revert "Initial changes: Refactor Attention" 2026-03-13 22:38:57 -07:00
models Fix VLM GRPO matmul shape mismatch in _get_per_token_logps_and_entropies (#4301) 2026-03-16 03:54:16 -07:00
registry Revert "[pre-commit.ci] auto fixes from pre-commit.com hooks" 2025-12-01 07:24:58 -08:00
utils Cache packed sequence metadata to reduce D2H syncs across layers (#4243) 2026-03-12 03:37:49 -07:00
__init__.py Bug fixed version 2026-03-08 06:39:19 -07:00
_auto_install.py Add PyTorch 2.10 and xformers 0.0.34 support (#3985) 2026-02-05 05:56:26 -08:00
chat_templates.py Fix regressions from security PRs #4042, #4044, and #4045 (#4062) 2026-02-15 23:16:17 -08:00
device_type.py Conditionally enable 4bit on CDNA for bitsandbytes>=v0.49.2 (#4161) 2026-03-07 01:33:40 -08:00
import_fixes.py Also patch accelerate's is_wandb_available for trl callbacks path (#4148) 2026-03-03 08:28:55 -08:00
ollama_template_mappers.py Refactor Ollama template wiring and harden packing helpers (#3890) 2026-02-09 04:04:48 -08:00
save.py Revert "Improve documentation on how to export model from Colab" 2026-03-13 22:38:41 -07:00
tokenizer_utils.py Add resilience to TRL internal API reclassification (#4111) 2026-02-25 06:34:21 -08:00
trainer.py Fix auto padding free logic to respect user passed False (#4128) 2026-03-01 19:30:47 -08:00