import math import torch import torch.nn as nn import torch.nn.functional as F # CMA is text-only. Disable Transformers' optional vision backend before it # imports PreTrainedModel: some cloud images expose a system torchvision built # against a different PyTorch, which otherwise crashes on torchvision::nms. import transformers as _transformers import transformers.utils as _transformers_utils from transformers.utils import import_utils as _transformers_import_utils def _cma_torchvision_unavailable(): return False _transformers_import_utils._torchvision_available = False _transformers_import_utils.is_torchvision_available = _cma_torchvision_unavailable _transformers_import_utils.is_torchvision_v2_available = _cma_torchvision_unavailable _transformers_utils.is_torchvision_available = _cma_torchvision_unavailable _transformers_utils.is_torchvision_v2_available = _cma_torchvision_unavailable _transformers.is_torchvision_available = _cma_torchvision_unavailable _transformers.is_torchvision_v2_available = _cma_torchvision_unavailable try: from transformers import GenerationMixin except ImportError: from transformers.generation import GenerationMixin from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast class CMAConfig(PretrainedConfig): model_type = "cma" def __init__( self, vocab_size=260, seq_len=2048, d_model=128, n_layers=6, n_heads=4, n_kv_heads=2, chunk=16, cma_heads=2, expand=2, cma_identity_prob=0.90, max_position_embeddings=None, n_positions=None, n_ctx=None, **kwargs, ): super().__init__(**kwargs) self.vocab_size = vocab_size self.seq_len = seq_len self.max_position_embeddings = max_position_embeddings or seq_len self.n_positions = n_positions or self.max_position_embeddings self.n_ctx = n_ctx or self.max_position_embeddings self.d_model = d_model self.n_layers = n_layers self.n_heads = n_heads self.n_kv_heads = n_kv_heads self.chunk = chunk self.cma_heads = cma_heads self.expand = expand self.cma_identity_prob = cma_identity_prob class RMSNorm(nn.Module): def __init__(self, d): super().__init__() self.w = nn.Parameter(torch.ones(d)) def forward(self, x): return F.rms_norm( x, (x.shape[-1],), self.w.to(dtype=x.dtype), eps=1e-6 ) def rope_cache(seq_len, head_dim, device, base=10000.0): inv = 1.0 / ( base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim) ) t = torch.arange(seq_len, device=device).float() freqs = torch.outer(t, inv) return torch.cos(freqs), torch.sin(freqs) def apply_rope(x, cos, sin): x1, x2 = x.chunk(2, dim=-1) return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) class TokenAttention(nn.Module): def __init__(self, d, n_heads, n_kv_heads): super().__init__() if d % n_heads or n_heads % n_kv_heads: raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero.") self.h, self.kv_h, self.hd = n_heads, n_kv_heads, d // n_heads self.q = nn.Linear(d, d, bias=False) self.k = nn.Linear(d, n_kv_heads * self.hd, bias=False) self.v = nn.Linear(d, n_kv_heads * self.hd, bias=False) self.o = nn.Linear(d, d, bias=False) self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd) def forward(self, x, cos, sin): B, T, d = x.shape q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2) k = self.k(x).view(B, T, self.kv_h, self.hd).transpose(1, 2) v = self.v(x).view(B, T, self.kv_h, self.hd).transpose(1, 2) q, k = self.qn(q), self.kn(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # Native GQA keeps compact K/V heads and lets SDPA select a fused CUDA # implementation without materializing repeated K/V activations. y = F.scaled_dot_product_attention( q, k, v, is_causal=True, enable_gqa=self.h != self.kv_h, ) return self.o(y.transpose(1, 2).reshape(B, T, d)) class CMA(nn.Module): def __init__(self, d, chunk=16, heads=2, expand=2, identity_prob=0.90): super().__init__() if min(d, chunk, heads, expand) <= 0: raise ValueError("CMA dimensions and expansion must be positive.") if d % chunk or chunk % heads: raise ValueError("CMA requires d % chunk == 0 and chunk % heads == 0.") if d // chunk < 2: raise ValueError("CMA requires at least two channel slots.") if not 0.0 < identity_prob < 1.0: raise ValueError("cma_identity_prob must be between zero and one.") self.d, self.n, self.c, self.h = d, d // chunk, chunk, heads self.hd, self.expand = chunk // heads, expand self.chunk_emb = nn.Parameter(torch.randn(self.n, chunk) * 0.02) self.wqk = nn.Parameter(torch.randn(self.n, chunk, 2 * chunk) * 0.02) self.global_proj = nn.Linear(d, chunk, bias=False) self.wv = nn.Linear(d, d * expand, bias=False) self.bias = nn.Parameter(torch.zeros(heads, self.n, self.n)) self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd) self.logit_scale = nn.Parameter(torch.zeros(heads)) self.layer_gain = nn.Parameter(torch.zeros(heads)) self.route_gate_weight = nn.Parameter(torch.randn(heads, self.hd) * 0.02) self.route_gate_bias = nn.Parameter(torch.zeros(heads)) self.gate = nn.Linear(d, d * expand, bias=False) self.o = nn.Linear(d * expand, d, bias=False) nn.init.zeros_(self.o.weight) diagonal_bias = math.log( (self.n - 1) * identity_prob / (1.0 - identity_prob) ) with torch.no_grad(): self.bias.add_(torch.eye(self.n) * diagonal_bias) def _routing(self, x): batch_tokens = x.numel() // self.d xc = x.reshape(batch_tokens, self.n, self.c) global_state = self.global_proj(x).reshape(batch_tokens, 1, self.c) q_input = xc + self.chunk_emb + global_state value_slots = self.wv(x).reshape( batch_tokens, self.n, self.h, self.hd, self.expand ) key_input = value_slots.mean(dim=-1).reshape(batch_tokens, self.n, self.c) key_input = key_input + self.chunk_emb q = torch.einsum("bnc,nco->bno", q_input, self.wqk[..., : self.c]) k = torch.einsum("bnc,nco->bno", key_input, self.wqk[..., self.c :]) v = value_slots.reshape( batch_tokens, self.n, self.h, self.hd * self.expand ).transpose(1, 2) q = self.qn(q.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2) k = self.kn(k.reshape(batch_tokens, self.n, self.h, self.hd)).transpose(1, 2) q = F.normalize(q.float(), dim=-1).to(v.dtype) k = F.normalize(k.float(), dim=-1).to(v.dtype) scale = self.logit_scale.clamp(max=math.log(100.0)).exp() logits = (q * scale.view(1, self.h, 1, 1).to(q.dtype)) @ k.transpose(-2, -1) logits = logits + self.bias.unsqueeze(0).to(logits.dtype) attn = F.softmax(logits, dim=-1, dtype=torch.float32).to(v.dtype) routed = attn @ v route_signal = ( (q.float() * self.route_gate_weight.float().view(1, self.h, 1, self.hd)).sum(-1) + self.route_gate_bias.float().view(1, self.h, 1) + self.layer_gain.float().view(1, self.h, 1) ) route_coeff = torch.tanh(route_signal).to(v.dtype) contribution = route_coeff.unsqueeze(-1) * (routed - v) return v, attn, route_coeff, contribution, logits def forward(self, x): B, T, _ = x.shape v, _, _, contribution, _ = self._routing(x) y = (v + contribution).transpose(1, 2).reshape(B, T, self.d * self.expand) return self.o(y * F.silu(self.gate(x))) @torch.no_grad() def diagnostic_stats(self, x): probe_count = x.shape[0] if x.ndim > 2 else 1 v, attn, route_coeff, contribution, logits = self._routing(x) B, T, _ = x.shape base = v.transpose(1, 2).reshape(B, T, self.d * self.expand) routed = (v + contribution).transpose(1, 2).reshape( B, T, self.d * self.expand ) gate = F.silu(self.gate(x)) base_output = self.o(base * gate) routed_output = self.o(routed * gate) output_delta = (routed_output - base_output).reshape(-1, self.d) token_effect = ( output_delta.float().norm(dim=-1) / routed_output.reshape(-1, self.d).float().norm(dim=-1).clamp_min(1e-12) ) probe_effect = token_effect.reshape(probe_count, -1).mean(dim=-1) probs = attn.float().clamp_min(1e-9) entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n) diagonal_mass = probs.diagonal(dim1=-2, dim2=-1).mean() if self.h > 1: flattened = probs.transpose(0, 1).reshape(self.h, -1) normalized = F.normalize(flattened, dim=-1) similarity = normalized @ normalized.mT head_similarity = ( similarity.sum() - similarity.diagonal().sum() ) / (self.h * (self.h - 1)) else: head_similarity = probs.new_zeros(()) return { "token_effect": token_effect, "probe_effect": probe_effect, "entropy": entropy.mean().item(), "diagonal_mass": diagonal_mass.item(), "head_similarity": head_similarity.item(), "gate_abs": route_coeff.float().abs().mean().item(), "gate_saturation": (route_coeff.float().abs() > 0.95).float().mean().item(), "logit_max": logits.float().abs().max().item(), } class Block(nn.Module): def __init__(self, config): super().__init__() d = config.d_model self.n1, self.n2 = RMSNorm(d), RMSNorm(d) self.attn = TokenAttention(d, config.n_heads, config.n_kv_heads) self.mix = CMA( d, config.chunk, config.cma_heads, config.expand, config.cma_identity_prob, ) def forward(self, x, cos, sin): x = x + self.attn(self.n1(x), cos, sin) x = x + self.mix(self.n2(x)) return x class CMAModel(nn.Module): def __init__(self, config): super().__init__() d = config.d_model self.config = config self.emb = nn.Embedding(config.vocab_size, d) self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers)) self.norm = RMSNorm(d) hd = d // config.n_heads cos, sin = rope_cache(config.seq_len, hd, "cpu") self.register_buffer("cos", cos) self.register_buffer("sin", sin) def forward_hidden(self, idx): if idx.size(1) > self.config.seq_len: idx = idx[:, -self.config.seq_len :] x = self.emb(idx) device_type = x.device.type rope_dtype = ( torch.get_autocast_dtype(device_type) if torch.is_autocast_enabled(device_type) else x.dtype ) cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=rope_dtype) sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=rope_dtype) for b in self.blocks: x = b(x, cos, sin) return self.norm(x) class CMAForCausalLM(PreTrainedModel, GenerationMixin): config_class = CMAConfig base_model_prefix = "model" _tied_weights_keys = {"head.weight": "model.emb.weight"} all_tied_weights_keys = {"head.weight": "model.emb.weight"} def __init__(self, config): super().__init__(config) self.model = CMAModel(config) self.head = nn.Linear(config.d_model, config.vocab_size, bias=False) # Modern Transformers creates loader metadata and performs configured # tying in post_init(); omitting it leaves all_tied_weights_keys absent. self.post_init() self.head.weight = self.model.emb.weight def get_input_embeddings(self): return self.model.emb def set_input_embeddings(self, value): self.model.emb = value def get_output_embeddings(self): return self.head def set_output_embeddings(self, new_embeddings): self.head = new_embeddings def raw_logits(self, idx): return self.head(self.model.forward_hidden(idx)) def logits(self, idx): return self.raw_logits(idx) def _masked_logits(self, input_ids, attention_mask): if attention_mask is None or bool(attention_mask.all()): return self.logits(input_ids) B, T = input_ids.shape out = None for i in range(B): keep = attention_mask[i].bool().nonzero(as_tuple=False).flatten() if keep.numel() == 0: keep = torch.tensor([T - 1], device=input_ids.device) trimmed = input_ids[i, keep].unsqueeze(0) logits_i = self.logits(trimmed) if out is None: out = logits_i.new_zeros(B, T, logits_i.size(-1)) out[i, keep, :] = logits_i[0, -keep.numel() :, :] return out def forward( self, input_ids=None, attention_mask=None, labels=None, use_cache=False, past_key_values=None, **kwargs, ): if input_ids.size(1) > self.config.seq_len: input_ids = input_ids[:, -self.config.seq_len :] if attention_mask is not None: attention_mask = attention_mask[:, -self.config.seq_len :] if labels is not None: labels = labels[:, -self.config.seq_len :] logits = self._masked_logits(input_ids, attention_mask) loss = None if labels is not None: shift_logits = logits[:, :-1, :].contiguous() shift_labels = labels[:, 1:].contiguous() loss = F.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)).float(), shift_labels.view(-1), ignore_index=-100, ) return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None) def prepare_inputs_for_generation(self, input_ids, **kwargs): attention_mask = kwargs.get("attention_mask") result = {"input_ids": input_ids[:, -self.config.seq_len :]} if attention_mask is not None: result["attention_mask"] = attention_mask[:, -self.config.seq_len :] return result