Coverage for transformer_lens/model_bridge/supported_architectures/openelm.py: 50%
79 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""OpenELM architecture adapter."""
3import sys
4from typing import Any
6import torch
8from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
9from transformer_lens.model_bridge.compat import patch_dynamic_cache_v5
10from transformer_lens.model_bridge.generalized_components import (
11 BlockBridge,
12 EmbeddingBridge,
13 JointGateUpMLPBridge,
14 LinearBridge,
15 RMSNormalizationBridge,
16 UnembeddingBridge,
17)
18from transformer_lens.model_bridge.generalized_components.attention import (
19 AttentionBridge,
20)
23class OpenElmArchitectureAdapter(ArchitectureAdapter):
24 """Architecture adapter for Apple OpenELM models.
26 OpenELM uses a unique architecture with per-layer varying head counts and FFN
27 dimensions. Key characteristics:
29 - Combined QKV projection (qkv_proj) with per-layer varying Q/KV head counts
30 - Gated MLP with combined gate+up projection (proj_1) and per-layer FFN sizes
31 - RMSNorm normalization
32 - Full rotary embeddings (per-layer, not shared)
33 - Optional Q/K RMSNorm (normalize_qk_projections=True)
34 - Weight tying (share_input_output_layers=True typically)
35 - Model root is 'transformer' (not 'model')
36 - Requires trust_remote_code=True (custom HF code)
38 The native HF attention handles all per-layer dimension variations, RoPE,
39 GQA group repeat, and Q/K normalization internally. The bridge delegates
40 to the native forward for correct computation.
42 Note: Individual Q/K/V hooks are not available since the model uses a combined
43 QKV projection. Attention-level hooks (hook_attn_in, hook_attn_out) are provided.
44 """
46 def __init__(self, cfg: Any) -> None:
47 """Initialize the OpenELM architecture adapter."""
48 super().__init__(cfg)
50 self._set_rms_rotary_defaults()
52 # OpenELM doesn't ship its own tokenizer — uses LLaMA tokenizer.
53 # Use NousResearch mirror (ungated) to avoid access restrictions.
54 self.cfg.tokenizer_name = "NousResearch/Llama-2-7b-hf"
56 # No weight processing conversions needed - native attention handles all
57 # per-layer dimension variations internally
58 self.weight_processing_conversions = {}
60 # Store reference for RoPE patching
61 self._original_rope_compute = None
62 self._rope_class = None
64 self.component_mapping = {
65 "embed": EmbeddingBridge(name="transformer.token_embeddings"),
66 "blocks": BlockBridge(
67 name="transformer.layers",
68 submodules={
69 "ln1": RMSNormalizationBridge(name="attn_norm", config=self.cfg),
70 "ln2": RMSNormalizationBridge(name="ffn_norm", config=self.cfg),
71 "attn": AttentionBridge(
72 name="attn",
73 config=self.cfg,
74 submodules={
75 "qkv": LinearBridge(name="qkv_proj"),
76 "o": LinearBridge(name="out_proj"),
77 },
78 maintain_native_attention=True,
79 requires_attention_mask=True,
80 # Fused qkv_proj (and per-layer head counts): the q/k/v
81 # aliases cannot resolve; hook_z stays via out_proj.
82 fused_qkv=True,
83 ),
84 # proj_1 is a FUSED gate+up projection: under a plain
85 # MLPBridge, hook_pre was the concatenated pre-GLU tensor.
86 "mlp": JointGateUpMLPBridge(
87 name="ffn",
88 config=self.cfg,
89 fused_attr="proj_1",
90 submodules={
91 "out": LinearBridge(name="proj_2"),
92 },
93 ),
94 },
95 ),
96 "ln_final": RMSNormalizationBridge(name="transformer.norm", config=self.cfg),
97 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg),
98 }
100 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
101 """Patch OpenELM for compatibility with transformers v5.
103 apple's forward also calls DynamicCache.from_legacy_cache /
104 to_legacy_cache (removed in v5) — restored by the shared
105 patch_dynamic_cache_v5, as phi3/internlm2/baichuan do.
107 Two module patches are needed besides that:
108 1. RotaryEmbedding: Custom _compute_sin_cos_embeddings fails on meta device
109 because it calls .cos() on meta tensors. We wrap it to catch NotImplementedError.
110 2. Weight re-initialization: OpenELM's _init_weights re-randomizes ALL weights
111 after they've been loaded from safetensors because transformers v5's
112 _finalize_load_state_dict calls initialize_weights() on modules lacking the
113 _is_hf_initialized flag. We patch _init_weights to skip real (non-meta) tensors.
115 Args:
116 model_name: The HuggingFace model name/path
117 model_kwargs: The kwargs dict for from_pretrained()
118 """
119 patch_dynamic_cache_v5()
121 # Force-import the modeling module so we can patch it
122 try:
123 from transformers.dynamic_module_utils import get_class_from_dynamic_module
125 get_class_from_dynamic_module(
126 "modeling_openelm.OpenELMForCausalLM",
127 model_name,
128 )
129 except Exception:
130 return
132 # Find ALL imported OpenELM modules and apply patches.
133 # Each model variant (e.g., OpenELM-1_1B vs OpenELM-1_1B-Instruct) gets its own
134 # module in sys.modules with a different cache path, so we patch all of them.
135 for key in list(sys.modules.keys()):
136 if "openelm" in key.lower() and "modeling" in key.lower():
137 module = sys.modules[key]
138 if hasattr(module, "OpenELMRotaryEmbedding"): 138 ↛ 163line 138 didn't jump to line 163 because the condition on line 138 was always true
139 rope_class = module.OpenELMRotaryEmbedding
140 # Skip if already patched (avoid wrapping safe_compute in safe_compute)
141 if getattr(rope_class, "_tl_patched", False): 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true
142 continue
143 # Patch 1: RoPE meta device fix
144 original_compute = rope_class._compute_sin_cos_embeddings
146 def safe_compute(
147 self,
148 key_len,
149 key_device="cpu",
150 key_dtype=torch.float32,
151 _original=original_compute,
152 ):
153 try:
154 _original(self, key_len, key_device, key_dtype)
155 except NotImplementedError:
156 pass # Deferred: re-initialized in prepare_model()
158 rope_class._compute_sin_cos_embeddings = safe_compute
159 rope_class._tl_patched = True
160 self._original_rope_compute = original_compute
161 self._rope_class = rope_class
163 if hasattr(module, "OpenELMPreTrainedModel"): 163 ↛ 135line 163 didn't jump to line 135 because the condition on line 163 was always true
164 pretrained_class = module.OpenELMPreTrainedModel
165 if getattr(pretrained_class, "_tl_patched", False): 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true
166 continue
167 # Patch 2: Prevent _init_weights from re-randomizing loaded weights.
168 # transformers v5 calls _init_weights on all modules after weight
169 # materialization. For modules with real (non-meta) tensors, we must
170 # skip re-initialization to preserve the loaded checkpoint values.
171 original_init_weights = pretrained_class._init_weights
173 def safe_init_weights(
174 self,
175 mod,
176 _original=original_init_weights,
177 ):
178 # Only initialize modules still on meta device (pre-loading)
179 first_param = next(mod.parameters(), None)
180 if first_param is not None and first_param.device.type != "meta":
181 return # Already loaded from checkpoint — don't re-randomize
182 _original(self, mod)
184 pretrained_class._init_weights = safe_init_weights
185 pretrained_class._tl_patched = True
187 def prepare_model(self, hf_model: Any) -> None:
188 """Post-load fixes for non-persistent buffers zeroed during meta materialization.
190 Transformers v5 creates models on meta device then materializes weights from
191 checkpoint. Non-persistent buffers (registered with persistent=False) are NOT
192 in the checkpoint, so they materialize as zeros. OpenELM has two critical
193 non-persistent buffers that must be recomputed:
195 1. RoPE inv_freq — zeroed inv_freq produces cos=1, sin=0 for all positions,
196 destroying positional information entirely.
197 2. causal_mask — zeroed mask means no causal masking, allowing all positions
198 to attend to future tokens. Single forward passes appear correct (no future
199 tokens to leak) but autoregressive generation degenerates immediately.
201 We also create a synthetic lm_head for weight-tied models.
203 Note: We intentionally do NOT restore the original _compute_sin_cos_embeddings.
204 The safe_compute wrapper is functionally equivalent for real (non-meta) tensors,
205 and keeping it avoids issues when multiple models are loaded in the same process
206 (e.g., benchmark suite loading both HF reference and bridge models).
208 Args:
209 hf_model: The loaded HuggingFace OpenELM model
210 """
211 # ffn_with_glu=False (ungated FFN) would make the fused gate/up split
212 # halve a projection that is not fused: boot succeeds, then the first
213 # forward dies in a mat-mul far from the cause. All published OpenELM
214 # checkpoints use GLU; refuse the config loudly at boot instead.
215 for module_name, module in hf_model.named_modules(): 215 ↛ 224line 215 didn't jump to line 224 because the loop on line 215 didn't complete
216 if getattr(module, "ffn_with_glu", True) is False:
217 raise NotImplementedError(
218 f"OpenELM ffn_with_glu=False on {module_name}: the adapter "
219 "splits proj_1 as a fused gate+up projection, which an "
220 "ungated FFN does not have."
221 )
222 # Ensure use_cache is set on config (transformers v5 raises AttributeError
223 # for missing config attributes, and OpenELM's custom config omits use_cache)
224 if not hasattr(hf_model.config, "use_cache") or "use_cache" not in hf_model.config.__dict__:
225 hf_model.config.use_cache = False
227 # Fix 1: Always recompute causal_mask (non-persistent buffer).
228 # After meta→real materialization, the buffer may contain garbage values
229 # (not all zeros) depending on the materializer's memory state. The old
230 # check `not cm.any()` only recomputed when all zeros, missing cases where
231 # garbage values are non-zero. Always recompute to guarantee correctness.
232 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "causal_mask"):
233 cm = hf_model.transformer.causal_mask
234 if cm is not None:
235 seq_len = cm.shape[-1]
236 correct_mask = torch.triu(
237 torch.ones(seq_len, seq_len, dtype=cm.dtype, device=cm.device),
238 diagonal=1,
239 )
240 hf_model.transformer.causal_mask = correct_mask
242 # Fix 2: Recompute RoPE inv_freq on all layers (non-persistent buffer zeroed
243 # during materialization), then force-recompute sin/cos embeddings.
244 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "layers"):
245 rope_max = getattr(hf_model.config, "rope_max_length", 4096)
246 for layer in hf_model.transformer.layers:
247 if hasattr(layer, "attn") and hasattr(layer.attn, "pos_embedding"):
248 rope = layer.attn.pos_embedding
249 # Always recompute inv_freq (non-persistent buffer).
250 # Like causal_mask, inv_freq may contain garbage after meta
251 # materialization rather than clean zeros.
252 correct_inv_freq = 1.0 / (
253 rope.freq_constant
254 ** (
255 torch.arange(0, rope.model_dim, 2, dtype=torch.float32) / rope.model_dim
256 )
257 )
258 rope.inv_freq = correct_inv_freq.to(rope.inv_freq.device)
259 # Force-recompute sin/cos (may have been computed with zero inv_freq)
260 rope._cached_cos = None
261 rope._cached_sin = None
262 rope._compute_sin_cos_embeddings(rope_max)
264 # Create synthetic lm_head when embeddings are shared
265 if getattr(hf_model, "lm_head", None) is None and hasattr(hf_model, "transformer"):
266 embed = hf_model.transformer.token_embeddings
267 lm_head = torch.nn.Linear(embed.embedding_dim, embed.num_embeddings, bias=False)
268 lm_head.weight = embed.weight
269 hf_model.lm_head = lm_head