Fix Ling compat (#2786)

* Fix LazyTurtle checkpoint path aliases

* Patch remote flash attention compat

* Restore legacy default rope init compat

Author: Qubitium

gptqmodel/models/base.py

@@ -247,6 +247,9 @@ class BaseQModel(nn.Module):
     server = None
 
     support_offload_to_disk = True
+    # Optional runtime->checkpoint prefix overrides used by LazyTurtle when the
+    # execution shell inserts wrapper modules that are absent from the checkpoint.
+    checkpoint_path_aliases: Optional[tuple[tuple[str, str], ...]] = None
     # Optional runtime->checkpoint prefix overrides for LazyTurtle. When unset,
     # the loader derives them from Hugging Face conversion mappings.
     HF_CONVERSION_MAP_REVERSED: Optional[Dict[str, str]] = None

gptqmodel/models/loader.py

@@ -657,6 +657,7 @@ def skip(*args, **kwargs):
                     config=model.config,
                     model_init_kwargs=shell_model_init_kwargs,
                     module_tree=copy.deepcopy(getattr(cls, "module_tree", None)),
+                    checkpoint_path_aliases=copy.deepcopy(getattr(cls, "checkpoint_path_aliases", None)),
                     hf_conversion_map_reversed=copy.deepcopy(
                         cls.resolve_hf_conversion_map_reversed(target_model=model)
                     ),

gptqmodel/utils/hf.py

@@ -576,6 +576,11 @@ def _patch_transformers_remote_code_compat() -> None:
     except Exception:
         utils = None
 
+    try:
+        import transformers.modeling_rope_utils as rope_utils
+    except Exception:
+        rope_utils = None
+
     import transformers.utils.generic as generic
     with _MONKEY_PATCH_LOCK:
         if not hasattr(import_utils, "is_torch_fx_available"):
@@ -599,6 +604,51 @@ def is_flash_attn_greater_or_equal_2_10() -> bool:
 
                 utils.is_flash_attn_greater_or_equal_2_10 = is_flash_attn_greater_or_equal_2_10
 
+        if rope_utils is not None and "default" not in getattr(rope_utils, "ROPE_INIT_FUNCTIONS", {}):
+            # transformers 5.x removed the legacy `"default"` RoPE entrypoint,
+            # but older trust_remote_code model files still resolve it directly.
+            # Recreate the unscaled/base initializer instead of aliasing to
+            # `"linear"` so configs do not need an artificial `factor=1.0`.
+            def _compute_default_rope_parameters_compat(
+                config: Optional["PreTrainedConfig"] = None,
+                device: Optional["torch.device"] = None,
+                seq_len: int | None = None,
+                layer_type: str | None = None,
+            ) -> tuple["torch.Tensor", float]:
+                del seq_len
+                if config is None:
+                    raise ValueError("`config` is required to compute default RoPE parameters.")
+
+                standardize_rope_params = getattr(config, "standardize_rope_params", None)
+                if callable(standardize_rope_params):
+                    standardize_rope_params()
+
+                rope_parameters = getattr(config, "rope_parameters", None)
+                if layer_type is not None and isinstance(rope_parameters, dict):
+                    rope_parameters = rope_parameters.get(layer_type, rope_parameters)
+
+                rope_theta = None
+                partial_rotary_factor = 1.0
+                if isinstance(rope_parameters, dict):
+                    rope_theta = rope_parameters.get("rope_theta")
+                    partial_rotary_factor = rope_parameters.get("partial_rotary_factor", partial_rotary_factor)
+
+                if rope_theta is None:
+                    rope_theta = getattr(config, "rope_theta", None)
+                    if rope_theta is None:
+                        rope_theta = getattr(config, "default_theta", 10_000.0)
+
+                head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+                dim = int(head_dim * partial_rotary_factor)
+                attention_factor = 1.0
+                inv_freq = 1.0 / (
+                    rope_theta
+                    ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+                )
+                return inv_freq, attention_factor
+
+            rope_utils.ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters_compat
+
         if cache_utils is not None and not hasattr(cache_utils, "SlidingWindowCache") and hasattr(cache_utils, "StaticCache"):
             # transformers 5.x folds sliding-window behavior into StaticCache
             # layers, but older remote code still imports the legacy symbol.
@@ -1028,7 +1078,7 @@ def prepare_remote_model_init_compat(model_id_or_path: Optional[str], config: An
     input_mode_enum = getattr(remote_module, "InputMode", None) if remote_module is not None else None
 
     with _MONKEY_PATCH_LOCK:
-        if config.model_type == "minicpm" or config.model_type == "instella":
+        if outer_model_cls is not None:
             try_patch_legacy_flash_attn_flag(outer_model_cls)
 
         if config.model_type == "minicpmv" or config.model_type == "minicpmo":
@@ -1212,7 +1262,8 @@ def try_patch_legacy_flash_attn_flag(model_cls):
         if model_cls is None or not isinstance(model_cls, type):
             return
 
-        # Find the "source class" that defines _supports_flash_attn_2.
+        # Find the most specific class that explicitly declares the newer
+        # `_supports_flash_attn_2` flag used by newer transformers releases.
         base_with_flag = None
         for cls in model_cls.__mro__:
             if "_supports_flash_attn_2" in cls.__dict__:
@@ -1222,8 +1273,15 @@ def try_patch_legacy_flash_attn_flag(model_cls):
         if base_with_flag is None:
             return
 
+        # Respect remote models that already define the legacy flag themselves.
+        for cls in model_cls.__mro__:
+            if cls is base_with_flag:
+                break
+            if "_supports_flash_attn" in cls.__dict__:
+                return
+
         flash_attn_2_val = base_with_flag.__dict__["_supports_flash_attn_2"]
-        setattr(cls, "_supports_flash_attn", bool(flash_attn_2_val))
+        setattr(base_with_flag, "_supports_flash_attn", bool(flash_attn_2_val))
 
 
 def load_tokenizer(tokenizer_or_path, *, model_config: Any = None, **kwargs):

gptqmodel/utils/structure.py

@@ -590,6 +590,7 @@ def __init__(
         config: Any,
         model_init_kwargs: Optional[Dict[str, Any]] = None,
         module_tree: Optional[Any] = None,
+        checkpoint_path_aliases: Optional[Any] = None,
         hf_conversion_map_reversed: Optional[Any] = None,
         target_model: Optional[nn.Module] = None,
     ) -> None:
@@ -602,10 +603,11 @@ def __init__(
         self._module_tree_layer_prefix, self._moe_alias_specs = self._build_moe_alias_specs(self._module_tree)
         # Resolve runtime->checkpoint aliases once up front so per-tensor lookups
         # only need cheap prefix rewrites.
-        alias_items = self._normalize_runtime_to_checkpoint_aliases(
+        alias_items = self._merge_runtime_to_checkpoint_aliases(
+            checkpoint_path_aliases,
             hf_conversion_map_reversed
             if hf_conversion_map_reversed is not None
-            else self.infer_hf_conversion_map_reversed(target_model=target_model)
+            else self.infer_hf_conversion_map_reversed(target_model=target_model),
         )
         self._runtime_to_checkpoint_aliases = tuple(alias_items)
         self._lock = threading.RLock()
@@ -618,6 +620,7 @@ def maybe_create(
         config: Any,
         model_init_kwargs: Optional[Dict[str, Any]] = None,
         module_tree: Optional[Any] = None,
+        checkpoint_path_aliases: Optional[Any] = None,
         hf_conversion_map_reversed: Optional[Any] = None,
         target_model: Optional[nn.Module] = None,
     ) -> Optional["LazyTurtle"]:
@@ -630,6 +633,7 @@ def maybe_create(
                 config=config,
                 model_init_kwargs=model_init_kwargs,
                 module_tree=module_tree,
+                checkpoint_path_aliases=checkpoint_path_aliases,
                 hf_conversion_map_reversed=hf_conversion_map_reversed,
                 target_model=target_model,
             )
@@ -765,6 +769,19 @@ def _normalize_runtime_to_checkpoint_aliases(raw_aliases: Optional[Any]) -> tupl
         alias_items.sort(key=lambda item: len(item[0]), reverse=True)
         return tuple(alias_items)
 
+    @classmethod
+    def _merge_runtime_to_checkpoint_aliases(cls, *raw_alias_sets: Optional[Any]) -> tuple[tuple[str, str], ...]:
+        merged: list[tuple[str, str]] = []
+        seen: set[tuple[str, str]] = set()
+        for raw_aliases in raw_alias_sets:
+            for alias_item in cls._normalize_runtime_to_checkpoint_aliases(raw_aliases):
+                if alias_item in seen:
+                    continue
+                seen.add(alias_item)
+                merged.append(alias_item)
+        merged.sort(key=lambda item: len(item[0]), reverse=True)
+        return tuple(merged)
+
     @classmethod
     def _extract_hf_source_prefix(cls, pattern: Any) -> Optional[str]:
         if not isinstance(pattern, str):

tests/test_hf_config_compat.py

@@ -185,6 +185,27 @@ def test_normalize_hf_config_compat_restores_flash_attn_legacy_version_probe(mon
     assert transformers.utils.is_flash_attn_greater_or_equal_2_10() is True
 
 
+def test_normalize_hf_config_compat_restores_default_rope_init_alias(monkeypatch):
+    import transformers.modeling_rope_utils as rope_utils
+
+    monkeypatch.delitem(rope_utils.ROPE_INIT_FUNCTIONS, "default", raising=False)
+
+    config = SimpleNamespace(
+        rope_parameters={"rope_type": "default", "rope_theta": 10000.0},
+        rope_theta=10000.0,
+        hidden_size=64,
+        num_attention_heads=8,
+        head_dim=8,
+    )
+
+    normalize_hf_config_compat(config, trust_remote_code=True)
+
+    inv_freq, attention_factor = rope_utils.ROPE_INIT_FUNCTIONS["default"](config, None)
+
+    torch.testing.assert_close(inv_freq, torch.tensor([1.0, 0.1, 0.01, 0.001]))
+    assert attention_factor == 1.0
+
+
 def test_normalize_hf_config_compat_restores_legacy_cache_length_helpers(monkeypatch):
     monkeypatch.delattr(cache_utils.Cache, "get_max_length", raising=False)
     monkeypatch.delattr(cache_utils.Cache, "get_usable_length", raising=False)
@@ -377,6 +398,31 @@ def tie_weights(self):
     assert getattr(DummyRemoteModel, "_gptqmodel_tie_weights_kwargs_patch", False) is True
 
 
+def test_prepare_remote_model_init_compat_backfills_legacy_flash_attn_flag(monkeypatch):
+    class DummyRemoteBase:
+        _supports_flash_attn_2 = True
+
+    class DummyRemoteModel(DummyRemoteBase):
+        __module__ = "transformers_modules.fake_bailing.modeling_bailing_moe_v2"
+
+    monkeypatch.setattr(
+        "transformers.dynamic_module_utils.get_class_from_dynamic_module",
+        lambda class_ref, model_id_or_path, **kwargs: DummyRemoteModel,
+    )
+
+    config = SimpleNamespace(
+        model_type="bailing_moe",
+        auto_map={"AutoModelForCausalLM": "modeling_bailing_moe_v2.BailingMoeV2ForCausalLM"},
+    )
+
+    assert "_supports_flash_attn" not in DummyRemoteBase.__dict__
+
+    prepare_remote_model_init_compat("/tmp/ling", config)
+
+    assert DummyRemoteBase._supports_flash_attn is True
+    assert DummyRemoteModel._supports_flash_attn is True
+
+
 def test_prepare_remote_model_init_compat_accepts_tokenizers_backend_for_ovis(monkeypatch):
     class DummyRemoteModel:
         __module__ = "transformers_modules.fake_ovis.modeling_ovis"

tests/test_local_model_paths.py

@@ -496,6 +496,7 @@ class DummyQModel:
         require_pkgs = []
         supports_desc_act = [True, False]
         support_offload_to_disk = True
+        checkpoint_path_aliases = (("shell_model", "model"),)
         config_class = None
 
         @staticmethod
@@ -530,6 +531,7 @@ def __init__(self, model, **kwargs):
     assert shell_configs
     assert load_calls == []
     assert isinstance(instance.turtle_model, LazyTurtle)
+    assert instance.turtle_model._runtime_to_checkpoint_aliases == DummyQModel.checkpoint_path_aliases
     assert instance.turtle_model.config._experts_implementation == "linear_loop"
     assert instance.turtle_model.config is not instance.model.config