fix(vlm): preload falcon config for eager attention

Signed-off-by: Georg Heiler <georg.kf.heiler@gmail.com>

Author: geoHeil

docling/models/inference_engines/vlm/transformers_engine.py

@@ -11,6 +11,7 @@
 from packaging import version
 from PIL.Image import Image
 from transformers import (
+    AutoConfig,
     AutoModel,
     AutoModelForCausalLM,
     AutoModelForImageTextToText,
@@ -120,6 +121,23 @@ def _falcon_ocr_category_from_prompt(prompt: str) -> str:
     return "plain"
 
 
+def _force_falcon_ocr_eager_attention_config(config_obj: Any) -> None:
+    if config_obj is None:
+        return
+    if getattr(config_obj, "_attn_implementation", None) in {
+        None,
+        "sdpa",
+        "paged|sdpa",
+    }:
+        config_obj._attn_implementation = "eager"
+    if getattr(config_obj, "_attn_implementation_internal", None) in {
+        None,
+        "sdpa",
+        "paged|sdpa",
+    }:
+        config_obj._attn_implementation_internal = "eager"
+
+
 class TransformersVlmEngine(BaseVlmEngine, HuggingFaceModelDownloadMixin):
     """HuggingFace Transformers engine for VLM inference.
 
@@ -251,6 +269,17 @@ def download_wrapper(repo_id: str, revision: str) -> Path:
         elif model_type == TransformersModelType.AUTOMODEL_IMAGETEXTTOTEXT:
             model_cls = AutoModelForImageTextToText  # type: ignore[assignment]
 
+        attn_implementation = self._get_attn_implementation()
+        model_config = None
+        if _value_mentions_falcon_ocr(repo_id):
+            model_config = AutoConfig.from_pretrained(
+                artifacts_path,
+                trust_remote_code=self.options.trust_remote_code,
+                revision=revision,
+                attn_implementation=attn_implementation,
+            )
+            _force_falcon_ocr_eager_attention_config(model_config)
+
         # Load processor
         self.processor = AutoProcessor.from_pretrained(
             artifacts_path,
@@ -269,10 +298,11 @@ def download_wrapper(repo_id: str, revision: str) -> Path:
             artifacts_path,
             device_map=self.device,
             dtype=torch_dtype,
-            attn_implementation=self._get_attn_implementation(),
+            attn_implementation=attn_implementation,
             trust_remote_code=self.options.trust_remote_code,
             revision=revision,
             quantization_config=quantization_config,
+            config=model_config,
         )
 
         self.vlm_model.eval()

tests/test_transformers_vlm_engine.py

@@ -30,6 +30,15 @@ def fake_model_from_pretrained(*args, **kwargs):
         captured["model_kwargs"] = kwargs
         return FakeModel()
 
+    def fake_auto_config_from_pretrained(*args, **kwargs):
+        _ = args
+        captured["config_kwargs"] = kwargs
+        return SimpleNamespace(
+            model_type="falcon_ocr",
+            _attn_implementation=kwargs.get("attn_implementation"),
+            _attn_implementation_internal=kwargs.get("attn_implementation"),
+        )
+
     def fake_generation_config_from_pretrained(*args, **kwargs):
         return SimpleNamespace()
 
@@ -45,6 +54,10 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
         "docling.models.inference_engines.vlm.transformers_engine.AutoModelForCausalLM.from_pretrained",
         fake_model_from_pretrained,
     )
+    monkeypatch.setattr(
+        "docling.models.inference_engines.vlm.transformers_engine.AutoConfig.from_pretrained",
+        fake_auto_config_from_pretrained,
+    )
     monkeypatch.setattr(
         "docling.models.inference_engines.vlm.transformers_engine.GenerationConfig.from_pretrained",
         fake_generation_config_from_pretrained,
@@ -67,6 +80,8 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
     )
 
     assert captured["model_kwargs"]["attn_implementation"] == "eager"
+    assert captured["config_kwargs"]["attn_implementation"] == "eager"
+    assert captured["model_kwargs"]["config"]._attn_implementation == "eager"
 
 
 def test_transformers_engine_accepts_legacy_private_attn_implementation_key(
@@ -86,6 +101,15 @@ def fake_model_from_pretrained(*args, **kwargs):
         captured["model_kwargs"] = kwargs
         return FakeModel()
 
+    def fake_auto_config_from_pretrained(*args, **kwargs):
+        _ = args
+        captured["config_kwargs"] = kwargs
+        return SimpleNamespace(
+            model_type="falcon_ocr",
+            _attn_implementation=kwargs.get("attn_implementation"),
+            _attn_implementation_internal=kwargs.get("attn_implementation"),
+        )
+
     def fake_generation_config_from_pretrained(*args, **kwargs):
         return SimpleNamespace()
 
@@ -101,6 +125,10 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
         "docling.models.inference_engines.vlm.transformers_engine.AutoModelForCausalLM.from_pretrained",
         fake_model_from_pretrained,
     )
+    monkeypatch.setattr(
+        "docling.models.inference_engines.vlm.transformers_engine.AutoConfig.from_pretrained",
+        fake_auto_config_from_pretrained,
+    )
     monkeypatch.setattr(
         "docling.models.inference_engines.vlm.transformers_engine.GenerationConfig.from_pretrained",
         fake_generation_config_from_pretrained,
@@ -123,6 +151,8 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
     )
 
     assert captured["model_kwargs"]["attn_implementation"] == "eager"
+    assert captured["config_kwargs"]["attn_implementation"] == "eager"
+    assert captured["model_kwargs"]["config"]._attn_implementation == "eager"
 
 
 def test_transformers_engine_defaults_falcon_ocr_to_eager(
@@ -142,6 +172,15 @@ def fake_model_from_pretrained(*args, **kwargs):
         captured["model_kwargs"] = kwargs
         return FakeModel()
 
+    def fake_auto_config_from_pretrained(*args, **kwargs):
+        _ = args
+        captured["config_kwargs"] = kwargs
+        return SimpleNamespace(
+            model_type="falcon_ocr",
+            _attn_implementation=kwargs.get("attn_implementation"),
+            _attn_implementation_internal=kwargs.get("attn_implementation"),
+        )
+
     def fake_generation_config_from_pretrained(*args, **kwargs):
         return SimpleNamespace()
 
@@ -157,6 +196,10 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
         "docling.models.inference_engines.vlm.transformers_engine.AutoModelForCausalLM.from_pretrained",
         fake_model_from_pretrained,
     )
+    monkeypatch.setattr(
+        "docling.models.inference_engines.vlm.transformers_engine.AutoConfig.from_pretrained",
+        fake_auto_config_from_pretrained,
+    )
     monkeypatch.setattr(
         "docling.models.inference_engines.vlm.transformers_engine.GenerationConfig.from_pretrained",
         fake_generation_config_from_pretrained,
@@ -178,6 +221,8 @@ def fake_generation_config_from_pretrained(*args, **kwargs):
     )
 
     assert captured["model_kwargs"]["attn_implementation"] == "eager"
+    assert captured["config_kwargs"]["attn_implementation"] == "eager"
+    assert captured["model_kwargs"]["config"]._attn_implementation == "eager"
 
 
 def test_transformers_engine_falls_back_without_generation_config_file(
@@ -199,8 +244,18 @@ def fake_processor_from_pretrained(*args, **kwargs):
     def fake_model_from_pretrained(*args, **kwargs):
         model = FakeModel()
         captured["model"] = model
+        captured["model_kwargs"] = kwargs
         return model
 
+    def fake_auto_config_from_pretrained(*args, **kwargs):
+        _ = args
+        captured["config_kwargs"] = kwargs
+        return SimpleNamespace(
+            model_type="falcon_ocr",
+            _attn_implementation=kwargs.get("attn_implementation"),
+            _attn_implementation_internal=kwargs.get("attn_implementation"),
+        )
+
     def fake_generation_config_from_pretrained(*args, **kwargs):
         _ = (args, kwargs)
         raise OSError("missing file named generation_config.json")
@@ -221,6 +276,10 @@ def fake_generation_config_from_model_config(model_config):
         "docling.models.inference_engines.vlm.transformers_engine.AutoModelForCausalLM.from_pretrained",
         fake_model_from_pretrained,
     )
+    monkeypatch.setattr(
+        "docling.models.inference_engines.vlm.transformers_engine.AutoConfig.from_pretrained",
+        fake_auto_config_from_pretrained,
+    )
     monkeypatch.setattr(
         "docling.models.inference_engines.vlm.transformers_engine.GenerationConfig.from_pretrained",
         fake_generation_config_from_pretrained,
@@ -247,6 +306,8 @@ def fake_generation_config_from_model_config(model_config):
 
     assert engine.generation_config.source == "fallback"
     assert captured["fallback_model_config"] is captured["model"].config
+    assert captured["config_kwargs"]["attn_implementation"] == "eager"
+    assert captured["model_kwargs"]["config"]._attn_implementation == "eager"
 
 
 def test_transformers_engine_uses_falcon_native_generate_batch() -> None: