feat: support text-only loading of Gemma 3 27B via FastLanguageModel (skip SiglipVisionModel) (#5816)

* feat: support text-only loading of Gemma 3 27B via FastLanguageModel (skip SiglipVisionModel) * test: instantiate text-only Gemma3 model and assert no vision tower Existing tests were AST source-introspection plus a config-resolves-to- text-config check; none actually instantiated a model from the text-only config. Add a small integration test that builds a shrunken Gemma3TextConfig (CPU-cheap), instantiates the matching CausalLM class, and asserts the resulting model exposes the LM head and has no vision_tower or multi_modal_projector attribute. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Deduplicate _get_text_only_config into _utils for PR #5816 * Fall back to full model when a VLM has no text-only class for PR #5816 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Preserve quantization_config and clarify warning for text-only loading for PR #5816 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Only take text-only path when the VLM has its own text decoder for PR #5816 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Convert source string-match assertions to AST checks per Gemini review * Load real VLM text weights on transformers 5.x for text-only mode in PR #5816 transformers >=5 changed Gemma3ForCausalLM base_model_prefix from language_model to model, so a VLM checkpoint's text weights (gemma3: language_model.model.*, gemma3n: model.language_model.*) no longer auto-strip onto the text decoder and were silently initialized random. Add a version-gated key_mapping that remaps them onto the text keys, returning None on transformers <5 where the prefix still strips and a mapping would break the load. Apply the same family-guarded remap on the load_in_fp8 offline path and for direct FastBaseModel callers, and remap quantization llm_int8_skip_modules off the wrapper prefix after stripping. Add a regression test that loads real VLM checkpoint weights (the prior tests only instantiated a fresh model so they missed this) and drop the bitsandbytes dependency from the quantization-config test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Separate FP8 text-only cache and hoist the text-only guard for PR #5816 Address review of the text-only changes: (1) _offline_quantize_to_fp8 produced different artifacts for text-only vs full VLM but reused the same <name>-fp8-<mode> cache dir, so one mode could load the other's saved model; decide text-only before the cache name and add a -text-only suffix. (2) FastBaseModel.from_pretrained rewrote the VLM auto class to AutoModelForCausalLM before loading auto_config and before the family check, leaving is_vlm wrong for the fast_inference/vLLM block; hoist the family-guarded text-only decision above those checks and drop the redundant later block. (3) Wire the text-only regression test into the curated CPU pytest job so it runs in CI across the transformers matrix. * Trim text-only code comments for PR #5816 Shorten and de-duplicate the comments added for the text-only loading work; keep the non-obvious rationale (the transformers >=5 base_model_prefix change) and drop the obvious parts. Comments only, no code changes; AST-based tests still pass on transformers 4.57.6 and 5.4.0. * Make text-only loading opt-in via a public text_only argument for PR #5816 Rename the internal _force_text_only flag to a public text_only parameter on FastLanguageModel, FastModel and FastBaseModel (and the fp8 helper), defaulting False on all three. Text-only loading is now opt-in (text_only=True) instead of forced on by FastLanguageModel; the family guard and key remap are unchanged. Updated the AST tests for the new parameter and forwarding. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim text-only code comments for clarity --------- Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-09 22:52:39 -07:00 · 2026-06-09 22:52:39 -07:00 · 03349d1e05
commit 03349d1e05
parent 8848a310df
6 changed files with 647 additions and 29 deletions
--- a/tests/python/test_fast_language_model_text_only.py
+++ 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 = <value>` 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"