fix: harden quantized checkpoint integration

tests/test_loader_boundaries.py

@@ -23,6 +23,7 @@ def test_compatibility_shims_tolerate_independent_missing_or_broken_imports(monk
     original_import = builtins.__import__
     failure_keys = {
         ("transformers", ("AutoModelForImageTextToText",)): ImportError,
+        ("transformers", ("PreTrainedModel",)): RuntimeError,
         ("transformers.utils", ("output_capturing",)): ImportError,
         ("transformers.utils.import_utils", ()): RuntimeError,
         ("transformers.utils", ()): RuntimeError,
@@ -109,6 +110,72 @@ def test_compatibility_shims_tolerate_rejected_generic_attribute_patch(monkeypat
     assert _execute_loader_source()["TASK_MODEL_MAP"]
 
 
+def test_all_tied_weights_compat_aggregates_filters_and_caches():
+    class LegacyModel(torch.nn.Module):
+        pass
+
+    assert loader._install_all_tied_weights_keys_compat(LegacyModel) is True
+    assert loader._install_all_tied_weights_keys_compat(LegacyModel) is False
+
+    model = LegacyModel()
+    model.lm_head = torch.nn.Linear(2, 2)
+    model._tied_weights_keys = ["lm_head.weight", "missing.weight"]
+    child = torch.nn.Module()
+    child._tied_weights_keys = {"lm_head.bias": "lm_head.weight"}
+    model.child = child
+
+    assert model.all_tied_weights_keys == {
+        "lm_head.weight": "lm_head.weight",
+        "lm_head.bias": "lm_head.weight",
+    }
+    model.all_tied_weights_keys = {"cached.weight": "source.weight"}
+    assert model.all_tied_weights_keys == {"cached.weight": "source.weight"}
+
+
+def test_all_tied_weights_compat_prefers_transformers_expansion_contract():
+    class ExpandableModel(torch.nn.Module):
+        def get_expanded_tied_weights_keys(self, *, all_submodels):
+            assert all_submodels is True
+            return {"lm_head.weight": "model.embed_tokens.weight"}
+
+    assert loader._install_all_tied_weights_keys_compat(ExpandableModel) is True
+    model = ExpandableModel()
+
+    assert model.all_tied_weights_keys == {
+        "lm_head.weight": "model.embed_tokens.weight",
+    }
+    assert model.__dict__["_all_tied_weights_keys_store"] == {
+        "lm_head.weight": "model.embed_tokens.weight",
+    }
+
+
+def test_all_tied_weights_compat_falls_back_when_expansion_is_legacy():
+    class LegacyExpandableModel(torch.nn.Module):
+        def get_expanded_tied_weights_keys(self, *, all_submodels):
+            raise TypeError("legacy list contract")
+
+    assert loader._install_all_tied_weights_keys_compat(LegacyExpandableModel) is True
+    model = LegacyExpandableModel()
+    model.lm_head = torch.nn.Linear(2, 2)
+    model._tied_weights_keys = ["lm_head.weight"]
+
+    assert model.all_tied_weights_keys == {"lm_head.weight": "lm_head.weight"}
+
+
+def test_flash_attention_config_helpers_recurse_into_transformers_subconfigs():
+    from transformers import PretrainedConfig
+
+    outer = PretrainedConfig()
+    inner = PretrainedConfig()
+    outer.inner = inner
+    inner._attn_implementation = "flash_attention_2"
+
+    assert loader._config_uses_flash_attention_2(outer) is True
+    loader._force_eager_config(outer)
+    assert inner._attn_implementation == "eager"
+    assert loader._config_uses_flash_attention_2(outer) is False
+
+
 def _config(**overrides):
     values = {
         "model_type": "gpt2",
@@ -200,6 +267,158 @@ def test_revision_trust_and_offline_flags_reach_every_provider(loader_boundary):
     loader.dev.empty_cache.assert_called_once_with()
 
 
+def test_unavailable_flash_attention_falls_back_to_eager(
+    loader_boundary, monkeypatch,
+):
+    import transformers.utils
+
+    loader_boundary.config._attn_implementation = "flash_attention_2"
+    monkeypatch.setattr(
+        transformers.utils,
+        "is_flash_attn_2_available",
+        lambda: False,
+    )
+    monkeypatch.setattr(
+        loader.qd,
+        "detect_quant_scheme",
+        lambda *args, **kwargs: loader.qd.QuantDetection(loader.qd.QuantScheme.NONE),
+    )
+
+    loader.load_model("local/model", skip_snapshot=True)
+
+    assert loader_boundary.config._attn_implementation == "eager"
+
+
+def test_broken_flash_attention_probe_falls_back_to_eager(
+    loader_boundary, monkeypatch,
+):
+    loader_boundary.config._attn_implementation = "flash_attention_2"
+    real_import = builtins.__import__
+
+    def rejecting_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if (
+            name == "transformers.utils"
+            and "is_flash_attn_2_available" in tuple(fromlist or ())
+        ):
+            raise RuntimeError("flash-attention availability probe failed")
+        return real_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", rejecting_import)
+    loader.load_model("local/model", skip_snapshot=True)
+
+    assert loader_boundary.config._attn_implementation == "eager"
+
+
+def test_available_flash_attention_preserves_requested_implementation(
+    loader_boundary, monkeypatch,
+):
+    import transformers.utils
+
+    loader_boundary.config._attn_implementation = "flash_attention_2"
+    monkeypatch.setattr(
+        transformers.utils,
+        "is_flash_attn_2_available",
+        lambda: True,
+    )
+
+    loader.load_model("local/model", skip_snapshot=True)
+
+    assert loader_boundary.config._attn_implementation == "flash_attention_2"
+
+
+def test_multimodal_unwrap_generation_cache_and_legacy_ties(
+    loader_boundary, monkeypatch,
+):
+    class FakeLanguageModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.config = _config(tie_word_embeddings=False)
+            self._tied_weights_keys = ["lm_head.weight"]
+            self.last_cache_position = None
+
+        def prepare_inputs_for_generation(
+            self, input_ids, *args, cache_position=None, **kwargs,
+        ):
+            self.last_cache_position = cache_position
+            return {"input_ids": input_ids, "cache_position": cache_position}
+
+    class MultimodalWrapper(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = FakeLanguageModel()
+
+        def forward(self, input_ids, pixel_values):
+            return input_ids, pixel_values
+
+    wrapper = MultimodalWrapper()
+    loader_boundary.model_class.from_pretrained.return_value = wrapper
+    monkeypatch.setattr(loader, "PreTrainedModel", FakeLanguageModel)
+    monkeypatch.setattr(
+        loader.qd,
+        "detect_quant_scheme",
+        lambda *args, **kwargs: loader.qd.QuantDetection(loader.qd.QuantScheme.NONE),
+    )
+
+    handle = loader.load_model("local/model", skip_snapshot=True)
+
+    assert handle.model is wrapper.language_model
+    assert handle.model._tied_weights_keys == {}
+    input_ids = torch.ones(1, 2, dtype=torch.long)
+
+    class Cache:
+        @staticmethod
+        def get_seq_length():
+            return 4
+
+    prepared = handle.model.prepare_inputs_for_generation(
+        input_ids, past_key_values=Cache(),
+    )
+    assert prepared["cache_position"].tolist() == [4, 5]
+
+    legacy_cache = ((torch.zeros(1, 1, 3, 1),),)
+    prepared = handle.model.prepare_inputs_for_generation(
+        input_ids, past_key_values=legacy_cache,
+    )
+    assert prepared["cache_position"].tolist() == [3, 4]
+
+    prepared = handle.model.prepare_inputs_for_generation(
+        input_ids, past_key_values=object(),
+    )
+    assert prepared["cache_position"].tolist() == [0, 1]
+
+    prepared = handle.model.prepare_inputs_for_generation(input_ids)
+    assert prepared["cache_position"].tolist() == [0, 1]
+
+    explicit = torch.tensor([7, 8])
+    prepared = handle.model.prepare_inputs_for_generation(
+        input_ids, cache_position=explicit,
+    )
+    assert prepared["cache_position"] is explicit
+
+
+def test_multimodal_signature_probe_failure_preserves_wrapper(
+    loader_boundary, monkeypatch,
+):
+    class FakeLanguageModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.config = _config()
+
+    class WrapperWithUninspectableForward(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = FakeLanguageModel()
+            self.forward = object()
+
+    wrapper = WrapperWithUninspectableForward()
+    loader_boundary.model_class.from_pretrained.return_value = wrapper
+    monkeypatch.setattr(loader, "PreTrainedModel", FakeLanguageModel)
+
+    handle = loader.load_model("local/model", skip_snapshot=True)
+
+    assert handle.model is wrapper
+
+
 def test_hf_token_is_forwarded_without_logging_value(loader_boundary, monkeypatch):
     monkeypatch.setenv("HF_TOKEN", "secret-token")
     loader.load_model("x", skip_snapshot=True)