diff --git a/tests/python/test_pad_token_fix.py b/tests/python/test_pad_token_fix.py new file mode 100644 index 000000000..0366bee67 --- /dev/null +++ b/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({"": 1, "": 2, "": 0}, pad = "", eos = "") + ns["_fix_pad_token"](tok) + assert tok.pad_token == "" + + +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 diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index f454d5c50..3cec00201 100644 --- a/unsloth/tokenizer_utils.py +++ b/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; excluded since reusing it as pad masks real OOV tokens. +_SAFE_TEXT_PAD_TOKENS = ("<|endoftext|>", "", "[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)