From eae59b25b6d83977a20aaa1438b00c62832f40b8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jun 2026 01:26:55 -0700 Subject: [PATCH] fix: use EMPTY_LOGITS on the fused-CE not-return_dict path (#2068) (#6482) * fix: use EMPTY_LOGITS on the fused-CE not-return_dict path (#2068) CausalLM_fast_forward's fused cross-entropy path (small batch, labels set, UNSLOTH_RETURN_LOGITS off) computes the loss straight from hidden_states via unsloth_fused_ce_loss and never materializes `logits`. The return_dict=True branch returns EMPTY_LOGITS, but the `not return_dict` branch returned `(logits,) + outputs[1:]`, raising "UnboundLocalError: cannot access local variable 'logits'" whenever it ran (e.g. training with return_dict=False). Same bug in the llama and mistral fast-forward paths. Return EMPTY_LOGITS on that branch too, matching the adjacent return_dict output. Verified on GPU: a forward(return_dict=False, labels=...) that raised UnboundLocalError now returns (loss, EMPTY_LOGITS, ...) and backward() succeeds. Adds tests/test_fused_ce_not_return_dict_logits.py, a CPU source-drift guard (the fused path itself is GPU/triton only) asserting both fast-forward paths keep using EMPTY_LOGITS there. * Address review: parse the fused-CE drift line with whitespace-tolerant regexes The drift detector sliced the source with exact string matching (source.index("output = (") + the next newline), so a formatter respacing or rewrapping the assignment would break the parse. Switch to anchored regexes that tolerate whitespace and line wrapping, keeping the match anchored after the fused guard so it targets the fused-CE branch and not the normal output = (logits,) path. Behavior and the two drift assertions are unchanged. * Tighten code comments (no logic change) --------- Co-authored-by: Daniel Han --- tests/test_fused_ce_not_return_dict_logits.py | 62 +++++++++++++++++++ unsloth/models/llama.py | 4 +- unsloth/models/mistral.py | 4 +- 3 files changed, 68 insertions(+), 2 deletions(-) create mode 100644 tests/test_fused_ce_not_return_dict_logits.py diff --git a/tests/test_fused_ce_not_return_dict_logits.py b/tests/test_fused_ce_not_return_dict_logits.py new file mode 100644 index 0000000000..51a6ec4843 --- /dev/null +++ b/tests/test_fused_ce_not_return_dict_logits.py @@ -0,0 +1,62 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Drift detector for unsloth#2068. + +The fused-CE path never materializes ``logits``, so its ``not return_dict`` +return must use ``EMPTY_LOGITS`` (it once used ``logits`` -> UnboundLocalError). +Pure text inspection, so it runs GPU-free (the fused path is GPU/triton only).""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + + +_REPO = Path(__file__).resolve().parent.parent + + +def _fused_ce_not_return_dict_return(source: str) -> str: + """Return the `output = (...) + outputs[1:]` assignment from the first + `if not return_dict:` after the `unsloth_fused_ce_loss(` call. Whitespace- + tolerant regexes survive reformatting; the anchor targets the fused branch.""" + fused = re.search(r"unsloth_fused_ce_loss\s*\(", source) + if fused is None: + raise ValueError("unsloth_fused_ce_loss( call not found") + guard = re.search(r"if\s+not\s+return_dict\s*:", source[fused.end() :]) + if guard is None: + raise ValueError("`if not return_dict:` guard not found after fused-CE call") + guard_start = fused.end() + guard.start() + out = re.search(r"output\s*=\s*\(.*?outputs\s*\[\s*1\s*:\s*\]", source[guard_start:], re.DOTALL) + if out is None: + raise ValueError("`output = (...) + outputs[1:]` assignment not found") + return out.group(0) + + +@pytest.mark.parametrize("rel", ["unsloth/models/llama.py", "unsloth/models/mistral.py"]) +def test_fused_ce_not_return_dict_uses_empty_logits(rel): + path = _REPO / rel + source = path.read_text(encoding = "utf-8") + assert "unsloth_fused_ce_loss(" in source, f"{rel}: fused-CE call vanished" + + ret = _fused_ce_not_return_dict_return(source) + assert "EMPTY_LOGITS" in ret, ( + f"DRIFT (#2068): {rel} fused-CE `not return_dict` returns {ret!r}; it must " + f"use EMPTY_LOGITS since `logits` is never assigned on the fused-CE path." + ) + assert "(logits,)" not in ret, ( + f"DRIFT (#2068): {rel} fused-CE `not return_dict` references the unassigned " + f"`logits` ({ret!r}) -> UnboundLocalError at runtime." + ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 7b9cf3df0e..c61e64b41a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1486,7 +1486,9 @@ def CausalLM_fast_forward(fast_forward_inference): logit_softcapping = logit_softcapping, ) if not return_dict: - output = (logits,) + outputs[1:] + # Fused CE never materializes `logits`; use EMPTY_LOGITS + # like the return_dict branch below (fixes #2068). + output = (EMPTY_LOGITS,) + outputs[1:] return (loss,) + output if loss is not None else output output = CausalLMOutputWithPast( diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 74f7acbb03..df2a4de5bd 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -332,7 +332,9 @@ def MistralForCausalLM_fast_forward( logit_softcapping = logit_softcapping, ) if not return_dict: - output = (logits,) + outputs[1:] + # Fused CE never materializes `logits`; use EMPTY_LOGITS + # like the return_dict branch below (fixes #2068). + output = (EMPTY_LOGITS,) + outputs[1:] return (loss,) + output if loss is not None else output output = CausalLMOutputWithPast(