diff --git a/tests/test_uninitialized_position_ids.py b/tests/test_uninitialized_position_ids.py new file mode 100644 index 0000000000..76ff71d06c --- /dev/null +++ b/tests/test_uninitialized_position_ids.py @@ -0,0 +1,76 @@ +"""RaiseUninitialized must ignore a checkpoint that only re-initializes deterministic +position_ids buffers, but still raise when a real weight is missing -- even if the same +HF record also lists a benign position_ids buffer. +""" + +from __future__ import annotations + +import logging + +import pytest + +from unsloth.models._utils import ( + _all_missing_keys_are_position_ids, + _RaiseUninitialized, +) + +_TEMPLATE = ( + "Some weights of DeepseekOCRForCausalLM were not initialized from the model " + "checkpoint at unsloth/DeepSeek-OCR and are newly initialized: {keys}\n" + "You should probably TRAIN this model on a down-stream task." +) + + +def _record(keys_repr: str) -> logging.LogRecord: + return logging.LogRecord( + name = "transformers.modeling_utils", + level = logging.WARNING, + pathname = "modeling_utils.py", + lineno = 1, + msg = _TEMPLATE.format(keys = keys_repr), + args = None, + exc_info = None, + ) + + +@pytest.mark.parametrize( + "keys_repr, expected", + [ + ("['model.vision_model.embeddings.position_ids']", True), + ( + "['model.vision_model.embeddings.position_ids', " + "'vision_model.encoder.layers.0.position_ids']", + True, + ), + # A real missing weight alongside position_ids must NOT be suppressed. + ( + "['model.vision_model.embeddings.position_ids', 'model.layers.5.mlp.weight']", + False, + ), + ("['model.layers.5.mlp.weight']", False), + ("[]", False), + ], +) +def test_all_missing_keys_are_position_ids(keys_repr, expected): + assert _all_missing_keys_are_position_ids(_TEMPLATE.format(keys = keys_repr)) is expected + + +def test_emit_suppresses_position_ids_only_record(): + # A record listing only position_ids buffers loads cleanly (no raise). + handler = _RaiseUninitialized() + handler.emit(_record("['model.vision_model.embeddings.position_ids']")) + + +def test_emit_raises_when_real_weight_missing_alongside_position_ids(): + # The core fix: one benign position_ids key must not mask a real missing weight. + handler = _RaiseUninitialized() + with pytest.raises(Exception, match = "some weights are not initialized"): + handler.emit( + _record("['model.vision_model.embeddings.position_ids', 'model.layers.5.mlp.weight']") + ) + + +def test_emit_raises_on_real_missing_weight(): + handler = _RaiseUninitialized() + with pytest.raises(Exception, match = "some weights are not initialized"): + handler.emit(_record("['model.layers.5.mlp.weight']")) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 4b2d0acec9..a15320668c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -1015,18 +1015,44 @@ except: from transformers.modeling_utils import logger as transformers_logger +def _all_missing_keys_are_position_ids(record_str): + """True only when EVERY key in the 'newly initialized: [...]' list is a position_ids + buffer. + + transformers reports all missing keys in a single record, so a substring test would + wrongly suppress the warning when a real missing weight is listed alongside a benign + position_ids buffer. position_ids is a deterministic arange buffer that transformers + itself lists in _keys_to_ignore_on_load_missing (some VLMs, e.g. DeepSeek-OCR, ship it + non-persistently), so a record listing ONLY position_ids keys is safe to ignore; + anything else must still raise. + """ + import ast + import re + + match = re.search(r"newly initialized:\s*(\[[^\]]*\])", record_str) + if not match: + return False + try: + keys = ast.literal_eval(match.group(1)) + except Exception: + return False + return bool(keys) and all("position_ids" in str(key) for key in keys) + + class _RaiseUninitialized(logging.Handler): def __init__(self): super().__init__() def emit(self, record): - record_lower = str(record).lower() + record_str = str(record) + record_lower = record_str.lower() if ( ("some weights of" in record_lower) and ("score.weight" not in record_lower) and ("classifier.weight" not in record_lower) and ("cls.predictions" not in record_lower) and ("predictions.decoder" not in record_lower) + and not _all_missing_keys_are_position_ids(record_str) and (os.environ.get("UNSLOTH_WARN_UNINITIALIZED", "1") == "1") ): raise Exception( diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index cfcfcae505..459b6ed32d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -1576,12 +1576,23 @@ class FastModel(FastBaseModel): auto_model = AutoModelForSequenceClassification elif is_vlm: # Check if the model's auto_map supports the VLM auto class. - # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM - # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq. + # Some repo-code VL models register only a generic auto class and not + # AutoModelForImageTextToText/AutoModelForVision2Seq: Nemotron-VL uses + # AutoModelForCausalLM, DeepSeek-OCR uses AutoModel. Calling the VLM auto + # class on those raises "Unrecognized configuration class ... for + # AutoModelForImageTextToText", so fall back to whatever generic class the + # repo actually registered. Match the CONCRETE class name we would pass + # (AutoModelForVision2Seq aliases to AutoModelForImageTextToText on tf>=5), + # since transformers resolves remote code by that exact name -- a config + # that only registers the legacy key must still take the generic fallback. _auto_map = getattr(model_config, "auto_map", {}) or {} _vlm_class_name = AutoModelForVision2Seq.__name__ - if "AutoModelForCausalLM" in _auto_map and _vlm_class_name not in _auto_map: + _has_vlm_class = _vlm_class_name in _auto_map + if not _has_vlm_class and "AutoModelForCausalLM" in _auto_map: auto_model = AutoModelForCausalLM + elif not _has_vlm_class and "AutoModel" in _auto_map: + from transformers import AutoModel + auto_model = AutoModel else: auto_model = AutoModelForVision2Seq else: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 66f9cf3d1b..3f8d9e6ce0 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -635,6 +635,13 @@ class FastBaseModel: # Pure text model requested text-only with a VLM auto class. auto_model = AutoModelForCausalLM is_vlm = auto_model in [AutoModelForVision2Seq, AutoModelForImageTextToText] + # A repo-code VLM may register only AutoModel / AutoModelForCausalLM (e.g. + # DeepSeek-OCR, Nemotron-VL), so auto_model is not a VLM class even though the + # config is a vision model. Keep is_vlm (auto-class derived) for processor + # selection below -- these repos ship no AutoProcessor -- but treat the model as a + # VLM on the vLLM path so a vision_config model is never silently loaded/converted + # as text-only. text_only resolves the tower away first, so honour that here. + is_vlm_config = is_vlm or (not text_only and hasattr(auto_config, "vision_config")) is_whisper = whisper_language is not None and whisper_task is not None auto_processor = AutoProcessor if (is_vlm or is_whisper) else AutoTokenizer @@ -646,7 +653,7 @@ class FastBaseModel: vllm_enable_lora = True - if is_vlm and fast_inference: + if is_vlm_config and fast_inference: if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types): raise RuntimeError( f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. " @@ -1060,7 +1067,7 @@ class FastBaseModel: disable_log_stats = disable_log_stats, use_bitsandbytes = load_in_4bit, unsloth_vllm_standby = unsloth_vllm_standby, - is_vision_model = is_vlm, + is_vision_model = is_vlm_config, fp8_mode = fp8_mode, ) for allowed_arg in allowed_args: @@ -1074,7 +1081,7 @@ class FastBaseModel: _, quant_state_dict = get_vllm_state_dict( llm, config = model_config, - is_vision_model = is_vlm, + is_vision_model = is_vlm_config, load_in_fp8 = load_in_fp8, ) model = convert_vllm_to_huggingface( @@ -1082,7 +1089,7 @@ class FastBaseModel: model_config, dtype, bnb_config, - is_vision_model = is_vlm, + is_vision_model = is_vlm_config, ) model.vllm_engine = llm llm.shared_weights = True