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