Coverage for transformer_lens/model_bridge/supported_architectures/_remote_code_compat.py: 99%

1"""Shared transformers-v5 compatibility helpers for remote-code adapters. 

2 

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. 

11 

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""" 

16 

17import sys 

18from collections.abc import Iterator 

19from types import ModuleType 

20from typing import Any 

21 

22import torch 

23from transformers import PreTrainedModel 

24 

25 

26def iter_remote_modeling_modules(*name_fragments: str) -> Iterator[ModuleType]: 

27 """Yield imported modules whose name contains ``modeling`` and any fragment. 

28 

29 Each checkpoint revision of a remote-code repo imports its own copy of the 

30 modeling file (distinct ``transformers_modules.*`` entries), so class-level 

31 patches must be applied to every one. Matching is case-insensitive. 

32 """ 

33 fragments = tuple(fragment.lower() for fragment in name_fragments) 

34 for key in list(sys.modules.keys()): 

35 key_lower = key.lower() 

36 if "modeling" not in key_lower: 

37 continue 

38 if not any(fragment in key_lower for fragment in fragments): 

39 continue 

40 module = sys.modules.get(key) 

41 if module is not None: 41 ↛ 34line 41 didn't jump to line 34 because the condition on line 41 was always true

42 yield module 

43 

44 

45def force_import_remote_class(model_name: str, dotted_ref: str, **kwargs: Any) -> type | None: 

46 """Import a remote-code class so its module lands in ``sys.modules`` to patch. 

47 

48 ``dotted_ref`` is ``"<module_file>.<ClassName>"`` as understood by 

49 ``get_class_from_dynamic_module``; extra kwargs are forwarded to it. 

50 Returns the class, or None when the dynamic module is unavailable 

51 (offline, renamed, not a remote-code repo) — callers treat that as 

52 "nothing to patch". 

53 """ 

54 try: 

55 from transformers.dynamic_module_utils import get_class_from_dynamic_module 

56 

57 cls = get_class_from_dynamic_module(dotted_ref, model_name, **kwargs) 

58 assert isinstance(cls, type) 

59 return cls 

60 except Exception: 

61 return None 

62 

63 

64def patch_init_weights_skip_loaded(cls: type[Any]) -> None: 

65 """Wrap ``cls._init_weights`` to skip modules already loaded from checkpoint. 

66 

67 v5's meta-device flow calls ``_init_weights`` on modules that already hold 

68 checkpoint weights (remote-code modules lack ``_is_hf_initialized``), 

69 re-randomising them; the wrapper only lets initialisation through for 

70 modules still on the meta device. Idempotent via the ``_tl_patched`` 

71 sentinel. 

72 

73 Raises: 

74 ValueError: for ``transformers.PreTrainedModel`` itself and for classes 

75 without their own ``_init_weights``. Remote modules re-export the 

76 HF base under local names, and patching it would disable weight 

77 init — including HF's rotary-buffer restoration — for every model 

78 loaded later in the process. 

79 """ 

80 if getattr(cls, "_tl_patched", False): 

81 return 

82 if cls is PreTrainedModel: 

83 raise ValueError( 

84 "Refusing to patch transformers.PreTrainedModel itself: that would " 

85 "disable weight init for every model loaded later in this process. " 

86 "Pass the remote code's own PreTrainedModel subclass." 

87 ) 

88 if "_init_weights" not in cls.__dict__: 

89 raise ValueError( 

90 f"Refusing to patch {cls.__name__}: it defines no _init_weights of " 

91 "its own, so the wrap would capture an inherited (possibly HF base) " 

92 "implementation." 

93 ) 

94 

95 original_init_weights = cls._init_weights 

96 

97 def safe_init_weights(self: Any, mod: Any, _original: Any = original_init_weights) -> None: 

98 # Only initialise modules still on meta device (pre-loading); never 

99 # re-randomise weights already read from the checkpoint. 

100 first_param = next(mod.parameters(), None) 

101 if first_param is not None and first_param.device.type != "meta": 

102 return 

103 _original(self, mod) 

104 

105 cls._init_weights = safe_init_weights 

106 cls._tl_patched = True 

107 

108 

109def retie_weights_keys_v5(cls: Any, mapping: dict[str, str]) -> None: 

110 """Rewrite a 4.x list-form ``_tied_weights_keys`` to the v5 dict form. 

111 

112 v5's ``tie_weights`` -> ``get_expanded_tied_weights_keys`` calls ``.keys()`` 

113 on the attribute and raises ``AttributeError`` on the legacy list. No-op 

114 when ``cls`` is None or the attribute is already a dict (or absent). 

115 """ 

116 if cls is not None and isinstance(getattr(cls, "_tied_weights_keys", None), list): 

117 cls._tied_weights_keys = mapping 

118 

119 

120def disable_tied_weights_lookup(cls: type[Any]) -> None: 

121 """Give ``cls`` an empty ``all_tied_weights_keys``. 

122 

123 Some remote ``__getattr__`` implementations fail to delegate v5's 

124 ``all_tied_weights_keys`` lookup back to ``PreTrainedModel`` and raise; 

125 the affected checkpoints are untied anyway. 

126 """ 

127 setattr(cls, "all_tied_weights_keys", {}) 

128 

129 

130def compute_default_rope_inv_freq( 

131 config: Any = None, 

132 device: Any = None, 

133 seq_len: Any = None, 

134 **rope_kwargs: Any, 

135) -> tuple[torch.Tensor, float]: 

136 """transformers-v4 ``_compute_default_rope_parameters``, removed in v5. 

137 

138 Standard (unscaled) RoPE inverse frequencies with the v4 call contract the 

139 remote code targets — ``(config, device) -> (inv_freq, attention_scaling)`` 

140 — plus v4's kwargs-only fallback (``base``/``dim`` passed directly). 

141 Registration strategy stays with each adapter: dream re-registers the 

142 global ``ROPE_INIT_FUNCTIONS["default"]``, ouro deliberately patches only 

143 its own modeling module's copy. 

144 """ 

145 if config is not None: 

146 base = config.rope_theta 

147 partial_rotary_factor = getattr(config, "partial_rotary_factor", None) or 1.0 

148 head_dim = getattr(config, "head_dim", None) or ( 

149 config.hidden_size // config.num_attention_heads 

150 ) 

151 dim = int(head_dim * partial_rotary_factor) 

152 else: 

153 base = rope_kwargs["base"] 

154 dim = rope_kwargs["dim"] 

155 inv_freq = 1.0 / ( 

156 base 

157 ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) 

158 ) 

159 return inv_freq, 1.0