Fix Qwen3 NaN loss: delegate pad_token repair to shared unsloth_zoo.pad_token (#6524)

* Fix Qwen3 NaN: self-heal vision pad_token in load_correct_tokenizer

Text-only Qwen3 (and Qwen2.5) models share Qwen3-VL's vocab, so their
Hub tokenizer configs ship <|vision_pad|> as pad_token. Padding text-only
training with a vision token corrupts attention/loss and produces NaN
losses and gradients on affected stacks.

patch_tokenizer already heals this, but only when a model with config is
passed. The standalone load_correct_tokenizer path (and custom training
loops) still returned <|vision_pad|>. This adds a model-independent guard
in load_correct_tokenizer that replaces a vision pad_token on text-only
tokenizers with the first safe text token (<|endoftext|>, <pad>, [PAD],
<unk>), falling back to eos_token only if it differs from pad_token.

The result now matches upstream Qwen configs (pad_token <|endoftext|>,
id 151643) with no new token added. Vision processors (image_processor
present) and non-vision pad tokens (Llama, Qwen2) are left untouched.

Fixes #3155

* Format pad_token helper for ruff kwarg-spacing hook (pre-commit)

* Harden vision pad_token fix: drop unk candidate, guard get_vocab, skip vision eos

* Tighten code comments (no logic change)

* Delegate pad_token fix to shared unsloth_zoo.pad_token

Generalize the narrow _fix_vision_pad_token by delegating to unsloth_zoo's
shared fix_pad_token (single source of truth, AGPL-3.0), which scans the
reserved-token families instead of only the vision-pad case. A guarded import
keeps this working against an older unsloth_zoo that has not shipped the module
yet: on ImportError it falls back to _fix_vision_pad_token.

allow_add=False is passed so the early load_correct_tokenizer call stays
side-effect free (no model here to resize embeddings); the later model-aware
patch_tokenizer call finishes the job and is idempotent.

Adds tests/python/test_pad_token_fix.py covering both dispatch paths offline.

tests/python/test_pad_token_fix.py

@@ -0,0 +1,112 @@
+"""_fix_pad_token dispatch in unsloth/tokenizer_utils.py.
+
+It must delegate to unsloth_zoo's shared fix_pad_token when present (single
+source of truth), and fall back to the narrow vision-token swap against an
+older unsloth_zoo. Static + CPU-only: the two helpers are exec'd in isolation
+so the test never imports torch / transformers / unsloth.
+"""
+
+import ast
+import os
+import sys
+import types
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+TOK_PATH = os.path.join(REPO_ROOT, "unsloth", "tokenizer_utils.py")
+
+WANTED = {
+    "_VISION_PAD_TOKENS",
+    "_SAFE_TEXT_PAD_TOKENS",
+    "_fix_vision_pad_token",
+    "_fix_pad_token",
+}
+
+
+def _load_pad_helpers():
+    """Exec only the pad-token helpers with a stub logger (no heavy imports)."""
+    tree = ast.parse(open(TOK_PATH).read())
+    nodes = []
+    for node in tree.body:
+        if isinstance(node, ast.Assign):
+            names = {t.id for t in node.targets if isinstance(t, ast.Name)}
+            if names & WANTED:
+                nodes.append(node)
+        elif isinstance(node, ast.FunctionDef) and node.name in WANTED:
+            nodes.append(node)
+    module = ast.Module(body = nodes, type_ignores = [])
+    ast.fix_missing_locations(module)
+    ns = {"logger": types.SimpleNamespace(warning = lambda *a, **k: None)}
+    exec(compile(module, TOK_PATH, "exec"), ns)
+    return ns
+
+
+class FakeTok:
+    def __init__(self, vocab, pad, eos):
+        self._v = dict(vocab)
+        self.pad_token = pad
+        self.eos_token = eos
+
+    @property
+    def pad_token_id(self):
+        return self._v.get(self.pad_token)
+
+    @property
+    def eos_token_id(self):
+        return self._v.get(self.eos_token)
+
+    def get_vocab(self):
+        return dict(self._v)
+
+
+def _block_shared_module(monkeypatch):
+    # Stub parent so importing it is cheap, and mark the submodule absent.
+    monkeypatch.setitem(sys.modules, "unsloth_zoo", types.ModuleType("unsloth_zoo"))
+    monkeypatch.setitem(sys.modules, "unsloth_zoo.pad_token", None)
+
+
+def test_fix_pad_token_none_is_noop():
+    ns = _load_pad_helpers()
+    assert ns["_fix_pad_token"](None) is None
+
+
+def test_fix_pad_token_falls_back_without_shared_module(monkeypatch):
+    ns = _load_pad_helpers()
+    _block_shared_module(monkeypatch)
+    # Qwen3 text tokenizer shipping a vision pad_token -> narrow swap heals it.
+    tok = FakeTok(
+        {"<|endoftext|>": 1, "<|im_end|>": 2, "<|vision_pad|>": 3},
+        pad = "<|vision_pad|>",
+        eos = "<|im_end|>",
+    )
+    ns["_fix_pad_token"](tok)
+    assert tok.pad_token == "<|endoftext|>"
+    assert tok.pad_token != tok.eos_token
+
+
+def test_fallback_leaves_valid_pad_untouched(monkeypatch):
+    ns = _load_pad_helpers()
+    _block_shared_module(monkeypatch)
+    tok = FakeTok({"<s>": 1, "</s>": 2, "<pad>": 0}, pad = "<pad>", eos = "</s>")
+    ns["_fix_pad_token"](tok)
+    assert tok.pad_token == "<pad>"
+
+
+def test_fix_pad_token_delegates_to_shared_module(monkeypatch):
+    ns = _load_pad_helpers()
+    calls = {}
+
+    fake = types.ModuleType("unsloth_zoo.pad_token")
+
+    def fix_pad_token(tokenizer, allow_add = True):
+        calls["allow_add"] = allow_add
+        tokenizer.pad_token = "SHARED"
+
+    fake.fix_pad_token = fix_pad_token
+    monkeypatch.setitem(sys.modules, "unsloth_zoo", types.ModuleType("unsloth_zoo"))
+    monkeypatch.setitem(sys.modules, "unsloth_zoo.pad_token", fake)
+
+    tok = FakeTok({"a": 1}, pad = "x", eos = "y")
+    ns["_fix_pad_token"](tok)
+    # Delegated, and crucially with allow_add=False (no model here to resize).
+    assert tok.pad_token == "SHARED"
+    assert calls["allow_add"] is False

unsloth/tokenizer_utils.py

@@ -626,6 +626,76 @@ def _load_correct_tokenizer(
         return fast_tokenizer
 
 
+# Qwen3 text models share Qwen3-VL's vocab, so configs ship a vision pad_token;
+# padding text-only training with one yields NaN losses (#3155, #4104).
+_VISION_PAD_TOKENS = frozenset(
+    (
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>",
+        "<|audio_pad|>",
+    )
+)
+# Preference order; <unk> excluded since reusing it as pad masks real OOV tokens.
+_SAFE_TEXT_PAD_TOKENS = ("<|endoftext|>", "<pad>", "[PAD]")
+
+
+def _fix_vision_pad_token(tokenizer):
+    """Swap a vision pad_token on a text-only tokenizer for a safe text token (#3155)."""
+    if tokenizer is None:
+        return tokenizer
+    if hasattr(tokenizer, "image_processor"):
+        return tokenizer
+    pad_token = getattr(tokenizer, "pad_token", None)
+    if pad_token is None or pad_token not in _VISION_PAD_TOKENS:
+        return tokenizer
+
+    get_vocab = getattr(tokenizer, "get_vocab", None)
+    if get_vocab is None:
+        return tokenizer
+    vocab = get_vocab()
+    if not isinstance(vocab, dict):
+        return tokenizer
+    new_pad_token = None
+    for candidate in _SAFE_TEXT_PAD_TOKENS:
+        if candidate in vocab and candidate != getattr(tokenizer, "eos_token", None):
+            new_pad_token = candidate
+            break
+    # Fall back to eos_token only if it is a distinct, non-vision token.
+    if new_pad_token is None:
+        eos_token = getattr(tokenizer, "eos_token", None)
+        if eos_token is not None and eos_token != pad_token and eos_token not in _VISION_PAD_TOKENS:
+            new_pad_token = eos_token
+    if new_pad_token is None:
+        return tokenizer
+
+    tokenizer.pad_token = new_pad_token
+    logger.warning(
+        f"Unsloth: pad_token was a vision token ({pad_token}) on a text-only "
+        f"model. Replaced with {new_pad_token} to avoid NaN losses."
+    )
+    return tokenizer
+
+
+def _fix_pad_token(tokenizer):
+    """Heal a bad/missing pad_token before chat-template repair.
+
+    Delegates to unsloth_zoo's shared fix_pad_token (single source of truth) when
+    available, falling back to the narrow vision-token swap against an older
+    unsloth_zoo. allow_add=False keeps this side-effect free: there is no model
+    here to resize embeddings, so a brand new pad token is never added - the later
+    model-aware patch_tokenizer call finishes the job and is idempotent.
+    """
+    if tokenizer is None:
+        return tokenizer
+    try:
+        from unsloth_zoo.pad_token import fix_pad_token
+    except Exception:
+        return _fix_vision_pad_token(tokenizer)
+    fix_pad_token(tokenizer, allow_add = False)
+    return tokenizer
+
+
 def load_correct_tokenizer(
     tokenizer_name,
     model_max_length = None,
@@ -645,6 +715,9 @@ def load_correct_tokenizer(
         fix_tokenizer = fix_tokenizer,
     )
 
+    if fix_tokenizer:
+        _fix_pad_token(tokenizer)
+
     ### 1. Fixup tokenizer's chat_template
     old_chat_template = getattr(tokenizer, "chat_template", None)