Coverage for transformer_lens/model_bridge/supported_architectures/_remote_code_compat.py: 99%
66 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Shared transformers-v5 compatibility helpers for remote-code adapters.
3Modeling files loaded via ``trust_remote_code`` were mostly written against
4transformers 4.x and break under v5 in recurring ways: the meta-device
5load-then-materialise flow re-runs ``_init_weights`` over already-loaded
6modules, ``tie_weights`` calls ``.keys()`` on list-form ``_tied_weights_keys``,
7``all_tied_weights_keys`` lookups hit broken remote ``__getattr__``s, and the
8``ROPE_INIT_FUNCTIONS["default"]`` entry the remote code looks up was removed.
9The patch mechanics live here; each adapter's ``prepare_loading`` keeps the
10WHY — what its remote code gets wrong and which classes need patching.
12These paths only fire on real ``trust_remote_code`` loads (CI-invisible), so
13the helpers are unit-tested directly in
14``tests/unit/model_bridge/test_remote_code_compat.py``.
15"""
17import sys
18from collections.abc import Iterator
19from types import ModuleType
20from typing import Any
22import torch
23from transformers import PreTrainedModel
25# Architectures whose prepare_loading force-imports the remote modeling file at
26# the caller's pinned revision (code_revision, falling back to the weights
27# revision) so the patched module copy is the one from_pretrained executes.
28# Kept in sync with the adapter sources by a static scan in
29# tests/unit/model_bridge/test_prepare_loading_orchestration.py.
30REVISION_FORWARDING_ARCHITECTURES: frozenset[str] = frozenset(
31 {
32 "BD3LM",
33 "BaiChuanForCausalLM",
34 "BaichuanForCausalLM",
35 "DreamModel",
36 "GiddForDiffusionLM",
37 "InternLM2ForCausalLM",
38 "LLaDA2MoeModelLM",
39 "OpenELMForCausalLM",
40 "OuroForCausalLM",
41 "RavenForCausalLM",
42 "RWKV7ForCausalLM",
43 }
44)
47def iter_remote_modeling_modules(*name_fragments: str) -> Iterator[ModuleType]:
48 """Yield imported modules whose name contains ``modeling`` and any fragment.
50 Each checkpoint revision of a remote-code repo imports its own copy of the
51 modeling file (distinct ``transformers_modules.*`` entries), so class-level
52 patches must be applied to every one. Matching is case-insensitive.
53 """
54 fragments = tuple(fragment.lower() for fragment in name_fragments)
55 for key in list(sys.modules.keys()):
56 key_lower = key.lower()
57 if "modeling" not in key_lower:
58 continue
59 if not any(fragment in key_lower for fragment in fragments):
60 continue
61 module = sys.modules.get(key)
62 if module is not None: 62 ↛ 55line 62 didn't jump to line 55 because the condition on line 62 was always true
63 yield module
66def force_import_remote_class(model_name: str, dotted_ref: str, **kwargs: Any) -> type | None:
67 """Import a remote-code class so its module lands in ``sys.modules`` to patch.
69 ``dotted_ref`` is ``"<module_file>.<ClassName>"`` as understood by
70 ``get_class_from_dynamic_module``; extra kwargs are forwarded to it.
71 Returns the class, or None when the dynamic module is unavailable
72 (offline, renamed, not a remote-code repo) — callers treat that as
73 "nothing to patch".
74 """
75 try:
76 from transformers.dynamic_module_utils import get_class_from_dynamic_module
78 cls = get_class_from_dynamic_module(dotted_ref, model_name, **kwargs)
79 assert isinstance(cls, type)
80 return cls
81 except Exception:
82 return None
85def patch_init_weights_skip_loaded(cls: type[Any]) -> None:
86 """Wrap ``cls._init_weights`` to skip modules already loaded from checkpoint.
88 v5's meta-device flow calls ``_init_weights`` on modules that already hold
89 checkpoint weights (remote-code modules lack ``_is_hf_initialized``),
90 re-randomising them; the wrapper only lets initialisation through for
91 modules still on the meta device. Idempotent via the ``_tl_patched``
92 sentinel.
94 Raises:
95 ValueError: for ``transformers.PreTrainedModel`` itself and for classes
96 without their own ``_init_weights``. Remote modules re-export the
97 HF base under local names, and patching it would disable weight
98 init — including HF's rotary-buffer restoration — for every model
99 loaded later in the process.
100 """
101 if getattr(cls, "_tl_patched", False):
102 return
103 if cls is PreTrainedModel:
104 raise ValueError(
105 "Refusing to patch transformers.PreTrainedModel itself: that would "
106 "disable weight init for every model loaded later in this process. "
107 "Pass the remote code's own PreTrainedModel subclass."
108 )
109 if "_init_weights" not in cls.__dict__:
110 raise ValueError(
111 f"Refusing to patch {cls.__name__}: it defines no _init_weights of "
112 "its own, so the wrap would capture an inherited (possibly HF base) "
113 "implementation."
114 )
116 original_init_weights = cls._init_weights
118 def safe_init_weights(self: Any, mod: Any, _original: Any = original_init_weights) -> None:
119 # Only initialise modules still on meta device (pre-loading); never
120 # re-randomise weights already read from the checkpoint.
121 first_param = next(mod.parameters(), None)
122 if first_param is not None and first_param.device.type != "meta":
123 return
124 _original(self, mod)
126 cls._init_weights = safe_init_weights
127 cls._tl_patched = True
130def retie_weights_keys_v5(cls: Any, mapping: dict[str, str]) -> None:
131 """Rewrite a 4.x list-form ``_tied_weights_keys`` to the v5 dict form.
133 v5's ``tie_weights`` -> ``get_expanded_tied_weights_keys`` calls ``.keys()``
134 on the attribute and raises ``AttributeError`` on the legacy list. No-op
135 when ``cls`` is None or the attribute is already a dict (or absent).
136 """
137 if cls is not None and isinstance(getattr(cls, "_tied_weights_keys", None), list):
138 cls._tied_weights_keys = mapping
141def disable_tied_weights_lookup(cls: type[Any]) -> None:
142 """Give ``cls`` an empty ``all_tied_weights_keys``.
144 Some remote ``__getattr__`` implementations fail to delegate v5's
145 ``all_tied_weights_keys`` lookup back to ``PreTrainedModel`` and raise;
146 the affected checkpoints are untied anyway.
147 """
148 setattr(cls, "all_tied_weights_keys", {})
151def compute_default_rope_inv_freq(
152 config: Any = None,
153 device: Any = None,
154 seq_len: Any = None,
155 **rope_kwargs: Any,
156) -> tuple[torch.Tensor, float]:
157 """transformers-v4 ``_compute_default_rope_parameters``, removed in v5.
159 Standard (unscaled) RoPE inverse frequencies with the v4 call contract the
160 remote code targets — ``(config, device) -> (inv_freq, attention_scaling)``
161 — plus v4's kwargs-only fallback (``base``/``dim`` passed directly).
162 Registered per remote modeling module, never in the shared
163 ``ROPE_INIT_FUNCTIONS`` (see ``restore_default_rope_init``).
164 """
165 if config is not None:
166 base = config.rope_theta
167 partial_rotary_factor = getattr(config, "partial_rotary_factor", None) or 1.0
168 head_dim = getattr(config, "head_dim", None) or (
169 config.hidden_size // config.num_attention_heads
170 )
171 dim = int(head_dim * partial_rotary_factor)
172 else:
173 base = rope_kwargs["base"]
174 dim = rope_kwargs["dim"]
175 inv_freq = 1.0 / (
176 base
177 ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
178 )
179 return inv_freq, 1.0
182def restore_default_rope_init(
183 model_name: str,
184 dotted_ref: str,
185 rotary_class_name: str,
186 revision: str | None = None,
187) -> None:
188 """Restore v4's ``ROPE_INIT_FUNCTIONS["default"]`` for one remote-code model.
190 ``dotted_ref`` (``"<modeling_file>.<ClassName>"``) is force-imported at the
191 caller's ``revision`` (each revision is its own module copy); every
192 loaded copy of that modeling file gets its module-level
193 ``ROPE_INIT_FUNCTIONS`` rebound to a copy with "default" restored, and its
194 ``rotary_class_name`` gets ``compute_default_rope_parameters`` for v5's
195 ``_init_weights``. The shared transformers dict must stay untouched: from
196 transformers 5.17, ``_init_weights`` lets its entries override every
197 native model's own ``compute_default_rope_parameters``.
198 """
199 if force_import_remote_class(model_name, dotted_ref, revision=revision) is None:
200 return
201 for module in iter_remote_modeling_modules(dotted_ref.rsplit(".", 1)[0]):
202 rope_functions = getattr(module, "ROPE_INIT_FUNCTIONS", None)
203 if rope_functions is not None and "default" not in rope_functions:
204 setattr(
205 module,
206 "ROPE_INIT_FUNCTIONS",
207 {**rope_functions, "default": compute_default_rope_inv_freq},
208 )
209 rope_class = getattr(module, rotary_class_name, None)
210 if rope_class is not None and not hasattr(rope_class, "compute_default_rope_parameters"):
211 rope_class.compute_default_rope_parameters = staticmethod(compute_default_rope_inv_freq)