diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7e8d52525c..c6c2e1fc37 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -269,7 +269,8 @@ jobs: tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ - tests/utils/test_trunc_normal_patch.py + tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py python -m pytest --collect-only -q "$RUNNER_TEMP/unsloth-zoo/tests/" - name: import_fixes drift detectors (18 tests, HARD GATE) @@ -333,11 +334,9 @@ jobs: python -m pytest -v --tb=short tests/test_callback_signature_drift.py - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) - # 16 tests across 5 files. They live inside tests/saving/ and - # tests/utils/, both of which Repo tests (CPU) excludes via --ignore - # because their sibling files need real GPUs / real HF weights. - # The five files below are pure-Python + AST/protobuf/regex tests - # that run cleanly on CPU. Env inherited from the job block. + # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ + # that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model + # loads; run cleanly here (transformers/torch installed). run: | python -m pytest -q --tb=short \ tests/saving/test_save_shell_injection.py \ @@ -345,11 +344,12 @@ jobs: tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest - # runner does not have. The other 15 Bucket-A tests pass cleanly. + # runner does not have. The other Bucket-A tests pass cleanly. - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip diff --git a/tests/python/test_fast_language_model_text_only.py b/tests/python/test_fast_language_model_text_only.py new file mode 100644 index 0000000000..ce4dd74439 --- /dev/null +++ b/tests/python/test_fast_language_model_text_only.py @@ -0,0 +1,433 @@ +"""Text-only FastLanguageModel routing for vision-capable configs.""" + +import ast +import copy +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[2] +LOADER_PATH = REPO_ROOT / "unsloth" / "models" / "loader.py" +VISION_PATH = REPO_ROOT / "unsloth" / "models" / "vision.py" +UTILS_PATH = REPO_ROOT / "unsloth" / "models" / "_utils.py" + + +def _source(path): + return path.read_text() + + +def _class_method(tree, class_name, method_name): + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == method_name: + return item + raise AssertionError(f"{class_name}.{method_name} not found") + + +def _assigns_name(method, target_name, predicate): + """True when the method contains `target_name = ` and predicate(value).""" + for node in ast.walk(method): + if not isinstance(node, ast.Assign): + continue + for target in node.targets: + if isinstance(target, ast.Name) and target.id == target_name: + if predicate(node.value): + return True + return False + + +def _calls_function(method, func_name): + """True when the method calls `func_name(...)` (bare name, not attribute).""" + for node in ast.walk(method): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == func_name + ): + return True + return False + + +def _names_in(node): + return {n.id for n in ast.walk(node) if isinstance(n, ast.Name)} + + +def _param_default(method, name): + # Default-value AST node for a named parameter, or None. + args = method.args + params = list(args.args) + list(args.kwonlyargs) + defaults = list(args.defaults) + list(args.kw_defaults) + return dict(zip([p.arg for p in params][-len(defaults) :], defaults)).get(name) + + +def _load_text_only_namespace(): + # Exec the text-only helpers from _utils into one namespace (no unsloth import), + # in dependency order so cross-references resolve. + source = _source(UTILS_PATH) + import transformers + from packaging.version import Version + + ns = { + "copy": copy, + "Version": Version, + "transformers_version": transformers.__version__, + } + funcs = { + node.name: ast.get_source_segment(source, node) + for node in ast.parse(source).body + if isinstance(node, ast.FunctionDef) + } + for name in ( + "resolve_model_class", + "_is_family_text_decoder", + "_remap_text_only_skip_modules", + "_get_text_only_config", + "_get_text_only_key_mapping", + "_apply_text_only_key_mapping", + ): + if name in funcs: + exec(funcs[name], ns) + return ns + + +def _load_text_only_helper(): + return _load_text_only_namespace()["_get_text_only_config"] + + +def test_gemma3_vision_config_resolves_to_text_config(): + transformers = pytest.importorskip("transformers") + helper = _load_text_only_helper() + + config = transformers.Gemma3Config() + text_config = helper(config, "google/gemma-3-27b-it") + + assert isinstance(text_config, transformers.Gemma3TextConfig) + assert text_config.model_type == "gemma3_text" + model_class = transformers.AutoModelForCausalLM._model_mapping[type(text_config)] + assert model_class.__name__ == "Gemma3ForCausalLM" + + +def test_text_only_helper_rejects_configs_without_text_submodel(): + helper = _load_text_only_helper() + + class VisionOnlyConfig: + vision_config = object() + + with pytest.raises(ValueError, match = "Cannot load vision-only as text-only"): + helper(VisionOnlyConfig(), "vision-only") + + +def test_fast_language_model_forwards_text_only_to_fast_model(): + source = _source(LOADER_PATH) + method = _class_method(ast.parse(source), "FastLanguageModel", "from_pretrained") + + # text_only defaults False (opt-in, not forced True), and both FastModel + # delegations forward it. + text_only_default = _param_default(method, "text_only") + assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + + fast_model_calls = [ + node + for node in ast.walk(method) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "from_pretrained" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "FastModel" + ] + assert len(fast_model_calls) == 2 + for call in fast_model_calls: + kw = [k for k in call.keywords if k.arg == "text_only"] + assert len(kw) == 1 + assert isinstance(kw[0].value, ast.Name) and kw[0].value.id == "text_only" + + +def test_fast_model_text_only_does_not_override_explicit_auto_model(): + # AST-based so formatting/refactors that keep the structure do not break it. + source = _source(LOADER_PATH) + method = _class_method(ast.parse(source), "FastModel", "from_pretrained") + + text_only_default = _param_default(method, "text_only") + assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + + # load_text_only is text_only AND a check that the caller did not pass auto_model. + def _is_guarded_bool(value): + names = _names_in(value) + has_none_check = any( + isinstance(n, ast.Compare) and any(isinstance(op, (ast.Is, ast.IsNot)) for op in n.ops) + for n in ast.walk(value) + ) + return "text_only" in names and "auto_model" in names and has_none_check + + assert _assigns_name(method, "load_text_only", _is_guarded_bool) + + assert _calls_function(method, "_get_text_only_config") + + def _forwards_kwarg(node): + return any( + isinstance(n, ast.Call) + and any( + kw.arg == "text_only" + and isinstance(kw.value, ast.Name) + and kw.value.id == "load_text_only" + for kw in n.keywords + ) + for n in ast.walk(node) + ) + + assert _forwards_kwarg(method) + # Falls back to the full model unless the family has its own text decoder. + assert _calls_function(method, "_is_family_text_decoder") + assert _assigns_name( + method, + "load_text_only", + lambda v: isinstance(v, ast.Constant) and v.value is False, + ) + + +def test_fast_base_model_text_only_bypasses_vision_auto_model(): + source = _source(VISION_PATH) + method = _class_method(ast.parse(source), "FastBaseModel", "from_pretrained") + + text_only_default = _param_default(method, "text_only") + assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False + + assert _assigns_name( + method, + "auto_model", + lambda v: isinstance(v, ast.Name) and v.id == "AutoModelForCausalLM", + ) + # Text-only path: strip config, apply the family guard, inject the key remap. + assert _calls_function(method, "_get_text_only_config") + assert _calls_function(method, "_is_family_text_decoder") + assert _calls_function(method, "_apply_text_only_key_mapping") + + +def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower(): + """Tiny end-to-end: build a Gemma3 text-only config, instantiate the + matching model class with shrunken hidden sizes, assert it has the + text language model attributes and no vision tower attribute. + + This is the integration check the AST-only tests were missing -- it + proves the text-only routing actually produces a model that can be + instantiated and that the resulting model is purely text. We use + shrunken hidden sizes so the test is fast and CPU-only. + """ + transformers = pytest.importorskip("transformers") + helper = _load_text_only_helper() + + full_config = transformers.Gemma3Config() + text_config = helper(full_config, "google/gemma-3-27b-it") + + # Shrink for a cheap CPU instantiation; keep the shape attrs read at construction. + text_config.num_hidden_layers = 1 + text_config.hidden_size = 32 + text_config.intermediate_size = 32 + text_config.num_attention_heads = 2 + text_config.num_key_value_heads = 1 + text_config.head_dim = 16 + text_config.vocab_size = 128 + + model_class = transformers.AutoModelForCausalLM._model_mapping[type(text_config)] + model = model_class(text_config) + + # Positive checks: text language model surface is present. + assert hasattr(model, "lm_head"), "text-only Gemma3 model should expose lm_head" + + # Negative checks: no vision tower / multimodal projector remains. + assert not hasattr( + model, "vision_tower" + ), "text-only Gemma3 model should not have a vision_tower" + assert not hasattr( + model, "multi_modal_projector" + ), "text-only Gemma3 model should not have a multi_modal_projector" + + +def test_helper_defined_once_in_utils_and_imported(): + # _get_text_only_config is defined only in _utils and imported by loader + vision. + def _defines(path): + return any( + isinstance(n, ast.FunctionDef) and n.name == "_get_text_only_config" + for n in ast.parse(_source(path)).body + ) + + def _imports(path): + return any( + isinstance(n, ast.ImportFrom) + and n.module == "_utils" + and any(a.name == "_get_text_only_config" for a in n.names) + for n in ast.walk(ast.parse(_source(path))) + ) + + assert _defines(UTILS_PATH) + assert not _defines(LOADER_PATH) and _imports(LOADER_PATH) + assert not _defines(VISION_PATH) and _imports(VISION_PATH) + + +def _load_util_func(name): + ns = _load_text_only_namespace() + if name not in ns: + raise AssertionError(f"{name} not found") + return ns[name] + + +def test_text_only_guard_predicate_across_vlm_families(): + # Text-only is taken only when the resolved class remaps VLM weights. + transformers = pytest.importorskip("transformers") + from transformers import AutoModelForCausalLM + + resolve = _load_util_func("resolve_model_class") + is_family = _load_util_func("_is_family_text_decoder") + helper = _load_text_only_helper() + + def takes_text_only(cfg): + text = helper(cfg, "x") + return resolve(AutoModelForCausalLM, text) is not None and is_family( + getattr(cfg, "model_type", ""), getattr(text, "model_type", "") + ) + + # Dedicated text decoder remaps language_model.* -> strip vision. + assert takes_text_only(transformers.Gemma3Config()) is True + + # No text class (Qwen2-VL/Mllama) or a generic reused decoder that would + # load random weights (Llava/PaliGemma/Idefics3/InternVL) -> keep full model. + for name in [ + "Qwen2VLConfig", + "Qwen2_5_VLConfig", + "MllamaConfig", + "LlavaConfig", + "PaliGemmaConfig", + "Idefics3Config", + "InternVLConfig", + ]: + cfg_cls = getattr(transformers, name, None) + if cfg_cls is None: + continue + assert takes_text_only(cfg_cls()) is False, name + + +def test_text_only_helper_preserves_quantization_config(): + # quantization_config must survive the strip so pre-quantized repos still load. A + # sentinel object avoids a bitsandbytes dependency on transformers 4.51.3. + transformers = pytest.importorskip("transformers") + helper = _load_text_only_helper() + config = transformers.Gemma3Config() + sentinel = object() + config.quantization_config = sentinel + text_config = helper(config, "google/gemma-3-27b-it") + assert getattr(text_config, "quantization_config", None) is sentinel + # The parent's shared text sub-config must not be mutated by the carry-over. + assert getattr(config.get_text_config(), "quantization_config", None) is None + + +def test_text_only_key_mapping_targets_published_prefixes(): + # The mapping must remap the published VLM decoder prefixes and only apply on + # transformers >=5 (on 4.x base_model_prefix handles it and a mapping hurts). + transformers = pytest.importorskip("transformers") + get_key_mapping = _load_util_func("_get_text_only_key_mapping") + mapping = get_key_mapping(transformers.Gemma3Config(), transformers.Gemma3TextConfig()) + if int(transformers.__version__.split(".")[0]) < 5: + assert mapping is None + else: + assert isinstance(mapping, dict) + assert mapping.get(r"^language_model\.model\.") == "model." # gemma3 + assert mapping.get(r"^model\.language_model\.") == "model." # gemma3n + assert mapping.get(r"^language_model\.lm_head\.") == "lm_head." + + +def test_gemma3_text_only_loads_real_language_weights_from_vlm_checkpoint(tmp_path): + # Regression for PR #5816: text-only loading of a Gemma 3 VLM checkpoint must load the + # real language weights, not random ones. Fails on tf >=5 without the key_mapping fix. + transformers = pytest.importorskip("transformers") + torch = pytest.importorskip("torch") + import shutil + from safetensors.torch import load_file, save_file + + get_text_config = _load_text_only_helper() + get_key_mapping = _load_util_func("_get_text_only_key_mapping") + + sentinel = 0.1234 + text_cfg = transformers.Gemma3TextConfig( + hidden_size = 32, + intermediate_size = 64, + num_hidden_layers = 1, + num_attention_heads = 2, + num_key_value_heads = 1, + head_dim = 16, + vocab_size = 128, + max_position_embeddings = 128, + sliding_window = 64, + ) + vision_cfg = transformers.SiglipVisionConfig( + hidden_size = 32, + intermediate_size = 64, + num_hidden_layers = 1, + num_attention_heads = 2, + image_size = 16, + patch_size = 8, + num_channels = 3, + ) + full_config = transformers.Gemma3Config( + text_config = text_cfg.to_dict(), + vision_config = vision_cfg.to_dict(), + ) + full_model = transformers.Gemma3ForConditionalGeneration(full_config) + + state = full_model.state_dict() + text_q = [ + k + for k in state + if "language_model" in k + and "vision" not in k + and k.endswith("layers.0.self_attn.q_proj.weight") + ] + assert text_q, [k for k in state if "q_proj" in k][:5] + with torch.no_grad(): + for k in text_q: + state[k].fill_(sentinel) + + save_dir = tmp_path / "vlm" + full_model.save_pretrained(save_dir, safe_serialization = True) + + # tf >=5 saves under an outer "model." prefix; strip it to reproduce the real + # language_model.model.* layout the published Gemma 3 checkpoints use. + real_dir = tmp_path / "real" + real_dir.mkdir() + weights = {} + for f in save_dir.glob("*.safetensors"): + weights.update(load_file(str(f))) + for f in save_dir.glob("*.bin"): + weights.update(torch.load(f, map_location = "cpu", weights_only = True)) + weights = { + (k[len("model.") :] if k.startswith("model.") else k): v.contiguous() + for k, v in weights.items() + } + for p in save_dir.iterdir(): + if not p.name.endswith((".safetensors", ".bin", ".index.json")): + shutil.copy(p, real_dir / p.name) + save_file(weights, str(real_dir / "model.safetensors")) + + text_config = get_text_config(full_config, "google/gemma-3-27b-it") + load_kwargs = {} + key_mapping = get_key_mapping(full_config, text_config) + if key_mapping is not None: + load_kwargs["key_mapping"] = key_mapping + model = transformers.AutoModelForCausalLM.from_pretrained( + real_dir, + config = text_config, + dtype = torch.float32, + local_files_only = True, + **load_kwargs, + ) + + loaded = model.state_dict() + q_key = [k for k in loaded if k.endswith("model.layers.0.self_attn.q_proj.weight")] + assert q_key, "text decoder q_proj weight missing from the loaded model" + assert float(loaded[q_key[0]].flatten()[0]) == pytest.approx( + sentinel + ), "text weights were randomly initialized instead of loaded from the checkpoint" + assert not any( + "vision_tower" in n for n, _ in model.named_modules() + ), "vision tower should be skipped on the text-only path" diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 486a907478..23723e1d27 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -92,6 +92,7 @@ from platform import system as platform_system platform_system = platform_system() import numpy as np import contextlib +import copy import re from dataclasses import dataclass, field import functools @@ -472,6 +473,98 @@ def resolve_model_class(auto_model, config): return result[0] if isinstance(result, (list, tuple)) else result +def _is_family_text_decoder(parent_model_type, text_model_type): + # True only for the family's own text variant (gemma3 -> gemma3_text); a generic + # reused decoder (llava -> llama) would load random weights, so keep the full model. + return bool(parent_model_type) and str(text_model_type).startswith(parent_model_type) + + +def _get_text_only_config(model_config, model_name): + # Text sub-config of a vision-language config so FastLanguageModel skips the vision tower (PR #5816). + text_config = None + if hasattr(model_config, "get_text_config"): + text_config = model_config.get_text_config() + if text_config is None: + text_config = getattr(model_config, "text_config", None) + if text_config is None: + raise ValueError(f"Cannot load {model_name} as text-only; use FastVisionModel") + # Carry over quantization_config; copy first since get_text_config() shares the parent's object. + qc = getattr(model_config, "quantization_config", None) + if qc is not None and getattr(text_config, "quantization_config", None) is None: + text_config = copy.copy(text_config) + text_config.quantization_config = _remap_text_only_skip_modules(qc) + return text_config + + +def _remap_text_only_skip_modules(qc): + # Remap llm_int8_skip_modules off the VLM wrapper prefix (language_model.model.* -> + # model.*) after text-only stripping, and drop vision/audio entries. See PR #5816. + is_dict = isinstance(qc, dict) + skip = ( + qc.get("llm_int8_skip_modules") if is_dict else getattr(qc, "llm_int8_skip_modules", None) + ) + if not skip: + return qc + remapped = [] + for name in skip: + for pref in ( + "language_model.model.", + "model.language_model.", + "language_model.", + ): + if name.startswith(pref): + name = ( + ("model." + name[len(pref) :]) + if pref != "language_model." + else name[len(pref) :] + ) + break + if name.startswith( + ( + "vision_tower", + "multi_modal_projector", + "audio_tower", + "modality_projection", + ) + ): + continue + remapped.append(name) + remapped = list(dict.fromkeys(remapped)) + qc = dict(qc) if is_dict else copy.copy(qc) + if is_dict: + qc["llm_int8_skip_modules"] = remapped + else: + qc.llm_int8_skip_modules = remapped + return qc + + +def _get_text_only_key_mapping(parent_config, text_config): + # transformers >=5 stopped auto-stripping the VLM wrapper prefix (base_model_prefix + # changed language_model -> model), so remap the text weights onto the decoder keys. + # None on tf <5 (still strips; a mapping would break the load) or non-family. See PR #5816. + if Version(transformers_version) < Version("5.0.0"): + return None + if not _is_family_text_decoder( + getattr(parent_config, "model_type", ""), + getattr(text_config, "model_type", ""), + ): + return None + return { + r"^language_model\.model\.": "model.", + r"^model\.language_model\.": "model.", + r"^language_model\.lm_head\.": "lm_head.", + } + + +def _apply_text_only_key_mapping(kwargs, parent_config, text_config): + # Add the text-only key_mapping to from_pretrained kwargs, under any user mapping. + mapping = _get_text_only_key_mapping(parent_config, text_config) + if not mapping: + return + user_mapping = kwargs.get("key_mapping", None) + kwargs["key_mapping"] = {**mapping, **user_mapping} if user_mapping else mapping + + def resolve_attention_implementation( model_class, config, diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index b0ce28969e..67d1ce9b8d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -97,6 +97,10 @@ from ._utils import ( process_vision_info, unsloth_compile_transformers, fast_inference_setup, + _get_text_only_config, + resolve_model_class, + _is_family_text_decoder, + _apply_text_only_key_mapping, ) # Single source of truth is unsloth_zoo.model_lists. Re-exported so callers @@ -256,6 +260,7 @@ class FastLanguageModel(FastLlamaModel): qat_scheme = None, load_in_fp8 = False, # fp8 LoRA (True, False, 'block') unsloth_tiled_mlp = False, + text_only = False, # Skip vision/audio towers and load only the text decoder *args, **kwargs, ): @@ -342,6 +347,7 @@ class FastLanguageModel(FastLlamaModel): qat_scheme = qat_scheme, load_in_fp8 = load_in_fp8, unsloth_tiled_mlp = unsloth_tiled_mlp, + text_only = text_only, *args, **kwargs, ) @@ -400,7 +406,7 @@ class FastLanguageModel(FastLlamaModel): load_in_8bit, load_in_16bit, ) - model_name = _offline_quantize_to_fp8(model_name, fp8_mode) + model_name = _offline_quantize_to_fp8(model_name, fp8_mode, text_only = text_only) else: assert new_model_name is not None model_name = new_model_name @@ -687,6 +693,7 @@ class FastLanguageModel(FastLlamaModel): qat_scheme = qat_scheme, load_in_fp8 = load_in_fp8, unsloth_tiled_mlp = unsloth_tiled_mlp, + text_only = text_only, *args, **kwargs, ) @@ -868,6 +875,7 @@ class FastModel(FastBaseModel): load_in_fp8 = False, # fp8 LoRA (True, False, 'block') unsloth_tiled_mlp = False, target_parameters = None, # For MoE expert parameters + text_only = False, # Skip vision/audio towers and load only the text decoder *args, **kwargs, ): @@ -997,7 +1005,7 @@ class FastModel(FastBaseModel): load_in_8bit, load_in_16bit, ) - model_name = _offline_quantize_to_fp8(model_name, fp8_mode) + model_name = _offline_quantize_to_fp8(model_name, fp8_mode, text_only = text_only) else: assert new_model_name is not None model_name = new_model_name @@ -1408,6 +1416,29 @@ class FastModel(FastBaseModel): architectures = [] is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures) is_vlm = is_vlm or hasattr(model_config, "vision_config") + load_text_only = text_only and auto_model is None + if load_text_only: + if hasattr(model_config, "vision_config"): + text_config = _get_text_only_config(model_config, old_model_name) + # Skip the vision tower only for families with their own text decoder (Gemma 3); + # others would load random weights, so keep the full model (use FastVisionModel). + text_class = resolve_model_class(AutoModelForCausalLM, text_config) + if text_class is None or not _is_family_text_decoder( + getattr(model_config, "model_type", ""), + getattr(text_config, "model_type", ""), + ): + load_text_only = False + else: + logger.warning_once( + f"Loading {old_model_name} as text-only; vision/audio towers skipped. " + "Use FastVisionModel for multimodal inputs." + ) + # Remap VLM text weights (tf >=5) while model_config is still the parent. #5816 + _apply_text_only_key_mapping(kwargs, model_config, text_config) + model_config = text_config + is_vlm = False + else: + is_vlm = False # If num_labels is set, use AutoModelForSequenceClassification _num_labels = kwargs.get("num_labels", None) if auto_model is None: @@ -1463,6 +1494,7 @@ class FastModel(FastBaseModel): max_lora_rank = max_lora_rank, disable_log_stats = disable_log_stats, load_in_fp8 = load_in_fp8, + text_only = load_text_only, *args, **kwargs, ) diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 98458ad307..06937af7db 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -264,42 +264,73 @@ def get_model_name( return new_model_name -def _offline_quantize_to_fp8(model_name: str, fp8_mode: str) -> str: +def _offline_quantize_to_fp8( + model_name: str, + fp8_mode: str, + *, + text_only: bool = False, +) -> str: """Quantize the model to fp8 via torchao, save to a temp dir, return its path. For vllm >= 0.12.0, prefer dynamic quantization in vllm instead (via hf_overrides={"quantization_config_file": "torchao_config.json"}). """ + from transformers import ( + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + AutoConfig, + ) + + config = AutoConfig.from_pretrained(model_name) + is_vlm = any( + x.endswith(("ForConditionalGeneration", "ForVisionText2Text")) + for x in (getattr(config, "architectures", None) or []) + ) + is_vlm = is_vlm or hasattr(config, "vision_config") + # Decide text-only before the cache name so the fp8 artifact and its path stay in sync. #5816 + text_config = None + if text_only and hasattr(config, "vision_config"): + from ._utils import ( + _get_text_only_config, + resolve_model_class, + _is_family_text_decoder, + ) + + candidate = _get_text_only_config(config, model_name) + text_class = resolve_model_class(AutoModelForCausalLM, candidate) + if text_class is not None and _is_family_text_decoder( + getattr(config, "model_type", ""), + getattr(candidate, "model_type", ""), + ): + text_config = candidate + is_vlm = False + temp_dir = tempfile.gettempdir() - new_model_name = model_name.split("/")[-1] + "-fp8-" + fp8_mode - new_model_name = os.path.join(temp_dir, new_model_name) + # Cache text-only and full-VLM artifacts separately so neither reuses the other. #5816 + cache_name = model_name.split("/")[-1] + "-fp8-" + fp8_mode + if text_config is not None: + cache_name += "-text-only" + new_model_name = os.path.join(temp_dir, cache_name) print(f"Unsloth: Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead") if not os.path.isdir(new_model_name): - from transformers import ( - AutoModelForCausalLM, - AutoModelForImageTextToText, - AutoTokenizer, - AutoProcessor, - TorchAoConfig, - AutoConfig, - ) + from ._utils import _apply_text_only_key_mapping qconfig = _get_torchao_fp8_config(fp8_mode) qconfig = TorchAoConfig(qconfig) - config = AutoConfig.from_pretrained(model_name) - is_vlm = any( - x.endswith(("ForConditionalGeneration", "ForVisionText2Text")) - for x in config.architectures - ) - is_vlm = is_vlm or hasattr(config, "vision_config") + load_kwargs = dict(torch_dtype = "auto", device_map = "auto", quantization_config = qconfig) + if text_config is not None: + _apply_text_only_key_mapping(load_kwargs, config, text_config) + config = text_config auto_model = AutoModelForImageTextToText if is_vlm else AutoModelForCausalLM auto_processor = AutoProcessor if is_vlm else AutoTokenizer model = auto_model.from_pretrained( model_name, - torch_dtype = "auto", - device_map = "auto", - quantization_config = qconfig, + config = config, + **load_kwargs, ) tokenizer = auto_processor.from_pretrained(model_name) model.save_pretrained(new_model_name, safe_serialization = False) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 500c6509a4..daa4a4835c 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -34,6 +34,9 @@ from ._utils import ( _prepare_model_for_qat, resolve_model_class, resolve_attention_implementation, + _get_text_only_config, + _is_family_text_decoder, + _apply_text_only_key_mapping, ) from ._utils import * from .loader_utils import _get_fp8_mode_and_check_settings @@ -583,6 +586,7 @@ class FastBaseModel: disable_log_stats = False, unsloth_vllm_standby = False, load_in_fp8 = False, # fp8 LoRA (True, False, 'block') + text_only = False, **kwargs, ): if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": @@ -597,6 +601,31 @@ class FastBaseModel: if os.environ.get("UNSLOTH_MODEL_NAME", "") == "": os.environ["UNSLOTH_MODEL_NAME"] = model_name.lower() + # Resolve text-only before the is_vlm / vLLM checks so is_vlm stays consistent; + # skip the vision tower only for families with their own text decoder (Gemma 3). #5816 + if text_only and auto_config is None: + auto_config = AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = trust_remote_code, + ) + if text_only and hasattr(auto_config, "vision_config"): + parent_config = auto_config + text_config = _get_text_only_config(parent_config, model_name) + text_class = resolve_model_class(AutoModelForCausalLM, text_config) + if text_class is not None and _is_family_text_decoder( + getattr(parent_config, "model_type", ""), + getattr(text_config, "model_type", ""), + ): + auto_config = text_config + auto_model = AutoModelForCausalLM + _apply_text_only_key_mapping(kwargs, parent_config, text_config) + elif text_only and auto_model in [ + AutoModelForVision2Seq, + AutoModelForImageTextToText, + ]: + # Pure text model requested text-only with a VLM auto class. + auto_model = AutoModelForCausalLM is_vlm = auto_model in [AutoModelForVision2Seq, AutoModelForImageTextToText] is_whisper = whisper_language is not None and whisper_task is not None auto_processor = AutoProcessor if (is_vlm or is_whisper) else AutoTokenizer