Fix Qwen3.5 fp16 training dtype mismatches under UNSLOTH_FORCE_FLOAT32 - #978
Fix Qwen3.5 fp16 training dtype mismatches under UNSLOTH_FORCE_FLOAT32#978chakshu-dhannawat wants to merge 13 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2b01ce3247
ℹ️ 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".
| **kwargs, | ||
| ) | ||
|
|
||
| hidden_states = outputs.last_hidden_state |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 and a caller (or config.return_dict=False) requests tuple outputs, this wrapper forwards return_dict=False through **kwargs to self.model, so outputs is a tuple rather than a ModelOutput. The next dereference of outputs.last_hidden_state then raises before logits/loss are computed, regressing the standard Transformers return_dict=False path; either force return_dict=True for the internal call or handle tuple outputs and return the matching tuple form.
Useful? React with 👍 / 👎.
| return | ||
| try: | ||
| import transformers.models.qwen3_5.modeling_qwen3_5 | ||
| cls = transformers.models.qwen3_5.modeling_qwen3_5.Qwen3_5ForCausalLM |
There was a problem hiding this comment.
For Qwen3.5 vision/multimodal models, the loaded architecture is Qwen3_5ForConditionalGeneration, but this patch only installs the lm-head dtype alignment on Qwen3_5ForCausalLM. In the same UNSLOTH_FORCE_FLOAT32=1 fp16 fallback scenario, the text layers can now return activations in the caller dtype while the conditional-generation lm_head remains fp16, so those users still hit the final bf16/fp16 linear mismatch that this patch is meant to fix; the wrapper needs to cover the conditional-generation class as well.
Useful? React with 👍 / 👎.
| from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig # noqa: E402 | ||
| import transformers.models.qwen3_5.modeling_qwen3_5 as qwen # noqa: E402 |
There was a problem hiding this comment.
On CI environments with a Transformers build that is installed but does not yet ship transformers.models.qwen3_5, pytest.importorskip("transformers") succeeds and these top-level imports then fail during collection with ModuleNotFoundError instead of skipping the Qwen3.5-only regression tests. Use pytest.importorskip for the specific configuration_qwen3_5/modeling_qwen3_5 modules so older supported Transformers jobs keep running.
Useful? React with 👍 / 👎.
|
Thanks for the review. Pushed updates:
Re-ran pytest tests/test_temporary_patches_imports.py tests/test_temporary_patches_exhaustive.py tests/test_qwen3_5_float32.py locally: 143 passed, 12 skipped. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 572d83a543
ℹ️ 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".
| return CausalLMOutputWithPast( | ||
| loss=loss, | ||
| logits=logits, | ||
| past_key_values=getattr(outputs, "past_key_values", None), | ||
| hidden_states=getattr(outputs, "hidden_states", None), | ||
| attentions=getattr(outputs, "attentions", None), | ||
| rope_deltas=getattr(outputs, "rope_deltas", None), |
There was a problem hiding this comment.
When a caller sets return_dict=False (or the config default is false), the upstream Qwen3.5 forwards are decorated to return tuples. This replacement is not decorated, and in the conditional-generation path it forwards return_dict=False into self.model, then these getattr(outputs, ...) reads all resolve to None because outputs is a tuple; requested cache, hidden states, attentions, and rope deltas are dropped while a ModelOutput is returned anyway. Fresh evidence in this revision is that the new wrapper now handles outputs[0] but still never converts the public return back to the tuple form; the CausalLM wrapper above has the same public-contract regression after popping return_dict.
Useful? React with 👍 / 👎.
| logits = self.lm_head(lm_input) | ||
|
|
||
| loss = None | ||
| if labels is not None: | ||
| loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 and training calls this path with labels, this replacement always materializes full-vocabulary logits before computing the loss. Because the temporary patch force-replaces Qwen3_5ForCausalLM.forward after the modeling import hook has had a chance to install the default fused lm-head/loss rewrite, Qwen3.5 fp16 fine-tuning loses the memory-saving unsloth_fused_lm_head_loss path and can OOM on the large-vocab/long-sequence cases that the fused forward is meant to protect. Keep the dtype alignment in the fused-loss path, or call the fused helper directly instead of allocating logits unconditionally.
Useful? React with 👍 / 👎.
| pixel_values=pixel_values, | ||
| pixel_values_videos=pixel_values_videos, | ||
| image_grid_thw=image_grid_thw, | ||
| video_grid_thw=video_grid_thw, | ||
| mm_token_type_ids=mm_token_type_ids, |
There was a problem hiding this comment.
In UNSLOTH_FORCE_FLOAT32=1 multimodal runs where pixel_values or pixel_values_videos are present, this still enters the Qwen3.5 vision path before any dtype guard other than the text-module wrappers added above. patch_model_and_tokenizer also downcasts the visual qkv/MLP/merger weights to fp16, so bf16 activations can still hit Qwen3_5VisionAttention.qkv or the vision MLP/merger linears and fail before the conditional-generation lm_head alignment is reached; add equivalent casts for the vision modules or keep the visual activations in the actual weight dtype.
Useful? React with 👍 / 👎.
|
Addressed the latest Codex points in 4f71efe:
Local run: pytest tests/test_temporary_patches_imports.py tests/test_temporary_patches_exhaustive.py tests/test_qwen3_5_float32.py -> 144 passed, 12 skipped. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4f71efe22a
ℹ️ 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".
| output_attentions = kwargs.pop("output_attentions", None) | ||
| output_hidden_states = kwargs.pop("output_hidden_states", None) |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 and a user enables config.output_hidden_states or config.output_attentions without passing the kwarg on every call, these pops produce None and the wrapper later forwards output_* = None into self.model. Qwen3.5's output-capture decorator treats a present None as the explicit value instead of falling back to the config default, so callers unexpectedly get no hidden states/attentions from this patched forward; only forward these keys when supplied, or substitute the config default.
Useful? React with 👍 / 👎.
| if target_dtype is not None and lm_input.dtype != target_dtype: | ||
| lm_input = lm_input.to(target_dtype) | ||
| loss = fused_loss(lm_input, self.lm_head, labels, vocab_size=self.config.vocab_size, **kwargs) | ||
| logits = EMPTY_LOGITS |
There was a problem hiding this comment.
In UNSLOTH_FORCE_FLOAT32=1 training/eval calls with labels, this branch always returns EMPTY_LOGITS whenever the fused loss helper is available, even if the user set UNSLOTH_RETURN_LOGITS=1. The existing fused-forward path uses that env var to materialize logits once for callbacks/metrics that need them, so Qwen3.5 users who opt in still receive an empty sentinel here; the conditional-generation wrapper has the same duplicated branch below.
Useful? React with 👍 / 👎.
| # Inference path: materialise logits, but align them first. | ||
| if target_dtype is not None and lm_input.dtype != target_dtype: | ||
| lm_input = lm_input.to(target_dtype) | ||
| logits = self.lm_head(lm_input) |
There was a problem hiding this comment.
When GRPO sets UNSLOTH_RETURN_HIDDEN_STATES=1 before calling the model, this replacement still runs the full lm_head and returns vocabulary logits instead of hidden-size activations. That defeats the memory-saving path in unsloth_zoo/rl_replacements.py, which is designed to receive hidden states and apply the LM head selectively; Qwen3.5 GRPO batches can therefore allocate full-vocab logits and OOM despite the env flag. The conditional-generation wrapper repeats the same behavior below.
Useful? React with 👍 / 👎.
|
Addressed the third Codex pass in b1bc535:
Local run remains 144 passed, 12 skipped. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b1bc535316
ℹ️ 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".
| attentions=outputs.attentions, | ||
| ) | ||
|
|
||
| target_dtype = getattr(getattr(self.lm_head, "weight", None), "dtype", None) |
There was a problem hiding this comment.
When a Qwen3.5 model is loaded with a quantized output head (for example a custom bnb config that does not skip lm_head), self.lm_head.weight.dtype is the packed storage dtype rather than the compute dtype. This branch then casts lm_input to that non-floating dtype before calling the head, so materialized-logit forwards either fail or produce invalid logits; the conditional-generation wrapper repeats the same logic. Please guard this like _unsloth_get_linear_weight_dtype does for quant_state/non-floating dtypes and let the quantized module choose its compute dtype.
Useful? React with 👍 / 👎.
|
|
||
| target_dtype = getattr(getattr(self.lm_head, "weight", None), "dtype", None) | ||
|
|
||
| if labels is not None and fused_loss is not None and EMPTY_LOGITS is not None and not RETURN_LOGITS: |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 and labels are passed, this branch now routes every non-UNSLOTH_RETURN_LOGITS training call through unsloth_fused_lm_head_loss as soon as that helper is importable, ignoring self.loss_function. If a user or config selected a nonstandard loss function, Qwen3.5 training silently falls back to the stock fused causal-LM CE objective; the conditional-generation wrapper repeats the same unconditional branch. Please keep the fused shortcut behind the same kind of self.loss_function compatibility guard used by the compiler rewrite, and otherwise call the configured loss.
Useful? React with 👍 / 👎.
|
Addressed the fourth Codex pass in 6a51a14:
Local run: pytest tests/test_temporary_patches_imports.py tests/test_temporary_patches_exhaustive.py tests/test_qwen3_5_float32.py -> 145 passed, 12 skipped. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6a51a14f61
ℹ️ 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".
| rather than being silently replaced by the fused CE kernel. | ||
| """ | ||
| name = getattr(loss_function, "__name__", "") | ||
| return name in ("ForCausalLMLoss",) |
There was a problem hiding this comment.
When Unsloth's loss patch is active in normal training, the default causal-LM loss is renamed UnslothForCausalLMLoss, and the existing compiler guard treats names ending in ForCausalLMLoss as compatible. This exact-name check returns false for that default, so Qwen3.5 calls with labels fall through to the logits branch and materialize full-vocabulary logits, reintroducing the OOM regression the fused path is meant to avoid. Fresh evidence beyond the earlier review is the new exact-name guard here; use the same suffix check or otherwise whitelist the Unsloth default.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aa29256401
ℹ️ 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".
|
|
||
| # Always work with ModelOutput internally; @can_return_tuple preserves | ||
| # the public tuple/return_dict contract. | ||
| kwargs["return_dict"] = True |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1, labels are supplied, and self.loss_function is a custom loss with the normal strict signature, this mutates the same kwargs dict that is later forwarded to self.loss_function, so the loss receives the synthetic return_dict=True kwarg even though upstream @can_return_tuple has already popped it before loss computation. The conditional-generation wrapper repeats the same pattern, so use a separate kwargs copy for the internal model call or remove return_dict before invoking the loss.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 94c9050333
ℹ️ 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".
| and _unsloth_is_default_causal_lm_loss(self.loss_function): | ||
| if target_dtype is not None and lm_input.dtype != target_dtype: | ||
| lm_input = lm_input.to(target_dtype) | ||
| loss = fused_loss(lm_input, self.lm_head, labels, vocab_size=self.config.text_config.vocab_size, **kwargs) |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 and Qwen3.5 VLM training is run through Unsloth/Trainer, num_items_in_batch can be present in kwargs; this line forwards it into unsloth_fused_lm_head_loss, which consumes it and divides the loss. The upstream Qwen3_5ForConditionalGeneration explicitly has accepts_loss_kwargs = False and does not pass **kwargs to its loss, so this dtype-only patch changes loss normalization for gradient accumulation; fresh evidence after the earlier return_dict fix is that the remaining kwargs are still forwarded here, so filter them or honor accepts_loss_kwargs before calling the fused/custom loss.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7017bb0afc
ℹ️ 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".
| weight = getattr(linear, "weight", None) | ||
| if weight is None: | ||
| continue |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1 is used for the normal LoRA fine-tuning path, projections such as q_proj or gate_proj can be PEFT wrapper modules whose fp16 base weight lives under base_layer/get_base_layer() rather than directly on linear.weight. In that case this helper sees no weight and returns None, so the wrapper skips the activation cast and the wrapped base linear still receives bf16 activations against fp16 weights, reproducing the dtype mismatch this patch is intended to fix; unwrap adapter modules before deciding the target dtype.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 113f4907ec
ℹ️ 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".
|
|
||
| def _unsloth_fused_loss_kwargs(kwargs): | ||
| """Return only the kwargs the fused loss kernel understands.""" | ||
| return {k: v for k, v in kwargs.items() if k in _FUSED_LOSS_KWARG_KEYS} |
There was a problem hiding this comment.
When UNSLOTH_FORCE_FLOAT32=1, labels are passed, and the fused branch is used, this filter drops supported loss kwargs such as ignore_index and label_smoothing. Upstream Qwen3_5ForCausalLM.forward forwards **kwargs to self.loss_function, and the local fused CE implementation already reads those kwargs, so calls like model(..., labels=..., ignore_index=pad_id) or label-smoothed training silently fall back to -100/0.0; include the supported loss kwargs in the allowlist instead of discarding them.
Useful? React with 👍 / 👎.
|
Confirmed the gap: qwen3_5 is in FORCE_FLOAT32 in model_lists.py but there is no qwen3_5 float32 patch alongside gemma4_float32.py and qwen3_moe_float32.py, so fp16 loads hit the mismatch. The direction is right, but this also reimplements the causal-LM and conditional-generation forwards with fused loss and vision-tower handling. Could you narrow it to the dtype normalization following the qwen3_moe_float32.py shape and leave the loss path alone? |
|
Rechecked and the missing qwen3_5 dtype normalization is still a real gap next to gemma4_float32.py and qwen3_moe_float32.py, but the branch still carries the causal-LM and conditional-generation forwards with fused loss and vision-tower handling. Could you push the narrowed version that only does the dtype casts in the qwen3_moe_float32.py shape and leaves the loss path untouched? |
|
Hi Daniel, the branch at bbfbe5b has been narrowed to component-level dtype wrappers only (Qwen3_5GatedDeltaNet, Qwen3_5Attention, Qwen3_5MLP). The causal-LM / conditional-generation forwards, fused-loss helpers, and vision-tower wrappers were removed. The PR body is also updated to match. Let me know if you'd like any further changes. |
bbfbe5b to
daa9d9b
Compare
GatedDeltaNet, Attention, MLP, and ForCausalLM wrappers align activations with the actual fp16 weight dtype when UNSLOTH_FORCE_FLOAT32=1, then restore the original dtype for the residual stream. Fixes unsloth#7506 mismatches like BFloat16 != Half. Includes regression tests.
…gen head - Force return_dict=True (and pass output_attentions/output_hidden_states explicitly) inside Qwen3_5ForCausalLM_dtype wrapper so outputs.last_hidden_state works regardless of the caller's return_dict value. - Add Qwen3_5ForConditionalGeneration wrapper covering vision/multimodal Qwen3.5 models, which use the same lm_head dtype issue. - Use pytest.importorskip for the specific qwen3_5 configuration/modeling modules so older Transformers CI jobs skip cleanly.
- Apply @can_return_tuple to Qwen3_5ForCausalLM / ForConditionalGeneration wrappers so return_dict=False still produces tuples. - On the training branch use unsloth_fused_lm_head_loss directly instead of materialising full fp32 logits, preserving the memory-saving fused path. - Add dtype guards for Qwen3.5 vision modules (VisionAttention, VisionMLP, VisionPatchMerger) so multimodal runs with UNSLOTH_FORCE_FLOAT32 do not fail before reaching the text head. - Extend _unsloth_get_linear_weight_dtype to cover qkv / linear_fc1. - Add CausalLM regression test covering both ModelOutput and tuple output.
…n states - Stop popping output_attentions/output_hidden_states; set return_dict=True and forward remaining kwargs, so config.output_hidden_states and config.output_attentions still take effect. - Honor UNSLOTH_RETURN_LOGITS=1 in the fused-loss branch by materialising logits alongside the fused CE loss. - Honor UNSLOTH_RETURN_HIDDEN_STATES=1 by returning the sliced hidden states in the logits field, matching qwen3_moe's GRPO helper pattern. - Add regression test that config.output_hidden_states propagates.
- Add _unsloth_weight_dtype helper that skips quantized/non-floating/fp8 weights, and use it for lm_head dtype alignment in both CausalLM and ConditionalGeneration wrappers. - Add _unsloth_is_default_causal_lm_loss guard so the fused CE shortcut is only used for the standard ForCausalLMLoss; custom self.loss_function configs fall back to materialised logits + their configured loss.
Remove the reimplemented Qwen3_5ForCausalLM and Qwen3_5ForConditionalGeneration forwards, the fused-loss helpers, and the vision-tower dtype wrappers. Keep only per-component dtype normalization for GatedDeltaNet, Attention and MLP, matching the qwen3_moe_float32.py component-level shape. Update regression tests to the narrower scope. Assisted-by: Claude Sonnet
daa9d9b to
750b97a
Compare
2 participants
Adds targeted dtype-normalization patches for Qwen3.5 when UNSLOTH_FORCE_FLOAT32=1.
Problem
Fix
Validation