from __future__ import annotations import json import math from dataclasses import asdict, dataclass from pathlib import Path from typing import Iterable import torch import torch.nn.functional as F from torch import Tensor, nn from .smollm2_amcenn import DEFAULT_SMOLLM2 from .smollm2_amcenn_v2 import AdaptivePositiveSoftmaxFeatures @dataclass(frozen=True) class LayerProfileV4: tier: str local_window: int feature_dim: int gate_init: float gate_cap: float def to_dict(self) -> dict: return asdict(self) @dataclass(frozen=True) class SmolAMCeNNV4Config: """Layer-adaptive exact-local/anchor + recurrent-global attention. The default profile comes directly from the v3 calibration pattern: easy layers keep a smaller exact window, difficult layers get a larger exact window and a larger AM-CeNN feature map, and the two most difficult layers get a 96-token exact window plus 512 recurrent features. """ easy_feature_dim: int = 256 medium_feature_dim: int = 320 hard_feature_dim: int = 384 critical_feature_dim: int = 512 easy_window: int = 32 medium_window: int = 48 hard_window: int = 64 critical_window: int = 96 anchor_tokens: int = 8 feature_seed: int = 4040 eps: float = 1e-6 antithetic_features: bool = True learnable_feature_correction: bool = True gate_init: float = 0.03 easy_gate_cap: float = 0.25 medium_gate_cap: float = 0.20 hard_gate_cap: float = 0.15 critical_gate_cap: float = 0.10 easy_layers: tuple[int, ...] = (0, 1, 2, 10, 11, 24, 29) hard_layers: tuple[int, ...] = (6, 7, 8, 14, 17, 23, 26) critical_layers: tuple[int, ...] = (18, 20) def validate(self, model_config) -> None: dims = ( self.easy_feature_dim, self.medium_feature_dim, self.hard_feature_dim, self.critical_feature_dim, ) if min(dims) < 16: raise ValueError("all v4 feature dimensions must be >= 16") if self.antithetic_features and any(dim % 2 for dim in dims): raise ValueError("antithetic v4 feature dimensions must be even") windows = ( self.easy_window, self.medium_window, self.hard_window, self.critical_window, ) if min(windows) < 1: raise ValueError("all v4 local windows must be >= 1") if self.anchor_tokens < 0: raise ValueError("anchor_tokens must be >= 0") caps = ( self.easy_gate_cap, self.medium_gate_cap, self.hard_gate_cap, self.critical_gate_cap, ) if not 0.0 < self.gate_init < 1.0: raise ValueError("gate_init must be in (0, 1)") if any(not 0.0 < cap < 1.0 for cap in caps): raise ValueError("all v4 gate caps must be in (0, 1)") if any(self.gate_init >= cap for cap in caps): raise ValueError("gate_init must be lower than every v4 gate cap") hidden = int(model_config.hidden_size) heads = int(model_config.num_attention_heads) if hidden % heads: raise ValueError("hidden_size must be divisible by num_attention_heads") def profile_for_layer(self, layer_idx: int) -> LayerProfileV4: idx = int(layer_idx) if idx in self.critical_layers: return LayerProfileV4( "critical", self.critical_window, self.critical_feature_dim, self.gate_init, self.critical_gate_cap, ) if idx in self.hard_layers: return LayerProfileV4( "hard", self.hard_window, self.hard_feature_dim, self.gate_init, self.hard_gate_cap, ) if idx in self.easy_layers: return LayerProfileV4( "easy", self.easy_window, self.easy_feature_dim, self.gate_init, self.easy_gate_cap, ) return LayerProfileV4( "medium", self.medium_window, self.medium_feature_dim, self.gate_init, self.medium_gate_cap, ) def to_dict(self) -> dict: return asdict(self) @classmethod def from_dict(cls, data: dict) -> "SmolAMCeNNV4Config": values = dict(data) for key in ("easy_layers", "hard_layers", "critical_layers"): if key in values: values[key] = tuple(int(v) for v in values[key]) return cls(**values) def _logit(probability: float) -> float: p = min(max(float(probability), 1e-6), 1.0 - 1e-6) return math.log(p / (1.0 - p)) class AdaptiveHybridAMCeNNAttentionV4(nn.Module): """Exact anchors + exact local window + token-dependent recurrent memory. For query position t, exact softmax covers: * the first ``anchor_tokens`` tokens (causally masked), and * the most recent ``local_window`` tokens. AM-CeNN summarizes only the older middle tokens not already covered exactly. A token-dependent, per-head gate chooses how much of that recurrent memory to use. The gate is intrinsically capped per layer tier, so difficult layers can never over-rely on a poor recurrent approximation. """ def __init__( self, original_attn: nn.Module, model_config, config: SmolAMCeNNV4Config, layer_idx: int, ) -> None: super().__init__() config.validate(model_config) self.hidden_size = int(model_config.hidden_size) self.num_heads = int(model_config.num_attention_heads) self.num_key_value_heads = int(model_config.num_key_value_heads) self.head_dim = self.hidden_size // self.num_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.layer_idx = int(layer_idx) self.anchor_tokens = int(config.anchor_tokens) self.eps = float(config.eps) self.profile = config.profile_for_layer(self.layer_idx) self.local_window = int(self.profile.local_window) self.feature_dim = int(self.profile.feature_dim) self.gate_cap = float(self.profile.gate_cap) self.tier = self.profile.tier self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) with torch.no_grad(): self.q_proj.weight.copy_(original_attn.q_proj.weight) self.k_proj.weight.copy_(original_attn.k_proj.weight) self.v_proj.weight.copy_(original_attn.v_proj.weight) self.o_proj.weight.copy_(original_attn.o_proj.weight) self.features = AdaptivePositiveSoftmaxFeatures( self.head_dim, self.feature_dim, seed=config.feature_seed + self.layer_idx, antithetic=config.antithetic_features, learnable_correction=config.learnable_feature_correction, ) # gate = gate_cap * sigmoid(W h + b). W starts at zero, therefore v4 # begins as a safe constant-gate model and can become token-adaptive. self.gate_proj = nn.Linear(self.hidden_size, self.num_heads, bias=True) nn.init.zeros_(self.gate_proj.weight) initial_fraction = float(self.profile.gate_init) / self.gate_cap nn.init.constant_(self.gate_proj.bias, _logit(initial_fraction)) self.last_mean_gate = torch.tensor(float(self.profile.gate_init)) self.last_max_gate = torch.tensor(float(self.profile.gate_init)) self.last_global_state_norm = torch.tensor(0.0) def _apply_rope(self, q: Tensor, k: Tensor, position_embeddings) -> tuple[Tensor, Tensor]: if position_embeddings is None: return q, k try: from transformers.models.llama.modeling_llama import apply_rotary_pos_emb cos, sin = position_embeddings return apply_rotary_pos_emb(q, k, cos, sin) except Exception: return q, k def _exact_attention(self, q: Tensor, k: Tensor, v: Tensor, attention_mask) -> Tensor: k_heads = k.repeat_interleave(self.num_key_value_groups, dim=1) v_heads = v.repeat_interleave(self.num_key_value_groups, dim=1) scores = torch.einsum("bhtd,bhsd->bhts", q.float(), k_heads.float()) scores = scores / math.sqrt(float(self.head_dim)) seq_len = q.shape[-2] positions = torch.arange(seq_len, device=q.device) qpos = positions[:, None] kpos = positions[None, :] causal = kpos <= qpos local = kpos >= (qpos - self.local_window + 1) anchors = kpos < min(self.anchor_tokens, seq_len) allowed = causal & (local | anchors) scores = scores.masked_fill(~allowed.view(1, 1, seq_len, seq_len), float("-inf")) if torch.is_tensor(attention_mask): mask = attention_mask try: if mask.ndim == 4 and mask.shape[-2:] == (seq_len, seq_len): scores = scores + mask.to(device=scores.device, dtype=scores.dtype) elif mask.ndim == 2 and mask.shape[-1] == seq_len: valid = mask.to(device=scores.device).bool().view(mask.shape[0], 1, 1, seq_len) scores = scores.masked_fill(~valid, float("-inf")) except Exception: pass probs = torch.softmax(scores, dim=-1, dtype=torch.float32) return torch.einsum("bhts,bhsd->bhtd", probs, v_heads.float()) def _global_middle_memory(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: phi_q = self.features(q).transpose(1, 2) # B,T,H,F or B,T,KV,F after repeat below phi_k = self.features(k).transpose(1, 2) # B,T,KV,F values = v.transpose(1, 2).float() # B,T,KV,D writes = torch.einsum("btkf,btkd->btkfd", phi_k, values) if self.anchor_tokens > 0: writes = writes.clone() phi_k = phi_k.clone() cutoff = min(self.anchor_tokens, writes.shape[1]) writes[:, :cutoff] = 0 phi_k[:, :cutoff] = 0 prefix_s = writes.cumsum(dim=1) prefix_z = phi_k.cumsum(dim=1) old_s = torch.zeros_like(prefix_s) old_z = torch.zeros_like(prefix_z) if q.shape[-2] > self.local_window: old_s[:, self.local_window:] = prefix_s[:, :-self.local_window] old_z[:, self.local_window:] = prefix_z[:, :-self.local_window] old_s_h = old_s.repeat_interleave(self.num_key_value_groups, dim=2) old_z_h = old_z.repeat_interleave(self.num_key_value_groups, dim=2) numerator = torch.einsum("bthf,bthfd->bthd", phi_q, old_s_h) denominator = torch.einsum("bthf,bthf->bth", phi_q, old_z_h).unsqueeze(-1) global_out = numerator / denominator.clamp_min(self.eps) valid = denominator > self.eps global_out = torch.where(valid, global_out, torch.zeros_like(global_out)) return global_out.transpose(1, 2), valid.transpose(1, 2) def _token_gate(self, hidden_states: Tensor, global_valid: Tensor) -> Tensor: logits = F.linear( hidden_states.float(), self.gate_proj.weight.float(), self.gate_proj.bias.float(), ) gate = self.gate_cap * torch.sigmoid(logits) gate = gate.transpose(1, 2).unsqueeze(-1) # B,H,T,1 return gate * global_valid.to(dtype=gate.dtype) def forward( self, hidden_states: Tensor, attention_mask=None, position_ids=None, past_key_values=None, use_cache: bool = False, cache_position=None, position_embeddings=None, **kwargs, ) -> tuple[Tensor, None]: if use_cache: raise RuntimeError("AM-CeNN v4 currently requires use_cache=False") bsz, seq_len, _ = hidden_states.shape q = self.q_proj(hidden_states).view( bsz, seq_len, self.num_heads, self.head_dim ).transpose(1, 2) k = self.k_proj(hidden_states).view( bsz, seq_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) v = self.v_proj(hidden_states).view( bsz, seq_len, self.num_key_value_heads, self.head_dim ).transpose(1, 2) q, k = self._apply_rope(q, k, position_embeddings) exact_out = self._exact_attention(q, k, v, attention_mask) global_out, global_valid = self._global_middle_memory(q, k, v) gate = self._token_gate(hidden_states, global_valid) mixed = exact_out.float() + gate * (global_out.float() - exact_out.float()) mixed = mixed.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size) out = self.o_proj(mixed.to(dtype=hidden_states.dtype)) valid_gates = gate.detach().float()[global_valid.expand_as(gate)] if valid_gates.numel(): self.last_mean_gate = valid_gates.mean().cpu() self.last_max_gate = valid_gates.max().cpu() else: self.last_mean_gate = torch.tensor(0.0) self.last_max_gate = torch.tensor(0.0) self.last_global_state_norm = global_out[:, :, -1].detach().float().norm().cpu() return out, None def replace_attention_layers_v4( model: nn.Module, config: SmolAMCeNNV4Config, layer_indices: Iterable[int], ) -> nn.Module: config.validate(model.config) for idx in layer_indices: layer = model.model.layers[int(idx)] if isinstance(layer.self_attn, AdaptiveHybridAMCeNNAttentionV4): continue old = layer.self_attn new = AdaptiveHybridAMCeNNAttentionV4(old, model.config, config, int(idx)) new.to(device=old.q_proj.weight.device, dtype=old.q_proj.weight.dtype) layer.self_attn = new model.config.use_cache = False if hasattr(model, "generation_config"): model.generation_config.use_cache = False return model def replace_all_attention_v4(model: nn.Module, config: SmolAMCeNNV4Config) -> nn.Module: return replace_attention_layers_v4(model, config, range(int(model.config.num_hidden_layers))) def freeze_for_v4_calibration(model: nn.Module, layer_indices: Iterable[int]) -> list[nn.Parameter]: selected = {int(i) for i in layer_indices} for p in model.parameters(): p.requires_grad = False trainable: list[nn.Parameter] = [] for module in model.modules(): if not isinstance(module, AdaptiveHybridAMCeNNAttentionV4) or module.layer_idx not in selected: continue if module.features.delta_projection is not None: module.features.delta_projection.requires_grad = True trainable.append(module.features.delta_projection) for p in module.gate_proj.parameters(): p.requires_grad = True trainable.append(p) return trainable def v4_global_parameter_groups( model: nn.Module, *, memory_lr: float, qkvo_lr: float, weight_decay: float = 0.01, ) -> tuple[list[dict], list[nn.Parameter]]: for p in model.parameters(): p.requires_grad = False memory: list[nn.Parameter] = [] qkvo: list[nn.Parameter] = [] for module in model.modules(): if not isinstance(module, AdaptiveHybridAMCeNNAttentionV4): continue if module.features.delta_projection is not None: module.features.delta_projection.requires_grad = True memory.append(module.features.delta_projection) for p in module.gate_proj.parameters(): p.requires_grad = True memory.append(p) for projection in (module.q_proj, module.k_proj, module.v_proj, module.o_proj): projection.weight.requires_grad = True qkvo.append(projection.weight) groups = [ {"params": memory, "lr": float(memory_lr), "weight_decay": float(weight_decay)}, {"params": qkvo, "lr": float(qkvo_lr), "weight_decay": float(weight_decay)}, ] return groups, [*memory, *qkvo] def v4_parameter_summary(model: nn.Module) -> dict[str, int | float]: total = sum(p.numel() for p in model.parameters()) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) modules = [m for m in model.modules() if isinstance(m, AdaptiveHybridAMCeNNAttentionV4)] feature_delta = sum( m.features.delta_projection.numel() for m in modules if m.features.delta_projection is not None ) gate_params = sum(p.numel() for m in modules for p in m.gate_proj.parameters()) return { "total": total, "trainable": trainable, "hybrid_attention": sum(p.numel() for m in modules for p in m.parameters()), "feature_delta": feature_delta, "token_gate": gate_params, "trainable_percent": 100.0 * trainable / max(total, 1), } def v4_attention_stats(model: nn.Module) -> dict: modules = [m for m in model.modules() if isinstance(m, AdaptiveHybridAMCeNNAttentionV4)] if not modules: raise RuntimeError("AM-CeNN v4 attention modules unavailable") means = [float(m.last_mean_gate) for m in modules] maxima = [float(m.last_max_gate) for m in modules] by_tier: dict[str, list[float]] = {} layer_profiles = [] for m, mean_gate in zip(modules, means): by_tier.setdefault(m.tier, []).append(mean_gate) layer_profiles.append({ "layer": m.layer_idx, "tier": m.tier, "local_window": m.local_window, "feature_dim": m.feature_dim, "gate_cap": m.gate_cap, "mean_gate": mean_gate, }) return { "mean_global_gate": sum(means) / len(means), "max_observed_gate": max(maxima), "mean_gate_by_tier": {k: sum(v) / len(v) for k, v in by_tier.items()}, "layer_profiles": layer_profiles, } def _v4_state(model: nn.Module) -> dict[str, Tensor]: state: dict[str, Tensor] = {} for name, tensor in model.state_dict().items(): if ".self_attn." in name: state[name] = tensor.detach().cpu() if not state: raise RuntimeError("no AM-CeNN v4 state found") return state def save_smollm2_amcenn_v4( model: nn.Module, output_dir: str | Path, *, config: SmolAMCeNNV4Config, base_model: str = DEFAULT_SMOLLM2, extra_metadata: dict | None = None, ) -> Path: output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) torch.save(_v4_state(model), output_dir / "smollm2_amcenn_v4.pt") metadata = { "format_version": 4, "architecture": "smollm2-amcenn-adaptive-v4", "base_model": base_model, "amcenn_v4": config.to_dict(), } if extra_metadata: metadata["training"] = extra_metadata (output_dir / "smollm2_amcenn_v4_config.json").write_text( json.dumps(metadata, indent=2), encoding="utf-8" ) return output_dir def load_smollm2_amcenn_v4_weights(model: nn.Module, student_dir: str | Path) -> nn.Module: state = torch.load( Path(student_dir) / "smollm2_amcenn_v4.pt", map_location="cpu", weights_only=True, ) incompatible = model.load_state_dict(state, strict=False) expected = set(_v4_state(model)) missing = [k for k in incompatible.missing_keys if k in expected] unexpected = [k for k in incompatible.unexpected_keys if k not in expected] if missing: raise RuntimeError(f"missing AM-CeNN v4 keys: {missing[:8]}") if unexpected: raise RuntimeError(f"unexpected AM-CeNN v4 keys: {unexpected[:8]}") return model def build_smollm2_amcenn_v4(student_dir: str | Path, *, device=None, dtype=None): from transformers import AutoModelForCausalLM student_dir = Path(student_dir) metadata = json.loads((student_dir / "smollm2_amcenn_v4_config.json").read_text()) if metadata.get("architecture") != "smollm2-amcenn-adaptive-v4": raise ValueError("checkpoint is not SmolLM2 AM-CeNN adaptive v4") kwargs = {} if dtype is not None: kwargs["dtype"] = dtype model = AutoModelForCausalLM.from_pretrained(metadata["base_model"], **kwargs) config = SmolAMCeNNV4Config.from_dict(metadata["amcenn_v4"]) replace_all_attention_v4(model, config) load_smollm2_amcenn_v4_weights(model, student_dir) if device is not None: model.to(device) if dtype is not None: model.to(dtype=dtype) model.config.use_cache = False if hasattr(model, "generation_config"): model.generation_config.use_cache = False return model