ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Load DeepSeek-OCR and other VLMs that register AutoModel in auto_map (#6421)

Browse source 
* Load repo-code VLMs that register AutoModel in auto_map

FastModel.from_pretrained already falls back from the VLM auto class to
AutoModelForCausalLM for repo-code VL models that register only that class
in their auto_map (e.g. Nemotron-VL). Models like DeepSeek-OCR and
DeepSeek-OCR-2 instead register their architecture under AutoModel, so they
fell through to AutoModelForImageTextToText and raised "Unrecognized
configuration class ... for AutoModelForImageTextToText".

Generalize the guard: when neither vision auto class is registered, fall
back to whichever generic auto class the repo actually registered
(AutoModelForCausalLM, else AutoModel).

* Do not hard-error on a newly initialized position_ids buffer

RaiseUninitialized turns transformers' "some weights of ... were not
initialized" warning into a hard error. position_ids is a deterministic
arange buffer that transformers itself lists in
_keys_to_ignore_on_load_missing, so re-initializing it is correct rather
than a sign of a corrupt checkpoint. Some VLMs (e.g. DeepSeek-OCR) ship it
non-persistently, which tripped the guard. Allowlist position_ids alongside
the existing classifier/predictions head weights.

* Only ignore missing-weight records that are exclusively position_ids

The previous substring check skipped the whole "Some weights of ..." record
whenever position_ids appeared anywhere in it. Transformers reports every
missing key in one record, so a corrupt or incompatible checkpoint missing a
real parameter could load with randomly initialized weights as long as one
missing key contained position_ids. Parse the "newly initialized: [...]" list
and suppress only when every listed key is a position_ids buffer; otherwise
raise as before.

* Match the concrete VLM auto class name when checking auto_map

Transformers resolves remote code by the exact auto class name being called,
and AutoModelForVision2Seq aliases to AutoModelForImageTextToText on
transformers >= 5. Checking for both spellings treated a config that only
registers the legacy AutoModelForVision2Seq key as having a supported VLM
class, skipping the AutoModelForCausalLM fallback that used to load it and
failing as an unrecognized config under AutoModelForImageTextToText. Match
only the concrete class name we would actually pass, keeping the AutoModel
and AutoModelForCausalLM fallbacks.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Preserve VLM mode on the vLLM path when falling back to AutoModel

A repo-code VLM that registers only AutoModel or AutoModelForCausalLM (DeepSeek-OCR, Nemotron-VL) routes to that generic class, so is_vlm, derived from the resolved auto class, is False. That is correct for processor selection (these repos ship no AutoProcessor) but wrong for the vLLM path, where is_vision_model=is_vlm made vLLM treat a vision_config model as text-only and skip the VLM guard and conversion.

Add is_vlm_config, derived from the config vision_config (and gated on not text_only so a text-only resolve still wins), and use it for the fast_inference VLM guard and the is_vision_model flags passed to load_vllm, get_vllm_state_dict and convert_vllm_to_huggingface. Processor selection still uses is_vlm, so DeepSeek-OCR keeps loading via its tokenizer. DeepSeek-OCR with fast_inference now raises the clear 'Fast inference is only supported for ...' error instead of being mishandled as text-only.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-18 07:03:43 -07:00 • committed by GitHub
commit 18f8869829
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 128 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

76
tests/test_uninitialized_position_ids.py Normal file
View file

@ -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']"))

28
unsloth/models/_utils.py
View file

@ -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(

17
unsloth/models/loader.py
View file

@ -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:

15
unsloth/models/vision.py
View file

@ -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