* 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>
eae59b25b6