Guard GRPO text branch against models returning logits

[Rorical]

Problem

In grpo_accumulated_loss (unsloth_zoo/rl_replacements.py), the branch handling text-only inputs (if pixel_values is None:) passes the model's .logits output unconditionally into efficient_log_softmax, which assumes hidden states and computes hidden_states @ lm_head.t(). When a model's forward actually returns real logits (last dim = vocab) rather than pre-lm_head hidden states, this produces a shape-mismatched matmul and crashes the RL step.

The vision branch (else:, lines 986–1001) already dispatches based on a shape check:

if new_hidden_states_chunk.shape[-1] == lm_head.shape[1]:
    logprobs_chunk = efficient_log_softmax(...)
else:
    logprobs_chunk = chunked_selective_log_softmax(new_hidden_states_chunk, completion_ids, temperature)

The text branch is missing the same guard.

Repro

Training nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 (hybrid Mamba+attention MoE, text-only) with Unsloth PatchFastRL("GRPO", ...) and trl.GRPOTrainer. The model's remote modeling code returns genuine logits from .logits, and the text branch crashes with a dimension mismatch inside the compiled chunked_hidden_states_selective_log_softmax. Other text models that ship with custom modeling code can hit the same path.

Workaround in the wild: monkey-patching the compiled cache file to add the guard.

Fix

Mirror the existing vision-branch guard into the text branch. When the model returns actual logits, skip the lm_head matmul and go directly to chunked_selective_log_softmax. No behavioral change for models that follow Unsloth's usual pattern of returning hidden states through .logits.

Diff

Pure addition of the shape-check branch in the text path; no logic changes elsewhere.

[gemini-code-assist]

Code Review

This pull request introduces a guard in unsloth_zoo/rl_replacements.py to handle cases where the model returns logits directly instead of hidden states. A review comment points out that the logic for handling direct logits uses a fixed chunk size which may cause memory issues for large vocabularies and notes that scaling and softcapping parameters are currently ignored in that execution path.

[gemini-code-assist]

While this guard mirrors the logic in the vision branch, chunked_selective_log_softmax uses a fixed number of chunks (4), whereas efficient_log_softmax uses a dynamic number of chunks based on the multiplier. For models with very large vocabularies that return logits directly, the else branch might be more susceptible to OOM errors. Additionally, note that logit_scale_multiply, logit_scale_divide, and logit_softcapping are ignored in the else branch. If these parameters are expected to be applied to the logits returned by the model, they should be handled before or within the chunked_selective_log_softmax call.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c76e33eee0

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Skip scaling when logits already come from model forward

In the direct-logits fallback (else branch), this now passes logit_scale_multiply/divide/softcapping into chunked_selective_log_softmax, which re-applies those transforms to logits that may already have had them applied by the model forward. For models that ignore UNSLOTH_RETURN_HIDDEN_STATES and return post-head logits (including softcapped/scaled logits), this silently changes token logprobs and therefore GRPO loss/KL, corrupting training rather than just fixing the shape mismatch. The pre-change vision fallback explicitly avoided this double application, so this is a behavior regression for direct-logits models.

Useful? React with 👍 / 👎.

[Datta0]

Nice find. Maybe we should consolidate the two into a function and reuse?

[Rorical]

Great suggestion! I've merged them into a shared helper function

[Datta0]

LGTM

[bi1101]

Hi @Datta0 , may I ask if this PR is going to be merged soon? This is blocking text-only GRPO for Qwen 3.5

[lastrei]

Hi @Datta0 , may I ask if this PR is going to be merged soon? This is blocking text-only GRPO for Qwen 3.5

Yes，This update breaks the text-only GRPO for Qwen 3.5. Once I load the model using FastVisionModel, it still expects images even when my input data contains only text and no images at all, triggering RuntimeError: The size of tensor a (2) must match the size of tensor b (0) at non-singleton dimension 1.