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 <michaelhan2050@gmail.com>

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."
+    )

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(

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(