Coverage for transformer_lens/model_bridge/supported_architectures/openelm.py: 48%
63 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 force_import_remote_class(model_name, "modeling_openelm.OpenELMForCausalLM") is None: 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true
127 return
129 # Each model variant (e.g., OpenELM-1_1B vs OpenELM-1_1B-Instruct) gets its
130 # own module in sys.modules with a different cache path; patch all of them.
131 for module in iter_remote_modeling_modules("openelm"):
132 rope_class = getattr(module, "OpenELMRotaryEmbedding", None)
133 # Skip if already patched (avoid wrapping safe_compute in safe_compute)
134 if rope_class is not None and not getattr(rope_class, "_tl_patched", False): 134 ↛ 156line 134 didn't jump to line 156 because the condition on line 134 was always true
135 # Patch 1: RoPE meta device fix
136 original_compute = rope_class._compute_sin_cos_embeddings
138 def safe_compute(
139 self,
140 key_len,
141 key_device="cpu",
142 key_dtype=torch.float32,
143 _original=original_compute,
144 ):
145 try:
146 _original(self, key_len, key_device, key_dtype)
147 except NotImplementedError:
148 pass # Deferred: re-initialized in prepare_model()
150 rope_class._compute_sin_cos_embeddings = safe_compute
151 rope_class._tl_patched = True
152 self._original_rope_compute = original_compute
153 self._rope_class = rope_class
155 # Patch 2: don't let _init_weights re-randomize loaded weights.
156 pretrained_class = getattr(module, "OpenELMPreTrainedModel", None)
157 if pretrained_class is not None: 157 ↛ 131line 157 didn't jump to line 131 because the condition on line 157 was always true
158 patch_init_weights_skip_loaded(pretrained_class)
160 def prepare_model(self, hf_model: Any) -> None:
161 """Post-load fixes for non-persistent buffers zeroed during meta materialization.
163 Transformers v5 creates models on meta device then materializes weights from
164 checkpoint. Non-persistent buffers (registered with persistent=False) are NOT
165 in the checkpoint, so they materialize as zeros. OpenELM has two critical
166 non-persistent buffers that must be recomputed:
168 1. RoPE inv_freq — zeroed inv_freq produces cos=1, sin=0 for all positions,
169 destroying positional information entirely.
170 2. causal_mask — zeroed mask means no causal masking, allowing all positions
171 to attend to future tokens. Single forward passes appear correct (no future
172 tokens to leak) but autoregressive generation degenerates immediately.
174 We also create a synthetic lm_head for weight-tied models.
176 Note: We intentionally do NOT restore the original _compute_sin_cos_embeddings.
177 The safe_compute wrapper is functionally equivalent for real (non-meta) tensors,
178 and keeping it avoids issues when multiple models are loaded in the same process
179 (e.g., benchmark suite loading both HF reference and bridge models).
181 Args:
182 hf_model: The loaded HuggingFace OpenELM model
183 """
184 # ffn_with_glu=False (ungated FFN) would make the fused gate/up split
185 # halve a projection that is not fused: boot succeeds, then the first
186 # forward dies in a mat-mul far from the cause. All published OpenELM
187 # checkpoints use GLU; refuse the config loudly at boot instead.
188 for module_name, module in hf_model.named_modules(): 188 ↛ 197line 188 didn't jump to line 197 because the loop on line 188 didn't complete
189 if getattr(module, "ffn_with_glu", True) is False:
190 raise NotImplementedError(
191 f"OpenELM ffn_with_glu=False on {module_name}: the adapter "
192 "splits proj_1 as a fused gate+up projection, which an "
193 "ungated FFN does not have."
194 )
195 # Ensure use_cache is set on config (transformers v5 raises AttributeError
196 # for missing config attributes, and OpenELM's custom config omits use_cache)
197 if not hasattr(hf_model.config, "use_cache") or "use_cache" not in hf_model.config.__dict__:
198 hf_model.config.use_cache = False
200 # Fix 1: Always recompute causal_mask (non-persistent buffer).
201 # After meta→real materialization, the buffer may contain garbage values
202 # (not all zeros) depending on the materializer's memory state. The old
203 # check `not cm.any()` only recomputed when all zeros, missing cases where
204 # garbage values are non-zero. Always recompute to guarantee correctness.
205 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "causal_mask"):
206 cm = hf_model.transformer.causal_mask
207 if cm is not None:
208 seq_len = cm.shape[-1]
209 correct_mask = torch.triu(
210 torch.ones(seq_len, seq_len, dtype=cm.dtype, device=cm.device),
211 diagonal=1,
212 )
213 hf_model.transformer.causal_mask = correct_mask
215 # Fix 2: Recompute RoPE inv_freq on all layers (non-persistent buffer zeroed
216 # during materialization), then force-recompute sin/cos embeddings.
217 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "layers"):
218 rope_max = getattr(hf_model.config, "rope_max_length", 4096)
219 for layer in hf_model.transformer.layers:
220 if hasattr(layer, "attn") and hasattr(layer.attn, "pos_embedding"):
221 rope = layer.attn.pos_embedding
222 # Always recompute inv_freq (non-persistent buffer).
223 # Like causal_mask, inv_freq may contain garbage after meta
224 # materialization rather than clean zeros.
225 correct_inv_freq = 1.0 / (
226 rope.freq_constant
227 ** (
228 torch.arange(0, rope.model_dim, 2, dtype=torch.float32) / rope.model_dim
229 )
230 )
231 rope.inv_freq = correct_inv_freq.to(rope.inv_freq.device)
232 # Force-recompute sin/cos (may have been computed with zero inv_freq)
233 rope._cached_cos = None
234 rope._cached_sin = None
235 rope._compute_sin_cos_embeddings(rope_max)
237 # Create synthetic lm_head when embeddings are shared
238 if getattr(hf_model, "lm_head", None) is None and hasattr(hf_model, "transformer"):
239 embed = hf_model.transformer.token_embeddings
240 lm_head = torch.nn.Linear(embed.embedding_dim, embed.num_embeddings, bias=False)
241 lm_head.weight = embed.weight
242 hf_model.lm_head = lm_head